{"id":"169cb33f-54ac-4ece-bb8e-48631e33992e","arxiv_id":"2504.18574","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":0,"one_line_summary":"In both Transformers and SSMs, in-context retrieval is performed by a small set of Gather and Aggregate heads, and the Transformer-SSM gap can be reduced to how well these few heads implement this mechanism.","lead":"The paper shows that both Transformer and recurrent state-space language models retrieve information from context using the same two-step mechanism, concentrated in a few attention-like heads. It argues that the Transformer-SSM performance gap comes from these few heads, not from general language modeling ability.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The dramatic 66%→25% single-head MMLU cliff is demonstrated only in a heavily pruned model; the paper never shows the same head is necessary in the intact full model.","rationale":"The reader identified the KV-retrieval diagnostic's format dependence as the weakest assumption. That is a real concern, but I find an earlier and more central one: the quantitative 'single head is necessary' result, which anchors the abstract and the 'few heads' framing, is obtained only in the pruned minimal model. The full-model data in Table 2 show gradual degradation over tens of heads, not a cliff at one head. The paper never bridges this gap by showing that L16H22 or L17H24 alone are critical in the unpruned model. This omission matters because pruning can create artificial bottlenecks: by removing layers and zeroing other heads, the experiment may be selecting the last remaining pathway rather than the original circuit. If a single-head ablation in the full model does not reproduce the 66%→25% drop, then the strongest evidence for the central claim is conditional on the pruning procedure, and the paper should soften the claim accordingly. I still think the overall story is plausible and the verdict should remain conditional, but the single-head necessity claim needs this additional check before it can be taken as evidence about intact models.","tokens_in":160,"tokens_out":6564,"duration_ms":70894,"concrete_test":"Run the identical MMLU evaluation on the intact, unpruned Llama-3.1-8B-Instruct with only L16H22 ablated, and separately with only L17H24 ablated, using the same eval harness as Table 1. Also run the analogous single-head ablation in full Llamba-8B. If MMLU stays within, say, 10 points of the 66–68% baseline, the single-head necessity is an artifact of the pruning procedure and Section 3.4's cliff does not support the 'few heads' claim. If MMLU collapses to ~25% in the unpruned model, the pruning concern is resolved.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's headline evidence—that disabling a single Gather or Aggregate Head drops MMLU from 66% to 25%—is measured exclusively in the heavily pruned 'minimal model' of Section 3.3, where all layers above L17 are removed and all noncritical heads in the final two layers are zeroed (Table 1 and Table 4). In this residual network, almost all redundant pathways are already gone, so the surviving head's contribution may be an artifact of the pruning intervention rather than a pre-existing bottleneck of the intact model. The full-model ablations in Table 2 tell a different story: removing 10 heads from Llama-3.1-8B drops MMLU by only about 9%, and even 20 heads leave accuracy at 38%, not at random. The paper never reports the effect of ablating L16H22 or L17H24 alone in the unpruned full model. If those single-head ablations cause only a small drop in the full model, then the claim that the Transformer-SSM retrieval gap 'exists in just a few heads' is overstated; the dramatic single-head cliff would be created by pruning away the model's redundancy, not by the intrinsic architecture.","agreement_with_reader":"disagree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes a unified mechanistic account of in-context retrieval in Transformer- and SSM-based language models. It identifies a two-stage Gather-and-Aggregate (G&A) mechanism: a Gather Head compresses each answer segment into a summary token, and an Aggregate Head later combines these summaries to select the correct label. The authors argue that this mechanism is concentrated in a small number of heads in both architecture families, that SSMs implement it less sharply (smoother attention patterns, more redundant heads), and that the resulting retrieval bottleneck can explain much of the Transformer–SSM performance gap on tasks like MMLU, GSM8K, BBH, and dialogue comprehension. Evidence comes from layer/head pruning and zeroing in Llama-3.1-8B, Falcon-Mamba-7B, Llamba-8B, and Zamba2, a synthetic KV-retrieval diagnostic, masking interventions that restrict heads to the hypothesized pattern, and hybrid layer-replacement experiments on a distilled SSM. The central claim is that the retrieval gap 'exists in just a few heads' rather than reflecting a broad architectural deficiency, with implications for hybrid model design.","tokens_in":23974,"tokens_out":6825,"duration_ms":60397,"significance":"If the central claim holds, the paper makes a valuable contribution: it extends the 'retrieval heads' literature from Transformers to SSMs, refines the single-head view by emphasizing a coordinated two-stage mechanism, and offers a concrete design rule for hybrid models (place or retain attention near the layers where G&A heads emerge). The paper's strengths include multiple independent interventions (layer pruning, head zeroing, runtime masking, hybrid replacement), a synthetic diagnostic intended to isolate retrieval, and honest caveats in Section 3.6 about cases where the two layers are non-adjacent. The findings are falsifiable and the experiments are publicly released. However, as detailed below, the strongest quantitative claim (a single-head MMLU cliff) is demonstrated only in a heavily pruned model, and the cross-task generalization from the synthetic diagnostic to natural benchmarks is not fully validated. These issues are load-bearing for the paper's headline conclusion, so the manuscript needs substantial revision rather than acceptance in its current form.","major_comments":[{"comment":"","section":"§3.4, Table 1 and §6.2, Table 2"},{"comment":"","section":"§6.1, Appendix A.5, Table 5"},{"comment":"","section":"§5.4, Hybrid Replacements"}],"minor_comments":[{"comment":"","section":"Abstract and §1"},{"comment":"","section":"§2"},{"comment":"","section":"Appendix A.4"},{"comment":"","section":"§3.6"},{"comment":"","section":"Throughout"}],"recommendation":"major_revision","confidential_remarks":"The paper is a serious mechanistic-interpretability study with a strong empirical component, but its headline claim is more sweeping than the evidence supports. The two major issues are (1) the single-head MMLU cliff is only shown in a heavily pruned model, and (2) the synthetic KV-retrieval diagnostic is not shown to select the same heads that drive natural-task behavior. Both are fixable with additional experiments or by qualifying the claims. I recommend major revision rather than rejection because the core G&A framework is well-motivated and the interventions are creative. The authors should also consider whether the abstract's 'just a few heads' framing overstates the results, as Table 2 indicates that even in Transformers, meaningful degradation requires ablating tens of heads in the full model."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Worth your time. The paper's core finding—that both Transformers and SSM-based language models implement in-context retrieval through a coordinated Gather-and-Aggregate mechanism concentrated in a few heads—is genuinely new and surprisingly well supported. It extends the known content-gatherer/correct-letter heads from Transformer-only work to SSMs, and it does this with multiple independent interventions: layer pruning, head zeroing, masking heads to segment-restricted patterns, and layer replacement in a distilled SSM. The synthetic KV-retrieval task provides external validation of the head selection, and the code is public. That is real evidence, and it earns the central claim a fair hearing.\n\nThe soft spots are real but not fatal. The stress-test note is correct: the dramatic 66%→25% single-head MMLU cliff is demonstrated only in the heavily pruned minimal model, where redundancy has already been stripped away. In the intact full model (Table 2), removing 10 heads from Llama-3.1-8B drops MMLU by about nine points, not to random. The paper never reports a single-head ablation in the unpruned model, so the claim that the gap \"exists in just a few heads\" is overstated if it is meant to apply to unmodified models. That said, the KV-retrieval experiments do use the full model and show clear single-head sensitivity (e.g., 85.6% to 76% or 62% depending on which head), so the localization claim does not collapse—it is just less dramatic than the abstract implies. The paper should also acknowledge that the KV diagnostic was deliberately tuned (no trailing space) to maximize head-level signal, which raises a fair question about how cleanly it transfers to natural benchmarks. And the framing of MMLU as primarily a retrieval benchmark is too strong; the data show retrieval is a critical component, but knowledge still matters in the intact model.\n\nMethodologically, the head selection is post hoc and there are no error bars, which is standard for this kind of mechanistic interpretability work but still worth a caveat. The citation pattern including self-citations to Llamba is fine because those models are used and the results are reproducible from the released code. The paper even includes a caveats section noting that the two-layer mechanism sometimes spans non-adjacent layers, which is honest.\n\nWho is this for? Anyone working on SSM/hybrid architecture, mechanistic interpretability of attention heads, or benchmark design. It deserves a serious referee—the interventions are well chosen and the core finding is likely to hold up with revisions. I would send it to review and ask the authors to (a) report single-head ablations in the unpruned full model, (b) move the pruned-model cliff to a clearly labeled supporting result rather than the headline, and (c) temper the MMLU-as-retrieval language. With those fixes, it becomes a solid contribution.","headline":"A credible and well-triangulated mechanistic story that localizes the Transformer–SSM retrieval gap to a few Gather-and-Aggregate heads, but the headline single-head cliff is an artifact of heavy pruning and the broad claim needs reining in.","tokens_in":24528,"tokens_out":1753,"would_cite":true,"duration_ms":19092,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"The Transformer-SSM retrieval gap comes down to a few specialized Gather-and-Aggregate heads, not the whole model.","keywords":["in-context retrieval","state-space models","attention heads","Gather-and-Aggregate mechanism","mechanistic interpretability","MMLU","hybrid language models","Mamba"],"falsifier":"A direct falsifier would be to use activation patching on a natural retrieval-heavy benchmark to identify the heads that carry the correct-answer information and check whether they coincide with the G&A heads found by the synthetic diagnostic. If patching points to different heads, or if removing the identified G&A heads leaves some model's MMLU unaffected while removing a different head collapses it, the localization claim would be refuted. A simpler check is to run the same head-ablation ranking on the KV-retrieval task with a trailing space in the prompt and see whether the top-ranked heads change.","tokens_in":23553,"feed_emoji":"🧠","tokens_out":6138,"duration_ms":55383,"temperature":0.7,"pith_summary":"The paper argues that the performance gap between Transformer and state-space-model (SSM) language models on retrieval-heavy tasks is not a broad architectural deficit but a localized one: in both architectures, in-context retrieval is carried out by a small number of specialized heads that implement a two-stage Gather-and-Aggregate mechanism. A Gather Head condenses relevant segments of context into single summary tokens, and an Aggregate Head later combines those summaries to answer the query. The authors show that removing a single such head from a layer-pruned Llama-3.1-8B collapses MMLU accuracy from 66% to 25% (random guessing), while knowledge benchmarks hardly move. In SSMs the same mechanism exists but is implemented more diffusely and with smoother attention patterns, so more heads are needed to approach Transformer-level retrieval. This reframes MMLU as largely a retrieval benchmark and gives a concrete target for hybrid designs: attention layers succeed precisely because they take over the Aggregate Head role.","feed_headline":"Retrieval lives in a few heads in Transformers and SSMs","feed_subtitle":"Disabling one Gather-Aggregate head drops MMLU from 66% to random guessing; the gap is localized.","key_machinery":"The central object is the Gather-and-Aggregate (G&A) mechanism, a two-stage circuit for in-context retrieval. The first stage, a Gather Head, attends over a segment of input and compresses it into a summary token, typically the segment-final newline; the second stage, an Aggregate Head, attends from the query position to those summary tokens and effectively selects the correct one. The paper isolates this pair with layer pruning, head ablations, attention masking, and layer-replacement experiments, and shows that the Transformer-SSM gap in retrieval-heavy benchmarks tracks how sharply each architecture's G&A heads can implement these two operations.","core_discovery":"Both Transformer- and SSM-based language models retrieve information from context through the same two-stage Gather-and-Aggregate mechanism, and in both families the mechanism is concentrated in a very small number of heads. A Gather Head condenses each relevant segment of context into a single representative token; an Aggregate Head then attends to those summaries from the query position and selects the relevant one. In a minimal Llama-3.1-8B, two heads, L16H22 and L17H24, are jointly responsible for MMLU letter selection: removing either drops accuracy from 66% to about 25%, while retaining just these two heads in the final layers sustains performance. Falcon-Mamba-7B and Llamba-8B show the same pattern, with a few Mamba channels or heads recovering most of the full model's MMLU score. The paper further shows that SSM heads implement the mechanism with smoother, less sharply localized attention, so they need more heads to approach Transformer-level retrieval, and that replacing one SSM Aggregate Head layer with attention raises MMLU from 33% to 50%.","pith_inferences":["If this diagnosis is right, then adding parameters or training data to pure SSMs will not by itself fix retrieval; the cheaper lever is architectural, inserting a small number of attention slots at the layers where G&A heads emerge.","The diagnostic format matters: the paper deliberately uses answer scoring without a trailing space because that format concentrates retrieval in a few heads. A natural extension is to test whether the same head identities are recovered under more natural token-boundary formats or longer contexts, since the head ranking could shift with prompt format.","Because MMLU's letter-selection format is what makes retrieval decisive, leaderboard comparisons between architectures may partly be comparisons of retrieval-format handling rather than knowledge; reformatting benchmarks to free-form answers could change relative rankings.","Another testable extension is to run the same head-ablation protocol across models of different scale, pretraining objective, or depth to see whether G&A heads always settle in middle layers or whether their preferred depth shifts with architecture and data."],"forward_implications":["MMLU scores on these models should be read partly as retrieval ability: a model can score 66% on MMLU while its knowledge-task scores are already degraded, and can fall to 25% when one head is removed.","Hybrid models work because attention takes over aggregation: ablating six attention heads in Zamba2-7B drops MMLU from 64.3% to 34.9% while knowledge-task accuracy remains at 70%.","For hybrid design, attention is most valuable in middle layers where G&A heads emerge; when distilling a Transformer into an SSM, preserving attention at layers containing Aggregate Heads recovers most of the lost MMLU performance.","SSMs can compensate for weak per-head aggregation by using more heads, but this redundancy only partially closes the gap; replacing a single Aggregate Head layer with attention gave the largest measured gain.","The localized-retrieval account extends beyond MMLU to GSM8K, BBH, SWDE, and chat-formatted ARC, so the bottleneck is not specific to multiple-choice letter selection."],"supporting_citations":[{"why":"Identifies Content Gatherer and Correct Letter heads in Chinchilla-70B; this paper extends those into the Gather-and-Aggregate mechanism.","marker":"[Lieberum et al., 2023]"},{"why":"Introduces retrieval heads and shows few heads are critical; this paper refines that into coordinated Gather and Aggregate roles.","marker":"[Wu et al., 2024]"},{"why":"Documents that Mamba-based models need substantially longer training to perform well on MMLU, setting up the gap being explained.","marker":"[Waleffe et al., 2024]"},{"why":"Provides theoretical and empirical evidence that SSMs struggle with precise copying, a baseline limitation this paper localizes to aggregation heads.","marker":"[Jelassi et al., 2024]"},{"why":"Shows RNNs' key bottleneck is in-context retrieval, which this paper attributes to a small set of G&A heads.","marker":"[Wen et al., 2024]"},{"why":"Shows associative recall accounts for much of the attention-versus-gated-convolution performance gap, motivating recall as the core capability.","marker":"[Arora et al., 2023]"},{"why":"Supplies the Zamba hybrid models used to test the claim that hybrid architectures delegate aggregation to attention layers.","marker":"[Glorioso et al., 2024]"},{"why":"Supplies the Llamba distilled SSM models used in the layer-replacement experiments that isolate the Aggregate Head.","marker":"[Bick et al., 2025]"},{"why":"Defines the MMLU benchmark whose retrieval-versus-knowledge decomposition is the paper's central case study.","marker":"[Hendrycks et al., 2021]"}],"fun_headline_variants":["One head makes or breaks retrieval in Transformers and SSMs","Disabling a single head drops MMLU from 66% to guessing","The Transformer-SSM gap lives in a few heads, not the whole model","SSMs need more heads for retrieval; one attention head fixes it","Gather-and-Aggregate: the few-head bottleneck behind AI retrieval"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The cross-task conclusion rests on the assumption that the synthetic KV-retrieval diagnostic, run with answer scoring and deliberately without a trailing space, activates the same heads and the same retrieval process that drive performance on natural benchmarks such as MMLU, GSM8K, BBH, and dialogue comprehension.","fun_headline_variants_meta":{"raw":{"variants":["One head makes or breaks retrieval in Transformers and SSMs","Disabling a single head drops MMLU from 66% to guessing","The Transformer-SSM gap lives in a few heads, not the whole model","SSMs need more heads for retrieval; one attention head fixes it","Gather-and-Aggregate: the few-head bottleneck behind AI retrieval"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000261,"raw_usage":{"total_tokens":1670,"prompt_tokens":1102,"completion_tokens":568,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":718,"completion_tokens_details":{"reasoning_tokens":473}},"tokens_in":718,"tokens_out":568,"duration_ms":5413,"temperature":1.0,"reasoning_tokens":473,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T11:12:28.957205+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"A direct falsifier would be to use activation patching on a natural retrieval-heavy benchmark to identify the heads that carry the correct-answer information and check whether they coincide with the G&A heads found by the synthetic diagnostic. If patching points to different heads, or if removing the identified G&A heads leaves some model's MMLU unaffected while removing a different head collapses it, the localization claim would be refuted. A simpler check is to run the same head-ablation ranking on the KV-retrieval task with a trailing space in the prompt and see whether the top-ranked heads change.","supporting_citations":[],"review_version":1}