Universal YOCO for Efficient Depth Scaling

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

Source-linked AI summary

Universal YOCO for Efficient Depth Scaling

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

arXiv:2604.01220v1Apr 1, 2026cs.CL

TL;DR

Standard Transformer looping makes inference-time depth scaling expensive because computation and KV-cache memory grow with depth. YOCO-U combines YOCO’s shared-cache decoder architecture with parameter-shared recursion restricted to shallow efficient-attention layers, and evaluations report competitive or better benchmark performance while preserving efficient serving properties.

Problem
Standard Transformer looping can extend computational depth but incurs high computational cost and KV-cache memory that grows with depth.
Method
YOCO-U replaces YOCO’s Self-Decoder with a parameter-shared Universal Self-Decoder that recursively iterates shallow efficient-attention layers while the Cross-Decoder reuses shared global KV caches.
Results
YOCO-U achieves competitive or better performance across general and long-context benchmarks at favorable efficiency, with negligible KV-cache overhead and retained linear pre-filling.
Takeaways & Limitations
YOCO-U supports scalable depth and improved parameter utility without increasing KV-cache footprint, offering a cost-effective architecture for large language models.
Takeaways & Limitations
Experiments default to looping the Self-Decoder three times, producing 2× the total FLOPs of the non-recursive baseline.

Abstract

from arXiv · show

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.

1 Introduction

YOCO-U addresses the inefficiency of recursive computation in standard Transformers by combining YOCO’s shared-cache decoder architecture with parameter-shared recursion in shallow efficient-attention layers. This design improves capability and scaling while preserving efficient inference and long-context advantages.

Test-time scaling increases LLM reasoning and agentic capabilities, making efficient inference-time computation increasingly important.
Standard Transformer looping extends computational depth but incurs high computation and KV-cache memory that grows with depth.
YOCO-U combines YOCO’s Self-Decoder/Cross-Decoder design with recursive computation, producing benefits beyond either technique alone.
A Universal Self-Decoder reuses parameters across iterations, increasing representational depth without increasing parameter count.
YOCO-U confines recursion to shallow efficient-attention modules, avoiding full-attention looping costs while preserving a one-piece KV cache and long-context capability.
Across benchmarks, YOCO-U improves token utility and matches or exceeds performance at equal FLOPs with negligible KV-cache overhead, while retaining linear pre-filling and reduced memory use.

2 Related Work

Prior work scales computation through parameter sharing, selective depth scaling, and inference-time reasoning, but these approaches raise efficiency or scope considerations for scalable model computation.

Universal Transformer shares parameters across depth to reduce parameter costs and enhance model capacity.
Selective depth scaling can improve model performance, but depth scaling typically increases latency and KV-cache requirements.
Inference scaling improves complex reasoning, while its gains derive from capacity established during pre-training rather than directly improving pre-training.

3 Method

YOCO-U recursively applies a parameter-shared Universal Self-Decoder to shallow efficient-attention layers, then uses a Cross-Decoder with shared global KV caches for prediction. This layout preserves YOCO’s efficient serving properties while limiting recursive overhead.

YOCO-U Architecture: YOCO-U splits autoregressive modeling into a Self-Decoder and Cross-Decoder, replacing the non-recursive Self-Decoder with a Universal Self-Decoder iterated for T steps.
YOCO-U Architecture: Recursion is restricted to a subset of shallow layers, preserving linear-complexity pre-filling and a one-piece KV cache.
Universal Self-Decoder: The Universal Self-Decoder recursively refines representations and generates the global KV caches used by the Cross-Decoder.
Efficient Self-Attention: Efficient self-attention provides O(1) inference memory in the self-decoder layers, with sliding-window attention used as the default implementation.
Recursive Computation: Parameter sharing across recursive iterations enhances representation capability without adding parameters, while shallow recursion improves computational efficiency.
Cross-Decoder: The Cross-Decoder reuses the global KV caches across layers for attention-based next-token prediction.
Inference Advantages: YOCO-U materializes the KV cache once, yielding O(N) memory rather than O(NL) and preserving long-context serving under a fixed GPU-memory budget.
Inference Advantages: Unlike standard Transformer looping, YOCO-U leaves global attention unchanged and adds only the small overhead of efficient-attention recursion.

4 Experiments

YOCO-U combines recursive computation with efficient-attention shallow blocks to improve capability and token utility while preserving efficient inference. Across language modeling, reasoning, long-context, scaling, and efficiency evaluations, it delivers consistent gains with limited KV-cache and throughput overhead.

Setup: YOCO-U uses three Self-Decoder loops by default, with 2× the total FLOPs of the non-recursive baseline.The training setup uses 10B total parameters and 1.3B activated parameters, with 20 layers and sliding-window attention in the Self-Decoder.
Token Scaling: YOCO-U achieves lower loss at matched training FLOPs and matches non-recursive YOCO trained with 210B tokens using 80B tokens.The matched-compute loss improvement is ∆L=0.033.
Task Evaluation: YOCO-U surpasses its baseline by +4.45 average under equal FLOPs, while math accuracy improves 24.4% on average across 11 benchmarks.The math evaluation reports gains over YOCO on every benchmark.
Architecture Comparison: Recursive scaling outperforms parallel scaling under equal FLOPs, while efficient-attention recursion performs comparably to RINS without proportional computational or KV-cache growth.Scaling bottom blocks is more effective than scaling all blocks, and scaling efficient-attention blocks is reported as equally effective to standard early-layer recursion.
Long-Context Behavior: YOCO-U matches RINS on long-context utilization, maintains robust needle retrieval, and shows diminishing returns as recursive iterations approach a representation fixed point.Angular-distance patterns remain consistent across loops, while mean distance gradually decreases with more iterations.
Inference Efficiency: YOCO-U improves inference efficiency with 10× Transformer and 20× RINS prefilling throughput, 2.21× Transformer decoding acceleration at 256K, and negligible additional KV-cache cost.At 16K context, prefilling is 1.15× faster and decoding is 1.1× faster than non-recursive Transformer; RINS requires 38× more KV-cache memory than YOCO-U.

5 Conclusion

YOCO-U is a recursive architecture that decouples depth scaling from memory overhead. It improves parameter utility and scaling behavior while preserving linear pre-filling efficiency.

YOCO-U decouples depth scaling from memory overhead through parameter sharing and recursive computation.
YOCO-U significantly enhances parameter utility and exhibits superior scaling behaviors.
YOCO-U outperforms baselines across various benchmarks while preserving linear pre-filling efficiency.

A Hyperparameters for Section 4.1

The Section 4.1 hyperparameters distinguish non-recursive YOCO from YOCO-U by their loop counts. Table 6 provides the detailed hyperparameters for these models.

Non-recursive YOCO uses Loop = 1, whereas YOCO-U uses Loop = 3.
Table 6 presents detailed hyperparameters for the Section 4.1 models.

B Thinking SFT Evaluation Details for Section 4.1

The Thinking SFT evaluation measures mathematical reasoning with a consistent prompt format and decoding strategy. Evaluations use greedy decoding and allow up to 32,768 generated tokens.

YOCO-U’s mathematical reasoning is evaluated on standard benchmarks after Thinking SFT.
Prompt Template: The evaluation prompt asks the model to solve each math problem, explain its reasoning, and place the final answer in boxed form.
Decoding Strategy: Greedy decoding is used across reasoning benchmarks to support reproducibility and focus on the model’s primary reasoning path.
Decoding Strategy: The maximum generation length is 32,768 tokens to accommodate extended thinking and detailed solutions.

C Hyperparameters for Section 4.4

Section 4.4 specifies model dimensions, optimization settings, attention configuration, and training duration. Table 7 reports model sizes and associated hyperparameters.

Training uses 8192-token sequences, a 1.0M-token batch size, AdamW, β1 = 0.9, β2 = 0.95, and weight decay 0.1.
The Self-Decoder contains half the total layers by default and uses sliding-window attention with size 512.
The FFN size is 3 × dmodel, where dmodel denotes the hidden dimension.
Models are trained for 20k steps, corresponding to 20B tokens.
Table 7 reports hidden dimensions, layer counts, and head counts for different model sizes.

D Inference Details for Section 4.5

Inference-efficiency benchmarks evaluate deployment performance across sequence lengths using standardized GPU, framework, precision, and memory settings. The reported tables cover prefill throughput, decode throughput, and KV-cache occupancy.

Benchmark setup: Inference benchmarks use NVIDIA H100-80GB GPUs, Nano-vLLM, and BF16 precision to measure practical deployment performance.Models use Paged Attention with block size 256, while GPU memory utilization is fixed at 0.7.
Measured quantities: Prefill throughput is reported in tokens/s across different sequence lengths.
Measured quantities: Decode throughput is reported in tokens/s across different sequence lengths.

E Experimental Details for Section 4.2

The experiments standardize model size, training budget, and compute across recursive and non-recursive baselines, then evaluate them on language-modeling and reasoning benchmarks. The setup also reports KV-cache occupancy and uses explicit benchmark metrics.

Model and training setup: All baseline models use 1.3B parameters, 20 layers, and hidden dimension 2560 under approximately 2× standard-Transformer FLOPs.Each model is trained for 20k steps with a 1M batch size, totaling 20B tokens.
Baseline configurations: Universal Transformer loops the full 20-layer block twice, producing 40 layers of computation while preserving the parameter count.
Efficiency evaluation: KV-cache occupancy is measured in MB across different sequence lengths.
Baseline configurations: RINS combines a 10-layer non-recursive block with a 10-layer recurrent block looped three times, also yielding 40 computational layers.This layout aligns with YOCO-U for direct comparison of recursive strategies.
Evaluation benchmarks: Evaluation covers WikiText-103, LAMBADA, PIQA, OpenBookQA, HellaSwag, Winogrande, and ARC-E/ARC-C.These tasks span long-form language modeling, long-range prediction, commonsense, fact-based, continuation, coreference, and science-question evaluation.
Evaluation metrics: Perplexity measures next-token uncertainty, while accuracy and normalized accuracy evaluate prediction correctness, with normalized accuracy adjusting for answer length.