Hybrid Latent Attention for Looped Language Models

Hybrid Latent Attention (HLA) shrinks the per-token KV cache of looped language models by 10.7× while retaining over 96% of the original accuracy and raising peak decoding throughput by up to 7.4×. <ref:2610.07940#pg1,shrinks the cache by 10.7× per token>. The work targets a structural cost of looped models, which apply the same stack of layers T times to each token, deepening the model without adding parameters but multiplying its key–value (KV) cache by T. <ref:2610.07940#pg1,multiplies its key–value (KV) cache by T> That cache bounds how many sequences decode together: an 80 GB GPU holds only five 16K-token sequences of Ouro-1.4B. <ref:2610.07940#pg1,an 80 GB GPU holds only five 16K-token sequences>

The design

HLA 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. <ref:2610.07940#pg1,stores each older token as a compact latent> For each layer, a writer stores two latents per token: a main latent accumulating a linear map of the hidden state at every loop after the first, and a smaller loop-1 latent, since the first loop's keys and values are the hardest to predict. <ref:2610.07940#pg2,whose keys and values are the hardest to predict> Each loop maps its query into the latent space, scores it against the stored latents, and maps their weighted sum to the attention output; both maps act on the current token only, so no keys or values are rebuilt. <ref:2610.07940#pg2,both maps act on the current token only> The exact window is motivated by measurement: in a model that reads every past token through its latent, the 127 nearest tokens receive about a quarter of the attention but account for about half of its divergence from the original attention. <ref:2610.07940#pg2,account for about half of its divergence from the original attention> With W=0, the method reduces to pure latent attention. <ref:2610.07940#pg4,Eq. (4) reduces to pure latent attention>

Training

HLA adds 0.35B parameters for Ouro-1.4B and 0.70B for Ouro-2.6B, and uptraining trains only these maps with the pretrained weights frozen. <ref:2610.07940#pg5,uptraining trains only these maps with the pretrained weights frozen> The maps are initialized by a principal component analysis of the pretrained keys and values, computed separately for each rotary frequency across loops and heads. <ref:2610.07940#pg5,computed separately for each rotary frequency> Uptraining minimizes a loss combining the KL divergence of attention distributions with a relative squared error on attention outputs, in two phases: 1,000 steps on short sequences with shifted positions, then 200 steps adding 16K-token documents. <ref:2610.07940#pg5,The second phase runs 200 steps> For fine-tuning, a two-pass procedure recovers parallel training: the first pass writes the main latents, and the second runs all tokens in parallel under the decoding rule. <ref:2610.07940#pg6,We recover parallel training with two passes over the sequence>

Results

On Ouro-1.4B and Ouro-2.6B served with vLLM, HLA shrinks the cache from 768 to 72 KiB per token on the smaller model, and a GPU holds 4.0–8.8× as many sequences at contexts from 1K to 16K tokens. <ref:2610.07940#pg2,from 768 to 72 KiB on the smaller model> With frozen weights it retains over 97% of the original accuracy on MATH500, MMLU-Pro and BBH, and 96–100% on RULER up to 16K tokens. <ref:2610.07940#pg2,retains over 97% of the original accuracy> After fine-tuning it performs on par with the fine-tuned original models on competition-level math, reaching 16.4% against 16.2% on Ouro-1.4B and 28.0% against 27.2% on Ouro-2.6B. <ref:2610.07940#pg9> Baselines fare worse: single-loop sharing loses far more accuracy despite storing more, and LLA, which rebuilds keys and values each step, reaches only 7–18% of the throughput of the full cache. <ref:2610.07940#pg1,LLA reaches 7–18% of the throughput of the full cache>

Limitations

With frozen weights HLA does not fully match the full cache, and the gap grows with context length, reaching 3.1 points on RULER at 16K for Ouro-1.4B. <ref:2610.07940#pg11,the gap grows with context length> HLA speeds up generation only: prompts are processed with exact attention, and a single sequence with 1K tokens decodes 26% slower, with the two caches breaking even near 4K. <ref:2610.07940#pg11,the two caches break even near 4K> Evidence covers one family of looped models with T=4 loops at two sizes, and contexts beyond 16K are untested. <ref:2610.

1. Frequency-energy-aware latent allocation. Replace the uniform every frequency receives the same number m = rk/dh of latent pairs assignment with a budget allocated per rotary frequency proportional to the eigenvalues of Σf, which the paper itself flags as a limitation (a budget allocated per frequency by these eigenvalues could keep more of the original scores at the same rank). The improved system keeps the same 72 KiB/token cache but reduces key-projection error at long distances, directly improving the RULER 16K retention (currently 96.0% on Ouro-1.4B) without extra memory.

2. Split rotary/position-free latent with a learned gate. Implement the comparison the paper leaves to future work: a hybrid key latent with a rotated part and a position-free part, plus a learned per-frequency gate initialized to the rotated solution. Since a position-free part of the latent would be insensitive to distance by construction, this removes the failure mode where without shifted positions the divergence rises from 0.022 to 0.098 at a shift of 4K, making the model robust to unseen distances and reducing reliance on PoSE augmentation.

3. Fused window+latent attention kernel. Merge the two attention paths into a single kernel rather than attending to the window and the latent separately, which the paper attributes the fixed per-step cost to. This targets the measured 26% single-sequence slowdown at 1K tokens (17.1 ms vs 13.6 ms) and should push the break-even point below 4K, making HLA faster than the full cache at all practical context lengths.

4. Single-pass fine-tuning via latent warm-start. Replace the two-pass fine-tuning (which runs two forward passes per step where standard fine-tuning runs one) with a cached-latent scheme: write main latents once per optimizer step in a no-grad pass, then run the gradient pass reading those latents, refreshing every k steps. Table 5 shows two passes reach 5.8×10−4 KL at 16-bit, so a periodic refresh should hold accuracy while cutting fine-tuning wall-clock roughly in half.

5. Pretrain-from-scratch looped model with HLA. Train a looped model with HLA attention from initialization so its hidden states adapt to the latent, rather than uptraining frozen weights. This addresses the paper's stated limitation that the gap may widen beyond 16K and should let the latent absorb information the frozen hidden states never encoded, potentially closing the 3.1-point RULER-16K gap and enabling longer contexts where the smaller cache saves the most memory.

6. Per-loop adaptive latent rank. Give each loop its own latent dimension scaled to its measured predictability—Table 3 shows loop 3 keys have R2=0.93 from h j,4 versus 0.66 for loop 1—instead of one shared main latent for loops 2–T. Reallocating bits from the highly predictable later loops to the harder early loops at constant total size should reduce the loop-1/loop-2 divergence the paper identifies as the hardest to fit.

7. Long-context extrapolation via distance-conditioned reader. Add a small position-relative bias or distance-conditioning to the reader maps A t, trained on the shifted-position data, so learned scores stay calibrated beyond the 16K training range. This directly targets the untested regime the paper flags (Contexts beyond 16K... are untested, and the gap may widen there) and makes the 10.7× cache reduction safe for 32K–128K serving.

8. Compose HLA with cache quantization and token eviction. Since HLA changes what is stored per token, so it is complementary to methods that lower the precision of the cache or keep fewer tokens, quantize the 72 KiB latent to 4-bit and apply heavy-hitter eviction within the latent. The improved system could fit roughly 4× more sequences than HLA alone (approaching 30× the full cache) while the exact window preserves the tokens where the divergence is 0.040 and 0.049 nats and is highest.

9. Learned window size per layer. Replace the fixed W=128 with a per-layer window chosen from the measured divergence profile (Fig. 4 shows divergence per unit attention varies 0.009–0.049 by distance). Layers whose attention concentrates on recent tokens get smaller windows, freeing memory for larger latents in layers where distant tokens matter, improving the accuracy/memory Pareto frontier the paper fixes by hand (chosen on Ouro-1.4B and reused for Ouro-2.6B without tuning).

10. Speculative decoding over the latent cache. Exploit the cheap latent read to draft multiple tokens per step and verify against the exact window, amortizing the fixed per-step map cost. This converts the paper's 26% single-sequence latency penalty at 1K into a throughput gain, since the dominant cost is now the per-step A t/B t application rather than cache reads that scale with context.

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.