PaperlayerView exact revision on arXiv

Source-linked AI summary

Hybrid Latent Attention for Looped Language Models

Yuhan Chen, Siyuan Zhang, Nan Wang, Feiyang Kang, Ruoxi Jia

arXiv:2610.07940v1cs.CLcs.AI

TL;DR

Looped models’ KV caches grow with the number of loops, limiting concurrent decoding and requiring each decoding step to read the whole cache. HLA keeps recent KV exact and stores older tokens as latents read directly by each loop; on two Ouro models it shrinks cache use while retaining benchmark accuracy close to the original, and remains trainable.

  • Problem

    Looped models cache KV for every loop, so the cache grows with loop count and limits concurrent sequences while each decoding step reads the whole cache.

    Source: §1 · p. 1

  • Method

    HLA keeps exact KV for recent tokens and stores older tokens as compact latents that each loop reads directly, and it is uptrained with pretrained weights frozen.

    Source: §1 · pp. 1–2

  • Results

    10.7× smaller KV cache per token accompanies 2.5–7.4× higher peak decoding throughput, over 97% retained accuracy on math, knowledge, and reasoning benchmarks, and 96–100% on RULER up to 16K tokens.

    Source: §1 · p. 2

  • Takeaways & Limitations

    At the tested context lengths, latents read directly by each loop can replace most per-loop cache without reconstructing keys and values, while HLA remains trainable.

    Source: §7 · p. 11

    Evidence covers two Ouro sizes with T=4; contexts beyond 16K are untested, fine-tuning was tested only on math reasoning, and it uses two forward passes per step.

    Source: §6 · p. 11

Abstract

from arXiv · show

Looped language models apply the same stack of layers T times to each token, which deepens the model without adding parameters but multiplies its key-value (KV) cache by T. The larger cache limits how many sequences a GPU can decode at once and slows each decoding step, which reads the whole cache. We propose Hybrid Latent Attention (HLA), which keeps exact keys and values within a sliding window of W recent tokens and stores each older token as a compact latent that the query of each loop reads directly, without reconstructing keys and values. We uptrain HLA on Ouro looped models (T=4) with 1.4B and 2.6B parameters, keeping the pretrained weights frozen and training only the added parameters to reproduce the original attention. The cache shrinks by 10.7x per token, fitting 4.0-8.8x as many concurrent sequences per GPU, and decoding throughput improves by 2.5x at 1K-token contexts and by up to 7.4x at 16K. HLA retains over 97% of the original accuracy on math, knowledge and reasoning benchmarks, and 96-100% on long-context retrieval up to 16K tokens. After supervised fine-tuning, it performs on par with the fine-tuned original model on competition-level math.

1 Introduction

Looped models multiply their KV cache across loops, constraining concurrent decoding and increasing per-step reads. HLA replaces most of that cache with directly read latents while preserving recent exact KV, yielding substantial memory and throughput gains with accuracy close to the original models.

1 Introduction

  • Each loop attends to its own KV, making the looped model’s cache T times that of a non-looped model with the same layers.The cache limits concurrent sequences, and each decoding step reads the whole cache.Source: §1 · p. 1
  • Prior reuse of the last loop’s KV was found unsafe for long generations, while LLA reconstructs per-loop KV and reached 7–18% of full-cache throughput in the authors’ measurements.The last-loop reuse variant retained 60% of full-cache math accuracy at the authors’ memory budget on Ouro-1.4B.Source: §1 · p. 1
  • 10.7× smaller KV cache per token and 2.5–7.4× higher peak decoding throughput accompany over 97% retained benchmark accuracy and 96–100% on RULER up to 16K tokens.After two-pass fine-tuning, HLA performs on par with fine-tuned originals on AIME and HMMT problems.Source: §1 · p. 2
  • HLA keeps exact KV for the latest W tokens and stores older tokens as compact latents that each loop reads without reconstructing KV.A shared softmax combines latent and exact attention; the first loop uses a smaller latent.Source: §1 · p. 2, Figure 1 · p. 2
  • The authors uptrain added HLA maps with pretrained weights frozen, then use two passes to fine-tune because a single parallel pass does not reproduce decoding.The first pass writes latents and the second reads them.Source: §1 · p. 2

2 Related Work

Prior work reduces KV cache through sharing, quantization, eviction, and low-rank compression, but these approaches do not directly resolve the cache growth across loops. HLA adapts direct latent reading to looped models rather than reconstructing loop-specific KV on each read.

Looped language models

  • Looped-model research on quality per parameter leaves open the KV cache that grows with loop count, while Ouro keeps only the final loop’s KV for generated tokens and LLA offers a latent alternative [Geiping et al., 2025, Zhu et al., 2025, O’Neill and Reid, 2026].HLA carries MLA-style absorption across loops with loop-specific maps and an exact recent-token window; FlashLoop is described as complementary in principle.Source: §2 · p. 3

Key–value cache reduction

3 Hybrid Latent Attention

Hybrid Latent Attention replaces most per-loop KV for past tokens with latents read directly by each loop, while retaining exact KV for a recent window. Its writer compresses hidden states across loops, and its reader combines latent and exact attention.

3 Hybrid Latent Attention

  • HLA replaces past-token per-loop KV with latents read directly by each loop, retaining exact KV for the most recent W tokens.The sequence stores a latent for every token and the full cache for W tokens; the window adds a fixed per-sequence cost.Source: §3 · p. 3, §3.4 · p. 4

3.1 Preliminaries

  • Because hidden states and their keys and values differ across loops, the full cache stores each past token’s KV at every loop, totaling 2TLd numbers per token.The model applies the same L Transformer layers T times to each token.Source: §3.1 · p. 3

3.2 Latent writer

  • The writer uses a main latent for loops 2 to T and a separate, smaller latent for loop 1, whose KV are harder to predict from the final hidden state.The main latent accumulates linear maps of later-loop hidden states; the loop-1 latent maps the first hidden state.Source: §3.2 · pp. 3–4

3.3 Latent reader

  • Each loop maps its query into latent key space, scores stored key latents, and maps the attention-weighted value-latent sum back to the head dimension without forming per-loop KV.The key-latent rotation approximates original scores and is exact only when its rank covers all (T −1)H key directions per frequency; shifted-position training helps keep scores accurate at long distances.Source: §3.3 · p. 4

3.4 Exact sliding window

  • HLA reads tokens older than the W-token window through latents and recent tokens through exact KV, combining both with one softmax.With W=0, the attention reduces to pure latent attention.Source: §3.4 · pp. 4–5

4 Training

Uptraining fits newly added HLA maps to the frozen pretrained model’s attention, using initialization and training data designed to cover long distances. Fine-tuning uses two sequence passes to recover parallel training while approximating the token-by-token decoding order.

4.1 Uptraining

  • Uptraining adds HLA maps and freezes pretrained weights; PCA initialization uses keys and values from 128 calibration sequences, and full rank reproduces original attention exactly.At the chosen ranks, initialization only approximates attention, and uptraining reduces the remaining error.Source: §4.1 · p. 5
  • The uptraining loss fits both attention distributions and outputs, with latent supervision only on tokens served by latent attention at decoding.One reference forward pass supplies hidden states for parallel fitting; training runs 1,000 steps on sequences up to 2,048 tokens with shifted positions, then 200 steps adding eight 16K-token documents per batch.Source: §4.1 · p. 5

4.2 Fine-tuning

  • Fine-tuning uses two passes: the first writes latents and the second applies the decoding rule in parallel, with loss on the second pass and gradients through both.A single parallel pass does not reproduce decoding because a token’s latent is complete only after its last loop, although two passes come close to sequential decoding.Source: §4.2 · pp. 5–6

5 Experiments

Experiments evaluate HLA’s cache and decoding savings, retained accuracy, fine-tuning behavior, and design choices on Ouro-1.4B and Ouro-2.6B. The results show higher serving throughput with accuracy close to the full cache, including after fine-tuning.

5.1 Setup

  • On T=4 Ouro models, HLA stores 10.7× less cache per token; evaluation covers math, knowledge, reasoning, and RULER retrieval at contexts up to 16K.The experiments use Ouro-1.4B and Ouro-2.6B checkpoints and serve models with vLLM at 16-bit precision on 80 GB A100 GPUs.Source: §5.1 · pp. 6–7

5.2 Memory and decoding speed

  • HLA fits 4.0–8.8× as many sequences and raises peak throughput 2.5–7.4×; at 16K, single-sequence decoding is twice as fast as the full cache.At 1K, HLA is 26% slower per decoding step, while at 4K it is equal; throughput gains vary with context length.Source: §5.2 · p. 7

5.3 Uptraining

  • With pretrained weights unchanged, HLA retains 97.5–99.6% of full-cache accuracy on MATH500, MMLU-Pro, and BBH, and 96.0% on Ouro-1.4B RULER at 16K.On Ouro-1.4B RULER, the accuracy loss grows from 0.8 points at 4K to 3.1 points at 16K.Source: §5.3 · p. 8
  • Single-loop sharing loses substantially more accuracy than HLA, while LLA is less accurate on every Ouro-1.4B benchmark and delivers only 7–18% of full-cache peak throughput.LLA stores 2.7× as much per token as HLA; rebuilding keys and values for the full history accompanies its lower throughput.Source: §5.3 · p. 8, §5.2 · p. 8

5.4 Fine-tuning

  • After fine-tuning, HLA averages 16.4% versus 16.2% for the full cache on Ouro-1.4B and 28.0% versus 27.2% on Ouro-2.6B.The paired-bootstrap intervals include zero for both averages, and no competition difference exceeds 1.7 points.Source: Table 2 · p. 9, §5.4 · p. 9

5.5 Ablations

  • A separate smaller latent for loop 1 and an exact 128-token window address poor predictability and the disproportionate divergence of recent tokens.Without the window, accuracy falls 8.1 points on MATH500 and 11.5 points on RULER at 16K; recent distances account for a large share of attention divergence.Source: §5.5 · pp. 9–10
  • On an earlier W=32 model, shifted positions limit attention divergence at a 4K shift to 0.026 rather than 0.098 and raise RULER accuracy at 16K from 54.3 to 62.7.The RULER gain grows with context length; the ablation used an earlier model and training setup.Source: §5.5 · p. 10
  • On an earlier W=128 model, the long-document phase raises RULER accuracy at 16K from 69.6 to 71.2 under the original settings.Source: §5.5 · p. 10

6 Limitations and Future Work

HLA has accuracy and speed boundaries: its frozen-weight accuracy gap grows with context length, while single-sequence decoding benefits only beyond a few thousand tokens. The evidence also covers only T=4 Ouro models at two sizes, with fine-tuning tested only on math reasoning.

6 Limitations and Future Work

  • Beyond 16K tokens, HLA's accuracy gap is untested and may widen, while the reported gap grows with context length.On RULER, Ouro-1.4B loses 0.8 points at 4K and 3.1 at 16K, retaining 96.0% of original accuracy at 16K.Source: §6 · p. 11
  • At 1K context, single-sequence decoding is 26% slower; the caches break even near 4K, although larger batches gain throughput even at 1K.Prompt processing uses exact attention and therefore costs the same as with the full cache; HLA's fixed per-step cost is attributed to its maps and separate window and latent attention.Source: §6 · p. 11
  • Evidence covers T=4 Ouro models at two sizes; fine-tuning was tested only on math reasoning and requires two forward passes per step.The window and latent sizes were chosen on Ouro-1.4B and reused for Ouro-2.6B; combinations with cache-precision or token-eviction methods remain unevaluated.Source: §6 · p. 11

7 Conclusion

The conclusion presents HLA as a direct-reading latent alternative to most of the per-loop KV cache, with substantial cache and decoding reductions at tested context lengths. The initialization appendix constructs frequency-aligned projections from pretrained keys and values, treating loop 1 separately.

7 Conclusion

  • A 10.7× per-token cache reduction and up to 7.4× peak decoding-throughput increase accompany retention of over 96% of original accuracy in the reported HLA results.HLA keeps exact keys and values for recent tokens and stores older tokens as compact latents read directly by each loop, without reconstructing keys or values.Source: §7 · p. 11

A Initialization of the Latent

  • At initialization, P_f projects stacked keys onto m directions shared across heads and loops; when m=(T−1)H, the projection is identity and recovers the original score.Source: A Initialization of the Latent · p. 15
  • The key directions are leading eigenvectors of a calibration-set second-moment matrix, minimizing squared projection error over those tokens.The same projection is optimal at every position because the matrix is unchanged by rotating every pair by the same angle.Source: A Initialization of the Latent · p. 15
  • Initialization projects values from loops 2–T jointly, while loop 1 receives its own latent built from its keys and values.Source: A Initialization of the Latent · p. 15

B Rotating the Whole Key Latent

HLA rotates every key-latent pair to preserve positional dependence in attention scores, while sharing frequency-aligned combinations across loops. The choice has initialization and post-uptraining trade-offs, and the alternatives have not been compared after uptraining.

B Rotating the Whole Key Latent

  • After uptraining, latent pairs may mix key components from other frequencies, so shifted-position training helps keep learned scores accurate over long distances.Scores still depend on positions through the relative offset, but their distance dependence may differ from the original model.Source: B Rotating the Whole Key Latent · p. 16
  • Rotating every latent pair avoids the position-free component's score-change term; at initialization, the rotated latent's only approximation is the projection P_f.A position-free component changes its contribution unless the relative-position angle is near zero or the component has no energy.Source: B Rotating the Whole Key Latent · p. 16
  • Across loops, one rotary key from one hidden state approximates other loops' keys, while one per loop makes the cache grow with T; a combined rotary latent avoids both.Source: B Rotating the Whole Key Latent · p. 16
  • The design has not been compared with split rotary and position-free latents after uptraining; equal-cache comparisons and frequency-energy allocation remain future work.Rotating every pair also distributes latent pairs equally across frequencies even when their eigenvalues may be uneven.Source: B Rotating the Whole Key Latent · p. 16

C Training Details

The sections describe the uptraining data and schedule, analyze where latent attention diverges from the full cache, and assess a two-pass procedure for fine-tuning. The two-pass procedure approaches sequential decoding, while distance-based analysis motivates the exact attention window.

C Training Details

  • The first uptraining phase uses 64–2,048-token sequences mixing math traces and web text at 3:2; the second uses 1,950 16,384-token book and code windows.Both phases use AdamW, batch size 128, and cosine learning-rate schedules; the phase peak rates are 10−3 and 3 × 10−4.Source: C Training Details · p. 16

D Where the Latent Errs

  • Tokens at distances of 1,024 or more hold 29.5% of attention and 11.0% of divergence, while distances 1–127 account for 47%–51% of divergence in every loop.The analysis compares full-cache and HLA attention from the same reference hidden states, separating single-step error from errors accumulated across layers.Source: D Where the Latent Errs · p. 17

E Two-Pass Fine-Tuning

  • A response of m tokens is reproduced exactly after ⌈m/(W+1)⌉ passes under the decoding rule.Each repetition extends the exact response range by W+1 tokens.Source: E Two-Pass Fine-Tuning · p. 17
  • At 32-bit precision, two-pass decoding lowers KL divergence by a factor of 16 to 19, and one repetition reduces it by roughly another factor of ten.At 16-bit precision, the two-pass divergence is 5.8 × 10−4 for Ouro-1.4B and 4.0 × 10−4 for Ouro-2.6B; one repetition halves the latter.Source: E Two-Pass Fine-Tuning · p. 18, Table 5 · p. 18

F Shifted Positions

This section tests whether shifted positions address errors at long query–key distances and whether that change improves long-context retrieval. The reported gains increase with context length.

F Shifted Positions

  • At 14K, attention divergence is 0.032 with shifted positions versus about 0.09 without them; the short-sequence model’s error comes from distances it has not seen.The divergence is 0.022 for both continuations at D=0.Source: F Shifted Positions · p. 18
  • At 16K, shifted positions raise the average RULER score from 54.3 to 62.7 and the eleven retrieval and aggregation tasks’ average from 58.7 to 72.8.The RULER average also rises from 81.5 to 82.3 at 4K and from 71.2 to 73.7 at 8K.Source: F Shifted Positions · p. 18

G Latent and Window Sizes

The section examines how latent ranks and window size affect attention and task accuracy, then explains the selected configuration. A larger window addresses long-distance errors and prompt repetition in question answering.

G Latent and Window Sizes

  • Before training, attention divergence for loops 2–4 depends only on key rank, while attention-output error depends only on value rank; all three measured errors fall as rank increases.The measurements use initialized latents without a window and do not represent the accuracy after uptraining.Source: G Latent and Window Sizes · p. 19, Table 6 · p. 19
  • Halving every latent rank to 256 loses about 15 MATH500 points, while doubling the key rank or both ranks to 1,024 changes accuracy by at most 3.4 points as the cache grows by a quarter.A 32-token window on the 72 KiB latent scores 71.8, above the no-window model with a latent 1.7 times larger.Source: G Latent and Window Sizes · p. 19, Table 7 · p. 19
  • A 128-token window covers the prompt segment repeated in 27 of 100 W=32 HotpotQA examples and includes the 54–61-token median distance from the prompt end; it costs 96 MiB per sequence.The divergence per unit attention is 0.049 nats at distances 32–127 versus 0.040 at distances 1–31.Source: G Latent and Window Sizes · p. 19
  • At 16K, the W=32 model reaches 54.2 RULER accuracy versus 74.7 for the full cache; shifted positions raise it to 62.7, and the main configuration reaches 69.1.The same W=32 model scores 71.8 on MATH500 versus 75.9 for the full cache.Source: G Latent and Window Sizes · p. 19

H Evaluation Details

The section specifies how RULER question-answering tasks are prompted and scored, and describes the LLA setup used for comparison. It distinguishes those question-answering settings from the other RULER tasks.

H Evaluation Details

  • For Table 1, question-answering prompts use Ouro’s chat template and generate up to 2,048 tokens; the other 11 RULER tasks retain their original settings.Each Table 1 RULER entry averages those 11 tasks with the two question-answering tasks; other paper results use the original settings for all 13 tasks.Source: H Evaluation Details · p. 20
  • The LLA comparison uses 96 key and 160 value latent dimensions per head, trains projections with pretrained weights frozen, and decodes in two passes.The vLLM decoder implementation matched the PyTorch implementation’s chosen token at 99.6% of positions on eight MATH500 prompts.Source: H Evaluation Details · p. 20

I Serving Implementation

The serving implementation maps HLA onto vLLM’s paged cache and decoding kernels, and the section reports agreement with a PyTorch reference. It also describes how concurrent-sequence capacity is measured and identifies Table 8 as the source of Fig. 2’s decoding measurements.

I Serving Implementation

  • vLLM 0.26 runs HLA as an Ouro model class using its paged cache, scheduler, and attention kernels; the engine is unchanged except for preemption logging.The measurements use that logging.Source: I Serving Implementation · p. 20
  • HLA stores main and loop-1 latents in shared caches, while exact attention uses per-loop sliding-window caches of W+1 tokens.The engine frees blocks that fall outside the window, so that cache does not grow with sequence length.Source: I Serving Implementation · p. 20
  • All T loops run within one forward pass; latent and exact-window attention results are merged by log-sum-exp, and the step can be CUDA-graph replayed without host synchronization.The loop-1 latent is stored after the first loop and the main latent after the last.Source: I Serving Implementation · pp. 20–21
  • Prompts use exact attention over their own tokens, while prefix caching and chunked prefill are disabled for HLA and the full cache.With those features enabled, whether a prompt token is read exactly or through its latent would depend on how the scheduler splits the prompt.Source: I Serving Implementation · p. 21
  • Concurrent-sequence capacity is measured by increasing the batch until the scheduler cannot run every sequence together, counting only batches with no preemptions.The engine’s startup capacity estimate undercounts sequences because it budgets sliding-window caches pessimistically.Source: I Serving Implementation · p. 21
  • On RULER with the W=32 starting model, the serving and PyTorch implementations reached 81.5 vs 81.6 at 4K, 70.3 vs 70.4 at 8K, and 54.2 vs 54.2 at 16K.The comparison checks the serving path against the separate PyTorch implementation of Eq. (4).Source: I Serving Implementation · p. 21

J Decoding Measurements

  • Table 8 provides Fig. 2’s measurements; its LLA comparator uses a reconstruct path with rank-256 per-head latents, and ratios compare HLA with the full cache.The LLA latent uses rank 256 per head, with 96 for keys and 160 for values.Source: Table 8 · p. 21, J Decoding Measurements · p. 21

Related papers

Loading 2610.07940v1…