Some Attention is All You Need for Retrieval

TL;DR

This work shows that in hybrid SSM-Transformer models, retrieval is exclusively handled by self-attention, with SSMs contributing no retrieval capability. Using entropy-based attention sparsification and Just Read Twice prompting across RG-2B, RG-9B, and Jamba, the authors demonstrate that removing most attention heads catastrophically kills retrieval, while keeping as few as ~15% of heads preserves retrieval and much of MMLU performance. They further reveal that retrieval requires precise mechanistic conditions, including needle-token exposure during generation and sufficient context during prefill or generation. These results argue for strict specialization within hybrid architectures, offering actionable implications for architecture optimization and interpretability, such as targeted sparsification and retrieval-focused modules. The study thus provides a concrete, mechanistic view of how hybrid models organize computation, moving beyond assumptions of redundancy toward task-specific computational delegation.

Abstract

We demonstrate complete functional segregation in hybrid SSM-Transformer architectures: retrieval depends exclusively on self-attention layers. Across RecurrentGemma-2B/9B and Jamba-Mini-1.6, attention ablation causes catastrophic retrieval failure (0% accuracy), while SSM layers show no compensatory mechanisms even with improved prompting. Conversely, sparsifying attention to just 15% of heads maintains near-perfect retrieval while preserving 84% MMLU performance, suggesting self-attention specializes primarily for retrieval tasks. We identify precise mechanistic requirements for retrieval: needle tokens must be exposed during generation and sufficient context must be available during prefill or generation. This strict functional specialization challenges assumptions about redundancy in hybrid architectures and suggests these models operate as specialized modules rather than integrated systems, with immediate implications for architecture optimization and interpretability.

Figures (10)

• Figure 1: Simplified representation of manipulation methods. Original attention weights (dotted blue), modified weights (solid blue), and needle tokens (yellow highlight).
• Figure 2: Accuracy as a function of $k$ in top-$k$ sparsification on standard NIAH for generation and prefill versions. Accuracy was approximated by the average score across all 100 prompts, relative to the maximum score (5). Note that scoring was neither linear nor continuous (see Table \ref{['tab:scoring']} for interpretation guidance). Read from right to left for increasing sparsity.
• Figure 3: Left: Accuracy as a function of $k$ in top-$k$ sparsification on NIAH with and without JRT applied, for all models in the generation version. Accuracy was approximated by the average score across all 100 prompts, relative to the maximum score (5). Note that scoring was neither linear nor continuous (see Table \ref{['tab:scoring']} for interpretation guidance). Read from right to left for increasing sparsity. Right: Scores for Jamba-Mini-1.6, in an ablated configuration ($k=0$), a sparse configuration ($k=5$), and the base configuration, on GLUE and MMLU.
• Figure 4: Retrieval maps for RG-2B on NIAH with the manipulations Keep-Omit, Omit-Keep, and Keep-Keep applied. The x-axis shows the used prompt length, the y-axis the depth of the needle, 0% being the very end of the prompt and 100% being the very beginning of the prompt.
• Figure 5: Retrieval maps for RG-2B, RG-9B and Jamba at $k=0$ on NIAH with and without JRT applied.