Enabling Dynamic Computation in Looped LMs

Summary (Overview)

• Problem: Looped Language Models (LMs) promise parameter-efficient dynamic computation via early-exit, but state-of-the-art models (Ouro) fail to realize practical efficiency gains due to rigid KV-cache requirements at each loop depth and ineffective early-exit training objectives.

• Key Contribution 1 – Ragged KV: A novel "best-available" KV caching strategy that enables true early-exit by aliasing deeper blocks to the best-available cache, reducing FLOPs and KV memory by up to 30% while maintaining full-depth performance—and it works out-of-the-box on Ouro models without retraining.

• Key Contribution 2 – Global Prior Enforcement: A simple fix to Ouro's exit-depth regularization objective that replaces per-token uniform prior enforcement with a global prior constraint, enabling truly heterogeneous, difficulty-aware exit depths across tokens.

• Key Contribution 3 – Ragged-Aware Training: An alternate-step training procedure that exposes the model to Ragged KV during training at no additional FLOPs cost, significantly improving performance in dynamic computation mode.

• Validation: Results are demonstrated on Ouro models (1.4B–2.6B, including thinking variants) across multiple benchmarks (MMLU, BBH, GSM8K, AIME, GPQA, HumanEval+), plus smaller 160M models pre-trained from scratch.

---

Introduction and Theoretical Foundation

Background

Looped LMs reuse a block of transformer layers multiple times. A standard LM applies a block $\mathcal{B}_\theta$ once; a looped LM reuses it for $U$ recurrent steps:

$$H^{(0)} = \mathrm{Embed}(x_{1:n}), \qquad H^{(d)} = \mathcal{B}_\theta\Big(H^{(d-1)}\Big), \quad d = 1, \ldots, U. \tag{1}$$

This yields the computational depth of an $L \times U$ layer model with only $L$ layers of parameters. The two publicly available competitive looped models are Ouro ($U=4$) and Nanbeige ($U=2$), though Nanbeige lacks early-exit capability and was excluded from the study.

Early-Exit Mechanism

Ouro adds a learned gate $g_\phi$ at each recurrent depth, predicting the probability of stopping token $i$ at depth $d$:

$$\lambda_{i,d} = \sigma\left(g_\phi(h_i^{(d)})\right). \tag{2}$$

This defines a distribution over exit depths:

$$\pi_i(d) = \begin{cases} \lambda_{i,d} \prod_{r < d} (1 - \lambda_{i,r}), & d < U, \\ \prod_{r < U} (1 - \lambda_{i,r}), & d = U, \end{cases} \tag{3}$$

At inference, a deterministic threshold $\tau \in [0,1]$ determines exit:

$$e_i(\tau) = \min\left\{d: \sum_{r=1}^{d} \pi_i(r) \geq \tau \right\}. \tag{4}$$

Two Key Problems

P1: The Cache Problem. Ouro's KV cache is rigidly fixed to individual loop iterations. Since all loops require their own cache level, deep computation is forced for all tokens regardless of exit depth, negating practical utility.

P2: Which Tokens Should Exit Early? Ouro's training objective applies entropy regularization equally to all tokens, forcing the exit gate to behave similarly across all tokens regardless of difficulty—defeating the purpose of a difficulty-aware exit gate.

---

Methodology

Ragged KV Cache Strategy

The core insight is that since exit depths are the model's own estimate of maturation, each token's exit-depth KV entries can be treated as final. Future tokens needing deeper KV entries can use this aliased version:

$$h_{i,\ell}^{(d)} = \mathcal{T}_{\theta_\ell}\left(h_{i,\ell-1}^{(d)}; K_{<i,\ell}^{(\mathbf{b})}, V_{<i,\ell}^{(\mathbf{b})}\right); \quad K_{<i,\ell}^{(\mathbf{b})} = \left[K_{1,\ell}^{(b_1)}, \dots, K_{i-1,\ell}^{(b_{i-1})}\right] \quad \text{where} \quad b_i = \min(d, e_i). \tag{7}$$

This produces a ragged pattern of physical KV entries across tokens. Memory savings are realized via page tables (as in vLLM), and FLOPs savings via depth-batching.

Two inference modes:

• Dense-Ragged: Prompt prefilled at full-depth ($U=4$), ragged KV used only during decode
• ETE-Ragged: Ragged KV applied end-to-end (including prefill)

Global Prior Enforcement Fix

Ouro's original objective enforces a uniform prior on each token's exit distribution:

$$\mathcal{L}_{\text{ouro}} = \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{d=1}^{U}\pi_i(d)\ell_{i,d} + \beta \cdot \mathrm{KL}(\pi_i \| \pi^p)\right) \quad \text{where} \quad \pi^p = \pi^{\mathrm{u}}. \tag{9}$$

The proposed fix enforces the KL term on the global mean distribution:

$$\boxed{\mathcal{L}_{\text{ouro}} = \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{d=1}^{U}\pi_i(d)\ell_{i,d}\right) + \beta \cdot \mathrm{KL}(\bar{\pi} \| \pi^p)} \quad \text{where} \quad \bar{\pi} = \frac{1}{N}\sum_{i=1}^{N}\pi_i. \tag{10}$$

This relaxes the unnecessary constraint on individual tokens, allowing difficulty-based depth assignment while maintaining global distribution similarity to the prior.

Ragged-Aware Training

Alternate-step training: alternate between dense steps (following Eq. 9/10) and Ragged KV forward passes (sampling online from current $\pi$), back-propagating:

$$\mathcal{L}_{\text{rag}} = \frac{1}{N}\sum_{i=1}^{N}\ell_{i,b}$$

This requires no additional training FLOPs (only one logits calculation vs. four in dense steps).

Model Architecture

---

Empirical Validation / Results

Ragged KV on Ouro Models (Out-of-the-Box)

Key finding: Ouro models work with Ragged KV without any calibration, achieving comparable performance to full-depth decode at 60–70% of the average exit depth.

Ouro-1.4B Performance Highlights:

Ouro-2.6B Thinking Highlights (macro averages across AIME24/25/26, GPQA-Diamond, OlympiadBench, HumanEval+):

Critical insight: Ragged KV enables ~10 percentage point improvement over fixed-depth 2 at only slightly higher average exit depth, demonstrating that selective deeper processing of ~20% of tokens significantly improves overall performance.

Why Ragged KV Works: Analysis of Exit Behavior

Analysis of exit-depth distributions reveals that Ouro models predominantly exit at depth 2 or 3 (depending on τ), with:

• Almost no depth-1 exits
• Most tokens continuing at their predecessor's exit depth
• Large jumps ($|e_i - e_{i-1}| \geq 2$) essentially non-existent

This shows Ouro behaves like a fixed-lower-depth processing model, not a truly dynamic one—which is why Ragged KV aligns well with its training.

Local vs. Global Prior Enforcement (160M Models)

Local enforcement produces:

• τ-sensitive exit depths (behaves like fixed-depth 3 or 4 depending on threshold)
• Sudden flips between depth 4 and 3 during training
• No true dynamic computation

Global enforcement produces:

• Peaky, confident exit distributions per token
• τ-insensitive by default
• Nontrivial mass at lower depths (1, 2)
• Healthy proportion of large jumps ($|e_i - e_{i-1}| \geq 2$)
• Difficulty-aware exits (low-depth tokens are actually easy, e.g., commas, spaces)

Ragged-Aware Training Results

• Global + RA training: Significant drop in ragged-PPL, becoming better than fixed@4-PPL (since forcing all 4 depths is suboptimal for dynamic computation)
• Local + RA training: Minimal impact (already behaves like fixed-depth model, little KV mismatch)
• RA training also slightly reduces average exit depth, even without KL terms in Ragged steps

---

Theoretical and Practical Implications

Theoretical Contributions

1. Formalization of the KV mismatch problem: The paper rigorously defines why naive early-exit fails in looped LMs—Eq. 6 shows that deeper computations are forced because later tokens require KV entries from earlier tokens at all depths.

2. Understanding of regularization effects: The analysis reveals that per-token KL regularization with a uniform prior creates an information-theoretic constraint that prevents difficulty-aware computation, essentially forcing all tokens toward the same exit distribution regardless of context.

3. Training-inference mismatch: The work demonstrates that dense full-depth training induces a representation mismatch when using dynamic exits, and provides a principled fix via ragged-aware training.

Practical Implications

• Deployment efficiency: Enables up to 30% reduction in FLOPs and KV memory while maintaining full-depth performance, making looped LMs viable for edge devices and resource-constrained environments.

• True dynamic computation: The global prior enforcement fix enables genuinely heterogeneous compute allocation—easy tokens (punctuation, common words) exit early while difficult tokens (reasoning-heavy) receive deeper processing.

• Zero-cost inference improvement: Ragged KV works out-of-the-box on existing Ouro models without retraining, providing immediate efficiency gains.

• Training efficiency: Ragged-aware training adds no additional FLOPs cost while improving both performance and average exit depth.

---

Conclusion

This work addresses two fundamental barriers to practical dynamic computation in looped LMs:

1. Ragged KV solves the cache problem, enabling true early-exit with real memory and FLOPs savings (up to 30%) while maintaining or slightly improving performance. It works out-of-the-box on Ouro models and creates a new frontier in the performance-vs-depth space.

2. Global prior enforcement fixes the ill-specified early-exit objective that caused Ouro to behave like a static lower-depth model. This unlocks genuinely heterogeneous, difficulty-aware exits.

3. Ragged-aware training resolves the training-inference KV mismatch at no additional training cost, further improving dynamic computation performance.

Limitations and Future Directions

• Scale: Validation at 160M scale for pre-training experiments; applicability to larger scale models remains open
• Implementation: vLLM implementation not optimal; throughput in ragged mode can be further improved
• Open questions: Mechanistic aspects of prefill vs. decode ragged, the 2.6B Ouro's buggy ETE-Ragged performance, and aggressive threshold schedule optimization

The work breaks down existing barriers in making looped LMs a practical inference efficiency solution for LLMs, with exciting avenues for future research in larger-scale validation and optimized implementations.