Universal YOCO for Efficient Depth Scaling
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
| Model | KV Cache Memory | Prefilling Time | Decoding Time |
| Transformer | |||
| YOCO | |||
| Loop / Universal Transformer | |||
| YOCO-U |
| 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 |
| Model | Wiki. | LMB. | LMB. | PIQA | OBQA | Hella. | Wino. | ARC-E | ARC-C | Avg. |
| ppl | ppl | acc | acc | acc | acc_n | acc | acc_n | acc_n | acc | |
| 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 |
| 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 |
| Model | Wiki. | LMB. | LMB. | PIQA | OBQA | Hella. | Wino. | ARC-E | ARC-C | Avg. |
| ppl | ppl | acc | acc | acc | acc_n | acc | acc_n | acc_n | acc | |
| 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 |
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 |
| Size | Hidden Dim. | #Layers | #Heads | LR |
| 300M | 1536 | 12 | 12 | |
| 500M | 2048 | 12 | 16 | |
| 870M | 2304 | 16 | 18 | |
| 1.3B | 2560 | 20 | 20 | |
| 3.4B | 4096 | 20 | 32 | |
| 6.4B | 5120 | 24 | 40 |
| 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 |
| 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 |
| 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 |
Efficient Language Model Inference
Efficient language model inference reduces the compute, memory, latency, and energy needed to generate useful outputs.
See all 314 papers