REVIEW 3 major objections 4 minor 1 cited by
Attamba: Attending To Multi-Token States
T0 review · 3 major / 4 minor · reviewed 2026-08-12 · deepseek-v4-flash
Pith's one-line read State-space models can replace transformer key-value projections, compressing token chunks and improving perplexity at a matched memory footprint.
desk verdict A genuinely novel way to compress KV states, undone in the abstract by a baseline that is not actually footprint-matched. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
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.
What would settle it
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.
Extended reading notes
Core claim
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.
Load-bearing premise
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.
Editorial extensions
If this is right
- 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.
Reading between the lines
- 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.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
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.
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 (3)
- [§5.1, Figure 7, Appendix A Eqs. (9)-(11)] 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.
- [§4.1, Eq. (7)] 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.
- [§5, §5.2, Appendix B] 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.
minor comments (4)
- [Notation throughout] 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.
- [Appendix A.1.1, A.2] There are typographical errors: 'intorduce' should be 'introduce', 'Psuedo-chunking' should be 'Pseudo-chunking', and 'inStuffed Mamba' in Related Work is missing a space.
- [§4.1, Figure 13] 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.
- [§6, Figure 17] 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.
Circularity Check
No derivation-level circularity; the central quality claim is empirical. The only mild circularity is that the 'footprint-matched' transformer baselines are defined from Attamba's own chunking hyperparameters, making the comparison partly self-defined rather than an independent control.
-
self definitional
[Section 5.1 and Figure 7 caption; Appendix A, Eq. 9]
"Specifically, for transformers, we emulate smaller KV-Cache sizes by reducing the attention model dimension F such that F = E/P, and smaller attention maps by employing sliding window attention (SWA) during evaluation. ... to match the attention map size, we evaluate these models in Sliding Window attention with window size = L/P."
The 'matched' transformer baseline's footprint is defined directly from Attamba's own hyperparameters (F = E/P and window = L/P), so the 'similar KV-Cache and attention footprint' claim is partly fixed by definition rather than independently established. Attamba's actual test-time cache and attention map include the leading-token block, giving roughly n/P + L keys (Eq. 8), while the constructed sliding-window baseline uses only L/P keys and a reduced attention dimension F. Thus the comparison baseline is constructed to be smaller than Attamba, and the headline footprint match is a definitional artifact rather than a measured equality. The perplexity advantage itself is still an empirical result, so this is a mild baseline-construction circularity, not a derivation-level one.
full rationale
The paper does not derive its quality claim from an equation that reduces to a fitted constant. Attamba is an architecture proposal evaluated empirically: models are trained on a dclm-baseline subset and evaluated on WikiText2 test perplexity, with pseudo-chunking (Figure 16) and ablations (Figures 11, 12, 14, 15) providing independent internal evidence. The self-citations present in the paper (Akhauri et al., 2024 for token importance; Palu, which shares an author) are related-work or motivational statements, not load-bearing premises: the token-importance self-citation motivates chunk-boundary selection, but the paper's own experiments show FAttn and FSSM chunking do not help, so the final method does not depend on that cited result. There is no imported uniqueness theorem and no ansatz smuggled in solely via self-citation. The only mild circularity is the construction of 'fair' transformer baselines using Attamba's own P and L (F = E/P, sliding-window size = L/P), which makes the footprint match definitional and arguably understates the baseline's attention restriction. This is an experimental-design fairness concern rather than a case where a prediction reduces to its inputs by construction, so the circularity score is low.
Assumptions & free parameters
free parameters (4)
- Chunk size P =
4, 8, 64, 128 (default 4 or 8)
- Leading tokens L =
32 (for P=4), 8, 16, 64, 128, 1024 in ablations
- SSM state dimension Ds =
32 (default)
- Cyclic chunking layer offset =
layer index
assumptions (4)
- domain assumption Selective state-space models (Mamba) can compress a chunk of tokens into a single state that retains information useful for attention (Section 3.2, Eq. 3).
- domain assumption WikiText2 test perplexity is a valid proxy for language-model quality for the claims made (Section 5).
- ad hoc to paper The test-time attention mask (Mtest in Eq. 7) enforces causality when attending to compressed chunk states.
- ad hoc to paper The iso-KV/iso-FLOPs baseline formulas (Eqs. 9-11) capture the dominant costs of attention and KV-cache, making the reduced-dimension transformers fair comparisons.
Cite this review
Pith. "Pith review of Attamba: Attending To Multi-Token States." pith.science (2026). https://pith.science/paper/4QPZWGJL
@misc{pith2026241117685,
author = {Pith},
title = {Pith review of: Attamba: Attending To Multi-Token States},
year = {2026},
howpublished = {\url{https://pith.science/paper/4QPZWGJL}},
note = {Machine review of arXiv:2411.17685}
}
read the original abstract
When predicting the next token in a sequence, vanilla transformers compute attention over all previous tokens, resulting in quadratic scaling of compute with sequence length. State-space models compress the entire sequence of tokens into a fixed-dimensional representation to improve efficiency, while other architectures achieve sub-quadratic complexity via low-rank projections or sparse attention patterns over the sequence. In this paper, we introduce Attamba, a novel architecture that uses state-space models to compress chunks of tokens and applies attention on these compressed key-value representations. We find that replacing key and value projections in a transformer with SSMs can improve model quality and enable flexible token chunking, resulting in 24% improved perplexity with transformer of similar KV-Cache and attention footprint, and ~4 times smaller KV-Cache and Attention FLOPs for 5% perplexity trade-off. Attamba can perform attention on chunked-sequences of variable length, enabling a smooth transition between quadratic and linear scaling, offering adaptable efficiency gains.
Figures
Figures from the paper (13 more)
Forward citations
Cited by 1 Pith paper
-
DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling
DART augments Mamba-2 with attention over its own chunk state contributions, improving associative recall and retrieval with a smaller inference cache than token-level attention.
Reference graph
Works this paper leans on
-
[1]
F., Dotzel, J., Zhang, Z., Rush, A
Akhauri, Y., AbouElhamayed, A. F., Dotzel, J., Zhang, Z., Rush, A. M., Huda, S., and Abdelfattah, M. S. S hadow LLM : Predictor-based contextual sparsity for large language models. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp.\ 19154--19167, Miami, Florida...
work page 2024
-
[2]
Chang, C.-C., Lin, W.-C., Lin, C.-Y., Chen, C.-Y., Hu, Y.-F., Wang, P.-S., Huang, N.-C., Ceze, L., Abdelfattah, M. S., and Wu, K.-C. Palu: Compressing kv-cache with low-rank projection, 2024. URL https://arxiv.org/abs/2407.21118
arXiv 2024
-
[3]
Stuffed mamba: State collapse and state capacity of rnn-based long-context modeling
Chen, Y., Zhang, X., Hu, S., Han, X., Liu, Z., and Sun, M. Stuffed mamba: State collapse and state capacity of rnn-based long-context modeling. arXiv preprint arXiv:2410.07145, 2024
arXiv 2024
-
[4]
M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J
Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mohiuddin, A., Kaiser, L., et al. Rethinking attention with performers. In International Conference on Learning Representations
- [5]
-
[6]
De, S., Smith, S. L., Fernando, A., Botev, A., Cristian-Muraru, G., Gu, A., Haroun, R., Berrada, L., Chen, Y., Srinivasan, S., et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.19427, 2024
arXiv 2024
-
[7]
O., Bengio, Y., and Hajimirsadegh, H
Feng, L., Tung, F., Ahmed, M. O., Bengio, Y., and Hajimirsadegh, H. Were rnns all we needed?, 2024. URL https://arxiv.org/abs/2410.01201
arXiv 2024
-
[8]
Y., Roziere, B., Lopez-Paz, D., and Synnaeve, G
Gloeckle, F., Idrissi, B. Y., Roziere, B., Lopez-Paz, D., and Synnaeve, G. Better & faster large language models via multi-token prediction. In Forty-first International Conference on Machine Learning
Show all 24 references
-
[9]
and Dao, T
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023
2023 arXiv
-
[10]
Efficiently modeling long sequences with structured state spaces
Gu, A., Goel, K., and Re, C. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations
-
[11]
Hippo: Recurrent memory with optimal polynomial projections
Gu, A., Dao, T., Ermon, S., Rudra, A., and R \'e , C. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 0 1474--1487, 2020
2020
-
[12]
W., Jitsev, J., Kollar, T., Dimakis, A
Li, J., Fang, A., Smyrnis, G., Ivgi, M., Jordan, M., Gadre, S., Bansal, H., Guha, E., Keh, S., Arora, K., Garg, S., Xin, R., Muennighoff, N., Heckel, R., Mercat, J., Chen, M., Gururangan, S., Wortsman, M., Albalak, A., Bitton, Y., Nezhurina, M., Abbas, A., Hsieh, C.-Y., Ghosh,...
2024
-
[13]
Jamba: A hybrid transformer-mamba language model
Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev-Shwartz, S., et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024
2024 arXiv
-
[14]
Pointer sentinel mixture models
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference on Learning Representations, 2022
2022
-
[15]
The illusion of state in state-space models
Merrill, W., Petty, J., and Sabharwal, A. The illusion of state in state-space models. In Forty-first International Conference on Machine Learning
-
[16]
Samba: Simple hybrid state space models for efficient unlimited context language modeling
Ren, L., Liu, Y., Lu, Y., Shen, Y., Liang, C., and Chen, W. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522, 2024
2024 arXiv
-
[17]
Shadowkv: Kv cache in shadows for high-throughput long-context llm inference, 2024
Sun, H., Chang, L.-W., Bao, W., Zheng, S., Zheng, N., Liu, X., Dong, H., Chi, Y., and Chen, B. Shadowkv: Kv cache in shadows for high-throughput long-context llm inference, 2024. URL https://arxiv.org/abs/2410.21465
2024 arXiv
-
[18]
Y., Haziza, D., Wehrstedt, L., Copet, J., Teytaud, O., and Lopez-Paz, D
Videau, M., Idrissi, B. Y., Haziza, D., Wehrstedt, L., Copet, J., Teytaud, O., and Lopez-Paz, D. Meta lingua: A minimal PyTorch LLM training library, 2024. URL https://github.com/facebookresearch/lingua
2024
-
[19]
and Li, Q
Wang, S. and Li, Q. Stablessm: Alleviating the curse of memory in state-space models through stable reparameterization. In Forty-first International Conference on Machine Learning
-
[20]
Z., Khabsa, M., Fang, H., and Ma, H
Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020
2006 arXiv
-
[21]
Efficient streaming language models with attention sinks
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations
-
[22]
A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 0 17283--17297, 2020
2020
-
[23]
H2o: Heavy-hitter oracle for efficient generative inference of large language models
Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., R \'e , C., Barrett, C., et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36: 0 34661--34710, 2023
2023
-
[24]
write newline
" write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 gl...
Reviewed August 12, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.