Universal YOCO: Recursive Depth Scaling with Constant Memory

Paper: 2604.01220 Authors: Yutao Sun, Li Dong, Tianzhu Ye, Shaohan Huang, Jianyong Wang, Furu Wei Categories: cs.CL

The Gap

Test-time scaling has become the new frontier for LLM reasoning, but standard Transformers hit a wall: every additional reasoning step inflates the KV cache proportionally to model depth, and conventional looping strategies carry massive computational overhead. Prior work splits into two camps: YOCO architecture achieves efficient attention but lacks depth flexibility, while recursive models gain depth through iteration but suffer from quadratic attention costs and exploding memory. Neither approach delivers both efficient inference and scalable depth.

Problem: Test-time scaling needs depth, but depth kills efficiency
   |
   v
Assumption: Architecture + recursion can be synergistic, not additive
   |
   v
Method: YOCO-U = YOCO decoder-decoder + recursive shallow layers
   |
   v
Evidence: Constant KV cache + competitive benchmarks + better scaling
   |
   v
Conclusion: Efficient-attention + partial recursion > either alone

The Increment

One sentence: Before this paper, you chose between efficient inference (YOCO) or flexible depth (recursion); after, you get both through architectural synergy.

Core Mechanism

YOCO-U builds on the YOCO decoder-decoder architecture, which splits the model into a Self-Decoder (processes input with self-attention) and a Cross-Decoder (generates output attending to Self-Decoder states). The key innovation: make the Self-Decoder “Universal” by having it loop over itself multiple times through parameter sharing, but only in the shallow efficient-attention layers.

Here’s the data flow: input tokens first pass through the Self-Decoder, which now runs for N iterations instead of one. Each iteration refines the representation using the same parameters. Critically, the KV cache from the Self-Decoder stays constant across iterations because YOCO’s architecture already confines global attention to the Cross-Decoder. After N iterations, the enriched representations feed into the Cross-Decoder for final output generation.

The recursion happens only in the Self-Decoder’s efficient-attention layers, not the full stack. This “partial recursion” gives you depth (N iterations = N× effective layers) without the memory explosion of full-model recursion. The Cross-Decoder remains non-recursive, maintaining linear pre-filling and constant global KV cache.

Input tokens
    |
    v
+-------------------+
| Self-Decoder      |<----+
| (efficient attn)  |     |
+-------------------+     | N iterations
    |                     | (parameter sharing)
    +---------------------+
    |
    v (constant KV cache)
+-------------------+
| Cross-Decoder     |
| (global attn)     |
+-------------------+
    |
    v
Output tokens

Think of YOCO-U like a pottery wheel. The Self-Decoder is the wheel that spins the clay (input representations) multiple times—same wheel, same hands, but each rotation refines the shape. The Cross-Decoder is the kiln that fires the final piece. Traditional Transformers are like hand-building pottery: every layer is a new tool, a new workspace, and you need to store every intermediate form. YOCO’s original design was like a wheel that only spins once—efficient but limited refinement. YOCO-U lets the wheel spin N times without needing N different wheels or N different storage shelves for intermediate forms. The magic: the wheel (Self-Decoder parameters) stays the same, and you only store the final refined clay state before firing, not every rotation’s snapshot.

Key Concepts

• Decoder-Decoder Architecture (YOCO): Standard Transformers use encoder-decoder or decoder-only designs. YOCO introduces decoder-decoder: the first decoder (Self-Decoder) processes input with efficient attention (like linear attention or local windows), the second decoder (Cross-Decoder) generates output with full global attention to Self-Decoder states. Why this matters: the Self-Decoder’s KV cache doesn’t grow with output length because it only processes input once. The Cross-Decoder’s KV cache is constant because it only attends to fixed Self-Decoder states, not the growing output sequence. Concrete example: translating a 1000-token document. Self-Decoder processes all 1000 tokens once with efficient attention (small cache). Cross-Decoder generates translation attending to those 1000 compressed states (constant cache), not the growing translation itself.

• Partial Recursion: Instead of making the entire model recursive (which multiplies memory by iteration count), YOCO-U only loops the Self-Decoder’s efficient-attention layers. The Cross-Decoder and any global-attention components remain non-recursive. This asymmetry is crucial: you get depth scaling (N iterations = deeper representation) in the cheap part of the model, while keeping the expensive global-attention part single-pass. Think of it like image compression: you can iteratively refine a low-resolution preview (cheap) before doing one high-resolution render (expensive), rather than iteratively rendering at full resolution.

• Parameter Sharing in Recursion: When the Self-Decoder loops N times, it uses the same weights each iteration—this is parameter sharing. It’s not N different layers; it’s one layer applied N times. The benefit: model size stays constant while effective depth grows. The tradeoff: each iteration sees the same transformation, so you’re betting that iterative refinement with fixed parameters beats having N distinct transformations. Analogy: it’s like editing a photo with the same filter applied multiple times (parameter sharing) versus applying N different filters (distinct layers). Sometimes repeated application of one good filter (sharpen, sharpen, sharpen) works better than a sequence of mediocre filters.

Framework Shift

Before (Standard Transformer):       After (YOCO-U):

Input                                Input
  |                                    |
  v                                    v
[Layer 1] ---> KV cache 1          [Self-Decoder]<--+
  |                                    |   ^         |
  v                                    |   |         | N loops
[Layer 2] ---> KV cache 2              |   +---------+ (same params)
  |                                    |
  v                                    v (constant cache)
[Layer N] ---> KV cache N          [Cross-Decoder]
  |                                    |
  v                                    v
Output                               Output

Memory: O(N * L)                   Memory: O(1)
Depth: Fixed N layers              Depth: Flexible N iterations

From fixed-depth sequential layers with growing memory to flexible-depth recursive refinement with constant memory, the core shift is trading architectural complexity for temporal complexity.

Expert Assessment

Problem choice: Real gap. Test-time scaling is the current battleground, and the memory-depth tradeoff is a genuine bottleneck. The problem sits at the intersection of two active research threads (efficient attention and recursive models), which is a sweet spot.

Method maturity: Clever architectural insight, not brute force. The synergy between YOCO’s constant-cache property and partial recursion is non-obvious. However, the paper doesn’t deeply explore why parameter sharing works well here—feels like an empirical bet that paid off rather than a principled design. A simpler approach might be progressive layer dropping or early exiting, but those don’t address the depth-scaling angle.

Experimental integrity: Baselines are fair (standard Transformers, original YOCO, other efficient-attention models). Numbers look solid across general and long-context benchmarks. One yellow flag: the paper doesn’t show failure modes or tasks where YOCO-U underperforms. The “highly competitive” framing suggests it’s not dominating everywhere, but we don’t see where it struggles. Would benefit from ablations on iteration count vs performance saturation.

Writing quality: The abstract and intro are crisp. The method section could use a clearer walk-through of the forward pass with concrete tensor shapes. The related work section is thorough but reads like a literature dump—cutting it by 30% would sharpen the narrative. The results section is competent but doesn’t dig into why YOCO-U wins on specific tasks.

Verdict: weak accept — Solid architectural contribution with clear practical benefits, but lacks depth in understanding why the approach works and where it fails.

Takeaways

Practitioners can steal the “partial recursion” pattern: identify the cheap operations in your model (efficient attention, local processing, compression) and make only those recursive while keeping expensive operations (global attention, generation) single-pass. This asymmetry is broadly applicable beyond LLMs—think video processing (iteratively refine low-res frames, render high-res once) or hierarchical planning (iterate on abstract plans, execute concrete actions once).

The decoder-decoder split is also transferable: if your model has a “processing” phase and a “generation” phase, consider making their memory footprints independent. The processing phase can use efficient attention with constant cache, while generation attends to processed states rather than raw input.

Finally, the paper validates that parameter sharing in recursion isn’t just a memory hack—it can improve token utility and scaling behavior. If you’re building models where depth matters but memory is constrained, don’t default to stacking more layers; try looping fewer layers with shared parameters.