PaperlayerView exact revision on arXiv

Source-linked AI summary

Fixed-Point Reasoners: Stable and Adaptive Deep Looped Transformers

Sajad Movahedi, Vera Milovanović, Shlomo Libo Feigin, Alexander Theus, Thomas Hofmann, Valentina Boeva, T. Konstantin Rusch, Antonio Orvieto

arXiv:2606.18206v1cs.AI

TL;DR

Looped reasoning models need to allocate more computation to harder inputs, but deciding when to halt can be difficult without specialized traces or a separate halting module. FPRM combines pre-norm and residual scaling with fixed-point halting; the authors report competitive benchmark performance and compute adaptation to task difficulty.

  • Problem

    Reasoning models need flexible test-time computation and a way to halt adaptively, while CoT can require specialized training and traces and looped models face a nontrivial halting decision.

    Source: §1 · pp. 1–2

  • Method

    FPRM uses a looped Transformer with pre-norm and residual scaling, halting when its hidden state converges to a fixed-point.

    Source: §1 · p. 2

  • Results

    FPRM outperforms similar-size hierarchical baselines on common symbolic reasoning benchmarks and adapts compute to task difficulty on state-tracking and Sudoku-Extreme.

    Source: §1 · p. 3, §6 · p. 15

  • Takeaways & Limitations

    The authors conclude that fixed-point iterations can dynamically scale depth with task difficulty while architectural modifications improve signal propagation.

    Source: §6 · p. 15

    Experiments cover algorithmic tasks rather than natural language and are limited to Transformer architectures.

    Source: §5 · p. 15

Abstract

from arXiv · show

Looped architectures provide an inductive bias toward learning step-by-step procedures for tasks that require compositional reasoning. The number of effective layers reached by looping determines the quality of the solution these models find. Like deep architectures, looped architectures are prone to a signal propagation problem induced by depth as the halting decision is postponed. In this paper, we address this signal propagation issue using pre-norm layers and residual scaling. Building on these architectural modifications, we propose FPRM, a Transformer-based Fixed-Point Reasoning Model that uses fixed-point convergence as an end-to-end halting mechanism in a looped architecture. We show that fixed-point halting allows FPRM to adapt its compute to task difficulty. FPRM is effective on common reasoning benchmarks, namely Sudoku, Maze, state-tracking, and ARC-AGI.

1 Introduction

Reasoning models need flexible computation and a way to decide when to stop, but existing approaches can require hand-crafted traces or separate halting modules. FPRM combines fixed-point halting with pre-norm and residual scaling, and is reported to outperform baselines on several reasoning benchmarks.

1 Introduction

  • CoT scales compute through verbalization and a specialized halting token, but requires a special training regime and hand-crafted reasoning traces.The authors argue this complexity undermines end-to-end training.Source: §1 · p. 1
  • Looped architectures scale compute along depth, and can provide variable computation while supporting algorithmic and reasoning tasks.The cited examples include HRM and TRM, which use hierarchical looping for puzzle tasks such as Sudoku, maze, and ARC-AGI.Source: §1 · p. 2
  • Fixed-point convergence provides a halting signal without an external module, letting the model spend compute according to the input.The authors contrast this with fixed or randomly sampled loop counts, which eliminate adaptivity, and learned halting modules, which introduce optimization challenges.Source: §1 · p. 2
  • FPRM combines pre-norm with residual scaling to make deeper effective layers trainable while addressing signal propagation in looped architectures.The paper frames looping as increasing effective depth, which can unlock expressivity but also creates a stability challenge.Source: §1 · pp. 2–3, Figure 2 · p. 2
  • FPRM outperforms baselines on Sudoku-Extreme, Maze-Hard, ARC-AGI-1, and state-tracking A5 and S5 among 7M-parameter models, without hierarchical structure.The authors also describe it as the first Transformer-based reasoning model to exhibit compute adaptivity to task difficulty.Source: §1 · p. 3

2 Background and Related Works

Related work frames looping as a way to scale computation through depth, while leaving open how to allocate iterations and make deep effective layers usable. FPRM addresses these issues through fixed-point halting and signal-propagation modifications.

Looped models

  • Looped architectures decouple effective depth from parameter count, providing variable computation and an inductive bias toward repeatedly applied compositional procedures.Recent compact models such as HRM and TRM use hierarchical looping and have shown strong results on reasoning tasks.Source: §2 · p. 3

Adaptive computation

  • Adaptive computation methods learn explicit halting decisions, but add a separate process that is difficult to optimize because halting is discrete.Prior analyses report that ACT does not always scale inference compute with input difficulty; the authors say they observe and address this limitation.Source: §2 · p. 3

Deep equilibrium and fixed-point models

  • Fixed-point approaches halt when consecutive latent states become sufficiently close, whereas Deep Equilibrium Models typically find equilibria using quasi-Newton methods or Anderson acceleration.The paper distinguishes its fixed-point iterations from those solver-based approaches.Source: §2 · pp. 3–4

Score-based methods

  • Energy-based reasoning halts at a local energy minimum, while FPRM models generation through iterations of a learned operator.The authors say learning the operator directly matters to their theoretical analysis of architectural modifications.Source: §2 · p. 4

Signal propagation in looped models

  • In looped models, post-norm bounds activations but can impede signal propagation; pre-norm without further modification can diverge as iteration count grows.FPRM combines pre-norm with residual scaling to recover bounded, stable dynamics while allowing gradients and representations to propagate through many iterations.Source: §2 · p. 4

3 Method

The method turns looped computation into an adaptive fixed-point process while addressing deep-loop trainability and stability. FPRM combines pre-norm, residual scaling, damping, and truncated-BPTT.

3.1 Improving signal propagation with pre-norm

  • Post-norm bounds activations but impairs deep signal propagation; pre-norm improves signal propagation yet its residual activations grow with effective depth.In the state-tracking experiment, increasing effective layer did not improve post-norm expressivity, while pre-norm diverged at larger depth.Source: §3 · p. 4, §3.1 · p. 5

3.2 Recovering boundedness via residual scaling

  • Tied scalars scale residual and sub-layer outputs within layers and mix the input between iterations; coupling the two schemes yields bounded recurrence.The input-mixing scalars are shared across iterations, while the layer-wise scalars are shared across layers.Source: §3.2 · pp. 5–6
  • When α2λf < 1, the loop map is contractive and converges linearly to a unique fixed point; learnable scaling avoids excessive contraction.The authors report that initializing α2 small improves performance and that the pre-norm model maintains performance over longer looping horizons before saturating than the post-norm variant.Source: §3.2 · p. 6

3.3 Oscillation around the fixed-point

  • For oscillatory iterations, residual-based patience reduces the damping step size; under ℜ(λi) < 1, sufficiently small damping converges locally without changing fixed points.The theorem's condition concerns every eigenvalue of the Jacobian at the fixed point.Source: §3.3 · pp. 6–7

3.4 Optimization of fixed-point models

  • Under contractivity, truncated BPTT approximates fixed-point gradients with an error bound, trading computation against accuracy while keeping training memory fixed as loops grow.The truncation depth controls the computation-accuracy trade-off.Source: §3.4 · pp. 7–8

3.5 The fixed-point reasoning model

  • FPRM uses a classic looped Transformer, iterating to convergence for difficulty-scaled compute, with truncated BPTT and deep supervision during training.The authors report no experimental improvement from adding the hierarchical structure used in recent Looped Transformers.Source: §3.5 · p. 8

4 Experiments

The experiments compare FPRM with looped reasoning baselines on puzzle, adaptivity, and signal-propagation benchmarks. They examine benchmark performance, compute allocation, depth utilization, and residual-scale behavior.

4.1 Dataset description

  • The evaluation covers Sudoku-Extreme, Maze-Hard, ARC-1 and ARC-2 puzzles, plus A5 and S5 state-tracking tasks.The benchmarks use exact sequence accuracy or pass@2 and include state-tracking evaluation on lengths up to 128 after training on up to 32 updates.Source: §4.1 · pp. 8–9

4.2 Puzzle tasks

  • FPRM is the best-performing 7M-parameter model on Sudoku-Extreme, Maze-Hard, and ARC-1, and performs on par with TRM on ARC-2.On ARC-2, the comparison may be unfair because performance appears more sensitive to parameter count; FPRM also underperforms larger baselines there.Source: §4.2 · pp. 9–10, Table 1 · p. 9

4.3 Adaptivity

  • At sequence length 128, TRM and TRM with ACT reach 45.8% ± 3.9% on A5 and 39.4% ± 1.9% on S5; adding causal convolution raises these to 91.4% ± 2.3% and 97.2% ± 2.5%.FPRM reaches 98.1% ± 2.2% on A5 and 98.8% ± 0.9% on S5 while smoothly scaling compute with length; ACT with convolution is unreliable across seeds.Source: §4.3 · p. 10
  • On Sudoku-Extreme, both FPRM and TRM with ACT adapt inference compute to puzzle difficulty, while default TRM exhausts its budget; FPRM is more efficient and accurate.Difficulty is measured by the number of empty cells, and the models' maximum compute budgets are matched.Source: §4.3 · pp. 10–11, Figure 5 · p. 11

4.4.1 Boundedness of activation norms and trainability

  • Pre-norm without residual scaling has exponentially growing activations and cannot reach deep effective layers; post-norm and pre-norm with residual scaling keep activations bounded.The experiments identify bounded activations as a prerequisite for reaching deeper effective layers.Source: §4.4.1 · p. 12

4.4.2 Pre-norm with residual scaling enables higher depth utilization

  • Only pre-norm with residual scaling shows state-tracking expressivity tracking effective depth, with maximum sequence length solved above 90% accuracy following the identity line.Models were trained with loop counts tied to sequence lengths from 8 to 512, then tested at matched effective depth.Source: §4.4.2 · pp. 12–13
  • On Sudoku-Extreme, test-time scaling improves performance for both normalization variants, but pre-norm with residual scaling saturates at almost twice the compute of post-norm.FPRM also outperforms TRM across tested effective-layer depths, with the gap widening at higher compute budgets.Source: Figure 7 · p. 12, §4.4.2 · p. 13

4.4.3 The training dynamics of residual scales

  • The best residual-scale initialization uses high α1 and low α2; α1 increases accuracy slightly more than decreasing α2 in the reported ablation.After training, α2 values mostly become smaller than their initial value, while α1's distribution widens more substantially.Source: §4.4.3 · pp. 13–14, Table 2 · p. 13, Figure 9 · p. 13

5 Discussion

The discussion presents FPRM as adaptive: it uses fewer effective layers than TRM while achieving better performance, and its compute can be tuned through the decay rate and maximum patience. The authors also identify limits in evaluation scope, implementation efficiency, and the optimality of learned algorithms.

Adaptive compute

  • FPRM adapts more effectively than TRM, using fewer effective layers while achieving better performance by halting closer to accuracy saturation.TRM’s ACT mechanism may halt too early and lower performance or too late and use excessive compute.Source: §5 · p. 14
  • At inference, halted samples remain in the batch until the last sample halts, and an efficient implementation is left for future work.FPRM uses a residual threshold of 0.1 to record when samples halt.Source: §5 · p. 14

The role of hierarchy in HRM and TRM

  • The authors hypothesize that hierarchy may help HRM and TRM by improving signal propagation, rather than being essential to their success.Their evidence includes FPRM improving on these models without hierarchy and TRM improving when compute is reallocated to deep supervision.Source: §5 · p. 14

Scaling behavior of FPRM

  • The decay rate γ and maximum patience P let practitioners select a desired point on FPRM’s compute–performance Pareto front.Better signal propagation helps FPRM use compute more efficiently, and its improved test-time scaling is mostly visible on harder tasks.Source: §5 · p. 14
  • The optimality of algorithms learned by looped models is not guaranteed, and the authors identify logarithmic-complexity state tracking with Turing completeness as an open challenge.They observe suboptimal scaling in state tracking, as also reported for recurrent-in-depth state-space models.Source: §5 · p. 15

Limitations

  • Experiments cover algorithmic tasks, not natural language, and are limited to Transformers despite FPRM’s architecture potentially supporting other model types.Whether compositional reasoning on algorithmic tasks transfers to other domains remains an open challenge.Source: §5 · p. 15

6 Conclusion

The conclusion presents architectural modifications that enable pre-norm in looped fixed-point Transformers, improving their use of deeper effective layers. FPRM outperforms similarly sized hierarchical baselines on common symbolic reasoning benchmarks and adapts compute to task difficulty on state tracking and Sudoku-Extreme.

6 Conclusion

  • FPRM outperforms similarly sized HRM and TRM baselines on common symbolic reasoning benchmarks, and adapts compute to task difficulty on state tracking and Sudoku-Extreme.The conclusion attributes this capability to dynamically scaling depth through fixed-point iterations and improving signal propagation.Source: §6 · p. 15

A Proofs

The proofs analyze boundedness across fixed-point iterations, contraction of the looped map, and stability under damping. They establish convergence under stated conditions, including local asymptotic stability for sufficiently small damping.

A.1 Fixed-point iterations bound

  • With 0 ≤ α1, α2 < 1, the geometric series in the unrolled fixed-point iteration converges, supporting a boundedness argument as the iteration count tends to infinity.The proof derives 0 ≤ ρ < 1 before taking the limit as i → ∞.Source: §A.1 · p. 22

A.2 Contractive mapping

  • When 0 ≤ α2λf < 1, the looped map is contractive and its iteration converges to a unique fixed point.The contraction condition follows from the Lipschitz constant of the L-layer model and the loop’s α2 scaling.Source: §A.2 · p. 23

A.3 Mitigating oscillation through damping

  • For every η in (0, η0), the damped iteration is locally asymptotically stable and converges to z⋆ from sufficiently close initializations.The proof chooses η0 so the damped Jacobian’s eigenvalues satisfy the stability condition, assuming ℜ(λi) < 1.Source: §A.3 · p. 23

A.4 Error of truncated-BPTT

The section bounds the approximation error from truncating the implicit-gradient series and applies the result to truncated backpropagation through time. The approximation and gradient errors decay exponentially with the number of backward passes.

A.4 Error of truncated-BPTT

  • The truncated BPTT gradient error decays exponentially with the number of backward passes k.The preceding approximation error decays as O(σ^k).Source: §A.4 · p. 24

B A Toy Failure Mode for Recurrent Post-norm

A toy sequence model illustrates why recurrent post-norm can be bounded yet difficult to train. Pre-norm with residual scaling improves signal propagation while preserving bounded dynamics.

B A Toy Failure Mode for Recurrent Post-norm

  • The toy setup uses random input-output sequences with length 100 and hidden size 2, and evaluates loss at the 20th effective layer.The loss is defined over a two-dimensional parameter slice.Source: B A Toy Failure Mode for Recurrent… · pp. 24–25
  • Post-norm keeps recurrent states bounded by projecting rows onto the unit sphere, but produces narrow, poorly conditioned low-loss regions.In this two-parameter slice, the loss landscape has thin angular sectors and sharp ridges, making random initializations likely to start outside useful basins.Source: B A Toy Failure Mode for Recurrent… · p. 25
  • Pre-normalization improves signal propagation, while residual scaling replaces post-normalization’s boundedness mechanism and preserves a live residual stream.The resulting toy dynamics remain bounded.Source: B A Toy Failure Mode for Recurrent… · p. 25

C Further Details About the Architecture

FPRM measures per-sample fixed-point convergence using residuals, trains with truncated backpropagation at deep-supervision intervals, and applies depth-wise convolutions to latent representations.

C Further Details About the Architecture

  • FPRM measures convergence per sample using the relative L∞ norm of the residual and declares convergence when it falls below τ, set to 0.1 in practice.The loop also has a hard iteration cap and terminates early if the adaptive step size falls below a minimum.Source: C Further Details About the Architecture · p. 25
  • Deep supervision decodes intermediate activations every Tsup iterations, computes loss, and applies truncated backpropagation through the latest k iterations before detaching the graph.FPRM sets Tsup = k and uses fewer forward and backward passes during training than TRM in practice.Source: C Further Details About the Architecture · p. 25
  • FPRM applies depth-wise convolutions at the start of each loop; 2D convolutions work better on Sudoku and ARC, while 1D convolutions are essential for state-tracking.The passage reports that depth-wise convolutions seem detrimental to TRM performance.Source: C Further Details About the Architecture · p. 26

D Description of Figure 1

The figure groups Sudoku puzzles by difficulty, measured by their number of empty cells, and marks halting decisions for FPRM and TRM.

D Description of Figure 1

  • Puzzles are grouped as easy, medium, or hard according to the number of empty cells.Each difficulty group has around 1,000 samples.Source: D Description of Figure 1 · p. 26
  • FPRM marks a group as halted when its mean residual falls below 0.1, whereas TRM does so when its ACT module signals halting for more than half the samples.Source: D Description of Figure 1 · p. 26
  • The hardest puzzles are excluded from the displayed groups because they do not halt within the current 10,000-effective-layer maximum budget.Source: D Description of Figure 1 · p. 26

E Fixed-point Residuals and Halting

FPRM’s residuals decay more slowly on harder Sudoku-Extreme problems, and accuracy plateaus at roughly the same effective layer as the residuals. At a fixed inference budget, TRM performs better when more compute goes to outer deep-supervision steps.

E Fixed-point Residuals and Halting

  • On Sudoku-Extreme, harder problems’ residuals decay more slowly, and accuracy stops improving at roughly the effective layer where the residual plateaus.This supports fixed points as a halting criterion; Figure 12 shows accuracy and residuals as a function of effective layer.Source: E Fixed-point Residuals and Halting · p. 26, Figure 12 · p. 27
  • The Sudoku-Extreme dataset is imbalanced across difficulty levels, which are defined by the number of empty cells.Source: Figure 11 · p. 27
  • Table 3 evaluates pre-norm and residual-scaling modifications to TRM, individually and in combination, with and without a conv2d layer.Source: Table 3 · p. 27

F How to effectively spend loops in TRM?

  • At an approximately 1,040-step inference budget, TRM performs better when compute is spent on more outer deep-supervision steps rather than deeper inner recurrence.The tested configurations keep the L-to-H ratio at 2; total effective depth is approximately the same across TRM and FPRM configurations on Sudoku-Extreme.Source: F How to effectively spend loops in TRM? · p. 27, Figure 13 · p. 28

G Additional Experimental Details

The additional experiments report that truncated-normal initialization speeds convergence without materially changing final sequence accuracy, while longer maze training did not improve performance. The section also records task-specific training configurations and a signal-propagation variant sweep.

G Additional Experimental Details

  • Truncated-normal (LeCun) initialization accelerated convergence but made very little material difference to sequence accuracy after convergence.The passage describes this initialization as common practice in looped architectures.Source: G Additional Experimental Details · p. 28
  • Maze training converged on the training data, and extending training to as long as 7 days did not improve performance.The passage notes some evidence for grokking in looped architectures, but reports this maze-task outcome.Source: G Additional Experimental Details · p. 28
  • Sudoku-Extreme experiments shared one A100-40GB, batch 768, 60,000 epochs, a 2,000-step warm-up, EMA rate 0.999, and puzzle-embedding length 16.All models used AdamW; the learning rate was constant after warm-up.Source: Table 4 · p. 29
  • Maze-Hard experiments used maze-30x30-hard-1k without augmentation on four A100-80GB GPUs; FPRM trained for 60,000 epochs and TRM for 50,000.Both used a 2,000-step warm-up, EMA rate 0.999, and puzzle-embedding length 16; FPRM used Adam-Atan2 and TRM AdamW.Source: Table 5 · p. 30
  • The signal-propagation analysis swept post-norm, pre-norm, and pre-norm plus residual-scaling variants of the Looped Transformer.The experiments used one A100-80GB, global batch 1024, Adam-Atan2, no warm-up or EMA, and training at ktrain=32 with evaluation for k ∈ [2, 128].Source: G Additional Experimental Details · p. 31, Table 6 · p. 31
  • The ARC-AGI runs shared the same configuration and differed only in training corpus: ARC1-Concept versus ARC2-Concept, each with 1,000 augmentations per sample.Both used four A100-80GB GPUs, batch 768, 100,000 epochs, puzzle-embedding length 16, hidden size 512, eight heads, and Adam-Atan2.Source: Table 7 · p. 32
Loading 2606.18206v1…