PaperlayerView exact revision on arXiv

Source-linked AI summary

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

James O' Neill, Fergal Reid

arXiv:2607.15456v2cs.LGcs.CL

TL;DR

Looped Transformers share weights but still cache separate K/V states for every recurrence step, making cache size a serving constraint. LLA stores a low-rank cross-loop K/V trajectory and reconstructs loop-specific states on demand; it outperforms matched-budget cache alternatives and transfers across looped-model families.

  • Problem

    Weight-tied Transformers still store a separate K/V state for every recurrence step, so their per-request cache grows with context and generated tokens and can constrain serving capacity.

    Source: §1 · p. 1

  • Method

    LLA replaces per-loop K/V entries with compact cross-loop latents and reconstructs the loop-specific states needed by attention.

    Source: §1 · p. 1, §3 · p. 4

  • Results

    LLA outperforms matched-budget head, layer, precision, and endpoint baselines and transfers across two looped-model families.

    Source: §9 · p. 12

  • Takeaways & Limitations

    The results identify recurrence as a practical cache-compression axis for looped decoders and a potential target for native latent-cache training.

    Source: §9 · p. 12

    The latent-store path realizes memory savings but is compute-bound in the current kernel, so the reported systems result is a capacity gain rather than a decode-speed claim.

    Source: §7 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

Weight tying reduces model parameters but leaves a per-request KV cache that grows with context and every recurrent step. LLA targets the loop axis because its K/V states form low-rank trajectories, and the paper reports matched-budget and serving benefits.

Motivation

  • At inference, the private KV cache grows with context and generated tokens, so it can constrain batch capacity and serving economics despite shared model weights.With T = 4 loops and D = 24 layers, the cache is equivalent in size to that of a 96-layer decoder.Source: §1 · p. 1

Cache structure

  • For a fixed token, layer, and head, K/V vectors trace a low-rank trajectory across loops, while head and layer spectra are flatter.LLA replaces the per-loop vectors with a latent and loop-specific reconstruction maps.Source: §1 · p. 1

Axis comparison

  • At matched cache budget, loop-axis LLA is favored over head-axis MLA, cross-layer sharing, quantization, and final-loop reuse; the endpoint control collapses GSM8K generation to zero.The final-loop control and LLA both provide 4× reduction in the stated T = 4 comparison.Source: §1 · p. 2

Deployment

  • An explicit on-policy refinement objective uses student-generated prefixes to stabilize long rollouts, while exact scalar reduction raises measured Ouro-1.4B maximum batch from 32 to 768 sequences at 21.3× compression.Source: §1 · p. 2

2 Related Work

Prior work ties Transformer weights across recurrence or compresses caches along head, layer, precision, and sequence axes. LLA instead targets the loop axis and retains a compressed loop trajectory rather than collapsing it to one state.

Looped and Universal Transformers

  • Universal Transformers cast depth as recurrence, and looped language models extend the approach to billion-parameter LMs (Zhu et al., 2025), but still store K/V for every loop.Figure 1 contrasts that loop-indexed cache with cross-loop latents that reconstruct the needed K/V at each step.Source: §2 · p. 2, Figure 1 · p. 3

Latent and compressed KV caches

Memory-efficient looped Transformers

  • MELT distills one gated K/V state per layer, whereas LLA preserves a rank-compressed loop trajectory; a collapsed state fixes reduction near the loop count, but LLA’s rank is tunable.The paper contrasts MELT with a final-loop-reuse control that fails at the same cache budget.Source: §2 · p. 2

3 Looped Latent Attention

LLA replaces per-loop cache entries with compact cross-loop K and V latents, then reconstructs loop-specific entries when attention needs them. Its variants, initialization, and refinement preserve a trajectory while offering different compression regimes.

3 Looped Latent Attention

  • LLA stores compact cross-loop K and V latents instead of separate per-loop cache entries, reconstructing the K/V needed at each recurrence step.The codec changes the cache representation while leaving the recurrent model and attention computation otherwise unchanged.Source: §3 · p. 3
  • The uncompressed cache stores both K and V for every loop and head, costing 2THdhead scalars per token per layer; LLA uses separate latents and potentially different ranks for K and V.The fixed codec matrices are shared across tokens, contexts, and batches, so they do not scale with sequence length.Source: §3 · pp. 3–4
  • LLA-2D folds heads into one latent per token and layer, reducing cache size by a factor H versus per-head LLA but coupling heads through a shared code.Its down-projection input grows with head count and recurrence length, becoming a constraint at large T.Source: §3 · p. 4
  • The codec is initialized from top-r right singular vectors of centered teacher K/V activations and then fine-tuned with the teacher frozen, updating only codec parameters.Reconstruction is fit to pre-RoPE content projections; the shared token position lets RoPE be applied after reconstruction.Source: §3 · p. 4
  • At 4× compression, SVD initialization reaches held-out KL 0.105 after conversion, versus 0.771 with random initialization under the same budget.The fine-tuning objective combines forward KL from the teacher distribution with attention-output matching, with λattn = 0.5 in reported runs.Source: §3 · p. 4
  • Optional on-policy refinement addresses long-rollout distribution shift by training on prefixes sampled from the swapped model, which the frozen teacher scores.The stage updates only codec parameters, does not filter by reward or correctness, and stops the gradient through sampling.Source: §3 · p. 4
  • Matched-budget baselines test head-axis MLA, adjacent-layer sharing, low-bit quantization, and a zero-parameter final-loop K/V cache reused for every loop.The final-loop control directly tests whether the trajectory can be replaced by its endpoint.Source: §3 · p. 4

4 Cross-loop structure in the teacher cache

The frozen teacher’s cache is substantially lower-rank across recurrence than across heads or layers, and K/V trajectories stabilize without collapsing to one endpoint. The same slower-V convergence pattern appears in Huginn-3.5B.

Cross-loop structure

  • In Ouro-1.4B and Ouro-2.6B-Thinking, loop-axis normalized effective rank is about 0.61–0.63, versus nearly full-rank head and layer axes.A single direction carries roughly 60 percent of the loop-axis energy.Source: §4 · p. 5

Cross-loop dynamics

  • K and V move toward the final loop as step-to-step changes shrink, but V stabilizes more slowly, motivating a larger value rank than key rank.At matched scalar cost, rk/rv = 96/160 outperforms shared-rank 128/128.Source: §4 · p. 5

Trajectory is not collapse

  • Most cross-loop variation remains outside the leading direction, so the cache follows a short low-rank trajectory rather than a single shared endpoint.Only 6 percent of K variance and 17 percent of V variance lie on the loop axis.Source: §4 · p. 5, Figure 3 · p. 5

Transfer to Huginn-3.5B

  • Huginn-3.5B shows the same recurrence-axis pattern across up to 32 steps: states approach a fixed point, each step refines less, and V converges more slowly than K.Huginn uses a different recurrent-depth architecture than Ouro.Source: §4 · p. 6, Figure 4 · p. 6

5 Experiments

Experiments compare loop-axis compression with alternatives at matched cache budgets and test performance across tasks, compression levels, model scales, and a second model family. Per-head LLA leads the axis comparison, while LLA-2D extends the family to extreme compression.

Quality under compression

  • Through 4×, code and MMLU-Pro stay at or above the teacher; reasoning tasks degrade, with MATH-500 most sensitive and GSM8K and BBH dropping sharply beyond 2×.In the cold-start GSM8K sweep, strict-EM accuracy rises from about 0.58 at 4×–2.67× to 0.75 at 2× and 0.795 at 1.33×.Source: §5 · p. 6, Figure 5 · p. 7

Iso-cache axis comparison

  • At 4×, final-loop reuse scores GSM8K 0.000, while loop-axis LLA preserves performance against matched-budget head, layer, and precision baselines.Table 1 compares the compression axes; its endpoint-reuse control uses the same 4× cache budget as LLA.Source: Table 1 · p. 7, §5 · p. 7

LLA-2D trade-off

  • At 21.3×, LLA-2D trails per-head LLA on generation, but its H× smaller cache makes it an option for extreme compression.LLA-2D also spends rank on the flatter head axis.Source: §5 · pp. 7–8

Scale generalization

  • From 4× to 10.67×, LLA remains near-lossless on full GSM8K for Ouro-2.6B-Thinking, while the head baseline loses up to 0.25 absolute accuracy at 4×.The scale experiment preserves the axis-comparison ordering.Source: Table 2 · p. 8, §5 · p. 8

Design ablations

  • SVD initialization reduces held-out KL by an order of magnitude versus random initialization; asymmetric ranks help at matched scalar cost, with a smaller gain from attention-output matching.The ablations indicate that LLA primarily exploits teacher geometry.Source: §5 · pp. 8–9

Transfer to Huginn-3.5B

  • On Huginn-3.5B, SVD alone is near-lossless to 32×, and KL distillation recovers much of the 64×–128× range.The family-transfer evaluation uses 32 recurrence steps and decoder-independent fidelity.Source: Table 3 · p. 9, §5 · p. 9

6 Long-generation stability

On-policy refinement stabilizes long compressed-model rollouts by training on the prefixes the codec model itself visits. Passkey retrieval also reveals a compression–context trade-off whose operating range shifts with model size.

6 Long-generation stability

  • Off-policy accuracy falls below the teacher on long rollouts and no-answer rates spike, while on-policy distillation tracks the teacher across length bins.Teacher-prefix fitting faces exposure bias because the deployed codec model conditions on its own prefixes; on-policy training changes the prefixes used while retaining the teacher as the target.Source: §6 · pp. 9–10, Figure 7 · p. 10
  • On-policy distillation raises MATH-500 accuracy by 0.16 to 0.24 at 3.2 to 5.33× and reduces no-answer generations at every ratio.At 21.3× it improves termination but not accuracy, indicating that the smallest latent has already lost task-relevant information.Source: §6 · p. 10, Table 4 · p. 10
  • Reliable exact passkey retrieval (≥0.9) extends to ≥64k context at ≤5.33× compression, 16k at 10.67×, and 8k at 14.2×.At 10.67–14.2×, retrieval is perfect up to a length-dependent knee, then decays gracefully as farther keys are lost first.Source: §6 · p. 10
  • At 10.67×, the 2.6B model retains perfect retrieval to 32k, while the 1.4B model has fallen to 0.80.The 2.6B model’s 64k cell OOMs during prefill, a memory limit rather than an accuracy failure.Source: §6 · p. 10, Figure 8 · p. 11

7 Serving and compute

The latent-store path turns exact cache reduction into measured batch-capacity gains, but its current reconstruction-on-read implementation is compute-bound. The full-memory fast path is faster but does not realize the memory savings.

7 Serving and compute

  • At context length 4096 on one H200, Ouro-1.4B maximum batch rises from 32 to 128 at 4× and 768 at 21.3× compression.Ouro-2.6B-Thinking rises from 16 to 64 at 4×; the tests fabricate the decode cache, push batch to OOM, and validate latent attention against full reconstruction.Source: §7 · p. 11, Table 5 · p. 11
  • The latent-store path realizes the ρ× memory saving but is compute-bound because the current kernel reconstructs cached tokens on read.At the reported endpoints it delivers 108 versus 748 peak tokens/s.Source: §7 · p. 11, Table 5 · p. 11
  • The full-memory fast path runs at 74.1 versus 63.6 ms/token (1.16× teacher latency), but does not realize the memory reduction.It caches reconstructed K/V rather than the latent.Source: §7 · p. 11, Table 5 · p. 11
  • Per-head LLA favors fidelity, while LLA-2D trades accuracy for a smaller single latent.An absorbed-attention formulation could avoid explicit reconstruction, but post-hoc decoupled RoPE is less faithful than full-RoPE reconstruction.Source: §7 · pp. 11–12

8 Discussion

The cache’s strongest redundancy lies across recurrence steps: it is a low-rank trajectory rather than a safely reusable endpoint. Deployment suitability depends on compression rank, and the evaluated absorbed-attention path has a distinct fidelity boundary.

8 Discussion

  • The loop spectrum is steeper than the head or layer spectra, and matched-budget codecs on those axes lose more accuracy; the ordering transfers from Ouro to Huginn.Final-loop reuse collapses generation at the same 4× budget, so intermediate cache states remain functionally distinct despite recurrent convergence.Source: §8 · p. 12
  • At 1.33–2×, LLA is close to a drop-in replacement across the broad suite; at roughly 3–5×, long rollouts benefit from on-policy refinement.Beyond 10×, it is primarily a capacity tool as exact retrieval and long-form reasoning become increasingly rank-limited; at 21.3×, refinement improves termination but not accuracy.Source: §8 · p. 12
  • The recurrence axis is orthogonal to head sharing, cross-layer reuse, quantization, and token eviction, but this priority need not hold for single-pass Transformers.The suggested priority follows from weight-tied iteration.Source: §8 · p. 12
  • The absorbed-attention path is 2.3× faster than reconstruct-then-attend at 262k context in a shape-matched random-weight decode benchmark.A learned query adapter lowers KL from 0.55 to 0.34, while full-RoPE reconstruction reaches 0.059; the measurements establish low rank in the cache, not in the full recurrent computation.Source: §8 · p. 12

9 Conclusion

The paper concludes that recurrence is a practical compression axis for looped decoders because the recurrence-indexed cache is low-rank even though it is not collapsed to one state. This reframes the serving cost of weight tying and motivates native latent-cache training.

9 Conclusion

  • LLA replaces the recurrence-indexed cache with a low-rank cross-loop K/V trajectory and outperforms matched-budget head, layer, precision, and endpoint baselines.The result transfers across two looped-model families, and exact scalar reduction yields measured serving-capacity gains.Source: §9 · p. 12
  • Weight tying shares parameters across loops but leaves the KV cache untied and growing with loop count, making serving cost a limit at long context and large batch.The same tied computation drives per-loop states onto a low-rank trajectory, making recurrence a compression design lever.Source: §9 · p. 12

A Experimental details

The experiments evaluate post-training KV codecs on frozen recurrent Transformers, using controlled cache swaps and multiple evaluation protocols. The additional results report task-relative accuracy, 2.6B compression behavior, and serving-capacity scaling.

A Experimental details

  • Experiments use frozen Ouro-1.4B and Ouro-2.6B-Thinking teachers at four loops, plus Huginn-3.5B at 32 loops; only the K/V codec is trained.The teachers’ weights remain unchanged throughout codec training.Source: A Experimental details · p. 14
  • Codecs are initialized by SVD of teacher pre-RoPE activations from 255 prompts, then fine-tuned with teacher-forced KL and per-loop attention-output matching.The recipe uses 3,000 steps (300 for design ablations), and selects the best held-out-KL checkpoint.Source: A Experimental details · p. 14
  • Learned baselines share the codec-conversion recipe but change architecture, while quantization is inference-only and final-loop reuse is a zero-parameter control.The baselines target head compression, cross-layer sharing, low-bit values, or reuse of the last loop’s K/V.Source: A Experimental details · p. 14
  • Evaluation spans GSM8K, MATH-500, BBH, code, MMLU-Pro, commonsense and knowledge tasks, and passkey retrieval, with protocols and item counts varying by task and codec.Most generative evaluations use cached reconstruction; train-KL and the commonsense/knowledge suite are decoder-independent.Source: A Experimental details · p. 14

B Additional figures

  • Table 6 expresses broad-suite accuracy as compressed-to-teacher ratios, making task-dependent retention comparable on a common scale.Code and MMLU-Pro retain the most under compression, while GSM8K and BBH degrade fastest.Source: Table 6 · p. 15, A Experimental details · p. 15
  • At fixed 4,096-token context, measured maximum batch capacity follows ideal compression scaling; the 2.6B figure shows the loop-axis codec closest to the teacher across reported evaluations.The precision baseline cannot reach the largest ratios without collapse.Source: Figure 9 · p. 15, Figure 10 · p. 16

C Compression Across The Broad Suite

Across tasks, light compression is close to the teacher, while quality at higher ratios depends on generation length and task. The reported comparisons also show a sharp GSM8K rank transition, close composition with token eviction, and robust loop-axis behavior at 2.6B and in passkey retrieval.

C Compression Across The Broad Suite

  • Light compression matches the teacher across tasks, but heavier compression is task-dependent and long-form mathematical reasoning breaks first as generation length grows.Multiple-choice and short-generation tasks tolerate the middle compression regime better than longer generations.Source: Table 7 · p. 16, C Compression Across The Broad Suite · p. 16
  • 0.795 GSM8K strict-EM accuracy is reached at 1.33× compression, after accuracy stays near 0.58 from 4× to 2.67× and rises sharply at 2×.The passage characterizes this as a rank effect rather than a method ceiling.Source: C Compression Across The Broad Suite · p. 17, Table 8 · p. 17

D Token-eviction composition table

  • 36 MiB is the reported 1M-context Ouro-1.4B cache size with LLA-21.3× plus a 1,024-token eviction budget.The table describes recurrence compression and sequence eviction as composing with little additional quality cost beyond eviction.Source: Table 9 · p. 17

E On-policy rank Pareto (figure)

  • On reliably scored MATH-500, self-rollout distillation raises low-compression Pareto points and reduces no-answer generations across the Pareto.Hollow markers show independent re-evaluations with bootstrap 95% confidence intervals.Source: Figure 11 · p. 17

F Design-ablation detail

  • SVD initialization beats random initialization for the 4× codec: held-out KL falls quickly from the SVD prior, which random initialization does not match within the conversion budget.The design ablation uses a matched 300-step conversion and compares against split ranks and attention-output matching.Source: Figure 12 · p. 18, Table 10 · p. 18

G Full 2.6B commonsense and knowledge suite

  • Per-head LLA is near-lossless on every reported Ouro-2.6B-Thinking commonsense and knowledge budget and remains several points above head-axis, layer, and LLA-2D baselines.The suite uses 0-shot loglikelihood accuracy and is decoder-independent.Source: Table 11 · p. 19

H Cross-loop stabilization at 2.6B

I Long Context Retrieval Details

  • LLA passkey retrieval is near-lossless through 5.33× across tested contexts, except the 4×, 64k cell at 0.87; 10.67× degrades at 32–64k and 21.33× loses exact key identity.The uncompressed teacher retrieves perfectly to 64k; the codec results use the latent-store capacity decoder.Source: Table 12 · p. 20

J Additional serving and discussion details

The serving and discussion section distinguishes memory-saving latent storage from fidelity-oriented reconstruction, describes an absorbed-attention decode path, and bounds post-hoc decoupled RoPE results. It also identifies native latent-attention designs as future directions rather than part of the evaluated path.

J Additional serving and discussion details

  • LLA-21.3× with a 1,024-token eviction budget caps the 1M-context Ouro-1.4B cache at 36 MiB, while streamed layer-wise prefill enables reported 64k passkey cells.A one-shot prefill implementation exhausts memory beyond approximately 32k because it retains captures and full-width reconstructions.Source: J Additional serving and discussion… · p. 20
  • LLA-2D stores one latent for all heads, giving a factor-H smaller cache than per-head LLA at fixed ranks, while per-head LLA is more accurate at matched cache budget.Either variant can use full-memory reconstruction or latent-store capacity mode.Source: J Additional serving and discussion… · p. 20
  • Absorbing the K up-projection into the query lets attention operate directly on latents; the shape-matched 4× benchmark reports 2.3× faster decoding at 262k context than reconstruct-then-attend.The absorbed path reaches teacher decode latency with 4× less cache; RoPE still requires explicit handling of the positional component.Source: J Additional serving and discussion… · p. 20, Table 13 · p. 20
  • Post-hoc decoupled RoPE does not reach full-RoPE reconstruction, so the main accuracy results use full-RoPE reconstruction.The characterization reports a learned query adapter and extended training, but retains a reconstruction gap.Source: J Additional serving and discussion… · pp. 20–21, Table 14 · p. 21
  • Direct latent attention and shared-content computation with loop-specific residuals are proposed future architecture directions; reported results instead preserve teacher attention computation and change the cache representation.Source: J Additional serving and discussion… · p. 21

Related papers

Loading 2607.15456v2…