Hybrid Latent Attention for Looped Language Models

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

cs.CLOct 6, 2026

Hybrid Latent Attention for Looped Language Models

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

Organizations: Virginia Tech · Independent Researcher

Abstract

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.

Figures & tables

Figure 1: Hybrid Latent Attention in one layer of a looped model with

T

loops. (a) The full cache stores the keys and values (KV) of every loop for every token. (b) HLA keeps exact KV for the last

W

tokens and stores each older token as one latent shared by loops 2 to

T

and a smaller latent for loop 1. (c) Attention within one loop. Blue modules are added by HLA and white modules are the frozen pretrained projections. The query and output maps let latent attention read the latent cache directly, and latent and exact attention share one softmax.

Figure 2: Decoding on one A100 GPU with the full cache, LLA and HLA; each sequence has a prompt of the given length and generates 256 tokens. Left: the largest batch that fits in memory. Middle: generated tokens per second summed over the batch, at the best batch size. Right: time per generated token for a single sequence. Labels give the ratio of HLA to the full cache; the left and middle columns use logarithmic axes. Table 8 lists the values.

Figure 3: Uptraining curves on held-out data for Ouro-2.6B (Ouro-1.4B in Fig. 6 ): KL divergence between the attention of HLA and of the full cache at each loop. Phase 1 is evaluated on short sequences (logarithmic axis), phase 2 on 16K-token documents (linear axis, right).

Model	Method	token (KiB)	4K	8K	16K	MMLU-Pro	BBH	MATH500
		Cache per	RULER
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

Table 1: Accuracy (%) with the pretrained weights frozen. Single-loop sharing and HLA also keep exact keys and values for the last 128 tokens (96 MiB per sequence on Ouro-1.4B, 192 MiB on Ouro-2.6B), not counted in the per-token cache. Shading marks the accuracy that HLA retains relative to the full cache: at least 97% (green) and 90–97% (yellow). Bold marks HLA results above LLA. † One sample per problem instead of four, because LLA decodes slowly.

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

Table 2: Accuracy (%) on competition math after fine-tuning on the same data for the same number of steps, averaged over 16 samples per problem. The better value of each pair is in bold.

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

Table 3: Keys and values of the original Ouro-1.4B, averaged over the 24 layers. Left:

R^{2}

of a ridge regression from the hidden state of the last loop to the keys and values of each loop. Right: fraction of variance kept by a rank-512 projection fitted to one loop or to loops 2–4 jointly.

Figure 4: Attention divergence of HLA without a window (

W{=}0

) on Ouro-1.4B, by distance between the query and the attended token. (a) How the original attention (top) and the KL divergence to it (bottom) are distributed over six ranges of distance; each bar sums to 100%. Orange ranges lie within the window of

W{=}128

tokens used in the main experiments, blue ranges beyond it. (b) Divergence per unit of original attention in each range.

$W$	4K	8K	16K	MMLU-Pro	BBH	MATH500
	RULER
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

Table 4: Effect of the exact window on Ouro-1.4B with frozen weights. Both variants are uptrained with the full recipe and store a latent of 72 KiB per token. RULER uses the original settings, unlike Table 1 ( Appendix H ).

Figure 5: Shifted positions on an earlier HLA model with

W{=}32

on Ouro-1.4B, continued for 100 steps without (green) and with (blue) them. (a) KL divergence between the attention of HLA and of the full cache when the queries of a 2,048-token sequence are moved later by the given number of positions; thick lines average three RULER tasks, thin lines show each task. (b) RULER accuracy, with the full cache in orange.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Uptraining curves on held-out data for Ouro-1.4B, drawn as in Fig. 3 .

	Ouro-1.4B		Ouro-2.6B
Procedure	32-bit	16-bit	32-bit	16-bit
Single exact pass	56	58	38	40
Two passes	3.4	5.8	2.0	4.0
One repetition	0.29	2.5	0.19	2.0
Two repetitions	0.05	–	0.02	–

Appendix

Table 5: KL divergence (

\times 10^{-4}

) from sequential decoding to each training procedure, averaged over the 47,933 response tokens of 32 fine-tuning validation examples, for the HLA models of the main experiments (

W{=}128

). Two repetitions were measured at 32-bit precision only.

	Key rank $r_{k}$	Value rank $r_{v}$	Loop-1 rank $r_{1}$
Rank	KL, loops 2–4	Output error	KL, loop 1
128	1.04	0.42	1.04
256	0.51	0.37	0.44
512	0.29	0.30	0.23
1,024	0.11	0.22	0.08
2,048	0.03	0.16	–

Appendix

Table 6: Error of the initialized latent on Ouro-1.4B when one rank varies and the others stay at

r_{k}{=}r_{v}{=}512

r_{1}{=}256

. KL: divergence between the attention of HLA and of the full cache, in nats, averaged over the loops the latent serves. Output error: relative squared error of the attention output, averaged over loops 2–4.

$r_{k}$	$r_{v}$	$r_{1}$	Cache per token (KiB)	MATH500
Without a window
256	256	256	48	51.4
512	512	512	96	66.8
1,024	512	256	96	66.6
1,024	1,024	256	120	70.2
With a window of 32 tokens

Appendix

Table 7: MATH500 accuracy (%) of Ouro-1.4B after the first uptraining phase for different latent sizes, with one sample per problem.

		Concurrent sequences				Throughput (tokens/s)				Latency (ms/token)
Model	Context	Full	LLA	HLA	Ratio	Full	LLA	HLA	Ratio	Full	LLA	HLA
Ouro-1.4B	1K	68	264	298	$4.4\times$	1,305	93.6	3,231	$2.5\times$	13.6	28.5	17.1
	4K	20	77	139	$7.0\times$	373	27.8	1,429	$3.8\times$	20.6	54.1	20.8
	8K	10	40	84	$8.4\times$	192	14.3	797	$4.2\times$	30.1	90.1	24.2
	16K	5	19	40	$8.0\times$	97	7.2	428	$4.4\times$	49.3	158.7	25.1
Ouro-2.6B	1K	32	123	127	$4.0\times$	519	46.3	1,278	$2.5\times$	27.0	56.8	34.3

Appendix

Table 8: Values of Fig. 2 , measured with the same settings. LLA uses its reconstruct path with a latent of rank 256 per head, 96 for the keys and 160 for the values; since the codec does not affect the timing, it is a per-head PCA codec for Ouro-1.4B and a random one for Ouro-2.6B. Ratios compare HLA with the full cache; the best value of each group is in bold.

Explore similar work

Closing the Loop: Practical Training Recipes for Looped Language Models
Sep 30, 2026Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin +5Recurrent Model Loop
RND NLP, DAIMLD, Russian Federation
Learning When to Attend: Conditional Memory Access for Long-Context LLMs
Mar 18, 2026Sakshi Choudhary, Aditya Chattopadhyay, Luca Zancato +4Efficient Long-Context Inference Time-To-First-Token
Department of Electrical and Computer Engineering, Purdue University, USA · AWS Agentic AI, USA
LT2: Linear-Time Looped Transformers
May 20, 2026Chunyuan Deng, Yizhe Zhang, Rui-Jie Zhu +4Kimi Delta Attention Transformer Architectures
1Rice University · 2Apple · 3UC Santa Cruz +1

Sep 30, 2026, cs.CLJ/K move · Enter open · S save

Closing the Loop: Practical Training Recipes for Looped Language Models

Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin, Inessa Fedorova+4

RND NLP, DAIMLD, Russian Federation

Recurrent Model Loop Multi-Stage Training Dense Baseline Recipe

Open paper AskPDF arXiv

Looped language models increase effective depth by repeatedly applying a shared block of layers, but existing large-scale recipes require multi-stage training over trillions of tokens, while the benefits of recurrence remain difficult to separate from differences in data and training. In this work, we establish practical training recipes for looped language models, with three main results. (1) We develop a compute-efficient from-scratch pipeline that reduces the training budget from 7.7T tokens in Ouro to 310B tokens while retaining strong reasoning performance. Pretraining followed by high-quality mid-training, together with learning-rate warmup and stronger exit-gate regularization, enables stable recurrent training without prior multi-stage schedules. (2) Under controlled comparisons, our 1.4B LoopLM outperforms a parameter-matched dense model trained on the same data and token budget on all 12 evaluated benchmarks, including +14 points on GSM8K, +10 on MATH, and +22 on DROP. At matched inference compute, it approaches a 3.9B dense model on mathematical reasoning and reading comprehension while using only 36% as many parameters. (3) We introduce a minimal recipe for converting pretrained dense models into looped ones: a single learned input-mixing scalar and a smoothed exit loss, with no step-specific parameters. Applied to Qwen3-1.7B-Base, Looped Qwen improves over an identically continued dense baseline on every evaluated benchmark across two data regimes, with statistically clear gains on GSM8K, MATH, and MMLU-Pro on the curated mixture. Together, these results make looped language models substantially cheaper to train from scratch and practical to introduce into existing pretrained checkpoints, while isolating the gains due to recurrence itself.

Abstract

Explore similar work

Closing the Loop: Practical Training Recipes for Looped Language Models

Learning When to Attend: Conditional Memory Access for Long-Context LLMs

LT2: Linear-Time Looped Transformers