🤖 machine learning

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

This paper introduces Looped Latent Attention (LLA), a post-training codec that exploits the low-rank structure of loop-indexed KV caches in weight-tied Transformers to achieve extreme compression (up to 32x) and significantly increase batch capacity while maintaining near-lossless performance through SVD initialization and distillation.

Original authors: James O' Neill, Fergal Reid

Published 2026-07-20

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: James O' Neill, Fergal Reid

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Memory Problem in AI Brains

Imagine you are trying to teach a robot to write a story. The robot has a brain made of millions of tiny switches (parameters) that know how words fit together. But here's the catch: to write a long story, the robot can't just look at its brain; it needs a scratchpad. Every time it writes a new word, it has to remember everything it just wrote so it doesn't repeat itself or lose the plot. In the world of Artificial Intelligence, this scratchpad is called the KV cache (Key-Value cache). It's a private notebook for every single conversation, and it grows bigger with every word the robot speaks.

Now, imagine a special kind of robot designed to be super efficient. Instead of having a huge brain with many different rooms (layers) for different tasks, this robot has one small, clever room that it visits over and over again. It reuses the same set of instructions, looping through them to think deeper. This is called a Looped Transformer. It saves space on the brain itself, which is great. But there's a sneaky problem: even though the robot reuses the same room, it still writes a new page in its scratchpad every single time it loops. If the robot loops four times to solve a math problem, it ends up with four separate pages of notes for the same moment in the story. This means the "looped" robot still needs a massive amount of memory to run, defeating the purpose of being efficient. The big question for scientists is: Can we shrink that scratchpad without making the robot forget how to think?

The Paper's Discovery: A Secret Shortcut in the Notes

This paper introduces a clever trick called Looped Latent Attention (LLA) to solve that memory problem. The researchers discovered that the robot's notes aren't as messy as they look. When the robot loops through its thinking process, the notes it writes for the same word don't change wildly; they follow a smooth, predictable path, like a marble rolling down a gentle hill. Instead of saving every single page of the robot's scratchpad, the new method saves just a tiny "summary" of that path and a set of instructions on how to rebuild the specific pages only when the robot actually needs to read them.

Think of it like this: Imagine you are watching a movie where the main character walks down a hallway. In a normal computer, you would save a full, high-definition photo of the character at every single step they take. That takes up a huge amount of storage. But with LLA, you realize the character is just walking in a straight line. So, instead of saving thousands of photos, you save one photo of the character and a simple note saying, "Move 1 inch forward for step 2, 2 inches for step 3." When the computer needs to see the character at step 3, it quickly draws that picture using the note. This is exactly what the paper does for AI: it compresses the "loop" part of the memory, turning a stack of four pages into a single, tiny summary that can be expanded back into the full pages instantly.

The researchers tested this on several AI models, including one called Ouro-1.4B. They found that this method works incredibly well. In fact, on a powerful computer chip called an H200, they were able to fit 768 different conversations into the memory at once, whereas before they could only fit 32. That is a massive jump in how many people can use the AI at the same time. They also showed that this trick works even when the AI is solving very long, difficult math problems, though it does need a little extra training to make sure it doesn't get confused when it's writing its own story instead of just copying a teacher's.

What This Paper Says "No" To

It is important to know what this method is not. The researchers tested a very simple idea first: "What if we just throw away the first three pages of notes and only keep the very last one?" They thought maybe the robot just settles down to a final answer after a few loops, so the earlier notes didn't matter. They tried this, and it was a disaster. The AI completely forgot how to do math, scoring zero on tests. This proves that the "path" the notes take is important; you can't just grab the end point. You need to capture the journey.

They also compared their method to other ways of shrinking memory, like squishing the notes together (quantization) or sharing notes between different parts of the brain (head-axis compression). While those methods help a little, the paper shows that compressing the "loop" axis is the secret sauce. If you try to compress the other parts instead of the loops, the AI loses its ability to reason. The paper is very clear: the loop is the most redundant part of the memory, and that is the only place you should squeeze.

How Sure Are They?

The authors are quite confident in these results because they measured them directly. They didn't just guess; they ran the AI on real tasks like solving math problems (GSM8K) and writing code (HumanEval). They showed that with a 4x compression (making the memory 4 times smaller), the AI still performs almost as well as the original, un-compressed version. When they pushed it to extreme compression (21.3x), the AI could still handle many tasks, though it started to struggle with the very longest, most complex reasoning steps.

They also tested this on a different type of AI model called Huginn-3.5B, which has a different brain structure and loops up to 32 times. Even there, the method worked, proving that this isn't just a fluke for one specific robot. The paper suggests that this "low-rank trajectory" (the smooth path of the notes) is a fundamental property of how these looping robots think. While they have proven it works for saving memory and increasing capacity, they note that for the absolute best performance on very long tasks, the AI needs a second round of training where it practices with its own compressed notes, rather than just copying a teacher. This "on-policy" training fixes the small errors that happen when the robot starts writing its own story.

In short, this paper shows that we can make looping AI models much more efficient by realizing their memory is repetitive in a specific way. By storing a tiny summary of that repetition instead of the full notes, we can fit hundreds of times more conversations on a single computer chip, making powerful AI available to more people without needing super-expensive hardware.

Technical Summary: Looped Latent Attention (LLA)

Problem Statement

In Looped, weight-tied Transformers, a single block is applied repeatedly (recurrence) to reduce parameter count. While this effectively compresses the model's weights, the inference memory bottleneck remains the Key/Value (KV) cache. In standard looped inference, the model recomputes attention at every recurrence step, storing a separate KV cache for each step, layer, and head. Consequently, a model with $T$ loops over $D$ layers carries a KV cache equivalent to a $T \times D$ layer decoder, negating the memory benefits of weight tying.

The paper posits that this loop-indexed cache is highly structured. For a fixed token, layer, and head, the K/V vectors across recurrence steps trace a short, low-rank trajectory rather than exploring a high-dimensional space. The redundancy is specific to the loop axis; the spectra along the head and layer axes remain significantly flatter (higher rank). The challenge is to compress this recurrence axis without collapsing the trajectory into a single fixed point, which would degrade performance.

Methodology: Looped Latent Attention (LLA)

The authors introduce Looped Latent Attention (LLA), a post-training cache codec that replaces the $T$ per-loop K/V vectors with a compact latent representation and loop-specific reconstruction maps.

Core Mechanism

Latent Storage: Instead of storing $T$ $T$ vectors ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ ) for a specific token, layer, and head, LLA stores a single centered low-rank latent vector $c \in \mathbb{R}^r$ $c \in R^{r}$ .
- $c = (x - \mu)W_{\text{down}}$
- Here, $\mu$ is a calibration-set mean, and $W_{\text{down}}$ is a down-projection matrix.
Loop-Specific Reconstruction: The specific K/V vector for recurrence step $t$ $t$ is reconstructed using a loop-specific up-projection:
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
Asymmetric Ranks: The method uses separate ranks for Keys ( $r_k$ ) and Values ( $r_v$ ). Empirical analysis shows that Value vectors stabilize more slowly across loops than Key vectors, motivating $r_v > r_k$ .
LLA-2D Variant: For extreme compression, LLA-2D folds the head axis into the latent, storing a single latent per token and layer to reconstruct all $H$ heads. This reduces cache size by a factor of $H$ but couples the heads.

Training and Initialization

SVD Initialization: The codec is initialized via Singular Value Decomposition (SVD) of teacher activations on a calibration set. The down-projection is set to the top- $r$ right singular vectors. This initialization is critical; random initialization results in significantly higher KL divergence.
Distillation: The codec is refined by freezing the teacher model and minimizing the forward KL divergence between the teacher's distribution and the swapped model's distribution, augmented with an attention-output matching term.
On-Policy Refinement: To address exposure bias in long autoregressive rollouts (where the model conditions on its own generated prefixes rather than the teacher's), a second stage of on-policy distillation is employed. The codec is updated using prefixes sampled from the student model itself, stabilizing long-generation performance.

Key Contributions

Definition of LLA: A cross-loop latent K/V codec for weight-tied Transformers that achieves a compression ratio $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ . It stores $(r_k + r_v)H$ scalars per token per layer instead of $2Td_{\text{head}}$ .
Axis-Specific Compression: The paper demonstrates that the recurrence (loop) axis is the most compressible dimension in looped models (Ouro-1.4B, Ouro-2.6B-Thinking, and Huginn-3.5B). Compressing the loop axis outperforms compressing the head axis (MLA), layer axis (cross-layer sharing), or applying quantization at matched cache budgets.
Rejection of Fixed-Point Collapse: The authors show that the loop trajectory cannot be replaced by a single state (the final loop). A zero-parameter control that reuses only the final-loop K/V collapses generation performance (GSM8K accuracy drops to 0.000), proving that the "short low-rank trajectory" is necessary.
Deployment Impact: The exact scalar reduction translates directly to increased serving capacity. On a single H200 GPU, LLA increases the maximum batch size for Ouro-1.4B at 4k context from 32 to 768 sequences (at 21.3× compression).

Experimental Results

Compression vs. Accuracy

GSM8K Performance: On Ouro-1.4B, light compression (1.33× to 2×) matches the teacher's accuracy (approx. 0.79–0.80). At 4× compression, LLA maintains 0.59 accuracy, significantly outperforming head-axis compression (0.52) and final-loop reuse (0.00).
Scale Generalization: On Ouro-2.6B-Thinking, LLA remains near-lossless (0.81–0.83) across 4× to 10.67× compression, whereas head-axis baselines lose substantial accuracy.
Transfer to Huginn: The method transfers to Huginn-3.5B (32 recurrence steps). A training-free SVD codec achieves near-lossless fidelity up to 32× in decoder-independent metrics.

Long-Generation Stability

MATH-500: Long mathematical rollouts are sensitive to prefix drift. Off-policy training (using teacher prefixes) leads to accuracy degradation and high "no-answer" rates as generation length increases.
On-Policy Gain: Applying on-policy refinement raises MATH-500 accuracy from 0.43 to 0.66 at 4× compression and significantly reduces no-answer generations.

Serving Capacity

Memory Reduction: The cache reduction is exact. At 21.3× compression, the KV cache for Ouro-1.4B at 1M context drops from 768 GiB to 36 GiB.
Throughput: The "latent-store capacity path" (caching only latents) is compute-bound but enables batches and contexts previously impossible due to OOM (Out of Memory). The "fast path" (reconstructing and caching full K/V) preserves throughput but does not realize the memory saving.

Significance and Claims

The paper claims that recurrence is the most compressible cache axis in weight-tied looped Transformers. The significance lies in:

Mechanistic Insight: It identifies that the loop axis is spectrally distinct (lower rank) from head and layer axes, allowing for efficient compression without collapsing the model's dynamic behavior into a static fixed point.
Practical Deployment: It provides a post-training solution that directly increases batch size and context length capabilities on existing hardware without retraining the base model architecture.
Future Direction: The results suggest that native recurrent architectures could benefit from decoupled latent caches, where the recurrence computes shared content plus cheap per-loop deltas, rather than re-materializing full K/V states.

The authors emphasize that the method does not rely on absorbed attention implementations for its accuracy claims; the core result is the validity of the low-rank trajectory assumption and the efficacy of the codec in preserving the teacher's behavior under strict memory constraints.