Retrievit: In-context Retrieval Capabilities of Transformers, State Space Models, and Hybrid Architectures
作者: Georgios Pantazopoulos, Malvina Nikandrou, Ioannis Konstas, Alessandro Suglia
分类: cs.AI
发布日期: 2026-03-03
💡 一句话要点
研究混合架构Transformer与SSM在上下文检索中的能力，探索其在数据效率和泛化性上的优势。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 上下文检索 Transformer 状态空间模型 混合架构 长序列建模
📋 核心要点
- Transformer在上下文检索中表现优异，但计算复杂度随序列长度呈平方增长，限制了其应用。
- 论文提出结合Transformer和状态空间模型(SSM)的混合架构，旨在兼顾Transformer的检索能力和SSM的效率。
- 实验表明，混合模型在数据效率和外推方面优于SSM，并在信息密集检索任务中与Transformer性能相当。
📝 摘要（中文）
Transformer在上下文检索方面表现出色，但序列长度的二次复杂度是其瓶颈。状态空间模型(SSM)提供高效的线性时间处理，但检索能力有限。本文研究了结合Transformer和SSM的混合架构是否能在两个合成的上下文检索任务中兼具两者的优点。第一个任务是n-gram检索，要求模型识别并重现输入序列中查询后的n-gram。第二个任务是位置检索，模型接收单个查询token，需要执行两次关联查找：首先定位序列中的对应元素，然后输出其位置索引。在受控实验条件下，我们评估了Transformer、SSM和混合架构在数据效率、长度泛化、对领域外训练样本的鲁棒性以及学习到的表征方面的性能。结果表明，对于信息密集的上下文检索，混合模型优于SSM，并且在数据效率和外推方面与Transformer相当或超过Transformer。然而，Transformer在位置检索任务中保持优势。通过表征分析，我们发现基于SSM的模型发展出局部感知嵌入，其中代表相邻位置的token在嵌入空间中成为邻居，形成可解释的结构。这种在Transformer中不存在的涌现特性解释了SSM和混合模型在不同检索任务中的优势和局限性。我们的发现为基于任务需求选择架构提供了原则性指导，并揭示了Transformer、SSM和混合模型在学习位置关联方面的根本差异。
🔬 方法详解
问题定义：Transformer在处理长序列时，计算复杂度呈平方级增长，成为瓶颈。状态空间模型(SSM)虽然计算效率高，但上下文检索能力较弱。论文旨在探索混合架构能否结合两者的优点，在上下文检索任务中实现高性能和高效率。现有方法的痛点在于无法同时兼顾长序列处理能力和准确的上下文检索能力。
核心思路：论文的核心思路是将Transformer和SSM结合，利用Transformer强大的上下文建模能力和SSM的高效序列处理能力。通过混合架构，期望能够克服Transformer在长序列上的计算瓶颈，并提升SSM的检索性能。这种设计旨在利用两种模型的互补优势，实现更优的整体性能。
技术框架：整体架构包含Transformer和SSM两个主要模块。具体流程是：首先，输入序列经过Transformer模块进行初步的上下文编码；然后，Transformer的输出被送入SSM模块进行高效的序列建模和信息压缩；最后，模型根据任务需求，利用Transformer和SSM的输出进行预测。论文设计了两种合成的上下文检索任务：n-gram检索和位置检索，用于评估不同架构的性能。
关键创新：论文的关键创新在于提出了一种混合架构，将Transformer和SSM结合，并证明了这种架构在特定上下文检索任务中具有优势。此外，论文还通过表征分析，揭示了SSM模型学习到的局部感知嵌入，这种嵌入方式能够有效捕捉相邻位置之间的关系，从而提升检索性能。这种对模型内部表征的深入分析是理解模型行为的关键。
关键设计：论文使用了标准的Transformer和SSM模块，并针对不同的检索任务进行了微调。在n-gram检索任务中，模型需要预测查询后的n-gram序列。在位置检索任务中，模型需要预测查询token在序列中的位置索引。损失函数采用交叉熵损失，用于衡量预测结果与真实标签之间的差异。具体的网络结构和参数设置在论文中有详细描述，但这里无法提供具体数值。
🖼️ 关键图片
📊 实验亮点
实验结果表明，混合模型在n-gram检索任务中，数据效率和外推能力优于SSM，并与Transformer相当或超过Transformer。通过表征分析发现，SSM模型学习到局部感知嵌入，相邻位置的token在嵌入空间中更接近。Transformer在位置检索任务中表现更优，表明不同架构擅长不同类型的检索任务。
🎯 应用场景
该研究成果可应用于需要高效处理长序列并进行上下文检索的领域，例如：信息检索、机器翻译、文本摘要、对话系统等。混合架构的设计思路为解决长序列建模问题提供了新的方向，有助于提升相关应用在实际场景中的性能和效率，并可能推动新型神经网络架构的进一步发展。
📄 摘要（原文）
Transformers excel at in-context retrieval but suffer from quadratic complexity with sequence length, while State Space Models (SSMs) offer efficient linear-time processing but have limited retrieval capabilities. We investigate whether hybrid architectures combining Transformers and SSMs can achieve the best of both worlds on two synthetic in-context retrieval tasks. The first task, n-gram retrieval, requires the model to identify and reproduce an n-gram that succeeds the query within the input sequence. The second task, position retrieval, presents the model with a single query token and requires it to perform a two-hop associative lookup: first locating the corresponding element in the sequence, and then outputting its positional index. Under controlled experimental conditions, we assess data efficiency, length generalization, robustness to out of domain training examples, and learned representations across Transformers, SSMs, and hybrid architectures. We find that hybrid models outperform SSMs and match or exceed Transformers in data efficiency and extrapolation for information-dense context retrieval. However, Transformers maintain superiority in position retrieval tasks. Through representation analysis, we discover that SSM-based models develop locality-aware embeddings where tokens representing adjacent positions become neighbors in embedding space, forming interpretable structures. This emergent property, absent in Transformers, explains both the strengths and limitations of SSMs and hybrids for different retrieval tasks. Our findings provide principled guidance for architecture selection based on task requirements and reveal fundamental differences in how Transformers and SSMs, and hybrid models learn positional associations.