We consider the stability of multi-step reasoning processes, which have extensive applications in language models, including chain-of-thought and algorithmic reasoning. While longer sequences of reasoning can improve a model's generation capability at test time, the errors due to intermediate reasoning steps can accumulate in autoregressive generation, and thus grow substantially at the end. In this paper, we ask: What are the key factors determining the stability of multi-step reasoning? First, we show an inference error bound governed by the product of spectral norms of the Jacobians taken through the input space across generation steps. This product can be viewed as an error amplification factor, which could scale exponentially with the number of reasoning steps, serving as a quantitative measure of reasoning stability. Second, we analyze this measure in transformer models trained to predict simple tasks like linear and quadratic functions. We theoretically prove that the transformer model converges to a solution where the stability measure decays, thus yielding nearly zero inference loss over (arbitrarily) long steps. Finally, the stability analysis leads to several algorithmic implications for controlling the stability, through (i) chain-of-thought length compression that reduces the sensitivity of each step, and (ii) quantization-aware training that regularizes the input Jacobian norms. We validate the proposed algorithms by fine-tuning language models on graph-algorithmic reasoning tasks and symbolic state-tracking tasks. Across seven evaluations, our algorithms improve over baseline comparisons by 3.5% on average, and by 8.2% for longer-length inputs. Ablation analysis validates that the stability measure is drastically reduced by 3-8×, confirming the regularization effect on the spectral norms of the (input space) Jacobians.
Figures & tables
Figure 1 : We present a bound on the inference loss in multi-step reasoning governed by the product of spectral norms of Jacobians taken through the input space across generation steps. Figure 1 : Our bound (equation ( 4 )) tracks the inference loss on the Bellman-Ford algorithmic reasoning task, showing that both can scale exponentially with steps. Figure 1 : We design a training method that regularizes the Jacobian spectral norms, which improves the test performance over baselines. Figure 1 : Models fine-tuned with our method exhibit drastically reduced Jacobian spectral norms.
Figure 2 : We illustrate the loss gap and our bound in equation ( 4 ), which can scale exponentially across generation steps, tracking the growth of the inference loss. Results are evaluated on breadth-first search, Dijkstra, and Prim’s algorithm, using QWen-1.5B. The loss gap and the bound are plotted on separate axes; equation ( 4 ) is a worst-case bound and is numerically larger than the gap it bounds. Table 2 in Appendix B reports the corresponding bound-to-gap ratios and constants.
Figure 3 : We illustrate that, in training transformers on both linear and quadratic functions, the inference loss and the noise-perturbed loss converge close to zero, and become lower with longer steps. We set d=10 and train a one-head linear transformer on linear functions ( n=20 , η=0.4 ) and a three-head non-linear transformer on quadratic functions ( n=200 , η=10−3 ). For the random Gaussian noise, we set σ=0.002 . On the right, we consider a slightly different setting by first adding noise to the weights before computing the gradients.
Figure 4 : Figure 4 and 4 : We show that the models trained with reduced intermediate steps yield lower inference losses than the model trained with longer steps. Correspondingly, the models exhibit lower Jacobian-based error amplification factors. Figure 4 and 4 : We show that quantization-aware training with lower bit-widths further reduces the Jacobian-based error amplification factors, yielding lower inference losses.
Dijkstra
Bellman-Ford
Prim
BFS
DFS
Cyclic
Symmetry
No CoT
53.3 ± 2.0
69.9 ± 0.3
56.2 ± 0.1
85.9 ± 0.1
71.4 ± 1.4
66.8 ± 0.9
48.5 ± 3.1
SFT-CoT
73.1 ± 1.4
75.3 ± 1.5
62.3 ± 0.1
86.0 ± 0.4
54.6 ± 2.0
89.2 ± 1.2
55.8 ± 1.5
Implicit CoT
71.6 ± 0.5
72.3 ± 0.6
64.1 ± 0.3
85.5 ± 1.7
72.4 ± 0.4
75.4 ± 1.0
60.8 ± 3.6
Coconut
74.1 ± 0.3
75.6 ± 0.5
63.5 ± 1.2
85.7 ± 0.3
72.6 ± 0.7
79.5 ± 0.1
53.9 ± 0.5
NSO
74.5 ± 0.8
75.3 ± 0.0
63.2 ± 0.3
86.5 ± 0.8
55.8 ± 0.1
89.3 ± 0.8
57.2 ± 0.9
Algorithm 1
76.2 ± 1.7
80.3 ± 0.1
66.2 ± 0.3
88.3 ± 1.3
74.4 ± 1.5
94.1 ± 1.3
62.2 ± 0.8
Table 1 : We report the test accuracy (%) of our method, as compared to supervised fine-tuning, implicit reasoning, and random noise injection. We evaluate five algorithmic reasoning tasks and two symbolic state-tracking reasoning tasks. For length generalization, models are evaluated on inputs 10% longer than the training sequences. All results are averaged over three random seeds.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Terms of equation ( 4 )
Constants
Model
T
Loss gap
Bound
Ratio
First-order
Second-order
C1
C2
Qwen-1.5B
1
0
0
–
0
0
0.01
0
2
0.08
0.68
8.0
0.68
9.5×10−4
0.02
5.44
3
0.25
1.5
6.1
1.3
1.7 ×10−1
0.03
5.83
4
0.47
3.4
7.2
2.5
8.4 ×10−1
0.05
5.83
5
0.83
6.9
8.3
4.5
2.4 ×100
0.07
5.83
Appendix
Table 2 : Decomposition of the bound in equation ( 4 ) for three models trained on the Bellman-Ford task. Loss gap denotes ∣L(fW;D)−L⋆(fW;D)∣ , Bound denotes the right-hand side of equation ( 4 ), and Ratio denotes the bound value divided by the loss gap. The first- and second-order terms are C1∑t=1T−1ρt(T)εt and 2C1C2∑t=2Tρt(T)(∑i=1t−1ri)2 . The absolute scale of the bound varies by orders of magnitude across architectures. At T=1 there is no accumulated error, so the loss gap and the bound value are zero.
Task
Bit-width
Test accuracy (%)
Loss gap
Bound
∑tρt(T)
∑tεt
Bellman-Ford
1
69.6
2.0
21.8
143.6
3.6
2
77.3
1.5
12.1
157.2
1.9
3
79.5
1.4
10.8
159.6
1.9
4
80.3
0.6
7.3
166.7
1.8
32
75.3
1.8
13.4
177.0
2.3
Dijkstra
1
76.2
0.9
2.1
9.3
2.9
Appendix
Table 3 : Decomposition of equation ( 4 ) under quantization-aware training, for Qwen-1.5B at T=5 . Lowering the bit-width monotonically reduces the error amplification factor ∑t=1T−1ρt(T) and increases the per-step prediction errors, summarized by ∑t=1T−1εt ; the bound and the test accuracy are determined by the balance of the two. The best bit-width per task is in bold.
Terms of equation ( 4 )
Constants
Model
Method
Loss gap
Bound
First-order
Second-order
C1
C2
Llama-1B
SFT-CoT
1.84
4.3×106
5.0×104
4.2×106
0.050
3.21×103
Coconut
0.699
1.1×105
1.3×103
1.1×105
0.077
1.63×103
Alg. 1
0.491
1.3×104
7.4×102
1.2×104
0.091
1.79×103
Gemma-2B
SFT-CoT
2.11
3.2×105
2.7×103
3.2×105
0.060
1.37×103
Coconut
0.662
1.5×104
1.34×102
1.5×104
0.094
9.27×102
Appendix
Table 4 : The error amplification factor tracks the inference loss across base models. We fine-tune Llama-1B and Gemma-2B on Bellman-Ford ( T=5 ) and report both terms of equation ( 4 ) together with the estimated constants; columns are as in Table 2 . Within each model, the ranking of methods by our bound matches their ranking by the loss gap, even though the absolute scale of the bound differs by orders of magnitude between the two models.
Figure 5 : We illustrate the loss gap and our bound in equation ( 4 ), which can scale exponentially across generation steps, closely tracking the growth of the inference loss. Results are evaluated on Bellman-Ford and Prim’s algorithm, using Gemma-2B and Llama-1B models. The two quantities are plotted on separate axes and differ in absolute scale; see Table 2 for the bound-to-gap ratios.
Figure 6 : We illustrate how transformers realize multi-step gradient descent behavior when trained on linear and quadratic functions. We verify that the Jacobian matrix norm is a converging series on linear functions and quadratic functions. The Jacobian norm of linear functions is ∑i=1T∥(I−nηXX⊤)i∥22 and the Jacobian norm of quadratic functions is ∑i=1T∥(I−n2ηXD(θi)X⊤)i∥22 .
Task
Quantization
Optimizer
Learning rate
LoRA Rank
Sampling factor λ
Latent steps k
Dijkstra
1-bit
AdamW
2×10−5
16
0.1
1
Bellman-Ford
4-bit
AdamW
2×10−5
16
0.2
1
Prim
4-bit
AdamW
2×10−5
16
0.1
1
BFS
1-bit
AdamW
2×10−5
16
0.1
1
DFS
3-bit
AdamW
2×10−5
16
0.1
1
Cyclic
1-bit
AdamW
2×10−5
16
1.0
1
Appendix
Table 5 : Summary of the hyperparameters used for each reasoning task. The exact configurations are selected from the defined search spaces during supervised fine-tuning.