Universal YOCO (YOCO-U): Efficient Depth Scaling

• The paper introduces Universal YOCO (YOCO-U), which leverages recursive parameter sharing in its self-decoder to efficiently scale inference depth while maintaining a constant global KV cache.
• YOCO-U confines recursion to shallow, efficient-attention layers, allowing the Cross-Decoder to perform global retrieval with linear prefilling and reduced memory overhead.
• Empirical results demonstrate that YOCO-U achieves lower validation loss, improved long-context performance, and significant speedups compared to standard Transformer models.

Universal YOCO (YOCO-U) is a decoder-only language-model architecture for efficient depth scaling at test time. It extends the YOCO (“You Only Cache Once”) decoder-decoder design by replacing the static Self-Decoder with a recursive Universal Self-Decoder, so that effective depth can be increased through parameter sharing while the model retains a constant global KV cache, linear prefilling, and efficient long-context inference (Sun et al., 1 Apr 2026). In this formulation, recursion is confined to shallow, efficient-attention layers, while the globally attentive Cross-Decoder remains non-recursive and continues to read from a single shared KV cache generated by the Self-Decoder, preserving the core invariants of the original YOCO architecture (Sun et al., 2024).

1. Scope, terminology, and lineage

YOCO-U belongs to the large-language-model line of work built around YOCO, “You Only Cache Once,” a decoder-decoder architecture introduced for autoregressive language modeling (Sun et al., 2024). Its immediate purpose is efficient test-time depth scaling: increasing inference-time computation by recursive reuse of a shallow block, rather than by looping a full Transformer stack whose KV cache and latency scale directly with depth (Sun et al., 1 Apr 2026).

The acronym “YOCO” is overloaded across several unrelated literatures. In medical imaging, YOCO in MediAug denotes a single-image spatial augmentation strategy that cuts an image into regions, applies different augmentations to each region, and reassembles the result without label mixing (Qi et al., 26 Apr 2025). In autonomous driving, YoCo refers to a weakly supervised 3D instance segmentation framework using single bird’s-eye-view clicks, vision foundation models, temporal-spatial label updating, and IoU-guided enhancement (Jiang et al., 27 Feb 2025). In planetary landing, YOCO abbreviates “You Only Crash Once,” a YOLOv3-based hazardous-terrain detector with unsupervised domain adaptation and visual-similarity clustering (Jr et al., 2023). Universal YOCO, by contrast, denotes specifically the efficient depth-scaling LLM architecture introduced in 2026 (Sun et al., 1 Apr 2026).

Historically, YOCO-U is best understood as an overview of two architectural ideas. From the original YOCO it inherits the decoder-decoder split, the single global KV cache, and the early-exit prefilling property (Sun et al., 2024). From the Universal Transformer tradition it inherits recursive parameter sharing, but only within the Self-Decoder, thereby avoiding the cost profile of full-stack looping (Sun et al., 1 Apr 2026).

2. Architectural formulation

Given token embeddings $X^0 \in \mathbb{R}^{|x|\times d_{\text{model}}}$, baseline YOCO divides a decoder-only model into two halves. The first half is a Self-Decoder with efficient self-attention; the second half is a Cross-Decoder that performs global cross-attention to a shared KV cache produced by the Self-Decoder. Concretely, YOCO computes

$H=\operatorname{Self\mbox{-}Decoder}(X^0),$

then forms a single global cache

$$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$

and applies Cross-Decoder layers whose queries are layer-specific but whose keys and values are always $(\hat K,\hat V)$ (Sun et al., 2024).

YOCO-U changes only the Self-Decoder. It introduces a Universal Self-Decoder (USD),

$\operatorname{USD}(X)=\underbrace{\operatorname{Self\mbox{-}Decoder}^{\,L/2}\circ\cdots\circ \operatorname{Self\mbox{-}Decoder}^{\,L/2}}_{T\ \text{iterations}}(X),$

so that the same shallow Self-Decoder block is applied $T$ times with shared parameters (Sun et al., 1 Apr 2026). The Cross-Decoder is then run once, using the global KV cache generated from the recursively refined hidden states. This yields an effective Self-Decoder depth of $T\cdot \frac{L}{2}$ without increasing parameter count.

Each Self-Decoder layer in YOCO-U has the form

$$Y^{l}=\operatorname{ESA}(\operatorname{LN}(X^{l})) + X^{l},\qquad X^{l+1}=\operatorname{SwiGLU}(\operatorname{LN}(Y^{l})) + Y^{l},$$

where $\operatorname{ESA}$ is an efficient self-attention module and SwiGLU is used for the feed-forward sublayer (Sun et al., 1 Apr 2026). The default ESA in the reported experiments is sliding-window attention, but the architecture is described as agnostic to the specific efficient-attention choice, with SWA, RetNet-style retention, Mamba SSMs, and gated DeltaNet cited as admissible examples (Sun et al., 1 Apr 2026).

The design is therefore “partial recursion in shallow, efficient-attention layers.” The recursion does not encompass the entire network. The heavy global-attention block is not looped; only the cheaper Self-Decoder is iterated. This structural asymmetry is the central technical move in YOCO-U.

3. KV cache structure and complexity

YOCO-U is motivated by the observation that naïvely looping a standard Transformer stack multiplies not only compute but also KV cache memory. For a standard decoder with sequence length $N$, total depth $H=\operatorname{Self\mbox{-}Decoder}(X^0),$0, hidden dimension $H=\operatorname{Self\mbox{-}Decoder}(X^0),$1, and loop count $H=\operatorname{Self\mbox{-}Decoder}(X^0),$2, full-stack recursion yields KV memory $H=\operatorname{Self\mbox{-}Decoder}(X^0),$3, prefill cost $H=\operatorname{Self\mbox{-}Decoder}(X^0),$4, and decoding cost $H=\operatorname{Self\mbox{-}Decoder}(X^0),$5. YOCO-U changes this scaling by keeping a single global KV cache for the Cross-Decoder and allowing only small efficient-attention states in the recursive Self-Decoder (Sun et al., 1 Apr 2026).

Architecture	KV cache memory	Prefill / decoding complexity
Standard Transformer	$H=\operatorname{Self\mbox{-}Decoder}(X^0),$6	$H=\operatorname{Self\mbox{-}Decoder}(X^0),$7 / $H=\operatorname{Self\mbox{-}Decoder}(X^0),$8
Loop / Universal Transformer	$H=\operatorname{Self\mbox{-}Decoder}(X^0),$9	$$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$0 / $$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$1
YOCO	$$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$2	$$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$3 / $$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$4
YOCO-U	$$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$5	$$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$6 / $$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$7

The crucial point is that the global KV cache in YOCO-U is constant with respect to recursive depth $$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$8. The $$\hat K = \operatorname{LN}(H) W_K,\qquad \hat V = \operatorname{LN}(H) W_V,$$9-dependent term appears only in the local efficient-attention states, represented as $(\hat K,\hat V)$0, and the paper characterizes this overhead as negligible for long contexts because $(\hat K,\hat V)$1 in the intended regime (Sun et al., 1 Apr 2026). A plausible implication is that recursion becomes primarily a compute knob rather than a memory knob.

This property is tightly linked to the original YOCO design. In the 2024 architecture, the Self-Decoder already used bounded-memory kernels such as sliding-window attention or gated retention, while the Cross-Decoder read from one global cache rather than maintaining per-layer KV states (Sun et al., 2024). YOCO-U preserves that interface and injects additional representational depth only where the cache discipline remains favorable.

4. Empirical behavior and scaling results

The main YOCO-U experiments include a 10B total-parameter MoE configuration with 1.3B activated parameters, 20 layers split into 10 Self-Decoder and 10 Cross-Decoder layers, 64 experts with 8 active plus 1 shared, sequence length 8192, SWA window 512, and training up to 300B tokens (Sun et al., 1 Apr 2026). In this setting, YOCO-U at equal FLOPs achieves lower validation loss than YOCO by $(\hat K,\hat V)$2. At equal training tokens, YOCO-U trained on 80B tokens is reported as comparable to non-recursive YOCO trained on 210B tokens, which the paper describes as roughly 62% fewer tokens to reach similar loss (Sun et al., 1 Apr 2026).

On LM Evaluation Harness benchmarks, the average score is 41.78 for YOCO, 46.23 for YOCO-U at equal FLOPs, and 47.08 for YOCO-U at equal steps (Sun et al., 1 Apr 2026). The gains are broad rather than isolated. For example, GSM8K rises from 38.06 to 50.49 and MMLU from 49.59 to 54.63 when comparing YOCO to YOCO-U at equal FLOPs. After additional supervised fine-tuning on math reasoning, YOCO-U outperforms YOCO on all 11 reported math benchmarks, with an average accuracy improvement of $(\hat K,\hat V)$3 percentage points (Sun et al., 1 Apr 2026).

The architecture comparison on dense 1.3B models trained on 20B tokens is also informative. Reported average accuracy is 47.1 for a standard Transformer, 47.0 for YOCO, 47.8 for a looped Universal Transformer, 46.8 for ParScale, 48.3 for RINS, and 48.3 for YOCO-U (Sun et al., 1 Apr 2026). This supports the paper’s claim that recursive compute is more productively spent in early or shallow blocks than in the full stack, and that efficient-attention recursion can match the task performance of recursive full-attention alternatives.

Long-context behavior is a central part of the evidence. On Needle-In-A-Haystack retrieval, YOCO-U records 1.00 accuracy for one needle and 0.95 for two needles, compared with 0.87 and 0.82 for a standard Transformer, and 0.99 and 0.91 for RINS (Sun et al., 1 Apr 2026). Long-context perplexity on books and repository-level code improves relative to non-recursive YOCO and standard Transformer as prefix length grows. The paper therefore presents YOCO-U not as an efficiency-only model, but as one that preserves or improves long-context utilization.

The systems results follow the theoretical picture. On H100-80GB with nano-vLLM, batch size 32, generating 128 tokens, prefill throughput at 16K context is 66,342 tokens/s for Transformer, 219,734 for YOCO, 33,276 for RINS, and 75,694 for YOCO-U (Sun et al., 1 Apr 2026). At 256K context, YOCO-U is reported as 10× faster than Transformer and 20× faster than RINS in prefill. Decoding throughput at 16K is 1,795 tokens/s for Transformer, 2,539 for YOCO, 580 for RINS, and 1,966 for YOCO-U; at 256K it is 137, 318, 56, and 303 respectively. KV cache memory per sequence at 256K is 10,240 MB for Transformer, 20,480 MB for RINS, 522 MB for YOCO, and 542 MB for YOCO-U (Sun et al., 1 Apr 2026). The empirical conclusion is that YOCO-U stays close to YOCO’s memory profile while adding a recursive compute path that is much cheaper than looping a full-attention model.

5. Comparative position within the YOCO family

Within the language-model lineage, YOCO-U is best viewed as an extension of the original YOCO rather than a replacement for it. The 2024 YOCO paper established three architectural claims: a self-decoder/cross-decoder split, a single global KV cache reused across the upper half of the network, and prefilling that can early-exit after the self-decoder without changing the final output (Sun et al., 2024). YOCO-U preserves all three, but repurposes the Self-Decoder as a recurrent computational substrate (Sun et al., 1 Apr 2026).

The paper’s own ablations indicate that loop placement matters. Looping the Self-Decoder gives an average accuracy of 48.25, whereas looping the Cross-Decoder gives 47.34, and looping the Cross-Decoder without shared KV gives 46.41 (Sun et al., 1 Apr 2026). This is presented as evidence that deeper-layer recursion exhibits diminishing returns and that the shared-KV regime is part of the reason the shallow-loop design is effective. In the paper’s framing, the Self-Decoder behaves like an iterative “encoding/thinking” block, while the Cross-Decoder performs global retrieval and final decoding.

YOCO-U is distinct from YOCO++, which addresses a different efficiency problem. YOCO++ targets cross-layer KV compression and introduces weighted KV residual connections from the bottom layer to each bottom-half layer, preserving a 50% KV cache compression rate while improving performance among cross-layer KV compression methods (Wu et al., 15 Apr 2026). It reports a zero-shot average of 48.99, higher than the standard Transformer’s 48.37 and YOCO’s 47.98 at the reported 1.1B scale, but its objective is not test-time depth scaling via recursive efficient-attention layers (Wu et al., 15 Apr 2026). In other words, YOCO-U and YOCO++ are adjacent but not interchangeable: one is primarily a recursion-and-scaling architecture, the other a cross-layer KV-compression refinement.

A broader misconception is that “YOCO” denotes a single unified family across all fields. It does not. The LLM architecture, the medical-image augmentation operator, the weakly supervised LiDAR segmentation framework, and the hazardous-terrain detector share an acronym but not a methodology (Qi et al., 26 Apr 2025).

6. Limitations, design implications, and prospective extensions

YOCO-U’s reported strengths come with explicit architectural commitments. Training with $(\hat K,\hat V)$4 increases Self-Decoder compute, so training is heavier than non-recursive YOCO even if token efficiency improves at scale (Sun et al., 1 Apr 2026). The decoder-decoder split is structural rather than cosmetic: the paper states that migrating an existing model to YOCO-U is non-trivial and that one effectively must build and pretrain the architecture from scratch, or at minimum undertake substantial finetuning (Sun et al., 1 Apr 2026). The main tunable design variables are the Self-/Cross-Decoder split, the efficient-attention mechanism, the window size $(\hat K,\hat V)$5, and the loop count $(\hat K,\hat V)$6.

The ablation literature in the paper suggests that recursion is a smooth scaling knob. For dense 1.3B models trained on 20B tokens, increasing $(\hat K,\hat V)$7 from 1 to 2 to 3 to 5 yields monotonic improvements in validation loss (Sun et al., 1 Apr 2026). This suggests a compute-quality trade-off analogous to external test-time scaling, but internalized into the backbone. The paper explicitly characterizes latent recursive reasoning and explicit test-time scaling methods such as chain-of-thought or multi-step reasoning as orthogonal and complementary.

Several future directions are described or directly implied. YOCO-U is presented as compatible with adaptive recursion, mixture-of-recurrences-style scheduling, and multimodal or audio extensions (Sun et al., 1 Apr 2026). The original YOCO paper had already argued that the decoder-decoder layout is general to multiple self-decoders and that cross-attention is natural for multimodal fusion, particularly in streaming settings (Sun et al., 2024). This suggests that “universal” in YOCO-U is partly architectural rather than merely nominal: the universalized component is a reusable recursive Self-Decoder that can, in principle, sit beneath a broader family of efficient cross-decoding systems.

Taken together, the available evidence positions Universal YOCO as a specific answer to a narrow but consequential systems problem: how to increase inference-time depth without incurring the KV-cache growth and quadratic prefilling costs of full-stack looping. Its contribution is not the invention of recursion, nor the invention of shared KV, but their combination in a decoder-decoder architecture that preserves long-context capability while making depth scaling computationally usable at inference time (Sun et al., 1 Apr 2026).