Retrievit: In-context Retrieval Capabilities of Transformers, State Space Models, and Hybrid Architectures

TL;DR

This paper compares Transformers, State Space Models, and hybrid architectures for in-context retrieval tasks, revealing their strengths, limitations, and emergent properties through controlled experiments and analysis.

Contribution

It introduces a comprehensive evaluation of hybrid models combining Transformers and SSMs for retrieval, highlighting their performance and representation differences.

Findings

Hybrid models outperform SSMs and match/exceed Transformers in data efficiency.

Transformers excel in position retrieval tasks.

SSMs develop locality-aware embeddings, forming interpretable structures.

Abstract

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.