PaperlayerView exact revision on arXiv

Source-linked AI summary

Quantifying the Stability of Multi-Step Reasoning via Error Amplification

Dongyue Li, Ziniu Zhang, Minxuan Duan, Hongyang R. Zhang

arXiv:2610.06404v1cs.LGcs.AIstat.ML

TL;DR

Intermediate errors can accumulate during multi-step reasoning, while standard fine-tuning does not explicitly control that accumulation. The paper bounds inference error using input-Jacobian norms, develops methods to control them, and reports improved accuracy across seven tasks.

  • Problem

    Intermediate errors can accumulate in autoregressive reasoning, but standard supervised fine-tuning does not explicitly control this accumulation.

    Source: §1 · p. 1

  • Method

    The paper bounds error amplification by products of input-Jacobian spectral norms and proposes length compression and quantization-aware training to control that factor.

    Source: §1 · p. 3, §7 · p. 10

  • Results

    3.5% average test accuracy improvement over baselines across seven tasks, with an 8.2% improvement on larger-input length generalization evaluations.

    Source: §1 · p. 3

  • Takeaways & Limitations

    For linear and quadratic weight-prediction tasks, transformers converge to solutions with Jacobian norms below one and polynomially small test loss at T = Θ(log d) steps.

    Source: §1 · p. 2

    The error bound is a worst-case upper bound that may not be numerically tight, although its values correlate consistently with inference losses across models.

    Source: §3.1 · p. 5

Abstract

from arXiv · show

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$\times$, confirming the regularization effect on the spectral norms of the (input space) Jacobians.

1 Introduction

The paper studies how intermediate prediction errors affect multi-step reasoning stability, then analyzes their amplification and proposes training methods to control it. The results combine theoretical guarantees for linear and quadratic prediction tasks with empirical improvements across seven reasoning evaluations.

1 Introduction

  • Autoregressive errors can destabilize long reasoning, while standard supervised fine-tuning does not explicitly control their accumulation.The motivation includes empirical super-linear growth of test error with steps in fine-tuned language models.Source: §1 · pp. 1–2
  • Prior theoretical analyses focus on task-specific transformer structures or whole-class VC-dimension bounds, motivating a computable analysis of multi-step reasoning.Source: §1 · p. 2
  • The inference-loss bound is governed by products of input-space Jacobian spectral norms, and empirical bounds track losses that can grow exponentially with reasoning steps.The analysis frames this product as an error amplification factor.Source: §1 · p. 2
  • For linear and quadratic weight-prediction tasks, transformers converge to a regime with decaying amplification and test loss O(1/poly(d)) at T = Θ(log d).The quadratic analysis uses a three-head transformer construction that exactly implements gradient descent.Source: §1 · p. 2
  • The proposed training method combines chain-of-thought length compression with quantization-aware training to regularize Jacobian-based error amplification.Quantized training uses gradients computed on the quantized model and regularizes the Hessian trace of the model loss.Source: §1 · pp. 2–3
  • 3.5% average test accuracy improvement over baselines across seven tasks, and 8.2% over baselines on larger inputs, were reported for the proposed method.The evaluations cover five algorithmic reasoning tasks and two state-tracking tasks.Source: §1 · p. 3

2 Preliminaries

The preliminaries define multi-step prediction as autoregressive generation over a ground-truth sequence. They distinguish training on correct intermediate steps from inference on model-generated steps and state the stability objective.

2 Preliminaries

  • A next-token predictor takes a sequence of vectors and outputs one vector, which is appended before the next prediction; task sequences have input-dependent step counts.Source: §2 · p. 3
  • Training conditions predictions on preceding ground-truth steps, whereas inference feeds back previous model predictions and evaluates loss at the final step.Source: §2 · p. 3
  • Because inference feeds earlier prediction errors back into later steps, errors can accumulate in long generation; the stated goal is to bound inference loss and improve stability.Source: §2 · p. 4

3 Theoretical Analysis

The analysis bounds inference loss using an error amplification factor built from input-space Jacobian spectral norms. In linear and quadratic function tasks, transformers converge to solutions where this factor decays, yielding low inference loss.

3.2 Learning Dynamics in Quadratic Models

  • For linear and quadratic tasks, transformers converge to solutions where the Jacobian-based stability measure decays, yielding test loss O(1/poly(d)) at T = Θ(log d) steps.The analysis attributes this to accumulated errors forming a converging geometric series.Source: §3 · p. 4, §3.2 · p. 6

3.1 Error Amplification in Multi-Step Reasoning

  • The inference-loss bound is governed by products of input-Jacobian spectral norms, which form an error amplification factor that can grow exponentially with reasoning steps.The factor sums products along propagation paths from earlier steps to later ones.Source: §3.1 · pp. 4–5, §3 · p. 4
  • Across graph-algorithm reasoning tasks, inference loss and the bound grow with step count, and the bound closely tracks that growth.The evaluation includes Qwen-1.5B, Llama-1B, and Gemma-2B models.Source: §3.1 · p. 5
  • The bound is a worst-case upper bound that correlates with inference loss but can be numerically loose, with bound-to-gap ratios varying substantially across models.Reported ratios are around 8 for Qwen-1.5B and around 105–106 for Llama-1B and Gemma-2B.Source: §3.1 · p. 5

3.2 Learning Dynamics in Quadratic Models

  • For linear and quadratic tasks, numerical experiments show inference and noise-perturbed losses converging near zero and decreasing with longer steps.The quadratic setting also tests adding noise to weights before computing gradients.Source: Figure 3 · p. 6, §3.2 · p. 6
  • A three-head transformer with cubic and squared activations can exactly implement gradient descent updates for quadratic functions, and training converges to this construction.The converged solution yields near-zero inference loss.Source: §3.2 · p. 6

4 Proposed Methods

The proposed training method combines chain-of-thought length compression and quantization-aware training to control accumulated reasoning error by reducing Jacobian-based error amplification.

4 Proposed Methods

  • Length compression uniformly subsamples T′ = ⌊λ · T⌋ intermediate steps and minimizes training loss over the resulting sequence.The sampling factor λ sets the compressed sequence length.Source: §4 · p. 7
  • Subsampling reduces error amplification by 92× and 181× versus all-step training; intermediate step counts yield the lowest inference loss on Dijkstra and Prim.The figure compares reduced-step and longer-step models, with inference loss and Jacobian-based amplification moving together.Source: §4 · p. 7, Figure 4 · p. 8
  • Subsampling lowers amplification but raises per-step prediction errors, so the best λ balances these effects.The method combines length compression and quantization-aware training in its training algorithm.Source: §4 · p. 7
  • Quantization-aware training applies uniform quantization to weights during fine-tuning to regularize Jacobian spectral norms.The proposed rationale links minimizing expected loss under randomized-rounding noise with penalizing the loss Hessian trace and input-space Jacobians.Source: §4 · p. 7
  • With 1-bit quantization, error amplification falls 3.6× on Dijkstra and 5.8× on Prim compared with full precision.Lower bit-widths also yield lower inference losses in the reported comparison.Source: §4 · p. 7, Figure 4 · p. 8
  • For out-of-distribution length generalization, covariate shift can raise per-step errors, and extra steps geometrically multiply errors when step-to-step Jacobian norms exceed one.The authors identify length generalization as a future research direction.Source: §4 · p. 8

5 Experiments

Across seven graph-algorithmic and symbolic reasoning tasks, the method improves test accuracy and performance on longer inputs. Experiments also report reduced error amplification, component ablations, and a boundary on longer-length evaluation.

5.1 Experimental Setup

  • Evaluation used five CLRS-Text graph algorithms and two LEGO state-tracking tasks, with Qwen-1.5B fine-tuning and 2,000 training examples per task.Baselines included No CoT, SFT-CoT, implicit CoT, Coconut, and noise stability optimization.Source: §5.1 · pp. 8–9

5.2 Experimental Results

  • 3.5% average test-accuracy improvement across seven tasks and 8.2% average improvement on inputs 10% longer than training sequences.Table 1 covers five graph-algorithmic and two symbolic state-tracking tasks; results are averaged over three random seeds.Source: Table 1 · p. 9, §5.2 · p. 9

5.1 Experimental Setup

  • At 20% longer inputs, accuracy fell below 25% for every method, so the paper reports the 10% longer-input setting as its main result.The authors leave further improvement on longer sequences for future work.Source: §5.1 · p. 9

5.2 Experimental Results

  • 3–8× lower error-amplification factors than baselines; on Qwen-1.5B, the factor fell 43% and inference loss fell 72% on average.The Qwen-1.5B comparison covers three tasks; the passage does not specify the metric associated with the 43% figure beyond the factor reduction.Source: §5.2 · p. 9
  • 77% lower bound and 58% lower empirical test-loss gap than Coconut on average across Llama-1B and Gemma-2B.The ordering of methods by bound value matched their ordering by empirical loss gap on both models; absolute bound scales differed sharply between architectures.Source: §5.2 · pp. 9–10
  • Removing length compression or quantization-aware training degraded performance by 4.3% and 3.2%, respectively; quantization noise outperformed Gaussian noise by 9%.Random sampling of intermediate steps performed slightly better than removing the first, last, or equally spaced steps.Source: §5.2 · p. 10

6 Related Work

Related work spans theoretical analyses of transformer reasoning and generalization, studies of chain-of-thought expressiveness, and methods for improving multi-step reasoning. The cited approaches differ in their tasks, architectures, and analytical focus.

Theoretical understanding of multi-step reasoning

  • Transformer analyses study task-specific settings including parity checks, state tracking, tree path-finding, and in-context linear regression [13] [30] [14] [34] [12].The cited analyses examine different attention architectures and reasoning tasks.Source: §6 · p. 10

Generalization of language models

  • PAC-learning analyses bound sample complexity using hypothesis-class VC dimension and CoT length [15], while CoT-information yields faster rates [1].Other language-model generalization work studies compression, subspace training, martingale-based bounds, and PAC-Bayes bounds tied to loss Hessians [23] [24] [7] [17] [16].Source: §6 · p. 10

Expressiveness and optimization methods

  • Prior work examines CoT expressiveness and scaling [22] [25] [27], latent-variable and implicit-reasoning methods [35] [5] [10] [11], and how training computation and reward shaping relate to long chains [36].Multi-step traces also appear in experimental-design settings where trajectory data is revealed over time [33].Source: §6 · p. 10

7 Conclusion

The paper concludes that input-Jacobian spectral norms govern multi-step error amplification, and reports theoretical and empirical evidence for controlling it. The appendix outlines proofs of the error bound and transformer convergence results, including a recursive Taylor-expansion argument for the bound.

7 Conclusion

  • The paper bounds inference errors using products of input-Jacobian spectral norms, proves decay of the amplification factor for in-context weight prediction, and introduces compression plus quantization-aware training.Across seven reasoning tasks, the approach improves test accuracy and length generalization while reducing empirical amplification factors.Source: §7 · p. 10

A Complete Proofs

  • The appendix supplies proofs for the error bound and for linear- and quadratic-function weight-prediction analyses, including a three-head construction implementing gradient descent on quadratic losses.It derives chain-of-thought training dynamics and proves convergence in Theorem 3.4.Source: A Complete Proofs · p. 13

A.1 Proof of Lemma 3.1

  • The proof distinguishes predictions conditioned on ground-truth preceding steps from predictions using autoregressively generated prefixes, defining their deviations from each step's label.These deviations provide the basis for bounding the prediction gap caused by autoregressive accumulation.Source: §A.1 · p. 13
  • The remainder bound assumes a C2-Lipschitz Jacobian, while the loss-gap argument assumes the loss is C1-Lipschitz in its first argument.The Jacobian smoothness assumption bounds the Taylor remainder quadratically in the input deviation.Source: §A.1 · pp. 13–14
  • The Lemma 3.1 proof expands prediction gaps with first-order Taylor approximation and recursively combines Jacobian spectral norms along step-to-step paths into a cumulative amplification factor.The step-to-step Jacobian is evaluated at the ground-truth context, and the expansion includes a higher-order remainder.Source: §A.1 · p. 14

A.2 Proof for the Linear Case

The linear-case analysis studies a one-layer linear self-attention transformer trained to predict weights of linear functions, with gradient-descent iterates forming the chain of thought. Its main result bounds the training-time population loss under specified initialization and training conditions.

A.2 Proof for the Linear Case

  • The task samples a Gaussian ground-truth weight vector and Gaussian inputs, with labels given by their inner products.The transformer predicts the weight vector from the in-context examples.Source: §A.2 · p. 15
  • The chain of thought consists of gradient-descent iterates on the average squared loss over the in-context examples, starting from θ0 = 0.The final step is labeled with the ground-truth weight vector θ⋆.Source: §A.2 · p. 15
  • Theorem A.5 gives a population-loss bound after gradient flow under specified sample-size, learning-rate, step-count, and initialization conditions.The theorem states that the loss can be made at most ϵ1 for an indicated range of ϵ1.Source: §A.2 · p. 17

A.2.1 Proof of Theorem A.5

For the linear case, the proof decomposes inference error into stepwise prediction errors propagated by powers of a single Jacobian, then bounds the Jacobian and residuals at convergence. The request also includes the quadratic-case setup, where weight prediction uses labels that are squared inner products.

A.2.1 Proof of Theorem A.5

  • The proof writes inference error as stepwise prediction errors carried to the final step by powers of the Jacobian; the linear model's affine map has zero second-order remainder.The Jacobian is shown close to that of the gradient-descent update, I−ηS.Source: §A.2.1 · pp. 17–18
  • For the trained transformer produced by gradient flow on population loss over chain-of-thought sequences, J=J0+E yields ∥J∥₂ ≤ ∥J0∥₂+∥E∥₂ ≤ 1−η+O(log^(−5/2)d)+Θ(d^(−c log(1/(1−η))+1)).Source: §A.2.1 · p. 18
  • At convergence, the linear transformer's Jacobian is contractive and its per-step residuals are small, so summing their propagated contributions bounds inference loss.The proof bounds the residual expectation by O(d^(2−a)) and the Jacobian norm by 1−η/2 on the stated event.Source: §A.2.1 · pp. 17, 19

A.3 Proof for the Quadratic Case

  • For the quadratic case, the task predicts a Gaussian weight vector from examples labeled by squared inner products, with the initial estimate set to a scalar multiple of the target.The initialization uses θ0 = α0θ⋆ for α0 in [3/4, 1).Source: §A.3 · p. 19

A.3.1 Proof of Lemma 3.3

This section develops the proof of Lemma 3.3 by analyzing gradient-computation invariances, reducing the transformer to an equivalent model, and deriving its gradients. It constructs a three-head nonlinear self-attention transformer that reproduces gradient-descent updates and establishes local contraction under stated conditions, including that initialized zero blocks remain zero under gradient flow.

A.3.1 Proof of Lemma 3.3

  • A constructed three-head nonlinear self-attention transformer reproduces each gradient-descent parameter update, so its generated output after step i equals the parameter after i steps.The construction performs a gradient-descent step for quadratic squared loss.Source: §A.3.1 · p. 20
  • Under the stated basin, sample-size, and stepsize conditions, gradient descent contracts distance to {±θ⋆}, with the required local regularity holding with probability at least 1 − d^-10.The proof uses a uniformly well-conditioned empirical Hessian near ±θ⋆; its descent argument keeps the iterates in the local basin.Source: §A.3.1 · p. 22

A.3.2 Gradient Computation

  • The subsection derives gradients for the full model’s value and key-query matrices by applying the chain rule to the loss at each chain-of-thought step.The derivation sums the per-step contributions and takes expectations.Source: §A.3.2 · p. 23
  • Under Assumption A.8, gradient flow preserves the initialized zero blocks, and the expected gradients of irrelevant blocks are zero by Gaussian sign-flip symmetry.The proof establishes the zero-gradient property for irrelevant blocks and states that the other cases follow similarly.Source: §A.3.2 · pp. 24, 25, 28
  • Under T, θ_i remains invariant, g_i changes sign, and the resulting gradient quantity is invariant by distributional invariance.The proof uses the deterministic gradient-descent construction of θ_i and the target increment, along with sign-flipping and evenness properties.Source: §A.3.2 · p. 28
  • Because final loss depends only on output entries d + 2 through 2d + 1, pruning irrelevant zero blocks yields an equivalent transformer with the same dynamics.The proof then writes the model in reduced form using those entries of the last column of Z_i.Source: §A.3.2 · pp. 28–29
  • For the reduced model, the proof expresses the CoT training loss as separate last-step-output and ground-truth-weight terms, then derives gradients using the weight-sharing mechanism.It presents expectation-form gradients for W_PV and begins the corresponding gradient decomposition for W_KQ.Source: §A.3.2 · pp. 29–30

A.3.3 Proof of Theorem 3.4

The quadratic-case proof reduces training dynamics to ordinary differential equations for two eigenvalues of the transformer weights. It establishes convergence to the construction and relates that convergence to the inference-loss rate.

A.3.3 Proof of Theorem 3.4

  • Symmetry and isotropic data reduce the weight dynamics to ordinary differential equations for eigenvalues λV(t) and λK(t).The proof characterizes the gradient of the chain-of-thought objective as main terms plus interaction terms before deriving the eigenvalue dynamics.Source: §A.3.3 · pp. 31–32
  • Outside a small neighborhood of (1, 1), a lower bound on the gradient norm yields exponential decrease of the chain-of-thought loss until the trajectory enters that neighborhood.The proof gives an entry time of O(log d).Source: §A.3.3 · p. 33
  • Inside the neighborhood, smoothness and a positive-definite Hessian yield a local PL condition for the chain-of-thought loss.This supports the second stage of the convergence argument.Source: §A.3.3 · p. 33
  • The quadratic-case proof establishes convergence to the construction, then applies the linear-case analysis to obtain the inference-loss rate at training convergence.The proof analyzes convergence in two stages, first entering a neighborhood of the construction and then converging locally.Source: §A.3.3 · p. 33

B.1 Evaluations of the Stability Measure

The section evaluates the stability bound across models, tasks, and training settings, finding that it tracks inference-loss behavior while exposing trade-offs in bound tightness and quantization.

B.1 Evaluations of the Stability Measure

  • The bound-to-gap ratio stays between 6 and 9 for Qwen-1.5B but reaches 2.3 × 10^6 for Llama-1B and 1.5 × 10^5 for Gemma-2B at T = 5.C2 differs by roughly three orders of magnitude across models; the second-order term dominates for Llama-1B and Gemma-2B from T = 3, while the first-order term remains larger for Qwen-1.5B.Source: §B.1 · p. 33, Table 2 · p. 34
  • For Bellman-Ford and Prim’s, the error amplification factor scales exponentially with generation steps and closely tracks inference-loss growth, although the plotted quantities differ in absolute scale.Figure 5 evaluates Gemma-2B and Llama-1B and plots the quantities on separate axes.Source: §B.1 · p. 33, Figure 5 · p. 35
  • On Dijkstra and Prim, 1-bit quantization lowers the bound by 6.8× and 2.5× and improves accuracy by 3.1% and 4.0%; on Bellman-Ford, both are best at 4 bits.Lower bit-widths reduce the amplification factor but increase per-step prediction errors; the best bit-width depends on the task and is selected on validation data.Source: §B.1 · pp. 33–35
  • On Bellman-Ford at T = 5, the method attains the lowest loss gap and bound on Llama-1B and Gemma-2B; versus Coconut, it reduces them by 58% and 78% on average.Compared with SFT-CoT, it lowers the bound by 330× and 65×.Source: Table 4 · p. 35, §B.1 · p. 35
  • For linear and quadratic functions, the Jacobian-product norm converges as reasoning length increases, consistent with the analyzed gradient-descent update rule.The numerical evaluation reports this behavior for both settings.Source: §B.1 · p. 35, Figure 6 · p. 36

B.2 Sample Complexity Comparison

The appendix reports that more training samples and intermediate reasoning steps improve measured robustness or sample complexity in the evaluated algorithmic tasks. It also notes that formal extension of the bounds to reinforcement learning and broader language-model tasks remains open.

B.2 Sample Complexity Comparison

  • Increasing training samples improves noise stability across evaluated tasks; Dijkstra models have more than 2× lower Hessian traces than Bellman-Ford models, and Kruskal models have 15× lower traces than Prim models.The stability evaluation injects isotropic Gaussian weight noise with σ = 0.005 and averages the loss increase over 10 random seeds.Source: §B.2 · p. 36
  • Using intermediate steps consistently achieves better sample complexity than training without them, and a larger number of steps substantially reduces sample complexity.The study subsamples T to 20, 10, 5, 2, and 0, with T = 0 denoting training without intermediate steps.Source: §B.2 · p. 36

B.3 Discussions

  • The analysis applies to reinforcement-learning-trained models because test-time generation is independent of the training objective, but formal extension of the bounds to reinforcement learning is left for future work.The authors propose studying whether policy-gradient training implicitly regularizes input Jacobian norms along high-reward trajectories.Source: §B.3 · pp. 36–37
  • The method regularizes Jacobian norms during training; input-adaptive test-time regularization, stability proxies for black-box models, and analysis of broader language-model tasks remain future directions.The broader tasks named include tool use, generation, and question answering.Source: §B.3 · p. 37
Loading 2610.06404v1…