REVIEW 4 major objections 5 minor 13 references
The central claim: in looped transformers, the recurrence axis of the KV cache is the most compressible one, and this structure is intrinsic to weight-tied iteration, not a quirk of a single architecture.
Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →
T0 review · deepseek-v4-flash
2026-08-01 23:17 UTC pith:G3W5GTDY
load-bearing objection First real compression of the loop axis in looped transformers; spectral and matched-budget evidence holds up, but the global-subspace assumption and loose evaluation protocols are the soft spots. the 4 major comments →
Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
The central claim: in looped transformers, the KV cache is low-rank along the recurrence axis specifically. Across two looped families, cross-loop K/V spectra have effective normalized rank ~0.6 with one direction carrying ~60% of energy, while head/layer axes are near full rank. LLA stores a centered latent per layer/head/axis, SVD-initialized and distillation-refined, reconstructing loop-specific K/V via per-loop up-projections. At matched cache budget it preserves GSM8K at 4x where final-loop reuse collapses to zero and head/layer/precision baselines lose substantial points; ordering holds at 2.6B scale and transfers to a 32-loop architecture. The exact scalar reduction raises batch capac
What carries the argument
The key machinery is the cross-loop low-rank codec. For each layer, head, and axis (K or V), the T per-loop vectors are stacked and centered, then projected to a latent c=(x−μ)W_down of rank r; loop t is reconstructed as x̂_t = c W_up,t + μ_t. The codec is initialized from the top right singular vectors of teacher activations and refined with forward-KL plus attention-output matching, with separate ranks for K and V (r_v>r_k). Its work is to replace the recurrence-indexed cache with a compact latent while leaving the attention computation and RoPE unchanged, yielding exact compression ρ = 2T d_head/(r_k + r_v) per head per layer.
Load-bearing premise
The codec assumes that every token's K/V trajectory across loops lies in one fixed linear subspace per layer/head/axis, estimated from a 255-prompt calibration set, and that this subspace generalizes to all contexts and tokens.
What would settle it
Take the trained LLA codec and a set of prompts deliberately outside the calibration distribution, then compute the per-token reconstruction error of the SVD codec against the teacher's true K/V trajectories; if the residual norm along the loop axis approaches the cross-loop variance for a substantial fraction of tokens, the single-subspace model fails. Alternatively, run the 21.3x codec on a passkey task with contexts in a domain unseen in the 255 prompts; a retrieval collapse far below the paper's in-distribution frontier would indicate context-dependent subspaces.
If this is right
- By making recurrence the first-class compression axis, looped transformers can serve longer contexts and larger batches at fixed memory: the paper measures a batch capacity increase from 32 to 768 sequences at 21.3x compression on one GPU at 4k context.
- The loop-axis advantage transfers across model families, suggesting the same codec will work for other weight-tied or iterative architectures without retraining, just from an SVD of teacher activations.
- At matched cache budget, loop-axis compression beats head/layer/precision compression and final-loop reuse, so deployment should compress recurrence first and stack other methods only if the quality cost is justified.
- Long reasoning rollouts require on-policy refinement (training on the model's own sampled prefixes) to prevent compounding drift; with it, MATH-500 accuracy at 4x rises from 0.43 to 0.66 and no-answer generations drop.
- Aggressive compression has a cost: exact passkey retrieval degrades past ~5.3x, with a predictable reliable-retrieval length that shrinks as compression rises, bounding the usable operating point.
Where Pith is reading between the lines
- If the shared-subspace assumption holds, the same cross-loop latent construction could compress caches in other iteration-indexed computations (e.g., diffusion denoising steps, iterative refinement loops) where a tied operator produces an untied but low-rank state.
- The paper's absorbed-attention analysis suggests that native training with a latent cache could close the decoupled-RoPE fidelity gap; if that holds, memory savings could convert into compute savings, making looped models practical for very long context serving.
- The sharp accuracy break at 2x compression on GSM8K implies a phase transition in what the latent encodes; a testable extension is per-token or per-head rank allocation that could push the break point without increasing total memory.
- Because the codec is fit to only 255 prompts, a natural stress test is whether the single global subspace survives out-of-distribution domains or long-tail tokens; if not, a small mixture of codecs would be needed.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes Looped Latent Attention (LLA), a post-training codec for the KV cache of looped/weight-tied transformers. For each token, layer, and head, LLA stores a low-rank latent along the loop axis and reconstructs loop-specific K/V vectors with per-loop up-projections, giving an exact scalar compression ratio (Eq. 4). The paper claims that the loop axis is the most compressible cache axis, based on spectral analysis of teacher activations, matched-budget comparisons against head-axis MLA, cross-layer sharing, KV quantization, and final-loop reuse, and transfer to a second model family (Huginn). It also reports serving-capacity gains on one H200 and an on-policy refinement stage that improves long-rollout MATH-500 accuracy.
Significance. If the central claims hold, the paper identifies a new, recurrence-specific KV-cache compression axis for looped models and gives a practically useful exact memory reduction. The derivation of Eq. (4) is exact and simple; the spectral diagnostics in Fig. 2 directly test the low-rank premise; and the matched-budget ablations across four axes are the right experimental design. The measured H200 batch-capacity numbers and the on-policy refinement objective are concrete additional contributions. The main weaknesses are evaluation-protocol inconsistencies and an untested global-subspace assumption, which prevent the results from being accepted at the reported strength without revision.
major comments (4)
- [Appendix A / Tables 1, 7, 8] The evaluation protocol is not apples-to-apples. Appendix A states that the 1.4B teacher is scored on the full 1,319-item GSM8K test set, while iso-cache and broad-suite codecs use a 100-item subset, and the Table 8 cold-start sweep uses a 200-item subset with a different decoder. This is visible in the tables: teacher GSM8K is 0.794 in Table 7 but 0.820 in Table 8, and LLA 4× is 0.800 in Table 1 but 0.590 in Table 8. The claim that per-head LLA 'preserves GSM8K' at 4× therefore rests on comparing a 100-item codec score to a full-set teacher score. Please re-score all arms on the same items (or full set) and report confidence intervals; otherwise the matched-budget ordering is not established.
- [§3, Eqs. (2)–(3); §4, Fig. 2] The codec assumes a single global linear subspace per layer/head/axis, estimated from 255 teacher prompts, but the spectral evidence in Fig. 2 is an aggregate property of the calibration set. The paper does not test whether the per-token T-step trajectories lie in the same shared subspace; if different tokens or contexts have different subspaces, the fixed W_down/W_up matrices will under-fit outside the calibration distribution. The passkey collapse at 21.3× (Table 12) and Huginn 128× degradation (Table 3) are consistent with this failure. Please report a per-token subspace-alignment diagnostic (e.g., fraction of per-token cross-loop variance captured by the global basis vs. a per-token basis) and a calibration-set-size sensitivity sweep.
- [§5, Table 3 (Huginn transfer)] The text states that on Huginn 'head-axis SVD reconstruction is 20 to 30 times worse than LLA at matched budget' and that MC-avg is unchanged at 16×/32×. Table 3 reports only the LLA arm; no Huginn baseline numbers or matched-budget comparisons are shown anywhere in the manuscript. Since the transfer claim is that the axis ordering holds in a second family, the missing baseline data are load-bearing. Please add the head-axis/layer/precision comparisons for Huginn at the same budgets.
- [Appendix A (held-out KL)] Held-out KL is evaluated every 200 steps on n=32 prompts, and the checkpoint with the best held-out KL is selected and reported. This makes the reported train-KL in Tables 1, 2, and 10 a selected statistic rather than the behavior of a fixed procedure, and with n=32 the selection noise is substantial. Please report final-step KL, seed variance, or confidence intervals, and justify that the n=32 held-out set is large enough for stable checkpoint selection.
minor comments (5)
- [§4, Fig. 3] The sentence 'only 6 percent of K variance and 17 percent of V variance lie on the loop axis' appears to contradict Fig. 2's statement that a single direction carries roughly 60 percent of the energy. Clarify whether the percentages refer to content variance rather than cross-loop variance.
- [Table 13] The header 'context Lteacher (ms/tok)' appears to be a typo; it should likely be 'context' and 'teacher (ms/tok)'.
- [Fig. 5 caption] The caption has missing punctuation: 'GSM8K is 5-shot strict EM HumanEval/MBPP are EvalPlus pass@1 (base); MATH-500...' Please clean up.
- [Tables 1 and 8] Explain in the main text why the same 4× LLA codec reports 0.800 in Table 1 and 0.590 in Table 8; the protocol note is buried in Appendix A and the discrepancy is confusing.
- [Eq. (5)] The attention-matching weight λ_attn is fixed at 0.5 in the reported runs. Consider reporting a small sensitivity sweep over λ_attn to show that the main ordering is robust to this choice.
Circularity Check
No significant circularity: the central derivation is self-contained and validated on external benchmarks.
full rationale
The paper's core claims are (1) that loop-indexed K/V caches are low-rank along the recurrence axis, (2) that compressing this axis with a per-head latent codec preserves downstream accuracy better than compressing other axes, and (3) that the resulting scalar reduction translates into measured serving-capacity gains. The low-rank observation is read directly from the frozen teacher's activations (Fig. 2, Fig. 4) and is an empirical measurement, not a fitted prediction. The LLA codec is initialized from the SVD of teacher activations (Eqs. 2–3), but the downstream evaluations are external: GSM8K, MATH-500, passkey retrieval, broad-suite benchmarks, and measured batch capacity. Held-out KL is computed on a separate set of n=32 prompts, and the best-KL checkpoint is selected on that held-out data, so it is not a train-set reconstruction. The matched-budget baselines (head-axis MLA, cross-layer sharing, quantization, final-loop reuse) are trained and evaluated under the same protocol, making the axis comparison fair rather than circular. The only self-referential metric is train-KL / decoder-independent KL, which is explicitly a reconstruction measure, not a claimed prediction. There is no load-bearing self-citation: the cited prior work (Dehghani, Zhu, Geiping, Blayney, DeepSeek-AI, etc.) is external and used for context, not to establish the paper's own results. No equation reduces a prediction to a fitted parameter by construction, and no uniqueness theorem or ansatz is smuggled in via self-citation. The paper is self-contained against external benchmarks, so the circularity score is 0.
Axiom & Free-Parameter Ledger
free parameters (4)
- Codec ranks (r_k, r_v) =
e.g., 96/160 at 4×; 120/200 at 3.2×; 144/240 at 2.67×; 192/320 at 2×; 240/400 at 1.6×; 288/480 at 1.33× (Table 8)
- Attention-matching weight λ_attn =
0.5
- Calibration set =
255 prompts
- Distillation training budget =
3000 steps, lr 1e-3, batch 8, seq 512
axioms (5)
- domain assumption The looped model reuses the same block and projections at every recurrence step, so recurrence is a well-defined cache axis.
- domain assumption A token has the same position at every recurrence step, so RoPE can be applied after reconstruction.
- domain assumption Teacher pre-RoPE K/V activations on a 255-prompt calibration set are representative of all contexts and lengths up to 64k.
- domain assumption The value cache stabilizes more slowly than the key cache, motivating r_v>r_k.
- domain assumption Low-rank loop-axis structure is a property of weight-tied iteration and transfers across architectures.
read the original abstract
Looped, weight-tied Transformers reduce parameters by reusing a single block, but decoding still stores a separate K/V cache for every recurrence step. We show that this loop-indexed cache is highly structured. For a fixed token, layer and head, K/V vectors trace a short low-rank trajectory across loops, while the head and layer axes remain much flatter. We introduce Looped Latent Attention (\lla{}), a post-training cache codec that stores compact K and V latents and reconstructs loop-specific K/V vectors only when attention reads them. The default per-head codec compresses recurrence, while \lla{}-2D also folds heads into one latent for the extreme-compression regime. The codec is initialized from the SVD of teacher activations and refined with logit and attention-output distillation. At matched cache budget, per-head \lla{} outperforms head-axis MLA, cross-layer sharing, KV quantization and final-loop reuse, showing that the recurrent cache is low-rank but not safely collapsible to a single state. The same axis advantage holds on Ouro-2.6B-Thinking and transfers to Huginn-3.5B, where an SVD codec remains near-lossless to $32\times$ compression in decoder-independent evaluation. The cache reduction is exact. On one H200, the latent-store path increases measured Ouro-1.4B batch capacity at 4k context from 32 to 768 sequences at $21.3\times$ compression. Lastly, for long reasoning rollouts such as in MATH-500, on-policy refinement on student-generated prefixes raises accuracy at $4\times$ compression from 0.43 to 0.66 and reduces no-answer generations when compared to token-level off-policy distillation.
Figures
Reference graph
Works this paper leans on
-
[1]
GQA: Training generalized multi-query transformer models from multi-head checkpoints
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP),
2023
-
[6]
The precision baseline cannot reach the largest ratios without collapse
The loop-axis codec remains closest to the teacher on train-KL, GSM8K and the commonsense/knowledge suite. The precision baseline cannot reach the largest ratios without collapse. 15 Figure 10:The latent cache gives an exact memory reduction and measured serving capacity. Left: cache memory scales linearly with context and is divided by the compression ra...
arXiv 2024
-
[7]
Fast transformer decoding: One write-head is all you need.arXiv preprint arXiv:1911.02150,
Noam Shazeer. Fast transformer decoding: One write-head is all you need.arXiv preprint arXiv:1911.02150,
Pith/arXiv arXiv 1911
-
[8]
RoFormer: Enhanced transformer with rotary position embedding.arXiv preprint arXiv:2104.09864,
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced transformer with rotary position embedding.arXiv preprint arXiv:2104.09864,
-
[10]
Victor Conchello Vendrell, Arnau Padres Masdemont, Niccolo Grillo, Jordi Ros-Giralt, Arash Behboodi, and Fabio Valerio Massoli. Memory-efficient looped transformer: Decoupling compute from memory in looped language models.arXiv preprint arXiv:2605.07721,
-
[11]
Scaling latent reasoning via looped language models.arXiv preprint arXiv:2510.25741,
Rui-Jie Zhu, Zixuan Wang, Kai Hua, Tianyu Zhang, Ziniu Li, Haoran Que, Boyi Wei, Zixin Wen, Fan Yin, Ge Zhang, Wenhao Huang, Yoshua Bengio, and Jason Eshraghian. Scaling latent reasoning via looped language models.arXiv preprint arXiv:2510.25741,
-
[12]
uses a200-item subset with the earlier streaming reconstruction decoder and a100-item teacher reference. At 2.6B, scored with the faithful decoder, the teacher and the per-head and head-axis arms use the full1319items and the LLA-2D arms use a100-item subset.MATH-500: 4-shot,max_gen_toks=2048, batch 1, faithful decoder, scored bymath_verify symbolic equiv...
2048
-
[2019]
Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach.arXiv preprint arXiv:2502.05171,
-
[2021]
Yutao Sun et al. You only cache once: Decoder-decoder architectures for language models.arXiv preprint arXiv:2405.05254,
-
[2023]
Akide Liu et al. Minicache: Kv cache compression in depth dimension for large language models.arXiv preprint arXiv:2405.14366,
-
[2024]
DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.arXiv preprint arXiv:2405.04434,
-
[2025]
Kvquant: Towards 10 million context length llm inference with kv cache quantization
Coleman Hooper et al. Kvquant: Towards 10 million context length llm inference with kv cache quantization. arXiv preprint arXiv:2401.18079,
-
[2026]
URL https://arxiv.org/abs/2604.11791. William Brandon et al. Reducing transformer key-value cache size with cross-layer attention.arXiv preprint arXiv:2405.12981,
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.