# Retrieval-Aware Distillation for Transformer-SSM Hybrids

arXiv: https://arxiv.org/abs/2602.11374 · HTML: https://mlsanity.com/paper/2602.11374/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 27 references that are in the map.
- First posted on arXiv 11 Feb 2026.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Repeat After Me: Transformers are Better than State Space Models at Copying](https://mlsanity.com/paper/2402.01032/) (2024)
- [Simple linear attention language models balance the recall-throughput tradeoff](https://mlsanity.com/paper/2402.18668/) (2024)
- [Zoology: Measuring and Improving Recall in Efficient Language Models](https://mlsanity.com/paper/2312.04927/) (2023)
- [Zamba: A Compact 7B SSM Hybrid Model](https://mlsanity.com/paper/2405.16712/) (2024)
- [The Mamba in the Llama: Distilling and Accelerating Hybrid Models](https://mlsanity.com/paper/2408.15237/) (2024)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks](https://mlsanity.com/paper/2402.04248/) (2024)
- [Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models](https://mlsanity.com/paper/2408.10189/) (2024)
- [RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval](https://mlsanity.com/paper/2402.18510/) (2024)
- [Llamba: Scaling Distilled Recurrent Models for Efficient Language Processing](https://mlsanity.com/paper/2502.14458/) (2025)
- [Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners](https://mlsanity.com/paper/2502.20339/) (2025)
- [M1: Towards Scalable Test-Time Compute with Mamba Reasoning Models](https://mlsanity.com/paper/2504.10449/) (2025)
- [Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism](https://mlsanity.com/paper/2504.18574/) (2025)

## Built on by (most cited first, 6 of 6)

- [Raven: High-Recall Sequence Modeling with Sparse Memory Routing](https://mlsanity.com/paper/2607.25357/) (2026)
- [Morphing into Hybrid Attention Models](https://mlsanity.com/paper/2606.30562/) (2026)
- [Taylor-Calibrate: Principled Initialization for Hybrid Linear Attention Distillation](https://mlsanity.com/paper/2606.16429/) (2026)
- [Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling](https://mlsanity.com/paper/2604.24715/) (2026)
- [On-Policy Attention Linearization](https://mlsanity.com/paper/2609.31947/) (2026)
- [Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs](https://mlsanity.com/paper/2610.06750/) (2026)

## Abstract (arXiv)

State-space models (SSMs) offer efficient sequence modeling but lag behind Transformers on benchmarks that require in-context retrieval. Prior work links this gap to a small set of attention heads, termed Gather-and-Aggregate (G&A), which SSMs struggle to reproduce. We propose *retrieval-aware distillation*, which converts a pretrained Transformer into a hybrid student by preserving only these retrieval-critical heads and distilling the rest into recurrent heads. We identify the essential heads via ablation on a synthetic retrieval task, producing a hybrid with sparse, non-uniform attention placement. We show that preserving **just 2% of attention heads recovers over 95% of teacher performance on retrieval-heavy tasks** (10 heads in a 1B model), requiring far fewer heads than hybrids that retain at least 25%. We further find that large recurrent states often compensate for missing retrieval: once retrieval is handled by these heads, the SSM backbone can be simplified with limited loss, even with an 8× reduction in state dimension. By reducing both the attention cache and the SSM state, the resulting hybrid is 5--6× more memory-efficient than comparable hybrids, closing the Transformer--SSM gap at a fraction of the memory cost.