Scheduling Recursive Reasoning in Looped Transformers

AI summary

Overview

Research area: Efficient inference for recurrent ("looped") transformer models that scale test-time computation by repeatedly applying shared parameters to a latent state.

Technical level: Advanced. The paper combines a first-order sensitivity analysis of the terminal loss (state Jacobians and their adjoints), an exact decomposition of that sensitivity, exponential-moving-average estimators, and three theorems, alongside a broad empirical study across several model families and benchmarks.

Scope in one sentence: The paper proposes, analyzes, and empirically validates a method for adaptively scheduling the scale of each recurrent update — rather than using the standard unit step — using only information observable along the evolving trajectory.

What This Paper Is About

Recurrent reasoning models improve their predictions by iterating a learned update many times, but every iteration is applied with a fixed unit scale, which can be too timid when updates make steady progress and too aggressive when they fluctuate. The authors ask how large each step along the recurrent trajectory should be, and whether the trajectory itself contains enough information to choose that scale online without access to the target label. Their answer is TAPS (Trajectory Adaptive Progress–Fluctuation Scheduler), a controller that tracks the balance between persistent progress and centered fluctuation and adapts the step size accordingly.

Key Contributions

1. An exact decomposition of terminal-loss sensitivity. The authors show that the temporally averaged gradient of the terminal loss with respect to the relaxation factors, denoted $\bar{A}_k^w$, decomposes exactly as $\bar{A}_k^w = P_k^w - S_k^w$, where $P_k^w$ captures persistent progress and $S_k^w$ captures centered fluctuations (Propositions 1 and 2). A remark gives the perturbed-schedule interpretation of $\bar{A}_k^w$ as a first-order terminal-loss advantage.

2. TAPS, a trajectory-only online scheduler. TAPS estimates persistent and fluctuation energies with exponential moving averages of the observed latent updates, forms the normalized balance $\widehat{B}_k = (\widehat{P}_k - \gamma \widehat{S}_k)/(\widehat{P}k + \gamma \widehat{S}k + \epsilon{\mathrm{num}}) \in [-1,1]$, and sets $\eta_k = \mathrm{clip}(1 + \rho \widehat{B}k, \eta{\min}, \eta{\max})$ after a warmup period. The paper provides six instantiations (TAPS (Adam), (GD), (BB), (P/S-Sign), (Momentum), (RMSProp)), varying how progress and fluctuation are estimated and mapped to the relaxation factor.

3. Theoretical guarantees. Under conditional comparability assumptions relating observable energies to oracle contributions, plus a drift bound, Theorem 3 shows the sign of $\widehat{B}_k$ correctly indicates whether $\eta_k$ should increase or decrease relative to $\eta_k = 1$. Theorem 4 gives a finite one-factor gain bound on the expected terminal-loss change, and Theorem 5 shows that when a cumulative gain lower bound $G_K^{\mathrm{lb}}$ is large enough, the policy reaches a target tolerance in fewer loops than the unit-relaxation policy.

4. Training–inference co-design and broad empirical validation. The same progress–fluctuation principle is folded into training through the penalty $\mathcal{L}{\mathrm{co}} = \mathcal{L}{\mathrm{task}} + \mathcal{L}_{\mathrm{ACT}} + \mu \sum_k [\widehat{S}_k - \kappa \widehat{P}k]+$, which suppresses fluctuation only when it exceeds persistent progress. TAPS is tested across latent-state recurrence (TRM), language-model recurrence (Ouro-1.4B, Huginn-0125, and Ouro-2.6B), and intermediate-layer recurrence (Qwen3-4B-Instruct), and combined with adaptive exit, hierarchical recurrence, fixed-point inference, and parallel recurrent inference.

Main Findings

• Inference-only gains across all six controllers. On Sudoku ($H=32$) and Maze ($H=16$), every TAPS variant improves terminal accuracy over the standard unit-step baseline and reaches the baseline's terminal accuracy in less wall-clock time. For example, TAPS (BB) reaches 90.19% on Sudoku and 78.90% on Maze inference-only, versus 89.67% and 78.80% for the standard schedule.

• Co-design delivers the largest gains. After training–inference co-design, the best controller reaches 91.39% accuracy at 1.300× speedup on Sudoku and 79.90% accuracy at 1.561× speedup on Maze, compared with 89.67% and 78.80% under unit-step inference; the paper reports up to 1.56× wall-clock speedup at matched baseline accuracy.

• Fixed constant step sizes are unreliable. Constant $\eta = 0.8$ gives 89.75% on Sudoku (0.990×) and 78.40% on Maze where the unit-step target is not reached within the loop budget (N/R); constant $\eta = 1.2$ is worse still at 85.43% and 77.00%, both N/R. Fixed scales also degrade language-model recurrence substantially (constant $\eta = 1.2$ drops Ouro-1.4B average accuracy to 32.85 and Huginn-0125 to 32.83).

• Difficulty-adaptive compute allocation. On Sudoku-Extreme, the TAPS-based adaptive-exit rule reaches 91.2% Exact with 316.5 effective updates, using 4.7× less compute than FPRM at comparable accuracy (91.1%), and allocates more compute to harder puzzles measured by empty-cell count.

• Fewer updates at fixed-point targets. The accuracy that fixed-point inference attains after 1,000 updates is matched by TAPS within roughly 260.

• Inner-loop control beats outer-loop control in hierarchical recurrence. On Sudoku under hierarchical recurrent inference, $\Delta_{\mathrm{I-O}} > 0$ for all adaptive controllers, with a larger gap after co-design (up to +0.67 percentage points for TAPS (RMSProp)).

• Non-monotonic controller refresh trade-off. Varying the refresh interval $q \in {1, 2, 4}$ shows that less frequent controller updates reduce evaluations but coarsen trajectory tracking, potentially requiring more recurrent updates to reach the same accuracy. Co-design improves the trade-off across refresh intervals.

• Gains transfer to language-model recurrence. All six TAPS variants improve average accuracy over unit-step inference on both Ouro-1.4B (best average 62.86 with TAPS (BB) versus 61.10 standard) and Huginn-0125 (best average 39.15 versus 38.87 standard), though gains vary across MMLU, ARC-C, HellaSwag, and GSM8K. The same trend is reported for the larger Ouro-2.6B model in Table 12.

• Intermediate-layer recurrence favors mild over-relaxation. Applied to frozen Qwen3-4B-Instruct with $K = 3$, all controllers select a mean relaxation factor $\bar{\eta} > 1$, and TAPS (BB) achieves the best average accuracy of 59.37 versus 59.22 for the unit-step Loop and 56.77 for the Base model.

• Nothing is reported about accuracy below the baseline. The paper defines Efficiency as speedup to reach the pretrained unit-step baseline's terminal performance, and reports N/R where that target is not reached within the loop budget; no other failure-mode statistics are given in the provided text.

Methodology in Plain English

The starting observation is that the standard recurrent loop can be rewritten as $\mathbf{X}_{k+1} = \mathbf{X}k + \boldsymbol{\Delta}\theta(\mathbf{X}k; u)$, which implicitly applies a scale of 1 to every learned update. The authors introduce a per-loop relaxation factor $\eta_k$ (allowed to exceed 1) so the loop becomes $\mathbf{X}{k+1} = \mathbf{X}k + \eta_k \boldsymbol{\Delta}\theta(\mathbf{X}_k; u)$.

To decide what $\eta_k$ should be, they first work out how sensitive the final task loss is to each $\eta_k$, using the chain rule and a backward propagation of gradients through the state Jacobians $D_j = I + \eta_j D_{\mathbf{X}}\boldsymbol{\Delta}_\theta(\mathbf{X}_j; u)$. That sensitivity is exact but depends on the label and on future gradients, so it cannot be used at inference time. The key step is to average this sensitivity over a temporal window: the average splits cleanly into two terms, one measuring how consistently the updates move in a common direction (persistent progress) and one measuring how much they deviate from their own average (centered fluctuation).

TAPS then replaces the unavailable oracle quantities with things it can actually observe: exponential moving averages of the latent updates themselves, with the squared norm of the average update serving as the progress energy and the average squared deviation from that mean serving as the fluctuation energy. Their difference, normalized, becomes a score in $[-1, 1]$. If progress dominates, the score is positive and the step size is increased above 1; if fluctuation dominates, the step size is damped below 1. Both are clipped to a safe interval, and the first few warmup loops keep $\eta_k = 1$ while the statistics accumulate.

The authors then prove, under stated assumptions, that this observable score points in the same direction as the true terminal-loss sensitivity and that the chosen step size reduces expected terminal loss. They extend this to a cumulative bound, which converts per-step gains into a statement about reaching a target performance level with fewer loops. Finally, they run inference-only experiments (no retraining) across many model families and inference strategies, and then retrain with an auxiliary penalty that discourages fluctuation relative to progress, so that the learned dynamics better suit adaptive inference.

Why This Matters

Impact on research. The paper reframes recurrent inference as having three separable design axes: what computation is repeated (architecture, $\boldsymbol{\Delta}_\theta$), how long it is repeated (depth, $K$), and how strongly each step is applied (scale, $\eta_k$). Prior work has focused largely on the first two. By providing both a theory and practical controllers for the third, the work opens a control knob that is orthogonal to architecture and horizon choices and can be combined with existing adaptive-depth and convergence-based halting methods.

Real-world applications (potential, inferred from the settings studied):

• Constraint-satisfaction and logical reasoning systems, such as Sudoku and Maze solvers, where the paper reports accuracy and inference-speed improvements without retraining.
• Deployed language-model assistants, where recurrent language models such as Ouro and Huginn-0125 are evaluated on MMLU, ARC-C, HellaSwag, and GSM8K, and where reducing wall-clock time at matched accuracy lowers serving cost.
• Resource-constrained or on-device inference, where the reported up-to-1.56× wall-clock speedup and reduced effective update counts translate into less compute per query.
• Mixed-difficulty workloads, where the adaptive-exit rule on Sudoku-Extreme allocates more compute to harder inputs and stops early on easy ones.

Industry relevance. Any production system that trades test-time compute for accuracy — chain-of-thought-style reasoning services, iterative solvers, or looped-transformer deployments — stands to benefit from a controller that can be applied to pretrained checkpoints without retraining, and that composes with parallel recurrent sampling and fixed-point halting already used in practice.

Future Directions

• Recovering a per-sample rather than per-batch scale. The provided text describes a single relaxation factor per loop; whether finer-grained, token-level or instance-level scaling would help is not addressed in the content provided.
• Tightening the assumptions. Theorems 3–5 rely on conditional comparability bounds between observable energies and oracle contributions plus a drift term $\delta_k^{\mathrm{drift}}$; characterizing when these hold in practice, and how loose they are for real recurrent architectures, remains open.
• Choosing controller hyperparameters without tuning. The memory parameter $\beta$, the fluctuation weight $\gamma$, the adaptation magnitude $\rho$, the warmup length $K_{\mathrm{warm}}$, and the clipping bounds are all set by the user; the paper does not report a principled selection procedure in the provided text.
• Extending beyond the studied settings. The empirical claims cover latent-state recurrence, language-model recurrence, intermediate-layer recurrence, adaptive exit, hierarchical recurrence, fixed-point inference, and parallel recurrent inference; whether the progress–fluctuation principle extends to other recurrent or iterative-optimization settings is not reported.

Target Audience

This paper is most useful to researchers and engineers working on recurrent or looped transformer architectures, test-time compute scaling, and adaptive inference. It will also interest readers of optimization literature, since the design is explicitly inspired by trajectory-based step-size adaptation (the Barzilai–Borwein and related lines of work), and readers who want a worked example of connecting a first-order sensitivity decomposition to a deployable, training-free controller. A solid background in calculus, linear algebra, and transformer-based language modeling is assumed; the empirical sections are accessible to practitioners who want to apply the controllers without following the derivations.

Authors’ abstract

Recurrent reasoning models have attracted growing attention for scaling test-time computation, typically by iteratively refining latent states with shared parameters. However, these models apply each learned update with a fixed unit scale, which can be conservative when updates make persistent progress and overly aggressive when they fluctuate, limiting the benefit of additional loops. To understand how the scale should vary along the trajectory, we first analyze the sensitivity of terminal loss to recurrent update scale. We show that its temporal average admits an exact decomposition into persistent-progress and centered-fluctuation contributions. Based on this, we introduce the Trajectory Adaptive Progress-Fluctuation Scheduler (TAPS), which tracks their balance across recurrent updates and adapts the step size online. Theoretically, we establish sufficient conditions under which TAPS reduces expected terminal loss and reaches a target quality in fewer recurrent loops. Empirically, we show that TAPS improves terminal accuracy across structured reasoning tasks without retraining. By further incorporating the progress-fluctuation principle into training, TAPS yields additional accuracy gains with up to 1.56 times wall-clock speedup at matched baseline accuracy. The broad applicability of TAPS is supported by its effectiveness across diverse recurrent architectures and inference strategies. Together, these results establish update scale as complementary control axis of recurrent inference alongside architecture and depth.