Retrieval-Aware Distillation for Transformer-SSM Hybrids

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\times$ reduction in state dimension. By reducing both the attention cache and the SSM state, the resulting hybrid is $5$–$6\times$ more memory-efficient than comparable hybrids, closing the Transformer–SSM gap at a fraction of the memory cost.

💡 Research Summary

The paper addresses the well‑known gap between Transformer language models and state‑space models (SSMs) on tasks that require in‑context retrieval. Recent work has identified that this gap is not a global deficiency of SSMs but is caused by a small subset of attention heads in Transformers, termed Gather‑and‑Aggregate (G&A) heads, which are responsible for the precise “gather” of key‑value pairs and the subsequent “aggregate” of relevant information across long contexts. Existing hybrid architectures that combine SSMs with attention typically allocate attention heads in a fixed, layer‑wise pattern (e.g., a constant ratio of attention to SSM layers). Such designs often retain far more heads than are truly needed for retrieval, leading to unnecessary memory overhead.

The authors propose retrieval‑aware distillation, a two‑stage method that converts a pretrained Transformer into an efficient hybrid student by (1) identifying the retrieval‑critical heads through ablation on a synthetic KV‑retrieval probe, and (2) preserving only those heads while replacing all others with recurrent SSM‑based heads. The KV‑retrieval probe measures the drop in accuracy when each head is removed; heads that cause the largest drop are deemed essential G&A heads. This ranking yields a highly non‑uniform attention layout: in a 1 B‑parameter model with 512 heads, the top 10 heads (≈2 % of the total) capture virtually all retrieval capability.

Once the critical heads are selected, the remaining attention slots are filled with SSM mixers based on the Discrete‑Mamba‑2 (DM2) variant. To keep the hybrid’s token‑mixing statistics stable, the concatenated outputs of the retained heads are normalized to match the mean and variance of the SSM output using a parameter‑free LayerNorm step. The hybrid is then trained with the MOHA‑WK distillation framework, which aligns teacher and student token‑mixing matrices (matrix orientation loss), hidden‑state representations (L2 alignment), and finally performs logit‑level knowledge distillation.

Experiments are conducted on two large‑scale models: Llama‑3.2‑1B and Qwen2.5‑1.5B. For each model, the authors evaluate a spectrum of head budgets (0, 5, 10, 20, 30, 40, 512 heads) on two groups of benchmarks: (i) knowledge‑focused tasks (PIQA, Winogrande, OpenBookQA, etc.) that require factual reasoning but little retrieval, and (ii) retrieval‑heavy tasks (Lambada, MMLU, GSM8K, SWDE, and the synthetic KV‑retrieval). Coverage is defined as the percentage of the teacher’s average score retained by the student.

Key findings include:

1. Retrieval performance recovery with minimal heads – Retaining only 10 heads restores over 95 % of the teacher’s performance on retrieval‑heavy tasks (e.g., KV‑Retrieval accuracy jumps from 13 % to 99 %). The same 10‑head hybrid also matches the teacher on most knowledge‑focused tasks, showing that the remaining heads contribute little beyond retrieval.

2. Efficiency gains over prior hybrids – Earlier hybrid distillation methods required at least 25 % of the attention heads to achieve comparable performance. The proposed method reduces the attention‑head budget by more than an order of magnitude.

3. State‑dimension reduction – With retrieval handled by the preserved heads, the SSM backbone can be dramatically simplified. The authors shrink the SSM state dimension from 64 to 8 (an 8× reduction) with only modest performance loss, indicating that large SSM states in pure SSMs primarily compensate for missing retrieval capability.

4. Memory and throughput improvements – By cutting both the KV‑cache (fewer attention heads) and the SSM state, the hybrid achieves 5× memory savings on short sequences (128 tokens) and up to 6× on long sequences (4 K tokens). Corresponding gains are observed in pre‑fill latency and decoding throughput.

5. Qualitative analysis – Attention‑map visualizations confirm that the retained heads exhibit strong G&A patterns, while the SSM heads provide smooth recurrent mixing. The non‑uniform placement of attention heads avoids the redundancy inherent in fixed‑ratio hybrids.

In summary, the paper introduces a principled, data‑driven approach to hybrid model design: by explicitly preserving only the attention heads that are indispensable for in‑context retrieval, and distilling the rest into efficient recurrent SSM components, it delivers a hybrid that is both memory‑lean and performance‑robust. This work opens a practical pathway for deploying large language models in resource‑constrained settings where long‑range retrieval is essential, and suggests future extensions to multimodal or even larger-scale models where similar head‑selection strategies could yield further efficiency gains.