Summary (Overview)

• Hybrid Latent Attention (HLA) is a novel cache compression method for looped language models that replaces per-loop key-value (KV) caches with compact latents that attention reads directly, without reconstructing keys and values.
• HLA reduces the KV cache by 10.7× per token (from 768 to 72 KiB on Ouro-1.4B), enabling 4.0–8.8× more concurrent sequences per GPU and 2.5–7.4× higher decoding throughput at 1K–16K token contexts.
• With pretrained weights frozen, HLA retains over 97% accuracy on MATH500, MMLU-Pro, and BBH, and 96–100% on RULER long-context retrieval up to 16K tokens.
• HLA keeps exact keys/values for the most recent 128 tokens (sliding window) while storing older tokens as compact latents, addressing the finding that nearby tokens contribute disproportionately to attention divergence.
• After supervised fine-tuning with a novel two-pass procedure, HLA performs on par with or better than the fine-tuned original model on competition-level math (AIME, HMMT).

Introduction and Theoretical Foundation

Background

Looped language models apply the same stack of Transformer layers $T$ times to each token, gaining depth without adding parameters. The Ouro models (1.4B and 2.6B parameters, $T=4$) match standard Transformers up to 12B parameters on benchmarks. However, this parameter efficiency is offset at inference time: the KV cache grows by a factor of $T$, limiting batch sizes and slowing decoding.

Key Problem

For a looped model, each loop attends to its own keys and values, so the KV cache is $T$ times larger than a non-looped model with the same layers. An 80 GB GPU holds only five 16K-token sequences of Ouro-1.4B, and per-step latency grows from 13.6 ms at 1K tokens to 49.3 ms at 16K.

Prior Approaches and Their Limitations

1. Single-loop sharing (suggested with Ouro): Keeps only the last loop's KV for generated tokens. Found to be unsafe—retains only 60% of math accuracy on long generations.
2. Looped Latent Attention (LLA): Stores a compact latent per token but reconstructs keys/values at every decoding step, achieving only 7–18% of full-cache throughput.

Key Insight

The paper asks whether attention can read a latent as it is, without reconstruction. This requires:

• A latent that captures the low-rank trajectory of keys/values across loops
• Loop-specific linear maps that project queries into latent space and outputs back to head dimension
• A sliding window of exact tokens to handle the most attention-critical recent context

Methodology

Latent Writer

For each layer, HLA stores two latents per token. The main latent accumulates a linear map of hidden states from loops 2 to $T$, and a separate smaller latent handles loop 1:

$$\begin{array}{r} c _ {j} = \sum_ {t = 2} ^ {T} E _ {t} h _ {j, t} = \left[ \begin{array}{c} c _ {j} ^ {K} \\ c _ {j} ^ {V} \end{array} \right], \\ c _ {j} ^ {(1)} = E _ {1} h _ {j, 1} = \left[ \begin{array}{c} c _ {j} ^ {K, 1} \\ c _ {j} ^ {V, 1} \end{array} \right]. \end{array}\tag{2}$$

The key part $c_j^K \in \mathbb{R}^{r_k}$ and value part $c_j^V \in \mathbb{R}^{r_v}$ serve loops 2 to $T$; $c_j^{K,1}, c_j^{V,1} \in \mathbb{R}^{r_1}$ serve loop 1. Total storage per token: $r_k + r_v + 2r_1$ numbers per layer (vs. $2Td$ for full cache).

Latent Reader

Each loop $t$ uses two linear maps per head, $A_t \in \mathbb{R}^{d_h \times r_k}$ and $B_t \in \mathbb{R}^{r_v \times d_h}$:

$$\hat {s} _ {i j, t} = \frac {(\tilde {R} _ {i} A _ {t} ^ {\top} q _ {i , t}) ^ {\top} (\tilde {R} _ {j} c _ {j} ^ {K})}{\sqrt {d _ {h}}}.\tag{3}$$

The output combines latent and exact attention:

$$o _ {i, t} = B _ {t} ^ {\top} \sum_ {j < i - W} a _ {i j, t} c _ {j} ^ {V} + \sum_ {i - W \leq j \leq i} a _ {i j, t} v _ {j, t}.\tag{4}$$

Rotary Embedding Alignment

Key latents are organized in pairs with $m = r_k/d_h$ pairs per frequency, allowing rotation $\tilde{R}_j$ to commute with linear combinations of key pairs sharing a frequency—extending TransMLA's approach across loops.

Training Procedure

Uptraining (frozen pretrained weights):

• Initializes $E_t, A_t, B_t$ via PCA of keys/values from 128 calibration sequences
• Minimizes KL divergence between HLA and original attention distributions plus output MSE:

$$\begin{array}{l} \mathcal {L} = \frac {1}{L T} \sum_ {\text { layers }} \sum_ {t = 1} ^ {T} \Big [ \mathrm{KL} \big (a _ {i,. t} ^ {\star} \| \hat {a} _ {i.. t} \big) \\ \qquad + \frac {\| \hat {o} _ {i , t} - o _ {i , t} ^ {\star} \| ^ {2}}{\| o _ {i , t} ^ {\star} \| ^ {2}} \Big ], \end{array}\tag{5}$$

• Phase 1: 1,000 steps on sequences up to 2,048 tokens with shifted positions (PoSE) for distance up to 16K
• Phase 2: 200 steps adding 16K-token documents

Fine-tuning: Uses a two-pass procedure—first pass writes latents with exact attention, second pass runs HLA attention in parallel for training.

Empirical Validation / Results

Cache Size

Accuracy with Frozen Weights

Table 1: Accuracy (%) with pretrained weights frozen

†One sample per problem instead of four (LLA decodes slowly).

Decoding Performance

• Throughput: HLA achieves 2.5× full-cache throughput at 1K contexts, up to 7.4× at 16K
• Latency: Single-sequence latency grows ~0.5 ms per 1K tokens (vs. 2.4 ms for full cache); HLA is 26% slower at 1K, equal at 4K, 2× faster at 16K
• LLA comparison: HLA reaches 28–59× the throughput of LLA, which reconstructs keys/values at every step

Fine-tuning Results

Table 2: Accuracy (%) on competition math after fine-tuning

Key Ablation Findings

Table 3: Predictability of keys/values across loops (Ouro-1.4B)

Table 4: Effect of exact window (W) on Ouro-1.4B

Theoretical and Practical Implications

Theoretical Contributions

1. Direct latent reading: Demonstrates that attention can consume compressed latents directly through loop-specific linear maps, eliminating the reconstruction bottleneck that limits LLA and Palu-style approaches.

2. Frequency-aligned latent design: Extends TransMLA's cross-head key combination to cross-loop combination, maintaining rotary embedding compatibility in latent space.

3. Window-latent hybrid: Provides empirical evidence that recent tokens (distances 1–127) carry ~48.7% of attention divergence while receiving only 26.6% of attention, justifying the hybrid design.

4. Two-pass fine-tuning: Introduces a practical parallel training procedure for architectures where latent completeness depends on sequential token processing.

Practical Implications

• Memory efficiency: 10.7× cache reduction enables serving 4–8.8× more concurrent sequences per GPU, directly reducing hardware requirements
• Throughput gains: Up to 7.4× faster decoding at 16K contexts makes long-context looped models practical
• Accuracy preservation: Over 96% retention of original accuracy with frozen weights; competitive after fine-tuning
• Complementary to other methods: HLA is orthogonal to quantization, token eviction, and sparsity approaches

Conclusion

HLA successfully replaces most of the per-loop KV cache in looped language models with compact latents that attention reads directly. Key takeaways:

1. Design principle: A latent written from all loops (with a separate smaller latent for the unpredictable first loop) captures the low-rank trajectory of keys/values effectively.

2. The exact window is crucial: The 128-token sliding window recovers 8+ points of accuracy across benchmarks by covering the most attention-critical recent context.

3. Shifted positions matter: Training with shifted positions is essential for accurate long-distance latent attention.

4. Future directions:

   • Pretraining looped models with HLA from scratch (separating imitation gap from design limits)
   • Testing with more loops ($T > 4$), where savings grow proportionally
   • Kernel fusion to reduce the fixed per-step cost of the maps
   • Combining with quantization and token eviction methods

The authors conclude that "at the context lengths we test, a latent that each loop reads directly can thus replace most of the per-loop cache without reconstructing keys or values."