Universal YOCO for Efficient Depth Scaling

Yutao Sun; Li Dong; Tianzhu Ye; Shaohan Huang; Jianyong Wang; Furu Wei

cs.CLApr 1, 2026

Universal YOCO for Efficient Depth Scaling

Authors: Yutao Sun, Li Dong, Tianzhu Ye, Shaohan Huang, Jianyong Wang, Furu Wei

Organizations: Microsoft Research · Tsinghua University

Abstract

The rise of test-time scaling has remarkably boosted the reasoning and agentic proficiency of Large Language Models (LLMs). Yet, standard Transformers struggle to scale inference-time compute efficiently, as conventional looping strategies suffer from high computational overhead and a KV cache that inflates alongside model depth. We present Universal YOCO (YOCO-U), which combines the YOCO decoder-decoder architecture with recursive computation to achieve a synergistic effect greater than either alone. Built on the YOCO framework, YOCO-U implements a Universal Self-Decoder that performs multiple iterations via parameter sharing, while confining the iterative process to shallow, efficient-attention layers. This combination yields a favorable capability-efficiency tradeoff that neither YOCO nor recursion achieves independently. The YOCO architecture provides a constant global KV cache and linear pre-filling, while partial recursion enhances representational depth with limited overhead. Together, YOCO-U improves token utility and scaling behavior while maintaining efficient inference. Empirical results confirm that YOCO-U remains highly competitive in general and long-context benchmarks, demonstrating that the integration of efficient-attention architectures and recursive computation is a promising direction for scalable LLMs.

Figures & tables

Figure 1: Overview of the YOCO-U architecture. Universal Self-Decoder (bottom) performs recursive computation (

T

iterations, indicated by the red dashed line) using efficient self-attention (local windows) to refine representations. Crucially, the global KV cache produced for cross-attention is generated once and remains constant regardless of

T

, while only the local window-based KV caches grow with iterations, keeping overall cache overhead negligible. Cross-Decoder (top) reuses this shared global KV cache via cross-attention for autoregressive token prediction.

Model	KV Cache Memory	Prefilling Time	Decoding Time
Transformer	$\mathcal{O}(LND)$	$\mathcal{O}(LN^{2}D)$	$\mathcal{O}(LND)$
YOCO	$\mathcal{O}((N+WL)D)$	$\mathcal{O}(\frac{L}{2}ND)$	$\mathcal{O}(\frac{L}{2}(N+W)D)$
Loop / Universal Transformer	$\mathcal{O}(LTND)$	$\mathcal{O}(LTN^{2}D)$	$\mathcal{O}(LTND)$
YOCO-U	$\mathcal{O}((N+WTL)D)$	$\mathcal{O}(\frac{L}{2}TND)$	$\mathcal{O}(\frac{L}{2}(N+WT)D)$

Table 1: Inference complexity comparison.

N,L,D,T

denote sequence length, number of layers, hidden dimension, and loop iterations.

W

is the local window size in efficient self-attention. In YOCO-U, the global KV cache (

\mathcal{O}(ND)

) is independent of

T

, and only the local efficient-attention cache (

\mathcal{O}(WTLD)

) scales with iterations. Since

W\ll N

, the additional overhead is negligible for long sequences.

Figure 2: Scaling behavior of language modeling loss with the same model size. (Left) Loss versus training FLOPs: YOCO-U achieves competitive or lower loss (

\Delta L{=}0.033

) at the same FLOPs budget, while incurring negligible KV cache overhead. (Right) Loss versus training tokens: YOCO-U also improves data efficiency, requiring approximately 62% fewer tokens to reach comparable performance.

Model	ARC-C	Winogrande	HellaSwag	MMLU	BBH	GSM8K	Humaneval	DROP	Average
YOCO	46.50	61.72	63.44	49.59	33.13	38.06	9.15	32.62	41.78
YOCO-U (Equal FLOPS)	47.87	68.67	66.80	54.63	35.49	50.49	10.98	34.94	46.23
YOCO-U (Equal Steps)	48.72	69.85	67.12	55.63	36.31	50.57	10.37	38.07	47.08

Table 2: Performance comparison of recursive and non-recursive YOCO. YOCO-U converts additional compute in recursive shallow blocks into consistent gains across tasks, with negligible KV cache overhead.

Figure 3: Accuracy comparison on 11 math benchmarks. YOCO-U consistently outperforms the YOCO baseline across all tasks, achieving a significant boost in average accuracy.

Model	Wiki.	LMB.	LMB.	PIQA	OBQA	Hella.	Wino.	ARC-E	ARC-C	Avg.
	ppl $\downarrow$	ppl $\downarrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc_n $\uparrow$	acc $\uparrow$
Non-Recursive
Transformer	22.52	22.26	38.4	69.6	22.6	45.7	57.1	59.6	36.6	47.1
YOCO	22.25	18.30	41.2	67.9	23.8	45.6	54.3	59.2	36.6	47.0
Recursive
Loop/Universal TRM	21.56	22.56	37.7	69.0	23.6	47.7	56.3	62.4	38.1	47.8

Table 3: Performance comparison among different architectures. RINS is a standard decoder-only Transformer with early-layer recursion. YOCO-U achieves comparable results to RINS while maintaining significantly lower KV cache and inference cost.

Figure 4: Long sequence perplexity decreases along with the increasing input length on Book (top) and Code (bottom) data. YOCO-U consistently achieves lower perplexity compared to non-recursive baselines, i.e., Transformer and YOCO. YOCO-U also maintains parity with the heavier recursive baseline (RINS), indicating effective utilization of long-range context.

Model Name	S-NIAH-1	S-NIAH-2
Transformer	0.87	0.82
YOCO	1.00	0.86
RINS	0.99	0.91
YOCO-U	1.00	0.95

Table 4: Accuracy on the Needle In-A-Haystack [ 17 ] test.

N

represents the number of needles.

Model	Wiki.	LMB.	LMB.	PIQA	OBQA	Hella.	Wino.	ARC-E	ARC-C	Avg.
	ppl $\downarrow$	ppl $\downarrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc_n $\uparrow$	acc $\uparrow$
Non-Recursive
YOCO	22.25	18.30	41.16	67.90	23.80	45.64	54.30	59.22	36.60	46.95
Deep (Instead of Wide)	22.04	21.76	37.67	68.50	23.80	45.84	56.12	61.03	35.15	46.87
Recursive
YOCO-U	21.01	18.32	41.18	68.66	24.60	48.89	55.33	62.16	36.95	48.25

Table 5: Performance comparison among different design choices. “Deep (Instead of Wide)”: double depth while using the same model size. “Upper Loop”: loop cross-decoder instead of self-decoder. “Upper Loop w/o Shared KV”: loop cross-decoder without using shared KV cache, i.e., employing self attention for loop cross-decoder.

Figure 5: Parameter scaling properties. We keep the training steps the same. Left: YOCO-U achieves comparable performance with 50% fewer parameters than YOCO. Right: YOCO-U demonstrates scalable parameter utility as the activated parameter count increases.

Figure 6: Loop scaling property. YOCO-U shows consistent improvements when scaling up loop numbers (

1,2,3,5

Panel of Figure 7 — Figure 7: Overall efficiency comparison across varying sequence lengths. The models have the same number of parameters, where YOCO-U achieves better performance. YOCO-U achieves competitive prefilling and decoding throughput (i.e., tokens/second), as well as KV cache memory consumption, compared with Transformer and vanilla loop Transformer (RINS). The curves of YOCO-U and YOCO in Figure 7(c) are overlapped.

Figure 8: Angular distances between consecutive layers.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Params	Values
Layers	20
Hidden Size	2560
Expert Number	64
Expert Topk	8
Expert FFN Size	1024
Shared Expert FFN Size	1024

Appendix

Table 6: Hyperparamters used for the model in Section 4.1 .

Size	Hidden Dim.	#Layers	#Heads	LR
300M	1536	12	12	$6\times 10^{-4}$
500M	2048	12	16	$6\times 10^{-4}$
870M	2304	16	18	$3\times 10^{-4}$
1.3B	2560	20	20	$3\times 10^{-4}$
3.4B	4096	20	32	$3\times 10^{-4}$
6.4B	5120	24	40	$3\times 10^{-4}$

Appendix

Table 7: Model size and hyperparameters used in Section 4.4 .

Model	8K	16K	32K	64K	128K	256K
Transformer	85707	66342	45566	27276	14987	7475
YOCO	220662	219734	218247	219106	220416	220407
RINS	42905	33276	22767	13630	7397	3739
YOCO-U	75637	75694	76202	76148	75534	76301

Appendix

Table 8: Prefill Throughput (tokens/s) across different sequence lengths.

Model	8K	16K	32K	64K	128K	256K
Transformer	2712	1795	638	450	220	137
YOCO	3356	2539	1722	975	556	318
RINS	1582	580	234	118	96	56
YOCO-U	2410	1966	1419	865	514	303

Appendix

Table 9: Decode Throughput (tokens/s) across different sequence lengths.

Model	8K	16K	32K	64K	128K	256K
Transformer	320	640	1280	2560	5120	10240
RINS	640	1280	2560	5120	10240	20480
YOCO	26	42	74	138	266	522
YOCO-U	46	62	94	158	286	542

Appendix

Table 10: KV Cache Occupancy (MB) across different sequence lengths.

Efficient Language Model Inference

Efficient language model inference reduces the compute, memory, latency, and energy needed to generate useful outputs.

See all 317 papers

Hidden Decoding at Scale: Latent Computation Scaling for Large Language Models
Jul 9, 2026Aiwei Liu, Cheng Shi, Chuhan Wu +45Language Model Scaling Laws Transformer Attention
UniScale: Adaptive Unified Inference Scaling via Online Joint Optimization of Model Routing and Test-Time Scaling
May 29, 2026Kaiyu Huang, Xingyu Wang, Mingze Kong +6Cost-Aware Inference Inference-Time Scaling
School of Computer Science and Technology, Tongji University · Shenzhen Research Institute of Big Data, The Chinese University of Hong Kong, Shenzhen · School of Data Science, The Chinese University of Hong Kong, Shenzhen +2
YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference
Apr 15, 2026You Wu, Ziheng Chen, Yizhen Zhang +6LLM Inference Acceleration KV-Cache Compression
ShanghaiTech University · Alibaba Group · Zhejiang University

Jul 9, 2026J/K move · Enter open · S save

Hidden Decoding at Scale: Latent Computation Scaling for Large Language Models

Aiwei Liu, Cheng Shi, Chuhan Wu, Ci Lei+44

Language Model Scaling Laws Transformer Attention Language Model Scaling LLM Training

Open paper AskPDF arXiv

Scaling Large Language Models (LLMs) has been driven mainly by enlarging the Transformer backbone, but for an already-strong model this requires another round of costly pretraining. We study whether an existing backbone can keep improving by allocating more computation to each token while leaving the Transformer backbone fixed. Depth-recurrent (looped) Transformers pursue this goal but are hard to scale, because looped computation does not fit naturally with the pipeline parallelism used to train the largest models. We add computation along the sequence-length dimension, where the extra computation is simply a longer input and stays compatible with standard large-model training. We propose Hidden Decoding, a sequence-length scaling method applied during continued pretraining (CPT). It expands each token into n streams with independent embedding tables and keeps the intermediate streams' key-value cache as context, so each token performs more internal computation without adding or widening Transformer layers. To keep this affordable at scale, we introduce Stream-Factorized Attention, in which most layers attend only within each stream and only a few layers mix across streams, reducing the attention cost from quadratic to roughly linear in n. Experiments support two scaling results. At frontier scale, we train WeLM-HD4-80B and WeLM-HD4-617B at n=4 and improve their matched non-HD baselines, making Hidden Decoding the first demonstrated sequence-length scaling method at the 100B+ MoE scale. Across expansion factors, the gains grow as n increases, showing that sequence-length expansion is a practical fixed-backbone scaling path for frontier-scale LLMs.

Abstract

Efficient Language Model Inference

Explore similar work

Hidden Decoding at Scale: Latent Computation Scaling for Large Language Models

UniScale: Adaptive Unified Inference Scaling via Online Joint Optimization of Model Routing and Test-Time Scaling

YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference