Paper deep dive

Quantifying the Stability of Multi-Step Reasoning via Error Amplification

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

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 177

Intelligence

Status: not_run | Model: - | Prompt: - | Confidence: 0%

Entities (0)

No extracted entities yet.

Relation Signals (0)

No relation signals yet.

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract: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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

173,625 characters extracted from source content.

Expand or collapse full text

Quantifying the Stability of Multi-Step Reasoning via Error Amplification Dongyue Li Affiliation: Northeastern University Affiliation: Boston, MA Email: li.dongyu@northeastern.edu Ziniu Zhang Affiliation: Northeastern University Affiliation: Boston, MA Email: zhang.zini@northeastern.edu Minxuan Duan Affiliation: Northeastern University Affiliation: Boston, MA Email: dminxuan@gmail.com Hongyang R. Zhang Affiliation: Northeastern University Affiliation: Boston, MA Email: ho.zhang@northeastern.edu Abstract 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 (i) 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. 1 Introduction Multi-step reasoning enhances language model capabilities by generating intermediate steps before producing a final answer. Specifically, for an input x, a model fWf_W generates a reasoning sequence (,y^(1),…,y^(T))(x, y^(1),…, y^(T)) autoregressively. As the error in previous outputs propagates over the steps, this can destabilize long reasoning and cause severe final prediction errors [57, 72]. However, standard training like supervised fine-tuning does not explicitly control this error accumulation. In this work, we theoretically analyze such error amplification across steps, aiming to identify the key factors governing reasoning stability and design training methods that control accumulated errors. Understanding such error amplification is valuable in many practical problems, especially for scaling test-time computation [62, 60]. One concrete setting is algorithmic reasoning, where models perform the step-by-step computation of classical algorithms like Bellman-Ford [70]. Models are expected to learn these underlying algorithms to generalize and extrapolate to larger, unseen inputs. However, a mistake in the early steps can corrupt the procedure. In our own experiments, we observe that the test error in a fine-tuned language model (e.g., Qwen-1.5B and Gemma-2B) can scale super-linearly with steps, as shown in Figure 1. Empirical methods, such as training models to generate fewer steps [51] or reason in the continuous feature space [63], have been proposed to reduce these errors. Yet, the reason for their empirical benefits remains underexplored. (a) Evaluation of our bound values (b) Inference losses (c) Spectral norm of Jacobians 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. Recent theoretical work has presented convergence and generalization guarantees of transformers trained on specific reasoning tasks, including linear regression [65, 78], state tracking [67], and path-finding in trees [75]. Their analyses focus on the task-specific transformer structures, leaving open the general cases. Studies based on PAC learning theory [68, 58] yield VC-dimension generalization bounds dependent on the entire hypothesis class. This motivates us to conduct a general analysis of multi-step reasoning with computable measures on data and fine-tuned models. We begin by analyzing the inference error amplification of multi-step reasoning. Our key finding is that the inference loss is governed by the products of the spectral norms of Jacobians taken through the input space across generation steps, which control reasoning stability. The idea is to bound the gap between the T-step inference loss, ℓ⁡(fW([,y^(1),…,y^(T−1)]),y(T)) (f_W([x, y^(1),…, y^(T-1)]),y^(T)), and the training loss, ℓ⁡(fW([,y(1),…,y(T−1)]),y(T)) (f_W([x,y^(1),…,y^(T-1)]),y^(T)), conditioned on the correct steps. Using a first-order approximation in the input space, we show that the prediction errors at step i are propagated forward by the Jacobian matrix Ji(t)J_i^(t) (evaluated at the t-th step output to the input y(i)y^(i)). Let ρi(T) _i^(T) be the sum of products of ‖Ji(t)‖\|J_i^(t)\| across all paths from step i to T. Then, recursively decomposing the errors yields a bound for the inference loss governed by ∑i=1T−1ρi(T) _i=1^T-1 _i^(T) (cf. Lemma 3.1), which we define as an error amplification factor. Empirically, this bound closely tracks the test losses on fine-tuned language models (including Qwen, Gemma, and Llama models), which can scale exponentially with T, as shown in Figure 1. Next, we instantiate this bound by analyzing one-layer multi-head transformers trained on weight-prediction tasks for linear and quadratic functions. In both settings, we prove that the transformer converges to a solution where ∑i=1T−1ρi(T) _i=1^T-1 _i^(T) decays with T, with the spectral norm of the Jacobians evaluated at each step input strictly less than one. Consequently, the trained transformer achieves a test loss of O(1/poly(d))O(1/poly(d)) with input dimension d and T=Θ⁡(log⁡d)T= ( d) steps (cf. Theorem 3.4). Extending prior work on linear functions [65], our analysis for quadratic functions introduces a novel three-head transformer construction that exactly implements gradient descent (cf. Lemma 3.3), a finding that may be of independent interest. Motivated by this analysis, we propose a training method to control accumulated error by regularizing the Jacobian-based error amplification factor. Our method incorporates chain-of-thought length compression that downsamples the reasoning steps during training, and quantization-aware training [69, 79] that trains weights with gradients computed on the quantized model. We show that our method reduces the Jacobian spectral norms by regularizing the Hessian trace of the model loss during quantized training. We evaluate this method on five algorithmic reasoning tasks [70] and two state-tracking tasks from LEGO [49]. First, across all seven tasks, our method yields an average test accuracy improvement of 3.5%\% over baselines, including implicit chain-of-thought [51] and Coconut [63], shown in Figure 1. Second, on evaluation of length generalization with larger inputs, our method enhances the test performance by 8.2%\% over baselines. Finally, ablation studies confirm that our method reduces the Jacobian-based error amplification factor by up to 8× compared to existing methods, as shown in Figure 1. The code for replicating these empirical findings is available at https://github.com/VirtuosoResearch/Multi-step-reasoning-experiments. Summary of contributions. First, our analysis reveals that the inference error in multi-step reasoning scales with the product of spectral norms of Jacobians taken through the input space. Second, we theoretically show that transformers trained on linear and quadratic functions converge to a regime where these norms are strictly less than one, yielding near-zero inference error. Finally, we propose a training method that regularizes these input Jacobian norms, empirically enhancing test performance and length generalization across diverse reasoning tasks. 2 Preliminaries We now describe the problem setup for multi-step reasoning. Let X denote the feature space. Let fW:∗→f_W:X denote a next-token predictor with parameters W∈ℝpW ^p, such as a language model. fWf_W takes a sequence of vectors in X of arbitrary length as input and outputs a single vector in X, which is then appended to the sequence before the next prediction during inference. For an input sequence ∈∗x , we assume that there exists an underlying rule generating a ground-truth sequence [,y(1),…,y(T⁡())][x,y^(1),…,y^(T(x))], where y(t)∈y^(t) is the ground-truth output at step t, and T⁡()T(x) denotes the total number of steps, which depends on the input x. We denote the maximum number of steps as T for a particular task and bounded input lengths. As a concrete example, suppose x contains a graph and a source node; then each y(i)y^(i) encodes an intermediate execution step of a graph traversal algorithm such as breadth-first search [70]. At inference time, the model fWf_W generates T⁡()T(x) steps for an input x autoregressively. We denote the prediction at step t as y^(t)=fW([,y^(1),…,y^(t−1)]), with y^(1)=fW(). y^(t)=f_W ([x, y^(1),…, y^(t-1)] ), with ~ y^(1)=f_W(x). For simplicity of notation, we write the generated sequence after t steps as ^(t)=[,y^(1),…,y^(t)], with ^(0)=[]. z^(t)=[x, y^(1),…, y^(t)], with ~ z^(0)=[x]. The corresponding ground-truth sequence with the correct intermediate steps is denoted as (t)=[,y(1),…,y(t)], for t=1,…,T(), and (0)=.z^(t)=[x,y^(1),…,y^(t)], for t=1,…,T(x), and z^(0)=x. During training, the model fWf_W predicts the next step conditioned on the preceding steps that encode the ground truth. Let S=(i,yi(1),…,yi(T⁡(i)))i=1nS= \(x_i,y_i^(1),…,y_i^(T(x_i))) \_i=1^n denote a training dataset of n examples, where every ix_i is drawn independently from an unknown distribution D. Given a loss function ℓ⁡(⋅,⋅) (·,·), the training loss is defined as: L^(fW,S)=1n∑i=1n∑t=0T⁡(i)−1ℓ⁡(fW(i(t)),yi(t+1)). L(f_W;S)= 1n _i=1^n _t=0^T(x_i)-1 (f_W (z_i^(t) ),y_i^(t+1) ). (1) At test time, the model predicts the next step conditioned on previous predictions. The outcome is evaluated by the loss function at the final step. We thus define this test loss over a random sample x drawn from the data distribution D as: L⁡(fW,)=∼[ℓ⁡(fW(^(T⁡()−1)),y(T⁡()))]. L(f_W;D)= E_x [ (f_W( z^(T(x)-1)),y^(T(x)) ) ]. (2) To be clear, we will call the above equation (2) the inference loss. At training time, minimizing equation (1) measures the loss relative to the ground truth. At inference time, the model assimilates prediction errors that occurred during previous steps back into the model. Thus, the prediction errors during inference can accumulate during long generation steps. Problem statement. We quantify how inference loss accumulates throughout generation. We study conditions under which multi-step reasoning is stable, where the inference loss L⁡(fW,)L(f_W;D) remains bounded in long chains. Additionally, we design training algorithms to mitigate error accumulation and improve the stability of multi-step reasoning. Notations. For any matrix X, let ‖X‖\|X\| denote the largest singular value (or spectral norm) of X. Let ‖X‖F\|X\|_F denote the Frobenius norm of X. f⁡(N)≲g⁡(N)f(N) g(N) indicates that there exists a fixed constant c that does not grow with N such that f⁡(N)≤c⋅g⁡(N)f(N)≤ c· g(N) for large enough integers of N. 3 Theoretical Analysis This section presents a theoretical analysis of multi-step reasoning in the above setup. First, we present a bound on the inference loss relative to the ground truth generation. We show that this bound depends on an error amplification factor defined by the product of spectral norms of Jacobian matrices measured at the input space. Second, we instantiate this bound in in-context prediction tasks for linear and quadratic functions. We show that transformer models converge to a regime where the error amplification factor decays across steps, thus yielding nearly zero inference loss. 3.1 Error Amplification in Multi-Step Reasoning We now quantify how the inference loss accumulates across generation steps. We write the number of steps as T when the input is clear from the context. For an input x, we compare the inference loss ℓ⁡(fW(^(T−1)),y(T)) (f_W( z^(T-1)),y^(T) ) of equation (2) with the loss ℓ⁡(fW((T−1)),y(T)) (f_W(z^(T-1)),y^(T) ) of the prediction conditioned on the correct steps, which is what training minimizes in equation (1). To measure the difference, let e(t)()=fW((t−1))−y(t)e^(t)(x)=f_W(z^(t-1))-y^(t) and e^(t)()=fW(^(t−1))−y(t) e^(t)(x)=f_W( z^(t-1))-y^(t) denote the prediction error at step t conditioned on the correct steps and on the model’s generations, respectively. Consider the first-order Taylor expansion in the input space, for every t, fW(^(t−1))−fW((t−1))=∑i=1t−1Ji(t)()e^(i)()+ℛ(t), where Ji(t)()=∂fW((t−1))∂y(i)f_W( z^(t-1))-f_W(z^(t-1))\;=\; _i=1^t-1J^(t)_i(x)\, e^(i)(x)\;+\;R^(t), where J^(t)_i(x)= ∂ f_W(z^(t-1))∂ y^(i) is the input Jacobian evaluated at step t, and ℛ(t)R^(t) is the second-order remainder error term. By recursively unrolling this expansion, an error at step i affects a later step t by first propagating to an intermediate step k and then from k to t. To capture the propagation path, we define a term: ρi(t)=sup∈∗‖Ji(t)()‖+∑k=i+1t−1(sup∈∗‖Jk(t)()‖⋅ρi(k)), with ρt(t)=1,ρ^(t)_i= _x \|J^(t)_i(x)\|+ _k=i+1^t-1 ( _x \|J^(t)_k(x)\|·ρ^(k)_i ), with ρ^(t)_t=1, which sums, over all paths from i to t, the product of the spectral norms of the input Jacobians along the path. Therefore, an error at step t reaches the final step amplified by a factor ρt(T)ρ^(T)_t. Our main result is to incorporate the spectral norm of Jacobians to quantify the accumulation of inference errors across generation steps. For every x in the support of D, we denote the maximum of the two errors, for t from 11 to T−1T-1, by: sup∈∗‖e(t)()‖=εt,sup∈∗‖e^(t)()‖=rt. _x \|e^(t)(x)\|= _t,~ _x \| e^(t)(x)\|=r_t. (3) We present a bound on the (expected) inference loss as follows. Lemma 3.1. Suppose that fWf_W is twice differentiable with respect to its input, the loss function ℓ⁡(⋅,⋅) (·,·) is C1C_1-Lipschitz continuous in its first argument, and the Jacobian of fWf_W with respect to its input is C2C_2-Lipschitz continuous. Let L⋆(fW,)=∼[ℓ⁡(fW((T−1),y(T)))]L (f_W;D)= E_x [ (f_W(z^(T-1),y^(T)) ) ] be the expected loss at the final step of the model prediction conditioned on correct steps. Then, for any fixed W∈ℝpW ^p that satisfies these conditions, and under equation (3), the inference loss is bounded away from the expected loss as follows: |L⁡(fW,)−L⋆(fW,)|≤C1∑t=1T−1ρt(T)εt+C1C22∑t=2Tρt(T)(∑i=1t−1ri)2. |L(f_W;D)-L (f_W;D) |≤C_1 _t=1^T-1ρ^(T)_t _t+ C_1C_22 _t=2^Tρ^(T)_t ( _i=1^t-1r_i )^2. (4) The proof involves a decomposition of the error between predictions on the model-generated sequence and the correct step sequence by applying the first-order expansion recursively. The full proof is in Appendix A.1. Lemma 3.1 highlights that the inference loss is governed by ∑i=1T−1ρi(T) _i=1^T-1ρ^(T)_i, which we term an error amplification factor. Even if the training error is small, this factor can grow exponentially with T, resulting in large inference errors on long reasoning sequences. 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. Examples. To empirically verify our analysis, we study the scaling of inference losses on graph algorithmic reasoning tasks [70]. We fine-tune Qwen-1.5B, Llama-1B, and Gemma-2B models on tasks with graphs of 1010 nodes. Inputs are text-based graph descriptions, and intermediate steps are encoded as node labels, such as the nodes along the shortest path. As a reasoning step in text-based datasets typically involves several tokens, we compute the loss, the ℓ2 _2 error, and the Jacobian matrices token-by-token within a step, and then aggregate them over the tokens within one step. For text-based reasoning tasks, we view X as the continuous space of input embeddings. The output of fWf_W at step t is then the expected embedding of the predicted next token, calculated as the weighted average of all token embeddings based on their predicted softmax probabilities. We compute the Jacobian of the output vector with respect to the input embeddings. We estimate its spectral norm via power iteration using Jacobian-vector products. To approximate sup‖Jk(t)‖ _x\|J_k^(t)\| in equation (4), we take the maximum norm across the dataset. After minimizing the training loss, we measure the loss gap and the empirical values of our bound in equation (4), for T ranging from 1 to 6. In Figure 2, the inference loss increases significantly with the number of steps. Our bound, which involves ∑i=1T−1ρi(T) _i=1^T-1ρ^(T)_i, scales exponentially with T, closely tracking the growth in inference loss. We also observe similar results on Llama-1B and Gemma-2B models, which are described in Appendix B. These results demonstrate that controlling input Jacobian norms during training is crucial for reducing accumulated errors in multi-step generation. Remark 3.2. Lemma 3.1 provides a worst-case upper bound that may not be numerically tight. Empirically, while we observe a consistent correlation between the bound values and inference losses across models, the numerical scales of the error amplification factor ∑i=1T−1ρi(T) _i=1^T-1ρ^(T)_i vary significantly. The bound-to-gap ratio is around 88 for Qwen-1.5B and around 10510^5–10610^6 for Llama-1B and Gemma-2B. We report the empirical loss gap, the bound, the bound-to-gap ratio, both terms of equation (4), and the estimated constants C1C_1 and C2C_2, for every model in Appendix B. 3.2 Learning Dynamics in Quadratic Models Next, we explicitly instantiate the bound in Lemma 3.1 in training transformers on an in-context weight prediction task. We show that if the spectral norm of the input Jacobian matrix is less than one, the inference loss can be reduced close to zero. Based on prior work on in-context learning [78], we study a task that predicts the weight vector θ⋆∼⁡(0,Idd)θ (0, _d) of a function, given n examples (xi,yi)(x_i,y_i) where xi∼⁡(0,Idd)x_i (0, _d) and yi=gθ⋆(xi)y_i=g_θ (x_i) for a known function g. We define the ground-truth steps as iterates by performing gradient descent on the regression examples θi+1=θi−ηn∑k=1n∇θgθi(xk,yk), _i+1= _i- ηn _k=1^n _θg_ _i(x_k,y_k), with a learning rate η and an initialization θ0 _0. We consider a one-layer transformer with h non-linear self-attention heads: fW(Z)=Z[:,−1]+∑i=1hWPViZ⋅ϕi(Z⊤WKQiZ[:,−1]n), and Z=[X0d⋯0dy0⋯00d×nθ0⋯θi01⋯1], f_W(Z)=Z_[:,-1]+ _i=1^hW^PV_iZ· _i ( Z W^KQ_iZ_[:,-1]n ), and Z= bmatrixX&0_d&·s&0_d\\ y&0&·s&0\\ 0_d× n& _0&·s& _i\\ 0&1&·s&1 bmatrix, (5) where Z∈(2d+2)×(n+i+1)Z∈ ^(2d+2)×(n+i+1) is the input embedding, X=[x1,…,xn]X=[x_1,…,x_n], and y=[y1,…,yn]y=[y_1,…,y_n]. The ones in the bottom are appended for formatting convenience. For step T, the last column contains the target vector θ⋆θ . Following prior work [65, 67, 75], WPViW^PV_i and WKQiW^KQ_i are the merged projection-value and key-query matrices, ϕi _i is a non-linear function, and Z[:,−1]Z_[:,-1] denotes the last column. We use the squared loss for this setting. Numerical example. We validate that the transformer can achieve near-zero inference loss in this setting, for both linear and quadratic function classes. The results are shown in Figure 3. The inference loss converges to near-zero, with longer steps yielding lower inference loss. (a) Linear functions (b) Quadratic functions (c) Linear functions (d) Quadratic functions 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=10d=10 and train a one-head linear transformer on linear functions (n=20n=20, η=0.4η=0.4) and a three-head non-linear transformer on quadratic functions (n=200n=200, η=10−3η=10^-3). For the random Gaussian noise, we set σ=0.002σ=0.002. On the right, we consider a slightly different setting by first adding noise to the weights before computing the gradients. We describe our main results. For linear functions y=θ⊤Xy=θ X, we prove that the model exhibits inference loss of O⁡(1/poly(d))O(1/poly(d)) with T=Θ⁡(log⁡d)T= ( d). This is because the transformer converges to a solution where Ji−1(i)≈I−ηnXX⊤J_i-1^(i)≈ I- ηnX and ρi(T)≈‖I−ηnXX⊤‖T−i _i^(T)≈\|I- ηnX \|^T-i. The accumulated errors form a converging geometric series. We present the complete proof in Appendix A.2. Next, we present results on quadratic functions y=(θ⊤X)⊙2y=(θ X) 2, which represent a fundamental yet tractable setting for studying gradient descent dynamics [45, 80]. We first establish that a three-head transformer can exactly implement the gradient update step using cubic and squared activations. Lemma 3.3. There exist weight matrices WPVh,WKQhh=13\W^PV_h,W^KQ_h\_h=1^3 and elementwise non-linear functions ϕhh=13\ _h\_h=1^3, such that a three-head transformer fW(Z^T−1)f_W( Z_T-1) outputs the T-step gradient descent solution with a learning rate η on the quadratic functions. We then show that the transformer model converges to this construction during training, yielding near-zero inference loss. Theorem 3.4. Let κ∈(0,1/24]κ∈(0,1/24], μ∈(0,4]μ∈(0,4], and α0∈[3/4,1) _0∈[3/4,1) be the constants. Suppose n=Θ⁡(d2log5d)n= (d^2 ^5d), η=κ/dη=κ/d, θ0=α0θ⋆ _0= _0θ , and T=⌈Clog⁡d⌉T= C d for some C>0C>0 with aC>1aC>1 for a=log⁡(11−μηd)a= ( 11-μη d ). There exists an initialization such that if we run gradient flow on the training loss, then after time t=O⁡(log⁡d+log⁡1ϵ0)t=O ( d+ 1 _0 ), for any ϵ0≥Θ⁡(log⁡d2aC−1) _0≥ ( dd^2aC-1 ), the trained transformer satisfies L⁡(fW,)≤ϵ1L(f_W;D)≤ _1 for any ϵ1∈(Θ⁡(d−2aC+1log⁡d), 1) _1∈ ( (d^-2aC+1 d),\,1 ). The proof tracks an ordinary differential equation for the eigenvalues of the weight matrices. In the converged solution, we have Ji−1(i)≈I−2ηnXD(θi)X⊤ and ρi(T)≈‖I−2ηnXD(θi)X⊤‖T−i,J_i-1^(i)≈ I- 2ηnXD( _i)X and _i^(T)≈\|I- 2ηnXD( _i)X \|^T-i, where D⁡(θi)=diag⁡(3(X⊤θi)⊙2−(X⊤θ⋆)⊙2)D( _i)=diag(3(X _i) 2-(X θ ) 2). Then, the accumulated errors again form a converging geometric series. We present the complete proof in Appendix A.3. Comparison with prior work. Prior analyses of transformers trained on multi-step reasoning differ from ours in their attention mechanisms, tasks, and tracked quantities. Previous works have examined one-head softmax attention on state-tracking tasks to demonstrate implicit biases and length generalization [66, 67]. [65] study one-head linear attention on in-context linear regression, where the trained transformer implements multi-step gradient descent. [75] study multi-head softmax attention on tree path-finding, where different heads specialize and coordinate across subtasks. In contrast, we study multi-head non-linear attention on quadratic regression. We identify the explicit form of Jacobians taken through the inputs at each step at convergence and connect it with the stability of the inference loss. 4 Proposed Methods Our analysis suggests two strategies to mitigate accumulated errors: reducing the sequence length to limit the error amplification factor ∑t=1T−1ρt(T) _t=1^T-1ρ^(T)_t and regularizing the input Jacobian spectral norms ‖Ji(t)‖\|J^(t)_i\|. Motivated by these insights, we design a training method combining two techniques. (1) Chain-of-thought length compression. First, we compress the reasoning sequence by subsampling intermediate steps. For a training example (,y(1),…,y(T))(x,y^(1),…,y^(T)) and a sampling factor λ∈(0,1]λ∈(0,1], we uniformly sample T′=⌊λ⋅T⌋T = λ· T intermediate steps to form a compressed sequence (,y(1),…,y(T′))(x,y^(1),…,y^(T )), and minimize the training loss over the sampled sequence. As shown in Figures 4 and 4, models trained on subsampled sequences consistently outperform those trained on full steps. Intermediate numbers of steps achieve the lowest inference loss on the Dijkstra and Prim tasks. Notably, subsampling steps reduces the error amplification factor by 92× and 181× compared to training with all steps. Subsampling steps reduces the amplification factor but increases the per-step errors εt _t. Thus, the best λ is the one that balances these two effects. (2) Regularization through quantization-aware training. Second, we incorporate quantization during supervised fine-tuning to explicitly regularize the Jacobian spectral norms. This is inspired by observations that quantized training smooths the loss landscape and provides a regularization effect [69, 79]. During the forward pass at training step t, we apply a standard uniform quantization function to the full-precision weights W(t)W^(t): Q⁡(W(t))=Δ⋅Round(Clamp(W(t)−b,Wmin,Wmax)/Δ)+b, Q(W^(t))= ·Round (Clamp(W^(t)-b,W_min,W_max)/ )+b, where Δ and b are bit-width-dependent scaling and bias parameters. We use the straight-through estimator [43] to approximate the gradient ∇WL^(fQ⁡(W))≈∇Q⁡(W)L^(fQ⁡(WCLOSE) _W L(f_Q(W))≈ _Q(W) L(f_Q(W). We hypothesize that quantization-aware training acts as a regularizer on the error amplification factor. Under the Pseudo-Quantization Noise model [41, 46], the quantized model loss can be viewed as the expected loss under mean-zero randomized-rounding noise. Minimizing this loss penalizes the trace of the loss Hessian, which can be decomposed into the Jacobians taken through the input space. Figures 4 and 4 demonstrate this regularization effect. Training with lower bit-widths yields a lower error amplification factor compared to full precision. When combining 1-bit quantization, ∑t=1T−1ρt(T) _t=1^T-1 _t^(T) is further reduced by 3.6×3.6× and 5.8×5.8×, compared to full precision, on the Dijkstra and Prim algorithmic tasks, respectively. We also observe that, while quantization increases the prediction error at each step, it reduces the error amplification factor much more significantly, thus leading to improved performance. We report evaluations under various bit-widths in Appendix B. Taken together, we summarize our approach in Algorithm 1. (a) Dijkstra (b) Prim’s algorithm (c) Dijkstra (d) Prim’s algorithm 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. Remark 4.1 (Length generalization). For out-of-distribution settings, the test error bound from Lemma 3.1 can increase in two aspects. First, the per-step prediction errors grow due to covariate shift. Second, the error amplification factor accumulates over the additional steps. If the step-to-step Jacobian norm exceeds one, the extra steps multiply the error geometrically, making our regularization approach crucial. Designing training algorithms for length generalization is an interesting direction for future work, and our analysis lays a concrete direction for this. Algorithm 1 Jacobian-Regularized Training through Length Compression and Quantization Input: n training examples with correct sequences (i,yi(1),…,yi(T))(x_i,y_i^(1),…,y_i^(T)) of length T Require: Pretrained language model fW(0)f_W^(0), number of epochs S, learning rate η, a sampling factor λ, quantization function Q⁡(⋅)Q(·) Output: Fine-tuned quantized model fQ⁡(W(S))f_Q(W^(S)) 1: for t=0,1,…,S−1t=0,1,…,S-1 do 2: (yi(1),…,yi(T′))←(y_i^(1),…,y_i^(T ))← Subsample T′=⌊λ⋅T⌋T = λ· T steps from example i, for i=1,…,ni=1,…,n 3: L^(fQ⁡(W(t)))← L(f_Q(W^(t)))← Compute the training loss of fQ⁡(W(t))f_Q(W^(t)) on the sampled sequences 4: W(t+1)←W(t)−η⋅∇Q⁡(W(t))L^(fQ⁡(W(t)))W^(t+1)← W^(t)-η· _Q(W^(t)) L(f_Q(W^(t))) 5: end for 6: return fQ⁡(W(S))f_Q(W^(S)) 5 Experiments We evaluate our proposed method in a range of algorithmic reasoning tasks. First, our method outperforms baselines by an average of 3.5% across five graph algorithmic and two symbolic reasoning tasks. Second, our method significantly enhances out-of-distribution length generalization, yielding an average 8.2% improvement. Third, empirical analysis confirms that our approach reduces the error amplification factor (measured via Jacobian spectral norms) by up to 8×. Lastly, we also conduct ablation studies for various design choices of our algorithm. 5.1 Experimental Setup Datasets and models. First, we evaluate our algorithm on graph algorithmic tasks from the CLRS-Text benchmark [70], which encodes graphs as flattened adjacency lists and represents the intermediate steps as lists of node labels. We use five classic graph algorithms, including Dijkstra, Bellman-Ford, Prim’s algorithm, breadth-first search, and depth-first search. For each task, we use 2,000 training examples generated on random graphs of 10 nodes and 200 examples for validation and testing. Second, we evaluate on symbolic state-tracking reasoning tasks, namely LEGO [49]. We use two task variants: transitive (cyclic) and symmetry group actions. For each task, we use a dataset of 2,000 training examples with k=8k=8 variables, and 200 examples each for validation and testing. We evaluate our method in fine-tuning the QWen-1.5B models. Baselines. We compare our method with three categories of baselines. First, we consider standard fine-tuning baselines that train without or with the provided chain-of-thought sequences, namely No CoT and SFT-CoT. Second, we consider existing chain-of-thought methods. This includes implicit CoT [51] that trains a model to directly predict the final output by gradually removing intermediate steps during training, and the chain of continuous thought (Coconut) [63] that trains a model to reason in a continuous latent space by using the hidden state as the next input embedding. Third, we consider noise stability optimization (NSO) [56] that regularizes the Hessian trace through random Gaussian noise injection. For each baseline, we evaluate the test accuracy on the final answer. Implementations. Our algorithm includes key hyperparameters: the sampling factor λ and the quantization bit-width. We vary λ across 0.1, 0.2, 0.5, and 1. The bit-width is varied across 1, 2, 3, and 4 bits. We use the quantization methods from existing methods [79]. For the base training method, we use LoRA fine-tuning with a rank between 16 and 64. The training uses the AdamW optimizer with a learning rate in the range of 1×10−51× 10^-5 to 4×10−54× 10^-5. We report the hyperparameters used for each result in Appendix B. For the length generalization experiment, we additionally evaluate at 20%20\% longer inputs on all five CLRS-Text tasks. Absolute accuracy degrades sharply for every method, falling below 25%25\%. While our method still leads, at this accuracy level, the input-length bottleneck dominates. We therefore report the 10%10\% setting as our main result. Further improving the length generalization on longer sequences is left for future work. 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. 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 Length generalization Dijkstra Bellman-Ford Prim BFS DFS Cyclic Symmetry No CoT 29.8 ± 0.4 42.9 ± 0.6 37.0 ± 0.5 48.5 ± 0.2 23.0 ± 0.5 55.4 ± 0.5 31.2 ± 0.7 SFT-CoT 42.1 ± 0.4 47.2 ± 0.8 34.1 ± 0.6 56.2 ± 0.6 21.5 ± 0.2 80.0 ± 0.4 48.8 ± 0.3 Implicit CoT 33.7 ± 0.3 46.5 ± 0.5 39.3 ± 0.4 54.6 ± 0.5 30.0 ± 0.3 82.5 ± 0.8 53.6 ± 1.0 Coconut 38.2 ± 0.6 38.9 ± 0.4 37.6 ± 0.9 53.9 ± 0.9 32.6 ± 0.4 80.4 ± 0.5 49.4 ± 0.2 NSO 36.4 ± 0.5 44.1 ± 0.8 38.4 ± 0.3 54.8 ± 0.2 28.2 ± 0.2 82.4 ± 0.4 51.2 ± 0.9 Algorithm 1 44.3 ± 0.8 50.1 ± 0.1 42.5 ± 0.6 59.6 ± 0.6 35.4 ± 0.8 87.2 ± 0.2 59.1 ± 0.7 5.2 Experimental Results Improved multi-step reasoning. Table 1 details the test accuracy of our fine-tuned models across seven reasoning datasets. Our method consistently outperforms all baselines, achieving an average improvement of 3.5% across five graph algorithmic and two symbolic state-tracking tasks. Specifically, by combining sequence compression with quantization-aware training, our algorithm surpasses latent-reasoning baselines (implicit reasoning and Coconut) by up to 6.5%. Enhanced length generalization. Next, we evaluate the trained models on examples with 10% longer input length than training examples in the lower half of Table 1. Our approach consistently exceeds baseline performance by an average of 8.2%, demonstrating its effectiveness in suppressing error amplification during extended generation. Reduction of the error amplification factor. Finally, to explain these performance gains, we compare the empirical error amplification factors across methods. Our algorithm drastically reduces this factor by 3–8× compared to baselines, validating our method’s regularization on Jacobian spectral norms. Quantitatively, on Qwen-1.5B across the three tasks in Figure 1, our method reduces ∑i=1T−1ρi(T) _i=1^T-1ρ^(T)_i by 43%43\% relative to baselines, with a 72%72\% average reduction in the inference loss. We further fine-tune Llama-1B and Gemma-2B on Bellman-Ford with SFT-CoT, Coconut, and Algorithm 1, and measure the full decomposition of equation (4) for each. We also find that the ordering of methods by bound value matches their ordering by empirical loss gap on both models. Relative to Coconut, our method reduces the bound by 77%77\% and the empirical test loss gap by 58%58\% on average. The absolute scale of the bound differs sharply between architectures. We report the full results in Table 4 of Appendix B. Ablation studies. We evaluate the individual components of our method through ablation studies. First, removing either length compression or quantization-aware training degrades performance by 4.3% and 3.2%, confirming both are essential. Second, we compare intermediate step sampling strategies for length compression. We consider random sampling, removing the first, last, or equally spaced steps. Among them, random sampling yields slightly better performance and is thus adopted. Finally, we find that injecting quantization noise during training noticeably outperforms standard random Gaussian noise by 9%. 6 Related Work Theoretical understanding of multi-step reasoning. Studying the training dynamics of transformers is key to understanding how multi-step reasoning improves generalization. Prior work analyzes single-head transformers trained on parity-check tasks [73, 66], and state tracking tasks [67], while [75] studies multi-head specialization on path-finding in binary trees. For the in-context weight prediction task, [65] establishes convergence for linear targets using single-head linear transformers. We extend the analysis to nonlinear multi-head transformers and quadratic functions. Complementary theoretical analyses include entropy-calibration analyses of error accumulation in autoregressive generation [59] and information-theoretic frameworks that quantify per-step information gain in multi-step reasoning [72]. [68] study learning general function classes of next-token generators with CoTs in the realizable setting, showing that the sample complexity scales with the VC dimension of the hypothesis class and log⁡(T) (T) of the CoT length T. [58] further improve the bound with a notion of CoT-information, which achieves faster rates than applying uniform convergence to training loss on chain-of-thought sequences. Prior work has also examined the expressiveness of transformers with chain-of-thought. [52] show that CoT enables transformers to perform computations beyond the expressiveness limits of models without CoT. Expressiveness has been further characterized through circuit complexity under logarithmic, linear, and polynomial generation lengths [54]. Increasing CoT length has been shown to outperform parallel scaling of multiple chains on graph connectivity tasks [71]. Generalization of language models. Recent work derives generalization bounds for LLMs via compression and subspace training [53, 48], with tighter martingale-based extensions that explain empirical pretraining scaling laws [61]. A separate line of work derives PAC-Bayes bounds for fine-tuning that depend on the loss Hessian [47, 50]. Optimization methods for multi-step reasoning. Several recent works aim to improve multi-step reasoning beyond supervised fine-tuning. [76] frame this process as a latent-variable problem, minimizing gradient variance to efficiently allocate sampling budgets. A complementary line of work explores implicit reasoning. Implicit chain-of-thought [51] gradually removes intermediate steps to internalize computation. Continuous chain-of-thought [63] shifts reasoning into a continuous latent space. Semantic chain-of-thought [64] aligns implicit and explicit steps via contrastive learning. Empirically, [77] investigate how long reasoning chains emerge in large language models, highlighting training computation and reward shaping as key factors. More broadly, multi-step reasoning traces also appear in experimental design settings [55, 74], where trajectory-level data is revealed over time. 7 Conclusion This work provides an analysis of multi-step reasoning by bounding its inference errors via the product of input Jacobian spectral norms across generation steps. We prove that transformers trained for in-context weight prediction converge to a multi-head construction where this error amplification factor decays, achieving polynomially small test loss. We introduce a training method that explicitly controls this factor using reasoning length compression and quantization-aware training. Across seven algorithmic tasks, our approach consistently improves test accuracy and length generalization, with reduced empirical amplification factors in fine-tuned models. Acknowledgement This work received financial support from NSF award IIS-2412008 and a startup fund from Northeastern University. D. Li is supported in part by a JPMorgan Chase PhD Fellowship. Z. Zhang and M. Duan received additional financial support from Khoury PhD fellowships. References [1] Awni Altabaa, Omar Montasser and John Lafferty “CoT Information: Improved Sample Complexity under Chain-of-Thought Supervision” In NeurIPS, 2025 [2] Yoshua Bengio, Nicholas Léonard and Aaron Courville “Estimating or propagating gradients through stochastic neurons for conditional computation” In arXiv preprint arXiv:1308.3432, 2013 [3] Emmanuel Candes, Xiaodong Li and Mahdi Soltanolkotabi “Phase retrieval via Wirtinger flow: Theory and algorithms” In IEEE Transactions on Information Theory IEEE, 2015 [4] Steven Cao, Gregory Valiant and Percy Liang “On the entropy calibration of language models” In NeurIPS, 2025 [5] Yuntian Deng, Yejin Choi and Stuart Shieber “From explicit cot to implicit cot: Learning to internalize cot step by step” In arXiv preprint arXiv:2405.14838, 2024 [6] Ryan Ehrlich, Bradley Brown, Jordan Juravsky, Ronald Clark, Christopher Ré and Azalia Mirhoseini “Codemonkeys: Scaling test-time compute for software engineering” In arXiv preprint arXiv:2501.14723, 2025 [7] Marc Finzi, Sanyam Kapoor, Diego Granziol, Anming Gu, Christopher De, J Kolter and Andrew Wilson “Compute-Optimal LLMs Provably Generalize Better With Scale” In ICLR, 2025 [8] Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian Bartoldson, Bhavya Kailkhura, Abhinav Bhatele and Tom Goldstein “Scaling up test-time compute with latent reasoning: A recurrent depth approach” In NeurIPS, 2025 [9] Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael Mahoney and Kurt Keutzer “A survey of quantization methods for efficient neural network inference” In Low-power computer vision, 2022, p. 291–326 [10] Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston and Yuandong Tian “Training large language models to reason in a continuous latent space” In COLM, 2025 [11] Yinhan He, Wendy Zheng, Yaochen Zhu, Zaiyi Zheng, Lin Su, Sriram Vasudevan, Qi Guo, Liangjie Hong and Jundong Li “SemCoT: Accelerating Chain-of-Thought Reasoning through Semantically-Aligned Implicit Tokens” In NeurIPS, 2025 [12] Jianhao Huang, Zixuan Wang and Jason Lee “Transformers learn to implement multi-step gradient descent with chain of thought” In ICLR, 2025 [13] Ruiquan Huang, Yingbin Liang and Jing Yang “How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias” In ICML, 2025 [14] Yu Huang, Zixin Wen, Aarti Singh, Yuejie Chi and Yuxin Chen “Transformers provably learn chain-of-thought reasoning with length generalization” In NeurIPS, 2025 [15] Nirmit Joshi, Gal Vardi, Adam Block, Surbhi Goel, Zhiyuan Li, Theodor Misiakiewicz and Nathan Srebro “A theory of learning with autoregressive chain of thought” In Annual Conference on Learning Theory (COLT), 2025 [16] Haotian Ju, Dongyue Li, Aneesh Sharma and Hongyang Zhang “Generalization in graph neural networks: Improved pac-bayesian bounds on graph diffusion” In International conference on artificial intelligence and statistics, 2023, p. 6314–6341 PMLR [17] Haotian Ju, Dongyue Li and Hongyang Zhang “Robust fine-tuning of deep neural networks with Hessian-based generalization guarantees” In International conference on machine learning (ICML), 2022 [18] Mujin Kwun, Depen Morwani, Chloe Su, Stephanie Gil, Nikhil Anand and Sham Kakade “LOTION: Smoothing the Optimization Landscape for Quantized Training” In arXiv preprint arXiv:2510.08757, 2025 [19] Dongyue Li, Zechun Liu, Kai Yi, Changsheng Zhao, Zhenshuo Zhang, Raghuraman Krishnamoorthi, Harshit Khaitan, Hongyang. Zhang and Steven Li “WinQ: Accelerating Quantization-Aware Training of Large Language Models around Saddle Points” In ICML, 2026 [20] Dongyue Li, Zhenshuo Zhang, Minxuan Duan, Edgar Dobriban and Hongyang. Zhang “Efficiently Learning Branching Networks for Multitask Algorithmic Reasoning” In arXiv preprint arXiv:2512.01113, 2025 [21] Yuanzhi Li, Tengyu Ma and Hongyang Zhang “Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations” In Conference on learning theory, 2018, p. 2–47 PMLR [22] Zhiyuan Li, Hong Liu, Denny Zhou and Tengyu Ma “Chain of thought empowers transformers to solve inherently serial problems” In ICLR, 2024 [23] Sanae Lotfi, Marc Finzi, Sanyam Kapoor, Andres Potapczynski, Micah Goldblum and Andrew Wilson “PAC-Bayes compression bounds so tight that they can explain generalization” In Advances in Neural Information Processing Systems (NeurIPS), 2022 [24] Sanae Lotfi, Marc Finzi, Yilun Kuang, Tim Rudner, Micah Goldblum and Andrew Wilson “Non-vacuous generalization bounds for large language models” In Proceedings of the 41st International Conference on Machine Learning, 2024, p. 32801–32818 [25] William Merrill and Ashish Sabharwal “The expressive power of transformers with chain of thought” In ICLR, 2024 [26] Alexandru Meterez, Pranav Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade and Alex Damian “A defense of the quadratic model” In arXiv preprint arXiv:2607.21716, 2026 [27] Parsa Mirtaheri, Ezra Edelman, Samy Jelassi, Eran Malach and Enric Boix-Adsera “Let Me Think! A Long Chain-of-Thought Can Be Worth Exponentially Many Short Ones” In NeurIPS, 2025 [28] Jean-Francois Ton, Muhammad Taufiq and Yang Liu “Understanding Chain-of-Thought in LLMs through Information Theory” In International Conference on Machine Learning, 2025, p. 59784–59811 PMLR [29] Roman Vershynin “Introduction to the non-asymptotic analysis of random matrices” In Compressed Sensing: Theory and Applications Cambridge University Press, 2012, p. 210–268 [30] Kaiyue Wen, Huaqing Zhang, Hongzhou Lin and Jingzhao Zhang “From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency” In ICLR, 2025 [31] Bernard Widrow, Istvan Kollar and Ming-Chang Liu “Statistical theory of quantization” In IEEE Transactions on instrumentation and measurement 45.2 IEEE, 1996, p. 353–361 [32] Ruoxuan Xiong “Automated experimental design with optimization from historical data simulations” In Available at SSRN 5126080, 2025 [33] Ruoxuan Xiong, Alex Chin and Sean Taylor “Data-driven switchback experiments: Theoretical tradeoffs and empirical bayes designs” In arXiv preprint arXiv:2406.06768, 2024 [34] Tong Yang, Yu Huang, Yingbin Liang and Yuejie Chi “Multi-head transformers provably learn symbolic multi-step reasoning via gradient descent” In NeurIPS, 2025 [35] Jiarui Yao, Yifan Hao, Hanning Zhang, Hanze Dong, Wei Xiong, Nan Jiang and Tong Zhang “Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL” In NeurIPS, 2025 [36] Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig and Xiang Yue “Demystifying long chain-of-thought reasoning in llms” In ICML, 2025 [37] Hongyang Zhang, Dongyue Li and Haotian Ju “Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approach” In Transactions on Machine Learning Research (TMLR) Journal of Machine Learning Research Inc., United States, 2024 [38] Muru Zhang, Ofir Press, William Merrill, Alisa Liu and Noah Smith “How language model hallucinations can snowball” In Proceedings of the 41st International Conference on Machine Learning, 2024, p. 59670–59684 [39] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar and Tal Wagner “Unveiling transformers with lego: a synthetic reasoning task” In arXiv preprint arXiv:2206.04301, 2022 [40] Ziniu Zhang, Zhenshuo Zhang, Dongyue Li, Lu Wang, Jennifer Dy and Hongyang Zhang “Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimation” In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, 2025, p. 16470–16488 References [41] Bernard Widrow, Istvan Kollar and Ming-Chang Liu “Statistical theory of quantization” In IEEE Transactions on instrumentation and measurement 45.2 IEEE, 1996, p. 353–361 [42] Roman Vershynin “Introduction to the non-asymptotic analysis of random matrices” In Compressed Sensing: Theory and Applications Cambridge University Press, 2012, p. 210–268 [43] Yoshua Bengio, Nicholas Léonard and Aaron Courville “Estimating or propagating gradients through stochastic neurons for conditional computation” In arXiv preprint arXiv:1308.3432, 2013 [44] Emmanuel Candes, Xiaodong Li and Mahdi Soltanolkotabi “Phase retrieval via Wirtinger flow: Theory and algorithms” In IEEE Transactions on Information Theory IEEE, 2015 [45] Yuanzhi Li, Tengyu Ma and Hongyang Zhang “Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations” In Conference on learning theory, 2018, p. 2–47 PMLR [46] Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael Mahoney and Kurt Keutzer “A survey of quantization methods for efficient neural network inference” In Low-power computer vision, 2022, p. 291–326 [47] Haotian Ju, Dongyue Li and Hongyang Zhang “Robust fine-tuning of deep neural networks with Hessian-based generalization guarantees” In International conference on machine learning (ICML), 2022 [48] Sanae Lotfi, Marc Finzi, Sanyam Kapoor, Andres Potapczynski, Micah Goldblum and Andrew Wilson “PAC-Bayes compression bounds so tight that they can explain generalization” In Advances in Neural Information Processing Systems (NeurIPS), 2022 [49] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar and Tal Wagner “Unveiling transformers with lego: a synthetic reasoning task” In arXiv preprint arXiv:2206.04301, 2022 [50] Haotian Ju, Dongyue Li, Aneesh Sharma and Hongyang Zhang “Generalization in graph neural networks: Improved pac-bayesian bounds on graph diffusion” In International conference on artificial intelligence and statistics, 2023, p. 6314–6341 PMLR [51] Yuntian Deng, Yejin Choi and Stuart Shieber “From explicit cot to implicit cot: Learning to internalize cot step by step” In arXiv preprint arXiv:2405.14838, 2024 [52] Zhiyuan Li, Hong Liu, Denny Zhou and Tengyu Ma “Chain of thought empowers transformers to solve inherently serial problems” In ICLR, 2024 [53] Sanae Lotfi, Marc Finzi, Yilun Kuang, Tim Rudner, Micah Goldblum and Andrew Wilson “Non-vacuous generalization bounds for large language models” In Proceedings of the 41st International Conference on Machine Learning, 2024, p. 32801–32818 [54] William Merrill and Ashish Sabharwal “The expressive power of transformers with chain of thought” In ICLR, 2024 [55] Ruoxuan Xiong, Alex Chin and Sean Taylor “Data-driven switchback experiments: Theoretical tradeoffs and empirical bayes designs” In arXiv preprint arXiv:2406.06768, 2024 [56] Hongyang Zhang, Dongyue Li and Haotian Ju “Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approach” In Transactions on Machine Learning Research (TMLR) Journal of Machine Learning Research Inc., United States, 2024 [57] Muru Zhang, Ofir Press, William Merrill, Alisa Liu and Noah Smith “How language model hallucinations can snowball” In Proceedings of the 41st International Conference on Machine Learning, 2024, p. 59670–59684 [58] Awni Altabaa, Omar Montasser and John Lafferty “CoT Information: Improved Sample Complexity under Chain-of-Thought Supervision” In NeurIPS, 2025 [59] Steven Cao, Gregory Valiant and Percy Liang “On the entropy calibration of language models” In NeurIPS, 2025 [60] Ryan Ehrlich, Bradley Brown, Jordan Juravsky, Ronald Clark, Christopher Ré and Azalia Mirhoseini “Codemonkeys: Scaling test-time compute for software engineering” In arXiv preprint arXiv:2501.14723, 2025 [61] Marc Finzi, Sanyam Kapoor, Diego Granziol, Anming Gu, Christopher De, J Kolter and Andrew Wilson “Compute-Optimal LLMs Provably Generalize Better With Scale” In ICLR, 2025 [62] Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian Bartoldson, Bhavya Kailkhura, Abhinav Bhatele and Tom Goldstein “Scaling up test-time compute with latent reasoning: A recurrent depth approach” In NeurIPS, 2025 [63] Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston and Yuandong Tian “Training large language models to reason in a continuous latent space” In COLM, 2025 [64] Yinhan He, Wendy Zheng, Yaochen Zhu, Zaiyi Zheng, Lin Su, Sriram Vasudevan, Qi Guo, Liangjie Hong and Jundong Li “SemCoT: Accelerating Chain-of-Thought Reasoning through Semantically-Aligned Implicit Tokens” In NeurIPS, 2025 [65] Jianhao Huang, Zixuan Wang and Jason Lee “Transformers learn to implement multi-step gradient descent with chain of thought” In ICLR, 2025 [66] Ruiquan Huang, Yingbin Liang and Jing Yang “How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias” In ICML, 2025 [67] Yu Huang, Zixin Wen, Aarti Singh, Yuejie Chi and Yuxin Chen “Transformers provably learn chain-of-thought reasoning with length generalization” In NeurIPS, 2025 [68] Nirmit Joshi, Gal Vardi, Adam Block, Surbhi Goel, Zhiyuan Li, Theodor Misiakiewicz and Nathan Srebro “A theory of learning with autoregressive chain of thought” In Annual Conference on Learning Theory (COLT), 2025 [69] Mujin Kwun, Depen Morwani, Chloe Su, Stephanie Gil, Nikhil Anand and Sham Kakade “LOTION: Smoothing the Optimization Landscape for Quantized Training” In arXiv preprint arXiv:2510.08757, 2025 [70] Dongyue Li, Zhenshuo Zhang, Minxuan Duan, Edgar Dobriban and Hongyang. Zhang “Efficiently Learning Branching Networks for Multitask Algorithmic Reasoning” In arXiv preprint arXiv:2512.01113, 2025 [71] Parsa Mirtaheri, Ezra Edelman, Samy Jelassi, Eran Malach and Enric Boix-Adsera “Let Me Think! A Long Chain-of-Thought Can Be Worth Exponentially Many Short Ones” In NeurIPS, 2025 [72] Jean-Francois Ton, Muhammad Taufiq and Yang Liu “Understanding Chain-of-Thought in LLMs through Information Theory” In International Conference on Machine Learning, 2025, p. 59784–59811 PMLR [73] Kaiyue Wen, Huaqing Zhang, Hongzhou Lin and Jingzhao Zhang “From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency” In ICLR, 2025 [74] Ruoxuan Xiong “Automated experimental design with optimization from historical data simulations” In Available at SSRN 5126080, 2025 [75] Tong Yang, Yu Huang, Yingbin Liang and Yuejie Chi “Multi-head transformers provably learn symbolic multi-step reasoning via gradient descent” In NeurIPS, 2025 [76] Jiarui Yao, Yifan Hao, Hanning Zhang, Hanze Dong, Wei Xiong, Nan Jiang and Tong Zhang “Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL” In NeurIPS, 2025 [77] Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig and Xiang Yue “Demystifying long chain-of-thought reasoning in llms” In ICML, 2025 [78] Ziniu Zhang, Zhenshuo Zhang, Dongyue Li, Lu Wang, Jennifer Dy and Hongyang Zhang “Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimation” In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, 2025, p. 16470–16488 [79] Dongyue Li, Zechun Liu, Kai Yi, Changsheng Zhao, Zhenshuo Zhang, Raghuraman Krishnamoorthi, Harshit Khaitan, Hongyang. Zhang and Steven Li “WinQ: Accelerating Quantization-Aware Training of Large Language Models around Saddle Points” In ICML, 2026 [80] Alexandru Meterez, Pranav Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade and Alex Damian “A defense of the quadratic model” In arXiv preprint arXiv:2607.21716, 2026 Appendix A Complete Proofs This appendix provides complete proofs for all theoretical results presented in the main paper. We organize the material in the following structure: • Appendix A.1 provides the full proof for Lemma 3.1, establishing the error bound between the inference loss and the training loss in T-step generation. The proof is via a recursive expansion of the prediction error using first-order Taylor expansion. Then, we derive the bound via spectral norms of Jacobians taken through the input space. • Appendix A.2 establishes the error bounds and noise stability analysis for linear self-attention transformers trained on the weight prediction task for linear functions. We begin with preliminaries that reformulate the training and inference losses, then present the complete proof of Theorem A.5 (Section A.2.1), where standard Gaussian concentration bounds are cited and applied directly. • Appendix A.3 extends the analysis to nonlinear self-attention transformers trained on quadratic functions. We first describe the problem setup, then provide the explicit construction of a three-head transformer that implements gradient descent on quadratic losses (Section A.3.1). We derive the gradient dynamics under chain-of-thought training (Section A.3.2) and prove the convergence in Theorem 3.4 (Section A.3.3). Notations. For any vector v, let ‖v‖\|v\| be the Euclidean norm of v. For any matrix X∈m×nX∈ ^m× n, let ∥X∥F\|X\|__F be the Frobenius norm and ‖X‖2 \|X \|_2 be the spectral norm of X. If X is a square matrix, then the trace of X, Tr⁡[X] [X ], equals the sum of X’s diagonal entries. For two functions f⁡(n)f(n) and g⁡(n)g(n), we write g⁡(n)=O⁡(f⁡(n))g(n)=O(f(n)) if there exists a fixed value C that does not grow with n such that g⁡(n)≤C⋅f⁡(n)g(n)≤ C· f(n) when n is large enough. We write g⁡(n)=Θ⁡(f⁡(n))g(n)= (f(n)) if there exists a fixed value C1C_1 and C2C_2 that do not grow with n such that C1f(n)≤g⁡(n)≤C2⋅f⁡(n)C_1f(n)≤ g(n)≤ C_2· f(n) when n is large enough. A.1 Proof of Lemma 3.1 In this section, we provide the full proof for Lemma 3.1, which shows a bound on the gap between the inference loss and the true prediction loss conditioned on the ground truth. For a given input ∼x , let y(t)y^(t) denote the ground-truth label at step t. We denote the model’s predictions under two regimes: 1. The prediction conditioned on the ground-truth preceding steps, as in the training loss, is denoted as fW((t−1))f_W(z^(t-1)). 2. the prediction based on t-step autoregressive generation in the inference loss is denoted as fW(^(t−1))f_W( z^(t-1)). We define their deviations from the label y(t)y^(t) at step t as e(t)()=fW((t−1))−y(t)e^(t)(x)=f_W(z^(t-1))-y^(t), and e^(t)()=fW(^(t−1))−y(t) e^(t)(x)=f_W( z^(t-1))-y^(t). The prediction gap caused by autoregressive accumulation is Δ(t)()=fW(^(t−1)())−fW((t−1)()). ^(t)(x)=f_W( z^(t-1)(x))-f_W(z^(t-1)(x)). By definition, we have the relation: e^(t)()=e(t)()+Δ(t)(). e^(t)(x)=e^(t)(x)+ ^(t)(x). For the Jacobian norms, we write ‖Ji(t)‖:=sup∈∗‖Ji(t)()‖\|J^(t)_i\|:= _x \|J^(t)_i(x)\| for the norm with the supremum taken, so that ρi(t)=‖Ji(t)‖+∑k=i+1t−1‖Jk(t)‖ρi(k)ρ^(t)_i=\|J^(t)_i\|+ _k=i+1^t-1\|J^(t)_k\|\,ρ^(k)_i with ρt(t)=1ρ^(t)_t=1. Next, we present the complete proof of Lemma 3.1, based on a recursive decomposition for Δ(t)() ^(t)(x). Proof of Lemma 3.1. First, we bound the loss gap via Lipschitz continuity. We evaluate the two losses at step T. By the assumption that the loss function ℓ⁡(⋅,⋅) (·,·) is C1C_1-Lipschitz continuous with respect to its first argument, the absolute gap between the inference loss and the reference loss for a specific input x is bounded by: |ℓ⁡(fW(^(T−1)()),y(T))−ℓ⁡(fW((T−1)()),y(T))|≤C1‖Δ(T)()‖. | (f_W( z^(T-1)(x)),y^(T) )- (f_W(z^(T-1)(x)),y^(T) ) |\;≤\;C_1\| ^(T)(x)\|. (6) To bound the loss gap, we bound the norm of the final prediction gap ‖Δ(T)()‖\| ^(T)(x)\|. We conduct a recursive expansion of the prediction gap ‖Δ(T)()‖\| ^(T)(x)\|. At the base case of t=1t=1, both predictions rely solely on x. Thus, Δ(1)()=0 ^(1)(x)=0 and ‖e^(1)()‖=‖e(1)()‖\| e^(1)(x)\|=\|e^(1)(x)\|. For an arbitrary step t>1t>1, we apply first-order Taylor’s expansion to expand the prediction on the generated prefix around the correct step sequence. Subtracting fW((t−1)())f_W(z^(t-1)(x)) from fW(^(t−1))f_W( z^(t-1)), we obtain the prediction gap: Δ(t)()=∑i=1t−1Ji(t)()e^(i)()+ℛ(t)(), ^(t)(x)= _i=1^t-1J^(t)_i(x)\, e^(i)(x)+R^(t)(x), (7) where Ji(t)()J^(t)_i(x) is the step-to-step Jacobian evaluated at the ground-truth context, and ℛ(t)()R^(t)(x) is the higher-order Taylor remainder. Under the assumption that the Jacobian of fWf_W is C2C_2-Lipschitz continuous, the norm of the remainder ℛ(t)()R^(t)(x) is bounded by a quadratic function of the input deviation, based on the mean value theorem: ‖ℛ(t)()‖≤C22(∑i=1t−1‖e^(i)()‖)2≤C22(∑i=1t−1ri)2,\|R^(t)(x)\|\;≤\; C_22 ( _i=1^t-1\| e^(i)(x)\| )^2≤ C_22 ( _i=1^t-1r_i )^2, (8) based on equation (3) (with ℛ(1)()=0R^(1)(x)=0). Applying the triangle inequality to equation (7), we have: ‖Δ(t)()‖≤∑i=1t−1‖Ji(t)‖‖e^(i)()‖+‖ℛ(t)()‖.\| ^(t)(x)\|\;≤\; _i=1^t-1\|J^(t)_i\|\| e^(i)(x)\|+\|R^(t)(x)\|. (9) Using e^(t)()=e(t)()+Δ(t)() e^(t)(x)=e^(t)(x)+ ^(t)(x) and applying the triangle inequality, the inference error norm is bounded by: ‖e^(t)()‖≤‖e(t)()‖+‖Δ(t)()‖.\| e^(t)(x)\|≤\|e^(t)(x)\|+\| ^(t)(x)\|. Substituting equation (9) above yields the full recurrence relation for the inference error: ‖e^(t)()‖≤st()+∑i=1t−1‖Ji(t)‖‖e^(i)()‖, where st():=‖e(t)()‖+‖ℛ(t)()‖.\| e^(t)(x)\|\;≤\;s_t(x)+ _i=1^t-1\|J^(t)_i\|\| e^(i)(x)\|, where s_t(x):=\|e^(t)(x)\|+\|R^(t)(x)\|. (10) To solve this recurrence, we recursively expand ‖e^(i)()‖\| e^(i)(x)\| in si()s_i(x). The first-order components become a sum of all possible directed step-to-step paths. The product of the Jacobian spectral norms along these paths forms the cumulative amplification factor ρi(t)ρ^(t)_i, defined recursively as: ρi(t)=sup∈∗‖Ji(t)()‖+∑k=i+1t−1(sup∈∗‖Jk(t)()‖⋅ρi(k)), with ρt(t)=1.ρ^(t)_i= _x \|J^(t)_i(x)\|+ _k=i+1^t-1 ( _x \|J^(t)_k(x)\|·ρ^(k)_i ), with ρ^(t)_t=1. Formally, we show that ‖e^(t)()‖≤∑i=1tρi(t)si(), for every t≥1, \| e^(t)(x)\|\;≤\; _i=1^tρ^(t)_i\,s_i(x), for every t≥ 1, (11) which holds at t=1t=1 by the base case. Assuming equation (11) up to step t−1t-1, substituting it into equation (10), we have ‖e^(t)()‖ \| e^(t)(x)\| ≤st()+∑k=1t−1(∑i=kt−1‖Ji(t)‖ρk(i))sk() ≤ s_t(x)+ _k=1^t-1 ( _i=k^t-1\|J^(t)_i\|ρ^(i)_k )s_k(x) ≤st()+∑k=1t−1(‖Jk(t)‖ρk(k)+∑i=k+1t−1‖Ji(t)‖ρk(i))sk() ≤ s_t(x)+ _k=1^t-1 (\|J^(t)_k\|\,ρ^(k)_k+ _i=k+1^t-1\|J^(t)_i\|\,ρ^(i)_k )s_k(x) ≤st()+∑k=1t−1ρk(t)sk(), ≤ s_t(x)+ _k=1^t-1ρ^(t)_ks_k(x), (12) which completes the induction. By replacing ρi(T)ρ^(T)_i and carrying each remainder to step T by its own factor ρt(T)ρ^(T)_t, the final prediction gap at step T is bounded by: ‖Δ(T)()‖≤∑i=1T−1ρi(T)‖e(i)()‖+∑t=2Tρt(T)‖ℛ(t)()‖.\| ^(T)(x)\|\;≤\; _i=1^T-1ρ^(T)_i\|e^(i)(x)\|\;+\; _t=2^Tρ^(T)_t\|R^(t)(x)\|. (13) Substituting equation (13) into equation (6), we obtain: |ℓ⁡(fW(^(T−1)()),y(T))−ℓ⁡(fW((T−1)()),y(T))| | (f_W( z^(T-1)(x)),y^(T) )- (f_W(z^(T-1)(x)),y^(T) ) | ≤ ≤ C1∑i=1T−1ρi(T)‖e(i)()‖+C1∑t=2Tρt(T)‖ℛ(t)()‖. C_1 _i=1^T-1ρ^(T)_i\|e^(i)(x)\|\;+\;C_1 _t=2^Tρ^(T)_t\|R^(t)(x)\|. (14) For every x in the support of D, recall from equation (3) that ‖e(i)()‖≤εi\|e^(i)(x)\|≤ _i and ‖ℛ(t)()‖≤C22(∑i<tri)2\|R^(t)(x)\|≤ C_22 ( _i<tr_i )^2 from equation (8), we have |ℓ⁡(fW(^(T−1)()),y(T))−ℓ⁡(fW((T−1)()),y(T))| | (f_W( z^(T-1)(x)),y^(T) )- (f_W(z^(T-1)(x)),y^(T) ) | ≤ ≤ C1∑i=1T−1ρi(T)εi+C1C22∑t=2Tρt(T)(∑i=1t−1ri)2. C_1 _i=1^T-1ρ^(T)_i _i\;+\; C_1C_22 _t=2^Tρ^(T)_t ( _i=1^t-1r_i )^2. (15) Taking the expectation over the data distribution ∼x , we have |L⁡(fW,)−L⋆(fW,)| |L(f_W;D)-L (f_W;D) | =|∼[ℓ⁡(fW(^(T−1)()),y(T))−ℓ⁡(fW((T−1)()),y(T))]| = | E_x [ (f_W( z^(T-1)(x)),y^(T) )- (f_W(z^(T-1)(x)),y^(T) ) ] | ≤∼[|ℓ⁡(fW(^(T−1)()),y(T))−ℓ⁡(fW((T−1)()),y(T))|] ≤ E_x [ | (f_W( z^(T-1)(x)),y^(T) )- (f_W(z^(T-1)(x)),y^(T) ) | ] ≤C1∑i=1T−1ρi(T)εi+C1C22∑t=2Tρt(T)(∑i=1t−1ri)2, ≤ C_1 _i=1^T-1ρ^(T)_i _i\;+\; C_1C_22 _t=2^Tρ^(T)_t ( _i=1^t-1r_i )^2, (16) which completes the proof of equation (4). ∎ A.2 Proof for the Linear Case We study the inference loss of a one-layer linear self-attention transformer trained to predict weight vectors of linear functions. Following the setup in Section 3.2, we describe the in-context weight prediction task with linear functions. The weight vector is θ⋆∼⁡(0,Idd)θ (0, _d), and the n in-context examples are xj∼⁡(0,Idd)x_j (0, _d) with labels yj=⟨θ⋆,xj⟩y_j= θ ,x_j for j∈[n]j∈[n], all independent. We write X=[x1,…,xn]∈d×nX=[x_1,…,x_n]∈ ^d× n, y=[y1,…,yn]=θ⋆⁣⊤Xy=[y_1,…,y_n]=θ X, and S:=1nXX⊤S:= 1nX . The chain of thought consists of the gradient descent iterates on the squared loss on average over n examples with learning rate η, started at θ0=0d _0=0_d: θi+1=θi−ηS(θi−θ⋆),so thatθi=(I−(I−ηS)i)θ⋆for i≥0. _i+1= _i-η S( _i-θ ), that _i= (I-(I-η S)^i )θ i≥ 0. (17) With de=2d+2d_e=2d+2, the input that contains the correct steps θ0,…,θi _0,…, _i is Zi=[X0d0d⋯0dy00⋯00d×nθ0θ1⋯θi01×n11⋯1]∈de×(n+i+1).Z_i= bmatrixX&0_d&0_d&·s&0_d\\ y&0&0&·s&0\\ 0_d× n& _0& _1&·s& _i\\ 0_1× n&1&1&·s&1 bmatrix∈ ^d_e×(n+i+1). The label of step i+1i+1 is θ¯i+1:=θi+1 θ_i+1:= _i+1 for i≤T−2i≤ T-2, and the label of the final step is θ¯T:=θ⋆ θ_T:=θ . We first rewrite the training and inference losses. During training, the network is trained to predict the next iterate given the preceding ground-truth sequences. We consider the gradient flow dynamics over the population loss: LCoT(fW)=X,θ⋆[12∑i=0T−1‖fW(Zi)−(0d,0,θ¯i+1,1)‖2]. L^CoT(f_W)= E_X,θ [ 12 _i=0^T-1\|f_W(Z_i)-(0_d,0, θ_i+1,1)\|^2 ]. At inference time, we apply the model to generate T steps before it outputs the final prediction. Starting from Z0Z_0, we define the model prediction at each step as Z^i+1=[Z^i,fW(Z^i)] Z_i+1= [ Z_i,f_W( Z_i) ], for i=0,1,…,T−1i=0,1,…,T-1, where the last column of Z^i Z_i encodes the estimate θ^i θ_i, and θ^0=θ0 θ_0= _0. Then, the inference loss, in equation (2), with the squared loss, can be written as: L⁡(fW,)=X,θ⋆[12‖fW(Z^T−1)−(0d,0,θ⋆,1)‖22]. L(f_W;D)= E_X,θ [ 12\|f_W( Z_T-1)-(0_d,0,θ ,1)\|_2^2 ]. (18) For linear functions, we consider a one-layer linear self-attention transformer with a residual connection. Given an input embedding matrix Z∈ℝde×mZ ^d_e× m of length m, we denote fW(Z)=Z[:,−1]+WPVZ⋅Z⊤WKQZ[:,−1]n, f_W(Z)=Z_[:,-1]+W^PVZ· Z W^KQZ_[:,-1]n, (19) We denote WPV∈ℝde×deW^PV ^d_e× d_e as the weight matrix that merges the projection and value matrix, and WKQ∈ℝde×deW^KQ ^d_e× d_e is the weight matrix that merges the key and query matrix. [:,−1][:,-1] indicates taking the last column of the matrix. n is the number of in-context examples. Next, we state results from the previous paper [65], which we will reuse in our statements. Under Assumption A.1, prior work [65] proves that, if we initialize the weight matrices with the block structure in Assumption A.1, the blocks that are initialized with 00 will remain 00 along the gradient flow trajectory. The initialization method is specified as follows to zero out the irrelevant blocks. Assumption A.1 (Initialization). Let σ0>0 _0>0 be a constant. We assume the parameters are initialized as WPV(0)=[0d×d0d×10d×d0d×101×d01×101×d01×1W31PV(0)0d×10d×d0d×101×d01×101×d01×1],WKQ(0)=[0d×d0d×1W13KQ(0)0d×101×d01×101×dW24KQ0d×d0d×10d×d0d×101×d01×101×d01×1]. W^PV(0)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0_1× 1&0_1× d&0_1× 1\\ W^PV_31(0)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0_1× 1&0_1× d&0_1× 1 bmatrix,W^KQ(0)= bmatrix0_d× d&0_d× 1&W^KQ_13(0)&0_d× 1\\ 0_1× d&0_1× 1&0_1× d&W^KQ_24\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0_1× 1&0_1× d&0_1× 1 bmatrix. Here, W31PV(0)=∑i=1dλiPVuiui⊤W_31^PV(0)= _i=1^d _i^PVu_iu_i and W13KQ(0)=∑i=1dλiKQuiui⊤W_13^KQ(0)= _i=1^d _i^KQu_iu_i , where uii=1d\u_i\_i=1^d is an orthonormal basis. The matrices W31PV(0)W_31^PV(0) and W13KQ(0)W_13^KQ(0) are symmetric and simultaneously diagonalizable, with eigenvalues satisfying λiPV≤−σ0 _i^PV≤- _0, λiKQ∈[σ0,1/2] _i^KQ∈[ _0,1/2], and for all i∈[d]i∈[d]. Finally, we fix W24KQ=−1W^KQ_24=-1 for all t≥0t≥ 0. Lemma A.2 (Block-wise invariance under gradient flow [65]). Under Assumption A.1, consider the gradient flow on the population CoT loss as dWdt=−∇LCoT(W). dWdt=-∇ L^CoT(W). Then, for all t≥0t≥ 0, the parameters WPV(t),WKQ(t)W^PV(t),W^KQ(t) take the following form: WPV(t)=[0d×d0d×10d×d0d×101×d001×d0W31PV(t)0d×10d×d0d×101×d001×d0],WKQ(t)=[0d×d0d×1W13KQ(t)0d×101×d001×dW24KQ0d×d0d×10d×d0d×101×d001×d0].W^PV(t)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_1× d&0\\ W^PV_31(t)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_1× d&0 bmatrix,W^KQ(t)= bmatrix0_d× d&0_d× 1&W^KQ_13(t)&0_d× 1\\ 0_1× d&0&0_1× d&W^KQ_24\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_1× d&0 bmatrix. Then, we can rewrite the linear self-attention model only in terms of W31PVW^PV_31 and W13KQW^KQ_13. Definition A.3 (Reduced transformer). Define S:=1nXX⊤S:= 1nX . Under Assumption A.1, for a given input ZiZ_i, we can rewrite the linear self-attention model by fW(Zi)=[0d,0,(θi+W31PVS(W13KQθi−θ⋆)),1]⊤.f_W(Z_i)= [0_d,0, ( _i+W^PV_31S (W^KQ_13 _i-θ ) ),1 ] . where θi _i is the last column of Zi[d+1:2d+1,−1]Z_i_[d+1:2d+1,-1] and θ⋆θ is the ground-truth weights of the linear function. Then, we can denote a reduced model by: f~W(Zi)=θi+W31PVS(W13KQθi−θ⋆). f_W(Z_i)= _i+W^PV_31S (W^KQ_13 _i-θ ). The right-hand side depends on ZiZ_i only through its last column. The same formula holds for the inputs during inference Z^i Z_i. Below, we write f~W(θ):=θ+W31PVS(W13KQθ−θ⋆) f_W(θ):=θ+W^PV_31S(W^KQ_13θ-θ ) for θ∈dθ∈ ^d, with X and θ⋆θ fixed. This map is affine, and its Jacobian is given by J:=∂f~W(θ)∂θ=I+W31PVSW13KQ, J:= ∂ f_W(θ)∂θ=I+W^PV_31SW^KQ_13, (20) which does not depend on θ. Theorem A.4 (Global convergence [65]). Suppose n=Θ⁡(dlog5d)n= (d ^5d), η∈(0.1,0.9)η∈(0.1,0.9), T=⌈clog⁡d⌉T= c d , and clog⁡(11−η)>2c ( 11-η)>2. Under Assumption A.1 with σ0>3(1−η)(2−η)T _0> 3(1-η)(2-η)T, if we run gradient flow on the population loss LCoT(fW)L^CoT(f_W), then after time t=O⁡(log⁡d+log⁡1ϵ)t=O( d+ 1ε), for any ϵ≥Θ⁡(log⁡ddclog⁡(11−η)−2)ε≥ ( dd^c ( 11-η)-2), we have LCoT(fW⁡(t))≤εL^CoT(f_W(t))≤ . Moreover, there exist matrices A and B, with ‖A‖op,‖B‖op≤Θ⁡(d−12clog⁡(11−η)+12)\|A\|_op,\|B\|_op≤ (d^- 12c ( 11-η)+ 12), such that W31PV(t)=(−ηI+A)W^PV_31(t)=(-η I+A) and W13KQ(t)=(I+B)W^KQ_13(t)=(I+B). We state Theorem A.5, which is the main result for the linear case. Theorem A.5 (Slightly adapted from [65]). Suppose n=Θ⁡(dlog5d)n= (d ^5d), η∈(0.1,0.9)η∈(0.1,0.9), and T=⌈clog⁡d⌉T= c d , with clog⁡(11−η)>2c ( 11-η )>2. Under the initialization in Assumption A.1 with σ0>3(1−η)(2−η)T _0> 3(1-η)(2-η)T, if we run gradient flow on the population loss LCoT(fW)L^CoT(f_W), after time t=O⁡(log⁡d+log⁡1ϵ0)t=O( d+ 1 _0), for any ϵ0=Θ⁡(log⁡ddclog⁡(11−η)−2) _0= ( dd^c ( 11-η)-2), we have L⁡(fW,)≤ϵ1L(f_W;D)≤ _1, for any ϵ1∈Θ⁡((d−clog⁡(11−η)+2log⁡d),1) _1∈ ((d^-c ( 11-η)+2 d),1 ). A.2.1 Proof of Theorem A.5 Next, we prove the bounds on inference loss and the noise stability measure of the linear self-attention model on the weight prediction task. Recall that, to evaluate the inference loss, we apply the model fWf_W to generate T steps before it outputs the final prediction. Starting from the input Z0Z_0, we denote the model’s generated output at step i (also serves as the model input at step i+1i+1) as follows: Z^i+1(W):=[Z^i(W),fW(Z^i(W))],i=0,1,…,T−1, Z_i+1(W):= [ Z_i(W),f_W( Z_i(W)) ], i=0,1,…,T-1, where we denote Z^0:=Z0 Z_0:=Z_0. Then, we denote the last column of Z^i(W) Z_i(W) as z^i(W):=Z^i(W)[:,−1] z_i(W):= Z_i(W)_[:,-1] and θ^i(W):=Z^i(W)[d+1:2d+1,−1] θ_i(W):= Z_i(W)_[d+1:2d+1,-1]. We can write down the block structure of Z^i Z_i as Z^i(W)=[X0d0d…0dy00…00d×nθ0θ^1(W)…θ^i(W)01×n11…1]. Z_i(W)= bmatrixX&0_d&0_d&…&0_d\\ y&0&0&…&0\\ 0_d× n& _0& θ_1(W)&…& θ_i(W)\\ 0_1× n&1&1&…&1 bmatrix. Correspondingly, we denote the ground-truth input with i intermediate steps as ZiZ_i, θi _i as ground-truth weight iterates, and we denote the final label by zT:=(0d,0,θ⋆,1)z_T:=(0_d,0,θ ,1). When W is fixed, we write θ^i θ_i for θ^i(W) θ_i(W). Proof of Theorem A.5. We will first write the inference error as a sum of prediction errors at each step, each carried to the final step by a power of the Jacobian of the trained transformer. This is a case of Lemma 3.1, where the Jacobian becomes a single matrix J, and the second-order remainder is zero. The inference loss is therefore small once J is contractive. At convergence, we will show that J is close to the Jacobian I−ηSI-η S, whose norm ‖J‖2<1\|J\|_2<1, and the errors are small, which yields the rate for the inference loss. Let W=W⁡(t)W=W(t) be the trained weights at time t given by Theorem A.4. By Definition A.3, we have the simplified form of the transformer: θ^i=f~W(θ^i−1)=θ^i−1+W31PVS(W13KQθ^i−1−θ⋆), θ_i= f_W( θ_i-1)= θ_i-1+W^PV_31S (W^KQ_13 θ_i-1-θ ), for i=1,…,Ti=1,…,T, and θ^T θ_T is the θ-block of the final output fW(Z^T−1)f_W( Z_T-1). Therefore, we can write L⁡(fW,)=12X,θ⋆[‖θ^T−θ⋆‖2].L(f_W;D)= 12 E_X,θ [\| θ_T-θ \|^2 ]. We compare the generated estimates θ^i θ_i with the gradient descent iterates θi _i. Let Δi:=f~W(θi−1)−θi,i=1,…,T, _i:= f_W( _i-1)- _i, i=1,…,T, be the prediction error of step i, when the model is given the correct previous step. Since f~W f_W is affine with the Jacobian J in equation (20), θ^i−θi=(f~W(θ^i−1)−f~W(θi−1))+Δi=J⁡(θ^i−1−θi−1)+Δi. θ_i- _i= ( f_W( θ_i-1)- f_W( _i-1) )+ _i=J( θ_i-1- _i-1)+ _i. Recursively decomposing the above from θ^0=θ0 θ_0= _0, we have θ^T−θ⋆=(θT−θ⋆)+∑i=1TJT−iΔi. θ_T-θ =( _T-θ )+ _i=1^TJ^T-i _i. (21) This equation is the form of Lemma 3.1 for the linear transformer. The first term, θT−θ⋆ _T-θ , is the error of T gradient descent steps on the given in-context examples, because the final label is θ⋆θ . Since ‖J‖\|J\| is not bounded uniformly over Gaussian inputs, we bound equation (21) in expectation rather than with the worst-case constants of Lemma 3.1. By the Cauchy–Schwarz inequality, L⁡(fW,)≤T+12(X,θ⋆[‖θT−θ⋆‖2]+∑i=1TX,θ⋆[‖JT−iΔi‖2]). L(f_W;D)≤ T+12 ( E_X,θ [\| _T-θ \|^2 ]+ _i=1^T E_X,θ [\|J^T-i _i\|^2 ] ). (22) It remains to show that J is contractive at convergence and that the residuals are small. (1) Upper bound of ‖J‖2E\|J\|_2. Now, we analyze the spectral norm of the Jacobian matrix of the linear transformer, J=I+W31PVSW13KQJ=I+W^PV_31SW^KQ_13. Based on the results in Theorem A.4, for the trained transformer by running the gradient flow on the population loss on the chain-of-thought sequences, there exists matrix A and B, such that W31=(−ηI+A)W_31=(-η I+A) and W13=(I+B)W_13=(I+B) with ‖A‖op,‖B‖op≤Θ⁡(d−12clog⁡(11−η)+12)\|A\|_op,\|B\|_op≤ (d^- 12c ( 11-η)+ 12). We denote β:=Θ⁡(d−12clog⁡(11−η)+12), and a:=clog⁡(11−η)>2.β:= (d^- 12c ( 11-η)+ 12), and a:=c ( 11-η)>2. In particular, β=Θ(d1−a2)≤d−1/2β= (d 1-a2)≤ d^-1/2 for a large enough d. Therefore, we have: J J =I+W31PVSW13KQ =I+W^PV_31SW^KQ_13 =I+(−ηI+A)S(I+B) =I+(-η I+A)S(I+B) =I−ηS+(AS−ηSB+ASB) =I-η S+(AS-η SB+ASB) =J0+E, =J_0+E, where J0=I−ηSJ_0=I-η S and E=(AS−ηSB+ASB)E=(AS-η SB+ASB). Recall that S=1nXX⊤S= 1nX where X=[x1,…,xn]X=[x_1,…,x_n] with xix_i sampled from ⁡(0,Idd)N(0, _d). By the Gaussian singular-value bound [42] applied to X, given n=Θ⁡(d2log5d)n= (d^2 ^5d), the event :=∥S−I∥2≤Cd/nG:=\\|S-I\|_2≤ C d/n\ has probability at least 1−2e−d/21-2e^-d/2 for a constant C. On G, we have that ∥J0∥2=∥I−ηS∥2≤1−η+Cηd/n=1−η+O(log−5/2d),and ∥S∥2≤2.\|J_0\|_2=\|I-η S\|_2≤ 1-η+Cη d/n=1-η+O( ^-5/2d),and \|S\|_2≤ 2. For matrix E, on G, submultiplicativity gives ‖E‖2≤(‖A‖2+η‖B‖2+‖A‖2‖B‖2)‖S‖2=Θ⁡(d−clog⁡11−η+1).\|E\|_2≤ (\|A\|_2+η\|B\|_2+\|A\|_2\|B\|_2 )\|S\|_2= (d^-c 11-η+1). Therefore, on G, ∥J∥2≤∥J0∥2+∥E∥2≤1−η+O(log−5/2d)+Θ(d−clog⁡11−η+1). \|J\|_2≤\|J_0\|_2+\|E\|_2≤ 1-η+O( ^-5/2d)+ (d^-c 11-η+1). (23) Thus, ‖J‖2≤1−η/2<1\|J\|_2≤ 1-η/2<1 for a large enough d. (2) Upper bound of ‖Δi‖2E\| _i\|^2. We have the one-step gradient descent update as θi=θi−1−ηS(θi−1−θ⋆) _i= _i-1-η S( _i-1-θ ). Then, the prediction error of step i conditioned on the correct steps is Δi _i =f~W(θi−1)−θi = f_W( _i-1)- _i =(W31SW13+ηS)θi−1−(W31+ηI)Sθ⋆=(AS−ηSB+ASB)θi−1−ASθ⋆. =(W_31SW_13+η S)\, _i-1-(W_31+η I)Sθ =(AS-η SB+ASB)\, _i-1-AS\,θ . Therefore, we have ‖Δi‖2 \| _i\|^2 ≤2(‖AS−ηSB+ASB‖22‖θi−1‖2+‖AS‖22‖θ⋆‖2). ≤ 2 (\|AS-η SB+ASB\|_2^2\| _i-1\|^2+\|AS\|_2^2\|θ \|^2 ). (24) Note that θi _i is the i-th iterate of the gradient descent applied to (X,y)(X,y). Given that θ⋆∼⁡(0,Idd)θ (0, _d) and S=1nXXTS= 1nX^T, with n=Θ⁡(d2log5d)n= (d^2 ^5d), we have X,θ⋆[‖θi‖2]≤O⁡(d) E_X,θ [\| _i\|^2 ]≤ O(d). Therefore, based on θi=(I−(I−ηS)i)θ⋆ _i=(I-(I-η S)^i)θ , on G, we have ‖I−ηS‖2≤1\|I-η S\|_2≤ 1, so ‖θi‖≤2‖θ⋆‖\| _i\|≤ 2\|θ \|. On G, we also have ‖AS−ηSB+ASB‖2≤6β\|AS-η SB+ASB\|_2≤ 6β, and ‖AS‖2≤2β\|AS\|_2≤ 2β. Therefore, plug those back to equation (24), for i=1,…,Ti=1,…,T, we have: ⁡[‖Δi‖22]≤2(36β2⋅4d+4β2d)=O⁡(dβ2)=O⁡(d2−a). [\| _i\|_2^21_G ]≤ 2 (36β^2· 4d+4β^2d )=O(dβ^2)=O(d^2-a). Upper bound of L⁡(fW,)L(f_W;D). Now, we have shown that ⁡[‖Δi‖2]=O⁡(d2−a)E[\| _i\|^21_G]=O(d^2-a) for i=1,…,Ti=1,…,T, and ‖J‖2≤1−η/2\|J\|_2≤ 1-η/2 on G. Then, summing the geometric series gives ∑i=1T⁡[‖JT−iΔi‖2]≤∑i=1T(1−η/2)2(T−i)[‖Δi‖2]≤O⁡(d2−a). _i=1^TE [\|J^T-i _i\|^21_G ]≤ _i=1^T(1-η/2)^2(T-i)\,E [\| _i\|^21_G ]≤ O(d^2-a). In addition, the first term of equation (22) is the error of exact gradient descent. By equation (17), θT−θ⋆=−(I−ηS)Tθ⋆ _T-θ =-(I-η S)^Tθ . On G, we have ‖I−ηS‖22T≤(1−η)2T(1+Cη1−ηd/n)2T≤2(1−η)2T.\|I-η S\|_2^2T≤(1-η)^2T (1+ Cη1-η d/n )^2T≤ 2(1-η)^2T. Since T≥clog⁡dT≥ c d, we have ⁡[‖θT−θ⋆‖2]≤2(1−η)2T‖θ⋆‖2≤2d1−2a.E [\| _T-θ \|^21_G ]≤ 2(1-η)^2T\,E\|θ \|^2≤ 2d^1-2a. Combining these results into equation (22), by the Cauchy–Schwarz inequality, we have L⁡(fW,)≤T+12(2d1−2a+O⁡(d2−a)+e−Ω⁡(d))=O⁡(d2−alog⁡d), L(f_W;D)≤ T+12 (2d^1-2a+O(d^2-a)+e^- (d) )=O(d^2-a d), (25) since T+1=O⁡(log⁡d)T+1=O( d). ∎ A.3 Proof for the Quadratic Case In this section, we study an in-context weight prediction task where the model is trained to predict the parameters of a quadratic function from n examples. Let the quadratic weight vector be θ⋆∈dθ ∈ ^d where θ⋆∼⁡(0,Idd)θ (0, _d). We sample n in-context examples xi∼⁡(0,Idd)x_i (0, _d) with the corresponding labels according to yi=(θ⋆⁣⊤xi)2y_i=(θ x_i)^2, for i=1,…,ni=1,…,n. The goal is to predict θ⋆θ given the n examples. Similar to weight prediction for the linear function, define the embedding dimension de=2d+2d_e=2d+2. The initial input sequence is constructed as Z0=[x1⋯xn0y1⋯yn00⋯0θ00⋯01]:=[X0dy00dθ001]∈ℝde×(n+1), Z_0= bmatrixx_1&·s&x_n&0\\ y_1&·s&y_n&0\\ 0&·s&0& _0\\ 0&·s&0&1 bmatrix:= bmatrixX&0_d\\ y&0\\ 0_d& _0\\ 0&1 bmatrix ^d_e×(n+1), (26) where X=[x1,…,xn]∈ℝd×nX=[x_1,…,x_n] ^d× n, y=[y1,…,yn]⊤∈ℝny=[y_1,…,y_n] ^n, and θ0∈ℝd _0 ^d is an initial estimate in the basin of θ⋆θ . We take θ0=α0θ⋆ _0= _0θ for a constant α0∈[3/4,1) _0∈[3/4,1). We also append a dummy constant 11 to each parameter token for formatting convenience. We construct the chain-of-thought using gradient descent on the squared loss, i.e., 12n∑j=1n((θ⊤xj)2−yj)2 12n _j=1^n ((θ x_j)^2-y_j )^2. Define a residual vector as e(θ)=y^(θ)−ye(θ)= y(θ)-y, where y^j(θ)=(θ⊤xj)2 y_j(θ)=(θ x_j)^2. The gradient of the loss is 2n∑j=1n((θ⊤xj)2−yj)xjxj⊤θ=2nXdiag⁡(e⁡(θ))X⊤θ. 2n _j=1^n ((θ x_j)^2-y_j )x_jx_j θ= 2nXdiag (e(θ) )X θ. Given a learning rate η, starting from θ0 _0, define the gradient descent iterates for i=1,…,Ti=1,…,T: θi=θi−1−η⋅2nXdiag⁡(e⁡(θi−1))X⊤θi−1=θi−1−η⋅2nXDi+Di−X⊤θi−1, _i= _i-1-η· 2nXdiag (e( _i-1) )X _i-1= _i-1-η· 2nXD_i^+D_i^-X _i-1, where Di+=diag⁡(X⊤θi−1+X⊤θ⋆)D_i^+=diag(X _i-1+X θ ) and Di−=diag⁡(X⊤θi−1−X⊤θ⋆)D_i^-=diag(X _i-1-X θ ). Using these iterates, we define the ground-truth sequences Zi=[x1⋯xn00⋯0y1⋯yn00⋯00⋯0θ0θ1⋯θi0⋯011⋯1]∈ℝde×(n+i+1), for i=1,2,…,T. Z_i= bmatrixx_1&·s&x_n&0&0&·s&0\\ y_1&·s&y_n&0&0&·s&0\\ 0&·s&0& _0& _1&·s& _i\\ 0&·s&0&1&1&·s&1 bmatrix ^d_e×(n+i+1), for i=1,2,…,T. (27) A.3.1 Proof of Lemma 3.3 Proof. We now construct a three-head nonlinear self-attention transformer that performs one gradient descent step for the quadratic squared loss. For head h∈1,2,3h∈\1,2,3\, we denote a three-head self-attention model as fW(Z)=Z[:,−1]+∑i=13WPViZ⋅ϕi(Z⊤WKQiZ[:,−1]n). f_W(Z)=Z_[:,-1]+ _i=1^3W^PV_iZ· _i ( Z W^KQ_iZ_[:,-1]n ). Here WVi,WKi,WQi∈ℝdi×deW^V_i,W^K_i,W^Q_i ^d_i× d_e for i=1,2,3i=1,2,3. One can set d1,d2d_1,d_2, and d3d_3 separately. We construct the transformer in the following way. Let ϕ⁡(⋅)φ(·) be applied elementwise as ϕ1(t)=t⊙3,ϕ2(t)=t⊙2,ϕ3(t)=t⊙2. _1(t)=t 3, _2(t)=t 2, _3(t)=t 2. (28) We then set WPV1=[0d×d0d×10d×d0d×101×d00d×10−2ηn2Idd0d×10d×d0d×101×d00d×10], W^PV_1= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ -2η n^2 _d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix, WKQ1=[0d×d0d×1Idd0d×101×d00d×100d×d0d×10d×d0d×101×d00d×10], W^KQ_1= bmatrix0_d× d&0_d× 1& _d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix, WPV2=[0d×d0d×10d×d0d×101×d00d×1012ηnIdd0d×10d×d0d×101×d00d×10], W^PV_2= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ 12η n _d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix, WKQ2=[0d×d0d×1Idd0d×101×d00d×110d×d0d×10d×d0d×101×d00d×10], W^KQ_2= bmatrix0_d× d&0_d× 1& _d&0_d× 1\\ 0_1× d&0&0_d× 1&1\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix, WPV3=[0d×d0d×10d×d0d×101×d00d×10−12ηnIdd0d×10d×d0d×101×d00d×10], W^PV_3= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ - 12η n _d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix, WKQ3=[0d×d0d×1Idd0d×101×d00d×1−10d×d0d×10d×d0d×101×d00d×10]. W^KQ_3= bmatrix0_d× d&0_d× 1& _d&0_d× 1\\ 0_1× d&0&0_d× 1&-1\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix. (29) Now, we can inductively calculate the i-th step of generation, showing that the output is exactly the parameter after the i-th gradient step. fW(Zi)= f_W(Z_i)= (0d,0,θi,1)+(0d,0,−2ηn2X⋅ϕ1(Z⊤WKQ1Z[:,−1]n) (0_d,0, _i,1)+ (0_d,0,-2η n^2X· _1 ( Z W^KQ_1Z_[:,-1]n ) +ηn2X⋅ϕ2(Z⊤WKQ2Z[:,−1]n)−ηn2X⋅ϕ3(Z⊤WKQ3Z[:,−1]n),0). + η n2X· _2 ( Z W^KQ_2Z_[:,-1]n )- η n2X· _3 ( Z W^KQ_3Z_[:,-1]n ),0 ). Based on the construction, we have the following: X⋅ϕ1(Z⊤WKQ1Z[:,−1])=X(X⊤θi)⊙3 X· _1 (Z W^KQ_1Z_[:,-1] )=X(X _i) 3 X⋅ϕ2(Z⊤WKQ2Z[:,−1])=X(X⊤θi+y⊤)⊙2 X· _2 (Z W^KQ_2Z_[:,-1] )=X(X _i+y ) 2 X⋅ϕ3(Z⊤WKQ3Z[:,−1])=X(X⊤θi−y⊤)⊙2 X· _3 (Z W^KQ_3Z_[:,-1] )=X(X _i-y ) 2 Therefore, we have fW(Zi) f_W(Z_i) =(0d,0,θi,1)+(0d,0,−2ηnX((X⊤θi)⊙3−(X⊤θi)⊙y⊤),0) =(0_d,0, _i,1)+(0_d,0,- 2ηnX ((X _i) 3-(X _i) y ),0) =(0d,0,θi−2ηnXdiag⁡(e⁡(θi))X⊤θi,1) =(0_d,0, _i- 2ηnXdiag (e( _i) )X _i,1) where we denote the residual error e(θ)=y^(θ)−ye(θ)= y(θ)-y, where y^j(θ)=(θ⊤xj)2 y_j(θ)=(θ x_j)^2. Then, after T steps of applying fWf_W, we obtain the T-step gradient descent solution using n examples of (x,y)(x,y). Denote the θT _T as the solution obtained by applying T-gradient step. Based on Lemma A.6, we have the test loss of fWf_W as LW=12X,θ⋆[‖fW(Z^T−1)−(0d,0,θ⋆,1)‖22]=12X,θ⋆[‖θT−θ⋆‖22]≤O⁡(d−2clog⁡(11−μηd)+1). L_W= 12 E_X,θ [ \|f_W( Z_T-1)-(0_d,0,θ ,1) \|_2^2 ]= 12 E_X,θ [ \| _T-θ \|_2^2 ]≤ O(d^-2c ( 11-μη d)+1). ∎ Lemma A.6. Let θ⋆∼⁡(0,Idd)θ (0, _d) and x1,…,xn∼i.i.d.⁡(0,Idd)x_1,…,x_n .i.d. N(0, _d). Define yj=(xj⊤θ⋆)2y_j=(x_j θ )^2 and the empirical loss ℒn(θ):=12n∑j=1n((xj⊤θ)2−yj)2=12n∑j=1n((xj⊤θ)2−(xj⊤θ⋆)2)2.L_n(θ):= 12n _j=1^n ((x_j θ)^2-y_j )^2= 12n _j=1^n ((x_j θ)^2-(x_j θ )^2 )^2. Let dist⁡(θ,±θ⋆):=min⁡‖θ−θ⋆‖2,‖θ+θ⋆‖2dist(θ,\±θ \):= \\|θ-θ \|_2,\ \|θ+θ \|_2\. Consider gradient descent θt+1=θt−η∇ℒn(θt),t=0,1,2,…,T _t+1= _t-η _n( _t),t=0,1,2,…,T. There exist absolute constants 0<m≤L<∞0<m≤ L<∞ and a radius ρ∈(0,1)ρ∈(0,1) such that for all θ satisfying dist⁡(θ,±θ⋆)≤ρ‖θ⋆‖2dist(θ,\±θ \)≤ρ\|θ \|_2, the Hessian of ℒnL_n obeys m‖θ⋆‖22Idd⪯∇2ℒn(θ)⪯L‖θ⋆‖22Idd.m\|θ \|_2^2 _d\ \ ∇^2L_n(θ)\ \ L\|θ \|_2^2 _d. Suppose the initialization lies in the basin, i.e., dist⁡(θ0,±θ⋆)≤ρ‖θ⋆‖2dist( _0,\±θ \)≤ρ\|θ \|_2. Assume that n≥Cd2log5dn≥ Cd^2 ^5d for a sufficiently large constant C>0C>0, 0<η≤12L‖θ⋆‖220<η≤ 12L\|θ \|_2^2, and T=⌈clog⁡d⌋T= c d . Then, with probability at least 1−d−101-d^-10 over (θ⋆,X)(θ ,X), the following holds dist(θT,±θ⋆)2≤dist(θ0,±θ⋆)2⋅d2clog⁡(1−mη‖θ⋆‖22). ( _T,\±θ \)^2\ ≤\ dist( _0,\±θ \)^2· d^2c (1-mη\|θ \|_2^2 ). Then, given η=Θ⁡(1/d)η= (1/d) and T=Θ⁡(log⁡d)T= ( d), we have X,θ⋆[dist(θT,±θ⋆)2]≤O⁡(d1−2clog⁡(11−mηd)). _X,θ [dist( _T,±θ )^2 ]≤ O (d^1-2c ( 11-mη d) ). Proof. We prove the claim on a high-probability event under which the empirical loss is uniformly strongly convex and smooth in a neighborhood of ±θ⋆\±θ \, and then apply a standard contraction argument for gradient descent. Let x∈ℝdx ^d and define ℓx(θ):=12((x⊤θ)2−(x⊤θ⋆)2)2. _x(θ):= 12 ((x θ)^2-(x θ )^2 )^2. Then ∇ℓx(θ)=((x⊤θ)2−(x⊤θ⋆)2)∇(x⊤θ)2=2((x⊤θ)2−(x⊤θ⋆)2)(x⊤θ)x,∇ _x(θ)= ((x θ)^2-(x θ )^2 )∇(x θ)^2=2 ((x θ)^2-(x θ )^2 )(x θ)x, and differentiating once more gives ∇2ℓx(θ) ∇^2 _x(θ) =2(∇((x⊤θ)2−(x⊤θ⋆)2)(x⊤θ)+((x⊤θ)2−(x⊤θ⋆)2)∇(x⊤θ))x⊤ =2 (∇ ((x θ)^2-(x θ )^2 )(x θ)+ ((x θ)^2-(x θ )^2 )∇(x θ) )x =2(2(x⊤θ)x(x⊤θ)+((x⊤θ)2−(x⊤θ⋆)2)x)x⊤ =2 (2(x θ)x(x θ)+ ((x θ)^2-(x θ )^2 )x )x =2(3(x⊤θ)2−(x⊤θ⋆)2)xx⊤. =2 (3(x θ)^2-(x θ )^2 )x . Therefore, the empirical Hessian is ∇2ℒn(θ)=1n∑j=1n∇2ℓxj(θ)=2n∑j=1n(3(xj⊤θ)2−(xj⊤θ⋆)2)xjxj⊤.∇^2L_n(θ)= 1n _j=1^n∇^2 _x_j(θ)= 2n _j=1^n (3(x_j θ)^2-(x_j θ )^2 )x_jx_j . Let ℬ:=θ∈ℝd:dist⁡(θ,±θ⋆)≤ρ‖θ⋆‖2B:=\θ ^d:\ dist(θ,\±θ \)≤ρ\|θ \|_2\. We use a claim that when n≥Cd2log5dn≥ Cd^2 ^5d (for C sufficiently large), with probability at least 1−d−101-d^-10 over (θ⋆,X)(θ ,X), the following event holds: ℰ:∀θ∈ℬ,m‖θ⋆‖22Idd⪯∇2ℒn(θ)⪯L‖θ⋆‖22Idd,E: ∀θ , m\|θ \|_2^2 _d ∇^2L_n(θ) L\|θ \|_2^2 _d, (30) for some absolute constants 0<m≤L<∞0<m≤ L<∞ and ρ∈(0,1)ρ∈(0,1). One can prove ℰE by a standard ε -net argument over ℬB together with Gaussian concentration for fourth-order polynomials. For example, see the local regularity condition established (cf. Section 7 of [44]), the empirical Hessian is uniformly well-conditioned in a neighborhood of ±θ⋆±θ . We condition on ℰE for the remainder of the proof. Assume θ0∈ℬ _0 . Since on ℰE we have ∇2ℒn(θ)⪯L‖θ⋆‖22Idd∇^2L_n(θ) L\|θ \|_2^2 _d for all θ∈ℬθ , the loss is L‖θ⋆‖22L\|θ \|_2^2-smooth on ℬB. With stepsize η≤1/(2L‖θ⋆‖22)η≤ 1/(2L\|θ \|_2^2), the standard descent lemma implies ℒn(θt+1)≤ℒn(θt)L_n( _t+1) _n( _t) as long as θt∈ℬ _t . Moreover, since ℒn(θ)L_n(θ) is strongly convex in a neighborhood of each minimizer ±θ⋆±θ on ℰE, and θ⋆θ is a minimizer with ∇ℒn(θ⋆)=0 _n(θ )=0, the GD iterates remain in the same local basin. Thus θt∈ℬ _t for all t≤Tt≤ T. Next, we show one-step contraction in distance to ±θ⋆\±θ \. Fix the sign s∈+1,−1s∈\+1,-1\ that minimizes ‖θ0−sθ⋆‖2\| _0-sθ \|_2 and define Δt:=θt−sθ⋆. _t:= _t-sθ . On ℰE, for all θ∈ℬθ we have m‖θ⋆‖22Idd⪯∇2ℒn(θ)⪯L‖θ⋆‖22Idd.m\|θ \|_2^2 _d ∇^2L_n(θ) L\|θ \|_2^2 _d. Since ∇ℒn(sθ⋆)=0 _n(sθ )=0 and ℒnL_n is twice differentiable, the mean value form of the gradient gives ∇ℒn(θt)−∇ℒn(sθ⋆)=(∫01∇2ℒn(sθ⋆+τ⁡(θt−sθ⋆))τ)(θt−sθ⋆)=:HtΔt, _n( _t)- _n(sθ )= ( _0^1∇^2L_n(sθ +τ( _t-sθ ))dτ )( _t-sθ )=:H_t _t, where HtH_t satisfies the same eigenvalue bounds: m‖θ⋆‖22Idd⪯Ht⪯L‖θ⋆‖22Idd.m\|θ \|_2^2 _d H_t L\|θ \|_2^2 _d. Therefore, the GD update can be written as Δt+1=θt+1−sθ⋆=θt−sθ⋆−η∇ℒn(θt)=Δt−ηHtΔt=(I−ηHt)Δt. _t+1= _t+1-sθ = _t-sθ -η _n( _t)= _t-η H_t _t=(I-η H_t) _t. Taking norms and using Ht⪰m‖θ⋆‖22IddH_t m\|θ \|_2^2 _d yields ‖Δt+1‖2≤‖I−ηHt‖op‖Δt‖2≤(1−ηm‖θ⋆‖22)‖Δt‖2,\| _t+1\|_2≤\|I-η H_t\|_op\ \| _t\|_2≤(1-η m\|θ \|_2^2)\| _t\|_2, where the last inequality uses η≤1/(2L‖θ⋆‖22)η≤ 1/(2L\|θ \|_2^2) so that all eigenvalues of I−ηHtI-η H_t lie in [0,1−ηm‖θ⋆‖22][0,1-η m\|θ \|_2^2]. Squaring gives the one-step contraction ‖Δt+1‖22≤(1−mη‖θ⋆‖22)2‖Δt‖22.\| _t+1\|_2^2≤(1-mη\|θ \|_2^2)^2\| _t\|_2^2. By induction, ‖ΔT−1‖22≤(1−mη‖θ⋆‖22)2T‖Δ0‖22.\| _T-1\|_2^2≤(1-mη\|θ \|_2^2)^2T\| _0\|_2^2. Since ‖Δt‖2=dist⁡(θt,±θ⋆)\| _t\|_2=dist( _t,\±θ \) by construction of s, and T=clog⁡dT=c d, dist(θT−1,±θ⋆)2 ( _T-1,\±θ \)^2 ≤dist(θ0,±θ⋆)2⋅(1−mη‖θ⋆‖22)2clog⁡d ( _0,\±θ \)^2·(1-mη\|θ \|_2^2)^2c d =dist(θ0,±θ⋆)2⋅d2clog⁡(1−mη‖θ⋆‖22). =dist( _0,\±θ \)^2· d^2c (1-mη\|θ \|_2^2). Finally, the probability statement follows since ℰE holds with probability at least 1−d−101-d^-10. ∎ A.3.2 Gradient Computation In this subsection, we compute the gradient of the full model and show that under certain initialization, running gradient flow on the training loss converges to the construction in Lemma 3.3. Recall that we train a three-head nonlinear self-attention transformer, denoted as fW(Z) f_W(Z) =Z[:,−1]+∑i=13WPViZ⋅ϕi(Z⊤WKQiZ[:,−1]n). =Z_[:,-1]+ _i=1^3W^PV_iZ· _i ( Z W^KQ_iZ_[:,-1]n ). Here, by the convention in prior work [65, 75], we merge WPiW^P_i and WViW^V_i into a single matrix and merge WKiW^K_i and WQiW^Q_i into a single matrix. WPViW^PV_i and WKQi∈ℝde×deW^KQ_i ^d_e× d_e for i=1,2,3i=1,2,3. By the construction, we set ϕi(⋅) _i(·) as elementwise non-linear functions ϕ1(t)=t⊙3,ϕ2(t)=t⊙2 _1(t)=t 3, _2(t)=t 2, and ϕ3(t)=t⊙2 _3(t)=t 2. The chain-of-thought training objective is written as LCoT(W)=xi∼⁡(0,Idd),θ⋆∼⁡(0,Idd)[12∑i=0T−1‖fW(Zi)−(0d,0,θ¯i+1,1)‖22]. L^CoT(W)= E_x_i (0, _d),θ (0, _d) [ 12 _i=0^T-1 \|f_W(Z_i)-(0_d,0, θ_i+1,1) \|_2^2 ]. Here, θ¯i+1:=θi+1 θ_i+1:= _i+1 for i≤T−2i≤ T-2 and θ¯T:=θ⋆ θ_T:=θ . As in the training loss equation (1), there are T steps. The label of the last step is the target θ⋆θ . The reduced loss and all gradients below use this convention. At inference time, we apply the model to generate T−1T-1 steps before it outputs the final prediction. The test loss is written as L⁡(W)=X,θ⋆[12‖fW(Z^T−1)−(0d,0,θ⋆,1)‖22]. L(W)= E_X,θ [ 12 \|f_W( Z_T-1)-(0_d,0,θ ,1) \|_2^2 ]. Now, we compute the gradients of WPViW^PV_i and WKQiW^KQ_i based on the loss of each CoT step. Lemma A.7 (Full model gradient). The gradients of WPViW^PV_i and WKQiW^KQ_i for i=1,2,3i=\1,2,3\ are given by the following. ∇WPV1LCoT(W)=X,θ⋆[∑i=0T−1ri(Ziϕ1(Zi⊤WKQ1Zi[:,−1]n))⊤] _W^PV_1L^CoT(W)= E_X,θ [ _i=0^T-1r_i (Z_i _1 ( Z_i W^KQ_1Z_i[:,-1]n ) ) ] ∇WPV2LCoT(W)=X,θ⋆[∑i=0T−1ri(Ziϕ2(Zi⊤WKQ2Zi[:,−1]n))⊤] _W^PV_2L^CoT(W)= E_X,θ [ _i=0^T-1r_i (Z_i _2 ( Z_i W^KQ_2Z_i[:,-1]n ) ) ] ∇WPV3LCoT(W)=X,θ⋆[∑i=0T−1ri(Ziϕ3(Zi⊤WKQ3Zi[:,−1]n))⊤] _W^PV_3L^CoT(W)= E_X,θ [ _i=0^T-1r_i (Z_i _3 ( Z_i W^KQ_3Z_i[:,-1]n ) ) ] ∇WKQ1LCoT(W)=X,θ⋆[1n∑i=0T−1Zi((Zi⊤WPV1⊤ri)⊙3(Zi⊤WKQ1Zi[:,−1]n)⊙2)Zi[:,−1]⊤] _W^KQ_1L^CoT(W)= E_X,θ [ 1n _i=0^T-1Z_i ( (Z_i W^PV_1 r_i ) 3( Z_i W^KQ_1Z_i[:,-1]n) 2 )Z_i[:,-1] ] ∇WKQ2LCoT(W)=X,θ⋆[1n∑i=0T−1Zi((Zi⊤WPV2⊤ri)⊙2(Zi⊤WKQ2Zi[:,−1]n))Zi[:,−1]⊤] _W^KQ_2L^CoT(W)= E_X,θ [ 1n _i=0^T-1Z_i ( (Z_i W^PV_2 r_i ) 2( Z_i W^KQ_2Z_i[:,-1]n) )Z_i[:,-1] ] ∇WKQ3LCoT(W)=X,θ⋆[1n∑i=0T−1Zi((Zi⊤WPV3⊤ri)⊙2(Zi⊤WKQ3Zi[:,−1]n))Zi[:,−1]⊤] _W^KQ_3L^CoT(W)= E_X,θ [ 1n _i=0^T-1Z_i ( (Z_i W^PV_3 r_i ) 2( Z_i W^KQ_3Z_i[:,-1]n) )Z_i[:,-1] ] where we define ri=∑h=13WPVhZi⋅ϕh(Zi⊤WKQhZi[:,−1]n)−(0d,0,θ¯i+1−θi,0)r_i= _h=1^3W^PV_hZ_i· _h ( Z_i W^KQ_hZ_i[:,-1]n )-(0_d,0, θ_i+1- _i,0) with θ¯i+1 θ_i+1 defined above. Proof. For a step index i=0,1,…,T−1i=0,1,…,T-1, write Z∈ℝde×mZ ^d_e× m with m=n+i+1m=n+i+1 and z=Z[:,−1]∈ℝdez=Z_[:,-1] ^d_e. We define sh:=Z⊤WKQhzn,vh:=ϕh(sh),uh:=WPVhZvh, for h=1,2,3.s_h:= Z W^KQ_hzn, v_h:= _h(s_h), u_h:=W^PV_hZv_h, for h=1,2,3. Then, we can write the squared loss in terms of the residual error in each step as ℓi(W):=12‖ri‖2, where ri=∑h=13uh−(0d,0,θi+1−θi,0), _i(W):= 12\|r_i\|^2, where r_i= _h=1^3u_h-(0_d,0, _i+1- _i,0), Then, we compute the gradients by the chain rule. For each head, the dependence of the loss on WPVhW^PV_h is only through uhu_h. Hence, ∇WPVhℓi(W)=ri(Zvh)⊤=ri(Zϕh(Z⊤WKQhzn))⊤. _W^PV_h _i(W)=r_i(Zv_h) =r_i (Z _h ( Z W^KQ_hzn ) ) . Summing over i and taking the expectation yields the stated formula. Then, we compute the gradient for WKQhW^KQ_h. Since vh=ϕh(sh)v_h= _h(s_h) is elementwise, we define qh:=∂ℓi∂sh=(Z⊤(WPVh)⊤ri)⊙ϕh′(sh).q_h:= ∂ _i∂ s_h= (Z (W^PV_h) r_i ) _h (s_h). Using the definition of shs_h, we have ∇WKQhℓi(W)=1nZqhz⊤=1nZ((Z⊤(WPVh)⊤ri)⊙ϕh′(sh))z⊤. _W^KQ_h _i(W)= 1nZq_hz = 1nZ ( (Z (W^PV_h) r_i ) _h (s_h) )z . Plugging in ϕ1′(t)=3t⊙2 _1 (t)=3t 2 and ϕ2′(t)=ϕ3′(t)=2t _2 (t)= _3 (t)=2t gives exactly the gradients after summing over i and taking the expectation. ∎ According to Lemma 3.3, we use the following specific initialization to zero out the irrelevant blocks. Assumption A.8 (Initialization). We assume the parameters are initialized as WPV1(0)=[0d×d0d×10d×d0d×101×d00d×10W31PV1(0)0d×10d×d0d×101×d00d×10],WKQ1(0)=[0d×d0d×1W13KQ1(0)W14KQ101×d00d×1W24KQ10d×d0d×10d×d0d×101×d00d×10] W^PV_1(0)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_1_31(0)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix,W^KQ_1(0)= bmatrix0_d× d&0_d× 1&W^KQ_1_13(0)&W^KQ_1_14\\ 0_1× d&0&0_d× 1&W^KQ_1_24\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix WPV2(0)=[0d×d0d×10d×d0d×101×d00d×10W31PV2(0)0d×10d×d0d×101×d00d×10],WKQ2(0)=[0d×d0d×1W13KQ2(0)W14KQ201×d00d×1W24KQ20d×d0d×10d×d0d×101×d00d×10] W^PV_2(0)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_2_31(0)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix,W^KQ_2(0)= bmatrix0_d× d&0_d× 1&W^KQ_2_13(0)&W^KQ_2_14\\ 0_1× d&0&0_d× 1&W^KQ_2_24\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix WPV3(0)=[0d×d0d×10d×d0d×101×d00d×10W31PV3(0)0d×10d×d0d×101×d00d×10],WKQ3(0)=[0d×d0d×1W13KQ3(0)W14KQ301×d00d×1W24KQ30d×d0d×10d×d0d×101×d00d×10] W^PV_3(0)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_3_31(0)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix,W^KQ_3(0)= bmatrix0_d× d&0_d× 1&W^KQ_3_13(0)&W^KQ_3_14\\ 0_1× d&0&0_d× 1&W^KQ_3_24\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix Additionally, we impose a weight-sharing mechanism that ties weights across WPViW^PV_i matrices and across WKQiW^KQ_i matrices. Let two underlying matrices be W31PVW^PV_31 and W13KQW^KQ_13. We represent the matrices in the transformer with the two matrices as follows: W31PV1(t)=−2ηn2W31PV(t),W31PV2(t)=ηn2W31PV(t),W31PV3(t)=−ηn2W31PV(t) W^PV_1_31(t)=-2η n^2W^PV_31(t), W^PV_2_31(t)= η n2W^PV_31(t), W^PV_3_31(t)=- η n2W^PV_31(t) W13KQ1(t)=W13KQ(t),W13KQ2(t)=W13KQ(t),W13KQ3(t)=W13KQ(t). W^KQ_1_13(t)=W^KQ_13(t), W^KQ_2_13(t)=W^KQ_13(t), W^KQ_3_13(t)=W^KQ_13(t). The weight matrices W31PV(t)W^PV_31(t) and W13KQ(t)W^KQ_13(t) are updated with the combined gradient of the corresponding three weights. Further, W31PV(0)=∑i=1dλiPVuiui⊤W_31^PV(0)= _i=1^d _i^PVu_iu_i and W13KQ(0)=∑i=1dλiKQuiui⊤W^KQ_13(0)= _i=1^d _i^KQu_iu_i share the same orthonormal basis, which is symmetric and simultaneously diagonalizable. We fix W24KQ1=0W^KQ_1_24=0, W24KQ2=1W^KQ_2_24=1, W24KQ3=−1W^KQ_3_24=-1, and W14KQ1=0d,W14KQ2=0d,W14KQ3=0dW^KQ_1_14=0_d,W^KQ_2_14=0_d,W^KQ_3_14=0_d for all t>0t>0. We take the initialization to be isotropic, λiPV=λV(0)λ^PV_i= _V(0) and λiKQ=λK(0)λ^KQ_i= _K(0) for all i∈[d]i∈[d], where for a constant σ∈(0,1/4)σ∈(0,1/4), λK(0)∈[13+σ, 1],λV(0)∈(0,1/σ],λV(0)2≥F⁡(λK(0)2)+σ2, _K(0)∈ [ 1 3+σ,\ 1 ], _V(0)∈(0,1/σ], _V(0)^2\ ≥\ F ( _K(0)^2 )+σ^2, where F⁡(x):=x−13+29log⁡23x−1F(x):= x-13+ 29 23x-1. We consider the gradient flow trajectory: dWdt= dWdt= −∇LCoT(W), where -∇ L^CoT(W), where W:= W:= [WKQ1/W24KQ1,W14KQ1,WKQ2/W24KQ2,W14KQ2,WKQ3/W24KQ3,W14KQ3, [W^KQ_1/\W^KQ_1_24,W^KQ_1_14\,W^KQ_2/\W^KQ_2_24,W^KQ_2_14\,W^KQ_3/\W^KQ_3_24,W^KQ_3_14\, WPV1,WPV2,WPV3]. W^PV_1,W^PV_2,W^PV_3]. Next, we prove that if we initialize those matrices as in Assumption A.8, then the zero blocks will stay at 0 along the gradient descent trajectory. Lemma A.9. Under Assumption A.8, when the transformer is trained under gradient flow, ∀t>0∀ t>0, WPV1(t)=[0d×d0d×10d×d0d×101×d00d×10W31PV1(t)0d×10d×d0d×101×d00d×10],WKQ1(t)=[0d×d0d×1W13KQ1(t)0d×101×d00d×100d×d0d×10d×d0d×101×d00d×10] W^PV_1(t)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_1_31(t)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix,W^KQ_1(t)= bmatrix0_d× d&0_d× 1&W^KQ_1_13(t)&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix WPV2(t)=[0d×d0d×10d×d0d×101×d00d×10W31PV2(t)0d×10d×d0d×101×d00d×10],WKQ2(t)=[0d×d0d×1W13KQ2(t)0d×101×d00d×110d×d0d×10d×d0d×101×d00d×10] W^PV_2(t)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_2_31(t)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix,W^KQ_2(t)= bmatrix0_d× d&0_d× 1&W^KQ_2_13(t)&0_d× 1\\ 0_1× d&0&0_d× 1&1\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix WPV3(t)=[0d×d0d×10d×d0d×101×d00d×10W31PV3(t)0d×10d×d0d×101×d00d×10],WKQ3(t)=[0d×d0d×1W13KQ3(t)0d×101×d00d×1−10d×d0d×10d×d0d×101×d00d×10] W^PV_3(t)= bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_3_31(t)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix,W^KQ_3(t)= bmatrix0_d× d&0_d× 1&W^KQ_3_13(t)&0_d× 1\\ 0_1× d&0&0_d× 1&-1\\ 0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix Proof. We show that when the irrelevant blocks are 0, the gradients for those blocks are always 0: WPV1Zi⋅ϕ1(Zi⊤WKQ1Zi[:,−1]n) W^PV_1Z_i· _1 ( Z_i W^KQ_1Z_i[:,-1]n ) = = 1n3[0d×d0d×10d×d0d×101×d00d×10W31PV1(t)0d×10d×d0d×101×d00d×10][X0⋯0y0⋯00dθ0⋯θi01⋯1](Zi⊤WKQ1[00θi1])⊙3 1n^3 bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_1_31(t)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix bmatrixX&0&·s&0\\ y&0&·s&0\\ 0_d& _0&·s& _i\\ 0&1&·s&1 bmatrix (Z_i W^KQ_1 bmatrix0\\ 0\\ _i\\ 1 bmatrix ) 3 = = 1n3[0d×d0d…0d01×d0…0W31PV1(t)X0d…0d01×d0…0]([X0⋯0y0⋯00dθ0⋯θi01⋯1]⊤[W13KQ1(t)θi00d0])⊙3 1n^3 bmatrix0_d× d&0_d&…&0_d\\ 0_1× d&0&…&0\\ W^PV_1_31(t)X&0_d&…&0_d\\ 0_1× d&0&…&0 bmatrix ( bmatrixX&0&·s&0\\ y&0&·s&0\\ 0_d& _0&·s& _i\\ 0&1&·s&1 bmatrix bmatrixW^KQ_1_13(t) _i\\ 0\\ 0_d\\ 0 bmatrix ) 3 = = 1n3[0d×d0d…0d01×d0…0W31PV1(t)X0d…0d01×d0…0][X⊤W13KQ1(t)θi0i+1]⊙3 1n^3 bmatrix0_d× d&0_d&…&0_d\\ 0_1× d&0&…&0\\ W^PV_1_31(t)X&0_d&…&0_d\\ 0_1× d&0&…&0 bmatrix bmatrixX W^KQ_1_13(t) _i\\ 0_i+1\\ bmatrix 3 = = 1n3[0d0W31PV1(t)X(X⊤W13KQ1(t)θi)⊙30]. 1n^3 bmatrix0_d\\ 0\\ W^PV_1_31(t)X (X W^KQ_1_13(t) _i ) 3\\ 0 bmatrix. For the second head, we have WPV2Zi⋅ϕ2(Zi⊤WKQ2Zi[:,−1]n) W^PV_2Z_i· _2 ( Z_i W^KQ_2Z_i[:,-1]n ) = = 1n2[0d×d0d×10d×d0d×101×d00d×10W31PV2(t)0d×10d×d0d×101×d00d×10][X0⋯0y0⋯00dθ0⋯θi01⋯1](Zi⊤WKQ2[00θi1])⊙2 1n^2 bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_2_31(t)&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix bmatrixX&0&·s&0\\ y&0&·s&0\\ 0_d& _0&·s& _i\\ 0&1&·s&1 bmatrix (Z_i W^KQ_2 bmatrix0\\ 0\\ _i\\ 1 bmatrix ) 2 = = 1n2[0d×d0d×10d×d0d×101×d00d×10W31PV2(t)X0d×10d×d0d×101×d00d×10]([X0⋯0y0⋯00dθ0⋯θi01⋯1]⊤[W13KQ2(t)θi10d0])⊙2 1n^2 bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_2_31(t)X&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix ( bmatrixX&0&·s&0\\ y&0&·s&0\\ 0_d& _0&·s& _i\\ 0&1&·s&1 bmatrix bmatrixW^KQ_2_13(t) _i\\ 1\\ 0_d\\ 0 bmatrix ) 2 = = 1n2[0d×d0d×10d×d0d×101×d00d×10W31PV2(t)X0d×10d×d0d×101×d00d×10][X⊤W13KQ2(t)θi+y⊤0i+1]⊙2 1n^2 bmatrix0_d× d&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0\\ W^PV_2_31(t)X&0_d× 1&0_d× d&0_d× 1\\ 0_1× d&0&0_d× 1&0 bmatrix bmatrixX W^KQ_2_13(t) _i+y \\ 0_i+1\\ bmatrix 2 = = 1n2[0d0W31PV2(t)X(X⊤W13KQ2(t)θi+y⊤)⊙20]. 1n^2 bmatrix0_d\\ 0\\ W^PV_2_31(t)X (X W^KQ_2_13(t) _i+y ) 2\\ 0 bmatrix. Similarly, for the third head, we have WPV3Zi⋅ϕ3(Zi⊤WKQ3Zi[:,−1]n)=1n2[0d0W31PV3(t)X(X⊤W13KQ3(t)θi−y⊤)⊙20]. W^PV_3Z_i· _3 ( Z_i W^KQ_3Z_i[:,-1]n )= 1n^2 bmatrix0_d\\ 0\\ W^PV_3_31(t)X (X W^KQ_3_13(t) _i-y ) 2\\ 0 bmatrix. Recall that we have the residual for a CoT step defined as ri=∑h=13WPVhZi⋅ϕh(Zi⊤WKQhZi[:,−1]n)−(0d,0,θi+1−θi,0)⊤=(0d,0,r^i−(θi+1−θi),0)⊤, r_i= _h=1^3W^PV_hZ_i· _h ( Z_i W^KQ_hZ_i[:,-1]n )-(0_d,0, _i+1- _i,0) =(0_d,0, r_i-( _i+1- _i),0) , where we define r^i r_i =1n3W31PV1X(X⊤W13KQ1θi)⊙3+1n2W31PV2X(X⊤W13KQ2θi+y⊤)⊙2 = 1n^3W^PV_1_31X (X W^KQ_1_13 _i ) 3+ 1n^2W^PV_2_31X (X W^KQ_2_13 _i+y ) 2 +1n2W31PV3X(X⊤W13KQ3θi−y⊤)⊙2 + 1n^2W^PV_3_31X (X W^KQ_3_13 _i-y ) 2 =1n3W31PV1X(X⊤W13KQ1θi)⊙3+4n2W31PV2X(X⊤W13KQ2θi⊙y⊤). = 1n^3W^PV_1_31X (X W^KQ_1_13 _i ) 3+ 4n^2W^PV_2_31X (X W^KQ_2_13 _i y ). (31) The second step is because of the symmetry in Assumption A.8 that W31PV2=−W31PV3W^PV_2_31=-W^PV_3_31 and W13KQ2=W13KQ3W^KQ_2_13=W^KQ_3_13. Now, we consider the gradient for WPVhW^PV_h. Denote the loss at step i as Li(W)L_i(W). For an X and θ⋆θ , ∇WPV1Li(W) _W^PV_1L_i(W) =ri(Ziϕ1(Zi⊤WKQ1Zi[:,−1]n))⊤=1n3ri(Zi[X⊤W13KQ1(t)θi0i+1]⊙3)⊤ =r_i (Z_i _1( Z_i W^KQ_1Z_i[:,-1]n) ) = 1n^3r_i (Z_i bmatrixX W^KQ_1_13(t) _i\\ 0_i+1 bmatrix 3 ) =1n3[0d0r^i−(θi+1−θi)0][X(X⊤W13KQ1(t)θi)⊙3y(X⊤W13KQ1(t)θi)⊙30d0]⊤ = 1n^3 bmatrix0_d\\ 0\\ r_i-( _i+1- _i)\\ 0 bmatrix bmatrixX (X W^KQ_1_13(t) _i ) 3\\ y (X W^KQ_1_13(t) _i ) 3\\ 0_d\\ 0 bmatrix =1n3[0d×d0d0d×d0d01×d001×d0∇W31PV1Li∇W32PV1Li0d×d0d01×d001×d0] = 1n^3 bmatrix0_d× d&0_d&0_d× d&0_d\\ 0_1× d&0&0_1× d&0\\ _W^PV_1_31L_i& _W^PV_1_32L_i&0_d× d&0_d\\ 0_1× d&0&0_1× d&0\\ bmatrix where ∇W31PV1Li=ri^(X(X⊤W13KQ1(t)θi)⊙3)⊤−(θi+1−θi)(X(X⊤W13KQ1(t)θi)⊙3)⊤ _W^PV_1_31L_i= r_i (X (X W^KQ_1_13(t) _i ) 3 ) -( _i+1- _i) (X (X W^KQ_1_13(t) _i ) 3 ) ∇W32PV1Li=ri^(y(X⊤W13KQ1(t)θi)⊙3)⊤−(θi+1−θi)(y(X⊤W13KQ1(t)θi)⊙3)⊤. _W^PV_1_32L_i= r_i (y (X W^KQ_1_13(t) _i ) 3 ) -( _i+1- _i) (y (X W^KQ_1_13(t) _i ) 3 ) . Similarly, we calculate the gradient of WPV2W^PV_2. We have ∇WPV2Li(W) _W^PV_2L_i(W) =ri(Ziϕ2(Zi⊤WKQ2Zi[:,−1]n))⊤=1n2ri(Zi[X⊤W13KQ2(t)θi+y⊤0i+1]⊙2)⊤ =r_i (Z_i _2( Z_i W^KQ_2Z_i[:,-1]n) ) = 1n^2r_i (Z_i bmatrixX W^KQ_2_13(t) _i+y \\ 0_i+1 bmatrix 2 ) =1n2[0d0r^i−(θi+1−θi)0][X(X⊤W13KQ2(t)θi+y⊤)⊙2y(X⊤W13KQ2(t)θi+y⊤)⊙20d0]⊤ = 1n^2 bmatrix0_d\\ 0\\ r_i-( _i+1- _i)\\ 0 bmatrix bmatrixX (X W^KQ_2_13(t) _i+y ) 2\\ y (X W^KQ_2_13(t) _i+y ) 2\\ 0_d\\ 0 bmatrix =1n2[0d×d0d0d×d0d01×d001×d0∇W31PV2Li∇W32PV2Li0d×d0d01×d001×d0], where = 1n^2 bmatrix0_d× d&0_d&0_d× d&0_d\\ 0_1× d&0&0_1× d&0\\ _W^PV_2_31L_i& _W^PV_2_32L_i&0_d× d&0_d\\ 0_1× d&0&0_1× d&0\\ bmatrix, where ∇W31PV2Li=r^i(X(X⊤W13KQ2(t)θi+y⊤)⊙2)⊤−(θi+1−θi)(X(X⊤W13KQ2(t)θi+y⊤)⊙2)⊤ _W^PV_2_31L_i= r_i (X (X W^KQ_2_13(t) _i+y ) 2 ) -( _i+1- _i) (X (X W^KQ_2_13(t) _i+y ) 2 ) ∇W32PV2Li=r^i(y(X⊤W13KQ2(t)θi+y⊤)⊙2)⊤−(θi+1−θi)(y(X⊤W13KQ2(t)θi+y⊤)⊙2)⊤. _W^PV_2_32L_i= r_i (y (X W^KQ_2_13(t) _i+y ) 2 ) -( _i+1- _i) (y (X W^KQ_2_13(t) _i+y ) 2 ) . Also, the gradient for WPV3W^PV_3 is as follows: ∇WPV3Li(W) _W^PV_3L_i(W) =ri(Ziϕ3(Zi⊤WKQ3Zi[:,−1]n))⊤=1n2ri(Zi[X⊤W13KQ3(t)θi−y⊤0i+1]⊙2)⊤ =r_i (Z_i _3( Z_i W^KQ_3Z_i[:,-1]n) ) = 1n^2r_i (Z_i bmatrixX W^KQ_3_13(t) _i-y \\ 0_i+1 bmatrix 2 ) =1n2[0d0r^i−(θi+1−θi)0][X(X⊤W13KQ3(t)θi−y⊤)⊙2y(X⊤W13KQ3(t)θi−y⊤)⊙20d0]⊤ = 1n^2 bmatrix0_d\\ 0\\ r_i-( _i+1- _i)\\ 0 bmatrix bmatrixX (X W^KQ_3_13(t) _i-y ) 2\\ y (X W^KQ_3_13(t) _i-y ) 2\\ 0_d\\ 0 bmatrix =1n2[0d×d0d0d×d0d01×d001×d0∇W31PV3Li∇W32PV3Li0d×d0d01×d001×d0], where = 1n^2 bmatrix0_d× d&0_d&0_d× d&0_d\\ 0_1× d&0&0_1× d&0\\ _W^PV_3_31L_i& _W^PV_3_32L_i&0_d× d&0_d\\ 0_1× d&0&0_1× d&0\\ bmatrix, where ∇W31PV3Li _W^PV_3_31L_i =r^i(X(X⊤W13KQ3(t)θi−y⊤)⊙2)⊤−(θi+1−θi)(X(X⊤W13KQ3(t)θi−y⊤)⊙2)⊤ = r_i (X (X W^KQ_3_13(t) _i-y ) 2 ) -( _i+1- _i) (X (X W^KQ_3_13(t) _i-y ) 2 ) ∇W32PV3Li _W^PV_3_32L_i =r^i(y(X⊤W13KQ3(t)θi−y⊤)⊙2)⊤−(θi+1−θi)(y(X⊤W13KQ3(t)θi−y⊤)⊙2)⊤. = r_i (y (X W^KQ_3_13(t) _i-y ) 2 ) -( _i+1- _i) (y (X W^KQ_3_13(t) _i-y ) 2 ) . Next, we calculate the gradient for WKQhW^KQ_h. We write riθ:=r^i−(θ¯i+1−θi)r^θ_i:= r_i-( θ_i+1- _i) for the θ-block of rir_i. Thus, ri=(0d,0,riθ,0)r_i=(0_d,0,r^θ_i,0), and Zi⊤WPVh⊤ri=[X⊤W31PVh⊤riθ; 0i+1]Z_i W^PV_h r_i=[X W^PV_h _31r^θ_i;\,0_i+1]. Then, we have ∇WKQ1Li _W^KQ_1L_i =1nZi((Zi⊤WPV1⊤ri)⊙3(Zi⊤WKQ1Zi[:,−1]n)⊙2)Zi[:,−1]⊤ = 1nZ_i ( (Z_i W^PV_1 r_i ) 3( Z_i W^KQ_1Z_i[:,-1]n) 2 )Z_i[:,-1] =1nZi([X⊤W31PV1⊤(t)riθ0i+1]⊙3n2[(X⊤W13KQ1(t)θi)⊙20i+1])Zi[:,−1]⊤ = 1nZ_i ( bmatrixX W^PV_1 _31(t)r^θ_i\\ 0_i+1 bmatrix 3n^2 bmatrix (X W^KQ_1_13(t) _i ) 2\\ 0_i+1 bmatrix )Z_i[:,-1] =3n3Zi[(X⊤W31PV1⊤(t)riθ⊙(X⊤W13KQ1(t)θi)⊙2)0i+1]Zi[:,−1]⊤ = 3n^3Z_i bmatrix (X W^PV_1 _31(t)r^θ_i (X W^KQ_1_13(t) _i ) 2 )\\ 0_i+1 bmatrixZ_i[:,-1] =3n3[X⁡(X⊤W31PV1⊤(t)riθ⊙(X⊤W13KQ1(t)θi)⊙2)y⁡(X⊤W31PV1⊤(t)riθ⊙(X⊤W13KQ1(t)θi)⊙2)0d0][0d0θi⊤1] = 3n^3 bmatrixX (X W^PV_1 _31(t)r^θ_i (X W^KQ_1_13(t) _i ) 2 )\\ y (X W^PV_1 _31(t)r^θ_i (X W^KQ_1_13(t) _i ) 2 )\\ 0_d\\ 0 bmatrix bmatrix0_d&0& _i &1 bmatrix =3n3[0d×d0d∇W13KQ1Li∇W14KQ1Li01×d0∇W23KQ1Li∇W24KQ1Li0d×d0d0d×d001×d001×d0], where = 3n^3 bmatrix0_d× d&0_d& _W^KQ_1_13L_i& _W^KQ_1_14L_i\\ 0_1× d&0& _W^KQ_1_23L_i& _W^KQ_1_24L_i\\ 0_d× d&0_d&0_d× d&0\\ 0_1× d&0&0_1× d&0\\ bmatrix, where ∇W13KQ1Li _W^KQ_1_13L_i =X⁡(X⊤W31PV1⊤(t)riθ⊙(X⊤W13KQ1(t)θi)⊙2)θi⊤ =X (X W^PV_1 _31(t)r^θ_i (X W^KQ_1_13(t) _i ) 2 ) _i ∇W23KQ1Li _W^KQ_1_23L_i =y⁡(X⊤W31PV1⊤(t)riθ⊙(X⊤W13KQ1(t)θi)⊙2)θi⊤, =y (X W^PV_1 _31(t)r^θ_i (X W^KQ_1_13(t) _i ) 2 ) _i , and we fix W24KQ1=0W^KQ_1_24=0 and W14KQ1=0dW^KQ_1_14=0_d for t>0t>0. Then, for head two, we have ∇WKQ2Li _W^KQ_2L_i =1nZi((Zi⊤WPV2⊤ri)⊙2(Zi⊤WKQ2Zi[:,−1]n))Zi[:,−1]⊤ = 1nZ_i ( (Z_i W^PV_2 r_i ) 2( Z_i W^KQ_2Z_i[:,-1]n) )Z_i[:,-1] =1nZi([X⊤W31PV2⊤(t)riθ0i+1]⊙2n[(X⊤W13KQ2(t)θi+y⊤)0i+1])Zi[:,−1]⊤ = 1nZ_i ( bmatrixX W^PV_2 _31(t)r^θ_i\\ 0_i+1 bmatrix 2n bmatrix (X W^KQ_2_13(t) _i+y )\\ 0_i+1 bmatrix )Z_i[:,-1] =2n2Zi[(X⊤W31PV2⊤(t)riθ⊙(X⊤W13KQ2(t)θi+y⊤))0i+1]Zi[:,−1]⊤ = 2n^2Z_i bmatrix (X W^PV_2 _31(t)r^θ_i (X W^KQ_2_13(t) _i+y ) )\\ 0_i+1 bmatrixZ_i[:,-1] =2n2[X⁡(X⊤W31PV2⊤(t)riθ⊙(X⊤W13KQ2(t)θi+y⊤))y⁡(X⊤W31PV2⊤(t)riθ⊙(X⊤W13KQ2(t)θi+y⊤))0d0][0d0θi⊤1] = 2n^2 bmatrixX (X W^PV_2 _31(t)r^θ_i (X W^KQ_2_13(t) _i+y ) )\\ y (X W^PV_2 _31(t)r^θ_i (X W^KQ_2_13(t) _i+y ) )\\ 0_d\\ 0 bmatrix bmatrix0_d&0& _i &1 bmatrix =2n2[0d×d0d∇W13KQ2Li∇W14KQ2Li01×d0∇W23KQ2Li∇W24KQ2Li0d×d0d0d×d001×d001×d0], where = 2n^2 bmatrix0_d× d&0_d& _W^KQ_2_13L_i& _W^KQ_2_14L_i\\ 0_1× d&0& _W^KQ_2_23L_i& _W^KQ_2_24L_i\\ 0_d× d&0_d&0_d× d&0\\ 0_1× d&0&0_1× d&0\\ bmatrix, where ∇W13KQ2Li _W^KQ_2_13L_i =X⁡(X⊤W31PV2⊤(t)riθ⊙(X⊤W13KQ2(t)θi+y⊤))θi⊤ =X (X W^PV_2 _31(t)r^θ_i (X W^KQ_2_13(t) _i+y ) ) _i ∇W23KQ2Li _W^KQ_2_23L_i =y⁡(X⊤W31PV2⊤(t)riθ⊙(X⊤W13KQ2(t)θi+y⊤))θi⊤, =y (X W^PV_2 _31(t)r^θ_i (X W^KQ_2_13(t) _i+y ) ) _i , and we fix W24KQ2=1W^KQ_2_24=1 and W14KQ2=0dW^KQ_2_14=0_d for t>0t>0. Similarly, we have the gradient for WKQ3W^KQ_3 as follows. ∇W13KQ3Li=X⁡(X⊤W31PV3⊤(t)riθ⊙(X⊤W13KQ3(t)θi−y⊤))θi⊤ _W^KQ_3_13L_i=X (X W^PV_3 _31(t)r^θ_i (X W^KQ_3_13(t) _i-y ) ) _i ∇W23KQ3Li=y⁡(X⊤W31PV3⊤(t)riθ⊙(X⊤W13KQ3(t)θi−y⊤))θi⊤, _W^KQ_3_23L_i=y (X W^PV_3 _31(t)r^θ_i (X W^KQ_3_13(t) _i-y ) ) _i , and we fix W24KQ3=−1W^KQ_3_24=-1 and W14KQ3=0dW^KQ_3_14=0_d for t>0t>0. Now, we are going to show that the expected gradients with respect to the irrelevant blocks are zero, i.e., ∇W32PV1Li _W^PV_1_32L_i, ∇W32PV2Li _W^PV_2_32L_i, ∇W32PV3Li _W^PV_3_32L_i, ∇W23KQ1Li _W^KQ_1_23L_i, ∇W23KQ2Li _W^KQ_2_23L_i, ∇W23KQ3Li _W^KQ_3_23L_i. We will show that X,θ⋆[∇W32PV1Li]=0 E_X,θ [ _W^PV_1_32L_i ]=0, and proofs for others follow similarly. We use the sign-flip symmetry of the Gaussian distribution to prove this fact. Denote gi:=y⊤ai⊙3g_i:=y a_i 3 and ai:=X⊤W13KQ1θia_i:=X W^KQ_1_13 _i. Define a sign-flip transformation :(X,θ⋆)↦(−X,θ⋆)T:(X,θ ) (-X,θ ). Since X has i.i.d. standard Gaussian columns, the distribution of (X,θ⋆)(X,θ ) is invariant under T. The parameter token θi _i and the target increment θi+1−θi _i+1- _i used in ZiZ_i are generated by the deterministic gradient descent on the quadratic squared loss. Thus, θi _i is invariant under T. Then, based on the definitions, we have y is invariant under T, and gig_i will be flipped to −gi-g_i under T. Moreover, we can show r^i r_i is even under T. Under the block expressions, we have riθ r^θ_i =r^i−(θi+1−θi) = r_i-( _i+1- _i) =1n3W31PV1X(X⊤W13KQ1θi)⊙3+4n2W31PV2X(X⊤W13KQ2θi⊙y⊤)−(θi+1−θi) = 1n^3W^PV_1_31X (X W^KQ_1_13 _i ) 3+ 4n^2W^PV_2_31X (X W^KQ_2_13 _i y )-( _i+1- _i) One can verify that riθr^θ_i is invariant under the transformation T. Thus, ∇W32PV1Li(W)=rigi _W^PV_1_32L_i(W)=r_ig_i is odd under T. Combining the above sign-flipping results and using the distributional invariance, we have X,θ⋆[∇W32PV1Li(W)]=⁡(X,θ⋆)[∇W32PV1Li(W)∘]=−X,θ⋆[∇W32PV1Li(W)], E_X,θ [ _W^PV_1_32L_i(W) ]= E_T(X,θ ) [ _W^PV_1_32L_i(W) ]=- E_X,θ [ _W^PV_1_32L_i(W) ], Therefore, X,θ⋆[∇W32PV1Li(W)]=0 E_X,θ [ _W^PV_1_32L_i(W) ]=0. ∎ Notice that the final loss only depends on the (d+2)(d+2) to (2d+1)(2d+1) entries of the transformer’s output. We can simplify the model by pruning the irrelevant and zero blocks. We can define a simplified transformer, where the dynamics of the equivalent model are exactly the same as the original model. Definition A.10. Under Assumption A.8, by Lemma A.9, for a given input ZiZ_i, we can rewrite the one-layer non-linear transformer by fW(Zi)=[0d,0,(θi−2ηnW31PVX(X⊤W13KQθi)⊙3+2ηnW31PVX(X⊤W13KQθi⊙y⊤)),1]⊤.f_W(Z_i)= [0_d,0, ( _i- 2ηnW^PV_31X (X W^KQ_13 _i ) 3+ 2ηnW^PV_31X (X W^KQ_13 _i y ) ),1 ] . where θi _i is the last column of Zi[d+2:2d+1,−1]Z_i_[d+2:2d+1,-1] and y⊤=(X⊤θ⋆)⊙2y =(X θ ) 2 where θ⋆θ is the ground-truth weights of the function. Then, we can denote a reduced model by: f~W(Zi)=θi−2ηnW31PVX(X⊤W13KQθi)⊙3+2ηnW31PVX(X⊤W13KQθi⊙y⊤). f_W(Z_i)= _i- 2ηnW^PV_31X (X W^KQ_13 _i ) 3+ 2ηnW^PV_31X (X W^KQ_13 _i y ). (32) For the ease of notations, we will write f~W(Zi)=f~(θi) f_W(Z_i)= f( _i), where θi _i denotes the (d+2)(d+2) to (2d+1)(2d+1) entries of the last column of ZiZ_i. In the following sections, we will consider the equivalent form of the transformer. Here we present the gradient with respect to the reduced model. Lemma A.11. The gradients of W31PVW^PV_31 and W13KQW^KQ_13 are given by the following expectations: ∇W31PVLCoT(W) _W^PV_31L^CoT(W) = = ∑i=0T−1X,θ⋆[−2ηn(f~W(θi)−θi+1)((X(X⊤W13KQθi)⊙3)⊤−(X⁡(X⊤W13KQθi⊙y⊤))⊤)] _i=0^T-1 E_X,θ [- 2ηn( f_W( _i)- _i+1) ( (X(X W^KQ_13 _i) 3 ) - (X (X W^KQ_13 _i y ) ) ) ] +X,θ⋆[−2ηn(θT−θ⋆)((X(X⊤W13KQθT−1)⊙3)⊤−(X(X⊤W13KQθT−1⊙y⊤))⊤)], and + E_X,θ [- 2ηn( _T-θ ) ( (X(X W^KQ_13 _T-1) 3 ) - (X (X W^KQ_13 _T-1 y ) ) ) ], and ∇W13KQLCoT(W) _W^KQ_13L^CoT(W) = = ∑i=0T−1X,θ⋆[−2ηnX(X⊤W31PV⊤(f~W(θi)−θi+1)⊙(3(X⊤W13KQθi)⊙2−y⊤))θi⊤] _i=0^T-1 E_X,θ [- 2ηnX (X W^PV _31( f_W( _i)- _i+1) (3 (X W^KQ_13 _i ) 2-y ) ) _i ] +X,θ⋆[−2ηnX(X⊤W31PV⊤(θT−θ⋆)⊙(3(X⊤W13KQθT−1)⊙2−y⊤))θT−1⊤]. + E_X,θ [- 2ηnX (X W^PV _31( _T-θ ) (3 (X W^KQ_13 _T-1 ) 2-y ) ) _T-1 ]. Proof. Based on the simplified form of the transformer in Definition A.10, we can rewrite the CoT training loss as follows. LCoT(W) L^CoT(W) =12X,θ⋆[∑i=0T−2‖fW(Zi)−(0d,0,θi+1,1)‖22]+12X,θ⋆[‖fW(ZT−1)−(0d,0,θ⋆,1)‖22] = 12 E_X,θ [ _i=0^T-2\|f_W(Z_i)-(0_d,0, _i+1,1)\|_2^2 ]+ 12 E_X,θ [\|f_W(Z_T-1)-(0_d,0,θ ,1)\|_2^2 ] =12X,θ⋆[∑i=0T−2‖f~W(θi)−θi+1‖22]+12X,θ⋆[‖f~W(θT−1)−θ⋆‖22]. = 12 E_X,θ [ _i=0^T-2\| f_W( _i)- _i+1\|_2^2 ]+ 12 E_X,θ [\| f_W( _T-1)-θ \|_2^2 ]. Here, we split the loss between the last step output and the ground-truth weights θ⋆θ . Then, based on the gradient computation in Lemma A.9, we can write out the gradient for the weight matrix WPVW^PV in the transformer. Recall that in Assumption A.8, we use a weight-sharing mechanism that defines W31PV1(t)=−2ηn2W31PV(t),W31PV2(t)=ηn2W31PV(t),W31PV3(t)=−ηn2W31PV(t)W^PV_1_31(t)=-2η n^2W^PV_31(t),W^PV_2_31(t)= η n2W^PV_31(t),W^PV_3_31(t)=- η n2W^PV_31(t). Therefore, ∇W31PVLCoT(W) _W^PV_31L^CoT(W) = = −2ηn2∇W31PV1LCoT(W)+ηn2∇W31PV2LCoT(W)−ηn2∇W31PV3LCoT(W) -2η n^2 _W^PV_1_31L^CoT(W)+ η n2 _W^PV_2_31L^CoT(W)- η n2 _W^PV_3_31L^CoT(W) = = ∑i=0T−2X,θ⋆[−2ηn(f~W(θi)−θi+1)(X(X⊤W13KQθi)⊙3)⊤] _i=0^T-2 E_X,θ [- 2ηn( f_W( _i)- _i+1) (X(X W^KQ_13 _i) 3 ) ] +X,θ⋆[−2ηn(f~W(θT−1)−θ⋆)(X(X⊤W13KQθT−1)⊙3)⊤] + E_X,θ [- 2ηn( f_W( _T-1)-θ ) (X(X W^KQ_13 _T-1) 3 ) ] −∑i=0T−2X,θ⋆[−2ηn(f~W(θi)−θi+1)(X(X⊤W13KQθi⊙y⊤))⊤] - _i=0^T-2 E_X,θ [- 2ηn( f_W( _i)- _i+1) (X (X W^KQ_13 _i y ) ) ] −X,θ⋆[−2ηn(f~W(θT−1)−θ⋆)(X(X⊤W13KQθT−1⊙y⊤))⊤] - E_X,θ [- 2ηn( f_W( _T-1)-θ ) (X (X W^KQ_13 _T-1 y ) ) ] = = ∑i=0T−1X,θ⋆[−2ηn(f~W(θi)−θi+1)((X(X⊤W13KQθi)⊙3)⊤−(X⁡(X⊤W13KQθi⊙y⊤))⊤)] _i=0^T-1 E_X,θ [- 2ηn( f_W( _i)- _i+1) ( (X(X W^KQ_13 _i) 3 ) - (X (X W^KQ_13 _i y ) ) ) ] +X,θ⋆[−2ηn(θT−θ⋆)(X(X⊤W13KQθT−1)⊙3)⊤+2ηn(θT−θ⋆)(X(X⊤W13KQθT−1⊙y⊤))⊤]. + E_X,θ [- 2ηn( _T-θ ) (X(X W^KQ_13 _T-1) 3 ) + 2ηn( _T-θ ) (X (X W^KQ_13 _T-1 y ) ) ]. The last step is because we subtract X,θ⋆[−2ηn(f~W(θT−1)−θT)(X(X⊤W13KQθT−1)⊙3])⊤ E_X,θ [- 2ηn( f_W( _T-1)- _T) (X(X W^KQ_13 _T-1) 3 ] ) from the second term and put it into the first term, then subtract X,θ⋆[−2ηn(f~W(θT−1)−θT)(X(X⊤W13KQθT−1⊙y⊤)] E_X,θ [- 2ηn( f_W( _T-1)- _T) (X (X W^KQ_13 _T-1 y ) ] from the fourth term and put it into the third term. Similarly, we can break down the gradient for W13KQW^KQ_13. Recall that in Assumption A.8, we use a weight-sharing mechanism that defines W13KQ1(t)=W13KQ(t),W13KQ2(t)=W13KQ(t),W13KQ3(t)=W13KQ(t)W^KQ_1_13(t)=W^KQ_13(t),W^KQ_2_13(t)=W^KQ_13(t),W^KQ_3_13(t)=W^KQ_13(t). Therefore, we have ∇W13KQLCoT(W) _W^KQ_13L^CoT(W) = = ∇W13KQ1LCoT(W)+∇W13KQ2LCoT(W)+∇W13KQ3LCoT(W) _W^KQ_1_13L^CoT(W)+ _W^KQ_2_13L^CoT(W)+ _W^KQ_3_13L^CoT(W) = = ∑i=0T−1X,θ⋆[−6ηnX(X⊤W31PV⊤(f~W(θi)−θi+1)⊙(X⊤W13KQθi)⊙2)θi⊤]+ _i=0^T-1 E_X,θ [- 6ηnX (X W^PV _31( f_W( _i)- _i+1) (X W^KQ_13 _i ) 2 ) _i ]+ X,θ⋆[−6ηnX(X⊤W31PV⊤(θT−θ⋆)⊙(X⊤W13KQθT−1)⊙2)θT−1⊤]− E_X,θ [- 6ηnX (X W^PV _31( _T-θ ) (X W^KQ_13 _T-1 ) 2 ) _T-1 ]- ∑i=0T−1X,θ⋆[−2ηnX(X⊤W31PV⊤(f~W(θi)−θi+1)⊙y⊤)θi⊤]− _i=0^T-1 E_X,θ [- 2ηnX (X W^PV _31( f_W( _i)- _i+1) y ) _i ]- X,θ⋆[−2ηnX(X⊤W31PV⊤(θT−θ⋆)⊙y⊤)θT−1⊤] E_X,θ [- 2ηnX (X W^PV _31( _T-θ ) y ) _T-1 ] = = ∑i=0T−1X,θ⋆[−2ηnX(X⊤W31PV⊤(f~W(θi)−θi+1)⊙(3(X⊤W13KQθi)⊙2−y⊤))θi⊤]+ _i=0^T-1 E_X,θ [- 2ηnX (X W^PV _31( f_W( _i)- _i+1) (3 (X W^KQ_13 _i ) 2-y ) ) _i ]+ X,θ⋆[−2ηnX(X⊤W31PV⊤(θT−θ⋆)⊙(3(X⊤W13KQθT−1)⊙2−y⊤))θT−1⊤]. E_X,θ [- 2ηnX (X W^PV _31( _T-θ ) (3 (X W^KQ_13 _T-1 ) 2-y ) ) _T-1 ]. ∎ A.3.3 Proof of Theorem 3.4 With the above computed gradients, we now prove the main theorem on the convergence of transformers trained on quadratic functions. We characterize the training dynamics using ordinary differential equations in terms of the eigenvalues of the transformer weights. We first characterize the gradient of the CoT objective as a main term plus interaction terms, which yields an ordinary differential equation for the eigenvalues of the two trained blocks. Then, we analyze the ordinary differential equation in stages. Proof. We first use the gradient computation in Lemma A.11 and derive ordinary differential equations for the eigenvalues of the weight matrices W31PVW^PV_31 and W13KQW^KQ_13. By Assumption A.8, we have weight matrices in the form for all t≥0t≥ 0: W31PV(t)=λV(t)Idd,W13KQ(t)=λK(t)Idd,W^PV_31(t)= _V(t) _d, W^KQ_13(t)= _K(t) _d, where the dynamics reduce to the two eigenvalues λV(t) _V(t) and λK(t) _K(t). Consider that the gradient descent initialization satisfies θ0=α0θ⋆ _0= _0θ with a constant α0∈[3/4,1) _0∈[3/4,1), and take n=Θ⁡(d2log5d)n= (d^2 ^5d). We will show that the gradients in the eigenbasis admit the following decomposition ∇W31PVLCoT(W)=(gV(λV,λK)+ΔV)Idd, _W^PV_31L^CoT(W)= (g_V( _V, _K)+ ^V ) _d, ∇W13KQLCoT(W)=(gK(λV,λK)+ΔK)Idd, _W^KQ_13L^CoT(W)= (g_K( _V, _K)+ ^K ) _d, where gV,gKg_V,g_K are explicit main terms, and ΔV,ΔK ^V, ^K are interaction terms. The induced eigenvalue ODEs can be represented by λ˙V(t)=−gV(λV(t),λK(t))−ΔV(t),λ˙K(t)=−gK(λV(t),λK(t))−ΔK(t), λ_V(t)=-g_V ( _V(t), _K(t) )- ^V(t), λ_K(t)=-g_K ( _V(t), _K(t) )- ^K(t), In particular, by expanding the main terms around the construction λV=λK=1 _V= _K=1, we will show that gV(λV,λK) g_V( _V, _K) =36κ2∑i=0T−1βi(1−βi)[(1−βi)(λV−1)+(1−3βi)(λK−1)] =36κ^2 _i=0^T-1 _i(1- _i) [(1- _i)( _V-1)+(1-3 _i)( _K-1) ] +O⁡(|λV−1|2+|λK−1|2), +O (| _V-1|^2+| _K-1|^2 ), gK(λV,λK) g_K( _V, _K) =36κ2∑i=0T−1βi(1−3βi)[(1−βi)(λV−1)+(1−3βi)(λK−1)] =36κ^2 _i=0^T-1 _i(1-3 _i) [(1- _i)( _V-1)+(1-3 _i)( _K-1) ] +O⁡(|λV−1|2+|λK−1|2), +O (| _V-1|^2+| _K-1|^2 ), where the coefficients depend only on (T,κ,α0)(T,κ, _0), and the weights βi∈[α02,1) _i∈[ _0^2,1). Now, we derive the explicit forms of the ODE. Based on Lemma A.9, if the parameters are initialized as in Assumption A.8, then under gradient flow all zero blocks remain zero for all t>0t>0. Hence, the model dynamics are fully characterized by the two matrices W31PV(t)W^PV_31(t) and W13KQ(t)W^KQ_13(t). Then, we write the reduced form of the transformer as in Equation 32. The chain-of-thought training loss can be written as LCoT(W)=12X,θ⋆[∑i=0T−2‖f~W(θi)−θi+1‖22+‖f~W(θT−1)−θ⋆‖22].L^CoT(W)= 12E_X,θ [ _i=0^T-2\| f_W( _i)- _i+1\|_2^2+\| f_W( _T-1)-θ \|_2^2 ]. with the residuals defined as ei:=f~W(θi)−θi+1,i=0,…,T−2,f~W(θT−1)−θ⋆,i=T−1.e_i:= cases f_W( _i)- _i+1,&i=0,…,T-2,\\ f_W( _T-1)-θ ,&i=T-1. cases Let Ψ⁡(θ):=X⁡((X⊤W13KQθ)⊙3−(X⊤W13KQθ)⊙y⊤). (θ):=X ((X W^KQ_13θ) 3-(X W^KQ_13θ) y ). Then the gradients of the reduced parameters are ∇W31PVLCoT(W) _W^PV_31L^CoT(W) =[∑i=0T−1−2ηneiΨ(θi)⊤], =E [ _i=0^T-1- 2ηne_i ( _i) ], (33) ∇W13KQLCoT(W) _W^KQ_13L^CoT(W) =[∑i=0T−1−2ηnX((X⊤W31PV⊤ei)⊙(3(X⊤W13KQθi)⊙2−y⊤))θi⊤]. =E [ _i=0^T-1- 2ηnX ((X W^PV _31e_i) (3(X W^KQ_13 _i) 2-y ) ) _i ]. (34) As in Assumption A.8, assume W31PV(t)W^PV_31(t) and W13KQ(t)W^KQ_13(t) are symmetric and share an orthonormal eigenbasis U. Rotating all variables by U⊤U leaves the data distribution invariant, so we can work in the basis where U=IddU= _d and both matrices are diagonal. Thus, the isotropy is preserved. The dynamics reduce to ODEs for the eigenvalues λV(t) _V(t) and λK(t) _K(t). Then, we further write the ODE as the expectation over θ⋆θ and X. Let θ⋆∼⁡(0,Idd)θ (0, _d) and initialize θ0=α0θ⋆ _0= _0θ with α0∈[3/4,1) _0∈[3/4,1). Then, gradient descent on the quadratic functions preserves collinearity, so θi=αiθ⋆,αi+1=αi−6η‖θ⋆‖22αi(αi2−1) _i= _iθ ,\ _i+1= _i-6η\|θ \|_2^2 _i( _i^2-1). Here, αi _i are scalars that only depend on the data ‖θ⋆‖2\|θ \|^2. Let α¯i α_i follow the same recursion with ‖θ⋆‖2\|θ \|^2 replaced by its mean d, and let βi:=α¯i2 _i:= α_i^2. The empirical iterates θi _i stay close to αiθ⋆ _iθ , and ‖θ⋆‖2\|θ \|^2 is close to d. Given W31PV=λVIddW^PV_31= _V _d and W13KQ=λKIddW^KQ_13= _K _d, let s:=1−λVλK3,τ:=1−λVλK,s:=1- _V _K^3, τ:=1- _V _K, which both vanish at the construction λV=λK=1 _V= _K=1. Then, we have Ψ⁡(θ)=λK3X(X⊤θ)⊙3−λKX((X⊤θ)⊙y⊤), (θ)= _K^3X(X θ) 3- _KX ((X θ) y ), for i≤T−2i≤ T-2, ei=2ηn(sX(X⊤θi)⊙3−τX((X⊤θi)⊙y⊤)),e_i= 2ηn (s\,X(X _i) 3-τ\,X ((X _i) y ) ), and eT−1e_T-1 is the same expression plus θT−θ⋆ _T-θ . As the other terms do not depend on (λV,λK)( _V, _K), LCoTL^CoT is exactly a quadratic function of (s,τ)(s,τ). Then, we apply the expectation to the Gaussian moment. for x∼⁡(0,Idd)x (0, _d) and fixed vectors v,w∈ℝdv,w ^d, we have ⁡[x(x⊤v)3]=3‖v‖2v,⁡[x⁡(x⊤v)(x⊤w)2]=‖w‖2v+2(v⊤w)w.E[x(x v)^3]=3\|v\|^2v, [x(x v)(x w)^2]=\|w\|^2v+2(v w)w. Applying these identities yields the population expectation X[Ψ(θ)∣θ,θ⋆]=n(3λK3∥θ∥2θ−λK(∥θ⋆∥2θ+2(θ⊤θ⋆)θ⋆)). _X[ (θ) θ,θ ]=n (3 _K^3\|θ\|^2θ- _K (\|θ \|^2θ+2(θ θ )θ ) ). (35) In particular, at θ=αθ⋆θ=αθ the expectations of 1nX(X⊤θ)⊙3 1nX(X θ) 3 and 1nX((X⊤θ)⊙y⊤) 1nX ((X θ) y ) are 3α3‖θ⋆‖2θ⋆3α^3\|θ \|^2θ and 3α‖θ⋆‖2θ⋆3α\|θ \|^2θ . We keep only the population parts and replace ‖θ⋆‖2\|θ \|^2 by d and αi _i by α¯i α_i, which defines the main terms. The remaining terms are the interaction terms. Then, taking the exact expectation over θ⋆θ and X, the population gradient-flow dynamics of the eigenvalues are given by λ˙V λ_V =36κ2λK∑i=0T−1βi(βis−τ)(λK2βi−1)−ΔV, =36κ^2 _K _i=0^T-1 _i( _is-τ)( _K^2 _i-1)- ^V, (36) λ˙K λ_K =36κ2λV∑i=0T−1βi(βis−τ)(3λK2βi−1)−ΔK. =36κ^2 _V _i=0^T-1 _i( _is-τ)(3 _K^2 _i-1)- ^K. (37) The interaction terms ΔV ^V and ΔK ^K are the differences between the exact expressions and the main terms in equation (36) and equation (37). On any bounded set of (λV,λK)( _V, _K), their orders are as follows |ΔV|+|ΔK|=O⁡(1log⁡d)(∑i=0T−1βi|βis−τ|+|s|+|τ|)+O⁡(d−aC).| ^V|+| ^K|=O ( 1 d ) ( _i=0^T-1 _i| _is-τ|+|s|+|τ| )+O (d^-aC ). The first part vanishes at the construction where s=τ=0s=τ=0. The second part is set by the length of training examples. Consequently, the population dynamics close exactly in terms of (λV,λK)( _V, _K) and the deterministic weights βi _i, up to the interaction terms, yielding a finite-dimensional ODE. Now, we prove the convergence of training transformers on quadratic functions based on the above forms of ODEs. By Lemma A.9, the block-sparse structure is preserved under gradient flow, so the dynamics are fully characterized by the two reduced matrices W31PV(t)W^PV_31(t) and W13KQ(t)W^KQ_13(t). We further use an isotropic symmetry: W13KQ(t)=λK(t)Idd,W31PV(t)=λV(t)Idd,∀t≥0W^KQ_13(t)= _K(t) _d,\ W^PV_31(t)= _V(t) _d,∀ t≥ 0, in order to identify the fixed point for the ODE. Then, the model reduces to scalars (λK(t),λV(t))( _K(t), _V(t)). We now show the convergence in two stages. In the first stage, using the structure of the ODE and the condition (here enforced by the colinear initialization θ0=α0θ⋆ _0= _0θ with α0∈[3/4,1) _0∈[3/4,1)), we show that whenever (λK,λV)( _K, _V) is outside a small neighborhood of (1,1)(1,1), the gradient norm is bounded below by a positive quantity depending at most polylogarithmically on d: ‖∇LCoT(λK,λV)‖22≥c1LCoT(λK,λV),when ‖(λK,λV)−(1,1)‖≥δ0\|∇ L^CoT( _K, _V)\|_2^2\;≥\;c_1L^CoT( _K, _V),when \|( _K, _V)-(1,1)\|≥ _0, for some constants c1>0c_1>0 and δ0=O⁡(1log⁡d) _0=O( 1 d). This yields an exponential decrease of LCoTL^CoT until the trajectory enters the δ0 _0-ball around (1,1)(1,1), in time t1=O⁡(log⁡d)t_1=O( d). In the second stage, inside the δ0 _0-ball of (1,1)(1,1), LCoTL^CoT is smooth, and its Hessian is positive definite (this follows from the identifiability of the construction in the assumption). Hence a local PL-condition inequality holds: ‖∇LCoT(λK,λV)‖22≥c2LCoT(λK,λV),when ‖(λK,λV)−(1,1)‖≤δ0,\|∇ L^CoT( _K, _V)\|_2^2\;≥\;c_2L^CoT( _K, _V),when \|( _K, _V)-(1,1)\|≤ _0,for some c2>0c_2>0. This implies LCoT(t)≤LCoT(t1)exp⁡(−c2(t−t1))L^CoT(t)≤ L^CoT(t_1) (-c_2(t-t_1) ), so reaching LCoT(t)≤εL^CoT(t)≤ takes an additional t2=O⁡(log⁡1ε)t_2=O( 1 ) time. Thus, the transformer converges to the construction as described in Lemma 3.3. Following Lemma 3.3 and Theorem A.5, we similarly obtain the rate of the inference loss when training converges. ∎ Appendix B Omitted Experimental Results B.1 Evaluations of the Stability Measure Recall that Lemma 3.1 is a worst-case upper bound, and the figures in the main text compare it to the measured loss gap on separate axes. We report the difference in scale explicitly. For each model and each reasoning length T, we report six quantities: • The loss gap |L⁡(fW,)−L⋆(fW,)||L(f_W;D)-L (f_W;D)| between the inference loss and the reference loss; • The value of our bound, i.e., the right-hand side of equation (4); • The bound-to-gap ratio, i.e., the ratio of the previous two quantities; • The first-order term C1∑t=1T−1ρt(T)εtC_1 _t=1^T-1ρ^(T)_t _t; • The second-order term C1C22∑t=2Tρt(T)(∑i=1t−1ri)2 C_1C_22 _t=2^Tρ^(T)_t ( _i=1^t-1r_i )^2; • The two constants entering the bound: the Lipschitz constant C1C_1 of the loss function, and the Lipschitz constant C2C_2 of the model’s input Jacobian. We estimate C1C_1 as the maximum norm of the loss gradient with respect to the model output over the evaluation set, and C2C_2 as the maximum finite-difference change in the input Jacobian over pairs of nearby inputs from the same set. Table 2 reports all six quantities for Qwen-1.5B, Llama-1B, and Gemma-2B trained on the Bellman-Ford task, for T=1,…,5T=1,…,5. At T=1T=1 there is no accumulated error, and both the gap and the bound are zero. We report the results in Table 2. First, the bound is relatively tight on Qwen-1.5B, where the bound-to-gap ratio stays between 66 and 99 across all T. On Llama-1B and Gemma-2B, the ratio reaches 2.3×1062.3× 10^6 and 1.5×1051.5× 10^5 at T=5T=5. Second, this difference is attributed to the constants. C1C_1 is small and comparable across all three models (between 0.010.01 and 0.070.07), while C2C_2 differs by roughly three orders of magnitude, from around 55 on Qwen-1.5B to above 3×1033× 10^3 on Llama-1B. Since C2C_2 multiplies the squared accumulated deviations in the second-order term, a large C2C_2 increases the bound value. Third, the second-order term grows faster in T than the first-order term. For Llama-1B and Gemma-2B, it dominates the bound from T=3T=3 onward. On Qwen-1.5B, the first-order term remains the larger of the two. Table 2: Decomposition of the bound in equation (4) for three models trained on the Bellman-Ford task. Loss gap denotes |L⁡(fW,)−L⋆(fW,)||L(f_W;D)-L (f_W;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)εtC_1 _t=1^T-1ρ^(T)_t _t and C1C22∑t=2Tρt(T)(∑i=1t−1ri)2 C_1C_22 _t=2^Tρ^(T)_t ( _i=1^t-1r_i )^2. The absolute scale of the bound varies by orders of magnitude across architectures. At T=1T=1 there is no accumulated error, so the loss gap and the bound value are zero. Terms of equation (4) Constants Model T Loss gap Bound Ratio First-order Second-order C1C_1 C2C_2 Qwen-1.5B 1 0 0 – 0 0 0.01 0 2 0.08 0.68 8.0 0.68 9.5×10−49.5× 10^-4 0.02 5.44 3 0.25 1.5 6.1 1.3 1.7 ×10−1× 10^-1 0.03 5.83 4 0.47 3.4 7.2 2.5 8.4 ×10−1× 10^-1 0.05 5.83 5 0.83 6.9 8.3 4.5 2.4 ×100× 10^0 0.07 5.83 Llama-1B 1 0 0 – 0 0 0.04 0 2 0.05 3.3×1033.3× 10^3 6.2×1046.2× 10^4 2.3×1032.3× 10^3 9.7×1029.7× 10^2 0.03 3.03×1033.03× 10^3 3 0.30 5.4×1045.4× 10^4 1.8×1051.8× 10^5 6.2×1036.2× 10^3 4.8×1044.8× 10^4 0.04 3.17×1033.17× 10^3 4 0.85 5.5×1055.5× 10^5 6.4×1056.4× 10^5 6.4×1036.4× 10^3 5.4×1055.4× 10^5 0.05 3.18×1033.18× 10^3 5 1.84 4.3×1064.3× 10^6 2.3×1062.3× 10^6 5.0×1045.0× 10^4 4.2×1064.2× 10^6 0.05 3.21×1033.21× 10^3 Gemma-2B 1 0 0 – 0 0 0.00 0 2 0.09 1.1×1031.1× 10^3 3.6×1043.6× 10^4 7.3×1027.3× 10^2 4.3×1024.3× 10^2 0.01 1.05×1031.05× 10^3 3 0.48 3.5×1043.5× 10^4 7.3×1047.3× 10^4 1.5×1031.5× 10^3 3.3×1043.3× 10^4 0.03 1.15×1031.15× 10^3 4 1.21 7.6×1047.6× 10^4 6.2×1046.2× 10^4 2.0×1032.0× 10^3 7.4×1047.4× 10^4 0.03 1.24×1031.24× 10^3 5 2.11 3.2×1053.2× 10^5 1.5×1051.5× 10^5 2.7×1032.7× 10^3 3.2×1053.2× 10^5 0.06 1.37×1031.37× 10^3 Table 3: Decomposition of equation (4) under quantization-aware training, for Qwen-1.5B at T=5T=5. Lowering the bit-width monotonically reduces the error amplification factor ∑t=1T−1ρt(T) _t=1^T-1ρ^(T)_t and increases the per-step prediction errors, summarized by ∑t=1T−1εt _t=1^T-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. Task Bit-width Test accuracy (%) Loss gap Bound ∑tρt(T) _tρ^(T)_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 2 75.3 0.9 2.4 9.4 2.9 3 74.2 1.1 3.3 13.9 2.7 4 73.8 1.6 5.2 22.6 1.9 32 73.1 2.3 14.2 34.4 1.8 Prim 1 66.3 1.1 1.1 11.9 3.0 2 65.6 1.3 1.4 14.7 3.0 3 64.7 1.7 1.5 20.1 2.8 4 63.5 2.0 2.0 30.1 2.5 32 62.3 2.1 2.7 69.8 2.3 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=5T=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. Terms of equation (4) Constants Model Method Loss gap Bound First-order Second-order C1C_1 C2C_2 Llama-1B SFT-CoT 1.84 4.3×1064.3× 10^6 5.0×1045.0× 10^4 4.2×1064.2× 10^6 0.050 3.21×1033.21× 10^3 Coconut 0.699 1.1×1051.1× 10^5 1.3×1031.3× 10^3 1.1×1051.1× 10^5 0.077 1.63×1031.63× 10^3 Alg. 1 0.491 1.3×1.3× 10^4 7.4×1027.4× 10^2 1.2×1041.2× 10^4 0.091 1.79×1031.79× 10^3 Gemma-2B SFT-CoT 2.11 3.2×1053.2× 10^5 2.7×1032.7× 10^3 3.2×1053.2× 10^5 0.060 1.37×1031.37× 10^3 Coconut 0.662 1.5×1041.5× 10^4 1.34×1021.34× 10^2 1.5×1041.5× 10^4 0.094 9.27×1029.27× 10^2 Alg. 1 0.091 4.9×4.9× 10^3 26.6 4.9×1034.9× 10^3 0.064 4.78×1024.78× 10^2 In Figure 5, we measure the loss gap and the error amplification factor ∑t=1T−1ρt(T) _t=1^T-1ρ^(T)_t for T ranging from 1 to 5. We report the results evaluated with Gemma-2B and Llama-1B models on the Bellman-Ford and Prim’s algorithms. The inference loss increases significantly with the number of steps. The error amplification factor scales exponentially with T, closely tracking the growth in the inference loss. 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. Evaluations of the bound under quantization. Table 3 decomposes equation (4) across bit-widths for Qwen-1.5B on three tasks. Lowering the bit-width has two opposing effects on the bound. First, it reduces the error amplification factor: from full precision to 11 bit, ∑tρt(T) _tρ^(T)_t decreases by 3.7×3.7× on the Dijkstra task, 5.9×5.9× on the Prim task, and 1.2×1.2× on the Bellman-Ford task. Second, the quantized model fits the training data less well, so the per-step prediction errors tend to grow. In practice, we observe that QAT reduces the error amplification factor more significantly than it increases prediction errors. As a result, lower bit-width quantization typically achieves a reduction in the bound, leading to improved test accuracy. On Dijkstra and Prim, the prediction errors increase by only 1.6×1.6× and 1.3×1.3×. The bound decreases monotonically as the bit-width is lowered, by 6.8×6.8× and 2.5×2.5× at 11 bit. The test accuracy improves by 3.1%3.1\% and 4.0%4.0\% over full precision. On Bellman-Ford, the reduction of ∑tρt(T) _tρ^(T)_t is small, while the prediction errors double from 44 bits to 11 bit. Therefore, the bound and the test accuracy are both best at 44 bits. The bit-width at which the growth of the error levels begins to dominate depends on the task, and we select it on the validation set. Evaluations of our bound across algorithms. To evaluate whether the relationship between our bound and the inference loss extends beyond Qwen-1.5B, we apply Algorithm 1 to Llama-1B and Gemma-2B on Bellman-Ford with T=5T=5, and compare it with standard chain-of-thought fine-tuning (SFT-CoT) and continuous chain-of-thought (Coconut) [63]. As shown in Table 4, our algorithm attains the lowest loss gap and the lowest bound on both models. Compared with Coconut, it reduces the loss gap by 58%58\% on average and the bound by 78%78\% on average. Compared with SFT-CoT, it lowers the bound by 330×330× and 65×65×. This confirms that the reduction in test loss correlates with the reduction of our bound components. Numerical evaluation of Jacobian norms. We further analyze the training dynamics by examining the Hessian trace of models fine-tuned with chain-of-thought on linear and quadratic functions. Specifically, we evaluate the norm of the Jacobian matrix product as a function of the reasoning length T and report results for both settings. As shown in Figure 6, as the reasoning length T increases, the norm of the Jacobian matrix product converges to a value. This behavior is consistent and aligned with our analysis of the gradient descent update rule for linear and quadratic functions. (a) Jacobian matrix norm on linear functions (b) Jacobian matrix norm in quadratic functions 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−ηnXX⊤)i‖22 _i=1^T\|(I- ηnX )^i\|_2^2 and the Jacobian norm of quadratic functions is ∑i=1T‖(I−2ηnXD(θi)X⊤)i‖22 _i=1^T\|(I- 2ηnXD( _i)X )^i\|_2^2. B.2 Sample Complexity Comparison We compare models fine-tuned with a varied number of training examples between 1000,20001000,2000, and 50005000. We also evaluate a model fine-tuned without intermediate steps. After training, we evaluate the noise stability using the fine-tuned model. We inject isotropic Gaussian noise with standard deviation σ=0.005σ=0.005 into the model weights and measure the increase in the inference loss. We compute the average loss increase with 1010 random seeds. We evaluate the noise stability of models fine-tuned with chain-of-thought on a range of algorithmic reasoning tasks, including DFS, Dijkstra, Prim, and DAG shortest path. First, we find that increasing the number of training samples leads to improved noise stability across all tasks. This suggests a strong correlation between generalization performance and robustness to weight noise. For the algorithms of the same task, we noted a difference in noise stability. Specifically, models fine-tuned with the Dijkstra algorithm have more than 2×2× lower Hessian traces than the Bellman-Ford algorithm. Additionally, models fine-tuned by Kruskal’s algorithm yield 15×15× lower Hessian trace values than Prim’s. Second, we study the effect of intermediate steps on the sample complexity. We randomly subsample the intermediate steps T to 2020, 1010, 55, 22, and 00 where T=0T=0 indicates training without intermediate steps. We find that a larger T substantially reduces sample complexity. Using intermediate steps consistently achieves better sample complexity compared to training without intermediate steps. Hyperparameters. We report the hyperparameters used in our experiments in Table 5. 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. Task Quantization Optimizer Learning rate LoRA Rank Sampling factor λ Latent steps k Dijkstra 1-bit AdamW 2×10−52× 10^-5 16 0.1 1 Bellman-Ford 4-bit AdamW 2×10−52× 10^-5 16 0.2 1 Prim 4-bit AdamW 2×10−52× 10^-5 16 0.1 1 BFS 1-bit AdamW 2×10−52× 10^-5 16 0.1 1 DFS 3-bit AdamW 2×10−52× 10^-5 16 0.1 1 Cyclic 1-bit AdamW 2×10−52× 10^-5 16 1.0 1 Symmetry 1-bit AdamW 2×10−52× 10^-5 16 1.0 1 B.3 Discussions Our formulation targets supervised fine-tuning on ground-truth steps. Reinforcement learning methods like Group Relative Policy Optimization or Proximal Policy Optimization optimize over the model’s own rollouts. Our analysis extends to reinforcement learning in two key ways. First, because test-time generation is independent of the training objective, our stability analysis applies equally to models trained with reinforcement learning. Second, our work can be applied to study whether policy-gradient training implicitly regularizes input Jacobian norms along high-reward trajectories. We leave the formal extension of our bounds to the reinforcement learning setting for future work. Lastly, we discuss several limitations of this work. First, our method regularizes these Jacobian norms during training. Future work could develop test-time regularization strategies that adapt to each input example during inference. Second, for black-box models, where gradients and Jacobians are not directly accessible, proxy stability measures based on noise perturbations may provide a practical alternative. Finally, extending the Jacobian-based error accumulation analysis to broader language-model tasks, including tool use, generation, and question answering, remains an interesting direction.

Instant research discovery