Training Nonlinear Transformers for Chain-of-Thought Inference:  A Theoretical Generalization Analysis
Authors: Hongkang Li
Songtao Lu
Pin-Yu Chen
Xiaodong Cui
Meng Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Theoretical Chain-of-Thought
generalization
deep learning theory
In-Context Learning
training dynamics of Transformers
Primary area: learning theory
TL;DR: We provide the first theoretical study of training nonlinear Transformers to obtain the Chain-of-Thought generalization ability.

Abstract: Chain-of-Thought (CoT) is an efficient prompting method that enables the reasoning ability of large language models by augmenting the query using multiple examples with multiple intermediate steps. Despite the empirical success, the theoretical understanding of how to train a Transformer to achieve the CoT ability remains less explored. This is primarily due to the technical challenges involved in analyzing the nonconvex optimization on nonlinear attention models. To the best of our knowledge, this work provides the first theoretical study of training Transformers with nonlinear attention to obtain the CoT generalization capability so that the resulting model can inference on unseen tasks when the input is augmented by examples of the new task. We first quantify the required training samples and iterations to train a Transformer model towards CoT ability.  We then prove the success of its CoT generalization on unseen tasks with distribution-shifted testing data. Moreover, we theoretically characterize the conditions for an accurate reasoning output by CoT even when the provided reasoning examples contain noises and are not always accurate. In contrast, in-context learning (ICL), which can be viewed as one-step CoT without intermediate steps, may fail to provide an accurate output when CoT does. These theoretical findings are justified through experiments.

Published as a conference paper at ICLR 2025
TRAINING NONLINEAR TRANSFORMERS FOR CHAIN- OF- -HOTHHTGHT INFERENCE: A THEORETICAL GENERAL- IZATION ANALYSIS
Pin-"Y Hongkang Li', Songtao Lu Pin-Yu Chen', Xiaodong Cui Meng Wang Rensselaer Polytechnic Institute, The Chinese University of Hong Kong, 'IBM Research
ABSTRACT
Chain-of-Thought (CoT) is an efficient prompting method that enables the rea- soning ability of large language models by augmenting the query using multiple examples with multiple intermediate steps. Despite the empirical success, the theoretical understanding of how train Transformer to achieve the CoT ability remains less explored. This is primarily due to the technical challenges involved analyzing the nonconvex optimization of our knowledge, this work provides the first theoretical study of training Trans- formers with nonlinear attention to obtain the CoT generalization capability so that the resulting model can inference on unseen tasks when the input is augmented by examples of the new task. We first quantify the required training samples and iterations to train a Transformer model towards CoT ability. We then prove the suc- cess of its CoT generalization on unseen tasks with distribution Moreover, we theoretically characterize the conditions for an accurate reasoning output by CoT even when the provided reasoning examples contain noises and are not always accurate. In contrast, in-context learning (ICL), which can be viewed one-step CoT without intermediate steps, may fail to provide an accurate output when CoT does. These theoretical findings are justified through experiments.
nonlinear attention models. To the best
testing data.
INTRODUCTION
Transformer- based large-scale foundation models, such as GPT-3 (Brown et al. 2020 GPT-4 (OpenAl (003). LLaMa (Touvron et al. (023a b). and Sora (Liu et al. 2024). have demonstrated
remarkable success across various tasks, including natural language processing (Brown et al. 2020
Touvron et al.] 2023b). multimodal learning (OpenAl 2023, Radford et al.) (021), and image/ video generation (OpenAl][023) [iu et al., 2024).. What more surprising is that large language models (LLMs) demonstrate reasoning ability through the so-called "Chain-of-Though
(CoT) method (Wei et al. 2022). The objective is to let a pre-trained LLM generate K steps of reasoning given input query Equery without any fine-tuning To achieve that, the input Equery is augumented with I examples {x,, [3 j.K K- of a certain K-step reasoning task, where each æ; is the input with y as the j-th reasoning step, and ys.k is the final output. A pre-traae pre-trained model then takes the resulting augmented input, referred to as a prompt, and outputs the corresponding reasoning steps {2, JK for guueryy or simply outputs zx- CoT can be viewed as extended and more intelligent method than the previous in-context learning (ICL) method, where only input- pairs {z,, Yr Y augmented in the prompt to predict z with the pre-trained model.
Inspired by the outstanding empirical performance of CoT in arithmetic reasoning (Wang al. 2023 Zhang et al. 2023b. Wang & Zhou 2024), symbolic reasoning (Zhang et al. 2023b Zhou
et al. 2023), and commonsense reasoning (Wang al., 2023) WWang & Zhou 2024), there have
been some recent works (Li et al., 2023c Feng et al. 2023 Li et al. 2024d ! Yang et al., 2024 Wen al. 2024) the theoretical understanding of COT. These works investigate CoT from the perspective of expressive power, i.c., they construct the Transformer architecture that is proven to have the CoT ability, They also demonstrate empirically that supervised training on pairs of CoT
"This work was done when Prof. Songtao Lu was at IBM Research. Corresponding author. Email: wangm7 @rpi.edu.

---

Published as a conference paper at ICLR 2025
prompts and corresponding outputs can lead to models with CoT ability. However, none of these results theoretically address the question of why a Transformer can obtain generalization cenoralizatioonguarattedd guaranteed CoT ability by training from data with gradient- basdd methods. Mcanwhile, another line of research (Zhang et al. 2023a Huang et al., 2023 Wu et al. 2023 Li et al. (2024)) aims to unveil the reasons behind the ICL ability of Transformers through characterizing the training dynamics of Transformer in the supervised setting. These analyses are specifically applicable to ICL. Therefore, a theoretical question still remains less explored, i.c.,
Why can a Transformer be trained to generalize on multi-step reasoning tasks via CoT?
1.1 MAJOR CONTRIBUTIONS
Following Li et al. 2023c : Feng et al. 2023 et al. (002dd Yang et al. (0024): Wen et al. (2024)
train the model in a supervised setting using prompt and label pairs. This paper provides the first theoretical analysis of the training dynamics of nonlinear Transformers to achieve CoT ability. We prove that the learned model has guaranteed CoT ability for new tasks with distribution shifts from the training tasks, even when there exist noisy and erroneous context examples in the prompt. We theoretically characterize the required number of training samples and iterations needed to train desirable model and the number of context examples required for successful CoT reasoning with a generalization guarantee. Moreover, we provide a theoretical explanation for why CoT outperforms ICL in some cases. Our main technical contributions are as follows:
quantitative analysis how the training can enable the CoT ability: We theoretically analyze the training dynamics required number of context examples in each training sample, the total number of training samples, and the number of training iterations needed to acquire CoT ability. We illustrate that the CoT ability results from the property that the attention values of the learned model are concentrated on testing context examples with the same input patterns as the testing query during each reasoning step.
one-layer single-head attention-only Transformer and quantify the
2. A quantitative analysis of how context examples affect CoT performance: We characterize
the required number of context examples in the testing prompt for successful CoT reasoning when noise and error exist in contexts. Our quantitative bounds are consistent with the intuition that more accurate context examples and more similar examples to the query improve CoT accuracy.
3. A theoretical characterization of why CoT outperforms ICL: We provide a quantitative analysis of the requirements for successful ICL reasoning with our studied trained model. We show that successful ICL requires an additional condition that the prompt has a dominant number of correct input-label examples, while the success of CoT does not depend on this condition. This can be viewed as one of the possible reasons why CoT outperforms ICL.
1.2 RELATED WORKS
Expressive power of CoT Li et al. 2023c proves the existence of a Transformer that can learn multi-layer perceptron (MLP). They interpret CoT as first filtering important tokens and then making predictions by ICL. They also establish the required number of context examples for a desired prediction with the constructed Transformer. |Fng al. (2023); Li et al. (2024d); Merrill & Sabharwal (2024) show that Transformers with CoT are more expressive than Transformers without COT. Yang al. (2024): : Wen et al. (2024) show the superiority of standard Transformers in some reasoning tasks compared with recurrent neural networks and linear Transformers.
Theoretical analysis of ICL As a simplified one-step version of COT, ICL has gained much attention from the theoretical community. Garg et al. (2022); Akyürek et al. (2023): Bai et al, (2023): ; Guo et al. (023)) demonstrate that Transformers are expressive conduct many machine learning algorithms in context. Akyürek et al. (2023); Von Oswald et al. (2023): ; Ahn et al. (2023); Cheng et al. (2023);
Ding et al. (2024) especially show the existence of Transformers to implement gradient descent and its variants with different input prompts. Zhang et al. (2023a); Huang et al. ); Wu et al.)
(2023); Li al. (2024a) explore the training dynamics and generalization of ICL on singlc-attention Transformers. Cui et al. (2024); Chen al. (2024) provably show the superiority of multi-head attention over single-head attention to achieve ICL ability.
Training and Generalization of Transformers There have been several recent works about the opti-
mization and generalization analysis of Transformers. Jelassi et al. (2022): Li et al. (2023d); Oymak
et al. (2023);
Li et al. (2023 b: 2024b): Luo 2003): Huang al. (2024): Zhang et al. (0224) study
2

---

Published as a conference paper at ICLR 2025
the generalization of one-layer Transformers by assuming spatial association, semantic//ontextneee structure, or the majority voting of tokens in the data. Oymak et al. (2023); Tarzanagh et al. (2023b a);
Tian et al. (2023a b): [i et al. 2024cc;; [idiz al. (2024); Nichamii et al. (2024); Makkuva et al. (2024b) investigate the training dynamics or loss landscape of Transformers for the next token prediction by assuming infinitely long input sequences, causal structure/Markov Chain of data, or proper prediction head. Deora et al.] ); [Cen & Li (2024) analyze the optimization and generalization
of multi-head attention networks.
PROBLEM FORMULATION
We study the problem of learning and generalization of K-steps reasoning tasks. Each task f fK o....... f2 a fi is a composition of functions {ft) and outputs labels ž1, 22, ... , 22 for the input Equery. During the k-th reasoning step, & E [K], the label is z fx(#x-1), where 2o == Bguery
2.1 TRAINING TO ACQUIRE THE CHAIN-OF THOUGHT ABILITY
Following theoretical analysis (Feng al.. 2023 Li et al. 2024d Wen et al. 2024) and empirical works like process supervision (Lightman et al. 2024D. we first investigate the training on a Trans- former model to obtain the CoT ability in evaluating new data and tasks. It is supervised learning setting on pairs of prompts and labels. Different from the testing prompt that includes examples and only Equery, the training prompt includes multiple K-steps reasoning examples and (k 1)-step reasoning of a Queryy for any & in [K] and the label for this prompt is Zk. Specifically,
Training Prompt and Label for COT. For every prompt and output pair from a task / = Sx f1, we construct prompt P that include the query input -k- by prepending 4y reasoning examples and the first k 1 steps of the reasoning query. The prompt P of the query input 1 is formulated as: P (E1, E2, E... Qk) E 2dxxx x
f2 ...
where E.
ya.1
yi.K- yi
QF =
20
21
2k-2
Zk-1
E [erl,
(I)
Vi
Vi.
21
Z2
2k-1 0
where E, is the i-th context example, and Qx is the first l steps of the reasoning query for any k in [K]. We have Vi = (u(Vikk-1) and z5- = (2x-1) ) for i E [err]. k E [K] with a notation 91o0 == x; Let ps and Pyuery be the s-th column and the last column of P. respectively, for E [rrK k 1]. x.. ki Zj Rdx for i [er]] and j, E [K]. We respectively call æ; and yi.k context inputs and outputs of the k-th step of the ith context example. For simplicity of presentation, we denote x as the label of P. which indeed 27 for [T. All the notations are summarized in Table
3 in Appendix.
The learning model is single-head, one-layer attention-only Transformer We consider positional
encoding ) E R22xx Following theoretical works (Jelassi et al. 2022 Huang et al. 2024 Ildiz
al. 2024), we add the positional encoding to each p; by p. = P: + Ci mod K). 2 E (4rr I)I.
Pouery is also defined by adding the corresponding Cg to Pquery- Mathematically, given a prompt P
defined in [I with len(P) (which is at most K (le, + 1)) denoting the number of columns, can be
written
len(P) P)
F(U:P)= = Σ
Wvp: softmax((WPp)) W WQPguery),
(2)
i=1
where Wo, WK E ( Wv E xdxxx(2dx) are the embedding matrices for queries, keys, and values, respectively. V = (We, WK. 3 is the set of all model weights|] [T Typically, m > 2dz
Here, softmax((Wxpp)) WoPpuery) (Wxp.) WaPseery |>)m(()) (Wxi.) Wap
The training problem to enhance the reasoning capability solves the empirical risk minimization. a
min (V)
N (";
using N prompt and label pairs (Pa, "n) J For the n-th sample, x Tueryy and the context input
æf are all sampled from an unknown distribution D. the training task fd is sampled from T, k is
randomly selected from to K, and pa is constructed following (T). The loss function is squared loss, i.e., E(V; Pm. 2" 1|2
where FV;; PT) is defined in (
WWe focus one-layer single-head Transformer motivated by recent advancements and current state in Transformer and CoT analysis. Please see Appendix B. for discussion.
3

---

Published as a conference paper at ICLR 2025
2.2 TRAINING ALGORITHM
For simplicity of analysis, we let W (Jelassi et al. 2022, Huang et al. 2023 Zhang et al. 2023a Huang et al. 2024). Let {ce]ke be set of orthonormal vectors. The model trained using stochastic gradient descent (SGD) with step size y with batch size B, summarized in Algorithm [ in Appendix C Each entry of W is generated from (O (²) for tiny { >> O. Wy is fixed during the training. The fraction of prompts with 2k-1 the query input is 1/K K by uniform sampling for any k [K] in each batch.
W K Wo and W (Oayxel Iay] Rdxx(2xx) as
2.3 CHAUGHHF-TOUGHT INFERENCE
We then consider another K -sttps reasoning task T whose target is to predict labels ( given the input query Equery. T' is the set of testing tasks, and T' # T.
Testing Prompt for Co The testing prompt P is composed of lts (s 42r) context examples of K steps plus query, which is constructed as
[E); E.,... EE Pquery) E (R(dd)) -queery = ( queryy OT
(4)
where E, follows the form in ① for i [[ts).
We follow the CoT-I/O scheme formulated in (Li al., 2023c Feng et al.
2023 Li al. 2024d
Yang et al. 2024 Park al. 2024) the inference method. Specifically, for a K -step CoT with les examples on a certain ʃ 7 given the testing prompt P defined in (4). let Pi = P and Po be the first K 44 columns of P. When use CoT prompting for prediction in the k-tth step, first generate the output Uk. & [K] via greedy decoding by feeding the k-th step prompt P; to the trained model v obtained from (3) The greedy decoding scheme means outputting the most probable
token from the discrete set y of all possible outputs, as stated in (5).
U2 = arg min us) 2 F(W; Pa)) - ull? (greedy decoding)
(5)
Then, we use the output 22 to update Ps and use ok as the query input to form the input prompt Pk+1 for the next step, which is computed as
PR = 1 qu) E xx(2d))x( +k)) Pa+1 = (Pr gk+1) E (K()))
where qs = (1) TT 4k+1
(v) TT)
where 9k is the k-th step reasoning column for the query. The model finally outputs "..... Uk as
CoT result for query Equery by (5). The CoT process is summarized in Algorithm of Appendix C
When K 2 2, following (ii et al. 2023c Feng et al., 2023 Li et al. 2024d Yang et al. 2024). the CoT generalization error given the testing query Qquery: the testing data distribution D', and the
labels { K steps testing task f E T is defined
K
RICOT Ae JET
VV)
1 + mk]
K
which measures the average eπror between the output and the label of each reasoning step. A zero CoT generalization error indicates correct generations in all steps.
2.4 IN-CONTEXT LEARNING INFERENCE
The ICL inference on steps reasoning task I T only predicts the final-step label by perpending examples of input and label pairs before the query. ICL can be viewed as one-step CoT without intermediate steps. Here, we evaluate the ICL performance the trained model.
Testing Prompt for ICL. Mathematically, ICL is implemented by constructing prompt P below,
P= EE ... Ei., Pquery), where Pquery =
TQquery
E.
D;
0
0
(8)
y:
0
E, E [( for E [4e.]. Note that in the ICL setting, E, only has input 2; and the K-step output yr.k but does not include any intermediate labels. We pad zeros in E, so that its dimension is the same as E, in (T) for the inference with the same model as for COT. The ICL u following (5) The ICL generalization error is
P E (x
output arg minney
2 [FV
Ru .ET" (T)
R
E Ea
[1[z 7 v]],
9
which measures the error between the one-step step reasoning output and the final step label.
4

---

Published as a conference paper at ICLR 2025
3 THEORETICAL RESULTS
We first summarize the main theoretical insights in Section 3.1 Then, we introduce the formulation of data and tasks in Section 3.2 Sections 3.3 3.4 and 3.5| respectively characterize the training
analysis of the Transformer and generalization using CoT and ICL with the trained model.
3.1 MAIN THEORETICAL INSIGITTS
We consider the setup that the model is trained using samples generated from tasks in T that operate on M orthonormal training-relevant (TRR) patterns, while both CoT and ICL are evaluated on tasks in T' that operate M' orthonormal testing-relevant (TSR) patterns. We consider the general setup that the context examples in the prompt for CoT and ICL testing are both noisy, i.e., TSR patterns with additive noise, and partially inaccurate, i.e., the reasoning in some examples contains incorrect steps. Our main insights are as follows.
P1. Training Dynamics of Nonlinear Transformer towards CoT. We theoretically analyze the training dynamics on a one-layer single-head attention-only Transformer to acquire the CoT general- ization ability and characterize the required number of training samples and iterations. Theorem [ shows that to learn model with guaranteed CoT ability, the required number of context examples in each training sample and the total number of training samples/iterations are linear in a and a=?, respectively, where a is the fraction of context examples with inputs that share the same TRR patterns as the query. This is consistent with the intuition that the CoT performance enhanced more context examples are similar to the query. Moreover, the attention values of the learned model are proved to be concentrated on testing context examples that share similar input TSR patterns as the testing query during each of the reasoning steps (Proposition D. which is an important property that leads to the success of the CoT generalization.
P2. Guaranteed CoT Generalization. To achieve zero CoT error on tasks in T' and data based on TSR patterns that contain a non-trivial component in the span of TRR patterns with the Icarned model, Theorem 2 shows that the required number of context examples, where noise and errors are present, for task / the testing prompt is proportional to (a'r'pp) -2 Here, a' is the fraction of context examples with inputs that share the same TSR patterns as the query. 7f in (0, 1) measures the fraction of accurate context examples, and a larger constant o in (0, 1) reflects a higher reasoning accuracy in each step of the examples. This result formally characterizes the intuition that more accurate context examples and more similar examples to the query improve the CoT accuracy.
P3. CoT outperforms ICL. In Theorem 3, context examples for ICL to be successful has a similar form to that for CoT in Theorem 2 but with an additional requirement (Condition 1) that the fraction of correct input-label examples in the testing prompt must be dominant. Because not all testing cases satisfy this requirement, our result provides one explanation for why CoT sometimes outperforms ICL.
We theoretically show that the required number of testing
3.2 THE FORMULATION OF DATA AND TASKS
Training data and tasks: Consider M training-relevan (TRR) patterns 44, u2..... HMM. which form an orthonormal set M { M = (d), M S d. (1, 0,2)) 1 Cg for E [M'), k E [K].
Every training prompt P in ( contains the query and training examples from the same training task f in the set of training tasks T. Specifically, each training task f is a composition of K functions f fx .... o f2 fi where each function fa belongs to a function set F. The k-th step label of the query is 2₁ = [x(2z=1) given the k-th step input zb-1 with z E M, k E [K]. Moreover, the k-th step label of the i-th (i [Err)) context example is y f& ( given the k Ith step input
Vi,k-1, k E [K] with æ;, Vi.k E M,, where yıoo
We assume that fa fa( * fie fax if and only
if either x + x² or fo + far-
Training prompt: Consider training prompt P task f T defined in T with the query input Zk-1, k E [K]] Let a E (0, 1 - c] for some constant e >
(F) denote the fraction of context examples
with input sharing the same TRR pattern as the query input.
The formulation of f is motivated by recent theoretical works on model training ICL with Transformers. Please see Appendix [.2 for details.
Thiis is to prevent the trivial case that the model only learns the positional encoding but not the TRR patterns when becomes arbitrarily close to 1.
5

---

Published as a conference paper at ICLR 2025
Testing task and query: Consider M' testing-relevant (TSR) patterns μ4, pl2, H'ss which form an orthonormal set M' { M M' I M. We also have μ L C₁ for E [AM'], E [K]. Let T' denote the set of testing tasks, which all operate on patterns in M' rather than M in training tasks in
T. Every testing task f fı o.... f2 f1 E T is a composition of K functions. The reasoning for the testing query is considered to be noiseless and accurate. That means,
Z E M' for all k E {0} U [K], and 2 fu( (1) Z0 3 Equery:
Testing prompt: We consider the general setup that testing examples are noisy and erroneous. By noisy examples, we mean all inputs and outputs of each step are noisy versions of TSR patterns, i.e.,
Di Y
{b R6 1e + j. [M"], 5 1 MM, || = 22/2),
(10)
with noise 6 * 0 for i E [K//] & E [K]. Denote TSR Rdd + Z+ as a function that outputs the index of the TSR pattern the noisy input. We consider the case that least an a' fraction of context examples where the TSR pattern of the input y, E [24.] is the same as Equery.
By erroneous examples, we mean that the reasoning steps in test examples may contain errors. To formally model this, we define the step-wise transition matrices {A J E g. xMM' such that A represents the reasoning probabilities of step l in test examples. Specifically, there exists some constant of in (0, 1) such that for all 8 E [1 k E [K], the j,h entry of A satisfies
((J) = Pr(TSR( x) = (/TSR(Be.AA))) = i):
(11)
and AL k(tsrr) 2 1/(1 - p²) - (1.5)) Vj E [M') where 1lg- = Fx(14),
Note that (1T) characterizes general case in inference that for any given k, in the k-th reasoning step of the test example, the k-th step output is a noisy version of the true label with the highest probability, which guarantees that the examples are overall informative in the k-th step. This requirement is intuitive because otherwise, these examples would overall provide inaccurate information on the k-th step reasoning. Moreover, (II) models the general case that, with some probability, the k-step reasoning is inaccurate in the examples. o is referred to as the primacy of the step-wise transition matrices. pf reflects the difference in the probability of correct reasoning and incorrect reasoning in each step, and a larger pf indicates a larger probability of accurate reasoning.
Let B I Al. be the K K-step transition matrix. Then B (1.3) is the probability that the K-th step output is noisy version of pzz, when the input a noisy version of p³ in the testing example. We similarly define oá in (0, 1) as the primacy of BJ, where
BB 2 1/(1 pb) Blaay Vj E [M') j arg max BB jE(M'
(12)
Example 1. Consider simple two-step inference example with K = 2, udy péy the TSR pattern, and 8 = O in inputs and outputs of every step, as shown Figure 7 The bluck solid arrows denote the correct inference process, where fū (1(() = ph fi(H') ulz f2(H() =Wz. and f2(m() uh- Hence, μn > pi - 2/2 and μl2 + 112 - μ1 are two inference trajectories under the function f. The testing examples contain errors and follow the transition matrices A and A (brown dashed
arrows). We let Af-
(0.6 0.4)
A =
0.44 0.6
which results in B 0.64
/0.56 0.44
0.4 0.6)
0.8 0.2).
0.36
3.3 THE SAMPLE COMPLEXITY ANALYSIS OF THE TRAINING STAGE
We first characterize the convergence and the testing perfor- mance of the model during the training stage with sample complexity analysis in Theorem
Step
Step 2
M1
Theorem 1. For any e > O, when (i) the number of context examples in every training sample is
0.6
0.4
0.4
0.4
ler 2 0(x 3)), (ii) the number of iterations satisfies K
(13)
0.6
0.6
0.2
M2
Figure I: An example of two-step inference
T 2 2nn a -2 -3 12310g
+ MM((a'++ +c")), (14)

---

Condition (v) indicates that, to achieve the desired CoT accuracy, the number of context examples should be proportional to -2 -=-² and - meaning it decreases as a', P rς increase. It can be interpreted as follows, if the number of context examples remains fixed, an increase in a', pl, Published as a conference paper at ICLR 2025
and (iii) the training tasks and samples are selected such that every TRR pattern is equally likely in
every inference step and in each training batel with batch size B 2 (mmax(c= .M} 2MM]. step size 1j < 1 and N BT samples, then with high probability, the returned model guarantees
-log M), the
E grrguery EM.JET [e(V; P.,) =)] S O(().
(15)
Theorem I indicates that with long enough training prompts and sufficient number of iterations and samples for training, a one-layer Transformer can achieve a diminishing loss of Oc() on data following the same distribution training examples. The results indicate that (i) the required number of context examples is proportional to a - (ii) the required number of iterations and samples increases M and a increases. As sanity check, these bounds are consistent with the intuition that it will make the training stage more time- and sample-consuning if the number of TRR patterns increases or the fraction of prompt examples that share the same TRR pattern as the query decreases.
3.4 COT GENERALIZATION GUARANTEE
In this section, first define two quantities, 7f and Tá for each testing task f T' based on the formulation testing data and tasks in Section 3.2 These two quantities are used to characterize the CoT and ICL generalization in Theorems 2 and respectively.
Definition 1. For fK f1 E T define the min-max trajectory transition probability
K
min (E(M))
II k=1
A (erp:l------oddt../rssee------- where fo(p)) Ui.ViE [M'),
(16)
which measures the minimum probability, over all the initial TSR patterns, of the K K-step reasoning trajectory that has the highest probability over all K K-step trajectories. We also define the min-max input-label transition probability as
min max B (E(M")
(17)
which measures the minimum probability, over all the initial TSR patterns, of the output that has the highest probability outputs.
For instance, in Example after 12), 75 min (0.36, 0.48) 0.36, T! min 0.64) 0.56.
Theorem 2 (CoT generalization), Given trained model, the training process of which satisfies conditions (i) to (ii) in Theorem then as long as
(iv) each TSR pattern pi in the orthonormal set {ur 1 satisfies
1 = }3 + j
(18)
. "pan(p)......ss)) and (v) the number of testing examples for any T 1T'1
where X span(
ر. E span [1) span(H1 ... HM). and || e((ogge-1)-
14 > D((a'r 7 -
(19)
have B B.oT..t.urre Mouery EM' JET (V) 0.
Remark Theorem 2 proves that a trained one-layer Transformer can generate all K K-steps reasoning correctly by CoT for a new task f in T' with two additional conditions. Condition (iv) means that each TSR pattern in the task set T' is the summation of a component that belongs to the span of the TRR patterns and component that is perpendicular to the span.
or Ti results improved CoT accuracy. This aligns with infuition, because a' represents the fraction of examples similar to the query, and o and 7 reflect the accuracy of the reasoning steps in the context examples.
'Our analysis assumes that the whole set of M is achievable uniformly in each step and training batch. This condition is to ensure balanced gradient update among all TRR patterns, as used in (Li et al. 2024al for ICI.
7

---

Published as a conference paper at ICLR 2025
3.5 ICL GENERALIZATION AND COMPARISON WITH CoT
Because only input-label pairs are used as context examples for ICL, the input-label pairs in context examples should be accurate overall be informative about the task. We formulate this requirement as Condition
Condition 1. For the testing task f = fK 0.... E TT, we have that for any i E [M'].
TSR( arg max B ((5))-
(20)
JE[M]
Condition D requires that context example, the input TSR is μ4 then f(4) is the output TSR pattern with the highest probability over all TSR patterns. Note that (IT) indicates that, for every k and i, when p is the k-th step input, fir (1(2 the step-k output with the highest probability over all TSR patterns. However, (11) does not necessarily imply (20). In Example given the input μ4, although the inference trajectory pt - μn -> uig under f has the highest probability over all 2-step trajectories, jú has the higher probability to be the final output than the correct output flz by the two-step transition matrix Bi, thus violating Condition
Our result of the ICL generalization is stated as follows.
Theorem 3 (ICL generalization), Given a trained model, the training process of which satisfies conditions (i) to (iii) of Theorem Aand 18 for the testing task ET',
Case A. if Condition does not hold, then RICLa EM' MET (V) 2 2(1). no matter how large the number of training samples lig is;
Case B. if Condition holds, then VICLaaper (V) 0, provided that
1/,, 2 20((''( 2 log M).
(21)
Remark (Comparison between CoT and ICL). Theorem Baa formally states that, Condition 7 is necessury for successful ICL generalization. Because Condition 7 is not required for CoT generalization, CoT performs better than ICL if Condition 1 fails Theorem33t characterizes that when Condition I holds, a desired ICL generalization needs a testing prompt length linear in a'
po-² 2 and To for the testing task f E T This result is the counterpart of the requirement ( 19) for the CoT generalization, indicating that more context examples with the same TSR pattern as the query and more accurate context examples improve ICL generalization.
Ref. et
et al. (2023c) also shows the advantage of CoT over ICL to learn MLP functions, but in
different setting from ours, where our studied tasks operate on patterns. More importantly, this paper characterizes the CoT and ICL performance theoretically when the testing task has a distribution shift from training tasks (TRR patterns to TSR patterns), and the testing examples contain errors, while Li et al. (223c) only empirically evaluates the CoT and ICL performance with noisy examples.
4
THE MECHANISM OF COT AND THE PROOF SKETCH
4.1
TRANSFORMERS IMPLEMENT COT BY ATTENDING TO THE MOST SIMILAR EXAMPLES EVERY STEP
largest fraction
We characterize the key mechanism of a properly trained one-layer Transformer to implement CoT on K-steps reasoning task via training dynamics amalysis of the atten- tion layer, as demonstrated in Figure 2 This is different from the mechanism study in (Li et al. 2023c] Feng et al. 2023) by constructing model that can conduct COT. We have the following proposition for the trained model.
greedy decoding
step
03
0.3
0.02
0.02
Wx
Wg
1
2
PL M4
Ps.
Pi Mi My
My PL
HaT query
Proposition 1. Let S denote the index set of the context columns of the testing prompt P in 4 that (a) correspond to the k-th step in a context example and (b) share the
Contoxt exomples of
Figure 2: Concentration of attention weights for CoT inference.
*Our insight of the comparison between CoT and ICL still holds when we evaluate CoT generalization only by the final step output. This is because successful CoT generalization in Theorem on all reasoning steps already ensures satisfactory CoT generalization the final step.
8

---

Published as a conference paper at ICLR 2025
same TSR pattern in the k-thh input the k-tth input Uk-1 of the query, k E [K] Given trained model that satisfies conditions (i) to (ii)) of Theorem 2 and (18) and (19) after T iterations, we have
Σ
softma((p. (P: ((T) q 2 - where p. Pi + C(i mod K); 9:- = qu + Ck.
(22)
VEST
with qle defined in 66. Moreover, for any f E T' the k-th step output d% given Equery = pl satisfies,
1% = J& o ... S (1()).
(23)
Proposition D first illustrates that, when conducting the k-th step reasoning of the query for any k E [K]], the trained model assigns dominant attention weights on the prompt columns that are also the le-th step reasoning of examples and share the same TSR pattern in the k-tth step input as the query. Then, given a sufficient number of testing context examples by (19). it is ensured that the fraction of the correct TSR pattern is the largest in the output of each step by (11). Subsequently, the generation by greedy decoding (5) is correct in each step, leading to a successtul CoTT' generalization.
4.2
AN OVERVIEW OF THE PROOF
The technical challenges of the proof are concentrated on Theorem I. where the property of the trained model is derived. The proof of Theorem i is built upon three Lemmas, which characterize the two stages of the training dynamics, i.e., Transformers first attend to tokens with the same step as the query and then, among them, further concentrate on tokens that share the same TSR pattern the query. Specifically, Lemmas and Show that if a training prompt P includes the first k steps of the reasoning query, then the attention weights on columns of P with a different step from the query decrease to be close to zero in the first stage. Lemma 3 computes the gradient updates in the second stage, where the attention weights on columns in P that correspond to the same step and have the same TRR pattern as the query gradually become dominant. Theorem Tunveils this training process by showing the required number of training iterations and sample complexity.
I
To prove Theorem 2 we first compute the required number of context examples for the new task E T' so that by concentration inequalities, the number of context examples with accurate TSR is larger than examples with inaccurate TSR patterns in all K reasoning steps with high probability. Then, by the correlation between TRR and TSR patterns (18). we also show that the trained Transformer can attend context columns with the same TSR pattern as the query. Therefore, the model can make the correct generation in each step. Theorem 33 follows a similar proof idea to Theorem 2 with the difference that the trained model predicts output directly from the input query following B' instead of A l E [K] in COT. Therefore, Condition is required for the success of ICL generalization.
NUMERICAL EXPERIMENTS
Data Generation and Model setup. We use synthetic data generated following Sections 2 and 3.2 Let dx = 30, M 20, M' 10, 0.4. We consider 3-steps tasks for training and testing, i.e.,
K = 3. A reasoning task f is generated by first sampling a set of numbers of permutations {p.) with p; E [M] and then let fk(()))) M)(ce=)) mod ar for i E [M], k, E [K]. The testing noise level set to be 0.2 for any examples and f T'. The learning model is one-layer single-head Transformer defined in (2) or three-layer two-head Transformer. We set 7¹ = 0.5, p = 0.8,
of 0.8 for CoT testing if not otherwise specified.
Experiments on the generalization of CoT. We first verify the required number of context examples for a desired CoT generalization on a one-layer Transformer. We investigate the impact of a', T and of by varying one and fixing the other two. Figure 3 illustrates that more testing examples are needed when a', -², or of is small, which verifies the trend of the lower bound of 1/, in (19).
Experiments on the generalization of ICL and a comparison with COT. We then verify the ICL generalization with the trained model. We vary T.. and p by changing 7- and p²_ Figure 3 indicates that more testing examples are required when a', T/. or p is small, which is consistent with our bound in (1)) We then consider the case where T! = 0.4 and Pf = 0.1 so that the generated testing prompt may not satisfy Condition depending on the specific choices of AC's. Figure 5 shows shows that when Condition Tholds, the ICL testing error decreases if the number of contexts increases. However, when Condition fails, the ICL testing error remains large, irrespective of the number of contexts.
9

---

Published as a conference paper at ICLR 2025
10a
10a
10"
Y-64
CoT Trrrrr
10-1
CaT error
10-1
(707
CaT error
10-1
10-2
10
10-2
deta
10
20 40
60 80 100
10
0
20 40
80 100
10
0
20 40
dett
60 80 100
context examples
context examples
context examples
(A)
(B)
(C)
Figure 3: CoT testing error with different (A) (B) r (C)
109
10%
109
ICL error
10-1
ICL error
10-1
CoT error
10-1
10
10
10
A--04
10
10-²
10 A004
60 80 100 of context examples (A)
20
100
60 80 100
# context examples
of context exampies
(B)
Figure 4: ICL testing error with different (A) a' (B) Tá (C) p!.
Experiments on the training dynamics of CoT. In Figure 6 we compute the total attention weights on four types of testing context columns along the training, which are contexts with the same (or dif- ferent) TSR pattern and in the same (or different) step as the query. The result shows that the atten- tion weights on contexts that share the same TSR pattern and in the same step as the query increase along the training and converge to around 1. This verifies the mechanism formulated in (22 Mean- while, Figure also justifies the two-stage training dynamics proposed in Section 4.2 where we add black vertical dashed line to demonstrate the stage transition boundary. We observe that the attention weights on context columns with different step, i.e., the red and yellow curves, decrease to zero in the first stage. Then, the attention weights on contexts with the same TSR pattern and the same step the query, i.e., the blue curve, increase mechanism of CoT on a three-layer two-head Transformer with a two-step reasoning task. Figure 7 shows that there exists at least one head in each layer of the Transformer that implements CoT as characterized in Proposition] D This indicates that the COT mechanism we characterize on one-layer Transformers can be extended to multi- multi-layer multi-head Transformers.
10"
Attention weight in average
1.0
Testing errar
10-1
same step same step
0.4
step
10-2
8
KLw Condiccn
wi Canditon
0 20 40 60 80 100
0.0
200
400 600 800 1000
context examples
Iterations
Figure 5: Comparison Figure 6: Training dy- between CoT and ICL namics of Transformers w./w.o. Condition
for CoT
in the second stage. We also justify the attention
1.0
1.0
1.0
TSA, det
15 sares
0.0
difl. TSAL
dil TSFI,
.ill
aff. 15tt (ll. Stap
dfl 1581 oll,
Artwottn weight
0.6
0.6
Attertion Weight
0.6
0.4-
0.2
0.2
0.2
0.0
HUG
Iterations
Iterations
(A)
(B)
(C)
Figure 7: Training dynamics of Transformers. (A) Layer 1, Head (B) Layer 2 Head 2 (C) Layer 3 Head 2.
6 CONCLUSION, LIMITATIONS, AND FUTURE WORKS
This paper theoretically analyzes the training dynamics of Transformers with nonlinear attention, together with the CoT generalization ability of the resulting model on new tasks with noisy and partially inaccurate context examples. We quantitatively characterize and compare the required conditions for the success of CoT and ICL. Although based on a simplified Transformer model and reasoning tasks operating patterns, this work deepens the theoretical understanding of the CoT mechanism. Future directions include designing efficient prompt- genertting methods for CoT and analyzing L.LM reasoning more complicated data model.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by National Science Foundation(NSF) #2430223, Army Research Office (ARO) W91 1 and the Rensselaer- Future of Computing Research Collaboration (http:/airre.ri..edd)) We also thank all anonymous reviewers for their constructive comments.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.00297 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.045: 046377 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1 33:1877-1901, 2020.
Sitan Chen and Yuanzhi Li. arXiv:2402.04084, 2024.
Provably learning a multi-head attention layer. arXiv preprint
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. 19442, 2024.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:23/2.06528, 2023.
Yingqian Cui, Jie Ren, Pengfei He, Jiliang Tang, and Yue Xing. Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401. 2024.
Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310. 12680, 2023.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. CausalLM is not op- timal for in-context leaminng In The Twelfth International Conference on Learning Representations, 2024. HRLLhtpss: htttttss///peeeeiieeeeeeee
net wwwpeee//ooommmnnnuneoeeeeeeeeee
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. Advances in Neural Information Processing Systems, 36, 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598 2022.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn mn-context beyond simple functions? a case study on learning with representa- 2023.
tions. arXiv preprint arXiv:2310.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In NewrIPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023.
Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn feature-position correlations in masked image modeling. arXiv preprint arXiv:2403.02233. 2024.
M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self- attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint arXiv:2402.1355 13512, 2024.
11

---

Published as a conference paper at ICLR 2025
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 15:37822-378366 2022.
Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference on Learning Representations, 2023a.
htttt::://ooeereeee....oo forum?
id [d=jcIGv3Q10)
Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, ZAIXI ZHANG, and Pin-Yu Chen. What improves the generalization of graph transformer? a theoretical dive into self-attention and positional encoding. In NeurlPS 2023 Workshop: New Frontiers in Graph Learning, 2023b. URL
TEttss://ppenevii.e net www.modddoaaaeeeeee.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear transformers learn and generalize in-context learning? In Forty-first International Conference on Machine Learning, 2024a. URL /tR/ttttt//////eeeeee//
net htt.:///ooooooia1111111aa
Hongkang Li, Meng Wang, Shuai Zhang. Sijia Liu, and Pin- Chen. Learning transformers is provable low-raw low-rank and sparse: one-layer analysis. arXiv preprint 17Xv:2240..1716, 2024b.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting thain-of-thoughtt Compositionality through in-context filtering and learning. Advances in Neural Information Processing Systems, 36, 2023c.
Yingcong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pp. 693. PMLR, 2024c.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. arXiv preprint arXiv:2303.044 04245,, 2023d.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transform- ers to solve inherently serial problems. In The Twelfth International Conference on Learning Representations, 2024d. URL http::ttp/eppe.eeeeeeeeeee
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ttp:ttt::///eeteeteetr
Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, Jianfeng Gao, et al. Sora: A review on background, technology. limitations, and opportunities of large vision models. arXiv preprint arXiv:2402.17177, 2024.
Yuankai Luo. Transformers for capturing multi-level graph structure using hierarchical distances. arXiv preprint arXiv:2308.1 // /29, 2023.
Ashok Vardhan Makkuva, Marco Bondaschi, Chanakya Ekbote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local global: Learning dynamics and effect of initialization for transformers. arXiv preprint arXiv:2406.03300 2024a.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.04100 2024b.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations, 2024. URL https /oopenevveww net/ pnppenrevree..ntffooooomaaaaaaad
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402. 14735, 2024.
12

---

Published as a conference paper at ICLR 2025
OpenAl. Gpt-4 technical report. OpenAl, 2023.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. tuning. arXiv preprint arXiv:2306.0343. 2023.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lec, Jaewoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how learn? a comparative study on in-context learning tasks. arXiv preprint arXiv: 2ee::2402.4444 4422044448, 2024.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference Machine Learning, pp. 8748-8763. PMLR, 2021.
Davoud Ataee Tärzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023a.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. CoRR, 2023b.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in I-layer transformer. arXiv preprint arXiv:2305.16388 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. Joma: Demystify- ing multilayer transformers via joint dynamics of mlp and attention. In Conference Parsimony and Learning (Recent Spotlight Track), 2023b.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint (:XXi:2302.13977 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Ilama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems. 30, 2017.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv: 101 1.3027, 2010.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174, PMLR, 2023.
Lei Wang, Wanyu Xu, Yihuai Lan, Zhiqiang Hu, Yunshi Lan, Roy Ka-Wei Lee, and Ee- Ee-Peng Lim. Plan-and-solve prompting: Improving zero-shot chain-of-thought reasoning by large language mod- els. In Proceedings of the 6t Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers), pp. 2609-2634 2023.
Xuezhi Wang and Denny Zhou. Chain-of-thouge reasoning without prompting. arXiv preprint 2xXiv:2402.10200 2024.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824 24837, 2022.
Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. Rons are not transformers (yet): The key bottleneck on in-context retrieval. arXiv preprint arXiv:2402.185100 2024.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu. and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint 1oX.::2310008391 2023.
13

---

Published as a conference paper at ICLR 2025
Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, and Liwei Wang. Do efficient transformers really save computation? arXiv preprint 2rXfv:2402.13934. 2024.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprint 7xXi:22306.09927, 2023a.
Yihua Zhang, Hongkang Li, Yuguang Yao, Aochuan Chen, Shuai Zhang, Pin- Yu Chen, Meng Wang, and Sijia Liu. Visual prompting reimagined: The power of activation prompts, 2024. URL
net nett:/oorummlddo03888
Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. Automatic chain of thought prompting in large language models. In The Eleventh International Conference on Learning Representations. 2023b.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoe V Le. et al. Least-to-most prompting enables complex reasoning in large language models. In The Eleventh International Conference Learning Representations, 2023.
14

---

Published as a conference paper at ICLR 2025
APPENDIX
A EXPERIMENTS ON REAL- WORLD DATA
We consider a simple arithmetic task that outputs ((Aio) A2 (22)))2))))) given A,, A2, As. Ai chosen from integers from to 9 as the input, where 01, 02, E = -+, x]. The CoT output follows the format of A = $1, A1 = S2. S203A4 = S3 and will be evaluated by whether all the three steps correct for the query (7). ICL directly outputs Ss, and the performance is evaluated by the prediction accuracy of Sq as
9 In the following experimental settings, the accuracy computed 50 prompts. Each prompt contains three context examples. The inference model is GPT-4 OOpenAl (023).
An increasing number of erroneous examples hurts the CoT generalization. To model the errors in the context examples in the testing prompt, we replace 0g with one operation 63 from Ov in the presentation of some of the context examples in the testing prompt. Note that the output values S3 are still correctly computed from S3 S203A4- Table I shows that when the total number of testing examples is fixed be three, with the increasing number of incorrect examples, the testing accuracy decreases. This is consistent with Remark [ Theorem 2
# of incorrect examples CoT accuracy
100%
100%
56%
0%
Table 1: The accuracy with different numbers of incorrect examples for Co'.. Errors in presenting 03-
CoT is more robust to erroneous examples with implementation error than ICL. In this setting, the error in a context examples is introduced by replacing o, with one operation o randomly and independently selected from 0) Hence, 8₁ A2, and the successive computation are based on the wrongly computed S The results in Table shows that when two incorrect examples exist, CoT performs better than ICL, which justifies Remark for Theorem
# incorrect examples
Co'l accuracy
100%
100%
100%
ICL accuracy
100%
100%
60%
Table 2: The accuracy with different numbers of incorrect examples for CoT and ICL. Errors in implementing 01.
B
ADDITIONAL DISCUSSIONS
B.1 THE MOTIVATION TO STUDY ONE-LAYER SINGLE-HEAD TRANSFORMERS
The reasons we study one-layer single-head attention-only nonlinear Transformers in this work are as follows.
First, it is much more challenging to theoretically analyze the training dynamics and generalization of layti-layyrrhead Transformers. This because the loss landscape for multi-layern laytiil Transformers is highly nonlinear and non-convex due to the interactions between multiple nonlinear functions. The simplified data helps to characterize the gradient updates in different directions for different patterns and steps. Non-orthogonal data make the updates less separable for different inputs, which is more challenging to analyze.
Second, the state-of-the-art theoretical works (Li et al. 2023a 2024a Huang al. 2023 Makkuva
et al. 2024a| Ildiz et al. 2024) on optimization and generalization also focus mainly o on one- one-layer Transformers. No existing works study the optimization and generalization of CoT even for one-layer Transformers. Therefore, we plan focus on the one-layer analysis to obtain more theoretical insights. We leave the theoretical analysis of the multi-layer case as future works.
15

---

Published as a conference paper at ICLR 2025
Third, although we admit the gap between theory and practice, theory still makes contributions under our settings. Our work is the first one to investigate the optimization and generalization of CoT and characterize the conditions when CoT is better than ICL. We establish the required number of context examples for a successful CoT in terms of how informative and erroneous the prompt is.
We also implement experiments the attention mechanism for three-layer two-head Transformers on two-step reasoning tasks. Please see Figure for details. The findings of all three layers are generally consistent with Proposition [ for the single-layer single-head case, which indicates that the CoT mechanism we characterize on one-layer Transformers can be extended to multi-layer multi-head Transformers.
B.2 THE MOTIVATION OF THE DATA AND TASK FORMULATION
There are several reasons for using such data formulation.
First, our data formulation of orthogonal patterns, on which the function is based, is widely used in the sfate-of-the-as theoretical study of model training or ICL on language and sequential data [[Tian) Huang et al., 2023; Li et al. 2024a] Chen et al. 2024). For example, (Huang al. Li et al. 2024a) study ICL on regression or classification tasks, which also use orthogonal patterns as data. Sections 2.1 and 2.2 in (Chen et al. 2024) consider learning n-gram data in ICL by
al.
2023a
2023
formulating transitions between orthogonal patterns. Section 3 of (Tian et al., (203a)) also assume orthogonal patterns in Transformer model training, and the generation comes from the orthogonal pattern set. The data formulation we use is consistent with the existing theoretical works.
Second, based on this formulation, one can characterize the gradient updates in different directions for different patterns and steps. This enables us to distinguish the impact of different patterns and steps in the convergence analysis of CoT using Transformers. Non-orthogonal data make the model updates less separable for different inputs, which is more challenging to analyze. Moreover, we would like to mention that during the inference, the tokens in testing prompts contain noises as defined in Equation 10. This makes the tokens of different TSR patterns not orthogonal to each other and relaxes our orthogonality condition to some degree.
B.3 THE DISCUSSION OF POSITIONAL ENCODING
The positional encoding (PE) we use is simplified for theoretical analysis. The formulation of PE we use is motivated by (Huang et al. 2024 Nichani et al., 2024), where each token is added with a PE represented by orthogonal vectors. These works formulate the distribution of the PE to be related to the structure of the data, such as patch-wise association (Huang et al., 2024). and sparse token selection (Nichani et al., 20224. Likewise, we follow their intuition make the PE vary in different steps of our reasoning tasks so that the Transformer can distinguish different steps when making inferences for the query.
Our analysis can be extended to study more general PEs with additional technical work in the future. One possible direction is studying the family of periodic and separable PE. For example, the absolute PE proposed by (Vaswani et al. 2017) considers PE as a sinusoid, which is periodic. Such analysis can be made by relaxing the "orthogonality" of PE vectors certain "separability" between PE vectors.
We also conduct experiments on three-layer single-head Transformer with the standard PE proposed in Section 3.5 of (Vaswani et al. 2017) for our problem. Figure shows that the blue curve increases to be the largest along the training, which means the attention weights on example steps that share the same TSR pattern and the same step as the query. This indicates that the CoT mechanism of using standard PE is the same as the one proposed in Proposition in our paper. One might note that the scores of the blue curve are not as high as Figure 6 in our paper. We guess the reason why the distinction in attention values is more significant in our PE may be the additional orthogonality of our PE and the property that its period is the same the reasoning length. Nevertheless. the strong similarity between the results on standard PE and our used PE shows the practical significance of our analysis.
16

---

Published as a conference paper at ICLR 2025
1.0
1.0
1.0
Suree 124, cif 20tp cff 1Se, cift chep
Sucte 126, cif 201p
12A, cif xtpp
Attemtion wsight. in aive
0.8
Attention weight In ave
0.8
130, cTf chi
Attention weight In aie
0.8
138, cift chep
0.0-
0.0
0.4
0.4
0.4
0.7
0.7
0.2
0.0
0.0
0.0
600
1200 1800 2400
600
1200 1800 3400
3000
600
1200 1800 2400 2400 3000 Iteretions
Iter sttretions
Itereti steretions
(A)
(C)
Figure 8: CoT mechanism with standard PE of (A) Layer 1 (B) Layer 2 (C) Layer 3.
C
ALGORITTMS
We first present the training algorithm introduced in Section 2.2
Algorithm Training with Stochastic Gradient Descent (SGD)
1: Hyperparameters: The step size 7j. the number of iterations T, batch size B.
2: Initialization: Let W W We and Wv (0ax xdx Iax Xdxxxd// ). Each entry of ( is generated from (00, <²) for a small constant { > 0. Wy and a are fixed during the training.
K
3: Training by SGD: For each iteration, we independently sample Qquery D, f Tir to form
batch of training prompt and labels (P", 22") IEE,, as introduced in Section 3.2
Each TRR
pattern is sampled equally likely in each batch. For each t 0,1, _T-1
(++)) = ( 7 B Σ
Vwool((0) ,",,""' x²
(24)
BEB,
4: Output: (T).
We then summarize the algorithm of the CoT inference introduced in Section 2.3 as follows.
Algorithm 2 Inference with Chain-of-Thought (CoT)
1: Input: 2g = 10 Bguery. Po. and Pt.
2: for k 1.... K- 1, do
Compute 22 by greedy decoding in (5. Then update Pa and Pa+1 by 6
(25)
3: end for
4: Output: 01, U2, "* UK- and UK by 5
D PRELIMINARIES
We first summarize the notations we use in this paper in Table
Lemma (Multiplicativee Chernoff bounds, Theorem D.4 of (Mohri et al. 2018 Let X1, Xm
be independent random variables drawn according to some distribution D with mean p and support included in [0, 1). Then, for any 7 E [0, - -1], the following inequality holds for ő 1 Σm X
mpn?
Pr(p 2 (1 + 7))) S
(26)
Pr(p < (1 7)p) Se-
(27)
( (Vershynin 2010 We say X is a Lab-Gaussian random variable with sub- Sab-Gaussian norm K > if (E]X K for for all p 2 1. In addition, the sab-Gaussian norm of X, denoted || is defined as || 102 supp21 (B)
Definition 2
Lemma 2 (Vershynin (010)) Proposition 5.1, Hoeffding's inequality). Let X1,X2, Xx be
independent centered sub- gaussian random variables, and let K = max, [| Then for every
a = (a), ,aN) E RN and every / 2 O, we have
N
ot2
Pr Σ
a; X 2t
S ee eeeexp
(28)
17

---

Published as a conference paper at ICLR 2025
Table 3: Summary of Notations
Notations
Annotation
Wi, yi Equerrys Zk
By is the input to the first step of a reasoning example. Vi.s is the k-th step output label of æ;;- Eqquery is the query input. 2/ the k-th step output label of Equuer- k E [K].
P. Pquerys Es, Qs. Uk
is training or testing prompt that consists of multiple training or testing examples and query. The last column of P is denoted by Piueer: which the query of P. E; is the i-th context example of
P. Qx is the first steps of the reasoning query. & E [K]. ty the k-th step generation by Co'T. & [K].
Cs. Pi, Pquery
C; is the positional encoding for the i-th column of the input sequence.
p; Ps Ct, where p; is the i-th column of P. Pquery is the p, of
the query column.
P"",""""""
F( v PT) the Transformer output for P with V the parameter. {(V; , ) is the loss function value given p and the correspond-
ing label z".
,4 M. E (SR()
M4 is the i-th training-relevant (TRR) pattern for i E [M] 14, is the i-th testing-relevant (TSR) pattern for E [M'. .M and M' are the
set of TRR and TSR patterns, respectively. TSR(-) is function that outputs the index of the TSR pattern of the noisy input.
Ju.
f is the task function with f fk f2 for -teps reasoning.
St is the k-th step task function.
T. T.
7 is the distribution of training tasks, while 7' is the distribution of testing tasks. D is the training data distribution. DD' is the testing data distribution.
a, a
a (or a') is the fraction of context examples with input sharing the same TRR (or TSR) pattern as the query.
A/e, B
A the step-wise transition matrix the k-th step for the task k E [K. B. is the Ktsteps transition matrix of the task f-
7, To To, p', .
+ is the min-max trajectory transition probability for task f. To is the min-max input-label transition probability for task f. of and p are primacy of the step-wise transition matrices and the K K-steps transition matrix, respectively.
The index set of context columns of the prompt that correspond to the k-th step of the example and share the same TSR pattern in the (k - 1)-th output as the (k- 1)-th output - of the query.
(t)
Pn is the summation of attention weights on context columns that share the same TRR/TSR pattern and in the same step the query.
B.
B, is the SGD batch at the b-th iteration.
ler
lor is the universal number training context examples. 4s is the number of testing context examples of the task f.
00, e()
We follow the convention that f(z) O(g(z)) (or Mg(+)),
(gg()))) means that f() increases at most, at least, or in the order of (x)), respectively.
2.5
f(x) Z g(z) (or f(x) S g(x)) means that f(x) 2 Mgg(z) (or
f(x) 5 O(g(z))).
where c >> is an absolute constant. Definition 3. Define that for ps that shares the same TRR/TSR pattern and in the same step as the query,
(t. = Σ softmax(p; ( Pquery):
(29)
Lemma 3. Given the SGD training scheme described in Section 2.2
B Q(M log M), and
4, 2 2(a-1), we have the following results. When - -23 log
2 l 12 for any p
18

---

Published as a conference paper at ICLR 2025
column of context examples in 1. we have
1
Øe(V; P" =")
p
B
OW
nEB
11 MM
a²
a²
2(K-
-1)
a²
(1- (ma)) +
K3 (1 - pn (1))))
k
K2
K
BEB
(30)
For any ÿ¹ that shares the same TRR pattern and a different positional encoding P, we have
a²
B
Σ
KM
- (3K 2)(1 Pm (1)(1+
K7
+ K3 (1 (n))))
REB
ØR(V; P", 2z)
p
(W p
R
a² (-4 - (3K - -2)(1 Pm (1)) + )p,(1)11 ())()())((())))))))) + PPr(()(1
1 Pn))
R
B
KM
BBB
(1+ K2))
(31)
For any d' that shares a different TRR pattern but the same positional encoding as p. we have
K2 + (K 11+ (2K
1)a²
a²
(x(())(
(P)())))(())----)))))))) Pm (t))
B BB
KM-
K2
K
a²
(1
- Pm (t)) Pm (t) (1--x((t)
K2
Sp'
B
1
Σ
0 PD,
( P
REB
(2K - 1)a²
a²
S1
Σ
+ (K - 1+
(pn(t))(1 (t))) - - (1))
B
KM
K2
K2
K3d
BEB
(32)
For any pl that shares a different TRR pattern and a different positional encoding from p. we have
(t) (1)) (2
K)aa
a²
7 B
Σ
PM
+(1 - Pt))))
KM
K2
K3
nEB,
PM .x")
OW())
(2 K)a²
) (z) + a²
Σ Sn B n
KM Pa (t)(1 (t)) - K+
k2)
K²
K
(33)
Lemma Given the SGD training scheme described in Section 2.2 4r 2 11)=1), and
2 2MM M), and
K 1 12 Th = l 2 log
(34)
we have that if Pquery is in the k-th step,
softmax(p, ((1) (1)))
Pquery) Se
(35)
NESTNTK
where SpRyy means the index set of context columns that are not in the k-th step.
Lemma 5. Given the SGD training scheme described in Section 2.2 B Q(M log M), and
4, 2 110-1), have the following results. When t 2 T " l 3 log for any p as
K
19

---

Published as a conference paper at ICLR 2025
column of context examples in 1. have V P", =x)
p"n 1l B nEBs
Σ
(1 - (o,(t))(xxx)))) (t))
OW
2MB nEB,
(36)
For any D' that shares the same TRR pattern and different positional encoding D, we have
T
Σ
P",";
OW x")
P S mJc
(37)
B
REB,
For any P' that shares different TRR pattern but the same positional encoding as P. have
T p'
ABB B
Σ
:E(V;
OW
2BM
Σ
(.( (P)())
(38)
REB,
NEB
For any P' that shares different TRR pattern and different positional encoding from d. we have
1
de(V; P ,",,"")) =²)
p'
OW P S npc.
(39)
B nEB,
E PROOF OF MAIN THEOREMS
E.1 PROOF OF THEOREM I
Proof. By the condition Lemma 3 we have that
2 Q(M log M).
(40)
We know that there exists gradient noise caused by imbalanced TRR patterns in each batch. Then, by Hoeffding's 's inequality (28
0 P", =")
OE(V; P", =²")
[ Pm 2²²)
Pr
E
2 E
(
B3
OW
OW
OW
(41)
NEB
Se e--B Be²
< A M-C
if B 2 E 2 log M. Therefore, require
B B max M) log M.
(42)
By Lemma and Definition 3 for p that share the same TRR pattern and the same positional encoding of Pquerg
(t. + 1) softmax |ST|
1
W ((+
p Pineery 2
(43)
+ (1 k) - ( - -ee
where by [66
Pm (b)) pm (b).
(44)
KM
For p that only share the same positional encoding of Pquery:
softmax(pp p (t+ Pyuery) 2 ee" + (1 1)) 6 x)
(45)
Therefore, make the attention weights between P Pquery and DÝ,' that share the same TRR pattern and the same positional encoding dominant, we need large enough u. When 1 (b) > 2(11,, we have
1 < T2 KMa
(46)
20

---

Published as a conference paper at ICLR 2025
When (b) z 011),
4
Pa (1+
K
(47)
e0
+
K
K
and
1 K
1
1 - pa(t+
(48)
2 e"
(1
Then, we prove that when t is large enough, u(t)= 2 log that the conclusion holds when to. then
n(1-a)tt We show it by induction. Suppose a²KM
to
u(t + 1 2 KM 2 (1 Pm. (m,(b))Pn(()) + 7
(1 - pP,(t))pp(())
KM
b=0
log
(K - a)²t
1)
(1 Pr. (t))
(49)
2
2a KM
KM
n(K () + 1)
log
a² L2KM
where the last step is by
1 log(1 (1
7)
1
7 KM
K
log
2
aM
(50)
To make (1 Pr (t)) < e, we need
K
a
2
2u c.
(51)
Then, we get
1
K - a
>
log
+ log
(52)
2
Therefore, by
1
K
1
K
log
+ log
2
log
log
(53)
2
KM
2
a
we finally obtain
2 T3
KM. K
(54)
For p`¹ that shares the same TSR pattern as the query, we have that when = Ti.
D ((1)) Pyuery 2 log
(55)
When / = Ti + T2 Ta.
P ( (t) pm Querry 2 O11) log 1 = O(log )
(56)
Then,
T := + +T2 T2 T
(57)
-1an-aa 2 log + MM(o +c))))
Therefore,
D.fE (, P, < O(e).
(58)
17
21

---

Published as a conference paper at ICLR 2025
E.2 PROOF OF THEOREM 2
Proof. We know that a' is the fraction of examples that share the same TSR pattern as the query. We need that in each step, the number of examples that share the same TSR pattern as the current step of the query is at least 1. Note that the probability of examples where each reasoning step produces the most probable output
II A (oN_:///iio--(((()1))sss/s 1=1
where fo(4)) == p4, Vi [M'],
(59)
where the input to the first step has the TSR pattern μ4g Define ()(a) as the TSR pattern in the k-th step output of the i-th context example by the transition matrix defined In Consider that the TSR pattern of the k-th step label of the testing query is 1 which is also the most probable k-th step output of the k-th step of certain with (2,) = "SR((1qurr)) = 90- Let the TSR pattern of another reasoning process, where for certain first-step input æ; with TSR az) TSR (Tguery) = qo- the k-th step output is the most probable for k E [K*]] (h), while the h-th step output is the second probable. Denote the TSR pattern of the k-th step output of æ; following this process as phu with uo qo. By the Chernoff bound of Bernoulli distribution in Lemma 1
we can obtain
le (MRC)))
K
Pr
Mo Vk E [K'] S (1 /2
A
As(ge-119u)
4s
(60)
P IT AbcuE-1 M-
Se
and by Lemma
les
K'
Pr
1 1mx()) 1/,, AK E [K']] (1 p2/2)aa II
A A 1.9
1ts i=1
1 4r lts i=1
K'
(61)
< Pr
1mmk()) ul. vK [K']] a' II
(uu + to 1=
A
Se -eal baly = -
for some c E (0, 1) and > O, where the first step is by the definition of pf in 11) and
K'
to S af k=1
II kigu--1-90)
A
(62)
Hence, with high probability,
K'
K'
les Z maxt(p! a II
A Abtan- - 'log log M, (o.a' II
A Ag
- log M}
kml
(63)
K
Zlofa'
Alte--1-g2) ) log M, 1=
II
AL
such that number of examples with the same TSR pattern as the query in each of the total K steps is at least 1. To make the above condition hold for any TSR pattern of the intermediate step of the query, we need
K'
lss 2 max (pfa' A
II
A
k(ge s.qu))
-log M
R=1
(64)
max (p/al HE|M
II
Al
ACTSR(A 10. (on 2og M
=(pfa'T!) 32 log M.
22

---

Published as a conference paper at ICLR 2025
Then, show the CoT testing error is zero by induction. In the first step, consider æ; [4 + ō; such that
Hi
6:
p. =
yi.1 [00
+
+ C2 mod K
(65)
Since that
(6 0 W
Pi S E,
(66)
by that each entry of W()) follows 5²) and
DE P". x²") OW(b)
8,, ,
Σ
Pquery 0,
(67)
B
REB. 6=0
we have that for p. that shares the same TSR pattern as the query,
P; W Pquery
Vl(V; P P=,,"))) ,=") @W
-p ( (0) 7
Pquery
(68)
BEB
b=0
=((rl) T ;v,,)) cf mod K (W
1
Ve( pm, p" OW
Pquery-
B
nEB 0=0
Since that X, orthogonal to all the p.ii E [M], we have similar conclusion for X, as d, i.e.
(2] 0 W
P. SE
(69)
and
V P 2n) 2") OW(
(1
7
Pquery = 0.
(70)
B
NEB b=U
ー Let 10/ = = + jlg p >3 + Eis M'
kg, Then, we have
(T P
p.
M'
2n) OW
= =(2 + Σ
u T 1. + cJ mnd x (W) +
[(2]
B
i=1
nEB,
b=
M'
+
Σ B ,0 C1)
M'
(HT . + c mod (W +
"
OE(V; P", ((H7,00) + +C)
B
dW (b)
nEB, h=0
de(V; Pm, 2z) OW())
Σ
Kasshig (uT, s0TT1 + cf mod (W + Σ 44
)(uT, 0-) CT)
nEB, 6=0
20. O(log 1) O(E) ((E)
=0(log )
(71)
where the second to last step is by Theorem The last step holds if C 2 Ollo (1/o).
Since the gradient updates for different TRR patterns are very close to each other, we have that Eigo (i < 1 and
00(V;P",, OW(
Ky.ks, ((NT v.TI) + c. mod x) (W iA1
+ Σ
)(H,)) +C)
REB, 6=0
Ps W Pquery
SE(1).
log
(72)
23

---

Published as a conference paper at ICLR 2025
where Pa shares the same TSR pattern and the same step as Pquery. Hence, for p; that shares different TSR pattern with Pquery.
pi W Pquery 2 ((1).
(73)
Therefore, we can derive that
softmax W Pquery) 2
(74)
where Sj is the set of the first step of examples that share the same TSR pattern as the query. Then, the first step leads correct prediction with zero testing error, since that maxje(Mr) AR(1o)) the largest to make the correct prediction for Qquery Equery ug i.e.,
11 ).
(75)
Suppose that the k-th step generates a zero testing error. Then, for the k + 1-th step, we know that there exists Pj that shares the same TSR pattern as Then, we can also derive that
Pj W ( ) + tT = = 1 E -
(76)
and
Σ softmax W ( (1]))0]) vT + JESt
1
(77)
Hence, the k + 1-th also makes the correct prediction, i.e.,
Uk+1 = fx+1 ..... fi (1(
(78)
where P/44 is the TSR pattern of the k + 1-th step input. Therefore, we show that CoT makes the correct prediction in each step as well as the final prediction, such that
R BCOTTEM 0.
(79)
E.3 PROOF THEOREM
Proof. We know that the positional encodings are the same for the ICL inference in all examples.
Hence, similar (74) we can derive that
softmax (p: T (T) Pquery) 2 1 -. C.
(80)
REST
where Si is the set of the last step output of examples that share the same TSR pattern as the last step output of the query. For Equery pg 9 E [K*), we know that the distribution of the corresponding label y of æ with TSR(x) = q follows the 9-th row the K-steps transition matrix Bf. Let F(V; ; P) DM' Xf" u1.. Hence, based the output scheme of ICL stated in Section 2.3 we have that
arg min 1
[FD; y aarg maxec XP- DEM' 2
(81)
Note that the probability of examples with the most probable final output with uy as the TSR pattern of the input is
B [(g,TRR(((r)))
(82)
To ensure that the number of examples with the same TSR pattern as the query that generates the most probable output is at least 1, we compute the following.
1
Pr
14,1 (1 P.. 22//2a ( B(o,TSR(((EE) 1m
115
(83)
1=1
-lup! 2tapf2at
Be M -C
<e
24

---

Published as a conference paper at ICLR 2025
for some (0, 1) and C 0 by the Chernoff bound of Bernoulli distribution in Lemma | Here, my defined as the TSR pattern in the final output of the i-th context example by the K ststeps transition matrix defined in [2] The TSR pattern of the most probable output of the testing query is 14 Similarly, let the TSR pattern of the second most probable output of the testing query be u9 We also have
Pr
1 mm = M4 2 (1 - p2./2///'' 2 (9.91)
(ea
(84)
Pr
1 1m, -1 a ((a,,)) + pe B (9.91)
i=1
lung ooa ((g M
Se
by Lemma 2 and 12] for some constant > O. Therefore, to make the number of examples with the same TSR pattern in the output as the label of the query be least for any TSR pattern the query and the output be the most probable one, need
1 ž max a' min ((ITTR(/(UN)) log M, (66o' min B(ITTRC((EE)) 22 log M] (E(M]]
(E[M]
(8.5)
=(pga'T!)~ log M).
addition, if Condition I holds such that the most probable output is the actual label, we can derive
SECT = 0.
(86)
When (85) holds but Condition does not, we know that ICL still always produces the most probable output by the K-steps transition matrix, but such an output not the label since Condition | fails. Hence,
RECL (V) 2 2(1).
(87)
When both Condition D and (85) do not hold, ICL can produce multiple possible outputs with non-trivial probability, which is decided by the distribution of the prompt instead the K K-steps transition matrix. This can be seen from that (83 and (84) both do not hold since (85) fails. Then, ICL can produce both the most probable and the second most probable output with constant probability. Let the TSR pattern of the r-th most probable output of the testing query be pt, Recall that F) Ei
M² XpP p/. we then have that for some small >> 0,
Hi [1/11 Vi μL, in P)
(88)
1/1
Then, by (8T. the output of the query is parg gax! A.- Since that (85) does not hold, there exists at least a constant probability of the prompt P' with the same query as such that
Xp' = l1 E [1 y: = n4, in P'H + + X.
(89)
1/2
for E [M']. Therefore, with constant probability, the output for the same testing query and the same testing task f varies. This leads to
R (V) ((1).
(90)
E.4 PROOF OF PROPOSITION
Proof. This proposition is derived from the proof of Theorem comes from 78). both by induction.
2 (22 comes from (77). while 23
25

---

Published as a conference paper at ICLR 2025
F PROOF OF LEMMAS
F.1 PROOF OF LEMMA 3
Proof.
VE(V; P ,²
B
OW
DEB
V P ) OF P)
Σ
1|
OF(V;
OW
(91)
= B
Σ
(P(V; P) -z")
W Wyp.softmax(p. T W Wiquery)
BEBs
(
Σ
softmax(p,
W P query
When 0, know that each entry of W is generated from the Gaussian distribution, (N,0 Then,
P
W (0)
Pqueryl =
Σ
Piuurry
Why SE.
(92)
RJ
Hence,
e O(E)
1
softmax (p; T (0) Pauery) 2
9(E) e O(E))
(93)
1 B(g)
[
softmax (0) Pquery) 5
(())
(94)
1. ((E)
We can obtain
A(E)
F(D:P)=
WvP:
(95)
Since that PE(-), and TRR(-) denote the positional encoding, and the TSR pattern of the input, respectively, we have that for p.
P Pquery 1 (T) TRR + 1[PE(p) = PE
(96)
Given lab(-) is the label embedding of the context as the input, we have that for p.
p p. = [[RR( = (P)) + 1 (lab(p) = lab(p.)) + (PE(p) = PB(p:)),
(97)
(Wvp) Wvp, llabbpp) = lab(p.)).
(98)
When 2 1, we first consider the case where p shares the same TRR pattern and the positional encoding Pquery. If and Pquery share the same TRR pattern, label pattern, and the positional encoding,
p
(p.
Σ softmax(p, WWPaery Pquery D Puerryp 22 (3 - 3pm(1) - (1 Pr)()))
(99)
= 4(1 - (pp()))t)).
and
p (p:
Σ
softmax(p,
Pee P 2. 3pn (t)) = 6(1 Pm. (t)).
(100)
When p and Pquery only share the same positional encoding or the same TRR pattern,
2 (1) 2 pT (Ø: Σ
softmax
W Pquery D Puery 2 (t).(t).
(101)
26

---

Published as a conference paper at ICLR 2025
When p and Pquery share both different positional encodings and TRR patterns,
(6pn(1) 2 P (p: softmax(p, W Powery)) Pr Preury p 2 -2 (t)(
(102)
Then, we consider the case where p only shares the same TRR pattern or the same positional encoding as p.. If p and Pquery share the same TRR pattern, label pattern, and the positional encoding,
3 - (t) 2 p (p.
T softmax(p, W Wiquery (eneery)prree Pqurry 21 (3--33-2) Pt. (1- 111-pp (t)))
22
(103)
When P and Pquery only share the same positional encoding or the same TRR pattern,
- (t())) 2 pT (p.
softmax (p. T W Wqqeery) ) Pueryy 2 0.
(104)
When p and Pquery only share both different positional encodings and TRR patterns.
PP (t) 2 2TT (P. Σ softmax Pr W P Pouery Pr
> Prueryp 2 -1.
(105)
Note that -(1 - Pm P..))) + (1 (/) < for (t) E Ja/K,ai a]. Then, when I Maa 1) and ý shares the same TRR pattern and the positional encoding pi.
softmax(p," T W Waquery) WvD: z²))
softmax(p: W Wiueery) Wvp:
p (p: Σ softmax p,," W Wqueeyy) Pr. P
5 4px - pn))) (1) (t)) a²
K2
+ 1 1 a
1 _
K マ (-4,
7 K )1
(1 Pm --2)))(22-)) 1)
a²
2
+ (1 - a- 1))
1) - (K + ())x,() +
1)).
=
K2
IK
(106)
We next consider the case where p shares the same TRR pattern and the different positional encoding as Pquery. Note that
2
a² )1 (x)(,)), (107)
KI (1 a)) K((KKII (t)))) 11(-(1 (1)p.. (1 Pn (t))
K2
if 2 M(a -). Then,
softmax (p, W Pyuerr)) Wyb - xn)
softmax (p: W W W
1=1
pT (p: Σ
softmax (p. T W Pquerr Pr) PP
(108)
5- 0 0..PP. (1) - (1)) + (1-pp(()) (1)) a
+ 1 1 (-(K-11)
(+2)
7 K R
R2
a²
K-1
-2(1 - pb.
K2
1 K
K
27

---

Published as a conference paper at ICLR 2025
We next consider the case where p shares the same positional encoding and the different TRR pattern as Pouery. Then,
softmax(p: (p: W Pqwery) P - xn) T softmax (p: Wiquery) Wvp:
1=1
p (p:
softmax T
W
Pquery )r)
(109)
50 (1 p., (1))
a²
1
x
(-(K
K
K
K
=- - - (1-
a²
K
1
1
a K
K²
K
Therefore, as long as
1 la )
(110)
we have
dE( ; P 2²)
B
OW
BE
-1 Σ
(F(V;P) P) - ==
Wv so T W Pquery)
REB
=p' (p:
softmax
W Puuery )r)
a²
2(K 1)a²
$1
(1
Pn (n))t) (t)
+
KM
K2
K2
REB,
1
1
(- Pr (1)) a² k2))
K
M
a
a
2(K -1)
a² (1- (1))))
Σ BEB
(1 pm (-))) KM
-1
K2
K
K3
(111)
We then consider the case where ÿ' shares a different positional encoding and the same TRR pattern as P. Let p share the same TRR pattern and the positional encoding as Pquery. p' and p, share the same TRR pattern, label pattern, and the positional encoding,
2(3 (() 2 p' (p.
softmax(p, TW Pquery Pr P Pruery 22 (3 - pml() (3-pp(() - (1 - pp(())
-4.
(112)
When p' and Pquery only share the same positional encoding or the same TRR pattern,
2(1 (t) p' (p. Σ softmax(p) (p, "WPueryy)P D Puerry? 2 0.
(113)
When p and p. only share both different positional encodings and TRR patterns,
(t 21))5'' (p. Σ softmax(p," TW Pquerry) (r p P p 2 -2.
(114)
Then, we consider the case where p only shares the same TRR pattern as Pquery. If p' and p. share the same TRR pattern, label pattern, and the positional encoding,
3 - (t)) 2p' (p: - Σ softmax (p. W Wyquery )p, P r=1
(115)
21 (3 - -3p - -1-18 (1-xx(t)) = 2(1 - Pm (t)).
28

---

Published as a conference paper at ICLR 2025
When p' and p: only share the same positional encoding the same TRR pattern.
1 - Pm (t) p' (p.
softmax W " Pr Pawery 2 (t),(t).
(116)
When pú and p. only share both different positional encodings and TRR patterns,
Puerry Pt((t) 2 p' (p: Σ softmax(p, W Pquery ) ) query 2 -1 (t).
(117)
Next, we consider the case where p only shares the same positional encoding as Pauery: p and p. share the same TRR pattern, label pattern, and the positional encoding,
32 p' (p: Σ softmax pr. T W Pquery) Prr Pouery? p
2 1. 1.:1--300 - (P)( = 2 Pt(t (118)
When p and p. only share the same positional encoding or the same TRR pattern;
1 12 P: Σ softmax(p, (p, Pguery (r P 2 (t)...
(119)
When p' and p; only share both different positional encodings and TRR patterns,
0 p (p softmax(p, W Pquery Pr) Pr) Puerry 2 -1 Pr.(t).
(120)
Then, when / l(aa1) 1) and P shares the same TRR pattern and the positional encoding Pauer
softmax (p: W Pguerry Wvp: 2n)
softmax (p; W WPguery) Wvp,
p' (p.
(p, W D Pquery Prr Pr) Purrry
(121)
(t) (1-))())) + 1 -
We next consider the case where p shares the same TRR pattern and the different positional encoding as P.uery. Then, by [177.
softmax/ p T W P qquery) Wvp: - zn)
Σ softmax(p, W Pquery) WvP:
p' (p: softmax (p, W Pquery) ) )P))) Pugery
(122)
< - (1) Pn))))² (p) pm (t))
a²
1 1// x - (t)
1
K2
a K²
a) (1-
+ (1
(t))
We next consider the case where p shares the same positional encoding and the different TRR pattern as Pquery- Then,
softmax (p, W Wyquery) W - xn
softmax (p: W Pquery) Wyp.
p' (p. Σ
softmax(p, T W Pquery pr )r PTerryp
(123)
(Ph(1)(1 (1)) + p.(())1 pn (1)) a
1 11 -)(-112p..())
K2
a² 2
1 (1 - a) + 2p. (t)).
(P) - (1))
K?
29

---

Published as a conference paper at ICLR 2025
Therefore, as long as
12 ()a
(124)
we have
1
Σ
de(V; p"",""
p' "B B
p
OW
nEBs
1 Σ
=I
(F(V: - 2n)
Wv Wypisoftmax W D (p:
BEB
softmax(p. W Pquery) P
1
1
(-4 1-4-2(( 1) a² (p)(((( PP)(t)
(125)
S1 B
KM
EEB
a² - Pm -P,(t)(-pp(())(( x2))
+ K
M
1
Σ
(MM( (-4 (3K - 2) (I()))(1+ + a²
K2 )pnn((() pn)())
= B
NEB,
a
(0((11 - (n)) R2 ))
and
V P", =z)
p
p
OW
BEB
1
(-4 (3K ()(())((t))11 + a²
K² )p,
27 B
Pt)(t)
KM
nEB,
Pm (1) - pr (1) + a² 1
+ (1 (1-pa (1- + (1 - (t) a²
+
K
K2) K
K2
=7
1
Σ REB,
KM 1 (-4 (-4- (3K 2)(1 - (px(t)( + L (111
- (m (t) +
a²
K3 (1
Pm 2))
B
(126)
We next consider the case where p' shares a different TRR pattern and the same positional encoding as P. Let share the same TRR pattern and the positional encoding as Pquery. If p' and p. share the same TRR pattern, label pattern, and positional encoding,
TW 2(3 (1)) 2 D' (p, Σ sof WWggerr)) Dr D Queryy 22. (3-Pm(1) - (1 - (p)()))
(127)
When p' and p; only share the same positional encoding or the same TRR pattern,
2(1 - Px)
(p.
softmax (p, Wiuery )rrr PD) Pqueryp 2 0.
(128)
When p' and p. only share both different positional encodings and TRR patterns,
22) 2 p. (p. Σ
softmax(p, "W Pquery PP)) Pueery p 2 -2.
(129)
Then, consider the case where only shares the same TRR pattern as Pquery- If D' and p. share the same TRR pattern, label pattern, and the positional encoding.
3 2 p (p. Σ softmax (p, T Waqeery) Dr, DD P D 1. - (1 Pm (t)) = 2 ++ Pt) (130)
30

---

Published as a conference paper at ICLR 2025
When p' and p; only share the same positional encoding the same TRR pattern.
12 p' (p:
softmax(p, T W Pqqerr (r) 2 (t).
(131)
When p' and p. only share both different positional encodings and TRR patterns,
02 028 P:
softmax(p, W Pquery )r, P.Pr)yuurrll 2 -1 Pra(t).
(132)
Next, we consider the case where p only shares the same positional encoding as Pauery. p' and p. share the same TRR pattern, label pattern, and the positional encoding.
3 (t) 2 p0 (p: Σ softmax(p, W Pquery (.. 1- (t) (1-Pm(t)) 2.
(133)
When p' and p; only share the same positional encoding the same TRR pattern,
1 (1( p (p:
softmax(pp. T Wiqeerylli Dr P 2 0.
(134)
When p' and p; only share both different positional encodings and TRR patterns,
(l) 2 p' T (p. Σ softmax (p, W yuuery (a )
Pueery 2 -1.
(135)
Then, when / > M(a and when shares the same TRR pattern and the positional encoding as Pouery. by (07).
softmax (p, W Pquery) W - x²) softmax W Pquery)
W WvP.
p' (p.
softmax(p, T
W Pquery ) )r) Puuery
(136)
<0- 2(1 (t))
1 R
a
)-2((K 1)).
K2
We next consider the case where P shares the same TRR pattern and the different positional encoding as P Pquery. Then,
softmax(p) (p: T W Pquery) Wvp. - - softmax TW P yowery W
(p: softmax(p. W Pyquer pp p
a
1
- C-P)())(1 Pn (t) Pm (t)) Pm. (t) Pm (t))
1 K
Pn(())
K2
(1) - (1))2 a²
+ 1) 7 (1 (-1++++ PP)())).
K²
(137)
31

---

Published as a conference paper at ICLR 2025
We next consider the case where p shares the same positional encoding and the different TRR pattern as Pouery. Then,
softmax(ps T
W Pqqeer)) P zn)
softmax (p: W Pquery) Wvp:
p' (p.
softmax(p, W
Pouery Prr Dr)
PqueryD
(138)
a
S (1 - P., (t))
- ++ 7 K * (-K
K2
a²
K -
= - (1- (t))
(1 a))
K2
Kl
Therefore, as long as
1 ()a
(139)
we have
(V: P", ,²")
p'
P
B
OW
BBB
=1
1
Σ
(F(T; -z")T
Σ
Wy pisommax
W Wygeey))
B
REB
p
(p:
softmax (p,
W Pquery (ineery)pr)) Puuery
1
a²
a² (t) - (1)
1
+ (K
1)(1+
B
KM
K2
K2
REB,
[1 ) - (t)) a²
M
K2
1
a²
(2K 1)a²
)p (1)(1 (n()))² (1 (P-))(()) a²
=1
(K 14
)p
RT)
KM
K2
K2
BEB.
(140)
and
1
Ø V
Pa, z))
p'
p
B
aw
D
21
Σ
a
+ (K 1
(2K 1)a²
(p,( - pn(())² - (1 - (1)) a²
B
KM-
K2
K2
K3
BEB,
a²
K (1 (1-P. (t) (-)))) + (1- (1-Pn(tt)
K2 ))
(141)
We next consider the case where p' shares a different TRR pattern and different positional encoding as P. Let share the same TRR pattern and the positional encoding Pquery. If p' and p, share the same TRR pattern, label pattern, and the positional encoding,
62 p (p. Σ softmax(p. W Pquery) (r Pr Puurry 2 -3- = -4 (t). (142)
r=1
When p' and p. only share the same positional encoding or the same TRR pattern,
2 2 p (p.
sofuma((p. T W Pyperry Dr DD Pueryy 2 (p.(!).
(143)
32

---

Published as a conference paper at ICLR 2025
When p' and p. only share both different positional encodings and TRR patterns,
0 p' (p.
softmax " Wyyeryy Pr) Puuery 2 - + (t).
(144)
Then, we consider the case where p only shares the same TRR pattern as Pquery- If p' and p. share the same TRR pattern, label pattern, and the positional encoding,
3 1 (p: Σ softmax(p, T W Pquery) pr Pquery? P 1 (3-3 (l-p((()) 2. (145)
When p' and p; only share the same positional encoding or the same TRR pattern,
1 12 p' (p.
Σ
softmax(p, T W Pquery Pr ( Pquery 2 O.
(146)
When p' and p. only share both different positional encodings and TRR patterns,
0 p' (p: Σ softmax(p, W Pyuery Pr ) pueeryp 2 -1.
(147)
Next, we consider the where p only shares the same positional encoding Pquery- If p' and p: share the same TRR pattern, label pattern, and the positional encoding.
32 p (p.
softmax (pr T W Pquery Pr P Pouery? 2 1. - (t)
(1)) = 2
(148)
When p' and p. only share the same positional encoding the same TRR pattern.
p
P: softmax T Pggeryy )
PP 2 Pm Pt.
(149)
When p' and p. only share both different positional encodings and TRR patterns,
0 p (p:
softmax(p, Wipeerypp query 2 -1 + (1).
(150)
Then, when 1 2 P(a 1), and when shares the same TRR pattern and the positional encoding as Pquery.
softmax(p. W Pqeery) Wvp: x²) =²)T²
Σ softmax T W Pguery) W
p'
(p:
softmax (p, W Pguuery )rr )P Tuerryp
a² K2 2pm(() + (1 a)(-2+2p.
S - Pm
2px(t)) - (t))
(1-p(tt))
(151)
We next consider the case where P shares the same TRR pattern and the different positional encoding as Pouery. Then,
softmax (p, W Wyuerry) W - xn) softmax (p. W Wyguery) Wvp:
(p. Σ
softmax(p, T W Pquery) Pr )r PTuery,
(152)
a (-1)
1 1
<0 10+p((1(11 (.((((1-pp())))
K2
7
7
K
(-K)
- (p)))) a
= =
+ (1-1))(11) a-(-
K2
33

---

Published as a conference paper at ICLR 2025
We next consider the case where p shares the same positional encoding and the different TRR pattern as P.uery. Then,
softmax(p, T W Piquer)) W WVD: zn)
softmax(p: W Wiquery) Wouery)WVp... WvP.
p' (p:
softmax p. W Pquery) ( Pupery
a² ())
5 - (1 -
P
7 K
(t))
K
a
+ 1 (1 a
= (1 (1)) +
K2
(153)
Therefore, as long as
12 1)
(154)
we have
VE(V; P=,"")
p
OW
BB
- Σ (F(V; - =²
Σ
Wv softmaaaa " Pquerg)
BEB
D (pi
softmax W Pqueryy Pr Pr T
(-Pm(t)(1 (x(())(-2++ + (3- K)(1 pn())) a²
K² Pn(()) (155)
$1
KM
1
1
(r)) (1)(1 + a² x²2)
+
(1
K
M
1
Σ
1
Pa (1) (x)))) - K (2 K)a2
=1 B
KM
K2
nEB,
+ (1 - Pm. (t)) Pz) (t)(1 + a² 1
k2) R)
and
1
8E(V; P" 2²)
p'
OW p
B
Pm (1) - (1) + (2-
K)a
+ (1 (1)) + a²
Σ
2n
B
KM
K2
R3
K
n
1 K
a²
(1 (-)) (1- p.))
(2)
=1 B Pu)((( - pn (1))
1
Σ
K ) 2
+ (1 - Pm (1)) a²
K2
R3
ncB.
(156)
F.2 PROOF OF LEMMA 4
Proof. We can derive that when 1 1 (t) 2 2(1). ( p increases if p and p' share the same positional encoding. Otherwise, p W(O decreases. We know that (!) > 2. Combining the
34

---

Published as a conference paper at ICLR 2025
results in Lemma 3 can derive that when t 2
w = w - 71 1
Σ
dE(V; V; ,",=")
(157)
@W()
neB,
Then, for p;" R that share the same TRR pattern and the same positional encoding of Pauerr
Pt +
1)
softmax (p:
W (1+1))
Pg Pquery)
S1
(158)
1
K
(K- 1)
K
c
* M( - 1)e + t
where
a²
a
(m())) = Σ (1 (b))2 2a²
$1 27
((1 Pn (b))
(159)
K3
K3
K3
b=0
6-0
21a²
82 2
Pn. (b))²
(160)
K3
b=
a²
+ a²
2(K - 1)
a
83 2 KM
(1 Pm (b) 4-4 (b)( +
K2
K
(1
+
K
+
K2
b=0
- (K- 1+ 2K- 2 Pm (b))
(161)
K2
a²
2K- -
(1 Pm (b)) (px)
(4+
)),
KM b=0
K2
K²
where the last step is by Kpn 2 4o 2/K2 when (b) 2 a/ K. For P. " that share the same TRR pattern and different positional encoding of Piuerr:
softmax T
(
P Pyuery)
1
(162)
a el (K-l)a
(1 K
a {)( -)ee 84 e0o)"
where
$4
2
17
((-4
(3K 2)(1 - (b) +
(),,(b(( - pr,(b))
M
K
b=
a²
(2 - KK)( +
K2 Pn (b) (b))²)
(163)
a²
17
(4+ 2K(1 Pm (b))(1
Pr pm (b)),
M
K²
1=0
2naa
85
(1 Pn (b))2
(164)
When M 2 2((+a= and t 2 s(n 3 log Ka -2),
(K - -0 + c0 > K.
(165)
M 2 11) and t < 0 K3 log Ka we cannot ensure
(K - 1)e + c0 K.
(166)
For pi that share different TRR pattern and the same positional encoding of P Pywery"
1 softmax(p) (++1) Pguery) 7
1
(167)
a e + K
e + C - ) [1+(KK
(K-KKK1) - 10e-
K
35

---

Published as a conference paper at ICLR 2025
where
20 Pp)
86
K3 (1
(168)
b=0
For P that share different TRR pattern and different positional encoding of Pquery
softmax (p: T
((+ en Pauery) = ee + (k x)(K - 1 0*0) + e
1
(169)
Note that when In K3 for Piueery in the k-th step, we have
softmax(p)" w(+1) Pouery) 2 21(1,
(170)
for P that share different positional encoding from Pauery To make the total softmax values on contexts that share different positional encoding and a different TRR pattern from the query smaller than e, we need
K $1 $2, 86 Z log
(171)
When further increases to larger than (1 a -2 2 K log ), we also have that the total softmax values contexts that share different positional encoding and the same TRR pattern from the query smaller than e. Therefore,
K
ス Ti := :
K log
(172)
n
F.3 PROOF OF LEMMA 5
Proof. We consider the case when Ti given Lemma When 1 2 9aa and when p shares the same TRR pattern and the positional encoding as Pquery.
softmax(p:
W Pqueery) WvD: x"))
Σ
softmax (p: W Pquery) Wrp:
(p: Σ softmax (p, T W Puuery )r) queryp
(173)
S 4p)(t)( (t)) +
$ - 4px (t))
We next consider the case where p shares the same TRR pattern and the different positional encoding as Pquery- Then,
softmax(p,, W ) Wyp: - 2n)²
Σ
softmax (p. W Wigerry) Wvp:
p (p: Σ softmax T W Pquery pr Puurry
(174)
$ - (1) - (t) + Sc.
We next consider the case where P shares the same positional encoding and the different TRR pattern as Pquery. Then,
softmax (p, W Wouerry W Wvp, - x²) Σ softmax (p, T Wygerry) Wvp.
(175)
T (p:
softmax (p, W Wyguery pr P
Se.
36

---

Published as a conference paper at ICLR 2025
Therefore,
V
P =n)
Σ
p 11 B
OW P
NEB,
Σ -
Wy softtatta T W Wyuerry)
1
P (p:
softmax
W Pquery prr PP Pqueryy.
(176)
-1
(pn(1)( - (n)) + +
M 6
2M
BEB,
40, - (.)
1 2M
B
nEB.
We then discuss if p and p' only share the same TRR pattern. When / 2 1)a="), and when p shares the same TRR pattern and the positional encoding as Pquery, obtain
softmax(pi (p: W Pquery) Wyp: - x²))
softmax p: W Wqueery) Wvp.
(177)
p' (p. Σ softmax(p, T W Pquery) ) Pr) T
Pqueryy
2-2(1 (())
We next consider the case where p shares the TRR pattern and the different positional encoding as Pquery. Then,
softmax (p, W Wgueery) W xn)
Σ softmax (p: W Pquery) Wvp:
(178)
p' (p. Σ softmax(p, p, W Pquery Pr Puprryy
- (1 Pa. (1)) Pm (t)))),,.tt)
We next consider the case where P shares the same positional encoding and the different TRR pattern as Pquery. Then,
spf W Pqwery) WVD. - =² Σ softmax(p, TW Piueery) Wvp:
(179)
p'
p
softmax(p, W Wguerry Dr P
Se.
Therefore,
D0(V;P". z")
p
11
OW
BEB,
1
Σ (F(D: 2n)
Wy p.stmmaxx W Pquery)p
B
(180)
nEBs
(p.
softmax(p, W Pquery pr (o PTurrppl
Snp..
37

---

Published as a conference paper at ICLR 2025
We next discuss when only shares the positional encoding as p'. When 2 M(aa and when p shares the same TRR pattern and the positional encoding as Pquery.
softmax (p; W Pguerr) Wvp:
2n))
Σ
softmax (p; W Pguery) Wvp:
(181)
p' (p.
softmax(p, W Pouery Pr (..))PPuuo Pauery
Se.
We next consider the case where p shares the same TRR pattern and the different positional encoding as P Pquery. Then,
softmax (p; W Piueery)W Wyp: 2n)
softmax W Pauery) WvP:
p (p. Σ softmax p.
W PqueeryPrr P
(182)
1
$ (L) (t)) Pm
M
-t)
We next consider the case where p shares the same positional encoding and different TRR pattern Pquery. Then,
softmax (p; W Pouery)) Wyp: 2²)
softmax (p: W Pquery) Wvp.
(183)
p'
softmax(p, T W Pquery) Pr )P Pquery
Se.
Therefore,
Øe(v; P²)
p' "1
Σ
OW
nEB
(F(V: - =²")
Wy W.p.softmaxpp. W WP PPqeery) P
B
n
(184)
(p:
(oftaxx(p, TW WygeeryPP P.)) P
2M Pb.( 2b))
BEB.
We then consider if p shares different TRR pattern and different positional encoding as ø'. When I 2 N(a ), and when p shares the same TRR pattern and the positional encoding as Pquery:
softmax(pp (p, W W Wyp: - x=) softmax (p, T Wggerry) Wvp.
(185)
(p. p'
softmax p, T W Pquery Pr P P
Ze.
38

---

Published as a conference paper at ICLR 2025
We next consider the case where p shares the same TRR pattern and the different positional encoding as Pquery. Then,
softmax(p:
W Pouery)) W x") T
softmax W Pquery) WvP:
(186)
p P
softmax(p, W Pyuery Prr Pr) Pqeryy
2 (1))...)))
We next consider the case where p shares the same positional encoding and the different TRR pattern as Pquery. Then,
softmax(p:
W Pquery) WvP. =²)²
Σ softmax T W Pquery) Wvp:
(187)
p'
Pp
softmax W Pouery) (r P
Sc.
Therefore,
VI(V; P 2n)
P
11
OW
nEB
1 Σ
(F(4; - =n) B
Wy po W "query)
(188)
nEB,
(p, Σ
softmax(p, T W WPquery) )r)) query.
Sme.
39


Reviews:
[{"review_id": "LiODM3JT8r", "reviewer": "Reviewer_LGXd", "summary": "This work presented a theoretical study on training Transformers to achieve Chain-of-Thought (CoT) reasoning capabilities. The authors explored the training dynamics of nonlinear Transformers and provided a quantitative analysis of the required training samples and iterations for developing CoT ability. They analyzed the model's success in generalizing CoT on new tasks with O.O.D. test data, even in noisy and inaccurate reasoning examples, furthermore, contrasted CoT with in-context learning (ICL) and theoretically characterizes scenarios where CoT outperforms ICL.", "questions": "1. For the CoT training set, this work constructed it with examples + query + (K-1)-reasoning steps, and the K-th step as the label. Why did not adopted more later steps as the label or more flexible label strategies?\n\n2. In the 3rd line of Example 1, should the last affinity function be ?\n\n3. How well do the theoretical results extend to more complex Transformer models with multiple layers and heads?\n\n4. Have the authors considered validating their theoretical findings with real-world datasets, and if so, what are the challenges and potential modifications needed?\n\n5. Could the authors elaborate on how the model's performance is affected by different levels of noise and inaccuracy in the context examples?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper provided a novel theoretical framework for understanding how a simplified Transformer can be trained to perform CoT reasoning, which is important for Transformers-based model reasoning.\n\n2. The authors offered a detailed quantitative analysis of the training process, including the number of context examples, training samples, and iterations required for successful CoT training. This level of detail helps in understanding the complexity and feasibility of training similar models.\n\n3. The paper addresses the generalization of CoT ability to new tasks with distribution shifts, which is a practical concern for real-world applications where training and testing data may differ.", "weaknesses": "1. The paper idealized the Transformer as a simplified one-layer single-head architecture. While such a simplification can aid in theoretical analysis, it may limit the generalizability of more complex Transformer architectures.\n\n2. The assumptions made about the data and tasks, such as the orthonormality of training-relevant and testing-relevant patterns, may not hold in all real-world scenarios, which could affect the applicability of the results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8a56B1uNrB", "reviewer": "Reviewer_y9x5", "summary": "The paper provides a theoretical characterization of three aspects of in-context learning with examples consisting of chain-of-thought using a simplified data model and a simplified single-layer transformer with non-linear attention. The central result of the paper is a quantitative characterization of the training dynamics that leads to CoT ability. Using this, the authors characterize how the noise and relevance of the in-context examples affect the CoT performance. The final result compares the error rates for ICL with CoT and ICL without CoT with the main conclusion that while the rates are similar, ICL without CoT needs the examples to be dominantly correct, while CoT does not need this. Experiments on synthetic data support the theoretical results.", "questions": "Are the results applicable, or can they be easily extended to the case where the reasoning steps $z_i, f(z_i)$ are not necessarily of the same sequence length? In other words, $z_1 \\in R^{1\\times d}$, $f(z_1) \\in \\mathbb R^{2\\times d}$, $z_2 \\in \\mathbb R^{2 \\times d}$, $f(z_2) \\in \\mathbb R^{1\\times d}$. Such a setting might be of greater interest because the steps are typically multi-token in applications involving natural language.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "**Originality and Significance**: While previous papers analyze ICL (with and without CoT), their assumptions do not match this paper's. Specifically, this paper does the analysis using a transformer with non-linear attention, with test data coming from a different distribution than the training data, making the setting more general and useful.\n\n**Quality and Clarity**: The theoretical results are presented clearly and are supported by relevant experiments. The discussion and remarks help with understanding the implications of the results.", "weaknesses": "It would increase the impact of the paper if the authors could include experiments on some real datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qClqz27INd", "reviewer": "Reviewer_txG6", "summary": "The paper provides a theoretical study on training Transformers with nonlinear attention to obtain CoT generalization capability, allowing the model to infer unseen tasks when the input is augmented with examples of the new task. They quantify the required training samples and iterations needed to train a Transformer model towards CoT ability and prove the success of its CoT generalization on unseen tasks with distribution-shifted testing data. The paper's major contributions include a quantitative analysis of training dynamics on a one-layer single-head attention-only Transformer, characterizing the conditions for accurate reasoning output by CoT even in the presence of noisy and inaccurate examples. It also provides a theoretical explanation for why CoT outperforms in-context learning (ICL) in certain cases, highlighting the robustness of CoT against noise in reasoning examples. Through experiments, the authors justify their theoretical findings, demonstrating the practical relevance of their theoretical analysis. This work advances the understanding of Transformer models' reasoning abilities and their potential applications in solving complex reasoning tasks.", "questions": "See the section on weakness. We will increase the score based on the answer to the question.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Theoretical Advancement: The paper makes a significant theoretical contribution by providing the first analysis of training dynamics for Transformers aimed at achieving chain-of-thought (CoT) generalization capability. This addresses a critical gap in the literature, as prior works have not fully explored how Transformers can be trained to generalize on multi-step reasoning tasks via CoT.\n﻿2. Deep Quantitative Analysis: The authors offer a detailed quantitative analysis of the training process, including the number of context examples required for successful CoT reasoning and the conditions under which CoT outperforms in-context learning (ICL). This level of detail is commendable and provides actionable insights for practitioners looking to implement CoT in Transformers.", "weaknesses": "1. **Limitations in Model Complexity**: The paper primarily analyzes a single-head attention Transformer, which may not encapsulate the full complexity and performance characteristics of multi-layer and multi-head attention models. The validation was confined to binary classification tasks, thereby restricting the generalizability of the theoretical findings.\n2. **Formula Accuracy**: The paper requires a meticulous review of its mathematical expressions. For instance, in Example 1, the function should be denoted as \\( f_2(\\mu_2') = \\mu_1' \\) instead of \\( f_1(\\mu_2') = \\mu_1' \\), and the second matrix \\( A_1^f \\) should be corrected to \\( A_2^f \\). It is essential to verify all formulas throughout the text to ensure accuracy.\n3. **Theorem Validity and Clarification**: The theorems presented in the article should be scrutinized for their validity, particularly the sections that substantiate reasoning, as they may induce some ambiguity. Reading the appendix is mandatory for a comprehensive understanding of the article; otherwise, it might lead to misconceptions.\n4. **Originality Concerns**: The article's reasoning and writing logic bear similarities to those found in \"How Do Nonlinear Transformers Learn and Generalize in In-Context Learning.\" It raises the question of whether this work is merely an extension of the previous study or if it introduces novel contributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yYTl4AjxbH", "reviewer": "Reviewer_6Hy5", "summary": "This paper presents a theoretical analysis of training Transformers with nonlinear attention mechanisms to develop Chain-of-Thought (CoT) reasoning abilities. The authors focus on quantifying the requirements (such as sample complexity and iterations) for training a Transformer model capable of CoT generalization on unseen tasks, even under conditions where examples might contain noise and inaccuracies. The study also compares CoT with In-Context Learning (ICL), showing that CoT can outperform ICL when intermediate reasoning steps are essential. The paper provides both theoretical insights and empirical evidence supporting the model's behavior during training and testing.", "questions": "No further issues.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "● This paper is one of the first to theoretically explore training dynamics of nonlinear Transformers with CoT capabilities, providing important insights into sample complexity, required iterations, and generalization guarantees.\n\n● The comparison between CoT and ICL is theoretically grounded, with a clear quantitative explanation of why CoT can succeed in contexts where ICL may fail.\n\n● This work addresses an understudied area in Transformer models, specifically how they generalize CoT reasoning abilities across unseen tasks, making it a relevant contribution to understanding advanced prompting methods in LLMs.", "weaknesses": "● The authors analyze attention-only models, aligning with previous work by [1] Olsson et al. (2022). Both studies observe that trained models’ attention values tend to focus on context examples with input patterns similar to the test query. While Olsson et al. attribute this to ICL (In-Context Learning) ability, the authors link it to CoT ability and explore how training may facilitate this behavior. A more detailed comparison of contributions between this study and Olsson et al. would be valuable. \n\n  ○ [1] In-context Learning and Induction Heads.  2022. Olsson et al.\n\n● The paper acknowledges the gap between the one-layer model studied and multi-layer Transformer models. As practical LLMs are complex, the findings might not directly extend to real-world Transformer architectures. Please discuss and analyze the applicability of these conclusions in more complex scenarios.\n\n● The theoretical results assume a linear span between training-relevant (TRR) and testing-relevant (TSR) patterns, limiting applicability to cases with greater diversity in test distributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)