{"id":"8d424e27-190c-4822-adcd-4ecb95278a48","arxiv_id":"2411.17685","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":4.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"Attamba compresses chunks of tokens with state-space models and runs attention over the compressed key-value states, improving WikiText2 perplexity by up to 24% over a reduced-footprint transformer baseline at small scale.","lead":"Attamba replaces the key and value projections in a transformer with state-space models (SSMs) that summarize small groups of tokens, so attention runs over compressed summaries instead of every token. The paper reports perplexity and efficiency trade-offs on small language models, claiming better quality for similar cache sizes and large cache savings for a small quality loss.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The efficiency-matched transformer baselines are not actually footprint-matched: Figure 7's sliding-window baseline uses window L/P and dimension E/P, giving far smaller attention maps and FLOPs than Attamba, so the 24% perplexity claim lacks a matched control.","rationale":"The reader identified the same weakest assumption: the baseline construction in Appendix A (Eqs. 9-11) and Figure 7 does not genuinely match KV-cache and attention footprint. My analysis confirms this is load-bearing. The iso-KV equation ignores the leading-token cache, and the attention-map matching uses a sliding window of L/P, which is drastically smaller than Attamba's actual key-value length of n/P + L_lead. For the 8B-token comparison in Figure 6, the '+SWA' transformer has E=128 and a window of 8, giving it 36× smaller attention map than Attamba. Therefore the reported 24% improvement may simply reflect that the baseline sees only 8 recent tokens while Attamba sees ~288. This is not a matched comparison. A corrected baseline with equal cache and attention map could drastically reduce or eliminate the claimed advantage. The pseudo-chunking result (Figure 16) is a cleaner iso-parameter comparison and suggests a modest real benefit, but it does not support the headline efficiency-quality trade-off. The test-time mask in Eq. 7 also deserves verification: under 0-indexing it would include the current chunk boundary before it is complete, but this appears to be a formulation issue that the implementation may handle correctly. The baseline mismatch is the primary concern, and the proposed recomputation would settle it.","tokens_in":16785,"tokens_out":13257,"duration_ms":109583,"concrete_test":"Recompute the iso-efficiency baselines in Fig. 7 and Sec. 5.2 using explicit footprint formulas. For each Attamba config (n=1024, P, L_lead, E=512), compute its actual per-layer KV-cache = 2B(n/P + L_lead)E and attention-map size = n(n/P + L_lead). Set the transformer baseline's attention dimension to F = E(n/P + L_lead)/n so its KV-cache matches, and set its sliding-window length to W = n/P + L_lead so its attention map matches. Train/evaluate these corrected baselines under the same 1B-token budget and compare final WikiText2 perplexity. If the gap between Attamba and this corrected transformer is below ~10% relative, the '24% with similar footprint' headline is unsupported. Also run at least 3 seeds for the main comparison to confirm the gap is not noise.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The headline '24% improved perplexity with similar KV-cache and attention footprint' rests on baselines that do not match Attamba's actual footprints. In the iso-KV construction (Appendix A, Eq. 9), the transformer's attention dimension is set to F = E/P to match 2BLF = 2BLE/P + 2BDS. This ignores the leading-token cache that Attamba stores: for typical settings (L_seq=1024, P=4, L_lead=32, E=512), Attamba's cache is 2*(1024/4 + 32)*512 entries per layer, while the transformer's is 2*1024*128 entries—already not equal. More importantly, the attention map is not matched at all. Attamba's query attends to n/P + L_lead keys (≈288 for the example), producing an L×288 map. The transformer baseline in Figure 7 uses a sliding window of size L_lead/P = 8, producing an L×8 map—36× smaller. Because the attention-dimension F is also reduced by 4×, its attention FLOPs are even smaller. Thus the '+KVC+SWA' baseline has a far smaller attention footprint than Attamba, not a similar one. The '24%' gap therefore likely reflects the baseline's severe attention restriction (only 8 recent tokens) rather than the benefit of SSM compression. To support the claim, the authors must construct a transformer baseline with (i) F such that 2BLF equals Attamba's complete cache including leading tokens, and (ii) sliding window W = n/P + L_lead. Under such a matched control, the perplexity advantage may shrink or disappear.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes Attamba, a transformer variant in which the key and value projections are replaced by SSM blocks that compress chunks of P tokens into single states, followed by attention over these compressed states plus a set of leading (uncompressed) tokens. The authors report WikiText2 perplexity for roughly 60M-parameter models trained on 1B or 8B tokens and claim a 24% perplexity improvement over a transformer with similar KV-cache and attention footprint, as well as a roughly 4x reduction in KV-cache and attention FLOPs for a 5% perplexity trade-off. The manuscript also studies chunk-boundary strategies, SSM state dimension, leading-token count, and pseudo-chunking, with code and training logs provided.","tokens_in":17219,"tokens_out":5877,"duration_ms":57209,"significance":"If the architecture is causally sound and the efficiency-matched baselines are genuine, Attamba is a useful hybrid that combines SSM-based local compression with global attention over compressed states, offering a controllable quality/efficiency knob and a path from quadratic to linear attention. The paper includes several strengths: the efficiency arithmetic is straightforward, the design choices are ablated, and the authors provide logs and code, which supports reproducibility. However, the central efficiency claim currently rests on transformer baselines whose attention footprint is much smaller than Attamba's, and the test-time masking rule appears to permit attention to future tokens within the current chunk; both issues must be resolved before the headline claims can be accepted.","major_comments":[{"comment":"The 'similar footprint' transformer baseline is not footprint-matched. The +KVC+SWA baseline uses attention dimension F = E/P = 128 and a sliding window of size L_lead/P = 8, whereas Attamba with P=4 and L_lead=32 attends to roughly n/P + L_lead = 288 key-value states of dimension E = 512. For n = 1024 the baseline attention map is therefore about 36x smaller, and its attention FLOPs are even smaller because the dimension is also reduced. The reported 24% perplexity gap (21.74 vs. 28.66 in Figure 6/20) may thus reflect the baseline's severe restriction to eight recent tokens rather than a benefit of SSM compression. Please construct a transformer baseline that matches Attamba's full cache, including leading tokens, and its attention map size, e.g., F satisfying 2BLF = 2B(n/P + L_lead)E and window W = n/P + L_lead, or clearly report the residual mismatch.","section":"§5.1, Figure 7, Appendix A Eqs. (9)-(11)"},{"comment":"The test-time mask Mtest appears to violate causality. For a query at position i inside chunk p = floor(i/P), the condition j <= floor(i/P) includes the compressed state K(p)[-1], which is the SSM output after processing the entire chunk, including tokens after i. At training time Eq. (5) avoids this by allowing only partial-chunk outputs for the current chunk, but at test time those partial outputs are discarded. If Attamba is evaluated with the stated mask and cached chunk-boundary states, the model can attend to future tokens within the current chunk. Please either restrict test-time attention to completed chunks only (with the current chunk handled by leading tokens), or maintain causally valid partial states; the current formulation needs clarification or correction.","section":"§4.1, Eq. (7)"},{"comment":"The training budget is stated inconsistently. Section 5 says 'Unless otherwise specified, we train on approximately 1B tokens', Appendix B repeats the 1B budget, but Section 5.2 and Figures 6/20 report training 'for 100,000 steps over 8 billion tokens'. Since the headline comparison with Mamba, minGRU, Hawk, and the Transformer baselines comes from the 8B run, the paper needs to state which budget applies to each figure and why the appendix describes a different setting; otherwise the results are not reproducible.","section":"§5, §5.2, Appendix B"}],"minor_comments":[{"comment":"The symbol L is used both for sequence length (Appendix A, Eqs. (9)-(11)) and for the number of leading tokens (main text, Figure 5, Figure 7 caption). This ambiguity makes the baseline construction in Eq. (9) and the caption 'window size = L/P' difficult to interpret; please use distinct symbols, e.g., N_seq and L_lead.","section":"Notation throughout"},{"comment":"There are typographical errors: 'intorduce' should be 'introduce', 'Psuedo-chunking' should be 'Pseudo-chunking', and 'inStuffed Mamba' in Related Work is missing a space.","section":"Appendix A.1.1, A.2"},{"comment":"The 'Attamba-Linear' variant and the claimed smooth transition between quadratic and linear scaling are motivated but not evaluated; a small experiment or explicit statement that this is future work would make the claim more precise.","section":"§4.1, Figure 13"},{"comment":"The limitations section is candid about the small-scale, local-task evaluation, and Figure 17 correctly shows that a true 128x KV-cache reduction with only one leading token performs poorly. These statements appropriately qualify the generality of the results.","section":"§6, Figure 17"}],"recommendation":"major_revision","confidential_remarks":"The architecture and ablation study are interesting and the code release is a plus. My main concern is that the headline efficiency claim is not yet supported by a genuinely matched baseline, and the test-time masking rule seems to have a causal leak that needs to be resolved. These are fixable within the scope of a revision: rerun or re-report the matched baseline and clarify or correct the inference-time mask. If the causal leak is confirmed rather than a notation issue, the experimental numbers would need to be re-examined before publication."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The core idea is worth a look: replace the key and value projections in a transformer with SSMs that compress chunks of tokens, then run attention over the compressed states. That combination is genuinely new relative to the cited related work, and the paper ships code and logs. The Limitations section is also unusually honest, conceding the 60M scale, the single benchmark, and the local nature of WikiText2.\n\nWhat the paper does well: the pseudo-chunking result (Figure 16) is a clean empirical finding—replacing KV projections with SSMs while keeping full attention slightly improves perplexity at iso-parameter count. Cyclic chunking is a sensible fix for boundary bias. The efficiency math in the appendix is internally consistent, and the authors correctly note that iso-parameter comparisons are inappropriate for KV-cache or attention footprint.\n\nThe soft spots are real and load-bearing. The headline claim of “24% improved perplexity with similar KV-Cache and attention footprint” rests on baselines that are not actually footprint-matched. The Appendix’s iso-KV construction (Eq. 9) sets F = E/P to match 2BLF = 2BLE/P + 2BDS, but it ignores the leading-token cache Attamba stores. More importantly, the Figure 7 baseline uses a sliding window of size L/P, so with P=4 and L=32 it attends to only 8 recent tokens, while Attamba attends to roughly 288 chunk states. That is a 36x smaller attention map, not a matched one. The perplexity gap almost certainly reflects the baseline’s severe window restriction rather than a benefit of SSM compression. The authors need a baseline with window W = n/P + L_lead and a properly matched F before the abstract claim is supportable.\n\nAlso, the test-time mask in Eq. 7 looks non-causal as written; it needs a corrected version. And there are no error bars or multiple seeds, so the smaller gaps (e.g., the 5% trade-off) are within noise.\n\nWho is this for? People working on hybrid attention-SSM architectures and KV-cache compression will find the mechanism worth engaging with. It deserves a serious referee, but the required revision is significant: corrected matched baselines, a clarified causal mask, and ideally seeds or confidence intervals. I would not desk-reject it, and I would not accept it as is.","headline":"A genuinely novel way to compress KV states, undone in the abstract by a baseline that is not actually footprint-matched.","tokens_in":17675,"tokens_out":2691,"would_cite":false,"duration_ms":24105,"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":"State-space models can replace transformer key-value projections, compressing token chunks and improving perplexity at a matched memory footprint.","keywords":["state-space models","attention","KV-cache compression","token chunking","transformer efficiency","language modeling","sequence compression"],"falsifier":"Measure peak GPU memory and attention FLOPs for Attamba and for the iso-KV-plus-sliding-window transformer on a long sequence (say 4096 tokens) at the same batch size, and retrain or evaluate both on a long-context benchmark; if the transformer matches Attamba's perplexity once the two are compared at equal measured footprint, the core efficiency-advantage claim collapses.","tokens_in":16589,"feed_emoji":"⚡","tokens_out":10430,"duration_ms":77402,"temperature":0.7,"pith_summary":"This paper proposes Attamba, a transformer variant that replaces the key and value projection matrices with state-space models, each of which compresses a chunk of P consecutive tokens into a single key or value state. Attention then runs over these compressed states rather than over every token, reducing both KV-cache memory and attention computation by roughly a factor of P. The authors report that on language modeling, Attamba improves perplexity by 24% relative to a transformer matched for KV-cache and attention footprint, and can trade about 5% perplexity for a roughly 4x smaller KV-cache and attention FLOP count. The broader claim is that SSM compression is not a lossy shortcut but a workable replacement for token-level attention, with flexible chunk boundaries that allow a smooth transition between quadratic and linear scaling.","feed_headline":"SSM compression in attention cuts KV cache 4x at 5% perplexity cost","feed_subtitle":"Replacing key/value projections with state-space models lets attention run on compressed chunks, saving cache and FLOPs.","key_machinery":"The KV-SSM block is the load-bearing component: a selective state-space model that reads a chunk of P tokens and outputs a single hidden state per chunk, replacing both the $W_K$ and $W_V$ projections. The chunk-aware causal mask $M_{\\text{train}}$ (Equation 5) then makes queries attend to the last token of each completed chunk and to tokens in the current partial chunk, while $M_{\\text{test}}$ (Equation 7) restricts inference attention to chunk boundaries; when a sliding window of L leading tokens is kept uncompressed (Equation 8), the mask combines boundary states with the most recent tokens. Cyclic chunking shifts the chunk offset by layer index to spread boundary bias across layers. Together these pieces turn the attention map from $n \\times n$ into $(n/P + L) \\times (n/P + L)$, which is where both the memory and FLOP savings come from.","core_discovery":"The paper's central claim is that a state-space model can serve as the key and value projector of a transformer, compressing every chunk of P tokens into one state that attention then treats as a single token. Formally, the key and value sequences are replaced by $K^{(p)} = \\mathrm{SSM}_K(X^{(p)})$ and $V^{(p)} = \\mathrm{SSM}_V(X^{(p)})$ for each chunk, and a causal mask restricts attention to the final state of completed chunks plus the current partial chunk. At inference only the boundary states are cached, so both KV-cache size and the attention map shrink by a factor of P. The authors report that this compressed attention is not simply a cheaper approximation: on WikiText-2, Attamba reaches 21.74 perplexity in an 8-billion-token training run, beating an iso-KV transformer (22.89) and an iso-KV-plus-sliding-window transformer (28.66), and the abstract's headline comparison states a 24% perplexity improvement over a transformer of similar KV-cache and attention footprint. They also show the SSM compresses robustly across chunking strategies, including randomized boundaries, which they take as evidence that attention over compressed states is a flexible alternative to token-level attention.","pith_inferences":["The paper evaluates only WikiText-2, a highly local task; a more demanding test would be long-context retrieval or needle-in-a-haystack, where compressed chunk states must preserve distal information rather than recent locality.","Robustness to randomized boundaries suggests a learned or importance-based boundary predictor could trade a small perplexity penalty for large compute savings at inference, extending the paper's fixed/cyclic chunking to adaptive chunking.","The claimed 24% gain rests on derived iso-KV and iso-FLOP baseline formulas (Appendix A); directly measuring peak memory and wall-clock time of Attamba against those baselines on the same hardware would make the efficiency comparison robust."],"forward_implications":["At inference, only the final SSM state of each chunk is cached, so KV-cache size and attention-map size shrink by a factor of P; with chunk size 8 and 64 leading tokens, the paper reports roughly 8x KV-cache compression at a 10% perplexity cost.","Variable-length chunk boundaries, implemented via cu_seqlens, let the same model move between quadratic and linear attention by adjusting how many tokens each SSM state compresses.","Cyclic chunking (shifting chunk boundaries by layer index) improves perplexity by about 5% over fixed boundaries, and randomized boundaries match uniform chunking, indicating the compression is not tied to a specific token grouping.","Removing the explicit key and value projection matrices costs only about 1% perplexity, so the SSM block can serve as a drop-in replacement for the projection weights."],"supporting_citations":[{"why":"Supplies the Mamba selective state-space model and its cu_seqlens support, which Attamba uses as the key/value chunk compressor.","marker":"(Gu & Dao, 2023)"},{"why":"Documents the state-collapse limitation of recurrent state models, the failure mode Attamba repurposes into chunk compression.","marker":"(Chen et al., 2024)"},{"why":"Provides dclm-baseline-1.0, the training corpus from which the 10% subset used in all experiments is drawn.","marker":"(Li et al., 2024)"},{"why":"Provides the training framework and the Xmer transformer baseline used for iso-parameter and footprint-matched comparisons.","marker":"(Videau et al., 2024)"},{"why":"Supplies evidence that token importance varies, motivating variable and importance-aware chunk boundaries.","marker":"(Zhang et al., 2023)"},{"why":"The Mamba2 structured state-space duality formulation, grounding the SSM family from which the compressors are taken.","marker":"(Dao & Gu)"}],"fun_headline_variants":["Attamba: compressed attention states slash KV cache 4x","SSM key/value projectors shrink attention cache and FLOPs","24% perplexity gain at same cache, or 4x cache cut","Attention on chunked states: variable-length, efficient","State-space compression in attention: better perplexity, less cache"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The footprint-matched comparison assumes that setting a transformer's attention dimension to $F = E/P$ and evaluating with a sliding window of $L/P$ tokens reproduces Attamba's KV-cache and attention-map cost; if those baselines are not truly equal in measured memory and FLOPs, the claimed 24% perplexity advantage is not established.","fun_headline_variants_meta":{"raw":{"variants":["Attamba: compressed attention states slash KV cache 4x","SSM key/value projectors shrink attention cache and FLOPs","24% perplexity gain at same cache, or 4x cache cut","Attention on chunked states: variable-length, efficient","State-space compression in attention: better perplexity, less cache"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000418,"raw_usage":{"total_tokens":2169,"prompt_tokens":979,"completion_tokens":1190,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":595,"completion_tokens_details":{"reasoning_tokens":1116}},"tokens_in":595,"tokens_out":1190,"duration_ms":10798,"temperature":1.0,"reasoning_tokens":1116,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-12T11:50:45.811456+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Measure peak GPU memory and attention FLOPs for Attamba and for the iso-KV-plus-sliding-window transformer on a long sequence (say 4096 tokens) at the same batch size, and retrain or evaluate both on a long-context benchmark; if the transformer matches Attamba's perplexity once the two are compared at equal measured footprint, the core efficiency-advantage claim collapses.","supporting_citations":[{"cited_title":"Y., Haziza, D., Wehrstedt, L., Copet, J., Teytaud, O., and Lopez-Paz, D","cited_arxiv_id":null,"evidence_quote":"Provides the training framework and the Xmer transformer baseline used for iso-parameter and footprint-matched comparisons."}],"review_version":1}