Scheduling Recursive Reasoning in Looped Transformers

1 Introduction

• Update scale complements recurrent architecture and depth as a way to control computation, alongside choices about what is repeated and how many times.Source: §1 · p. 2
• TAPS adapts recurrent update scale online by tracking persistence and fluctuation across updates and adjusting the step size according to their balance.The scheduler uses observed trajectory information rather than task-loss feedback.Source: §1 · p. 2
• The temporal average of terminal-loss sensitivity decomposes into contributions from persistent progress and centered fluctuations.This decomposition motivates larger steps when updates make consistent progress and smaller steps when fluctuations dominate.Source: §1 · p. 2
• TAPS improves terminal accuracy on Sudoku and Maze without retraining, while training with the same principle yields up to 1.56× wall-clock speedup at matched performance.The method also combines with several inference strategies and transfers across recurrent architectures.Source: §1 · p. 2

2 Problem Setup

• A frozen encoder initializes the latent representation, and the same frozen recurrent core is applied at each loop.Source: §2 · pp. 2–3
• The standard loop applies each learned update at unit scale, a fixed choice the paper exposes for adaptive control.The setup studies whether scale can vary across loops without sacrificing task performance.Source: §2 · p. 3
• A scheduled trajectory replaces unit updates with per-loop relaxation factors η_k, where η_k=1 recovers the standard loop, values above 1 amplify updates, and values below 1 damp them.Source: §2 · p. 3
• The scheduling policy chooses η_k using the input and information available through loop k, without access to the target or future states.Source: §2 · p. 3
• The objective is to reach target performance with fewer loops than the standard policy with η_k≡1.Source: §2 · p. 3

3 The Impact of Scheduling: A decomposition of the Gradient

• For a differentiable realized trajectory, each relaxation factor’s effect on fixed-horizon task loss is characterized by its coordinate sensitivity.The sensitivity gives the first-order terminal-loss advantage of a weighted relaxation perturbation.Source: §3 · p. 4
• The task-aligned quantity A_k indicates whether increasing or decreasing η_k favors lower final loss, but it depends on the target and downstream gradient.Those dependencies make A_k unavailable for designing the schedule.Source: §3 · p. 4
• Aggregating A_k over a temporal window connects the task-relevant sensitivity to trajectory information and informs its estimation.Source: §3 · p. 4
• The temporal average of terminal-loss sensitivity decomposes exactly into persistent temporal motion and centered-fluctuation contributions, whose signs are not fixed in general.The terms respectively capture persistent trajectory motion and centered fluctuations.Source: §3 · p. 4

4 TAPS: Trajectory-Adaptive Progress–Fluctuation Scheduling

• TAPS uses pseudo-estimates of persistent progress and fluctuation in learned latent updates to construct a causal scheduling policy.The section introduces the method and its estimates as the basis for the policy.Source: §4 · p. 4, §4.1 · pp. 4–5

4.1 Methodology

• During warmup, TAPS sets the relaxation factor to 1 while continuing to update its temporal statistics.The number of warmup loops is controlled by K_warm.Source: §4.1 · p. 5
• The Adam instantiation adapts first- and second-moment tracking; persistent-dominated motion is over-relaxed, while fluctuation-dominated motion is damped.TAPS (GD) and TAPS (BB) vary how the energies are estimated and mapped to the relaxation factor.Source: §4.1 · p. 5

4.2 Theoretical Property

• Under Eqs. (4) and (5), the sign of the balance statistic indicates whether to increase or decrease the relaxation factor and matches the expected oracle direction.The result concerns the proposed causal controller’s proxy statistics and an unobservable task-improving direction.Source: §4.2 · p. 6
• Under Eqs. (4) and (5), the selected factor has a nonpositive conditional expected terminal-loss difference; additional stated conditions make it strictly negative for over- or under-relaxation.The one-factor bound assumes a Lipschitz terminal-loss derivative and uses drift conditions for the two directions.Source: §4.2 · p. 6
• When cumulative gain meets the baseline-to-target gap, the proposed policy reaches target tolerance by loop K and can do so before unit relaxation.The strict fewer-loop conclusion applies when the unit-relaxation policy has not reached the target by K.Source: §4.2 · p. 6

5.2 Step-Size Schedule at Inference and Training

• All six inference-only TAPS variants improve terminal accuracy and reach the unit-step baseline’s terminal performance faster on both tasks; fixed scales do not consistently improve both.Different variants lead on different metrics.Source: §5.2 · p. 8
• 91.39% accuracy at 1.300× speedup on Sudoku and 79.90% at 1.561× on Maze, versus unit-step accuracies of 89.67% and 78.80%, respectively.With progress–fluctuation co-design, the best controller achieves both accuracy and speedup gains.Source: §5.2 · p. 8

5.3 Comparison with Different Inference Strategies

• On Sudoku-Extreme, adaptive exit reaches 91.2% Exact with 316.5 effective updates, using 4.7× less compute than FPRM at comparable 91.1% accuracy.TAPS allocates more compute to harder puzzles; gains emerge mainly at high accuracy after sufficient trajectory history enables reliable exit.Source: §5.3 · p. 8, Figure 2 · p. 8
• Inner-loop scheduling yields higher terminal accuracy than outer-loop scheduling, with a larger gap after co-design; less frequent controller updates save overhead but can reduce accuracy and slow progress.Across refresh intervals, co-design improves the accuracy–speedup trade-off, while lower-frequency control does not monotonically improve end-to-end efficiency.Source: §5.3 · p. 9, Figure 3 · p. 9
• Applying TAPS independently within PTRM’s parallel rollouts consistently improves performance across parallel-sampling budgets and remains robust to noise-scale variation.The parallel sampling procedure is unchanged.Source: §5.3 · p. 9
• TAPS reaches higher-accuracy fixed-point targets with substantially fewer recurrent updates, matching accuracy attained after 1,000 fixed-point updates within roughly 260.Source: §5.3 · p. 10

5.4 Extensions to Other Model Architectures

• Across Ouro-1.4B and Huginn-0125 on four language benchmarks, all six TAPS variants improve average accuracy over unit-step inference on both models, while fixed rescaling can substantially degrade performance depending on scale.Source: §5.4 · p. 10
• On frozen Qwen3-4B-Instruct intermediate-layer loops, TAPS (BB) achieves the best average accuracy, and all controllers select a mean step-size multiplier above 1.This indicates these loops generally favor mild over-relaxation.Source: §5.4 · p. 10

• Recurrent inference exposes three design axes: the update Δθ, iteration count K, and update scale ηk; prior work has extensively studied the first two, while this paper turns to scale.Source: §6 · pp. 10–11
• Stability-oriented methods use smaller fixed steps, while fixed-point methods damp updates or accelerate directions to recover or speed convergence.Source: §6 · p. 11
• TAPS instead adapts ηk online to balance progress and stability for task performance under a finite inference budget, while preserving the learned update direction.Source: §6 · p. 11

7 Conclusion

• TAPS adapts recurrent step sizes using local dynamics, and the theory characterizes how progress–instability balance guides beneficial adjustments.Experiments demonstrate improved, generalizable recurrent inference across architectures and inference strategies without modifying model weights.Source: §7 · p. 11
• The theory depends on conditions linking trajectory statistics to task contributions; richer observable signals and deeper training integration remain open directions.The conclusion also identifies broader recurrent dynamics and architectures as areas for future study.Source: §7 · p. 11

A.4 Proof of Theorem 4

• Theorem 4’s proof establishes two strict-improvement statements.Source: §A.4 · p. 17

A.5 Proof of Theorem 5

• The proof compares causal controller rollouts that share the same pre-decision trajectory statistics, with unit relaxation during warmup before applying the theorem’s expected-gain argument.For policies differing at a decision, the proof uses their shared state and controller factor at that point; the available passages do not specify the theorem’s full quantitative conditions.Source: §A.5 · pp. 17–18

B.1 Exact EMA Mean–Variance Representation

• The EMA lemma establishes an exact mean–variance representation using normalized weights w_k,j.The supplied passage names the representation but does not include its full identity.Source: §B.1 · p. 18

B.2 Limits of constant relaxation: a toy example

• When error modes have distinct response rates, a nonconstant schedule can target them separately, whereas a constant factor cannot match each mode’s preferred scale.In the toy model, η_k = 1/λ_i eliminates component i; the proposition considers distinct λ_i and a horizon K ≥ r.Source: §B.2 · p. 19

B.3 Discussion of assumptions

• The proxy statistics are constructed from recurrent updates, while the task-dependent oracle quantities depend on task loss, motivating explicit assumptions to connect them.The assumptions are discussed in terms of reference schedules.Source: §B.3 · p. 20

B.3.1 How the assumptions connect the proxy to the decision

• Proxy alignment and temporal alignment jointly connect persistence–fluctuation statistics to the current loss-improving direction, without requiring unbiased or pointwise estimates.A positive lower bound on expected sensitivity favors increasing η_k; a negative upper bound favors decreasing it, when the stated scale bounds hold.Source: §B.3.1 · p. 20

B.3.2 A two-mode fixed-point example

• The model corresponds to gradient descent on a quadratic whose unit-relaxation factors are r and −r, so increasing scale can help one mode while hurting the other.These factors arise from the balanced step for a quadratic with distinct extreme curvatures.Source: §B.3.2 · p. 21
• When the controller moves in the oracle-indicated direction under the stated conditions, the example establishes a positive loss decrease.With exact curvature, the theorem’s gain bound equals the actual loss decrease.Source: §B.3.2 · pp. 21–22
• The two-mode model captures competing recurrent behaviors: one mode contracts monotonically, while the other alternates in sign.It is framed as a canonical local behavior of recurrent fixed-point dynamics.Source: §B.3.2 · pp. 21–22
• At the controlled decision, the persistence–fluctuation statistics can match their task-dependent oracle counterparts exactly, with zero temporal mismatch.The proposition states this alignment for the two-mode setup and the specified initial-state conditions.Source: §B.3.2 · pp. 21–22

C.1 Experimental Setup

• Latent-state recurrence uses TRM on Sudoku-Extreme and Maze-Hard, with exact accuracy requiring the entire predicted grid to match the reference.Sudoku uses five disjoint 10,000-example test subsets; Maze uses its full 1,000-example test set.Source: §C.1 · p. 22
• Language-model recurrence evaluates Ouro and Huginn models on MMLU, ARC-Challenge, HellaSwag, and GSM8K using checkpoint-released recurrent depths.GSM8K uses exact match on the extracted final answer; other tasks use likelihood ranking over candidate labels.Source: §C.1 · pp. 22–23
• Intermediate-layer recurrence applies three damped forward Euler stages to selected Qwen3-4B-Instruct layers during prefill and decoding.The benchmarks are CommonsenseQA, GPQA-Main, and a 1,000-question MMLU-Pro subset; the Loop baseline uses η_k = 1, an Euler step of 1/3.Source: §C.1 · p. 23
• All six controllers share the stated γ, ρ, and relaxation bounds, with β = 0.8 where moving averages are used.Controllers are reinitialized for each forward call; the first stage is a unit-scale warmup whose statistics seed the adaptive stages.Source: §C.1 · p. 23

C.2 Additional Experiments on Training Inference Co-design

• 0.8 percentage points: P/S co-design improves validation exact accuracy at matched checkpoints, while reducing late-stage training-loss variability by 10.3%.At matched training loss, held-out terminal loss falls by 5.9%; the figure reports the validation-accuracy comparison.Source: Figure 4 · p. 24, §C.2 · p. 24
• 5.7%: P/S co-design reduces adjacent total variation in TAPS’s mean multiplier while changing the average multiplier only slightly.The average changes from 0.8713 to 0.8649, indicating that the smoother adaptation is not explained by uniformly smaller scales.Source: §C.2 · p. 24, Figure 5 · p. 24

C.3 Additional Experiments on Adaptive Exit

• Adaptive exit uses step-size-adjusted residual activity relative to its trajectory peak and exits after the prediction remains unchanged and activity stays below threshold for consecutive loops.The stated configuration uses k_min = 4, p = 2, and τ = 0.15; Table 7 reports accuracy and inference cost, but the supplied caption gives no outcome values.Source: §C.3 · pp. 25–26, Table 7 · p. 26

C.4 Additional Experiments on Update Frequency

• Inner-loop scheduling achieves higher Maze accuracy than outer-loop scheduling, while less frequent refresh reduces controller overhead but can offset savings through coarser tracking and slower progress.The frequency trade-off is non-monotonic, and co-design improves it across refresh intervals.Source: Figure 6 · p. 26, §C.4 · p. 27, Table 8 · p. 27

C.5 Additional Experiments on Parallel Sampling

• Under zL noise, adaptive controllers retain comparable or improved Best-Q accuracy, but zH noise can make them underperform unit-step.Table 11 evaluates 1,000 held-out Sudoku-Extreme examples with H = 16 and 25 parallel rollouts.Source: §C.5 · pp. 27–28, Table 11 · p. 28
• At K = 10, TAPS-Adam raises exact accuracy from 97.8% to 98.6%, while TAPS-BB reaches 99.2% at K = 100.The best TAPS variant improves over unit-step at every evaluated budget.Source: §C.5 · p. 27, Table 10 · p. 28

C.6 Additional Experiments on Recurrent Language Models

• On Ouro-2.6B, the best TAPS configuration raises four-task average accuracy from 71.35% to 71.70%, while fixed η = 0.8 or η = 1.2 substantially degrades performance.This contrast suggests gains come from adapting scale along the trajectory rather than uniformly changing its magnitude.Source: §C.6 · pp. 28–29, Table 12 · p. 29

C.7 Additional Experiments on Efficient TAPS Implementation

• Algebraic simplification preserves the controller output and recurrent trajectory, averaging 1.003× speedup and reaching 1.017× for TAPS (P/S-Sign).It reuses residual energies and inner products to avoid redundant tensor materialization and repeated reductions.Source: §C.7 · p. 29, Table 13 · p. 30
• Fusion yields consistent 1.033×–1.049× end-to-end speedups across TAPS variants, averaging 1.042×.The fused operator reduces kernel launches, global-memory traffic, and intermediate tensor materialization without changing controller equations.Source: §C.7 · pp. 29–30, Table 13 · p. 30

D Inference-Time Step-Size Schedulers

• The unified inference interface represents each scheduler as a stateful rule whose history determines the step size for adaptive schedules, while fixed schedules use no memory.Schedulers operate on the same pretrained recurrent model and differ only in how they choose each update's step size.Source: D Inference-Time Step-Size Schedulers · p. 30

D.1 TAPS with Adjacent-update Scores.

• Adjacent-update scores distinguish persistent update components from alternating ones, with a two-step warmup that sets ηk = 1 while storing required states.The scores are based on the sum and difference of consecutive updates.Source: §D.1 · p. 31

D.2 TAPS with Moment and Curvature Scores

• The moment- and curvature-score schedulers define progress and stability from temporal moments or local directional curvature.Source: §D.2 · p. 33