Retrievit: In-context Retrieval Capabilities of Transformers, State Space Models, and Hybrid Architectures
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 reproduce an n-gram that succeeds the query within the input sequence. The second task, position retrieval, presents the model with a query token and requires it to perform a two-hop lookup: first locating the corresponding element in the sequence, and then outputting its positional index. Under controlled 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 terms of data efficiency and extrapolation for tasks that require precise information retrieval from the input context. 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 property is absent in Transformers as causal attention is sufficient for acquiring positional associations, and the introduction of positional encoding amplifies this behavior, leading to improvement in data efficiency. SSMs on the other hand update their internal representations incrementally and without positional encodings, are required to learn these associations. Our findings reveal fundamental differences in how Transformers and SSMs, and hybrid models learn positional associations.
Code (0)
등록된 구현이 없습니다.
Tasks
Information RetrievalSimilar Papers 제목 키워드 기반
Birdie: Advancing State Space Models with Reward-Driven Objectives and Curricula
Efficient state space models (SSMs), such as linear recurrent neural networks and linear attention variants, offer computational advantages over Transformers but struggle with tasks requiring long-range in-context retrie…
Computational EfficiencyQuestion AnsweringRetrievalState Space ModelsDifferential Mamba
Sequence models like Transformers and RNNs often overallocate attention to irrelevant context, leading to noisy intermediate representations. This degrades LLM capabilities by promoting hallucinations, weakening long-ran…
Language ModelingLanguage ModellingMambaRetrievalContext-Gated Associative Retrieval: From Theory to Transformers
Hopfield networks and their generalizations have established deep connections among biological associative memories, statistical physics, and transformers. Yet most models treat retrieval as a fixed query-to-memory mappi…
Shaking Up VLMs: Comparing Transformers and Structured State Space Models for Vision & Language Modeling
This study explores replacing Transformers in Visual Language Models (VLMs) with Mamba, a recent structured state space model (SSM) that demonstrates promising performance in sequence modeling. We test models up to 3B pa…
Language ModelingLanguage ModellingMambaQuestion Answering+4Differential Multimodal Transformers
Small language models have gained significant popularity due to their efficiency and growing capabilities. However, incorporating additional modalities, such as vision, can exacerbate the challenge of limited context win…
Information Retrieval