# Hybrid Latent Attention for Looped Language Models

> Hybrid Latent Attention compresses looped language model KV caches 10.7x by having attention read compact latents directly, boosting decoding throughput up to 7.4x while retaining over 97% accuracy.

- **Source:** [arXiv](https://arxiv.org/abs/2610.07940)
- **Published:** 2026-10-10
- **Permalink:** https://picx.dev/p/a0K1YZ
- **Whiteboard:** https://picx.dev/p/a0K1YZ/image

## Summary

## 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

| Model | Full Cache (KiB/token) | HLA Cache (KiB/token) | Reduction |
|-------|------------------------|----------------------|-----------|
| Ouro-1.4B | 768 | 72 | 10.7× |
| Ouro-2.6B | 1,536 | 144 | 10.7× |

### Accuracy with Frozen Weights

**Table 1: Accuracy (%) with pretrained weights frozen**

| Model | Method | Cache (KiB) | RULER 4K | RULER 8K | RULER 16K | MMLU-Pro | BBH | MATH500 |
|-------|--------|-------------|----------|----------|-----------|----------|-----|---------|
| Ouro-1.4B | Full cache | 768 | 89.5 | 85.2 | 77.5 | 48.9 | 71.0 | 75.9 |
| | Single-loop sharing | 192 | 19.0 | 10.0 | 7.0 | 16.1 | 43.3 | 45.9 |
| | LLA | 192 | 87.4 | 82.7 | 71.1 | 46.6 | 69.1 | 73.4† |
| | **HLA** | **72** | **88.7** | **84.2** | **74.4** | **47.7** | **69.4** | **74.1** |
| Ouro-2.6B | Full cache | 1,536 | 92.3 | 86.3 | 82.9 | 56.0 | 80.1 | 81.1 |
| | Single-loop sharing | 384 | 71.4 | 64.5 | 52.6 | 52.6 | 75.8 | 73.9 |
| | LLA | 384 | 92.5 | 85.4 | 81.0 | 54.2 | 79.3 | 80.6† |
| | **HLA** | **144** | **92.2** | **85.0** | **81.5** | **55.5** | **79.8** | **79.4** |

†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**

| Model | Method | AIME 2024 | AIME 2025 | HMMT 2025 | Avg. |
|-------|--------|-----------|-----------|-----------|------|
| Ouro-1.4B | Full cache | 20.6 | 19.8 | 8.1 | 16.2 |
| | **HLA** | **19.2** | **20.2** | **9.8** | **16.4** |
| Ouro-2.6B | Full cache | 35.4 | 30.4 | 15.6 | 27.2 |
| | **HLA** | **35.6** | **31.9** | **16.5** | **28.0** |

### Key Ablation Findings

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

| Loop | $R^2$ from $h_{j,4}$ (Keys) | $R^2$ from $h_{j,4}$ (Values) | Var. at rank 512 (One) | Var. at rank 512 (2-4) |
|------|---------------------------|-----------------------------|----------------------|----------------------|
| 1 | 0.66 | 0.55 | 0.80 | 0.79 |
| 2 | 0.80 | 0.71 | — | — |
| 3 | 0.93 | 0.89 | — | — |
| 4 | 1.00 | 1.00 | 0.93 | 0.89 |

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

| W | RULER 4K | RULER 8K | RULER 16K | MMLU-Pro | BBH | MATH500 |
|---|----------|----------|-----------|----------|-----|---------|
| 0 | 78.3 | 71.9 | 57.6 | 39.8 | 61.0 | 66.0 |
| 128 | 83.2 | 78.7 | 69.1 | 47.7 | 69.4 | 74.1 |

## 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."

---

_Markdown view of https://picx.dev/p/a0K1YZ, served by PicX — AI-generated visual whiteboard summaries of research papers._
