混合架构 LM 过度依赖注意力,辅助损失可改善记忆利用
Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs
训练混合架构模型的朋友看看这篇:加了限制注意力访问历史的辅助损失后,长上下文任务提升明显,方法可直接套用。
研究针对交替使用注意力层与循环层的混合语言模型,发现模型对注意力通路的依赖远高于循环状态,标准监督微调也改变不了这种失衡。论文提出在训练中加入辅助损失,限制注意力访问更早上下文,迫使模型通过循环状态保留和利用信息。该方法在问答和 agentic 任务上提升了多个循环-注意力 LM 的整体表现,长上下文与信息聚合类任务收益尤其明显,对混合多种记忆形式的注意力 LM 同样有效。
Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs
Recurrent-attention hybrid language models (LMs), which interleave attention and recurrent layers, are increasingly used to combine the efficiency of the recurrent layers with the strong performance of attention layers. Prior work suggests that attention and recurrent layers offer complementary pathways to use past information: attention supports precise memory recall from earlier tokens, while recurrent layers support consolidation of disparate information over long contexts. However, we observe that simply having access to both pathways does not mean that hybrid LMs are effectively using them. We find that they rely substantially more on attention than on the recurrent state. Standard supervised fine-tuning improves overall performance but does not improve how the two memory pathways are coordinated: the model becomes more reliant on information propagated by attention layers, while its use of information propagated by recurrent layers remains limited. To encourage better coordination between the two memory pathways, we add an auxiliary loss that limits attention's access to earlier context while the recurrent state propagates through the full sequence. This objective encourages the model to retain and use information through the recurrent pathway alongside attention. It improves overall performance, with particularly strong gains on tasks involving longer contexts or requiring information aggregation, consistent with the strengths of recurrent layers observed in analysis. Crucially, this imbalance and the benefit of our auxiliary loss generalize: they apply to multiple recurrent-attention LMs in question-answering and agentic tasks, as well as to attention-based LMs that combine different forms of memory. Together, our findings show that simply providing multiple memory pathways does not ensure their effective use, and that targeted supervision is needed to better coordinate them.