From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency
Authors: Kaiyue Wen
Huaqing Zhang
Hongzhou Lin
Jingzhao Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: chain of thoughts
sample complexity
sparsity
sample efficiency
parity learning
Primary area: learning theory

Abstract: Chain-of-thought (CoT)  significantly enhances the reasoning performance of large language models (LLM). While current theoretical studies often attribute this improvement to increased expressiveness and computational capacity, we argue that expressiveness is not the primary limitation in the LLM regime, as current large models will fail on simple tasks. Using a parity-learning setup, we demonstrate that CoT can substantially improve sample efficiency even when the representation power is sufficient. Specifically, with CoT, a transformer can learn the function within polynomial samples, whereas without CoT, the required sample size is exponential. Additionally, we show that CoT simplifies the learning process by introducing sparse sequential dependencies among input tokens, and leads to a sparse and interpretable attention. We validate our theoretical analysis with both synthetic and real-world experiments, confirming that sparsity in attention layers is a key factor of the improvement induced by CoT.

Published as a conference paper at ICLR 2025
FROM SPARSE DEPENDENCE TO SPARSE ATTENTION: UNVEILING How CHAIN-OF- -FTTHOUGHT ENHANCES TRANSFORMER SAMPLE EFFICIENCY
Kaiyue Wen*
Huaqing Zhang'
Stanford University kaiyuew@stanford..ed
III,, Tsinghua University
2hanghg22/emill tsinghua. edu en
Hongzhou Lin
Jingzhao Zhang
Amazon
III, Tsinghua University
hongzhou. 1in8 1ennani@@mmlllmmmmmmmmmmm com
Shanghai AI Lab
Shanghai Qizhi Institute
jingzhao saiiaaz@@miii www.eemmaa...iiiiiieme tsinghua. edu. cn
ABSTRACT
Chain-of-thought (CoT) significantly enhances the reasoning performance of large language models (LLM). While current theoretical studies often attribute this im- provement to increased expressiveness and computational capacity, we argue that expressiveness is not the primary limitation in the LLM regime, as current large models will fail on simple tasks. Using parity-learning setup, we demonstrate that CoT can substantially improve sample efficiency even when the representa- tion power is sufficient. Specifically, with CoT, transformer can learn the fune- tion within polynomial samples, whereas without CoT, the required sample size is exponential. Additionally, show that CoT simplifies the learning process by introducing sparse sequential dependencies among input tokens, and leads to a sparse and interpretable attention. We validate our theoretical analysis with both synthetic and real-world experiments, confirming that sparsity in attention layers is key factor of the improvement induced by COT.
INTRODUCTION
Chain-of-thought (CoT) has proven to be a powerful technique for enhancing reasoning large language models Wei et al. (2022); Kojima al. (2022). By instructing the model break complex problems into smaller, manageable steps, CoT facilitates more efficient reasoning and better gener- alization, particularly in algorithmic and logical tasks Nyee al. (2022); Lewkowyez al. (202));
Wang et al. (2023b). Building on this, performance can be further improved through mu multi step prompting and multi-path sampling techniques Chowdhery et al. 2023 Wang et al. (2022b): Zhou
et al. 2023ay: /Zhang et al. (2023); Fu et al. (2023).
This focus on CoT within in-context learning has since expanded to more structured learning ap- proaches Yao et al.) (2024); [Besta et al.) (2024). By adding reasoning examples of CoT style to the instruction-tuning dataset, models enhance their problem-solving abilities more effectively than re- lying solely on CoT during prompting Zelikman al. (2022): Chung et al. (2024). As result, CoT now shaping a new paradigm in language model development, marking shift from simply scal- ing data Kaplan et al. (2020); Hooffmann et al. (2022) to focusing advanced reasoning strategies Lightman et al. (2024), which leads to more effective learning outcomes.
"These authors contributed equally. 1 A large part of this work was done while Kaiyue at Tsinghua University.
This work is independent of and outside of the work at Amazon.
Our code is available at wtthttta://// / www.tthhbb....... com/ http://ooo...ooooo.oaaaaaaaaaaaa

---

Published as a conference paper at ICLR 2025
Normalized ledd Attention Entropy
with
Normlllzed Attention Entropy
without
Sample Complexity
-10
without
with CoT
2.0
1.1
Evaluation Accuracy
1.0
0.8
0.8
10
0.6-
10
0.4
0.
0.5
10
0.0
Number of Secret Variables
Iterations x Batch Size
lttratioon x Batch Size
(a) Sample complexity of 4- layer 4-head transformer on parity with n 30.
(b)
Normalized attention entropy and evaluation accuracy
curves of 4-layer 4- 4-head transformer on (n = 30, k = 3) parity problem.
Figure 1: (a) We show that, without Chain-of-Thoughe Thought (CoT), the sample complexity for training transformers to learn the parity function grows exponentially with the hardness parameter k. In contrast, utilizing CoT significantly improves sample efficiency. (b) We also show that the sparsity of attention layers, measured by normalized entropy [ is crucial in the parity learning experiment. In both CoT and non-CoT scenarios, as the attention layers become sparser-indicater by a rapid decrease in normalized entropy a corresponding jump in evaluation accuracy occurs.
While COT's success is well-established, understanding why it works is still hotly debated topic Saparov & He (2023); Prystawski et al." (2024). Recent theoretical studies suggest that CoT en- hances a model's expressiveness, increasing representational capacity when the sequence is long enough[Feng ct al. 2023);; ;; Li et al. (2024c). However, expressivity alone does not guarantee success. Large language models often struggle with simple tasks like counting the number of 'r's in "straw- berry"- when not using CoT. Given the increasing model sizes, it seems unlikely that such tasks are inherently inexpressible. This discrepancy calls for a deeper study of generalization, hinting the true power of CoT may lie beyond the expressiveness.
In this paper, we study the benefit of CoT from the sample efficiency perspective. We provide con- crete examples where expressiveness is not the limiting factor; that is, the function can be expressed both with and without CoT by the transformers. We demonstrate, both in theory and in practice, that without CoT, the learning requires exponentially more samples comparing to with COT. Further, show that CoT sequences introduce sparse sequential dependence, thereby enhancing the sparsity in the attention layers. We then show that transformers efficiently optimize and generalize on such sequences. We summarize our contributions as follows.
1. Theoretically, show that while the parity problem can be expressed by I-layer transformer, learning requires exponentially many samples without CoT when the number of parameters is limited (Theorem I and 2). Meanwhile, for CoT data with sparse sequential dependence, a 1- layer transformer can learn the parity function and faithfully represent the sparse dependence in
its attention pattern with almost linear samples with respect to sequence length(Theoren) 3
Empirically, we verify our analysis that training on parity function with CoT data requires only polynomial samples and will induce sparse and interpretable attention (Figure B We further show that evaluating and training on CoT data will also induce sparser attention on real-world dataset GSM-8k on pretrained language models.
2 MOTIVATING EXAMPLES: LEARNING PARITY FUNCTIONS
We start with empirically exploring the sample efficiency of training transformers to learn the class of parity functions (Blum et al. 2003) with and without COT. The parity functions exhibits simple structure but yet hard to learn by traditional networks Abbe & Sandon (2018); Malach & Shalev- Shwartz (2022). Specifically, given set of binary variables by, b2 b2,,b..... ... b, parity function takes k secret variables bi..... bi, and outputs if the sum of these variables is odd, and 0 if it is even:
((b,...,be b.)) bi + big @ ... big.
where b; E {0, and is the XOR operator. For example, f(b,, b2, ba, ba, bb) bi @ b2 @ by is a 5-variable parity function with & = 3. The function of returns if the sum of b2 by and by is odd, and 0 otherwise, independently of the value of by and by. Intuitively, the parameter i controls the hardness of the problem. When increases, the number of possible subsets grows exponentially, making the identification of the correct secret more challenging.
2

---

Published as a conference paper at ICLR 2025
generate sequences of length + 2 auto-regressive manner: b1. ... bn. bn+1 := and bn+2 = f(])) To incorporate CoT. we break the sum XOR down into & steps and add all the intermediate steps into the sequence, i.. including b, b; + big ..., bi @ big .)... ..... big. As a result, the CoTT data has length n + k + 1. With the example function f(b1, b2, bs, ba, hs) b1 + b2 ba, one sampled sequence would be
Given a nt- naviiable parity function
No CoT
0, 1,0.
E (0,1)",
as b b2 ba 0
O.
inpat
[EOS] answer
With CoT
0, 1,000.0000
0,1 CoT
O {0,1)" b1 0,61 =2 =1
input
[EOS]
answer
Next, train transformer networks with and without CoT respectively, with common held-out test set, and compare their sample complexities. The sample complexity is defined as the amount of data model sees at the first time achieves perfect validation accuracy. More precisely, we halt the training process when the model reaches 100% validation accuracy and record the total number of training examples used up to this point as the empirical sample complexity.
In Figure we compare the empirical results of learning parity function with and without CoT. The results show that the sample complexity without CoT (orange curve) grows exponentially, while with CoT (blue curve), it increases linearly. In other words, incorporating CoT allows us to learn the task with exponentially fewer samples. Additionally, suggests that sparsity in the attention layers plays crucial role in the learning process. To explore why CoT enables transformers learn more efficiently, we formally define our problem setup below.
THEORETICAL ANALYSIS
In this section, we provide a formal analysis of the training dynamics of Transformers on the parity problem, both with and without Chain of Thought (CoT). We select the parity problem testbed because identifying a set of key variables amidst various confounding ones is a fundamental aspect of many reasoning tasks, and parity serves as abstraction of this process.
3.1 NOTATIONS AND DEFINITIONS
We represent a sequence of length T as b[l].... b[7]. A parity function is function of form parity s (b) = @esbbii, where S is the the XOR operator on binary
Parity Problem. Each token in our setting is either
set of secret indices that is fixed during training and testing, and variables. The cardinal of the secret set A = |S| controls the hardness of the problem.
Definition (Parity Problem (n, k) without CoT). Given a secret set S of cardinal
define the blb] U[O, 1}), vi E [1. b/ + 1] = 0, b[n + 2] parity g (b).
parity problem without CoT as the following distribution of sequences PNoCoT:
where b[]], [m]] are uniformly sampled from 0 and 1.
Definition 2 (Parity Problem (n, k) with CoT). Given a secret set S of cardinal k. we define the parity problem with CoT as the following distribution of augmented sequences PCoT:
b[i] ~ ({0, 1)), E [1, m), b(n + +1] 0,bli + 1] = b @ [[]] Vi E (n + I,..., n k). where S is any permutation of S, i.e. {S[]] = + 1
S.
When CoT is provided, the data includes a step-by-step computation of the desired parity function, adding one variable at a time. Note that CoT is not unique for given secret set since we can arbitrarily permute S, given the commutative property the XOR operator.
Transformer Architecture. To conduct the theoretical analysis, we simplify the Transformer archi- tecture similar to prior works (see e.g. Wang al. (202bb): Nichani al. (2024): al. (2023b). More precisely, we drop all the layer norms or batch normalization; simplify the positional embed- ding; use a square matrix to represent the attention layer and concatenate the residual branch a Densenet fashion Huang al. (2017). After simplification, the network still has strong expressive power like the standard Transformers.
In standard Transformer, embedding layer matches each vocabulary token into dense vector of size d and then adds a positional embedding. In our case, we only have boolean tokens and the sequence has constant length T. Hence, we simply freeze 27 unit vectors the embedding vectors:
3

---

Published as a conference paper at ICLR 2025
Definition (Embedding Module). For any position [T], sample two embedding vectors eto and 6 uniformly at random from unit hypercube U (1 t2)") These embedding vectors are frozen during training. Then for any binary sequence {0.1)", the embedding is defined as
E (b) (ecssts)lettt E B
Due to the properties of the hypercube, with high probability, the embedding vectors are near- orthogonal with each other.
Next, we define the attention layer. standard architecture, the attention layer is derived from the product of query and key matrices, i.e. QK X W (WK)"X. In our case, we directly use full matrix X AX to parameterize it, which has the same representation power.
Definition 4 (Attention Module). Given an imput matrix define the attention module
E RdxT and attention weight A, we
A(X) X So
where C' E Tx is the autoregressive mask with value oo on the lower triangular matrix.
The attention module is then followed by a fully connected layer with ReLU activation: Definition 5. We define the FFN function with width 2m over matrix X E R22xT as FFN(X) h22 ReLU(WX).
Here W E 2mx2d h E 222xx with o being the output dimension and ReLU is element-wise.
Finally, we concatenate the modules into simplified Transformer block.
Definition 6 (Simplified Transformer Block). We define simplified Transformer block as,
T(b) -FFN ((A
where is the embedding module, is the attention module and FFN the fully connected layer. An L-layer Transformer consists of a composition of L such blocks, with the embedding module E appearing only in the initial layer. The intermediate dimension is set to d, while the output dimension of the final layer is I.
Compared with the standard residual structure, we use Densenet structure concatenate the resid- ual branch with identity branch. Again, this transformation does not affect the representation power of the network, which is the standard practice in previous theoretical analysis (see e.g. Wang et al. 20244))(Nichan et al. )2024)).
Loss function. To simplify our analysis, the loss function. We define the next token prediction loss of boolean sequence as L(w) = (b) 1], bfn + 2)) and L(w) Ei (b) blū IJ) for with and without CoT setup respectively, where w denotes all the trainable parameters.
hinge loss E(j. y) max (oxx((1))0 + 1,0) as
3.2 EXPONENTIAL SAMPLE COMPLEXITY WITHOUT CoT
We now present analysis in the no CoT setup. First, we show that the parity problem is easy to represent with the simplified Transformer architecture:
Theorem (Easy to Represent). Consider the Transformer model defined in Definition for any 8 < 0.1 and large enough n, when the number of secret indices is in [n/ (og²(n/)), n/ log' (n/8)), with probability at least - 8 over the randomness of embedding e, there exists weight configu- ration of the Transformer with dimension
O tog(n/5) and width 2m O(k) with O n)
precision of the weights and activations, such that d2 + dm = o log n), and it achieves perfect accuracy on the parity problem (n, k) without CoT, i.e.
vb PNoCoT, sgn (T(b) (n 1)) =111)01+++11
4

---

Published as a conference paper at ICLR 2025
In other words, the model possesses sufficient expressive power to represent any parity function, even in the absence of the chain of thought. Therefore, we are in the representationalsufiicci regime where expressiveness is not the bottleneck. The proof of this statement is based on random matrix theory and concentration inequalities, which we defer to Appendix B.2
While the function is expressible, this does not guarantee that the solution be easily found. In fact, we show that achieving a perfectly accurate solution using a gradient- optimization method requires an exponential number of samples in & assuming the memory, i.e. space complexity, to perform optimization is bounded throughout the training.
Theorem (Hard to Learn). For any randomly initialized simplified Transformer model with con- stant layers, when the embedding dimension d and width 2m satisfies de + dm o(nk/ log(n)), (n)n)),,] for any constant number of passes q, when the model is trained with q 9-pass stochastic gradient descent with ((d dm) log n) memory, the sample complexity required to learn the parity problem (n, k) without CoT any nontrivial accuracy a 50% with nontrivial probability
50% is 20(k)
The proof is deferred to Appendix B.3 We use results from the classical online learning communi- ties (Lyu et al. 2023) to show that the regime where the memory required perform the training is less than nk, the model can't effectively store information about enough samples inside the pa- rameters during training, and hence can't It infer the secret indices effectively. Here the parameters denotes the size of the secret set S and the sample complexity grows exponentially when grows.
3.3 POLYNOMIAL SAMPLE COMPLEXITY WITH CoT
The result in Theorem presents seemingly negative outcome for learning algorithmic reasoning, the model requires exponentially many samples relative k. However, our main result, Theo- rem we will demonstrate that the sample complexity can be significantly reduced if a step- step-by-step derivation is provided. In other words, CoT much more sample efficient, where model can effec- tively learn complex relationships as long as each token depends on only few previous tokens.
We will use the following initialization of the 1-layer Transformer model. The attention is initialized to be uniform (A 0) and the contribution of attention output is initialized to be zero (Wr 11,21:12d 0). The parameter in the FFN and the word embedding E is fixed during training. Assumption (Initialization). initialization,
a+t Σ
Vr E [2m] A 0, W, W,4+1224 O, W,
Σ Vrsibei,b. him
2m hm+1:mm
2nn
Here Vr,i. is independent random variable sampled uniformly from {-
We can then train the model with stochastic gradient descent (SGD) and get the following theorem. Theorem 3 (Easy Learn with CoT). For any constant ó (0, 1), when the number of secret indices & is in [n/ log" (n/8), log"(n/)), n/ log" (n/8)), with probability 1 1-6 8, randomly initialized simplified Transformer (see Assumption 7 with o(nk) parameters trained for constant steps using mini-bare
SGD with (n) samples using appropriate hyperparameters (see Assumption 2 can reach perfect
accuracy on a parity problem (n, k) with CoT.
vb PCoT, i E (n 1,......,.. k). (T(b)li) bli
Furthermore, after training, the attention pattern is interpretable and one-hot in the sense that, for any b PCOT, E {n 1,, .n k),J S
softmax(C E(b) E() - < 1/n².
Moreover, the result still holds even if all the weights and activations are in O n) precision.
This theorem indicates that the attention module can successfully extract the sparse sequential de- pendencies in the CoT data and faithfully represent in the attention pattern, which validated in experiments (see Figure 3. It is also amongst the first optimization dynamics analysis of Transformers using finite-sampi gradients rather than population gradients. We delay the full proof of the theorem Appendix B.5
5

---

Published as a conference paper at ICLR 2025
Proof Sketch. In our analysis, the dynamics of the model includes three key phases. In the first phase, the weight of the FFN layers become correlated with the embeddings associated to the secret indices. In the second phase, this correlation caused the attention module to receive strong signal to amplify the attention weight on the corresponding secret indices at each position. The attention pattern will become one-hot after this step. Finally, in the third phase, the FFN layers learn the correct mapping to the output, utilizing both the embedding at the current token and the retrieved embedding from the secret indices as indicated by the attention.
Phase 1. Configuring the FFN At initialization, as the FFN weight corresponds to the attention output (Wr r,2+1:24) is initialized to be zero. The set of neurons activated each position E [n + 1, k] is solely determined by the word embedding ((t.b]t)) at the position. Because 6,, are nearly orthogonal, it holds that (ei,bft), Wr.1aa) R3 Vr,i, for all i and r with high probability. As a result, for a fixed and b[i], the set of activated neurons at position i is frr Vy,aa/j > 0}. Therefore, at initialization, the MLP can be viewed as an ensemble of multiple linear functions specialized each position and boolean value. We can show that when the learning rate is small, the set of activated neurons each position remains the same through the training process. Hence, we can conceptually view the FFN function as set of different linear functions applied independently at each position and value. We will denote this set of linear weight as Kg,i (formally defined in Lemma 18).
As the attention weight matrix A is initialized as zero, the attention pattern will be uniform at initialization. Notice that in the parity data with CoT, the only position whose value correlates with bli + 1] when conditioned on blt] is S[i]. This suggests that the linear weight (6.) will have a stronger correlation with the embedding ((s())b())) than other embeddings.
This step crucially relies on the sparse dependency in the CoT data. The strong linear correlation will not be present in the data without CoT when k >> 1, as every token b[i] for í E [n + 1] will be uncorrelated with the desired output b(n + 2) in such case.
Phase Learning the Sparse Attention At the second step, as the FFN weight corresponds to the attention output is no longer zero, the attention weight matrix A will receive a non-zero gradient. By the chain rule, the gradient corresponding to how the attention from i-th token attends to the j-th token softmax(C+*X"AX) of the j-th token to the output at the i-th token (K1,,.bfi) ((bb())))) ) is larger. This suggests that the attention will be amplified the index that has strong correlation with the output at the current token, which is the secret index Sli). This step will make the attention pattern approximately one-hot as in the theorem statement.
i] will be larger when the approximate contribution of the embedding
Phase 3. Learning the Output Mapping At the final step, as the attention pattern is one-hot, the FFN layer only needs to learn mapping from the embedding at the current token and the embedding at the secret index to the output. As the mapping is linear conditioned on the embedding the current token, the FFN layer can learn the correct mapping within a single step. This step will make the model reach perfect accuracy on the parity problem with COT.
4 EMPIRICAL EXPERIMENTS
In this section, we validate and extend our theoretical findings through comprehensive experiments in the following three aspects: First, we empirically confirm that CoT reduces the sample complexity of transformers in learning the parity problem. Second, we conduct multi-pass training of transform- ers without CoT, as complement to the established lower bound (Theorem ( which only applies to constant- pass training. The results indicate that multi-pass passs training indeed improves the models' ability learn parity problems. Third, we conduct experiments on the GSM8K dataset (Cobbe et al.) 2021) to show that CoT introduces sparse sequential dependence on real-world training data. Al- though necessary simplifications on transformer are made in Section develop theoretical results, we use the standard transformer architecture in the subsequent experiments following the GPT-2 architecture (Radford et al. 2019) with trainable position embeddings unless otherwise specified.
4.1 PARITY LEARNING WITH MULTI-LAYER TRANSFORMERS
In this section, we conduct an ablation study on the sample complexity of transformers with standard GPT-2 architectures and one-pass training, with and without COT
Ablation study on sample complexity (Figure 2). We train transformers on parity problem with 30 and k 1, 2, 3, 4, with and without CoT, varying layers and heads from 1 to 4. We choose

---

Published as a conference paper at ICLR 2025
k
with CoT Layer Number
Layer Number
Layer Number
Auuman
Javer
Sample Complexity
100
3
Head Number
Head Number
Head Number
Head Number
without CoT Layer Number
pqquee Lake
Layer Number
Layer Number
100
3
10a
Head Number
Head Number
Head Number
Head Number
(a) Sample complexity of transformers with varying numbers of layers and heads parity problem with n=30
Sample Complexity
with CoT
Sample Complexity
with CoT
104
10a
10F
10a
101
Number Secret Varlables
Number Secret Variables
(b) Sample complexity of 1-layer 1-head transforer on parity problem with n 30
(c) Sample complexity of 4-layer 4-head transformer on parity problem with 30
Attention Pattern
Sample Complexity (x10")
with CoT
0.91
0.9
0.7
0.6
20 30
40
50
60 70 80 90 100
Number of Secret Variables
k
(a) Sample complexity of 1-layer 1-head transformer on parity problem with 100 using COT.
(b) The attention pattem of 1-layer 1- head transformer trained on (n 40, 20) parity problem, detailed in Section 4.1.
Figure 3: (a) When n is fixed, the sample complexity of learning parity with CoT grows approx- imately linearly with (b) The attention pattern learned by the transformer with CoT inter- pretable, the i-th output token of CoT predominantly attends to secret index S(i).
the best performer across learning rates from 6 10-5 8 10-5, and x 10 10-4. At each step, fresh batch of training data is sampled, with a maximum budget of 10" samples. We record the number of samples seen by the model before reaching an evaluation accuracy of 1 as the sample complexity.
Results. Figure 2hooes that training with CoT [first row in Figure (a)] consistently achieves better sample efficiency than training without CoT [second row in Figure 2 (a)]. Moreover, for a fixed configuration of heads and layers, the sample complexity without COT grows exponentially with parameter contrast, CoT greatly reduces sample complexity, showing an exponential improvement across different numbers of heads and layers. These findings consistent with our theoretical analysis.
Training with CoT induce sparse and interpretable attentions (Figure We evaluate the sample complexity of a 1-layer, 1-head transformer on larger parity problems (n = 100 and 20 < k 100). Without CоТ, training exceeds the sample budget when & 4. However, with CoT, we can successfully train for any k. Figure Bp shows the attention pattern 1-layer, 1-head transformer trained parity problem with (n 40, k 20) using CoT random query. In

---

Published as a conference paper at ICLR 2025
1.0
Evaluation Accuracy
4
0.8
Layernumt
3
2
1
0.6
number
200
400
600
800
1000
epoch number
Figure 4: Evaluation accuracy of transformers on the (n = 20, k 6) parity problem without COT. The model trained on a dataset of 10, 000 samples for I, 000 epochs. Almost all layer-head configuration achieve perfect evaluation accuracy. Adding more heads is more effective than adding layers. The blue dashed line marks the with CoT setup, which achieves perfect accuracy in 5 epochs.
this case, CoT processes the secret variables in ascending order. The i-th row illustrates the atten- tion weight of the i-th token attending to the previous ones, with indices corresponding to secret variables highlighted red.
Results. In Figure 3h we show that training with CoT succeed with large and k. Moreover, the empirical result suggests that the sample complexity of learning parity with CoT grows approx- imately linearly with k. This shows clear contrast to without CoT setting where the required samples explodes exponentially. To validate our theorem on learning the sparse dependence, we show in Figure 3 that the attention pattern learned by the transformer with CoT is interpretable. Specifically, the i-th output token of CoT primarily focuses on the i-th secret variable, with minimal attention given other tokens.
4.2 MULTI-PASS TRAINING IMPROVES PARITY LEARNING WITHOUT CoT
In our theory, we only establish lower bounds for the sample complexity of learning parity problems with constant pass SGD training. To complement the established lower bounds for transformers learning parity (Theorem 2). we conduct empirical experiments by training transformers without CoT using multi-pass training. We make two observations below. First, the results demonstrate that training with repeated data can help Transformers learn the parity function, although this process still consumes significantly more computation (epochs) than trained on CoT data (Figure(4). Second, a key difference between one-pass and multi-pac culti-pass training is the development of sparse attention (see Figure 5 Right), similar the role of CoT shown in the previous section. This shows that the development of sparse attention is crucial even in when training without CoT on the parity data.
Experiment protocols. On (n = 20, 6) parity problem, we conduct multi- paltipppas training without COT.. In Figure the models are trained with 10a samples for 1000 epochs, with the number of layers and heads ranging from to 4. In Figure 5 we compare the training of 4-layer 4-head transformer on 5 1 and 10th samples respectively. The learning rate is initialized 10
Normalized attention entropy. To measure attention sparsity, we introduce the concept of normal- ized attention entropy for each attention head (illustrated in Figure 5 Right). Let (x)(()[i] denote the attention score distribution produced by the h-th head in the f-th layer at the i-th token of input The normalized attention entropy for the input is then defined as:
((PF((x)[i)) log
Ent (, min
[0. 1].
(1)
122
where HPP) = - (o is the entropy of distribution P. The normalization term log i rep- resents the entropy of uniform distribution over i tokens, account for the varying context length. The minimum is taken over different tokens in since attention heads may specialize in extracting information for specific tokens. As result, lower normalized attention entropy indicates a sparser attention pattern. To compute the normalized attention entropy for each head, we would average the normalized entropy across all question-answer pairs in the validation set.
when trained on 10 samples, which only accounts for a small por- tion (~ 1%) of all possible inputs (220 x 10%), most of the transformer architectures we examined achieve perfect evaluation accuracy given sufficient epochs. While the A = problem is intractable with one- training without COT, these results demonstrate that multi pass training can indeed en- hance learning in the no-CoT setup. However, CoT is by far the most effective accelerator, achieving perfect accuracy in just 5 epochs, significantly outperforming the multi-pass no-CoT training.
Results. As shown in Figure 4
8

---

Published as a conference paper at ICLR 2025
Norma Foee Attermon Entropy Normatzeed Anbermton Entray
Layer
Layer
4
50000 traiiing sampxs
Eowkuston Accurney 6.
14
0.5
om
Berations Batch
04
Peraaiona Batch Siz 141
Teraciana Batth
01
62
Batch Size
1000000 training cemples
Eowllation Rccurnecy
6.
-8-
231
6.7
Sigglnt
Eerations x Batch Szz
Btrations Batch Size
beracione Batthh
herations x Batch Size
Figure 5: 4-layer 4-head transformer trained on the (n = =00, 6) parity problem without CoT using multi-pass training, detailed in Section When trained on small dataset of 50000 samples, the model achieves perfect evaluation accuracy (Top), accompanied by significant decrease in en- tropy. Surprisingly, when trained on an even larger training set with 1000000 samples, the model fails to learn (Bottom), and both the training loss and the normalized attention entropy remain ele- vated. This shows that the development attention sparsity may improve optimization.
Although learning without CoT is less efficient, offers a "slow-mode" trajectory. In Figure 3 (Top), (Top), the no-CoT loss function initially plateaus before eventually dropping to zero. During this plateau, the entropy in the attention layers continues to decrease, indicating that feature learning is occurring gradually. When all the attention heads become sparse, a transition phase occurs in the loss and accuracy: the loss drops to zero, and evaluation accuracy jumps to 1. In contrast, in the failure case shown in Figure 5 (Bottom), the entropy of the attention layers remains high throughout.
This experiment on no-CoT confirms that sparse attention is crucial for parity learning- As previ- ously demonstrated in Figure CoT accelerates learning by quickly inducing sparsity. This suggests that CoT not only improves the sample efficiency but also improve the optimization landscape by facilitating sparsity, due to the introduction of sparse dependencies added the intermediate steps.
4.3 COT INDUCES SPARSITY ON REAL-WORLE DATA
Now we move from the synthetic parity problem to real world experiment on GSM8K dataset of grade-school math word problems (Cobbe et al. 2021). We observe that:
Real-world CoT data also exhibits sparse sequential dependence, leading to a sparser attention pattern in pre-trained models.
2. Fine-tuning on CoT data further enhances the models' attention sparsity on the input data.
Experiment protocols. In Figure 6 we examine data types: With CoT, where inputs from GSM8K dataset are concatenated with ground truth answers that include multiple reasoning steps, and No COT, where inputs are directly concatenated with the final answer. We evaluate two lan- guage models: the pre-trained model Qwen2-7B (Yang et al. (2224) and the specialized mathemat- ies model (wwe2-7B-Maae (Qwen] 2024) which is fine-tuned from Qwen2-7B on mathematics- specific corpus with CoT data. We plot the normalized attention entropy (Equation i across different heads. More details can be found in the Appendix Appendix C.T
Results. Unlike the synthetic parity problem. the sequential dependency of real-world data is hard to measure directly. However, it can inferred from the attention sparsity pre-trained models when they process such data as input. As shown in Figure 6 comparing the normalized attention entropy of the same pre-trained model Qwen2-7B on different types of data, we can see that the entropy is lower for With CoT data compared to No CoT data, indicating that real-world CoT data indeed exhibits a sparser structure. Furthermore, on the same With CoT data, Qwen2-Math-77 model demonstrates lower attention entropy compared Qwen2-7B model, suggesting that fine-tuning on CoT data promotes the development of sparser attention patterns in the model.
5 ADDITIONAL RELATED WORK
Parity Learning. The most relevant work to ours is Wies et al. (2023). ) which shows that subtask decomposition enables learning the parity problem with polynomial sample complexity in recurrent neural networks (RNNs). Their learnability results rely on techniques from WWng et al. (2022a), which operate within the NTK linearization regime of RNNs. In contrast, the Transformers

---

Published as a conference paper at ICLR 2025
Normaliea Abteritonn Entopy
Layer 1
Layer 11
Layer 21
Layer 28
0.4
0.3
0.4
0.3
0.1
Qwein278-- No CoT
0.2
0.2
Qwen2-78- With CoT
0.2
0.1-
With CoT
Owen2-Math-PP
0.0
28
0.0
28
0.0
28 0.0
head
head
head
head
Figure 6: We compare the normalized attention entropy of the pre-trained Qwen2-7B and math- specialized Qwen2-Math-7B models on the GSM8K dataset with and without CoT prompting (Sec- tion 4.3) Each bar represents the entropy of an attention head. The Qwen2-7B model exhibits sparser attention when processing CoT data, indicating that worldworld CoT data has a sparser struc- model suggests that
ture. The entropy difference between the Qwen2-7B and Qwen2-Math
fine-tuning on CoT data promotes the development of sparser attention patterns of the model.
alyzed in our work exhibit feature learning and identify the sparse secret set within the attention module. On the positive side, line of works that studies optimization dynamics of neural networks on parity (Kou et al.) 2024; Barak et al., 2023 Edelman et al. 2023 Daniely & Malach 2020 Abbe
et al., 2023; 2024b) show that (2) samples is sufficient to learn parity, which is close to the sta-
tistical query lower bound (Kearns 1998.. On the negative side, it is well-established that learning
parity using gradient descent (GD) requires ((x*) iterations due to the SQ lower bound (Kerns, 1998). For SGD learning, (Shalev-Shwartz et al. (2017); Abbe & Sandon (2020) showed that an exponential number of samples is necessary when SGD is hindered by additional noise or when the number of weights updated in each step is constrained. However, remains unelear whether standard mini-batch SGD will still need exponential sample complexity. In our paper, we take the alternative approach consider the training as an online algorithm with bounded memory (parame- ters) and utilize results from online learning literature to provide a rigorous exponential lower bound (Lyu et al. 2023; Kol et al., 2017). (Hahn& Rofin (2024) shows that while parity is easy to represent using Transformers, the sensitivity structure of the function will require the representation to have large weight norms. A recent work (Abbe et al.) 2024a) conjectured that a distribution is weakly learnable by a Transformer if and only if it has constant globality, which is supported by empirical experiments. The work also highlights the parity problem as a notable special case.
Chain-of-Thoughh (CoT). CoT is the technique let the model generate a reasoning process be- fore final answers. CoT prompting has proven effective in enhancing language models' reasoning capabilities (Wei al. 2022 Zhang et al., 2022 Wang et al., 2023a Zhou al., 2023b Wang & /Zhou 2024 Creswell et al. (003).. Training with CoT data has further improved the model 's capa-
bility performing complex reasoning (see Qwen| (2024): Yue et al. (2023); Yu et al. (2024): ) Kim
ct al. (202)) and reference therein). Different lines of work have explored the effect of COT. From the
representation theory perspective, works including Feng et al. (2024); Merrill Sabharwal (0024);
Li et al. (2024d); Nowak et al.] (2024); Wen et al. (2024) show that CoT can provably expand the representation power of neural architectures including Transformers and RNNs. From the statistical approximation level, prompting with CoT reduces the statistical error (Hu et al. 2024 Prystawski
al. 2023). Li et al. (2023a) studies how MIP learns with CoT data. [L et al. (2024a) studies how Transformers learn to perform CoT through prompting using gradient descent. We differ from this work as (1) we focus on how Transformer captures the sparse sequential dependency in CoT, which not reflected in their data modeling; (2) we study zero-shot CoT data directly. Dutta et al. (024); Wang et al.] (2024a) studies pre-trained Transformers' activation on CoT data. They highlight the attention head's 's role in moving the essential information from the context to the reasoning process, which is consistent with our theoretical insight. Concurrently with our work, Kim & Suzuki (2024a) theoretically demonstrated that CoT reduces the number of gradient steps required to solve the parity problem. While their analysis focuses on GD with a noisy oracle, our work investigates mini-batch SGD and provides thcoretical guarantees on the exponential reduction in sample complexity.
We defer other related works to Appendix A CONCLUSION AND FUTURE WORKS.
Our work demonstrates that Transformers trained on Ofain-of-Thoughhe (CoT) data with sparse se- quential dependencies can efficiently learn sparse attention mechanisms, accurately capturing these dependencies while requiring exponentially fewer samples than models trained without COT. Our current analysis of CoT training assumes that each subsequent token depends on exactly one token in the context and focuses solely on the parity function. A promising future direction is to explore more general scenarios where each token depends on multiple previous tokens and extends to func- tion classes beyond parity.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
Jingzhao Zhang is supported by National Key R&D Program of China 2024YFA1015800 and Shanghai Qi Zhi Institute Innovation Program.
REFERENCES
Emmanuel Abbe and Colin Sandon. Provable limitations of deep learning. trXiv:1812.063699 2018.
arXiv preprint
Emmanuel Abbe and Colin Sandon. Poly-time universality and limitations of deep learning, 2020.
URL http:///wwxi...or//////////// 02992
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: laap complexity and maddle-to-sadde dynamics, 2023. URL htttss://axiv..or org.v..rrraaaaa abs/
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can trans- formers reason? the locality barrier and inductive scratchpad. arXiv preprint arXiv:2406. 06467, 2024a.
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks, 2024b. URL http://xxvvv.//// wwt.o/////b..... 2202. 08658
Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit, 2023. URL
/ttt:s///araiv.orr////////2/2/// wttp://////pbbb//ooo/////// 08799
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gian- inazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, 2024.
Avrim Blum, Adam Kalai, and Hal Wasserman. Noise-tolerant learning, the parity problem, and the statistical query model. J. ACM, 50(4):506-519) jul 2003. doi: 00.1145/79258..795544
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learn- ing in transformers. Advances in Neural Information Processing Systems, 2022.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality, 2024. URL 19442
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsyyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito. David Luan, Hyeontack Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Heier-Hellsstem, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240), 2023.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu,

---

Published as a conference paper at ICLR 2025
Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pel- lat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoe V. Le, and Jason Wei. Scaling instruction-finetueed language models. Journal of Machine Learning Research, 25(70), 2024.
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Frank Cole, Yulong Lu, Riley O O'Neill, and Tianhao Zhang. Provable in-context learning of linear systems and linear elliptic pdes with transformers, 2024. URL htttt::/tsaaat http://axivvvoo/ org/ abs/
1,09.12299
Antonia Creswell, Murray Shanahan, and Irina Higgins. Selection-inferente Exploiting large lan- guage models for interpretable logical reasoning. In The Eleventh International Conference Learning Representations, 2023.
Amit Daniely and Eran Malach. Learning parities with neural networks, 2020. URL https:
/aax:vvcrr//bb/22000.0
Subhabrata Dutta, Joykirat Singh, Soumen Chakrabarti, and Tanmoy Chakraborty. How to think
step-by-step: mechanistic understanding of chain-of-thought reasoning, 2024. URL https:
Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Pareto frontiers in neural feature learning: Data, compute, width, and luck, 2023. URL https://arxiv.
org/ ht5/abb/220 2309/ 03800
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. In shirty-sevntth Conference on Neural Information Processing Systems, 2023.
Guhao Feng. Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. In Advances in Neural Information Processing Systems, 2024.
Yao Fu, Hao Peng, Ashish Sabharwal, Peter Clark, and Tushar Khot. Complexity-based prompting for multi-step reasoning. In The Eleventh International Conference on Learning Representationss 2023.
Michael Hahn and Mark Rofin. Why are sensitive functions hard for transformers?, 2024. URL https://arxiv. http://axxxv...r//baaa/aa
09963
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Henni- gan, Eric Noland, Katherine Millican, George den Driessche, Bogdan Damoe, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack William Rae, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Advances in Neural Information Processing Systems, 2022.
Xinyang Hu, Fengzhuo Zhang, Siyu Chen, and Zhuoran Yang. Unveiling the statistical foundations of maii--ofthhough prompting methods, 2024. URL
14511
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings of the IEFE conference on computer vision and pattern recognition, 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
12

---

Published as a conference paper at ICLR 2025
Michael Kearns. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM), 45(6):983-1066 1998.
Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought.
2024a. URL
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean- field dynamics on the attention landscape, 2024b. wtt.t//tpe////aa...ooro/o////b/a
01258
Seungone Kim, Se June Joo, Doyoung Kim, Joel Jang, Sconghyeon Ye, Jamin Shin, and Minjoon Seo. The cot collection: Improving zero-shot and few-shot learning of language models via chain-of-thought fine-tuning. 2023. httpt/tppp///a...oorp//re//
14045
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Gillat Kol, Ran Raz, and Avishay Tal. Time-space hardness of learning sparse parities. In Pro- ceedings of the 49th Annual ACM SIGACT Symposium Theory of Computing, pp. 1067-1080, 2017.
Yiwen Kou, Zixiang Chen, Quanquan Gu, and Sham M. Kakade. Matching the statistical query lower bound for k-sparse parity problems with stochastic gradient descent, 2024. URL https:
// //a.xii..rggab//////a/aaa
12376
Aitor Lewkowycz, Anders Johan Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Venkatesh Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with language models. In Advances in Neural Information Processing Systems, 2022.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear trans- formers acquire generalization -arraliiaatttogggaaateeeeedd cot ability? In High-dimensional Learning Dynamics 2024: The Emergence of Structure and Reasoning, 2024a. URL httpt::////// /tttps://openreviiee net Eorum/irummi 9pM8IrT6Xc
Hongkang Li, Meng Wang, Shuai Zhang, Sijia Liu, and Pin- Yu Chen. Learning on transformers is provable low-rank and sparse: one-layer analysis. In 2024 IEEE 13rd Sensor Array and Multichannel Signal Processing Workshop (SAM), pp. 1-5, 2024b. doi: 10.1109/SAM6025.. 2024. 10636506...9.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thought: Compositionality through in-context filtering and learning, 2023a. URL htt.tt:s//aaiiiv.... www.vvooorrraaaaoo abrr0/5.. 18869
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 19689-19729. PMLR, 23-29 Jul 2023b. URL v2C2/1123p html
mll. presss
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In The Twelfth International Conference on Learning Represen- lations, 2024c.
Zhiyuan Li, Hong Liu, Denny Zhou. and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems, 2024d. URL htt.hhttt////iiiiiieee http://///aiii.oo///aaaas//////
org hwwgg////////e argg/bb////2/2202/ 12875
13

---

Published as a conference paper at ICLR 2025
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, llya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference Learning Representations, 2024.
Xin Lyu, Avishay Tal, Hongxun Wu, and Junzhao Yang. Tight time- space lower bounds for constant- 08070
pass learning, 2023. URL https: http:/:://axxxv..rrg.
Eran Malach and Shai Shalev-Shwarde When hardness of approximation meets hardness of learning. Journal of Machine Learning Research, 23, 2022.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought, 2024. httpa///tts/////aaa.....//////// 07923
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers lear causal structure with gradient descent, 2024. / 14735
Franz Nowak, Anej Svete, Alexandra Butoi, and Ryan Cotterell. On the representationaal capacity of neural language models with chain-of-thought reasoning, 2024. http::sts://arxiv..o.
abs/2405 14197
14197
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. In Deep Learning for Code Workshop, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 2019.
Ben Prystawski. Michael Y. Li, and Noah D. Goodman. Why think step by step? reasoning emerges 03843
from the locality of experience, 2023.
Ben Prystawski. Michael Li, and Noah Goodman, Why think step by step? reasoning emerges from the locality of experience. Advances Neural Information Processing Systems, 36, 2024.
Team Qwen.
Introducing qwen2-math. gwen2-mathn mh/ 2024. Accessed: 2024-09-28
https:// www.enlm.ighhu..... github. bambbbi.o/oo/////
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Abulhair Saparov and He He. Language models are greedy reasoners: A systematic formal analysis of th-tugh... In The Eleventh International Conference on Learning Representations, 2023.
Shai Sawwar-, Ohad Shamir, and Shaked Shammah. Failures of gradient- baseed deep learning, 2017.
07950
Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison Zhou. Implicit regularization of gradient flow one-layer softmax attention, 2024. httpss::: // arrivv. wtgg/bb////3...... 08699
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understand- ing training dynamics and token composition in 1-layer transformer. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and Levine (eds.), Advances in Neural Information Processing Systems, volume 36. pp. 71911-71947. Curran Associates, Inc., 2023. URL |atpps://prcceedingg/ www.:/cceedinng..uuiiisss.... neurins cc wwtp:rppppscppppp....p.pppppppee wwtps//papeeeeeee////. Fil File/ (ale:33398ebba333c66722223333620
Roman Vershynin. High- Digeeniimms Probability: An Introduction with Applications in Data Science, volume 47 of Cambridge Series in Statistical and Probabilistic Mathematics. Cam-
bridge University Press, 2018 ISBN
doi: 10.1017/9781 118.779977811891596996 URL
httpss:/ddi.. .wreaooooorgg////o
wwwoorg//o.o /00. 1tt.g//....11/111111111
14

---

Published as a conference paper at ICLR 2025
Lei Wang, Wanyu Xu, Yihuai Lan, Zhiqiang Hu, Yunshi Lan, Roy Ka-Wei Lee, and Ee-Peng Lim. Plan-and-solve prompting: Improving sero oh reasoning by large language models, 2023a. URLttss:: wwt.:v/.rggg/bao. 2305. 04091
Lifu Wang, Bo Shen, Bo Hu, and Xing Cao. On the provable generalization of recurrent neural networks, 2022a.
Xuezhi Wang and Denny Zhou. Chain-of-thought reasoning without prompting, 2024. URL
https: wwg/ abtpa//s//22222 10200
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le. Ed Chi, and Denny Zhou. Rationale- augmented ensembles in language models. arXiv preprint TrXiv:2207.00747, 2022b.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Le. Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. selfsccons improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023b.
Zhiwei Wang, Yunji Wang, Zhongwang Zhang, Zhangchen Zhou, Hui Jin, Tianyang Hu, Jiacheng Sun, Zhenguo Li, Yaoyu Zhang, and Zhi-Qin John Xu. Towards understanding how transformer perform multi-step reasoning with matching operation, 2024a.
abs/ 0as5/2005.1300
Zixuan Wang. Stanley Wei, Daniel Hsu, and Jason D. Lee. Transformers provably learn sparse token selection while cunnecco nets cannot, 2024b. URL WRLLttts:///xxii.oor /bs/
2406. 06893
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. Rnns are not transformers (yet): The key bottleneck on in-context retrieval, 2024.
Kevin Christian Wibisono and Yixin Wang. In-context learning from training on unstructured data: The role of co-occurrence, positional information, and training data structure. In ICMI, 2024 Workshop on Theoretical Foundations of Foundation Models.
Kevin Christian Wibisono and Yixin Wang. How in-context learning emerges from training on unstructured data: The role of co-occurrence, positional information, and noise structures. In ICML 2024 Workshop on Theoretical Foundations of Foundation Models, 2024. URL https: //openzeview net wwtp://eetteddddndnemmama
Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in quence to sequence tasks. In The Eleventh International Conference on Learning Representations, 2023.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, et al. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances Neural Information Processing Systems, 36, 2024.
Longhui Yu, Weisen Jiang, Han Shi, Jineheng Yu, Zhengying Liu, Yu Zhang, James T. Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models, 2024. URL
http:/awwwi....aaa/aaaaaaaaaseaa 2309. 11288
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. 2023, URL
https:://ars /ttpv//rxvvvvv//v abg//b///// (tt./////3390...
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. STar: Bootstrapping reasoning with reasoning. In Advances in Neural Information Processing Systems, 2022.
15

---

Published as a conference paper at ICLR 2025
Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. Automatic chain of thought prompting in large language models, 2022. URL /ttE:p::/arxiv.og/// wtt.a//...///abb/21...o. 03493
Zhuosheng Zhang, Aston Zhang, Mu and Alex Smola. Automatic chain of thought prompting in large language models. In The Eleventh International Conference on Learning Representations, 2023.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoc V Le, and Ed H. Chi. Least-to-most prompting enables complex reasoning in large language models. In The Eleventh International Conference on Learn- ing Representations, 2023a.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoc Le, and Ed Chi. Least-to-most prompiing enables com- plex reasoning in large language models, 2023b. URL /ttpss//aaxx.v..org//////////
10625
16

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORK
Transformer Optimization Dynamics. Our works fall in the line of works that studies the opti-
mization dynamics of Transformers synthetic datasets (Li al. 2023b: Chen al. 2024 Kim
& Suzuki 2024b] Wibisono & Wang Chan et al.] [202]] Wibisono & Wang 2024
Cole et al. 2024
Sheen al. 2024, Chen al. 2024 Tian
2023 Nichani al.,
2024: Li al. 2024b).
Similar to our works, Wang et al. (2024b) highlights that Transformers can learn to select sparse critical tokens from the context linear data. We differ from their work in studying the mini-batch
optimization dynamics on nonlinear data and hence establishing sample complexity bounds.
B OMITTED PROOF
B. NOTATION AND ASSUMPTIONS.
We will denote sequence of binary variables b[ to denote variables corresponding to the j-th token dimension without otherwise specified.
..... b[T] with b[i] E 10, We will use [)]
B.1.1
TRAINING SPECIFICATION
Training Update
We will consider the following 1 1-pass SGD with batch size B:
B n+
Σ
Σ
{T
(b
(2 1)).
s=1 i=
(1+1)
L
1ht B
Here w includes A and W. We will leave unchanged for the simplicity of analysis.
B.2 REPRESENTATION THEORY
Proof of Theorem ! We will
loogn/()) and k + We will use p(x, 8) as
shorthand for 2(x 1) +
Define embedding matrix M as follows
Mp(t,)) = T
for [S[1),.... S(p}} = S. Then M E 22xxd
Lemma 1. For large enough n, with probability 1- 8, it holds that
| llpp = O(1). Amin (MM") > 1/12.
Proof. The first result follows from Lemma 28 noted that K = =(V1//d The second result follows from Lemma 29
Lemma 2. With probability 8/4 over the randommess of embedding ej,b for 8,b (0,1), there exists Uj satisfying that Vj E S. E 1),
Vj E S,b' {0, 11, =
Further,
- = O(Vk)
(0.1)
Proof. We can simply choose
U M (MMM"")) -1 02(1-1)46-
with 02(1-1)+6 being the one-hot vector in 2k
17

---

Published as a conference paper at ICLR 2025
Lemma 3. When the event in Lemma
happens, there exists v/--
(1g, = (0,,, 1.
Aj * J.J E ES (oyy;es))
Further,
Σ 2/1l2 S O(vk). JES
Proof. We can choose vi = 4j,0 + u,,1-
Lemma 4. When vú defined in Lemma 2 exists, for all S,b (0, 1}, with probability -
holds that
Σ j'eS
uy, ez,6 S 1/2.
Proof. Noted that 9j,b independent from My., assuming V Eyes uy, then [V O(vk). using Avema-Hog Bound, with probability 1 - 5/22
ΣΣ ey, ey,) 0 E Eyes 1y vlog(2k/0)) O(V 1lo((2k/0) j'ES
1
5 2
dd
The proof is then completed.
Assuming the event in Lemmas and happens, defined V = Eyes by and U = Ej'es ug
40 10 Wr.1d a,ent1, Wr,d41:2d b.U.
for some scalar a,, b, to be determined. For an arbitrary holds that.
1. j E S, (,00] Ae = mo.
2. If $ S, (ND(] Aca41.0 - 20
If we define
2j Aen+i1enn+.0
The softmax value is given by:
E
softmax(2) (=) = Eg E35
1. For j E S:
c 40 = 11
2. For & S:
e < 20 log 20
With [S] k; the denominator the softmax becomes:
Ee'r = Σ 10g n +
Σ
e E
Ikn 40 kn + 211
8'
J'ES
J'ES
For j E S:
n 40
softmax(2)xxt11 = kn + (n)
1 O(n 19) = + O(n -19).
18

---

Published as a conference paper at ICLR 2025
For j E S:
e²1
no
softmax(=))..)
O(n 20).
kn
On220)
kn
We then have
1 E S)
(olmax(E(()AAE()))B..
0 (7))
k
As
(U,e,,)) = 10j 1).
We have that
(U:)(E(()) +
Σ
bli + 0
h:
n
jES
Now we only need to map from the summation
EeEs b[jl to the parity of the summation. We set
b, 2k,
4[5 +4
if m,
a
1
i m +1 ff 2m.
2 "1'"""+++ +
ifl m,
h
-1 if 2m,
Then
T(E(b)) + 1]
ReLU(a, +2 Σb())) 0 ReLU (a, - (U, 1}})
p=1
m
ReLU(a, +2 >>biil
bij)
1
JES n
JES
(-1) Ses bbl+1 bbi+1 + + n1
(-1)bf++2)+1 + 0 1
n
Thus sgn(T(E((b)) + 1)) = ())))(++) for large enough n.
Note that the range of parameters is polynomial with n, thus could be represented with O n) precision. With (log n) precision, the error of each activation can bounded by a small inverse polynomial of and hence we can show the same result.
B.3 DYNAMICS WITHOUT COT
Proof. We will utilize the Theorem 6 in Lyu et al. (2023). which shows that any branching programs with o(nk) memory will require exponential samples to learn sparse parities with constant passes. Here the frozen embedding matrix e, which will utilize naively O(nd) memory, can't 't be saved in the memory. However, generator with fixed seed on each step. This allow us to simulate standard SGD optimization with o(kk) memory, which is special case of branching programs.
take the alternative approach regenerate using random number
19

---

Published as a conference paper at ICLR 2025
B.4 DYNAMICS WITH CoT
Assumption 2. Consider the following conditions for sufficiently large n.
The secret Hamming weight satisfies log" log (m/6))
Set d ls log' (%) and m 2k. This implies molog o(nk).
Define the batch size as B Coglog 20 (*) for sufficiently large constant C2 1.28 10?
Set the learning rates to be
mEv B
4ke
100
100 log (8)
3
Lemma 5, Under Assumption the following conditions hold as n + 0o:
1. lim 0 and lim
B 0.
7216
2. limn-roo V
2k d log
O.
3. Each of the following expressions tends
Vk log
300mmn
(a)
(b) 2/ log(
400mm
(c) 2nk log( 300mm
/Bd
4. Y/o = ?h I min
3mcv B
80 log
300
120
5. limm log 2
300mnB
nv
0.
dv Bd
6. 256mn
40 log
7. limn-+00 90
122 0.
8. limm (70 + m/) < 5m122-
Proof. The conditions in Lemma 5 are satisfied based on the definitions provided Assumption Below, we outline the verification for each condition:
1. Limit of and B nk"
A log 1.1 (3)
1
+0 as 00.
n
n2
logg.. (n/6))
log 2.9 (ng))))))))
Similarly,
B
Can log 20 ni
(.)
C2 log
20
C2 log 20
(3)
log 25 (n/5) 0.
C2
nk
(n./())))
n
2. Limit of / log
Bmn
2k
Bmm =
2h:
log (Czn log (?) -mn/6).
log
d
V log
1.1 (g)
20

---

Published as a conference paper at ICLR 2025
Simplifying,
2k
Bmn
vlog
log
Ollogn) = 0
0.
d
log (
log
0.55 (n/0)
Limits of Sub-Inequalities (a), (b), and (c):
(a)
log
[300mn
Vk. Ollogn)
vh vlog n VV-llog 10 (n/8)
5
-0
VB
VCn log 20 (7)
Given k S
n (n/8)"
log [n/o) Vlgg
v k vlogn n
log (n)
- 0.
vn- log 10 (n/0)
Vn log 10 (n/o)
log 11 (n/8)
(b)
2k log
400mn
2k Ollog
Vlog n
+ 0.
vd
k log" 1.1
(7)
log
0.55
(n/8)
(c)
V 2nk log
300mn à
2nk Ologg))
Vnk log
0
VBd
C2n log
(3)
k kllgg"
1.1
(3)
Vnk
log 10.55 (n/5))
Vlogn
0
- 0.
log
10.35
(n/8)
4. Bound on Yoo and !/:
2kE C2n log
(#)
MEV V
O (nn/r)
To T 100 log (9)
100 (7)
Comparing to the minimum the two terms:
3mc /
(kev (n
80 log
300nB
and
min
Oknn)..
120
Since no scales similarly to the first term and /nlog 10 (n/8) < n for large n, the inequality "oo S min(.) holds.
Limit of log?
330mmBB
ay
Bd
300mnB
log²
Olloo2(n/8))
8
and
NV nk
ny
nk
dv Bd
1: klo 1.1 (?) V
Can log
20 (7) klog (7)
=0 klog (l (n/8)))kkk log
nynk
=
10.55 (n/6)
k log 11 (m//))))
21

---

Published as a conference paper at ICLR 2025
Given & 2 (n/6)"
71
71
1
log
11.65 (n/6)
log
11.05
(n/6)
log
5.66
(n/0)
log" (x2/0)
Therefore,
300mn
nV nynk
log" (n2)(1/8)
log"
0 (log 1log-.000
4.63 (n/5)) - 0 as n - oc.
dv dvBd Bd
5.66 5.6 (n/0)
6. Limit of Dazi 256mm
As
2
me B
m²=² B
7oo
100 log (3)
10 log" (?)
Substituting into the left term:
m BB2B
me 2 B
7o 256 m 2
256 x 10a log 2
(3)
mn
256 10ah log (4) n²
Substitute 2k and C2n 10 (?):
2ke² 2 log
10 (3)
2kC₂ log (7)
70
256 m n
256 10 log"
(7) n2
256 x 10tn
Setting C2
1.28 10",
2k: x 1.28 10 log (*)
log (7)
To
10 x
256 m
256 x 10 x n
Using k 2 log" (x)
from the assumption:
log (#)
7om
Tog
10 1og (7)
10
256 m
n
Since log (7) > 4 log for sufficiently large n, have:
10
(r)
40 log n.
Thus,
40 log n.
256 n2
7. We aim show that:
722
+
as
11 + 0o.
7/0
Here
2kr / C2 n log (#)
2kc Wice/C2/e /2t////llooooo log" (7)
kc/C2v/n log' (?)
770
100 log (?)
100 log (1)
50
Thus, the ratio is:
Too
722
3
4ke x 50
200
770
Rc /CvC / log'
Obte/Cy/ii log (3)
3/C2v/nlo'' (:)
50
22

---

Published as a conference paper at ICLR 2025
Substituting C2
1.28x10",
1/2
200
200c
170
3/
122x100 10 a²
Vog (*)
x 15800nlog (7)
As co, volog" 4 (3) + 00. hence:
7/2 0.
77
8. Limit of (no -m)= < 5mmz:
mE v B
mev B log (3)
yo + 7 =
100 log (3)
4he
5nmz 5n
Onke))
3
Given VB 0 O (Vn log
mevB
kc log" (n/0) O (oe//))
O
log (3)
log(n/o)
Since nıke grows faster than / logg (og))n/)) for large n, the inequality (no + m1) < 5rep holds.
This concludes the proof.
Proof. We will set our hyperparameters according to Assumption 2 We will first warmup our analysis on linear loss E(g. v)
(-1)"g. We can rewrite the training dynamics in this case
B a+
(1++) (t) - 11
Vwe(T(w (b.lt)))
(1,)) 1])
B
i=n#1
(((() 178 Σ
1
Σ
(-1)6(00)
V Tww (btt.)))
We will train the model end end in three steps and analyze the evolution of all the weight changes. The final results can be shown combining Lemmas 21 and [26] Lemma 27 generalizze the analysis hinge loss. also easy to verify that all the results hold if the parameters has only Oloo n) precision.
B.4.1 AUXILIARY STATISTICS
To simplify our calculations, we define several auxiliary statistics the data. Table 1 provides the definitions and rough orders of these statisties, including logarithmic terms.
Lemma 6. Fixing j. by. by let F(..t.) be the filtration generated by the random variables (b(t.e)
max(n + 1,j} < S * j,t' S t, S B} ordered in lexicographic order, the process
(-1) (666 = bi. (t.*) [1] b2)
a martingale with respect to {F.).
Proof. bli+1] b[i + b[S(i), and since S(i) # j. b[S[i] is independent from b[j],... [J]]]...,[[]]] Thus, E((-1)) (( F(r.,)))
provided ((+1,+))) b1) constant. Hence, this process is martingale.
23

---

Published as a conference paper at ICLR 2025
Statistic
Definition
Rough Order
B
81,i,
(b(c.)) =b)
n 8 B log
[+111
[b(c)))g] = b2, b(t:)) = by)
ai.j,bi
(V BBlog (m)
i2 n n +2 i= n bB
bi.j,
(-1"
(B) when = S(i].
2
0 otherwise
a+ Σ
nB Bl 0
Be.j.b
O
i=max n+ 1, b==0 s
n+* Σ
v Bb Bog log
mn 12
(t,r.j.b
(Vr,i,b2 > 0) ac,isbbib
f=max 1,j1 S(A))
1
(Rrag > 0) (Urar > 20)
v log (3) n
Y.b, b2
11+1 1 (1 =i; b bz)
2 m
2m
my kB log
mnn
(Vrii > 0) Vlr.
Table 1: Definitions and rough orders of auxiliary statistics, including logarithmic terms. Here,
t e 1,2), E [n+11n j [n+ k), b, bı, by {0, 1}. [B], and [2m].
This version maintains the essential details while focusing on the core aspects of the martingale property and independence conditions.
Lemma 7. For each batch, the data inside the batch is approximately balanced. With probability at least 5/10, the following bounds hold for all E [3], i,j E In + k), b, b,b1.2 E {0, 1}, and r E [m] for statistics defined in Table 7
15c,0,l
B log
300n
2B log
300m
(j = Sli))
2Bk log
300m
18r,5,b11 1
n
2,
Bk log
300mnB
n
2m log
200m
4m
kB log
300mn
5
Proof. Fix any t, i, and b. For each [B]. define X, (-1)6(10)(641))
(t,0) b). The
sequence 8 consists of independent random variables with |X,| and mean. By the
24

---

Published as a conference paper at ICLR 2025
Azuma-Hoefiding inequality,
R
Pr (162, R) S exp
2B
Setting R 2B log (300nB) get
Pr (1oc,bbl R) S 150n
Applying union bound over all t, i, and b, we ensure that with probability at least 1 - 8/50, the bound on di,, holds for all choices.
Similarly, for agiiiibb we fix 1, i, J. bi, and b2, and define
Y (-1) []] b,, (( =b2)
Again, {Y.) independent with |Y,| s and mcan b, Hoeffding incquality as before, we obtain the stated bound for at.dj,bb.
b/ /.B Applying the Azuma-
For Br.j and Vi.riid the bound can be derived similarly combining Lemma 6 and Azuma- Hoeffding bound.
For "/.. consider the sum 2m (Pr,,6 > 0) (1r,r b2 > 0). Each term is an indepen- dent Bernoulli random variable with mean (1+ [1+1( ( 1', b b= b2)). The variance of the sum
is bounded by m. Applying the Azuma- Azuma- Hoeffding inequality,
get
R²nn²
Pr (hrar > R) S 2exp
2m
200m
2m
Setting R
we obtain the desired bound.
For Sti,b.. we note that a sum over 2m terms, each involving Ve..g... which are bounded as
2 log 300mn
in Lemma The total number of terms is 2m, and each term is bounded by
B.4.2 FIRST STEP: CONFIGURING THE MIPSS
Lemma 8. For the followwig Equation attention weight stays constant in first step,
i.e., A = (0)
Proof. Because W Wdd+1.2= O, V (b) 0.
Lemma 9, For the dynamics following Equation attention output stays constant in the first step, and
Vt (0,1),i I (0.1),ie|+**+))41 Bbb
Proof. This follows from the definition of the attention module and Lemma
Lemma 10. For an input b. With probability 1 508m at initialization, whether neuron outputs nonzero value at position is determined by Vr.iM
(wo 1. E(b)(t)) > 0) (Vr,2o1) 00).
Further,
kwo Paaa E(b)l() Ur.sMal < <e/100.
25

---

Published as a conference paper at ICLR 2025
Proof. Under Assumption| |
W r.1:
b=U
Because,
V 2kee?
p- bm
By Lemma 30
with probability
50
4log( 100Bm
<w r,1:d: E(b)(t) Vr.i.b() 1
d
By Lemma 2,
4k
log( log(100Bm/0) < 0.01, can conclude that
n+
Pr
Σ 1'=++11
Dea 11' # peeeeteeeeeee [ 2 c/100
50Bm
The proof is then complete.
Lemma 11. With probability 1- - the following E, happens: for any input b (0.s) in the first batch, at initialization, for any r, whether a neuron outputs nonzero value at position is determined by Pr..
((WH) B(b)(() > >0) = [P > >0).
Further,
kw F. E(b)li) Mr.s.bll e/100.
Proof. Apply union bound to Lemma 10 over neuron dimension and data in the first batch.
Lemma 12. When En defined in Lemma gradient on W Vod satisfies that,
happens, for all i' E [ 1, n k 11, E (0, the
d.((0)
hr
n k imn--11
1 (2rasbb > 0)
dW, r,1id
B
b=0
Proof.
(C(())
(- b1bb0011
W. ((b)0)))
dW, WWd
B
i=n+1
h
B n+*
(-1) (1) 1 1 1 1V/,,,.0000001 > 0).
B
$= i=n+1
26

---

Published as a conference paper at ICLR 2025
We can then expand this formula to have
B
a+
Σ
(4 1 [4,,,000001
> 0)
B n+
1(b ((b 11000000
= b)e, 1 > 0)
1=+ 6=0
e.all 0) Σ
1) 0
[+1 1(b (0.) b)
i b=0
a+
Σ
Σ
ei (L) 0)00.in-
1m++1 b=0
The proof is then complete. Lemma 13. With probability 1 0.135, for all i' E Im+ 1, + 11,6% (0, 11, the gradient on Wr,1d satisfies that,
dWaddd" Gi,b )
d(
log (300mmB30))
mv B
Proof. Combining Lemmas 12 and 30 have that for any with probability 1 8/ 1000m.
h
1Vrra > 0700, 01.b )
2mB sup do,2. (1+ (1+1
200m 2k log
6
B
d
+=
0=
By Lemmas and 7
with probability - 115/100, for all r, holds that for all r,
hr
k
(Vr.i. > 0)00 C1,b )
v /log Mog(330mB/66
mB sup 6n..b
B
mvB B
imn+l
6=0
Combining with Lemma and applying union bound concludes the proof.
Lemma 14. For and to defined in Appendix B.4.1 When E defined in Lemma i' E [n + I, + k + 11,b6 E 13, the gradient on Wr.41:24 satisfies that,
11 happens, for all
d(
hr
n+
h
(Urri > (1
Vo Cj,b1
dW, [Wr.++1:2
B
B
1= b, = 62=0
j=1 0
Proof. With probability 1 8/50, &, defined in Lemma happens, and
d.
B
n
b1000)1)111111 V W. T (b000)(
1:2d
h
n+t
( Mi] (0 > 0)
B
i=++1
; 0) Σ ΣΣ1 1bbl = 6) ej,0
B
1=++1
j=1b
Rearranging the summation, and use a defined in Lemma 7.
()
h
n+tk
Σ
( > 0)a0,..,bb..
dW,
d4124
j=1
mn
27

---

Published as a conference paper at ICLR 2025
Breaking the summation,
dC
hr
atck
(Unid > 01
dW..441:2d
B
t=#t1 by =0 by=
n+ +
Σ
Σ
[Vnas,0 > S..0.
jjbb
7
1=1
b
iε(max +++ 1g) -k). b0=
This simplifies to dC
h,
n++
h
n+
1Pr,s,b2
Do.rd Cjp
dW, n.d+-222
b
The proof is then complete.
This leads to the following bound on the gradient, Lemma 15. With probability 1 - 0.166, for all i E [n+*+1],bb (0, 11. the gradient on W 1..41:22 satisfies that,
d
64, )
3 2mn
$
dW dW, rd4+122
Proof. Denote
n
hr B
hy
G.
1 (Ur,i,02 > 0),,,, (1 +
Vo..rjj Sj
B
1=n+1 b1 b1=0 b2=0
j=1
b1=0
By Lemma 30 for all r E [2m], i E [n + + 1].be E 1), with probability 5/25,
n+
(G,, e.,iil S
1
Σ
(Ursi > 0) [SS(V] =) + vo.n.e
2mB
V
2oog 400mm)
a+t
m+*
ΣΣΣ
1 (Vr,i > St
2
2mBvd
j.rij
By Lemma we have with probability 8/10,
2/
Bk
Bllog log
300mn
sup uo. S
sup S/ S B/2.
n
Combining with Lemma 5
n+
Σ
Σ
1Vrrib O) 1(S =
<
1
2mB
2mm
w=n+1 62=0
2 2//k log ( (300mBB]
2mB |Wo.m..l I
mnyB
200mn
2 log( 100mn
2k log( 100mm
Σ
ΣΣ
1 12
> 0)aa 0)
1
2mB
2mnyd
200mn
0b
2 log 100mm 5
n+k
2 2nkloog 300mn
2 Orrib.
2mBvd
mnv BBd
200mm
#=1 0
Further by Lemma proof.
14 we have that
dc
G, with probability - 55/50. This concludes the
dW. Wi14+
28

---

Published as a conference paper at ICLR 2025
B.4.3 SECOND STEP: CONFIGURING THE ATTENTION
Lemma 16. With probability 1 1-22,,, the following &2 happens: the change in the neuron outcome is small after the first step for any imput b,
sup
( W A(4A 3c/80.
This leads to the case that for any imput (l1,)) in the second batch, for any r, whether a neuron outputs nonzero value at position is determined by Vrri
((W7, r,lds B(b)(t) > 0) (Vr,i,00 > 0)
Proof. By Lemma with probability 0.136,
dC(0)
v vlog(300nB15)
sup (
1.ldd W ((b) = 170 dWW,1:a''
E(b)
mvB B
By Lemma 13 with probability 1 0.165.
dC
sup (W1
7:41:22
W
7. 1:20 A(A (EA)(() 1j0
dW r,4+1:2d
A[A (E()))
S 2mn
Hence, no0 I min
B
(Lemmal
3), conclude that
80 100(300mB/8)
sup
kW, W A(A))
3e/80.
The proof is then complete.
This shows that the gradient of MLP on the second and first steps is the same in distribution. Lemma 17. With probability 1 0.586, inequalities in Lemma 7 hold and for I E {0, 1]. Vi [n + k]. 4b E {0, 1}
dC
hr
mork
1(r.a > 0)0
dW..1.d r.1:d
B i===1
b=0
dL(())
hr
n+*
hy n
Σ
(Uraih 0) espd,, +
Σ
Derihhh Cj,b.
dW, 1d2222
B
1= 0,=0 by=0
j=1
b1=0
d((()
3
dW r..41:22d
2mn
Moreover for any input b,
sup <W; W., [VB < 3e/40.
Further,
(( We P.lid: E(b ((I( >0 = (1 >
Proof. Consider the backward caleulation of MI.,, it decided by the input and the corresponding activation of neurons, which jointly follows the same distribution in the first and second steps. The proof similar to Lemmas 12 14 and 16
We will now show that while MLP changes little, it is enough to provide high-quality gradient information to inform the attention layer. To this end, we will define the following terms,
2m
Σ 1 (Pr.L > 0) h, (t)) 1:24
Ke,6,
A (Ke,.... C3,,,)
Intuitively, the FFN with weight (1) maps attention output & to -K4,b(j at position i.
29

---

Published as a conference paper at ICLR 2025
Lemma 18. For a, 3,,, S defined in Appendix B.4.1 ? laddb b,
when the event in Lemmaa 17 happens, for any
K1,ib
7
br +
Σ
Σ
Yi,baa b2 (2, do (5(P9),
8Bmi
4Bm²
d'=++ b2=
n++
TJo
4Bm²
b,=
Proof. When the event in Lemma
happens, we have that
dL(0)
hr
ntck
h B
n
(v > by
+
Do,rij.b ojbb
dW, d,d+1:2dd
B
d'antib=00ze
j=1 b
Summing over the axis of r as in 1s,
7o
n+*
Σ
Σ
2m
(Uras 0) ( > >0]
1
K1,ib 4Bm²
v'=++1
bı 00 00
n+k
Σ 1 (Vrsi.b > 0) No.rsily ))
2m
To
+
ejub
4Bm²
01=0
Recall that Ti, have
E 2m
(Prsigh > 0) (Vr.3 > 0))
1-1(i=i', b=bb) m.
We then
7/0
n+k
1 62)
K1,.b
āv S[r) bo esp
4Bm
2
#=ntl 01=0 b==0
+
na,.s' by by 0s
4Bm²
Mmn++
bg 0
n
2m
(Mrai > 0) No.r.... )
Σ
Σ
Ejb
4Bm²
j=1 b,=0
The first term can be greatly simplified, as
04.8k/), + aze s(e)),, 0.
This concludes the proof.
Lemma 19. With probability 1 0.68, inequalities and equalities in Lemmas 17 and for all i, b2,j,bb by.
18 hold,
A Aliieas
8Bmi S1 1 = S(ii)
To
10rjo V mn Bd
log (300mn B/5))
Proof. Define
Yo
H
8Bmi
b
n+t
Bi (111),, este1,6
4Bm
F d'ntle by m0 00
TJo
Σ
Sobai ejbb
4Bm²
b=1
30

---

Published as a conference paper at ICLR 2025
As the inequalities in Lemma hoold, by Lemmas and 30
with probability 8/50, it holds that,
k
100
, V==++1 1 by =0 by=D
4 2
Sosupp suspusuppy
1 llg(2000/0)
1 1+ 1V
4Bm²
10
2k log
200n md
4mn
and,
n+t
90 4Bm²
1/0 S 4Bm²
log (200m//
1+41
d
1/0 mn
log(200n/8)
+4 dB
nk log (200m/5)
B
and,
2Vlog(200n/0)
110 8Bmi
110 8Bmi
a Sk..,, S(i))
8mnvd
With d and B nk (Lemma])
5 we conclude that,
10ro Vnk
[Hine
8Bmi C...
Sti)
log (0000mmmB/
mn
Bd
By Lemma 18 we have the desired result.
Lemma 20. Under the setting of Lemma 19 for every b E (b
s E [B]]
(d) bii 11) = (-1106++
dsoftmax (E(b)TAA(()) AE(b)) [..
(
dA
0A
with
Ōsoftmax (B(b)TAAE(b)) b. softmax(Z)]i.))
E(b)
Σ
softmax(22jp,
E(b
0A
Proof. Our goal to calculate the gradient of the loss with respect the attention layer. We will first calculate the gradient of the loss with respect to the output of the attention layer, and then use chain rule. When the event in Lemmal Thappens, for any binary sequence the second batch, we have that for any i,
dT (A(E(b))
2m
h w
((wL)
E A
(
[1) >0)
(b)
2m
h (1) Σ
r.d41:2d (Vriibt) >0)
This then implies that,
de (T(())i)bb
(11))((+11)
Σ
2m
h W
1.411201 (Mr > 00
dA(E(
31

---

Published as a conference paper at ICLR 2025
We now calculate the gradient of A (the attention matrix) to the attention output. (oofsso
[A(E(b))
((())))()))) [
0A
0A
Σ
EE(b) softmax (E(b)* b.
0A
Osoftmax (E(b)TAE(() (E(()) [i,
Σ E(b]i
04
Hence, the gradient of the loss with respect to the attention matrix is dE(T(b)(i) bli 1])
dA
de(T (b)
D 1)) dA E(b))
E(b @softmax (E(b)TAE(b) AE(b)) [i,
0A
h (M > 0) (W r.d+1:20 E(b) Osoftmax (E())*AEQ AE(b)) [i,
=( ()1))(+11
0A
f=1==1
Osoltmax (E(b) (B(()) [
((-1))(+1]++
(K1,ab]) 9ybb1)
0A
We will use 7 denote E(b) AE(b), and calculate the derivative of th+e softmax function applied to Z with respect to A.
Osoftmax(Z]) (Z)lj.d
0
e A
0A
0A
E p==
e 21AA
Ep
2lpiil
- 215.0 E/
BA
8A
E ebbo.l) 2
Given Z[p. i| [[" [E)) the derivative with respect is:
az pp. 4]
= B(b)(e) E(b)
0A
Substituting the derivatives into the quotient rule expression:
(ooftmax(2)]i.)
7bM
E(b)) EjE z'ail 2 EE Anill E(b) E(b
0A
EE cZ(o.s)
Thus, the derivative is:
Osoftmax (E(b) AE(b)) b..
B(b) softmax(2]lj.))
(otm p)] E(b)
0A
The proof is then complete.
Lemma 21. After two training steps, the attention layer will show the following structure. With probability 0.76, for all i, b, V. bz.
32

---

Published as a conference paper at ICLR 2025
1. Jf S(1".. then (ey,0z A ,b)
256mn
2. 1f + S then (e M A (2)) ex .B2)
1071
512mm
Further, events in Lemmas and hold.
Proof. By Lemma 20
@softmax (E(b)* (o(())((b))))) [5.
/ B(b)bb - Σ B(b)(p) B(b)a1
0A
can rewrite the gradient as
de( bii 11)
(-1)
Σ (iS) B(bL)))(() E(D
dA
8
1) b(++ (K..)()) B(b) (b) ((
12
p=1
we define
Σ
((b bi, ((.)) b2, bit.)) 'p] -ba)
Pt..i.p.hh
Summing over the second batch, we have that
(T
ΣΣΣ
Oniij Aragg dy Cy c.
dA B 1=
j=1 11:0000000 b2=0
n+k
Σ3ΣLEΣ
B i=
This implies (2) is updated as,
(1)
2 - 71 dA
n
Mii by Anssh ei,0
B i=n+l
=0 by=0
att
?
Σ
B
Recall that our goal is to calculate MM A (2) M. We can then calculate the contribution from each term by separating the calculation follows:
33

---

Published as a conference paper at ICLR 2025
Т₁ B
E
Σ
7
(eys Sg ev
m iE|maxx + nt by (0.1)..11(+b1
n+*
R
7
a1 a.s. Am.ss estjb,)) ex Cv,6,)-
B
i=++1 01=0 62=0
n+k
7
Σ
ΣΣΣ 12 M.j.bbbbb Arijb teg. Cxob (Cppba cp 62,6))
B
i=max (a+ w.t.spsssssllll........ ST
n
V
171
Σ
ΣΣ
2 Ut.... Ausii (ey eg (eniaa ev1
B
immax (x+ L,)
n++
W
71
Σ
ΣΣΣL
by Ae (3,,, CrM,))
B
i=ax i++ +.5) by 0bg 0b
n+
ΣEΣ
Σ
p M1.. bs.bs ba (e (er (()))
B
b1=0 by=0 by=
Then we have that
n
k n-th
A
W
Σ V +Y p=1
j=1 p=1
.B Lemma 31 with probability 8/100,
log( (8000dd/8)
Kei,b ey (g) 1(i &&& = 62) s 2
We will also assume the event in Lemma 19 holds.
We will now discuss each term separately
1. For Tj, by Lemma 6 and Azuma- Hoeffdngg bound, with probability 8/55,
Σ
Σ
0 F1Adbe A (egrog, ey (ea ev (r1)))
|E|max(n (nnaxx+++.j))) +1J) mt (0000..))) 60
100n
B log
Σ
A M.iiii (err,, (lerb) ey xz))
5
01=0
iE fE(max (n+ J. n+ m b (0.1).
2,
log( 100n
sup
A by
Σ
Key, G.aa))
Σ
((ei, ep va,))
V
b,=0
(1.1 ((1))
Now by Lemma 19 we have that
10nj V vk
sup A A
log (300330mn//5)
/ Bd
Further by Lemma 31
8k log( (og(100nd/6))
Σ
(ferbn Cv m)
+15 114. 4.
d
iE 1...
log log(100nd/5)
Keyy, ((bza)
$1(j = &&& = b1) + 22
d
34

---

Published as a conference paper at ICLR 2025
[Tyl -
40m Vnk
log (300mmB3))//)) (4V
logí 100nd/0)
+ +1 = 53)
mn?
Bvd
d
Summing over j. by Lemma
we have that the contribution bounded by
160m '70 mn2
n/ 71 2000mn
log" (300mnB/0 //
(3)
dv Bd
For R. we will directly put everything to an upper bound. We will discuss three cases,
Ar", SKr" = j. by Lemma 19
16km log(100nd/o) supa
"n S1 A (te)ss oge es()),)
R S
Bnd
16km log(100nd/0) B no
log(100nd/0)
711 B no
Bnd
4mn
Bn
4mn
Bonk log(100nd/0
Mom
log (000nd88))
70 2000mn
mnnd
4mn?
d
3i" S"t" = j + i', we can get similar bound,
707h
R
2000mm
S|1' = j', in this case, we can show that,
707 8i STP),b
Th 2000mn
R
This concludes that
no7h 2 mi² ( b,, by 1(S = 5¹) -
77
(4)
2000mn
3. Ug.p can be bounded in the same way as Ty and we have that,
ΣΣ UpP Ug
771
(5)
1000mm2
4. For Vp,, WG, Y, we will directly put everything to upper bound similar to the bound of Rt we have that
n+*
7h 7o 1000m
V
(6)
71 1000m
W
7) 9
Y
<
1000mn
Summing over Equations 3 to
8 we have that
Ker,M, 1() M/
707h
22
7/ 512mn2
8Bmi'
Further
7
nom
2
8 822mi
128mm
Hence, can show that
35

---

Published as a conference paper at ICLR 2025
1. j S[t'), then (ey,a) (Ree,
256mm
2. Jf + Skr then (eg.n A
512mn
The proof is complete.
Lemma 22. Under the setting of Lemma 21| the attention output is approximately one-hot after the second step, with
[ Mspi < 1/n"
Proof of Lemma 22
This is a direct combination with Lemmas 5 and 32
B.4.4
THIRD STEP: MOVING MLP LINEAR REGIME
Lemma 23. With probability 0.86, the gradient of the FFN layer on the third step can be written
as,
2
log(300nB/8)
Ei,
dW,
mvB
B
(.((2)
n+h hr
es + +0 (Vr,i,b >
dW,,2+1:2d
n
0b2
Further the event in Lemmas 27 and 22 hold.
Proof. The proof is similar to Lemmas and 14 switching the original output with the near one-hot output of the attention layer.
Lemma 24. With probability 0.96, the attention output is almost unchanged after the third step, with
A[A (E()))l
Est Essl 1//n"
Further, events Lemma 23 hold.
Proof. By Lemmas 5 20 and 22 A - (2) is of order 1/n²", this implies that the attention weight
and output is almost unchanged.
Lemma 25. With probability 1 0.918,
2c (wra rd41::2d A(A*) [E()][])) + sign(M,. (vra,o0y < 0) <
c
300
Further, events in Lemma 24 hold.
We will first calculate projection of the gradient on espaay When the event Lemma happens,
Proof.
23
2
n
= hy
Σ
Σ
(Vr,a,0 > 0))(1))(++b (esp)r espe,o) + +
dWyd4112d"
n
imn+ =0 00
By Lemmas and 30 with probability 1 0.016.
K dW :d41:2d (sp]) - sign(u h, (Ves <
(2)
<0)
200m
By Lemmas and 17 have that,
(2)
rd41:22 estR],M sign(Vr.s v me < <0)
(no m))
122
172
20mn
200m
100m
Combining with Lemma 24 and 1/2 3 have the result.
36

---

Published as a conference paper at ICLR 2025
Lemma 26. With probability 1 for all b {0,1) i [n 1, kj, we have that
T(E()))] <(-1))bi-1}+
3 3
Proof. By Lemma 23 the following holds,
(w(())
Filaa E(b)(() sign(ur...))))
Combining with Lemma 25
(
(w.e) 1:2+1:24 A(A*) [bb]]l + sign(yr.sMs1//h))) 1 < <0)<
2e
<
c
300
Hence, we still have that
(w
(3) 7,411:2:: ( A(A)) E(b]I))
5c
6
which implics that,
1( 1((w1)) r,l:d E(b)(a) > (( 0).
Consider the output contribution of the attention part,
2m
hp
W
(3) r.d11.
((A)) B(b]lit) 1 0)
A(AAAAAAA)
2hE sign (ign(( M,,,,, > >0| 300*
The later term with 8/100 satisfies,
2m
2h,
sinn(s..ossglhh: < O, UyH.bb > 0)
3
2m
c
sign (ggn < 0, (,M..ibi)) >
3m
2m
E (-1) b([ti] 1 <0)
3mm
log (00((00n//)
b| + O
vm
This shows that,
2m
Σ
hy (W (3)
A(4)) (
1 >0) -0)- (
<
150"
On the other hand, we have that
2m
hr
(W(3) m,la
B(b)[i) I (Vr.i.(()) > >0)
(W:(0)
E(b)(a)) (V > >0)
+ Σ
h,
(W-3)
T.1: B(b)(i) (660))((111))) 1( 0)
r.1:d
37

---

Published as a conference paper at ICLR 2025
The first term is bounded due standard concentration inequality over the axis of The second term is bounded by Lemmas 17 and 23 Combining the terms, we have that
c(-1) 1 +(([+++++
c(-1))(++1)+
T(E())
T(E(b))
3
3
3
This concludes that the model is able to predict the correct output.
B.5 FINAL PROOF
Lemma 27. The results in Lemmas 18 and 26 can be extended to hinge loss ((..v) v)
max((-1)"j 1.00 with the same probability.
Proof. Our Lemmas 17 and 23 shows that the output of the FFN layer bounded by c/3 throughout training. The hinge loss is linear when the output is in [-1,1]. Hence, the results can be directly extended to hinge loss.
Proof of Theorem The original theorem is combination of Lemmas 18 26 and 27
B.6 TECHNICAL LEMMA
Lemma 28 (Theorem 4.4.5 Vershynin (2018)). There exists universal constant C.Iet be an m x random matrix whose entries Aij are independent, mean zero, lub-Gaussian random variables. Then, for any >> O, we have
|||| CK (Vm vn vn++1)
with probability at least 2exp(-1²). Here = maxi,j || |1 is the maximum sub- Gaussian norm of Aij-
Lemma 29 (Theorem 4.6.1 Vershynin (2018)) Let A be x n matrix whose rows A, are independent, mean zero, sub- Gaussinnn isotropic random vectors in R Then, for any 2 0, we have
- mm-CK3(Vn+t)) (A) < s1() S vm + CK* with probability at least 1 - exp(-t)) Here, K = max; || Furthermore, a slightly stronger conclusion holds:
1 ATA A
In S 22 max(6.82).
m
where 8 = (V 1 Vm
m
Lemma 30. Define
Eb di.bii, then with probability 1 6,
2 log (2)
K) Av
d
(L
Proof. We aim to bound Ker,y, v)), where
m=n++1
hm
{-12) 1
Note that each 6g,b E IRd has entries that are independent random variables from
V2
val
First, observe that
(er,M.)) = Σ Neplev ex.)
Ne +
Σ
Acolerss,
(8.b)
(1,6)
38

---

Published as a conference paper at ICLR 2025
Since (er ev.v) = 1 and for (i, # (2', the inner products (ev,v, e.,b) sums of indepen- dent random variables with mean zero.
Define
X = u for (i, + (d',6))
Each X1 is sum of d independent random variables bounded in Az
because each
component j is 2 By the Azuma- Hooffdigg inequality, for any 0,
2dt
dt
Pr
Σ
21
< 2xxp
=exp
(i 4X2,
2 (1.B)
2
X2,
Setting the right-hand side equal to 5 and solving for I, we get
2 log
Σ N2.0-
d
(L,D)+(((
Therefore, with probability at least 1 - 8.
2 2 (3)
Key,,, Av
22.00
The proof is then complete. Lemma 31. With probability 1 5, it holds that
1(i = i'& = 6')ll2 - / Log((rud/))
vi, E in + h3,6,,b' E
d
Proof. This is direct consequence combining Lemma 30 and union bound.
Lemma 32. If for constant C, the attention score before softmax has the following property:
for each position the targes position j =S(i] satisfies:
(CA.,,, A19) 2C log n.
For all other positions j' + J, the attention weights satisfy:
A (9)
Keym A A
s Colonn.
The attention output is approximately one-hot, with
(E - essaasii 4/n
Proof of Lemma 22
The attention output for position is given by:
([()\( = X softmax (X AX) [+]
where the softmax is applied column wise.
Given the condition from Lemma
we set
= 20 2ogog
This implies that:
A/2 < CCC-Coogg Clogn
n
39

---

Published as a conference paper at ICLR 2025
Define 2y' = (e A e.,,),
The softmax for the target position Sli] is:
e
softmax(z2; = 6:S + Eg ex
Given that 2j 2 A and |zyr S 2
(since
7o 512mn2
A) ) we have:
2
A
2j 2j
2 2
Thus:
e
1
softmax(2); 2 c + (T - 1)
1 (T
1)e
A/2
Therefore with < 2n:
2
softmax(2); 2 1 (T -1-1). 1) n
2 1
n
Therefore
softmax(2)3- S nC-T
J'
The attention output for position i is:
([()((
Σ @:b)ssoftmax(2);
Substituting the bounds:
(Bbb)() este)))
ej My)soffaax(2)) + (ssttaass) - 1)
S
$ softmax(2); (2) + |(softmaxx))s5p - 1)/
s nd-1
3 S
We conclude our proof.
40

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENT RESULTS
In this section, we provide details of the experiment setup and present additional results. All training was conducted using PyTorch Paszke al. (2019) on NVIDIA RTX AI0 GPUs.
C.I EXPERIMENT DETAILS
The transformer architecture adopted in the experiment section is based on the GPT-2 model (Rad- ford et al.) 2019) with a hidden size of 720, an intermediate size of 3072, and trainable position embeddings. For all experiments, we use Adam. (Kingma| 2014) optimizer with random initializa- tion, using hyperparameters 31 = 0.9, 3- = 0.999, a weight decay of 0 and a linear decay learning rate schedule. The batch size is set to 512 throughout. The validation set contains 2048 samples which are nonintersecting with the training data. In all experiments regarding sample complexity, a test set of size 2048 which is non-intersected with the training data is used.
This section examines the normalized attention entropy of Qwen2- 7B (Yang et al. 2024) and Owen2-Math-77 models (Qwen 2024) on the GSM8K dataset (Cobbe et al. (021) with and without CoT prompting respectively. Both models consist of 28 layers, each with attention heads. The normalized attention entropy is computed as the average across the GSM8K test set. In Figure 6 the entropy values of each attention head in a layer are sorted separately under the three setups. While Figure 6 presents the normalized attention entropy for specific attention heads in the first, last, and two intermediate layers, Figure 7 shows the entropy across all layers.
Details of the experiments in Section 4.3
For With CoT data, the input is concatenated with the ground truth answer from the GSM8K dataset. For With COT data, we extract the final answer from the ground truth, and concatenate the input with the string "The answer is [Final Answer).".
C.2 ADDITIONAL RESULTS
In Figure [ggureBBiigee (iight), we present the attention pattern of a single-layer, single-head transformer trained on the (n 20, 40) parity problem with CoT data. Figure we show the attention patterns of a multi-layer, multi-head transformer trained on the same problem. We observe that in the first layer, the attention pattern is sparse and interpretable, with each secret variable attended to by at least one attention head. In contrast, the second layer exhibits almost uniform attention pattern. A possible explanation is that the first layer captures sufficient information, which is then transferred to subsequent layers via the residual connections.
In Section 4.2 we observe that training with repeated data can help Transformers learn the parity function, but it still requires significantly more computation compared to trained on CoT data (Fig- ure 4). To further substantiate this observation, we explore the training of models with and without CoT on the (n 20. k 12) parity problem. Compared with the (n 20, k 6) parity problem considered in Figure 4 this problem is harder for the models to learn as the number of secret vari- ables & is larger. We examine various configurations: the number of layers and heads ranges from 1, 3, 4, 6, and 8; the learning rate varies from 6 x 10 8 x 10-8, to 1 x 10-4; and the training dataset size varies from 10², 10 to 10", with corresponding epochs ranging from 1000, 100 to 10. Across all configurations we examined, training without CoT fails to achieve non non-trivial accuracy (Figure In contrast, models trained with CoT achieve perfect evaluation accuracy when trained on a dataset with 10, 000 samples for only 6 epochs, or with approximately 60, 000 fresh samples in one-pass training setting.
4-layer 4-head transformer achieves perfect evaluation accuracy (n 20, k 6) problem when trained on a small dataset of 50000 samples, but fails to achieve non-trivial accuracy when trained on a larger dataset. Furthermore, successful learning coincides with significant de- crease in attention entropy, indicating the development of sparse attention, while entropy remains high when trained on larger dataset. In Figure [00 we present more results across different chitectures (1-layer 1-head, 2-layer 3-head and 4-layer 4-head transformers) with dataset size of 5000, 10000, 50000, 100000, 1000000, and observe the same pattern.
In Figure 5
41

---

Published as a conference paper at ICLR 2025
Nomallzzd Altenton Entrop
Layer
Layer 4
0.8
0.4
0.6
0.6
0.4
0.4
0.2
0.2
0.2
0.2
0.0
28
0.0
28
0.0
0.0
28
head Layer
head Layer
head Layer 7
head Layer 8
Normalize Atenteen Entrogy
0.3
0.2
0.1
head Layer
head Layer 10
Normalzed Attention Entrapy
Layer 11
Layer 12
0.2
0.2
0.3
0.2
0.1
0.1
0.0
28
0.0
26
0.0
0.0
28
head Layer
head
head
head
Qwen2-78-
Normalize Nientoon Entrapy
Layer
Layer
Layer 16
0.4
0.3
0.3-
Qwen2-7B+
With CoT
0.2
Owenn--Mat-78
0.2
0.2
With CoT
0.2
0.1
0.1
28
0.0
head Layer 17
head Layer 18
head Layer 19
head Layer 20
Nom allaad Atrention Entropy
0.3
0.3
0.2
0.2
0.1:
0.1
0.0
28
0.0
28
0.0
0.0
28
head Layer
head Layer
head
head Layer 24
Nommiized Attention Lntropy
Layer
0.3
0.20
0.100
0.2
0.15
0.075
0.04
0.30
0.050
0.05
0.025
0.0
20
0.00
28 0.000
0.00
head Layer 25
head Layer 26
head Layer 27
head Layer 28
Nomalzed Abontten Ertooy
0.08
0.2
0.06
0.04
0.10
0.02
0.05
0.00
28
0.60
head
head
head
head
Figure 7: The normalized attention entropy of the pre-trained Qwen2-7B and math-specialized Owen2-Math-7B models on the GSM8K dataset with and without CoT prompting (Section 4.3) Each bar represents the entropy of an attention head.
42

---

Published as a conference paper at ICLR 2025
Attention Pattern
Layer 1
Layer
Head www.bhoooooooeeeeeooo.oooo
Head 2
Layer 2000
Layer
Head
Head 2
Figure 8: The attention pattern learned by 2-layer 2-head transformer on (n 20, 6) parity problem with CoT.
accuracy
10000 samples 1000 epochs 100000 samples x 100 epochs 1000000 samples 10 epochs
1.0
0.9
Layer Number
Layer Number
Layer Number
0.8
0.7
0.6
0.5
8
4 6
8
6
8
4
6
8
0.4
Head Number
Head Number
Head Number
Figure 9: Training with CoT on (n 20, k = 12) parity problem fails to achieve non-trivial accuracy under different configurations.
43

---

Published as a conference paper at ICLR 2025
layer heads
layers 3 heads
4 layers heads
5000 training samples
12
1.2
Accuracy
0.8
Accuracy
0,9
0.8
Accuracy
0.9
0.8
0.8
Loss
0.8
Los
0.8
0.6
0.7
0.4
0.6
0.2
0.6
0.2
0.6
0.2
0
10000
20000
10000
20000
10000
20000
Steps
Steps
Steps
10000 training samples
1.2
1.2
1.2
1.0
1.0
Accuracy
0.9
Accuracy
0.9
0.8
Accuracy Los3
0.9
0.8
LO5S
0.6
0.6
-0.6
0.7
0.7
0.7
0.2
0.2
0.2
0.5
0.0
0.5
0.0
10000
20000
0
10000
20000
10000
20000
Steps
Steps
Steps
50000 training samples
1.2
1.2
1.0
1.0
1.0
1.0
1.0
1.0
Accuracy
0.9
Accuracy
Accuracy
LOSS
0.8
0.8
0
0.6
0.6
0.6
0.7
0.7
0.7
0.4
0.4
0.5
0.0
0.5
0.0
0.5
0.0
10000
20000
10000
20000
10000
20000
Steps
Steps
Steps
100000 training samples
1.0
1.0
1.0
1.0
1.0
Accuracy
Accuracy
0.8
Accuracy
0.8
0.8
0.8
0.8
0.6
0.4
0.4
0.4
0.6
0.6
0.5
0.5
0.5
10000 Steps
20000
10000
20000
20000
Steps
Steps
1000000 training samples
1.2
1.2
1.0
1.0
0.9
0.9
0.9
Accuracy
0.8
Accuracy
0.8
Accuracy
0.8
Loss
0.8
0.8
0.6
0.7
0.7
0.4
0.4
0.6
0.2
0.6
0.2
0.6
0.2
0
10000
20000
10000 Steps
20000
10000 Steps
20000
Steps
Evaluation Accuracy
Training Loss Evaluation Loss
Figure 10: Transformer trained the (n 20, k 6) parity problem without CoT with different sizes of training dataset and fixed number of iterations.
44

---

Published as a conference paper at ICLR 2025
Layer 1
Layer
Layer 1
Norma lized Attention Entropy Normalized Attention Entropy
Normalized Attention Entropy Normalized Attention Entropy
1.0
1.0
1.0
50000 training samples
50000 training samples
0.8
0.6
0.6
0.6
0.2
0.2
0.2
0.0
0.00
0.25
0.50
0.75
1.00
0.0
0.00
0.25
0.50
0.75
1.00
0.0
0.00
0,25
0.50
0.75
1.00
Iterations
1e7
Iterations
Iterations
1000000 training samples
1.0
1.0
1000000 training samples
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.4
0.0
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
1.00 Iterationn x Batch Size le?
0.25
0.50 0.75
Iterations x Batch Size
1e7
Iterations x Batch Size 1e7
layer 3-head transformer
(b) 1-layer 2-head transformer
Figure 11: Normalized attention entropy curve of transformers training on the (n = 20, k = 6) parity problem without CoT. Each line the graph represents the attention entropy for head of certain layer.
45


Reviews:
[{"review_id": "dFGbULhlqW", "reviewer": "Reviewer_vjNa", "summary": "The paper studies how Transformers may or may not solve k-sparse parity problems. Without chain of thought, the authors provide a lower bound for one-pass algorithms with bounded memory, showing that exponential in k samples are needed, even though small transformers exist that solve the task. Then, the authors show positive results for learning with a CoT prompt using minibatch SGD, where order n samples are sufficient (where n is dimension, and up to log factors). The theory is complemented with extensive numerical experiments.", "questions": "* could you comment on how the lower bound compares to other known computational lower bounds for parities, e.g. based on statistical queries, or in the spirit of Abbe-Sandon?\n\n* could you clarify why the initialization on W_{r,1:d} in Assumption 1 is helpful, and whether it is needed? Is this what leads to \"stronger correlation with the embedding ...\" (L273) ? The phase 1 proof sketch would benefit from more clarity.\n\n* note that O(n) learning of parities with fully-connected networks (thus with no sparse attention) is also possible when the staircase property holds (e.g. Abbe et al 2023, 2024). Is there a link between the CoT structure you consider and the staircase property?\n\n* the proposed CoT prompt seems one of many possible choices, e.g. one could also have a length-n sequence that incrementally computes sparse parities up to each number. In a way, your way of encoding the problem almost requires that the user knows the secret set of locations, which seems pretty close to just giving the set as an input, which is a lot of information... Could you comment on any other possible choices, and why you chose this one?\n\nminor typos: L247, \"for any i\" should be for any i, j and b vector?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The results are interesting and novel, tackling the important question of training dynamics of transformers and the role of chain-of-thought prompting.", "weaknesses": "A couple of minor weaknesses (see also questions below):\n\n* the lower bound is claimed to be about sample complexity, though it feels more like a computational lower bound. Perhaps this should be clarified.\n\n* the initialization seems a bit non-standard, particularly W_{r,1:d}. Some more justification would be helpful.", "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": "cyOs8fstZ6", "reviewer": "Reviewer_P7hs", "summary": "This work aims to understand how Chain-of-Thought (CoT) improves sample complexity for transformers. In a setup of parity learning, it provides theoretical results of exponential sample complexity without CoT while the complexity of CoT is linear. Meanwhile, the training procedure with CoT induces sparse sequential patterns in attention, the rising of which is then empirically observed to be strongly correlated with perfect predicting in both settings of with and without CoT. The sparse attention patterns are also observed in finetuning on real-world data with CoT.", "questions": "Please see the above concerns.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The motivation is good: understanding CoT is a significant direction to understand pre-training and finetuning LLMs, especially for reasoning tasks.\n2. The theoretical setup is natural: parity learning is a popular setup for understanding neural networks and gradient methods, where some previous works were focusing on sampling complexity for staircase property in parity learning. Meanwhile, the sparse property of parity learning naturally brings the sparsity in learned attention.\n3. Both theory and experiments clearly provide a exponential gap between the sample complexity with and without CoT.\n4. The observation of sparsity in Qwen models on GSM8K is interesting, especially the math model finetuned with CoT.", "weaknesses": "1. In definition 2, the augmented sequence is with any permutation $S$. Is this permutation random or fixed during training? If it is random, it seems impossible to achieve perfect accuracy for next-token prediction in Theorem 3, and the monotone pattern in Figure 3(b) might also be questionable. If it is fixed, some clarification is helpful.\n2. It would be better to include discussion regarding a recent related work [1].\n\n[1] How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad. Abbe et al. June 2024.", "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": "N1c22v3y1l", "reviewer": "Reviewer_3yRE", "summary": "This paper theoretically studies the training dynamics of one-layer Transformers learning sparse parity problems using zero-shot CoT. The authors especially compare the sample complexity with or without CoT, which is interesting and verified by the experiments. The experiments and discussion on multi-pass training are also interesting.", "questions": "1. Can your analysis be extended to few-shot CoT?\n2. In lemma 22, line 1834, is $W^{(3)}$ in the third phase of the training? If so, then there is only one step in each phase? Does it imply that the training process is too fast or simple? This seems unrealistic.\n3. This work seems to remove positional encoding. Why? If I add positional encoding, can the sample complexity be reduced in both cases, i.e., with or without CoT?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and easy to follow.\nThe experiments are overall solid to support the theory.\nIt is novel to study parity problems using Transformers. The theoretical analysis is novel. The comparison between CoT and without CoT is quite impressive and significant in future research.", "weaknesses": "There are two major issues. \n1. The theoretical proof seems incomplete. I failed to find proof of Theorem 3. I can only find a proof sketch and some supportive lemmas. As a theoretical paper, this harms the soundness.\n2. If my understanding is correct, during the training and evaluation, the location to compute the parity in every sequence is the same, e.g., we always use the location 1,2, and 5 to compute the parity. I think it is a stronger and more reasonable result to make the location not fixed. For example, let $k=3$, then I can learn the sparse parity with any three locations. Especially during the testing, I can compute parity on a certain tuple of three locations that never appear during the training. This can make your work a stronger CoT inference rather than more of a training analysis.", "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": "nIX7AmrVft", "reviewer": "Reviewer_ezDX", "summary": "This paper presents a theory to explain why Chain-of-Thought (CoT) helps Transformer models to learn complex tasks. The paper first shows that a Transformer architecture with one layer and one head is expressive enough to represent the parity learning task. However, exponentially many samples are needed to solve this task with a subquadratic memory. Then, the authors demonstrate that a reformulation of the task with CoT can be solved with only a linear number of samples.", "questions": "In Definition 3, is there any advantage to taking random embeddings vectors, instead of fixed orthogonal vectors? From a modelling perspective, I don’t think random embeddings add much, and taking them fixed would greatly simplify the proofs by cancelling many terms, which you currently have to bound using random matrix theory and concentration inequalities. This contributes to making the proofs difficult to follow.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper studies the learning dynamics of a Transformer layer (with non-linear attention and MLP sublayers), which is crucial in order to understand better the properties of this architecture, and a technically difficult problem due to the highly non-linear parameterization of the architecture. The analysis of CoT is original and compelling, in that it gives a provable mechanism by which CoT simplifies the learning problem.", "weaknesses": "While the ideas in the paper are nice, the paper would highly benefit from polishing to improve the clarity and presentation of the message. In particular, the proofs are written in a telegraphic style, which makes it particularly difficult to check correctness. It also makes it hard to get a sense of how general the approach is, where the specific initialization scheme comes into play, etc. I give below more specific examples of things that confused me and would benefit from cleaning up.\n\nBesides, I am not convinced by the real-world experiments. The difference in terms of entropy of the attention layers does not seem very significant, so in my opinion it is difficult to reach any strong conclusion from the figures.\n\nAll in all, the rating reflects the fact that the ideas are interesting, but I believe polishing the paper would make it much better.\n\nExamples of confusing statements in the main text and the beginning of the proofs (the only point for which I would like an answer during the rebuttal is the first one, others are less important):\n- in Theorem 2, the condition on $d$, $m$, $n$ and $k$ is not satisfied for the dimensions taken in the proof of Theorem 1, so it is hard to relate both results.\n- line 296: “which only applies to one-pass training”, but Theorem 2 applies to multipass training?\n- line 831: why $m>100$?\n- line 841: I agree that the results should follow from random matrix theory, but you should give a proper statement and proof/reference. The result depends on the relative value of $d$ and $k$, and the current statement $d = \\tilde{\\Omega}(k)$ is too hand-wavy to conclude that the specific numerical values given in the Lemma are verified.\n- line 844: why $\\delta/4$?\n- lines 896: why $8$ instead of $10$?\n- line 901: the $\\mathcal{O}(1/n^7)$ needs to be explained.\n- line 913: “by choosing $a_r$ and $b_r$ iteratively”: this needs to be explained.\n- line 916: “the parameters can be encoded in $\\mathcal{O}(\\log n)$ precision”. Why?\n- line 960: this is not a valid proof of the Lemma, the derivation of the inequalities is non-trivial from the Assumption.\n\nMinor remarks (no impact on score):\n- line 46: the specific example of the number of letters in a word is not very relevant, because perhaps this is precisely an example of a task which may not be representable by a Transformer model due to the tokenisation mechanism.\n- line 77: $k$ is not introduced yet.\n- line 191: the bounds of the sum should depend whether it is a CoT or non-CoT task.\n- line 230: it should be $W_{r,1:m}$ instead of $W_{r,1:d}$.", "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)