arXiv:2504.18574v1 类型: cross 摘要: SSMs 提供了处理固定状态大小的长序列的有效方法，但难以应对如检索过去上下文这样的算法任务。在这项工作中，我们探讨了这种上下文内检索在基于Transformer和SSM的语言模型中的运作方式。我们发现，这两种架构都发展出了相同的基本Gather-and-Aggregate (G&A)机制。一个Gather头首先识别并提取上下文中的相关信息，然后由一个Aggregate头将这些信息整合进最终表示。在两种模型类型中，G&A机制集中在少数几个头中，即使是在需要基本检索能力的基准测试中，这些头也成为了关键的瓶颈。例如，禁用精简版Llama-3.1-8B中的单个Gather或Aggregate头会使其在MMLU中正确回答问题的能力大幅下降，准确率从66%降至25%。这一发现表明，上下文内检索可以掩盖某些任务的知识需求限度。尽管在检索功能保持完整的情况下，精简模型在MMLU方面的表现仍然很强，但在其他知识测试中却失败。G&A依赖性同样存在于GSM8K、BBH和对话任务中。鉴于G&A在性能中的重要性，我们展示了SSM中的检索挑战表现在它们实施G&A的方式上，导致注意力模式更加平滑，而不是有效的G&A所依赖的尖锐的标记转换。因此，虽然Transformer和SSM在实施上下文内检索方面存在差距，但这种差距仅限于少数几个头，而不是整模型。这一见解不仅提供了一种统一的解释，即Transformer和SSM之间性能差异的原因，同时也突显了结合两者的优点的方法。例如，在预训练混合同模型中，注意力组件自然承担了Aggregate头的角色。同样，在预训练纯SSM中，用基于注意力的变体替换单个G&A头能显著提高检索能力。