Scheduling Recursive Reasoning in Looped Transformers
Boyuan Wang, Chengyao Yu, Jiaxi Ren, Hongxin Wei, Bingyi Jing, Yuxin Tao
cs.LG
2026-09-29
TAPS adapts loop step sizes from progress vs. fluctuation; no-retrain gains on Sudoku/Maze, and up to 1.56× wall-clock speedup after co-training at matched accuracy.
Looped Transformers reuse one block of parameters on an evolving latent state, so reasoning depth is no longer tied to parameter count. Prior work mostly chose what to repeat (layer, MoE, or full-model recurrence) and how long to repeat it (fixed K, adaptive exit, fixed-point halt). The scale of each update stayed at one.
A unit step is too small when updates keep pointing the same way, and too large when they jitter. Extra loops then buy less than they should. The right η would come from the terminal task loss, but labels and loss gradients are missing at inference. This paper, from SUSTech, CUHK-Shenzhen, and Shenzhen Loop Area Institute, treats the trajectory as the only test-time signal and asks whether the balance of persistent progress and centered fluctuation is enough to set η online.
Write the frozen loop as X{k+1}=Xk+Δθ(Xk;u). TAPS inserts a relaxation factor: X{k+1}=Xk+ηk Δθ. η=1 recovers the original loop; η>1 over-relaxes along the learned update; η<1 damps it. The controller may use only the input and the trajectory so far.
The terminal loss's partial with respect to each ηk is the inner product of the back-propagated gradient and the current update. Its time average splits exactly into persistent progress P minus centered fluctuation S. Larger P argues for a bigger step; larger S argues for a smaller one. Those oracle terms are hidden at inference, so TAPS estimates them from exponential moving averages of Δ:
Warmup keeps η=1 while statistics fill in. This instantiation is TAPS (Adam), because the two moments match Adam. Swapping the estimator yields GD, sign, Momentum, RMSProp, and Barzilai-Borwein (BB) variants.
Theory gives sufficient conditions: if observable energies track task contributions and window drift is small, following B̂ lowers expected terminal loss and can hit a target quality in fewer loops. Training may add a hinge that penalizes fluctuation only when it exceeds progress, leaving the task loss and ACT in place so the loop stays persistent enough to push η above 1.
Primary numbers are on TRM latent recurrence, Sudoku with a 32-step budget and Maze with 16.
| Setting | Method | Accuracy | Speedup |
| Sudoku, inference-only | Standard η=1 | 89.67% | 1.000× |
| Sudoku, inference-only | TAPS (BB) | 90.19% | 1.039× |
| Sudoku, co-designed | TAPS (Momentum) | 91.39% | 1.300× |
| Maze, inference-only | Standard η=1 | 78.80% | 1.000× |
| Maze, co-designed | TAPS (P/S-Sign) | 79.90% | 1.561× |
With frozen weights, all six TAPS variants beat unit step on both tasks and reach the unit-step terminal accuracy in less wall-clock time. Fixed η=0.8 or 1.2 does not: Sudoku at η=1.2 falls to 85.43% and misses the unit-step target in budget; both constants are N/R on Maze.
After co-design, the best Sudoku run is TAPS (Momentum) at 91.39% and 1.300×; the best Maze run is TAPS (P/S-Sign) at 79.90% and 1.561×, against unit-step 89.67% and 78.80%. Co-designed Standard already reaches 90.60%/1.194× and 78.80%/1.249×, so TAPS is an extra increment on top.
Stacked with other recipes: adaptive exit on Sudoku-Extreme hits 91.2% Exact with 316.5 effective updates, about 4.7× less compute than FPRM at 91.1%; the accuracy fixed-point inference hits at 1,000 updates is matched in about 260; inner-loop control on hierarchical TRM/HRM beats outer-loop by 0.58–0.67 points after co-design. Parallel PTRM numbers sit in Table 10 and are not listed in the main text.
On recurrent LMs, Ouro-1.4B average accuracy moves from 61.10 to 62.86 with TAPS (BB). GSM8K carries most of it (51.78 to 58.98); MMLU dips from 68.23 to 67.95. Huginn-0125 average goes from 38.87 to 39.15. Fixed η=1.2 on Ouro collapses the average to 32.85. Frozen Qwen3-4B-Instruct with three intermediate-layer loops: base 56.77, unit-step Loop 59.22, TAPS (BB) 59.37. The gap over Loop is 0.15 points, and every controller picks mean η̄>1.
If you already run looped inference, this is a cheap extra knob: keep an EMA and stretch or shrink the learned update along the trajectory. On Sudoku and Maze, plug-in TAPS adds about 0.3–0.5 points and roughly 1.04× speed. After co-training, Maze hits the unit-step accuracy 1.56× faster in wall-clock. The 4.7× cut versus FPRM shows the same signal can also drive early exit.
If you train recurrent LMs or wrap intermediate-layer loops around a frozen LLM, keep expectations low. Ouro's 1.8 point average gain sits almost entirely on GSM8K; MMLU does not improve. On Qwen3, TAPS is nearly tied with unit-step Loop. Treating update scale as a third axis beside architecture and depth is fair. On general language tasks that axis is still thin.
The constant-η ablations make the problem real: a better fixed scale is not enough. The schedule has to move along the trajectory.
The authors flag three gaps. The theory needs trajectory statistics to line up with task contributions, and extending that to broader recurrent dynamics is open. TAPS uses only local trajectory proxies. Training-inference co-design is an early attempt.
Inference-only gains on the main table are small: 90.19% vs 89.67% on Sudoku, 79.10% vs 78.80% on Maze. The 1.56× figure is tied to co-training, Maze, and matched unit-step accuracy, not a drop-in speedup. Language results have no variance or multi-seed report, and TAPS (BB) on Ouro is below unit step on MMLU. Qwen3 runs only three loops, so 0.15 points is weak evidence of transfer. The 4.7× comparison is against FPRM's effective updates, not wall-clock versus a fixed unit-step budget. Parallel PTRM and Ouro-2.6B magnitudes are not expanded in the main text.