Fixed-Point Reasoners: Stable and Adaptive Deep Looped Transformers — part 2/6
Cleaned-text code points [15570, 31231). Bibliography omitted here; source.md includes it.
Consequently, in this section we propose to (a) switch to pre-norm to recover trainability at higher depth, while preserving boundedness by (b) scaling the residual stream and the sub-layer outputs (the attention and feed-forward maps). Together, these modifications yield a layer that is both stable under looping and trainable at large depth. Additionally, we introduce (c) a fixed-point halting mechanism that decides the effective layer adaptively per input. Interestingly, we find that with these changes we can (d) remove the hierarchy common in recent Looped Transformers ( Wang et al., 2025 ; Jolicoeur-Martineau, 2025 ) , resulting in a much simpler model. We provide an overview of our proposed model in Figure 3 .
### 3.1 Improving signal propagation with pre-norm
We start by introducing pre-norm and post-norm in a Transformer layer. A Transformer layer consists of two sub-layers (with f θ ℓ ℓ ( . ) f_{\theta^{\ell}}^{\ell}(.) denoting the ℓ t h ∈ { 1 , … , 2 L } \ell^{th}\in\{1,\dots,2L\} sub-layer): multi-head attention and a feed-forward network. We denote the Looped Transformer model as defined in Equation 1 , consisting of multiple Transformer layers, by f θ ( . ) f_{\theta}(.) . In the Transformer layer, the two sub-layers are interleaved with layer normalization ( Norm \mathrm{Norm} ). Two canonical placements of the normalization define two variants of the layer. The original post-norm formulation ( Norm post \mathrm{Norm}_{\mathrm{post}} ) ( Vaswani et al., 2017 ) applies normalization after the residual addition with the residual stream ( 𝐳 ℓ − 1 \mathbf{z}^{\ell-1} ), while the pre-norm variant ( Norm pre \mathrm{Norm}_{\mathrm{pre}} ) ( Xiong et al., 2020 ) applies normalization to the input of each sub-layer without modifying the residual stream: 𝐳 ℓ = Norm post ( 𝐳 ℓ − 1 + f θ ℓ ℓ ( Norm pre ( 𝐳 ℓ − 1 ) ) ) , ℓ = 1 , … , 2 L , \mathbf{z}^{\ell}=\mathrm{Norm}_{\mathrm{post}}\!\left(\mathbf{z}^{\ell-1}+f^{\ell}_{\theta^{\ell}}\!\left(\mathrm{Norm}_{\mathrm{pre}}\!\left(\mathbf{z}^{\ell-1}\right)\right)\right),\qquad\ell=1,\ldots,2L, where L L is the number of layers (Transformer blocks). In fixed-depth models, both normalization placements have been linked to training issues at large depth: post-norm bounds activation magnitudes but induces a signal propagation problem ( Kim et al., 2025 ; Noci et al., 2022 ) , while pre-norm improves signal propagation but causes exponential growth in residual magnitude ( Kim et al., 2025 ; Xiong et al., 2020 ) . Therefore, while using pre-norm in a deep neural network is desirable from the signal propagation standpoint, it can introduce instability due to unbounded activations.
Motivated by these observations, we investigate both normalization placements in Looped Transformers ( Saunshi et al., 2025 ; Dehghani et al., 2019 ) . In Figure 2(a) , we test the positive correlation between the effective layer and the expressivity (maximum sequence length with > 90 % >\!90\% test accuracy) of a Looped Transformer for the state-tracking task A 5 A_{5} ( Merrill et al., 2024 ) . We observe that increasing the effective layer of a Looped Transformer with post-norm does not translate into improved expressivity, while a pre-norm variant diverges at larger depth. This divergence can be attributed to the exponential growth of the activations, apparent in Figure 2(b) . Therefore, to use pre-norm and improve signal propagation in deep looped models, we must first stabilize it, which we do in the following.
### 3.2 Recovering boundedness via residual scaling
While pre-norm mitigates trainability problems in looped architectures, it removes the necessary boundedness condition that motivated the use of post-norm in them. This effect can be observed in Figure 2(b) , where the activations of the pre-norm model grow with deeper effective layer, causing trainability issues observable in Figure 2(a) . Consequently, we propose to restore boundedness by introducing scaling parameters applied at two different scales: one over each sub-layer of the network, and one across the iterations.
##### Layer-wise residual scaling.
Within a single application of f θ ( 𝐳 , 𝐱 ) f_{\theta}\left(\mathbf{z};\mathbf{x}\right) , the residual stream and sub-layer output f θ ℓ ℓ ( 𝐳 ℓ − 1 ) f_{\theta^{\ell}}^{\ell}(\mathbf{z}^{\ell-1}) are weighted by tied scalars ( α 1 , β 1 ) (\alpha_{1},\beta_{1}) shared across all L L layers:
##### Iteration-wise input mixing.
Between consecutive applications of f θ ( 𝐳 , 𝐱 ) f_{\theta}\left(\mathbf{z};\mathbf{x}\right) , we re-inject the input 𝐱 \mathbf{x} with tied scalars ( α 2 , β 2 ) (\alpha_{2},\beta_{2}) shared across all iterations ( Bai et al., 2019 ) :
The two scaling schemes are not independent. With an appropriate coupling between them, the resulting recurrence is bounded for any input , resulting in a stable looping ( Orvieto et al., 2023 ) . In the following statement, we formalize this claim:
Note that the boundedness condition of the sequence model in Theorem 1 is satisfied when using pre-norm, as shown by Kim et al. (2021) . However, boundedness still does not guarantee that the looping to converge to a fixed-point, which we propose to utilize for adaptivity. In the following, we show that there exists a choice of α 2 \alpha_{2} that satisfies this requirement. Consequently, as empirically shown by Bansal et al. (2022) , the model may become locally contractive during training.
###### Theorem 2 (Small α 2 \alpha_{2} implies convergence) .
Let λ f \lambda_{f} be the Lipschitz constant of the L L -layer model defined by Equation 2 , i.e. the map 𝐳 0 ↦ 𝐳 2 L \mathbf{z}^{0}\mapsto\mathbf{z}^{2L} . Then the looped step f θ ( ⋅ , 𝐱 ) f_{\theta}(\,\cdot\,;\mathbf{x}) of Equations 2 – 3 is Lipschitz in its first argument with constant α 2 λ f \alpha_{2}\lambda_{f} . In particular, if α 2 λ f < 1 \alpha_{2}\lambda_{f}<1 , then f θ ( ⋅ , 𝐱 ) f_{\theta}(\,\cdot\,;\mathbf{x}) is a contraction and the iteration 𝐳 i + 1 = f θ ( 𝐳 i , 𝐱 ) \mathbf{z}_{i+1}=f_{\theta}(\mathbf{z}_{i};\mathbf{x}) converges to a unique fixed-point 𝐳 ⋆ = f θ ( 𝐳 ∗ , 𝐱 ) \mathbf{z}^{\star}=f_{\theta}(\mathbf{z}^{*};\mathbf{x}) at a linear rate: ‖ f θ ( 𝐳 i , 𝐱 ) − 𝐳 i ‖ ≤ ( α 2 λ f ) i ‖ f θ ( 𝐳 0 , 𝐱 ) − 𝐳 0 ‖ . \left\|f_{\theta}(\mathbf{z}_{i};\mathbf{x})-\mathbf{z}_{i}\right\|\leq\left(\alpha_{2}\lambda_{f}\right)^{i}\left\|f_{\theta}(\mathbf{z}_{0};\mathbf{x})-\mathbf{z}_{0}\right\|.
The proof is in Section A.2 .
While a contractive map is needed for convergence, an excessively contractive one severely limits expressivity ( Bai et al., 2019 ; Anil et al., 2022 ) . We avoid this by making α 1 \alpha_{1} and α 2 \alpha_{2} learnable. In practice, we find that initializing the network to be more contractive by setting α 2 \alpha_{2} to be small yields better performance ( Section 4.4.2 ). In Figure 9 , we observe that after training, the distribution of α \alpha values widens, with the median very close to the initial point. Intriguingly, these observations are also in line with the common solutions to signal propagation and rank-collapse problems in deep neural networks ( Noci et al., 2022 ; Sun et al., 2025 ) , suggesting a connection between rank-collapse and the signal propagation issue in looped architectures. Together, these modifications yield a pre-norm Looped Transformer that maintains performance over longer looping horizons before saturating, compared to the post-norm variant (see Figure 3 ).
### 3.3 Oscillation around the fixed-point
So far, we have been able to establish that, given a small enough α 2 \alpha_{2} , the model introduced in Equation 3 becomes locally contractive and converges to a fixed-point. However, in practice the contraction factor of Theorem 2 is not itself guaranteed. For some inputs, we observe that the model often descends into an oscillatory behavior, causing the iteration to stay in a small region of latent space without converging. This non-convergent behavior is not in tension with Theorem 2 , as the theorem gives a sufficient condition for convergence, not a complete characterization of the iteration’s behavior. In fact, oscillation around the fixed-point can happen when the Jacobian satisfies certain conditions.
Linearizing the iteration near a fixed-point 𝐳 ⋆ \mathbf{z}^{\star} gives 𝐳 i + 1 − 𝐳 ⋆ ≈ 𝐉 ( 𝐳 i − 𝐳 ⋆ ) \mathbf{z}_{i+1}-\mathbf{z}^{\star}\approx\mathbf{J}\,(\mathbf{z}_{i}-\mathbf{z}^{\star}) , where 𝐉 = ∂ f θ / ∂ 𝐳 | 𝐳 ⋆ \mathbf{J}=\partial f_{\theta}/\partial\mathbf{z}\,|_{\mathbf{z}^{\star}} . Oscillation around 𝐳 ⋆ \mathbf{z}^{\star} arises when 𝐉 \mathbf{J} has an eigenvalue with ℜ ( λ i ) < 1 \Re(\lambda_{i})<1 but | λ i | ≥ 1 |\lambda_{i}|\geq 1 , in which case the iteration spirals around 𝐳 ⋆ \mathbf{z}^{\star} rather than contracting toward it. The half-plane condition ℜ ( λ i ) < 1 \Re(\lambda_{i})<1 is exactly what licenses a runtime fix that does not require modifying f θ f_{\theta} :
Theorem 3 shows that a suitable damping factor η \eta eliminates the oscillations while preserving the fixed-points of f θ ( 𝐳 , 𝐱 ) f_{\theta}(\mathbf{z};\mathbf{x}) . We measure convergence to a fixed-point at iteration i i as r i = ‖ 𝐳 i − f θ ( 𝐳 i , 𝐱 ) ‖ ∞ ‖ f θ ( 𝐳 i , 𝐱 ) ‖ ∞ + ϵ , r_{i}\;=\;\frac{\|\mathbf{z}_{i}-f_{\theta}(\mathbf{z}_{i};\mathbf{x})\|_{\infty}}{\|f_{\theta}(\mathbf{z}_{i};\mathbf{x})\|_{\infty}+\epsilon}, which serves as the halting signal: the iteration stops once r i r_{i} falls below a tolerance τ \tau . To choose η \eta adaptively at inference time, we use a patience mechanism that decreases η \eta whenever this residual stops improving. We track the smallest residual observed so far, r ⋆ = min j ≤ i r j r^{\star}=\min_{j\leq i}r_{j} . A geometric decay η ← γ η \eta\leftarrow\gamma\,\eta with γ ∈ ( 0 , 1 ) \gamma\in(0,1) is applied to the step-size after P P consecutive iterations with no improvement in the residuals. The full procedure is given in Algorithm 1 , which is based on the implementation provided by Movahedi et al. (2025) .
### 3.4 Optimization of fixed-point models
One advantage of contractive fixed-point models is that they can be trained using truncated back-propagation through time (BPTT) ( Geng et al., 2021 ) . Let 𝐉 = ∂ f θ ∂ 𝐳 ( 𝐳 ⋆ , 𝐱 ) \mathbf{J}=\frac{\partial f_{\theta}}{\partial\mathbf{z}}(\mathbf{z}^{\star};\mathbf{x}) and 𝐏 = ∂ f θ ∂ θ ( 𝐳 ⋆ , 𝐱 ) \mathbf{P}=\frac{\partial f_{\theta}}{\partial\theta}(\mathbf{z}^{\star};\mathbf{x}) denote the Jacobians of f θ f_{\theta} at the fixed-point with respect to the state and the parameters, respectively. Following the implicit function theorem we can write the gradient w.r.t. the parameters of the model as ( Bai et al., 2019 ) :
d 𝐳 ⋆ d θ = ( 𝐈 − 𝐉 ) − 1 𝐏 . \frac{d\mathbf{z}^{\star}}{d\theta}\;=\;(\mathbf{I}-\mathbf{J})^{-1}\,\mathbf{P}. (4)
The Neumann series ( 𝐈 − 𝐉 ) − 1 = ∑ j ≥ 0 𝐉 j (\mathbf{I}-\mathbf{J})^{-1}=\sum_{j\geq 0}\mathbf{J}^{\,j} converges, assuming f θ ( 𝐳 , 𝐱 ) f_{\theta}(\mathbf{z};\mathbf{x}) is contractive. Therefore, truncating the series at depth k k yields the estimate:
d 𝐳 ⋆ d θ ≈ ∑ j = 0 k − 1 𝐉 j 𝐏 , \frac{d\mathbf{z}^{\star}}{d\theta}\;\approx\;{\textstyle\sum_{j=0}^{k-1}}\mathbf{J}^{\,j}\,\mathbf{P}, (5) which is closely related to Jacobian-free backpropagation, where the full implicit linear solve is replaced by cheaper approximate gradients ( Fung et al., 2022 ) . The truncation depth k k trades off computation against accuracy. The following proposition bounds the resulting error under contractivity.
###### Proposition 1 (Exponential decay of truncated-BPTT error) .
Let 𝐳 ⋆ = f θ ( 𝐳 ⋆ , 𝐱 ) \mathbf{z}^{\star}=f_{\theta}(\mathbf{z}^{\star};\mathbf{x}) and 𝐉 = ∂ f θ ∂ 𝐳 ( 𝐳 ⋆ , 𝐱 ) ∈ ℝ D × D \mathbf{J}=\frac{\partial f_{\theta}}{\partial\mathbf{z}}(\mathbf{z}^{\star};\mathbf{x})\in\mathbb{R}^{D\times D} . If 𝐉 \mathbf{J} is contractive in spectral norm, ‖ 𝐉 ‖ 2 = σ < 1 \|\mathbf{J}\|_{2}=\sigma<1 , then for every k ≥ 0 k\geq 0 , ‖ ( 𝐈 − 𝐉 ) − 1 − ∑ j = 0 k − 1 𝐉 j ‖ F ≤ D σ k 1 − σ . \left\|(\mathbf{I}-\mathbf{J})^{-1}-{\textstyle\sum_{j=0}^{k-1}}\mathbf{J}^{\,j}\right\|_{F}\;\leq\;\sqrt{D}\,\frac{\sigma^{k}}{1-\sigma}.
The proof is in Section A.4 . An essentially equivalent result appears in the proof of Theorem 2 of Geng et al. (2021) .
Proposition 1 allows for a fixed memory footprint during training, essentially decoupling the number of loops from the memory complexity of the model. In the same spirit, HRM ( Wang et al., 2025 ) and TRM ( Jolicoeur-Martineau, 2025 ) approximate the gradient with a small number of backward passes at the fixed-point, although TRM argues that the fixed-point condition is unnecessary in practice.
### 3.5 The fixed-point reasoning model
So far, we introduced modifications to improve signal propagation in Looped Transformers ( Section 3.1 ) without sacrificing stability ( Section 3.2 ). These modifications yield fixed-point iterations that can be made non-oscillatory ( Section 3.3 ) and trainable through truncated-BPTT ( Section 3.4 ), making adaptivity through fixed-points reliable. We assemble these components into FPRM, summarized in Algorithm 2 and illustrated in Figure 3 . The result is a Looped Transformer that iterates until its hidden state converges, spending compute proportional to each input’s difficulty—the adaptivity our halting mechanism was designed to provide. Since we observed no improvements in our experiments by using the hierarchical structure of Wang et al. (2025) , we opt for a classic looped architecture ( Dehghani et al., 2019 ) instead. An overview of the framework is available in Algorithm 2 . In Appendix C we provide further details about FPRM.
During training, the model performs looping in windows of k k iterations, with k k being a hyperparameter, which determines the truncated-BPTT value. During inference, we set k = 1 k=1 . At each forward pass, the fixed-point optimizer introduced in Algorithm 1 is called to dampen the fixed-point iteration steps. Then, we get a prediction from the current state 𝐳 \mathbf{z} of the model and perform a deep supervision step, following Wang et al. (2025) ; Jolicoeur-Martineau (2025) . To truncate the computation graph between deep supervision steps during training, we detach the state 𝐳 \mathbf{z} from the graph and stop when the optimizer detects fixed-points. This happens when the residual falls below the tolerance, or the step-size becomes too small.
## 4 Experiments
We evaluate FPRM against looped reasoning models on puzzle, adaptivity, and signal-propagation benchmarks. Our implementation builds on the public TRM codebase ( Jolicoeur-Martineau, 2025 ) and adopts its deep-supervision training procedure. Additional experimental details are provided in Appendix G . We first describe the datasets, then evaluate puzzle-solving performance, adaptivity to task difficulty, and depth-induced signal-propagation effects.
### 4.1 Dataset description
In the following, we provide a brief description for each dataset used in this paper. We refer the reader to the cited literature for more information.
##### Sudoku-Extreme.
The task consists of exceptionally challenging, partially filled 9 × 9 9\times 9 Sudoku puzzles with unique solutions, introduced by Wang et al. (2025) . Each sample is flattened into a sequence of length 81 81 . The train data consists of 1000 1000 unique samples, each augmented 1000 1000 times, giving a total of ~ 1 1 M train samples. The test data contains 422,786 422{,}786 samples. The evaluation metric is exact (sequence) accuracy.
##### Maze-Hard.