# Enabling Dynamic Computation in Looped LMs

> Ragged KV caching enables true early-exit in looped language models, cutting FLOPs and KV memory by up to 30% without retraining or performance loss.

- **Source:** [arXiv](https://arxiv.org/abs/2610.09013)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/ZYcSo7
- **Whiteboard:** https://picx.dev/p/ZYcSo7/image

## Summary

# 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

| Model | Params | $L$ | $d_{model}$ | FFN | $H_q/H_{kv}$ | $D$ | $U$ | vocab |
|-------|--------|-----|-------------|-----|--------------|-----|-----|-------|
| Ouro-160M | 160.47M | 12 | 768 | 2048 | 12/12 | 64 | 4 | 49,152 |
| Ouro-1.4B | 1.44B | 24 | 2048 | 5632 | 16/16 | 128 | 4 | 49,152 |

---

## 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:**

| Mode | Avg. Depth | MMLU | BBH pass@1 | GSM8K pass@1 | Macro pass@1 |
|------|-----------|------|------------|--------------|--------------|
| Fixed@2 | 2.00 | 60.4 | 32.0 | 58.8 | 35.8 |
| Fixed@3 | 3.00 | 66.7 | 62.5 | 74.2 | 60.5 |
| Fixed@4 | 4.00 | 67.5 | 69.0 | 76.6 | 64.2 |
| Dense-Ragged τ=0.3 | 2.39 | — | 61.6 | 72.4 | 59.5 |
| Dense-Ragged τ=0.5 | 3.05 | — | 67.5 | 75.3 | 62.9 |
| ETE-Ragged τ=0.3 | 2.63 | 61.9 | 59.8 | 72.0 | 58.0 |
| ETE-Ragged τ=0.5 | 3.21 | 66.2 | 66.4 | 75.5 | 62.4 |

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

| Mode | Avg. Depth | Macro pass@1 | Macro pass@4 |
|------|-----------|--------------|--------------|
| Fixed@2 | 2.00 | 48.9 | 67.1 |
| Fixed@3 | 3.00 | 61.8 | 77.5 |
| Fixed@4 | 4.00 | 62.1 | 78.0 |
| Dense-Ragged τ=0.5 | 2.83 | **62.3** | **79.4** |
| ETE-Ragged τ=0.5 | 3.00 | 62.0 | 79.4 |

**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.

---

_Markdown view of https://picx.dev/p/ZYcSo7, served by PicX — AI-generated visual whiteboard summaries of research papers._