Hybrid Latent Attention for Looped Language Models

summary

Video file (mp4)

The gist

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×.

In short

Looped language models reuse the same layers T times, which multiplies their KV cache and limits how many sequences a GPU can decode. Hybrid Latent Attention stores recent tokens exactly and older ones as compact latents read directly by each loop's query, shrinking the cache 10.7x, boosting throughput up to 7.4x, and keeping over 96% of accuracy.

Key concepts

Looped language models
Models that apply the same stack of layers T times to each token. This deepens the model without adding parameters, but it multiplies the key-value cache by T, which limits how many sequences a GPU can decode at once and slows each decoding step.
KV cache
Stored keys and values for every past token, so a model does not recompute them at each decoding step. It grows with context length and sequence count, and decoding reads the whole cache each step, so a larger cache means fewer concurrent sequences and slower decoding.
Latent attention
Instead of storing exact keys and values for old tokens, the model stores a compact latent per token. Each loop's query is mapped into that latent space, scored against the latents, and the weighted sum is mapped to the attention output, so no keys or values are reconstructed.
Uptraining
Training only newly added parameters while keeping pretrained weights frozen. Here the new maps are initialized from a principal component analysis of the original keys and values, then trained to reproduce the original attention distributions and outputs, so the model keeps its prior abilities.

Terminology used across episodes

This episode discusses

The paper

Hybrid Latent Attention for Looped Language Models · Read on arXiv

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

Virginia Tech · Independent Researcher

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.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Hybrid Latent Attention for Looped Language Models".

Jane: 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…

Tom: First, who's behind it and why it matters.

Paper summary: Tom: We're looking at a paper called Hybrid Latent Attention for Looped Language Models, from Yuhan Chen, Siyuan Zhang, Nan Wang, Feiyang Kang and Ruoxi Jia. The problem it starts from is simple to state: a looped model runs the same stack of layers T times on every token, so it gets deeper without new parameters, but its key-value cache gets multiplied by T. That bigger cache is what limits how many sequences a GPU can decode at once.

Jane: So explain the looped part to me, because a listener hearing "the same layers four times" might wonder why anyone does that. The paper says this deepens the model without adding parameters, and it works on Ouro models with four loops at one point four and two point six billion parameters. The catch is the cache, which is the memory the model keeps of everything it has already read.

Lu: And the cache is the thing that actually gets in the way. The paper's own example is that an eighty gigabyte GPU holds only five sequences of sixteen thousand tokens for Ouro-1 point 4B. That's the motivation for the whole method, and it's a real constraint, not a hypothetical one.

Meng: Five sequences. That's the number that would make me read the paper, because if you're serving this to users, that's your entire capacity on one card. So what does their fix actually keep?

Tom: It keeps exact keys and values for a sliding window of the most recent tokens, and for anything older it stores one compact latent per token, which each loop's query reads directly. Nothing gets rebuilt into keys and values. The cache shrinks by ten point seven times per token <ref:2610.07940#pg1,The cache shrinks by 10.7>.

Lalam: The claim that matters is the combination: four point zero to eight point eight times as many concurrent sequences per GPU, and throughput up two point five times at one thousand tokens and up to seven point four times at sixteen thousand. And accuracy stays above ninety-seven percent on math, knowledge and reasoning, and ninety-six to one hundred percent on long-context retrieval <ref:2610.07940#pg1,on math, knowledge and reasoning>. I'd want to see the baselines before believing the accuracy half.

Jane: Which is exactly what the next part of the show is for. We'll get into how they train it and where it stops working.

Conclusion: Tom: We've covered what Hybrid Latent Attention for Looped Language Models claims and roughly how it works. What's left is what it means once the numbers settle.

Jane: Start with the title, because it's doing real work. "Hybrid" is the sliding window of exact keys and values plus the latents for everything older. "Latent attention" is reading those stored latents directly instead of reconstructing keys and values. And "looped language models" is the family it's aimed at.

Lu: The interesting part for me is that they don't retrain the model. They freeze the pretrained weights and train only the added maps, zero point three five billion parameters for the 1 point 4B model and zero point seven zero billion for the 2 point 6B. That's a retrofit, and retrofits are how research actually reaches deployed systems.

Meng: Retrofits are also where the caveats live. The paper says with frozen weights it doesn't fully match the full cache, and the gap grows with context, three point one points on RULER at sixteen thousand tokens for the smaller model. And it only speeds up generation. Prompts still go through exact attention, and a single one thousand token sequence decodes twenty-six percent slower.

Lalam: That last number is the honest one. The two caches break even near four thousand tokens, so this pays off when you're serving many long sequences at once, not when one person chats with a short prompt. I'm fairly confident that's the right reading, and less confident it holds past sixteen thousand tokens, which they don't test.

Jane: So for a listener: this is a memory trick for a specific kind of model, and it buys capacity on a GPU at the cost of a small accuracy gap. After fine-tuning they're on par with the original on competition math, sixteen point four against sixteen point two percent on the smaller model <ref:2610.07940#pg1>.

Tom: And it's one family, four loops, two sizes. That's the scope. Thanks for listening.

More episodes

← Home