Source-linked AI summary
Enabling Dynamic Computation in Looped LMs
Aayush Mishra, Arnau Padrés Masdemont, Victor Conchello Vendrell, Jordi Ros Giralt, Arash Behboodi, Fabio Valerio Massoli
TL;DR
Looped LMs promise parameter-efficient dynamic computation, but rigid KV caches and equal per-token exit-prior enforcement limit practical difficulty-aware early exit. The paper introduces Ragged KV and global prior enforcement, reporting up to 30% lower average exit depth while maintaining or slightly improving performance, and notes that larger-scale applicability remains open.
Problem
Method
Results
Takeaways & Limitations
Ragged KV makes early exit practical on Ouro without additional training, while global prior enforcement supports heterogeneous exits based on context difficulty.
Source: §1 · p. 2The experiments do not establish applicability to larger-scale models because compute constraints prevented LLM-scale pre-training.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Looped LMs are parameter efficient and promise dynamic computation (saving memory and FLOPs on easy tokens). However, state-of-the-art open Looped LMs trained with this dynamic computation capability (Ouro models) do not realize it in practice as each loop iteration (depth) requires its own level of KV-cache, necessitating all loop computations. Moreover, Ouro's early-exit prior is enforced on each token equally, which results in static lower-depth like processing of all tokens regardless of difficulty. In this work, we propose a simple "best-available" KV caching strategy that works out-of-the-box, creating a new frontier in the performance vs depth space. Our approach enables up to 30% reduction in FLOPs and KV memory while retaining full-depth performance, showing the true flexibility of Looped LMs. Furthermore, training looped LMs with awareness about this KV caching strategy improves performance and efficiency. Finally, we apply a small but effective fix to the early-exit prior enforcement objective that makes tokens exit at truly heterogeneous depths based on effort. Our findings are validated on Ouro models as well as smaller looped LMs pre-trained from scratch.
1 INTRODUCTION
The introduction identifies two obstacles to difficulty-aware computation in Looped LMs: rigid per-depth KV caches force full computation, while per-token prior enforcement produces similar exit behavior across tokens. It proposes Ragged KV and global prior enforcement to make early exit practical and difficulty-aware.
The cache problem
- Rigid per-depth KV caches force deep computation for every token, negating early exit’s practical utility.Using a fixed lower loop count for all tokens hurts performance and lacks difficulty-aware flexibility.Source: §1 · p. 1
Which tokens should exit early?
- Equal per-token prior regularization makes exit behavior similar across tokens regardless of context, undermining difficulty-aware exits.The prior is intended to distribute exit depth across loop iterations.Source: §1 · p. 1
S1: Ragged KV
- On BigBenchHard, Ragged KV is compared with fixed-depth decoding across loop counts and exit thresholds, with higher performance at low average decode depths.The figure caption describes fixed-depth settings of 1–4 loops and Ragged KV at different thresholds.Source: Figure 1 · p. 2
- Ragged KV reduces Ouro models’ average exit depth by up to 30% while maintaining or slightly improving performance.It works out-of-the-box, without additional training; Ragged-KV-aware training is proposed to address the training–inference cache mismatch.Source: §1 · p. 2
S2: Difficulty-aware exits
- Global prior enforcement enables exit depths to vary with context difficulty while keeping the overall distribution close to the prior.Unlike equal per-token enforcement, it applies the prior globally.Source: §1 · p. 2
2 BACKGROUND
The background describes looped models as reusing a transformer block across recurrent steps and outlines Ouro’s token exit mechanism. The paper limits its study to publicly available Ouro models with early-exit capability.
Looped LMs
- A looped LM reuses the same transformer block for U recurrent steps, while a standard LM applies each block once.Looping gives an L-layer block the computational depth of an L × U-layer model while retaining L layers of parameters.Source: §2 · p. 2
Early-Exit
- Ouro adds a learned gate at each recurrent depth to predict each token’s probability of stopping there.The language-model head can read out a next-token distribution from any recurrent depth.Source: §2 · p. 2
- At inference, Ouro exits a token at the first depth where cumulative exit probability crosses threshold τ; changing τ changes average loop applications.The final depth receives the remaining probability mass to form a discrete exit-depth distribution.Source: §2 · pp. 2–3
Evaluation scope
- The study is restricted to publicly available Ouro models equipped with early-exit capability.Nanbeige-4.2-3B lacks early exit, and fixed-depth-1 generations became degenerate with performance collapsing to zero.Source: §2 · p. 3
3 RAGGED KV: A SOLUTION TO OURO’S KV CACHE PROBLEM
Ragged KV replaces depth-specific cache rigidity with best-available cached entries, letting tokens stop early while later tokens retain needed deeper-depth entries. On Ouro, it supports full-depth-comparable performance with lower compute and memory.
3 RAGGED KV: A SOLUTION TO OURO’S KV CACHE PROBLEM
- Final-depth-only KV caching saves memory but not FLOPs, and performance drops to 0 on many benchmarks, so it does not work universally.The paper reports confirming the benchmark failure.Source: §3 · p. 3
- Ragged KV aliases each token’s exit-depth cache as final for later tokens needing deeper entries, avoiding forced full-depth computations while retaining causal depth-local access.Ouro otherwise forces all tokens through U = 4 depth iterations, with each level computing its own KV entries.Source: §3 · p. 3
- Using only 75% of FLOPs+memory, Dense-Ragged and ETE-Ragged perform on par with full-depth.The reported results include Ouro-1.4B and Ouro-2.6B Thinking.Source: Table 1 · p. 4, §3 · p. 4, Table 2 · p. 5
- At 60–70% average exit depth, Ouro-1.4B and 2.6B non-thinking models in Dense-Ragged achieve performance comparable to full-depth decode on BigBenchHard.The paper reports this frontier from its exit-threshold sweep.Source: §3 · p. 4
- Most tokens exit at depth 2 or 3 and continue at their predecessor’s depth; 20% of decode tokens go deeper than 2, improving pass@(1, 4) by 10pp.The reported depth pattern undermines the claim that Ouro’s uniform prior produces strongly difficulty-dependent computation.Source: §3 · p. 5
4 A SIMPLE FIX TO OURO’S EXIT DEPTH REGULARIZATION PROBLEM
The section argues that enforcing an exit-depth prior independently on every token makes exits similarly high-entropy, limiting difficulty-aware computation. It proposes applying the prior to the global exit distribution instead, then evaluates this change in smaller Ouro models trained from scratch.
4 A SIMPLE FIX TO OURO’S EXIT DEPTH REGULARIZATION PROBLEM
- Local prior enforcement constrains each token toward a similar high-entropy exit distribution, regardless of difficulty.The section links this constraint to optimizing all depths for all tokens rather than dynamically allocating computation.Source: §4 · p. 6, Figure 2 · p. 6
- Global prior enforcement applies the KL term to the mean exit distribution, allowing peaky per-token exits while keeping the aggregate close to the prior.The mean is practically constrained to the training batch.Source: §4 · p. 6, Figure 3 · p. 7
4.1 PRETRAINING 160M OURO MODELS FROM SCRATCH
- The authors pre-trained scaled-down Ouro models from scratch on 20B FineWeb-edu tokens, using the same data and optimizer settings for local and global modes.They tested Uniform, Geometric, and Poisson priors; a conservative β schedule meant the models did not fully enforce the prior by the end of pre-training.Source: §4.1 · pp. 6–7
- Global enforcement produces heterogeneous, difficulty-aware exits, while local enforcement behaves like fixed-depth 3/4 and shifts token exits together.In the reported evaluation, global assigns nontrivial mass to depths 1 and 2 and has large jumps; local is τ-sensitive, while global is τ-insensitive.Source: §4.1 · p. 8
- Global training’s more dynamic exits incur a ragged-perplexity penalty relative to fixed@4, whereas local training stays closer to fixed@4 because it behaves like a fixed-lower-depth model.The section identifies ragged-aware training as a way to address the global model’s penalty.Source: §4.1 · p. 8
5 RAGGED AWARE TRAINING
Ragged-aware training targets the representation mismatch caused by Ragged KV’s changed post-exit cache entries. The authors interleave dense and Ragged KV training steps, improving ragged performance for global models.
5 RAGGED AWARE TRAINING
- Alternate-step training interleaves a dense step with a Ragged KV forward pass that samples from the current exit distribution and backpropagates exit-depth prediction loss.Dense steps align language-modeling and prior-enforcement objectives; Ragged steps compute logits once rather than four times for the expected LM loss.Source: §5 · p. 8
- Ragged-aware training significantly improves global models’ ragged perplexity, surpassing fixed@4 performance while slightly reducing average exit depth.The reported training does not materially affect local models, which behave like fixed-lower-depth models and have less KV mismatch.Source: §5 · p. 8
- The authors report that alternate-step training adds no training FLOPs, while counterfactual and REINFORCE alternatives cost more and perform no better.The simple alternate-step method is therefore used as the default.Source: §5 · pp. 8–9
6 RELATED WORK
Related work traces looped models from weight-shared recurrent transformers to parameter-efficient modern language models, alongside methods for dynamic inference. The section positions Ragged KV near prior KV-selection and early-exit approaches.
6 RELATED WORK
- Universal Transformers combined Transformer parallelization with RNN expressivity through weight sharing and supported dynamic computation; later work applied looping to modern decoder-only LLMs.The section cites Dehghani et al. (2018) and Saunshi et al. (2025).Source: §6 · p. 9
- Recent looped-model work reports parameter efficiency and latent-reasoning advantages, while Ouro competed with larger Qwen, Gemma, and Llama models.The section also reports Nanbeige’s 4B model, with 3B looped parameters, beating Qwen3.5-9B on intelligence index.Source: §6 · p. 9
- Dynamic-inference approaches include speculative decoding, layer skipping, and routing; the section says Ouro early exits remained ineffective because of KV mismatch.Ragged KV is described as inspired by Fu et al. (2025) and closest to AdaPonderLM, whose torch.where KV selection yields only logical memory and FLOPs savings.Source: §6 · p. 9
7 LIMITATIONS AND CONCLUSION
The paper presents Ragged KV and an early-exit objective fix as steps toward practical dynamic computation, while noting scale and implementation limits. Its systems implementation uses compact page aliases and mechanisms to avoid later work for exited tokens.
7 LIMITATIONS AND CONCLUSION
- Ragged KV and the objective fix target practically useful early exits; the authors report savings up to 30% and leave larger-scale applicability open.They also identify a non-optimal vLLM implementation and possible throughput improvements as future work.Source: §7 · p. 9
A SYSTEMS IMPLEMENTATION OF RAGGED KV
- Ragged KV needs compact aliases to save memory and removal of exited tokens from later iterations to save FLOPs.The implementation uses compact aliases in native vLLM and depth-batched prototypes for token removal.Source: A SYSTEMS IMPLEMENTATION OF RAGGED KV · p. 13
A.1 PAGE-TABLE ALIASES FOR COMPACT K/V STORAGE
- After exit depth ei, later reads reuse the terminal page without copying its K/V entry, while attention computation remains unchanged.The scheduler retains pages through ei, returns deeper pages to the allocator, and blocks writes through terminal aliases.Source: §A.1 · p. 13
A.2 DEPTH BATCHING FOR INCREASED THROUGHPUT
Depth batching groups work by recurrent depth so exited tokens need not continue through later passes. The current nested scheduler reduces work relative to full depth, but its throughput depends on batch size and remains short of the work-conserving lower bound.
A.2 DEPTH BATCHING FOR INCREASED THROUGHPUT
- Depth queues batch entries requiring the same recurrent-block application; after a pass, continuing tokens advance to the next queue and exited tokens restart at Q1.The final depth is forced to exit, and autoregressive order within each request is preserved.Source: §A.2 · pp. 13–14
- The reported ratio counts useful model FLOPs, but real throughput also depends on microbatch width and launch cost.The current nested runner can still perform unnecessary work inside a step.Source: §A.2 · pp. 13–14
- 64.0% of dense cache was retained; the nested runner executed 823 depth passes versus 574.2 ideal, and beat full-depth 1,020.For matched eager H100 runs, Ragged KV was faster at batch sizes 1, 4, and 8, but slower at batch size 32.Source: Table 3 · p. 14, §A.2 · p. 14
B DETAILS OF OURO-160M PRE-TRAINING EXPERIMENTS
The pretraining experiments use a fixed 20B-token corpus and compare loop-normalization choices and exit-depth priors. A slowly increasing β schedule avoids the collapse seen with stronger early prior enforcement, though the prior remains under-enforced at training end.
B DETAILS OF OURO-160M PRE-TRAINING EXPERIMENTS
- Training uses 20B FineWeb-Edu tokens in randomized 2,048-token windows, with an approximately 5M-token validation subset for plots.The batch size is 524,288 tokens, and the learning rate follows warmup and cosine decay.Source: B DETAILS OF OURO-160M PRE-TRAINING… · pp. 14–15
- Using raw loop outputs as the next-loop input learned faster than normalizing each loop’s last hidden representation, so the experiments removed per-loop normalization.This was the main architecture difference from original Ouro reported in the passage.Source: B DETAILS OF OURO-160M PRE-TRAINING… · p. 15
- A flat β = 0.1 schedule caused fixed-depth-1-like behavior, with language modeling in the first loop and near-identity behavior in later loops.The authors report that an earlier ramping schedule also collapsed when β first increased.Source: B DETAILS OF OURO-160M PRE-TRAINING… · p. 15, Figure 7 · p. 15
- The selected schedule raises β from 0 to 0.01 over 0–5B tokens, to 0.02 over 5B–10B, and to 0.05 over 10B–20B.It is intended to let parameters mature for looped application before stronger prior enforcement.Source: B DETAILS OF OURO-160M PRE-TRAINING… · p. 15
- The authors chose conservative β values to protect language-modeling capability, so the prior was not fully enforced by the end of 20B training tokens.They note that the schedule could be improved.Source: B DETAILS OF OURO-160M PRE-TRAINING… · p. 15
- For U = 4, the experiments compare a uniform prior with mean depth 2.5/4 and geometric and Poisson priors with mean depths 2.07/4 and 2.16/4.The geometric prior uses λ = 0.3; the Poisson prior uses λ = 2.1.Source: B DETAILS OF OURO-160M PRE-TRAINING… · p. 16
B.1 COMPLETE LOCAL VS GLOBA L RESULTS
Global prior enforcement develops heterogeneous exits across tokens, while local enforcement tends toward concentrated exit depths. The reported comparisons show that global exits vary across a prompt and are qualitatively associated with token difficulty.
B.1 COMPLETE LOCAL VS GLOBA L RESULTS
- Global exit fractions evolve gradually toward lower depths, while local enforcement can flip depth-4 tokens to depth 3 and behave like a fixed-depth-3 model.The local flip also appears earlier at lower exit thresholds; Figure 11 indicates another possible flip between depths 2 and 3 under stronger or longer enforcement.Source: §B.1 · p. 16, Figure 11 · p. 17
- Across the three priors, global enforcement mixes hard exit depths within a prompt, whereas paired local enforcement remains concentrated around one depth.The authors qualitatively describe global low-depth assignments as meaningfully related to difficulty.Source: §B.1 · p. 16
- 19.2%, 20.9%, and 13.2% global large-jump rates for Poisson, Geometric, and Uniform priors exceed paired local rates of 0.2%, 2.5%, and 1.9%.A large jump is an adjacent hard-depth change of at least two.Source: Figure 12 · p. 18
- Ragged-aware alternate-step training branches greatly improve ragged performance for global enforcement, but not as much for local enforcement.Figure 13 compares training progression, average exit depth, and depth fractions across the three priors.Source: Figure 13 · p. 19
C RAGGED-AWARE TRAINING OBJECTIVES
The section compares training objectives that expose looped models to Ragged-KV cache states, finding direct alternate-step training improves Ragged performance while retaining standard training cost.
Direct alternate-step training (RA)
- Ragged-aware alternate-step training greatly improves performance across priors, especially in global mode.Each optimizer step alternates the dense objective and a sampled Ragged-KV objective; ragged steps contain no KL term.Source: C RAGGED-AWARE TRAINING OBJECTIVES · p. 19
- RA alternates dense and sampled-ragged objectives, detaching sampled exits so ragged steps train on inference-time cache states while dense steps enforce the prior.Each step runs one branch, preserving standard dense training cost.Source: C RAGGED-AWARE TRAINING OBJECTIVES · p. 19
Counterfactual training (CF)
- Counterfactual training privately evaluates each token at candidate depths while keeping preceding committed states fixed and exposing only its committed path to future tokens.It asks which depth would have produced the best prediction for each token while preserving one deployable Ragged-KV history.Source: C RAGGED-AWARE TRAINING OBJECTIVES · pp. 19–20
REINFORCE with causal cache credit (RL)
- REINFORCE uses suffix loss to credit an exit for effects on later tokens’ K/V histories, while counterfactual training directly supervises only the current token.The method samples complete Ragged-KV trajectories; its score term makes trajectories with positive, higher-than-baseline future-loss advantage less likely.Source: C RAGGED-AWARE TRAINING OBJECTIVES · p. 20
Training compute
- Direct RA alone retains standard training cost, while counterfactual and REINFORCE variants add computation without clear improvement over direct alternate-step training.Matched Ouro-160M experiments put CF/RL within roughly 0.02 perplexity of direct alternate-step training; REINFORCE costs about 2D for K=1 and 5D for K=4.Source: Figure 14 · p. 20, C RAGGED-AWARE TRAINING OBJECTIVES · p. 21
D BENCHMARKS, TRAINING DATASETS AND MODELS USED
The section specifies benchmark settings, model-evaluation protocols, judge procedures, and sources, then details the BBH frontier sweep used for additional experiments.
D.1 NON-THINKING MODELS
- Non-thinking Ouro-1.4B and Ouro-2.6B were evaluated without chat templates, using deterministic log-likelihood scoring for multiple-choice tasks and 16 rollouts for generated-answer tasks.MMLU and WinoGrande use raw accuracy; ARC-Challenge and HellaSwag use length-normalized accuracy, with generated answers judged by Qwen3.8-27B.Source: §D.1 · p. 21, Table 4 · p. 21, Table 5 · p. 21
D.2 THINKING MODELS
- Thinking Ouro models were evaluated zero-shot with their chat templates and thinking enabled, using Qwen3.8-27B judgments or executable EvalPlus tests.OlympiadBench uses a fixed 581-question subset; fewer than 0.1% of Qwen judgments were faulty and counted as wrong.Source: §D.2 · pp. 21–22, Table 6 · p. 21
D.3 QWEN JUDGE DETAILS.
- The Qwen judge protocols extract an explicit answer from the model response and compare it with the reference, with separate rules for thinking traces and incomplete generations.The non-thinking protocol treats responses as untrusted and rejects answers that cannot be extracted; the thinking protocol permits a supported answer inside an unterminated trace.Source: §D.3 · pp. 22–24
D.4 DATASET AND MODEL SOURCES
- The dataset and model-source tables list evaluation benchmarks, training datasets, and models used in the work.The reported table descriptions identify the lists but do not provide their entries in these passages.Source: §D.4 · p. 24, Table 7 · p. 25, Table 8 · p. 25
E ADDITIONAL EXPERIMENTAL DETAILS AND RESULTS
- All experiments used H100 GPUs.The additional experimental section reports this hardware setting for both training and evaluation.Source: E ADDITIONAL EXPERIMENTAL DETAILS AND… · p. 24
E.1 BBH FRONTIER EVALUATIONS
- The BBH frontier sweep evaluated 135 questions (five from each of 27 subtasks; seed 1234), with four rollouts per question, comparing fixed@1–4 against thresholded Dense-Ragged and ETE-Ragged settings.Dense-Ragged and ETE-Ragged thresholds were .20, .25, .28, .30, .35, .40, .45, .50, .60, and .75.Source: §E.1 · pp. 24–25
E.2 RAGGED-KV COMPREHENSIVE EVALUATION
The evaluation compares Ragged-KV inference modes and explores alternative cache strategies; native ragged prefill is unstable for non-thinking Ouro-2.6B, while fixing its prompt prefix recovers performance.
E.2 RAGGED-KV COMPREHENSIVE EVALUATION
- The comprehensive evaluation contrasts fixed-depth inference with Dense-Ragged and ETE-Ragged, reporting decode depth for Dense-Ragged and overall depth for ETE-Ragged.Dense-Ragged uses full-depth prefill with decode early exit; ETE-Ragged exits early during both prefill and decode.Source: §E.2 · pp. 26–27, Figure 15 · p. 27, Figure 16 · p. 27
- Cosine and fixed-threshold prefill schedules slightly lower average prefill depth and improve performance over fixed prefill/decode thresholds, but not significantly.The authors leave these schedules as exploration rather than a first-class inference mode.Source: §E.2 · p. 27, Table 14 · p. 33, Table 15 · p. 34
- Keeping the first two prompt tokens at fixed depth four restores Ouro-2.6B non-thinking performance close to the fixed-depth frontier across exit thresholds.The native ETE-Ragged prefill failure is concentrated at the prompt beginning; early prefix-state errors otherwise propagate through the autoregressive KV cache.Source: §E.2 · p. 27, Figure 21 · p. 31
E.3 OTHER DYNAMIC KV CACHES
- The study reports exit-depth and immediate-transition jump fractions across benchmarks, but omits fractions for buggy non-thinking Ouro-2.6B.Figures 19 and 20 show these token fractions for Ragged-KV modes.Source: §E.2 · p. 27, Figure 19 · p. 29, Figure 20 · p. 30
- Using only final-depth prompt K/V as shared decode cache saves no FLOPs and collapsed performance to zero in the tested Ouro models.The authors note prompt K/V must remain parallelizable, unlike autoregressively built decode K/V.Source: §E.3 · p. 28
- Alternative latest-only cache variants perform poorly on Ouro-1.4B: Ragged-Decode-Shared-KV is essentially zero, while Ragged-Collapse-KV improves slightly but remains unusable.The reported cache-baseline table uses greedy decoding and a raw strict answer-marker parser, unlike the headline generation tables.Source: §E.3 · p. 28, Table 17 · p. 35