From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency

## Abstract

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. 1 1 1 Our code is available at https://github.com/zhqwqwq/Learning-Parity-with-CoT .

## 1 Introduction

Chain-of-thought (CoT) has proven to be a powerful technique for enhancing reasoning in large language models [, ]. By instructing the model to break complex problems into smaller, manageable steps, CoT facilitates more efficient reasoning and better generalization, particularly in algorithmic and logical tasks [, , ].
Building on this, performance can be further improved through multi-step prompting and multi-path sampling techniques [, , , , ].

This focus on CoT within in-context learning has since expanded to more structured learning approaches [, ]. By adding reasoning examples of CoT style to the instruction-tuning dataset, models enhance their problem-solving abilities more effectively than relying solely on CoT during prompting [, ]. As a result, CoT is now shaping a new paradigm in language model development, marking a shift from simply scaling data [, ] to focusing on advanced reasoning strategies [], which leads to more effective learning outcomes.

While CoT’s success is well-established, understanding why it works is still a hotly debated topic [, ]. Recent theoretical studies suggest that CoT enhances a model’s expressiveness, increasing its representational capacity when the sequence is long enough [, ]. However, expressivity alone does not guarantee success. Large language models often struggle with simple tasks—like counting the number of ’r’s in "strawberry"—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 concrete 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, we show that CoT sequences introduce sparse sequential dependence, thereby enhancing the sparsity in the attention layers. We then show that transformers can efficiently optimize and generalize on such sequences. We summarize our contributions as follows.

- 1.
Theoretically, we show that while the parity problem can be expressed by a 1-layer transformer, learning requires exponential to $k$ samples without CoT when the number of parameters is limited (Theorem 1 and Theorem 2). Meanwhile, for CoT data with sparse sequential dependency, a 1-layer transformer can learn the parity function and faithfully represent the sparse dependency in its attention pattern with almost linear samples with respect to $k$ and $n$ (Theorem 3).
- 2.
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 3). 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

Figure: Figure 1: (a) We show that, without Chain-of-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 (1), is crucial in the parity learning experiment. In both CoT and non-CoT scenarios, as the attention layers become sparser—indicated by a rapid decrease in normalized entropy—a corresponding jump in evaluation accuracy occurs.
Refer to caption: x1.png

We start with empirically exploring the sample efficiency of training transformers to learn the class of parity functions [] with and without CoT. The parity functions exhibits simple structure but yet hard to learn by traditional networks [, ]. Specifically, given a set of $n$ binary variables $b_{1},b_{2},\dots,b_{n}$, a parity function takes $k$ secret variables $b_{i_{1}},\dots b_{i_{k}}$ and outputs $1$ if the sum of these $k$ variables is odd, and 0 0 if it is even:

$$ $f(b_{1},...,b_{n})=b_{i_{1}}\oplus b_{i_{2}}\oplus\cdots\oplus b_{i_{k}},$ $$

where $b_{i}\in\{0,1\}$ and $\oplus$ is the XOR operator.
For example, $f(b_{1},b_{2},b_{3},b_{4},b_{5})=b_{1}\oplus b_{2}\oplus b_{4}$ is a $5$-variable parity function with $k=3$. The function $f$ returns 1 if the sum of $b_{1}$, $b_{2}$ and $b_{4}$ is odd, and 0 0 otherwise, independently of the value of $b_{3}$ and $b_{5}$. Intuitively, the parameter $k$ controls the hardness of the problem. When $k$ increases, the number of possible subsets grows exponentially, making the identification of the correct secret set more challenging.

Given a $n$-variable parity function $f$, we generate sequences of length $n+2$ as auto-regressive manner: $b_{1},\cdots b_{n},b_{n+1}:=0$ and $b_{n+2}=f(b_{[1:n]})$. To incorporate CoT, we break the sum of XOR down into $k$ steps and add all the intermediate steps into the sequence, i.e. including $b_{i_{1}}$, $b_{i_{1}}\oplus b_{i_{2}}$, $\cdots$, $b_{i_{1}}\oplus b_{i_{2}}\oplus\cdots\oplus b_{i_{k}}$. As a result, the CoT data has length $n+k+1$.
With the example function $f(b_{1},b_{2},b_{3},b_{4},b_{5})=b_{1}\oplus b_{2}\oplus b_{4}$, one sampled sequence would be

$$ No CoT $\displaystyle\underbrace{0,1,0,1,0}_{\mathrm{input}},\underbrace{0}_{\mathrm{[ EOS]}},\underbrace{0}_{\mathrm{answer}}\in\{0,1\}^{7},\text{ as }b_{1}\oplus b _{2}\oplus b_{4}=0\oplus 1\oplus 1=0.$ With CoT $\displaystyle\underbrace{0,1,0,1,0}_{\mathrm{input}},\underbrace{0}_{\mathrm{[ EOS]}}\underbrace{0,1}_{\mathrm{CoT}},\underbrace{0}_{\mathrm{answer}}\in\{0,1 \}^{9},\text{as }b_{1}=0,b_{1}\oplus b_{2}=1.$ $$

Next, we train transformer networks with and without CoT respectively, with a common held-out test set, and compare their sample complexities. The sample complexity is defined as the amount of data a model sees at the first time it 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 1, 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, it suggests that sparsity in the attention layers plays a crucial role in the learning process. To explore why CoT enables transformers to learn more efficiently, we formally define our problem setup below.

## 3 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 as our testbed because identifying a set of key variables amidst various confounding ones is a fundamental aspect of many reasoning tasks, and parity serves as an abstraction of this process.

### 3.1 Notations and Definitions

Parity Problem. Each token in our setting is either 0 0 or $1$. We represent a sequence of length $T$ as $\textbf{b}[1],\ldots,\textbf{b}[T]$. A parity function is a function of form $\mathrm{parity}_{S}(\textbf{b})=\oplus_{j\in\mathcal{S}}\textbf{b}[j]$, where $\mathcal{S}$ is the set of secret indices and $\oplus$ is the XOR operator on binary variables. The cardinal of the secret set $k=|S|$ controls the hardness of the problem.

###### Definition 1 (Parity Problem $(n,k)$ without CoT) .

Given a secret set $\mathcal{S}$ of cardinal $k$, we define the parity problem without CoT as the following distribution of sequences $p_{\rm{NoCoT}}$:

$$ $\displaystyle\textbf{b}[i]\sim U(\{0,1\}),\forall i\in[1,n],\textbf{b}[n+1]=0, \textbf{b}[n+2]=\mathrm{parity}_{S}(\textbf{b}).$ $$

where $\textbf{b}[1],...,\textbf{b}[n]$ are uniformly sampled from 0 0 and $1$.

###### Definition 2 (Parity Problem $(n,k)$ with CoT) .

Given a secret set $\mathcal{S}$ of cardinal $k$,
we define the parity problem with CoT as the following distribution of augmented sequences $p_{\rm{CoT}}$:

$$ $\displaystyle\textbf{b}[i]\sim U(\{0,1\}),\forall i\in[1,n],\textbf{b}[n+1]=0, \textbf{b}[i+1]=\textbf{b}[i]\oplus\textbf{b}[S[i]],\forall i\in\{n+1,\ldots,n +k\}.$ $$

where $S$ is any permutation of $\mathcal{S}$, i.e. $\{S[i]\mid i=n+1,\ldots,n+k\}=\mathcal{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 a given secret set since we can arbitrarily permute $\mathcal{S}$, given the commutative property of the XOR operator.

Transformer Architecture. To conduct the theoretical analysis, we simplify the Transformer architecture similar to prior works (see e.g. [, , ]). More precisely, we drop all the layer norms or batch normalization; simplify the positional embedding; use a square matrix to represent the attention layer and concatenate the residual branch in a Densenet fashion []. After simplification, the network still has strong expressive power like the standard Transformers.

In a standard Transformer, an embedding layer matches each vocabulary token into a 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 $2T$ unit vectors as the embedding vectors:

###### Definition 3 (Embedding Module) .

For any position $i\in[T]$, we sample two embedding vectors $e_{i,0}$ and $e_{i,1}$ uniformly at random from unit hypercube $\mathrm{U}\left(\{\frac{1}{\sqrt{d}},-\frac{1}{\sqrt{d}}\}^{d}\right)$. These embedding vectors are frozen during training.
Then for any binary sequence $\textbf{b}\in\{0,1\}^{T}$, the embedding is defined as

$$ $\displaystyle\mathbf{E}\left(\textbf{b}\right)=[e_{(i,\textbf{b}[i])}]_{i\in[T ]}\in\mathbb{R}^{d\times T}.$ $$

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. In standard architecture, the attention layer is derived from the product of query and key matrices, i.e. $QK^{T}=\mathrm{X}^{T}W^{Q}(W^{K})^{T}\mathrm{X}$. In our case, we directly use a full matrix $\mathrm{X}^{T}A\mathrm{X}$ to parameterize it, which has the same representation power.

###### Definition 4 (Attention Module) .

Given an input matrix $\mathrm{X}\in\mathbb{R}^{d\times T}$ and attention weight $A$, we define the attention module as

$$ $\displaystyle\mathcal{A}(\mathrm{X})=\mathrm{X}\ \mathrm{softmax}(C+\mathrm{X} ^{T}A\mathrm{X}),$ $$

where $C\in\mathbb{R}^{T\times T}$ is the autoregressive mask with value $-\infty$ 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 a matrix $\mathrm{X}\in\mathbb{R}^{2d\times T}$ as

$$ $\displaystyle\mathrm{FFN}(\mathrm{X})=h^{T}\mathrm{ReLU}(WX).$ $$

Here $W\in\mathbb{R}^{2m\times 2d},h\in\mathbb{R}^{2m\times o}$ with $o$ being the output dimension and $\mathrm{ReLU}$ is element-wise.

Finally, we concatenate the modules into a simplified Transformer block.

###### Definition 6 (Simplified Transformer Block) .

We define a simplified Transformer block as,

$$ $\displaystyle\mathcal{T}(\textbf{b})=\mathrm{FFN}\left(\begin{bmatrix}\mathbf{ E}\left(\textbf{b}\right)\\ \mathcal{A}\left(\mathbf{E}\left(\textbf{b}\right)\right)\end{bmatrix}\right),$ $$

where $\mathbf{E}$ is the embedding module, $\mathcal{A}$ is the attention module and $\mathrm{FFN}$ the fully connected layer. An $L$-layer Transformer consists of a composition of $L$ such blocks, with the embedding module $\mathbf{E}$ appearing only in the initial layer. The intermediate dimension is set to $d$, while the output dimension of the final layer is 1.

Compared with the standard residual structure, we use a Densenet structure to concatenate the residual 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. [], Nichani et al. []).

Loss function. To simplify our analysis, we use hinge loss $\ell(\hat{y},y)=\max\{(-1)^{y}\hat{y}+1,0\}$ as the loss function. Consequently, we define the next token prediction loss of a boolean sequence b as
$L(w)=\sum_{i=n+1}^{n+k}\ell(\mathcal{T}^{(L)}(\textbf{b})[i],\textbf{b}[i+1])$,
where $w$ denotes all the trainable parameters $[A,W]$ in the attention module and the FFN map across all layers.

###### Definition 1 (Parity Problem ( n , k ) 𝑛 𝑘 (n,k) ( italic_n , italic_k ) without CoT) .

###### Definition 2 (Parity Problem ( n , k ) 𝑛 𝑘 (n,k) ( italic_n , italic_k ) with CoT) .

###### Definition 3 (Embedding Module) .

###### Definition 4 (Attention Module) .

###### Definition 5 .

###### Definition 6 (Simplified Transformer Block) .

### 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 1 (Easy to Represent) .

Consider the Transformer model defined in Definition 6, for any $\delta<0.1$ and large enough $n$, with probability at least $1-\delta$ over the randomness of embedding $e$, there exists a weight configuration of the Transformer with dimension $d=\tilde{\Omega}(k)$ and width $2m=\Omega(k)$, such that it can achieve perfect accuracy on the parity problem $(n,k)$ without CoT, i.e.

$$ $\displaystyle\forall\textbf{b}\sim p_{\rm{NoCoT}},\mathrm{sgn}\left(\mathcal{T }(\textbf{b})[n+1]\right)=\textbf{b}[n+2].$ $$

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 *representational-sufficient* 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 Section A.2.

While the function is expressible, this does not guarantee that the solution can be easily found. In fact, we show that achieving a perfectly accurate solution using a gradient-based optimization method requires an exponential number of samples in $k$ assuming the memory, i.e. space complexity, to perform optimization is bounded throughout the training.

###### Theorem 2 (Hard to Learn) .

For any randomly initialized simplified Transformer model with constant layers, when the embedding dimension $d$ and width $2m$ satisfies $d^{2}+dm=o(nk/\log(n))$, for any constant number of passes $q$, when the model is trained with $q-$pass stochastic gradient descent with $O((d^{2}+dm)\log n)$ memory, the sample complexity required to learn the parity problem $(n,k)$ without CoT to any nontrivial accuracy $a>50\%$ with nontrivial probability $p>50\%$ is $2^{\Omega(k)}$.

The proof is deferred to Section A.3. We use results from the classical online learning communities [] to show that in the regime where the memory required to perform the training is less than $nk$, the model can’t effectively store information about enough samples inside the parameters during training, and hence can’t infer the secret indices effectively. Here the parameters $k$ denotes the size of the secret set $\mathcal{S}$ and the sample complexity grows exponentially when $k$ grows.

###### Theorem 1 (Easy to Represent) .

###### Theorem 2 (Hard to Learn) .

### 3.3 Polynomial Sample Complexity with CoT

The result in Theorem 2 presents a seemingly negative outcome for learning algorithmic reasoning, as the model requires exponentially many samples relative to $k$. However, in our main result, Theorem 3, we will demonstrate that the sample complexity can be significantly reduced if a step-by-step derivation is provided. In other words, CoT is much more sample efficient, where model can effectively learn complex relationships as long as each token depends on only a 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 ($W_{r,1:m}=0$). The parameter $h$ in the FFN and the word embedding $\mathbf{E}$ is fixed during training.

###### Assumption 1 (Initialization) .

At initialization,

$$ $\displaystyle\forall r\in[2m],A$ $\displaystyle=0,W_{r,d+1:2d}=0,W_{r,1:d}=\sum_{i=n+1}^{n+k}\sum_{b=0}^{1}\nu_{ r,i,b}e_{i,b},h_{1:m}=\frac{1}{2m},h_{m+1:2m}=-\frac{1}{2m}.$ $$

Here $\nu_{r,i,b}$ is independent random variable sampled uniformly from $\{-\epsilon,\epsilon\}$.

We can then train the model with stochastic gradient descent (SGD) and get the following theorem.

###### Theorem 3 (Easy to Learn with CoT) .

For any constant $\delta\in(0,1)$, when the number of secret indices $k$ is in $[n/\log^{5}(n/\delta),n/\log^{4}(n/\delta)]$, with probability $1-\delta$, a randomly initialized simplified Transformer (see 1) with $o(nk)$ parameters trained for constant steps using mini-batch SGD with $\tilde{O}(n)$ samples using appropriate hyperparameters (see 2) can reach perfect accuracy on a parity problem $(n,k)$ with CoT,

$$ $\displaystyle\forall\textbf{b}\sim p_{\rm{CoT}},i\in\{n+1,\ldots,n+k\},\mathrm {sgn}\left(\mathcal{T}(\textbf{b})[i]\right)=\textbf{b}[i+1].$ $$

Furthermore, after training, the attention pattern is interpretable and one-hot in the sense that, for any $i\in[n+1,n+k]$,

$$ $\displaystyle\big{|}\mathrm{softmax}(C+\mathbf{E}(\textbf{b})A\mathbf{E}( \textbf{b}))[j,i]-\bm{1}(j=S[i])\big{|}<1/n^{8}.$ $$

Moreover, the result still holds even if all the weights and activations are in $\Theta(\log n)$ precision.

This theorem indicates that the attention module can successfully extract the sparse sequential dependencies in the CoT data and faithfully represent it in the attention pattern, which is validated in our experiments (see Figure 3). It is also amongst the first optimization dynamics analysis of Transformers using finite-sample gradients rather than population gradients.

*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 a 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 ($W_{r,d+1:2d}$) is initialized to be zero. The set of neurons activated at each position $i\in[n+1,n+k]$ is solely determined by the word embedding $e_{(i,\textbf{b}[i])}$ at the position. Hence, at initialization, the MLP can be viewed as an ensemble of multiple linear functions specialized to each position and boolean value. We can show that when the learning rate is small, the set of activated neurons at each position remains the same through the training process. Hence, we can conceptually view the FFN function as a set of different linear functions applied independently at each position and value. We will denote this set of linear weight as $\kappa_{t,i,b}$ (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 $\textbf{b}[i+1]$ when conditioned on $\textbf{b}[i]$ is $S[i]$. This suggests that the linear weight $\kappa_{1,i,\textbf{b}[i]}$ will have a stronger correlation with the embedding $e_{(S[i],\textbf{b}[j])}$ 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 $\textbf{b}[i]$ for $i\in[n+1]$ will be uncorrelated with the desired output $\textbf{b}[n+2]$ in such case.

Phase 2. 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 $\mathrm{softmax}(C+\mathrm{X}^{T}A\mathrm{X})[j,i]$ will be larger when the approximate contribution of the embedding of the $j-$th token to the output at the $i-$th token $\langle\kappa_{1,i,\textbf{b}[i]},e_{(j,\textbf{b}[j])}\rangle$ is larger. This suggests that the attention will be amplified on the index that has a strong correlation with the output at the current token, which is the secret index $S[i]$. This step will make the attention pattern approximately one-hot as in the theorem statement.

Phase 3. Learning the Output Mapping
At the final step, as the attention pattern is one-hot, the FFN layer only needs to learn a 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 at 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.

###### Assumption 1 (Initialization) .

###### Theorem 3 (Easy to Learn 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 transformers without CoT, as a complement to the established lower bound (Theorem 2) which only applies to one-pass training. The results indicate that multi-pass training indeed improves the models’ ability to learn parity problems.
Third, we conduct experiments on the GSM8K dataset [] to show that CoT introduces sparse sequential dependence on real-world training data.
Although necessary simplifications on transformer are made in Section 3 to develop theoretical results, we use the standard transformer architecture in the subsequent experiments following the GPT-2 architecture [] 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 $n=30$ and $k=1,2,3,4$, with and without CoT, varying layers and heads from $1$ to $4$. We choose the best performer across learning rates from $6\times 10^{-5}$, $8\times 10^{-5}$, and $1\times 10^{-4}$.
At each step, a fresh batch of training data is sampled, with a maximum budget of $10^{7}$ samples. We record the number of samples seen by the model before reaching an evaluation accuracy of $1$ as the sample complexity.

Figure: Figure 2: The sample complexity for learning parity without CoT increases exponentially with $k$. CoT significantly reduces the sample complexity, demonstrating exponential improvement across varying numbers of heads and layers.
Refer to caption: x2.png

Figure: Figure 3: (a) When $n$ is fixed, the sample complexity of learning parity with CoT grows approximately linearly with $k$. (b) The attention pattern learned by the transformer with CoT is interpretable, as the $i$-th output token of CoT predominantly attends to secret index $S[i]$.
Refer to caption: x3.png

Figure: Figure 4: Evaluation accuracy of transformers on the $(n=20,k=6)$ parity problem without CoT. The model is trained on a dataset of $10,000$ samples for $1,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.
Refer to caption: x4.png

Results. Figure 2 shows that training with CoT [first row in Figure 2 (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 the parameter $k$. In contrast, CoT greatly reduces the sample complexity, showing an exponential improvement across different numbers of heads and layers. These findings are consistent with our theoretical analysis.

Training with CoT induce sparse and interpretable attentions (Figure 3).  We evaluate the sample complexity of a $1$-layer, $1$-head transformer on larger parity problems ($n=100$ and $20\leq k\leq 100$). Without CoT, training exceeds the sample budget when $k\geq 4$. However, with CoT, we can successfully train for any $k$. Figure 3b shows the attention pattern of a $1$-layer, $1$-head transformer trained on a parity problem with $(n=40,k=20)$ using CoT on a random query. In this case, CoT processes the secret variables in ascending order. The $i$-th row illustrates the attention weight of the $i$-th token attending to the previous ones, with indices corresponding to secret variables highlighted in red.

Results. In Figure 3a, we show that training with CoT succeed with large $n$ and $k$. Moreover, the empirical result suggests that the sample complexity of learning parity with CoT grows approximately linearly with $k$. This shows a 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 3b 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 to other tokens.

### 4.2 Multi-pass Training Improves Parity Learning without CoT

Figure: Figure 5: $4$-layer $4$-head transformer trained on the $(n=20,k=6)$ parity problem without CoT using multi-pass training, detailed in Section 4.2. When trained on a small dataset of $50000$ samples, the model achieves perfect evaluation accuracy (Top), accompanied by a significant decrease in entropy. 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 elevated. This shows that the development of attention sparsity may improve optimization.
Refer to caption: x5.png

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-pass training is the development of sparse attention (see Figure 5 Right), similar to 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,k=6)$ parity problem, we conduct multi-pass training without CoT. In Figure 4, the models are trained with $10^{4}$ samples for $1000$ epochs, with the number of layers and heads ranging from $1$ to $4$. In Figure 5, we compare the training of $4$-layer $4$-head transformer on $5\cdot 10^{4}$ and $10^{6}$ samples respectively. The learning rate is initialized $10^{-4}$.

Normalized attention entropy. To measure attention sparsity, we introduce the concept of *normalized attention entropy* for each attention head (illustrated in Figure 5 Right). Let $\mathcal{P}^{\ell,h}(\mathbf{x})[i]$ denote the attention score distribution produced by the $h$-th head in the $\ell$-th layer at the $i$-th token of input $\mathbf{x}$. The normalized attention entropy for the input $\mathbf{x}$ is then defined as:

$$ $\displaystyle\mathrm{Ent}(\mathbf{x};\ell,h)=\min_{i\geq 2}\frac{H(\mathcal{P} ^{\ell,h}(\mathbf{x})[i])}{\log i}\in[0,1],$ (1) $$

where $H(\mathcal{P})=-\int\log(\mathcal{P})d\mathcal{P}$ is the entropy of distribution $\mathcal{P}$. The normalization term $\log i$ represents the entropy of a uniform distribution over $i$ tokens, account for the varying context length. The minimum is taken over different tokens in $\mathbf{x}$ since attention heads may specialize in extracting information for specific tokens. As a result, a 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.

Results. As shown in Figure 4, when trained on $10^{4}$ samples, which only accounts for a small portion ($\sim 1\%$) of all possible inputs ($2^{20}\approx 10^{6}$), most of the transformer architectures we examined achieve perfect evaluation accuracy given sufficient epochs. While the $k=6$ problem is intractable with one-pass training without CoT, these results demonstrate that multi-pass training can indeed enhance 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.

Although learning without CoT is less efficient, it offers a "slow-mode" trajectory. In Figure 5 (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 previously demonstrated in Figure 1, 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 in the intermediate steps.

### 4.3 CoT Induces Sparsity on Real-World Data

Figure: 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 (Section 4.3). Each bar represents the entropy of an attention head. The Qwen2-7B model exhibits sparser attention when processing CoT data, indicating that real-world CoT data has a sparser structure. The entropy difference between the Qwen2-7B and Qwen2-Math-7B model suggests that fine-tuning on CoT data promotes the development of sparser attention patterns of the model.
Refer to caption: x6.png

Now we move from the synthetic parity problem to real world experiment on GSM8K dataset of grade-school math word problems []. We observe that:

- 1.
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 two 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 language models: the pre-trained model Qwen2-7B [] and the specialized mathematics model Qwen2-7B-Math [] which is fine-tuned from Qwen2-7B on a mathematics-specific corpus with CoT data.
We plot the normalized attention entropy (Equation 1) across different heads. More details can be found in the Appendix Section B.1.

Results.
Unlike the synthetic parity problem, the sequential dependency of real-world data is hard to measure directly. However, it can be inferred from the attention sparsity of 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-7B model demonstrates lower attention entropy compared to 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 [], 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 [], which operate within the NTK linearization regime of RNNs. In contrast, the Transformers analyzed in our work exhibit feature learning and identify the sparse secret set within the attention module. On the positive side, a line of works that studies optimization dynamics of neural networks on parity [, , , , , ] show that $d^{\Omega(k)}$ samples is sufficient to learn parity, which is close to the statistical query lower bound. On the negative side, it has been shown that in SGD learning with additional noise, $d^{\Omega(k)}$ samples are indeed needed [, ]. However, it remains unclear whether standard mini-batch SGD will still need exponential sample complexity. In our paper, we take the alternative approach to consider the training as an online algorithm with bounded memory (parameters) and utilize results from online learning literature to provide a rigorous exponential lower bound [, ]. [] 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.

Chain-of-Thoughts (CoT). Chain-of-Thoughts (CoT) is the technique to let the model generate a reasoning process before final answers.
CoT prompting is an effective method in practice to improve language model’s reasoning capabilities [, , , , , ]. Training with CoT data has further improved the model’s capability in performing complex reasoning (see Qwen [], Yue et al. [], Yu et al. [], Kim et al. [] and reference therein).
Different lines of work aimed at understanding the effect of CoT. From the representation theory perspective, works including [, , , , ] show that CoT can provably expand the representation power of different neural architectures including Transformers and RNNs. From the statistical approximation level, prompting with CoT has improved the statistical error [, ]. [] studies how MLP models learn with CoT data. [] 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 Chain of Thought, which is not reflected in their data modeling; (2) we study zero-shot CoT data directly. Finally, Dutta et al. [], Wang et al. [] studies pretrained Transformers’ activation on CoT data. They highlight the attention head’s role in moving the essential information from the context to the reasoning process, which is consistent with our theoretical insight.

Transformer Optimization Dynamics. Our works fall in the line of works that studies the optimization dynamics of Transformers on synthetic datasets [, , , , , , , , , , , ]. Similar to our works, [] highlights that Transformers can learn to select sparse critical tokens from the context on linear data. We differ from their work in studying the mini-batch optimization dynamics on nonlinear data and hence establishing sample complexity bounds.

## 6 Conclusion and Future Works.

Our work demonstrates that Transformers trained on Chain-of-Thought (CoT) data with sparse sequential 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 function classes beyond parity.

## Appendix A Omitted Proof

### A.1 Notation and Assumptions.

We will denote a sequence of binary variables as $\textbf{b}[1],\ldots,\textbf{b}[T]$ with $\textbf{b}[i]\in\{0,1\}$. We will use $[j]$ to denote variables corresponding to the $j-$th token dimension without otherwise specified.

#### A.1.1 Training Specification

##### Training Update

We will consider the following 1-pass SGD with batch size $B$:

$$ $\displaystyle L^{(t)}$ $\displaystyle=\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}\ell(\mathcal{T}[w^{(t)}]( \textbf{b}^{(t,s)})[i],\textbf{b}^{(t,s)}[i+1]).$ $\displaystyle w^{(t+1)}$ $\displaystyle=w^{(t)}-\eta_{t}\frac{1}{B}\nabla_{w}L^{(t)}.$ $$

Here $w$ includes $A$ and $W$. We will leave $h$ unchanged for the simplicity of analysis.

### A.2 Representation Theory

###### Proof of Theorem 1 .

We will set $d=\Theta(k\log(n/\delta))$ and $m=\Theta(k)$. We will use $p(x,y)$ as shorthand for $2(x-1)+y+1$.

Define embedding matrix $M$ as follows

$$ $\displaystyle M_{p(i,b)}=e_{S[i],b}^{T}.$ $$

for $\{S[1],\ldots,S[p]\}=\mathcal{S}$. Then $M\in\mathbb{R}^{2k\times d}$.

###### Lemma 1 .

For large enough $n$, with probability $1-\delta$, it holds that

$$ $\displaystyle\|M^{T}\|_{\mathrm{op}}=O(1).\quad\lambda_{\mathrm{min}}(MM^{T})> 1/12.$ $$

###### Proof.

The first result follows from Lemma 28, noted that $K=\Theta(\sqrt{1/d})$. The second result follows from Lemma 29.
∎

###### Lemma 2 .

With probability $1-\delta/4$ over the randomness of embedding $e_{j,b}$ for $j\in\mathcal{S},b\in\{0,1\}$, there exists $u_{j,b}$, satisfying that $\forall j\in\mathcal{S},b\in\{0,1\}$,

$$ $\displaystyle\forall j^{\prime}\in\mathcal{S},b^{\prime}\in\{0,1\},\langle u_{ j,b},e_{j^{\prime},b^{\prime}}\rangle$ $\displaystyle=\bm{1}((j,b)=(j^{\prime},b^{\prime})).$ $$

Further,

$$ $\displaystyle\|\sum_{j\in\mathcal{S},b\in\{0,1\}}{u_{j,b}}\|_{2}=O(\sqrt{k}).$ $$

###### Proof.

We can simply choose

$$ $\displaystyle u_{S[i],b}$ $\displaystyle=M^{T}(MM^{T})^{-1}o_{2(i-1)+b}.$ $$

with $o_{2(i-1)+b}$ being the one-hot vector in $\mathbb{R}^{2k}$.
∎

###### Lemma 3 .

When the event in Lemma 2 happens, there exists $v_{j}$,

$$ $\displaystyle\langle v_{j},e_{j,0}\rangle=\langle v_{j},e_{j,1}\rangle$ $\displaystyle=1.$ $\displaystyle\forall j^{\prime}\neq j,j\in S,\langle v_{j^{\prime}},e_{j,b}\rangle$ $\displaystyle=0.$ $$

Further,

$$ $\displaystyle\|\sum_{j\in\mathcal{S}}{v_{j}}\|_{2}\leq O(\sqrt{k}).$ $$

###### Proof.

We can choose $v_{j}=u_{j,0}+u_{j,1}$.
∎

###### Lemma 4 .

When $v_{i}$ defined in Lemma 2 exists, for all $j\not\in\mathcal{S},b\in\{0,1\}$, with probability $1-\delta$, it holds that

$$ $\displaystyle|\langle\sum_{j^{\prime}\in\mathcal{S}}v_{j^{\prime}},e_{j,b} \rangle|\leq 1/2.$ $$

###### Proof.

Noted that $e_{j,b}$ is independent from $v_{j^{\prime}}$, assuming $V=\sum_{j^{\prime}\in\mathcal{S}}v_{j^{\prime}}$, then $\|V\|_{2}=O(\sqrt{k})$, using Azuma-Hoeffding Bound, with probability $1-\delta/2$

$$ $\displaystyle|\langle\sum_{j^{\prime}\in\mathcal{S}}v_{j^{\prime}},e_{j,b} \rangle|=O(\frac{\|\sum_{j^{\prime}\in\mathcal{S}}v_{j^{\prime}}\|_{2}}{\sqrt{ d}}\sqrt{\log(2k/\delta)})=O(\sqrt{\frac{k\log(2k/\delta)}{d}})\leq\frac{1}{2}.$ $$

The proof is then completed.
∎

Assuming the event in Lemmas 2 and 4 happens, defined $V=\sum_{j^{\prime}\in\mathcal{S}}v_{j^{\prime}}$ and $U=\sum_{j^{\prime}\in\mathcal{S}}u_{j^{\prime},1}$

$$ $\displaystyle A=40\log(n)Ve_{n+1,0}^{T}\quad W_{r,1:d}=-a_{r}e_{n+1,0}\quad W_ {r,d+1:2d}=b_{r}U.$ $$

for some scalar $a_{r},b_{r}$ to be determined.

For an arbitrary b, it holds that,

- 1.
If $j\in\mathcal{S}$, $e_{j,\textbf{b}[j]}^{T}Ae_{n+1,0}=40\log n$.
- 2.
If $j\not\in\mathcal{S}$, $e_{j,\textbf{b}[j]}^{T}Ae_{n+1,0}\leq 20\log n$.

We then have

$$ $\displaystyle\mathrm{softmax}(\mathbf{E}(\textbf{b})A\mathbf{E}(\textbf{b}))[j ,n+1]$ $\displaystyle=\frac{\bm{1}(j\in\mathcal{S})}{k}+O(\frac{1}{n^{7}}).$ $$

As

$$ $\displaystyle\langle u,e_{j,b}\rangle=\bm{1}(j\in\mathcal{S},b=1).$ $$

We have that

$$ $\displaystyle\langle u,\mathcal{A}(\mathbf{E}(\textbf{b}))[n+1]\rangle=\frac{1 }{k}\sum_{j\in\mathcal{S}}\textbf{b}[j]+O(\frac{1}{n^{6}}).$ $$

Now we only need to map from the summation $\sum_{j\in\mathcal{S}}\textbf{b}[j]$ to the parity of the summation. This can be done by choosing $a_{r}$ and $b_{r}$ iteratively. The proof is then complete.
∎

###### Proof of Theorem 1 .

###### Lemma 1 .

###### Proof.

###### Lemma 2 .

###### Proof.

###### Lemma 3 .

###### Proof.

###### Lemma 4 .

###### Proof.

### A.3 Dynamics without CoT

###### Proof.

We will utilize the Theorem 6 in Lyu et al. [], 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 be saved in the memory. However, we can take the alternative approach to regenerate $e$ using a random number generator with a fixed seed on each step. This allow us to simulate standard SGD optimization with $o(nk)$ memory, which is a special case of branching programs.
∎

###### Proof.

### A.4 Dynamics with CoT

###### Assumption 2 .

Consider the following conditions for sufficiently large $n$:

- 1.
The secret Hamming weight satisfies $k\in\left[\frac{n}{\log^{3}(n/\delta)},\frac{n}{\log^{2}(n/\delta)}\right]$.
- 2.
Set $d=k\log^{1.1}\left(\frac{n}{\delta}\right)$ and $m=2k$. This implies $md\log n=o(nk)$.
- 3.
Define the batch size as $B=C_{2}n\log^{20}\left(\frac{n}{\delta}\right)$ for a sufficiently large constant $C_{2}$.
- 4.
Set the learning rates to be
$\eta_{0}=\eta_{1}=\frac{m\epsilon\sqrt{B}}{100\log\left(\frac{n}{\delta}\right
)},\quad\eta_{2}=\frac{4k\epsilon}{3}.$

###### Lemma 5 .

Under Assumption 2, the following conditions hold as $n\to\infty$:

- 1.
$\lim_{n\to\infty}\frac{d}{n}=0$ and $\lim_{n\to\infty}\frac{B}{nk}=0$.
- 2.
$\lim_{n\to\infty}\sqrt{\frac{2k}{d}\log\left(\frac{Bmn}{\delta}\right)}=0$.
- 3.
Each of the following expressions tends to zero:
(a)
$\frac{\sqrt{k\log\left(\frac{300mnB}{\delta}\right)}}{\sqrt{B}}$.
(b)
$\frac{\sqrt{2k\log\left(\frac{400mn}{\delta}\right)}}{\sqrt{d}}$.
(c)
$\frac{\sqrt{2nk\log\left(\frac{300mnB}{\delta}\right)}}{\sqrt{Bd}}$.
- 4.
$\eta_{0}=\eta_{1}\leq\min\left\{\frac{3m\epsilon\sqrt{B}}{80\sqrt{\log\left(
\frac{300nB}{\delta}\right)}},\frac{mn\epsilon}{120}\right\}$.
- 5.
$\lim_{n\to\infty}\log^{2}\left(\frac{300mnB}{\delta}\right)\frac{n\sqrt{nk}}{d
\sqrt{Bd}}=0$.
- 6.
$\frac{\eta_{0}\eta_{1}}{256mn^{2}}>10\log n$.
- 7.
$\lim_{n\to\infty}\frac{\eta_{2}}{\eta_{0}}=0.$
- 8.
$\lim_{n\to\infty}(\eta_{0}+\eta_{1})<5n\eta_{2}$.

###### Proof.

The conditions in Lemma 5 are satisfied based on the definitions provided in Assumption LABEL:assum:order_concise. Below, we outline the verification for each condition:

- 1.
Limit of $\frac{d}{n}$ and $\frac{B}{nk}$:
$\frac{d}{n}=\frac{k\log^{1.1}\left(\frac{n}{\delta}\right)}{n}\in\left[\frac{1
}{\log^{1.9}(n/\delta)},\frac{1}{\log^{0.9}(n/\delta)}\right]\to 0\text{ as }n
\to\infty.$
Similarly,
$\frac{B}{nk}=\frac{C_{2}n\log^{20}\left(\frac{n}{\delta}\right)}{nk}=\frac{C_{
2}\log^{20}\left(\frac{n}{\delta}\right)}{k}\leq\frac{C_{2}\log^{20}\left(
\frac{n}{\delta}\right)}{\frac{n}{\log^{3}(n/\delta)}}=C_{2}\frac{\log^{23}(n/
\delta)}{n}\to 0.$
- 2.
Limit of $\sqrt{\frac{2k}{d}\log\left(\frac{Bmn}{\delta}\right)}$:
$\sqrt{\frac{2k}{d}\log\left(\frac{Bmn}{\delta}\right)}=\sqrt{\frac{2k}{k\log^{
1.1}\left(\frac{n}{\delta}\right)}\log\left(C_{2}n\log^{20}\left(\frac{n}{
\delta}\right)\cdot mn/\delta\right)}.$
Simplifying,
$=\sqrt{\frac{2}{\log^{1.1}\left(\frac{n}{\delta}\right)}\cdot O(\log n)}=O
\left(\frac{\sqrt{\log n}}{\log^{0.55}(n/\delta)}\right)\to 0.$
- 3.
Limits of Sub-Inequalities (a), (b), and (c):
(a)
$\frac{\sqrt{k\log\left(\frac{300mnB}{\delta}\right)}}{\sqrt{B}}=\frac{\sqrt{k
\cdot O(\log n)}}{\sqrt{C_{2}n\log^{20}\left(\frac{n}{\delta}\right)}}=O\left(
\frac{\sqrt{k}\cdot\sqrt{\log n}}{\sqrt{n}\cdot\log^{10}(n/\delta)}\right).$
Given $k\leq\frac{n}{\log^{2}(n/\delta)}$,
$\frac{\sqrt{k}\cdot\sqrt{\log n}}{\sqrt{n}\cdot\log^{10}(n/\delta)}\leq\frac{
\sqrt{\frac{n}{\log^{2}(n/\delta)}}\sqrt{\log n}}{\sqrt{n}\cdot\log^{10}(n/
\delta)}=\frac{\log^{0.5}(n)}{\log^{11}(n/\delta)}\to 0.$
(b)
$\frac{\sqrt{2k\log\left(\frac{400mn}{\delta}\right)}}{\sqrt{d}}=\frac{\sqrt{2k
\cdot O(\log n)}}{\sqrt{k\log^{1.1}\left(\frac{n}{\delta}\right)}}=O\left(
\frac{\sqrt{\log n}}{\log^{0.55}(n/\delta)}\right)\to 0.$
(c)
$\frac{\sqrt{2nk\log\left(\frac{300mnB}{\delta}\right)}}{\sqrt{Bd}}=\frac{\sqrt
{2nk\cdot O(\log n)}}{\sqrt{C_{2}n\log^{20}\left(\frac{n}{\delta}\right)\cdot k
\log^{1.1}\left(\frac{n}{\delta}\right)}}=O\left(\frac{\sqrt{nk\log n}}{\sqrt{
nk}\cdot\log^{10.55}(n/\delta)}\right)=O\left(\frac{\sqrt{\log n}}{\log^{10.55
}(n/\delta)}\right)\to 0.$
- 4.
Bound on $\eta_{0}$ and $\eta_{1}$:
$\eta_{0}=\eta_{1}=\frac{m\epsilon\sqrt{B}}{100\log\left(\frac{n}{\delta}\right
)}=\frac{2k\epsilon\sqrt{C_{2}n\log^{20}\left(\frac{n}{\delta}\right)}}{100
\log\left(\frac{n}{\delta}\right)}=O\left({k\epsilon\sqrt{n}\log^{9}(n/\delta)
}\right).$
Comparing to the minimum of the two terms:
$\frac{3m\epsilon\sqrt{B}}{80\sqrt{\log\left(\frac{300nB}{\delta}\right)}}=O
\left(k\epsilon\sqrt{n}\log^{10}(n/\delta)\right),$
and
$\frac{mn\epsilon}{120}=O(kn\epsilon).$
Since $\eta_{0}$ scales similarly to the first term and $\sqrt{n}\log^{10}(n/\delta)\ll n$ for large $n$, the inequality $\eta_{0}\leq\min\{\cdot\}$ holds.
- 5.
Limit of $\log^{2}\left(\frac{300mnB}{\delta}\right)\frac{n\sqrt{nk}}{d\sqrt{Bd}}$:
$\log^{2}\left(\frac{300mnB}{\delta}\right)=O(\log^{2}(n/\delta)),$
and
$\displaystyle\frac{n\sqrt{nk}}{d\sqrt{Bd}}$
$\displaystyle=\frac{n\sqrt{nk}}{k\log^{1.1}\left(\frac{n}{\delta}\right)\cdot
\sqrt{C_{2}n\log^{20}\left(\frac{n}{\delta}\right)\cdot k\log^{1.1}\left(\frac
{n}{\delta}\right)}}$
$\displaystyle=O\left(\frac{n\sqrt{nk}}{k\log^{1.1}(n/\delta)\sqrt{nk}\log^{10.
55}(n/\delta)}\right)=O\left(\frac{n}{k\log^{11.65}(n/\delta)}\right).$
Given $k\geq\frac{n}{\log^{3}(n/\delta)}$,
$\frac{n}{k\log^{11.65}(n/\delta)}\leq\frac{n}{\frac{n}{\log^{3}(n/\delta)}\log
^{11.65}(n/\delta)}=\frac{1}{\log^{8.65}(n/\delta)}.$
Therefore,
$\log^{2}\left(\frac{300mnB}{\delta}\right)\frac{n\sqrt{nk}}{d\sqrt{Bd}}=O\left
(\frac{\log^{2}(n/\delta)}{\log^{8.65}(n/\delta)}\right)=O\left(\log^{-6.65}(n
/\delta)\right)\to 0\text{ as }n\to\infty.$
- 6.
Limit of $\frac{\eta_{0}\eta_{1}}{256mn^{2}}$:
As
$\eta_{0}\eta_{1}=\left(\frac{m\epsilon\sqrt{B}}{100\log\left(\frac{n}{\delta}
\right)}\right)^{2}=\frac{m^{2}\epsilon^{2}B}{10^{4}\log^{2}\left(\frac{n}{
\delta}\right)}.$
Substituting into the left term:
$\frac{\eta_{0}\eta_{1}}{256\,m\,n^{2}}=\frac{m^{2}\epsilon^{2}B}{256\times 10^
{4}\log^{2}\left(\frac{n}{\delta}\right)\,m\,n^{2}}=\frac{m\epsilon^{2}B}{256
\times 10^{4}\log^{2}\left(\frac{n}{\delta}\right)\,n^{2}}.$
Substitute $m=2k$ and $B=C_{2}n\log^{10}\left(\frac{n}{\delta}\right)$:
$\frac{\eta_{0}\eta_{1}}{256\,m\,n^{2}}=\frac{2k\epsilon^{2}C_{2}n\log^{10}
\left(\frac{n}{\delta}\right)}{256\times 10^{4}\log^{2}\left(\frac{n}{\delta}
\right)n^{2}}=\frac{2kC_{2}\log^{8}\left(\frac{n}{\delta}\right)}{256\times 10
^{4}n}.$
Substituting $C_{2}=\frac{1.28\times 10^{7}}{\epsilon^{2}}$:
$\frac{\eta_{0}\eta_{1}}{256\,m\,n^{2}}=\frac{2k\times 1.28\times 10^{7}\log^{8
}\left(\frac{n}{\delta}\right)}{256\times 10^{4}\times n}=10\times\frac{k\log^
{8}\left(\frac{n}{\delta}\right)}{n}.$
Using $k\geq\frac{n}{\log^{3}\left(\frac{n}{\delta}\right)}$ from the assumption:
$\frac{\eta_{0}\eta_{1}}{256\,m\,n^{2}}\geq 10\times\frac{\frac{n}{\log^{3}
\left(\frac{n}{\delta}\right)}\log^{8}\left(\frac{n}{\delta}\right)}{n}=10\log
^{5}\left(\frac{n}{\delta}\right).$
Since $\log^{5}\left(\frac{n}{\delta}\right)>\log n$ for sufficiently large $n$, we have:
$10\log^{5}\left(\frac{n}{\delta}\right)>10\log n.$
Thus,
$\frac{\eta_{0}\eta_{1}}{256\,m\,n^{2}}>10\log n.$
- 7.
We aim to show that:
$\frac{\eta_{2}}{\eta_{0}}\to 0\quad\text{as}\quad n\to\infty.$
Here
$\eta_{0}=\frac{2k\epsilon\sqrt{C_{2}\,n\,\log^{10}\left(\frac{n}{\delta}\right
)}}{100\log\left(\frac{n}{\delta}\right)}=\frac{2k\epsilon\sqrt{C_{2}}\sqrt{n}
\log^{5}\left(\frac{n}{\delta}\right)}{100\log\left(\frac{n}{\delta}\right)}=
\frac{k\epsilon\sqrt{C_{2}}\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)}{50}.$
Thus, the ratio $\frac{\eta_{2}}{\eta_{0}}$ is:
$\frac{\eta_{2}}{\eta_{0}}=\frac{\frac{4k\epsilon}{3}}{\frac{k\epsilon\sqrt{C_{
2}}\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)}{50}}=\frac{4k\epsilon\times 5
0}{3k\epsilon\sqrt{C_{2}}\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)}=\frac{
200}{3\sqrt{C_{2}}\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)}.$
Substituting $C_{2}=\frac{1.28\times 10^{7}}{\epsilon^{2}}$:
$\frac{\eta_{2}}{\eta_{0}}=\frac{200}{3\sqrt{\frac{1.28\times 10^{7}}{\epsilon^
{2}}}\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)}=\frac{200\epsilon}{3\times
3
580\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)}.$
As $n\to\infty$, $\sqrt{n}\log^{4}\left(\frac{n}{\delta}\right)\to\infty$, hence:
$\frac{\eta_{2}}{\eta_{0}}\to 0.$
- 8.
Limit of $(\eta_{0}+\eta_{1})<5n\eta_{2}$:
$\eta_{0}+\eta_{1}=2\cdot\frac{m\epsilon\sqrt{B}}{100\log\left(\frac{n}{\delta}
\right)}=O\left(\frac{m\epsilon\sqrt{B}}{\log\left(\frac{n}{\delta}\right)}
\right).$
$5n\eta_{2}=5n\cdot\frac{4k\epsilon}{3}=O(nk\epsilon).$
Given $\sqrt{B}=O\left(\sqrt{n}\log^{10}(n/\delta)\right)$,
$\frac{m\epsilon\sqrt{B}}{\log\left(\frac{n}{\delta}\right)}=O\left(\frac{k
\epsilon\sqrt{n}\log^{10}(n/\delta)}{\log(n/\delta)}\right)=O\left(k\epsilon
\sqrt{n}\log^{9}(n/\delta)\right).$
Since $nk\epsilon$ grows faster than $k\epsilon\sqrt{n}\log^{9}(n/\delta)$ for large $n$, the inequality $(\eta_{0}+\eta_{1})<5n\eta_{2}$ holds.

∎

###### Proof.

We will set our hyperparameters according to 2.
We will first warmup our analysis on linear loss $\ell(\hat{y},y)=(-1)^{y}\hat{y}$.
We can rewrite the training dynamics in this case as

$$ $\displaystyle w^{(t+1)}$ $\displaystyle=w^{(t)}-\eta_{t}\frac{1}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k} \nabla_{w}\ell(\mathcal{T}[w^{(t)}](\textbf{b}^{(t,s)})[i],\textbf{b}^{(t,s)}[ i+1])$ $\displaystyle=w^{(t)}-\eta_{t}\frac{1}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^ {\textbf{b}^{(t,s)}[i+1]}\nabla_{w}\mathcal{T}[w^{(t)}](\textbf{b}^{(t,s)})[i].$ (2) $$

We will train the model end to end in three steps and analyze the evolution of all the weight changes. The final results can be shown combining Lemmas 26 and 21. Lemma 27 generalizes the analysis to hinge loss. It is also easy to verify that all the results hold if the parameters has only $O(\log n)$ precision.
∎

###### Assumption 2 .

###### Lemma 5 .

###### Proof.

###### Proof.

#### A.4.1 Auxiliary Statistics

To simplify our calculations, we define several auxiliary statistics on the data. Table 1 provides the definitions and rough orders of these statistics, including logarithmic terms.

**Table 1: Definitions and rough orders of auxiliary statistics, including logarithmic terms. Here, $t\in\{0,1,2\}$, $i^{\prime}\in[n+1,n+k+1]$,$i,j\in[n+k]$, $b,b_{1},b_{2}\in\{0,1\}$, $s\in[B]$, and $r\in[2m]$.**
| Statistic | Definition | Rough Order |
| --- | --- | --- |
| $\delta_{t,i,b}$ | $\displaystyle\sum_{s=1}^{B}(-1)^{\textbf{b}^{(t,s)}[i+1]}\bm{1}\left(\textbf{b<br>}^{(t,s)}[i]=b\right)$ | $\mathcal{O}\left(\sqrt{B\log\left(\tfrac{nB}{\delta}\right)}\right)$ |
| $\alpha_{t,i,j,b_{1},b_{2}}$ | $\displaystyle\sum_{s=1}^{B}(-1)^{\textbf{b}^{(t,s)}[i+1]}\bm{1}\left(\textbf{b<br>}^{(t,s)}[j]=b_{1},\textbf{b}^{(t,s)}[i]=b_{2}\right)$ | $\bar{\alpha}_{i,j,b_{1},b_{2}}$+ $\mathcal{O}\left(\sqrt{B\log\left(\tfrac{nB}{\delta}\right)}\right)$ |
| $\bar{\alpha}_{i,j,b_{1},b_{2}}$ | $\displaystyle\begin{cases}\frac{(-1)^{b_{1}+b_{2}}\bm{1}\left(j=S[i]\right)B}{<br>4},&i\geq n+2\\<br>\frac{(-1)^{b_{1}}\bm{1}\left(j=S[i],\ b_{2}=0\right)B}{2},&i=n+1\end{cases}$ | $\mathcal{O}\left(B\right)$ when $j=S[i]$, 0 otherwise |
| $\beta_{t,j,b_{1}}$ | $\displaystyle\sum_{\begin{subarray}{c}i=\max\{n+1,j\}\\<br>S[i]\neq j\end{subarray}}^{n+k}\sum_{b_{2}=0}^{1}\frac{1}{i}\alpha_{t,i,j,b_{1<br>},b_{2}}$ | $\mathcal{O}\left(\frac{\sqrt{Bk\log\left(\tfrac{nB}{\delta}\right)}}{n}\right)$ |
| $\psi_{t,r,j,b_{1}}$ | $\displaystyle\sum_{\begin{subarray}{c}i=\max\{n+1,j\}\\<br>S[i]\neq j\end{subarray}}^{n+k}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}\left(\nu_{r<br>,i,b_{2}}>0\right)\alpha_{t,i,j,b_{1},b_{2}}$ | $\mathcal{O}\left(\frac{\sqrt{Bk\log\left(\tfrac{mnB}{\delta}\right)}}{n}\right)$ |
| $\gamma_{i,b,i^{\prime},b_{2}}$ | $\displaystyle\frac{1}{i^{\prime}}\left(\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}><br>0\right)\bm{1}\left(\nu_{r,i^{\prime},b_{2}}>0\right)\right)-\frac{1}{i^{<br>\prime}}\frac{1+\bm{1}\left(i=i^{\prime},\ b=b_{2}\right)}{2}m$ | $\mathcal{O}\left(\frac{\sqrt{m\log\left(\tfrac{n}{\delta}\right)}}{n}\right)$ |
| $\zeta_{t,i,b,j,b_{1}}$ | $\displaystyle\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}>0\right)\psi_{t,r,j,b_{1}}$ | $\mathcal{O}\left(\frac{m\sqrt{kB}\log\left(\tfrac{mnB}{\delta}\right)}{n}\right)$ |

###### Lemma 6 .

Fixing $j,b_{1},b_{2}$, let $\mathcal{F}_{(i,t,s)}$ be the filtration generated by the random variables $\{\textbf{b}^{(t,s)}[i]\mid\max\{n+1,j\}\leq i,S[i]\neq j,t^{\prime}\leq t,s
\leq B\}$ ordered in lexicographic order, the process

$$ $(-1)^{\textbf{b}^{(t,s)}[i+1]}\bm{1}(\textbf{b}^{(t,s)}[j]=b_{1},\textbf{b}^{( t,s)}[i]=b_{2})$ $$

is a martingale with respect to $\{\mathcal{F}_{t}\}_{t}$.

###### Proof.

$\textbf{b}[i+1]=\textbf{b}[i]\oplus\textbf{b}[S[i]]$, and since $S[i]\neq j$, $\textbf{b}[S[i]]$ is independent from $\textbf{b}[j],\ldots,\textbf{b}[i]$. Thus,

$$ $\mathbb{E}[(-1)^{\textbf{b}^{(t,s)}[i+1]}\mid\mathcal{F}_{(i,t,s)}]=0,$ $$

provided $\mathbf{1}(\textbf{b}^{(t+1,s)}[j]=b_{1})$ is constant. Hence, this process is a martingale.
∎

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 $1-\delta/10$, the following bounds hold for all $t\in[3]$, $i,j\in[n+k]$, $b,b_{1},b_{2}\in\{0,1\}$, and $r\in[m]$:

$$ $\displaystyle|\delta_{t,i,b}|$ $\displaystyle\leq\sqrt{B\log\left(\tfrac{300nB}{\delta}\right)}$ $\displaystyle|\alpha_{t,i,j,b_{1},b_{2}}-\bar{\alpha}_{i,j,b_{1},b_{2}}|$ $\displaystyle\leq\sqrt{2B\log\left(\tfrac{300nB}{\delta}\right)}\bm{1}\left(j= S[i]\right)$ $\displaystyle|\beta_{t,j,b_{1}}|$ $\displaystyle\leq\frac{\sqrt{2Bk\log\left(\tfrac{300nB}{\delta}\right)}}{n}$ $\displaystyle|\psi_{t,r,j,b_{1}}|$ $\displaystyle\leq\frac{2\sqrt{Bk\log\left(\tfrac{300mnB}{\delta}\right)}}{n}$ $\displaystyle\left|\gamma_{i,b,i^{\prime},b_{2}}\right|$ $\displaystyle\leq\frac{\sqrt{2m\log\left(\tfrac{200n}{\delta}\right)}}{n}$ $\displaystyle\left|\zeta_{t,i,b,j,b_{1}}\right|$ $\displaystyle\leq\frac{4m\sqrt{kB\log\left(\tfrac{300mnB}{\delta}\right)}}{n}$ $$

###### Proof.

Fix any $t$, $i$, and $b$. For each $s\in[B]$, define $X_{s}=(-1)^{\textbf{b}^{(t,s)}[i+1]}\bm{1}\left(\textbf{b}^{(t,s)}[i]=b\right)$. The sequence $\{X_{s}\}_{s=1}^{B}$ consists of independent random variables with $|X_{s}|\leq 1$ and zero mean. By the Azuma-Hoeffding inequality,

$$ $\Pr\left(\left|\delta_{t,i,b}\right|>R\right)\leq 2\exp\left(-\frac{R^{2}}{2B} \right).$ $$

Setting $R=\sqrt{2B\log\left(\tfrac{300nB}{\delta}\right)}$, we get

$$ $\Pr\left(\left|\delta_{t,i,b}\right|>R\right)\leq\frac{\delta}{150nB}.$ $$

Applying a union bound over all $t$, $i$, and $b$, we ensure that with probability at least $1-\delta/50$, the bound on $\delta_{t,i,b}$ holds for all choices.

Similarly, for $\alpha_{t,i,j,b_{1},b_{2}}$, we fix $t$, $i$, $j$, $b_{1}$, and $b_{2}$, and define

$$ $Y_{s}=(-1)^{\textbf{b}^{(t,s)}[i+1]}\bm{1}\left(\textbf{b}^{(t,s)}[j]=b_{1}, \textbf{b}^{(t,s)}[i]=b_{2}\right).$ $$

Again, $\{Y_{s}\}_{s=1}^{B}$ are independent with $|Y_{s}|\leq 1$ and mean $\bar{\alpha}_{i,j,b_{1},b_{2}}/B$. Applying the Azuma-Hoeffding inequality as before, we obtain the stated bound for $\alpha_{t,i,j,b_{1},b_{2}}$.

For $\beta_{t,j,b_{1}}$ and $\psi_{t,r,j,b_{1}}$, the bound can be derived similarly combining Lemma 6 and Azuma-Hoeffding bound.

For $\gamma_{i,b,i^{\prime},b_{2}}$, consider the sum $\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}>0\right)\bm{1}\left(\nu_{r,i^{\prime},b
_{2}}>0\right)$. Each term is an independent Bernoulli random variable with mean $\frac{1}{2}\left(1+\bm{1}\left(i=i^{\prime},\ b=b_{2}\right)\right)$. The variance of the sum is bounded by $m$. Applying the Azuma-Hoeffding inequality, we get

$$ $\Pr\left(\left|\gamma_{i,b,i^{\prime},b_{2}}\right|>R\right)\leq 2\exp\left(- \frac{R^{2}n^{2}}{2m}\right).$ $$

Setting $R=\frac{\sqrt{2m\log\left(\tfrac{200n}{\delta}\right)}}{n}$, we obtain the desired bound.

For $\zeta_{t,i,b,j,b_{1}}$, we note that it is a sum over $2m$ terms, each involving $\psi_{t,r,j,b_{1}}$, which are bounded as in Lemma 7. The total number of terms is $2m$, and each term is bounded by $\frac{2\sqrt{Bk\log\left(\tfrac{300mnB}{\delta}\right)}}{n}$.
∎

###### Lemma 6 .

###### Proof.

###### Lemma 7 .

###### Proof.

#### A.4.2 First Step: Configuring the MLPs

###### Lemma 8 .

For the dynamics following Equation 2, attention weight stays constant in the first step, i.e., $A^{(1)}=A^{(0)}=0$.

###### Proof.

Because $W_{:,d+1:2d}=0$, $\nabla_{A}\mathcal{T}[w^{(0)}](\textbf{b})[i]=0$.
∎

###### Lemma 9 .

For the dynamics following Equation 2, attention output stays constant in the first step, and

$$ $\displaystyle\forall t\in\{0,1\},i\in[n+k+1],\mathcal{A}[A^{(t)}](\mathbf{E}( \textbf{b}))[i]=\frac{1}{i}\sum_{j=1}^{i}e_{i,\textbf{b}[i]}.$ $$

###### Proof.

This follows from the definition of the attention module and Lemma 8.
∎

###### Lemma 10 .

For an input b, With probability $1-\frac{\delta}{50Bm}$ , at initialization, whether a neuron $r$ outputs nonzero value at position $i$ is determined by $\nu_{r,i,\textbf{b}[i]}$.

$$ $\displaystyle\bm{1}\left(\left\langle W^{0}_{r,1:d},\mathbf{E}(\textbf{b})[i] \right\rangle>0\right)=\bm{1}\left(\nu_{r,i,\textbf{b}[i]}>0\right).$ $$

Further,

$$ $\displaystyle\left|\left\langle W^{0}_{r,1:d},\mathbf{E}(\textbf{b})[i]\right \rangle-\nu_{r,i,\textbf{b}[i]}\right|<\epsilon/100.$ $$

###### Proof.

Under 1,

$$ $\displaystyle W^{0}_{r,1:d}$ $\displaystyle=\sum_{i=n+1}^{n+k}\nu_{r,i,\textbf{b}[i]}e_{i,\textbf{b}[i]}$ $$

Because,

$$ $\displaystyle\sum_{i=n+1}^{n+k}\nu_{r,i,\textbf{b}[i]}^{2}=k\epsilon^{2}.$ $$

By Lemma 30, with probability $1-\frac{\delta}{50Bm}$,

$$ $\displaystyle\left|\left\langle W^{0}_{r,1:d},\mathbf{E}(\textbf{b})[i]\right \rangle-\nu_{r,i,\textbf{b}[i]}\right|\leq\sqrt{\frac{2\log(\frac{100Bm}{ \delta})k}{d}}\epsilon$ $$

By Lemma 5.1, $\sqrt{\frac{2k}{d}\log(Bm/100\delta)}<0.01$, we can conclude that

$$ $\displaystyle\Pr\left(\left|\sum_{i^{\prime}=n+1}^{n+k}\sum_{l=1}^{d}\nu_{r,i^ {\prime},b^{\prime}}\bm{1}(i^{\prime}\neq i)e_{i^{\prime},\textbf{b}[i^{\prime }]}[l]e_{i,\textbf{b}[i]}[l]\right|\geq\epsilon/100\right)\leq\frac{\delta}{50 Bm}.$ $$

The proof is then complete.
∎

###### Lemma 11 .

With probability $1-\frac{\delta}{50}$ , the following $\mathcal{E}_{1}$ happens: for any input $\textbf{b}^{(0,s)}$ in the first batch, at initialization, for any $r$, whether a neuron $r$ outputs nonzero value at position $i$ is determined by $\nu_{r,i,\textbf{b}[i]}$.

$$ $\displaystyle\bm{1}\left(\left\langle W^{0}_{r,1:d},\mathbf{E}(\textbf{b})[i] \right\rangle>0\right)=\bm{1}\left(\nu_{r,i,\textbf{b}[i]}>0\right).$ $$

Further,

$$ $\displaystyle\left|\left\langle W^{0}_{r,1:d},\mathbf{E}(\textbf{b})[i]\right \rangle-\nu_{r,i,\textbf{b}[i]}\right|<\epsilon/100.$ $$

###### Proof.

Apply union bound to Lemma 10 over neuron dimension and data in the first batch.
∎

###### Lemma 12 .

When $\mathcal{E}_{1}$ defined in Lemma 11 happens, for all $i^{\prime}\in[n+1,n+k+1],b^{\prime}\in\{0,1\}$, the gradient on $W_{r,1:d}$ satisfies that,

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,1:d}}=\frac{h_{r}}{B}\sum_{i=n+1} ^{n+k}\sum_{b=0}^{1}\bm{1}(\nu_{r,i,b}>0)\delta_{0,i,b}e_{i,b}.$ $$

###### Proof.

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,1:d}}=$ $\displaystyle\frac{1}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^{\textbf{b}^{(t,s )}[i+1]}\nabla_{W_{r,1:d}}\mathcal{T}[w^{(t)}](\textbf{b}^{(t,s)})[i]$ $\displaystyle=$ $\displaystyle\frac{h_{r}}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^{\textbf{b}^{ (t,s)}[i+1]}e_{i,\textbf{b}^{(t,s)}[i]}\bm{1}(\nu_{r,i,\textbf{b}^{(t,s)}[i]}> 0).$ $$

We can then expand this formula to have

$$ $\displaystyle\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^{\textbf{b}^{(t,s)}[i+1]}e_{ i,\textbf{b}^{(t,s)}[i]}\bm{1}(\nu_{r,i,\textbf{b}^{(t,s)}[i]}>0)$ $\displaystyle=$ $\displaystyle\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}\sum_{b=0}^{1}(-1)^{\textbf{b}^{( t,s)}[i+1]}\bm{1}(\textbf{b}^{(t,s)}[i]=b)e_{i,b}\bm{1}(\nu_{r,i,b}>0)$ $\displaystyle=$ $\displaystyle\sum_{i=n+1}^{n+k}\sum_{b=0}^{1}e_{i,b}\bm{1}(\nu_{r,i,b}>0)\sum_ {s=1}^{B}(-1)^{\textbf{b}^{(t,s)}[i+1]}\bm{1}(\textbf{b}^{(t,s)}[i]=b)$ $\displaystyle=$ $\displaystyle\sum_{i=n+1}^{n+k}\sum_{b=0}^{1}e_{i,b}\bm{1}(\nu_{r,i,b}>0) \delta_{0,i,b}.$ $$

The proof is then complete.
∎

###### Lemma 13 .

With probability $1-0.13\delta$, for all $i^{\prime}\in[n+1,n+k+1],b^{\prime}\in\{0,1\}$, the gradient on $W_{r,1:d}$ satisfies that,

$$ $\displaystyle\left|\left\langle\frac{d\mathcal{L}^{(0)}}{dW_{r,1:d}},e_{i,b} \right\rangle\right|\leq\frac{\sqrt{\log(300nB/\delta)}}{m\sqrt{B}}.$ $$

###### Proof.

Combining Lemmas 30 and 12, we have that for any $r$, with probability $1-\delta/100m$,

$$ $\displaystyle\left|\left\langle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b=0}^{1} \bm{1}(\nu_{r,i,b}>0)\delta_{0,i,b}e_{i,b},e_{i,b}\right\rangle\right|\leq \frac{1}{2mB}\sup\delta_{t,i,b}\left(1+\sqrt{\frac{2k\log\left(\frac{200m}{ \delta}\right)}{d}}\right)$ $$

By Lemmas 7 and 5, with probability $1-11\delta/100$, for all $r$, it holds that for all $r,i^{\prime},b^{\prime}$

$$ $\displaystyle\left|\left\langle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b=0}^{1} \bm{1}(\nu_{r,i,b}>0)\delta_{0,i,b}e_{i,b},e_{i,b}\right\rangle\right|<\frac{1 }{mB}\sup\delta_{0,i,b}<\frac{\sqrt{\log(300nB/\delta)}}{m\sqrt{B}}.$ $$

Combining with Lemma 12 and applying union bound concludes the proof.
∎

###### Lemma 14 .

For $\bar{\alpha}$ and $\psi$ defined in Section A.4.1, When $\mathcal{E}_{1}$ defined in Lemma 11 happens, for all $i^{\prime}\in[n+1,n+k+1],b^{\prime}\in\{0,1\}$, the gradient on $W_{r,d+1:2d}$ satisfies that,

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=$ $\displaystyle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0} ^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}e_{S[ i],b_{1}}+\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}\psi_{0,r,j,b_{1}}e _{j,b_{1}}.$ $$

###### Proof.

With probability $1-\delta/50$, $\mathcal{E}_{1}$ defined in Lemma 11 happens, and

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=$ $\displaystyle\frac{1}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^{\textbf{b}^{(t,s )}[i+1]}\nabla_{W_{r,1:d}}\mathcal{T}[w^{(t)}](\textbf{b}^{(t,s)})[i]$ $\displaystyle=$ $\displaystyle\frac{h_{r}}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^{\textbf{b}^{ (t,s)}[i+1]}\mathcal{A}[A^{(t)}](\mathbf{E}(\textbf{b}^{(t,s)}))[i]\bm{1}(\nu_ {r,i,\textbf{b}^{(t,s)}[i]}>0)$ $\displaystyle=$ $\displaystyle\frac{h_{r}}{B}\sum_{s=1}^{B}\sum_{i=n+1}^{n+k}(-1)^{\textbf{b}^{ (t,s)}[i+1]}\bm{1}(\nu_{r,i,\textbf{b}^{(t,s)}[i]}>0)\frac{1}{i}\sum_{j=1}^{i} \sum_{b_{1}=0}^{1}\bm{1}\left(\textbf{b}[j]=b\right)e_{j,b_{1}}$ $$

Rearranging the summation, and use $\alpha$ defined in Lemma 7,,

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=$ $\displaystyle\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}e_{j,b_{1}}\sum_ {i=\max\{n+1,j\}}^{n+k}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0) \alpha_{0,i,j,b_{1},b_{2}}.$ $$

Breaking the summation,

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=$ $\displaystyle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0} ^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0)\alpha_{0,i,S[i],b_{1},b_{2}}e_{S[i],b _{1}}$ $\displaystyle+\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}e_{j,b_{1}}\sum _{i\in[\max\{n+1,j\}.n+k],S[i]\neq j}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}(\nu_{ r,i,b_{2}}>0)\alpha_{0,i,j,b_{1},b_{2}}$ $$

This simplifies to

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=$ $\displaystyle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0} ^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}e_{S[ i],b_{1}}$ $\displaystyle+\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}\psi_{0,r,j,b_{ 1}}e_{j,b_{1}}.$ $$

The proof is then complete.
∎

This leads to the following bound on the gradient,

###### Lemma 15 .

With probability $1-0.16\delta$, for all $i\in[n+k+1],b\in\{0,1\}$,
the gradient on $W_{r,d+1:2d}$ satisfies that,

$$ $\displaystyle\left|\left\langle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}},e_{i,b }\right\rangle\right|\leq\frac{3}{2mn}.$ $$

###### Proof.

Denote

$$ $\displaystyle G_{r}=$ $\displaystyle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0} ^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}e_{S[ i],b_{1}}+\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}\psi_{0,r,j,b_{1}}e _{j,b_{1}}.$ $$

By Lemma 30, for all $r\in[2m],i\in[n+k+1],b\in\{0,1\}$, with probability $1-\delta/25$,

$$ $\displaystyle\left|\left\langle G_{r},e_{i,b}\right\rangle\right|\leq\frac{1}{ 2mB}\left|\sum_{i^{\prime}=n+1}^{n+k}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}(\nu_{ r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b,b_{2}}\bm{1}(S[i^{\prime}]=i)+\psi_{0,r,i, b}\right|$ $\displaystyle+\frac{\sqrt{2\log(\frac{400mn}{\delta})}}{2mB\sqrt{d}}\left( \sqrt{\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}( \nu_{r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}^{2}}+\sqrt{\sum_{j=1}^{n+k }\sum_{b_{1}=0}^{1}\psi_{0,r,j,b_{1}}^{2}}\right)$ $$

By Lemma 7, we have with probability $1-\delta/10$,

$$ $\displaystyle\sup\psi_{0,r,i,b}\leq\frac{2\sqrt{Bk\log\left(\tfrac{300mnB}{ \delta}\right)}}{n},\quad\sup\bar{\alpha}_{i,S[i],b_{1},b_{2}}\leq B/2.$ $$

Combining with Lemma 5,

$$ $\displaystyle\frac{1}{2mB}\left|\sum_{i^{\prime}=n+1}^{n+k}\sum_{b_{2}=0}^{1} \frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b,b_{2}}\bm{1}(S[i^{ \prime}]=i)\right|\leq\frac{1}{2mn}$ $\displaystyle\frac{1}{2mB}\left|\psi_{0,r,i,b}\right|\leq\frac{2\sqrt{k\log \left(\tfrac{300mnB}{\delta}\right)}}{mn\sqrt{B}}\leq\frac{1}{200mn}$ $\displaystyle\frac{\sqrt{2\log(\frac{400mn}{\delta})}}{2mB\sqrt{d}}\sqrt{\sum_ {i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_ {2}}>0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}^{2}}\leq\frac{\sqrt{2k\log(\frac{400 mn}{\delta})}}{2mn\sqrt{d}}\leq\frac{1}{200mn},$ $\displaystyle\frac{\sqrt{2\log(\frac{400mn}{\delta})}}{2mB\sqrt{d}}\sqrt{\sum_ {j=1}^{n+k}\sum_{b_{1}=0}^{1}\psi_{0,r,j,b_{1}}^{2}}\leq\frac{2\sqrt{2nk}{\log (\frac{300mnB}{\delta})}}{mn\sqrt{Bd}}\leq\frac{1}{200mn}.$ $$

Further by Lemma 14, we have that $\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=G_{r}$ with probability $1-\delta/50$. This concludes the proof.
∎

###### Lemma 8 .

###### Proof.

###### Lemma 9 .

###### Proof.

###### Lemma 10 .

###### Proof.

###### Lemma 11 .

###### Proof.

###### Lemma 12 .

###### Proof.

###### Lemma 13 .

###### Proof.

###### Lemma 14 .

###### Proof.

###### Lemma 15 .

###### Proof.

#### A.4.3 Step 2: Configuring the Attention

###### Lemma 16 .

With probability $1-0.29{\delta}$ , the following $\mathcal{E}_{2}$ happens: the change in the neuron outcome is small after the first step for any input b,

$$ $\displaystyle\sup_{r}\left|\left\langle W^{1}_{r}-W^{0}_{r},[\mathbf{E}( \textbf{b})[i],\mathcal{A}[A^{1}](\mathbf{E}(\textbf{b}))[i]\right\rangle \right|<3\epsilon/80.$ $$

This leads to the case that for any input $\textbf{b}^{(1,s)}$ in the second batch, for any $r$, whether a neuron $r$ outputs nonzero value at position $i$ is determined by $\nu_{r,i,\textbf{b}[i]}$.

$$ $\displaystyle\bm{1}\left(\left\langle W^{1}_{r,1:d},\mathbf{E}(\textbf{b})[i] \right\rangle>0\right)=\bm{1}\left(\nu_{r,i,\textbf{b}[i]}>0\right).$ $$

###### Proof.

By Lemma 13, with probability $1-0.13\delta$,

$$ $\displaystyle\sup_{r}|\left\langle W^{1}_{r,1:d}-W^{0}_{r,1:d},\mathbf{E}( \textbf{b})[i]\right\rangle|=\eta_{0}\left|\left\langle\frac{d\mathcal{L}^{(0) }}{dW_{r,1:d}},\mathbf{E}(\textbf{b})[i]\right\rangle\right|\leq\eta_{0}\frac{ \sqrt{\log(300nB/\delta)}}{m\sqrt{B}}.$ $$

By Lemma 15,with probability $1-0.16\delta$,

$$ $\displaystyle\sup_{r}|\left\langle W^{1}_{r,d+1:2d}-W^{0}_{r,d+1:2d},\mathcal{ A}[A^{1}](\mathbf{E}(\textbf{b}))[i]\right\rangle|=\eta_{0}\left|\left\langle \frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}},\mathcal{A}[A^{1}](\mathbf{E}(\textbf {b}))[i]\right\rangle\right|\leq\eta_{0}\frac{3}{2mn}.$ $$

Hence, as $\eta_{0}\leq\min\{\frac{3m\epsilon\sqrt{B}}{80\sqrt{\log(300nB/\delta)}},\frac
{mn\epsilon}{120}\}$ (Lemma 5.3), we conclude that

$$ $\displaystyle\sup_{r}\left|\left\langle W^{1}_{r}-W^{0}_{r},[\mathbf{E}( \textbf{b})[i],\mathcal{A}[A^{1}](\mathbf{E}(\textbf{b}))[i]\right\rangle \right|<3\epsilon/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.58\delta$, inequalities in Lemma 7 hold and for $t\in\{0,1\},\forall i\in[n+k],\forall b\in\{0,1\}$,

$$ $\displaystyle\frac{d\mathcal{L}^{(t)}}{dW_{r,1:d}}=$ $\displaystyle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b=0}^{1}\bm{1}(\nu_{r,i,b} >0)\delta_{t,i,b}e_{i,b}.$ $\displaystyle\frac{d\mathcal{L}^{(t)}}{dW_{r,d+1:2d}}=$ $\displaystyle\frac{h_{r}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0} ^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}}>0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}e_{S[ i],b_{1}}+\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}\psi_{t,r,j,b_{1}}e _{j,b_{1}}.$ $\displaystyle\left|\left\langle\frac{d\mathcal{L}^{(t)}}{dW_{r,d+1:2d}},e_{i,b }\right\rangle\right|\leq$ $\displaystyle\frac{3}{2mn}.$ $$

Moreover for any input b,

$$ $\displaystyle\sup_{r}\left|\left\langle W^{t}_{r}-W^{0}_{r},[\mathbf{E}( \textbf{b})[i],\mathcal{A}[A^{1}](\mathbf{E}(\textbf{b}))[i]\right\rangle \right|<3\epsilon/40.$ $$

Further,

$$ $\displaystyle\bm{1}\left(\left\langle W^{t}_{r,1:d},\mathbf{E}(\textbf{b}^{(t, s)})[i]\right\rangle>0\right)=\bm{1}\left(\nu_{r,i,\textbf{b}^{(t,s)}[i]}>0 \right).$ $$

###### Proof.

Consider the backward calculation of MLP, it is decided by the input and the corresponding activation of neurons, which jointly follows the same distribution in the first and second steps. The proof is similar to Lemmas 14, 12 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,

$$ $\displaystyle\kappa_{t,i,b}$ $\displaystyle=-\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}>0\right)h_{r}W^{(t)}_{r, d+1:2d}$ $\displaystyle\Delta_{t,i,b,j,b_{1}}$ $\displaystyle=\langle\kappa_{t,i,b},e_{j,b_{1}}\rangle$ $$

Intuitively, the FFN with weight $W^{(1)}$ maps attention output $x$ to $-\kappa_{i,\textbf{b}[i]}x$ at position $i$.

###### Lemma 18 .

For $\alpha,\beta,\gamma,\zeta$ defined in Section A.4.1,
when the event in Lemma 17 happens, for any $i$ and $b$,

$$ $\displaystyle\kappa_{1,i,b}=$ $\displaystyle\frac{\eta_{0}}{8Bmi}\sum_{b_{1}=0}^{1}\bar{\alpha}_{i,S[i],b_{1} ,b}e_{S[i],b_{1}}+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{i^{\prime}=n+1}^{n+k} \sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\gamma_{i,b,i^{\prime},b_{2}}\bar{\alpha}_ {i^{\prime},S[i^{\prime}],b_{1},b_{2}}e_{S[i],b_{1}}\right)$ $\displaystyle+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1} \zeta_{0,i,b,j,b_{1}}e_{j,b_{1}}\right).$ $$

###### Proof.

When the event in Lemma 17 happens, we have that

$$ $\displaystyle\frac{d\mathcal{L}^{(0)}}{dW_{r,d+1:2d}}=\frac{h_{r}}{B}\sum_{i=n +1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\frac{1}{i}\bm{1}(\nu_{r,i,b_{2}} >0)\bar{\alpha}_{i,S[i],b_{1},b_{2}}e_{S[i^{\prime}],b_{1}}$ $\displaystyle+\frac{h_{r}}{B}\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1}\psi_{0,r,j,b_{ 1}}e_{j,b_{1}}.$ $$

Summing over the axis of $r$ as in $\kappa$,

$$ $\displaystyle\kappa_{1,i,b}=$ $\displaystyle\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{i^{\prime}=n+1}^{n+k}\sum_{b_ {1}=0}^{1}\sum_{b_{2}=0}^{1}\left(\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}>0 \right)\bm{1}(\nu_{r,i^{\prime},b_{2}}>0)\right)\frac{1}{i^{\prime}}\bar{ \alpha}_{i^{\prime},S[i^{\prime}],b_{1},b_{2}}e_{S[i],b_{1}}\right)$ $\displaystyle+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1} e_{j,b_{1}}\left(\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}>0\right)\psi_{0,r,j,b_ {1}}\right)\right).$ $$

Recall that
$\gamma_{i,b,i^{\prime},b_{2}}=\frac{1}{i^{\prime}}\left(\sum_{r=1}^{2m}\bm{1}
\left(\nu_{r,i,b}>0\right)\bm{1}\left(\nu_{r,i^{\prime},b_{2}}>0\right)\right)
-\frac{1}{i^{\prime}}\frac{1+\bm{1}\left(i=i^{\prime},\ b=b_{2}\right)}{2}m$. We then have

$$ $\displaystyle\kappa_{1,i,b}=$ $\displaystyle\frac{\eta_{0}}{4Bm}\left(\sum_{i^{\prime}=n+1}^{n+k}\sum_{b_{1}= 0}^{1}\sum_{b_{2}=0}^{1}\frac{1}{i^{\prime}}\frac{1+\bm{1}\left(i=i^{\prime}, \ b=b_{2}\right)}{2}\bar{\alpha}_{i^{\prime},S[i^{\prime}],b_{1},b_{2}}e_{S[i^ {\prime}],b_{1}}\right)$ $\displaystyle+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{i^{\prime}=n+1}^{n+k}\sum_{b _{1}=0}^{1}\sum_{b_{2}=0}^{1}\gamma_{i,b,i^{\prime},b_{2}}\bar{\alpha}_{i^{ \prime},S[i^{\prime}],b_{1},b_{2}}e_{S[i],b_{1}}\right)$ $\displaystyle+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1} e_{j,b_{1}}\left(\sum_{r=1}^{2m}\bm{1}\left(\nu_{r,i,b}>0\right)\psi_{0,r,j,b_ {1}}\right)\right).$ $$

The first term can be greatly simplified, as

$$ $\displaystyle\alpha_{i^{\prime},S[i^{\prime}],b_{1},0}+\alpha_{i^{\prime},S[i^ {\prime}],b_{1},1}=0.$ $$

This concludes the proof.
∎

###### Lemma 19 .

With probability $1-0.6\delta$, inequalities and equalities in Lemmas 7, 17 and 18 hold, and for all $i,b_{2},j,b_{1}$,

$$ $\displaystyle\left|\Delta_{1,i,b_{2},j,b_{1}}-\frac{\eta_{0}}{8Bmi}\bar{\alpha }_{i,S[i],b_{1},b_{2}}\bm{1}(j=S[i])\right|\leq\frac{10\eta_{0}}{mn}\frac{ \sqrt{nk}}{\sqrt{Bd}}\log(300mnB/\delta).$ $$

###### Proof.

Define

$$ $\displaystyle H_{i,b}=$ $\displaystyle\frac{\eta_{0}}{8Bmi}\sum_{b_{1}=0}^{1}\bar{\alpha}_{i,S[i],b_{1} ,b}e_{S[i],b_{1}}$ $\displaystyle+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{i^{\prime}=n+1}^{n+k}\sum_{b _{1}=0}^{1}\sum_{b_{2}=0}^{1}\gamma_{i,b,i^{\prime},b_{2}}\bar{\alpha}_{i^{ \prime},S[i^{\prime}],b_{1},b_{2}}e_{S[i^{\prime}],b_{1}}\right)$ $\displaystyle+\frac{\eta_{0}}{4Bm^{2}}\left(\sum_{j=1}^{n+k}\sum_{b_{1}=0}^{1} \zeta_{i,b,j,b_{1}}e_{j,b_{1}}\right)$ $$

As the inequalities in Lemma 7 hold, by Lemmas 30 and 5, with probability $1-\delta/50$, it holds that,

$$ $\displaystyle\left|\frac{\eta_{0}}{4Bm^{2}}\left\langle\left(\sum_{i^{\prime}= n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\gamma_{i,b,i^{\prime},b_{2}} \bar{\alpha}_{i^{\prime},S[i^{\prime}],b_{1},b_{2}}e_{S[i^{\prime}],b_{1}} \right),e_{j,b_{1}}\right\rangle\right|$ $\displaystyle\leq$ $\displaystyle\frac{\eta_{0}\sup\bar{\alpha}\sup\gamma}{4Bm^{2}}\left(1+{4\sqrt {\frac{k\log(200n/\delta)}{d}}}\right)$ $\displaystyle\leq$ $\displaystyle\frac{\eta_{0}}{4mn}\sqrt{\frac{2k\log\left(\frac{200n}{\delta} \right)}{{md}}}.$ $$

and,

$$ $\displaystyle\left|\frac{\eta_{0}}{4Bm^{2}}\left\langle\left(\sum_{j=1}^{n+k} \sum_{b_{1}=0}^{1}\zeta_{i,b,j,b_{1}}e_{j,b_{1}}\right),e_{j,b_{1}}\right \rangle\right|$ $\displaystyle\leq$ $\displaystyle\frac{\eta_{0}\sup\zeta}{4Bm^{2}}\left(1+{4\sqrt{\frac{n\log(200n /\delta)}{d}}}\right)$ $\displaystyle\leq$ $\displaystyle\frac{\eta_{0}}{mn}\left(\sqrt{\frac{k}{B}\log(200n/\delta)}+{4 \sqrt{\frac{nk}{dB}}}\log(200n/\delta)\right).$ $$

and,

$$ $\displaystyle\left|\frac{\eta_{0}}{8Bmi}\left\langle\sum_{b_{1}=0}^{1}\bar{ \alpha}_{i,S[i],b_{1},b}e_{S[i],b_{1}}e_{j,b_{1}}\right\rangle-\frac{\eta_{0}} {8Bmi}\bar{\alpha}_{i,S[i],b_{1},b_{2}}\bm{1}(j=S[i])\right|\leq$ $\displaystyle\frac{2\sqrt{\log(200n/\delta)}}{8mn\sqrt{d}}.$ $$

With $d<n$ and $B<nk$ (Lemma 5), we conclude that,

$$ $\displaystyle\left|\langle H_{i,b},e_{j,b_{1}}\rangle-\frac{\eta_{0}}{8Bmi} \bar{\alpha}_{i,S[i],b_{1},b_{2}}\bm{1}(j=S[i])\right|\leq\frac{10\eta_{0}}{mn }\frac{\sqrt{nk}}{\sqrt{Bd}}\log(300mnB/\delta).$ $$

By Lemma 18, we have the desired result.
∎

###### Lemma 20 .

Under the setting of Lemma 19, for every $\textbf{b}\in\{\textbf{b}^{(1,s)}\mid s\in[B]\}$,

$$ $\frac{d\ell(\mathcal{T}(\textbf{b})[i],\textbf{b}[i+1])}{dA}=(-1)^{\textbf{b}[ i+1]+1}\sum_{j=1}^{i}\langle\kappa_{1,i,\textbf{b}[i]},e_{j,\textbf{b}[j]} \rangle\frac{\partial\mathrm{softmax}\left(\mathbf{E}(\textbf{b})^{T}A\mathbf{ E}(\textbf{b})\right)[j,i]}{\partial A}.$ $$

with,

$$ $\frac{\partial\mathrm{softmax}\left(\mathbf{E}(\textbf{b})^{T}A\mathbf{E}( \textbf{b})\right)[j,i]}{\partial A}=\mathrm{softmax}(Z)[j,i]\cdot\left( \mathbf{E}(\textbf{b})[j]-\sum_{p=1}^{i}\mathrm{softmax}(Z)[p,i]\mathbf{E}( \textbf{b})[p]\right)\mathbf{E}(\textbf{b})[i]^{T}$ $$

###### Proof.

Our goal is to calculate the gradient of the loss with respect to the attention layer. We will first calculate the gradient of the loss with respect to the output of the attention layer, and then use the chain rule. When the event in Lemma 17 happens, for any binary sequence b in the second batch, we have that for any $i$,

$$ $\displaystyle\frac{d\mathcal{T}(\textbf{b})[i]}{d\mathcal{A}(\mathbf{E}( \textbf{b}))[i]}$ $\displaystyle=\sum_{r=1}^{2m}h_{r}W^{(1)}_{r}\bm{1}\left(\left(W^{(1)}\begin{ bmatrix}\mathbf{E}\left(\textbf{b}\right)\\ \mathcal{A}\left(\mathbf{E}\left(\textbf{b}\right)\right)\end{bmatrix}\right) \left[i\right]_{r}>0\right)$ $\displaystyle=\sum_{r=1}^{2m}h_{r}W^{(1)}_{r}\bm{1}\left(\nu_{r,i,\textbf{b}[i ]}>0\right)$ $$

This then implies that,

$$ $\displaystyle\frac{d\ell(\mathcal{T}(\textbf{b})[i],\textbf{b}[i+1])}{d \mathcal{A}(\mathbf{E}(\textbf{b}))[i]}$ $\displaystyle=(-1)^{\textbf{b}[i+1]}\left(\sum_{r=1}^{2m}h_{r}W^{(1)}_{r}\bm{1 }\left(\nu_{r,i,\textbf{b}[i]}>0\right)\right)$ $$

We now calculate the gradient of $A$ (the attention matrix) to the attention output.

$$ $\displaystyle\frac{\partial\mathcal{A}(\mathbf{E}(\textbf{b}))[i]}{\partial A}$ $\displaystyle=\frac{\partial\left(\mathbf{E}(\textbf{b})\mathrm{softmax}\left( \mathbf{E}(\textbf{b})^{T}A\mathbf{E}(\textbf{b})\right)\right)[i]}{\partial A}$ $\displaystyle=\frac{\sum_{j=1}^{i}\partial\mathbf{E}(\textbf{b})[j]\mathrm{ softmax}\left(\mathbf{E}(\textbf{b})^{T}A\mathbf{E}(\textbf{b})\right)[j,i]}{ \partial A}$ $\displaystyle=\sum_{j=1}^{i}\mathbf{E}(\textbf{b})[j]\frac{\partial\mathrm{ softmax}\left(\mathbf{E}(\textbf{b})^{T}A\mathbf{E}(\textbf{b})\right)[j,i]}{ \partial A}$ $$

Hence, the gradient of the loss with respect to the attention matrix is

$$ $\displaystyle\frac{d\ell(\mathcal{T}(\textbf{b})[i],\textbf{b}[i+1])}{dA}$ $\displaystyle=$ $\displaystyle\sum_{j=1}^{i}\frac{d\ell(\mathcal{T}(\textbf{b})[i],\textbf{b}[i +1])}{d\mathcal{A}(\mathbf{E}(\textbf{b}))[i]}\mathbf{E}(\textbf{b})[j]\frac{ \partial\mathrm{softmax}\left(\mathbf{E}(\textbf{b})^{T}A\mathbf{E}(\textbf{b} )\right)[j,i]}{\partial A}$ $\displaystyle=$ $\displaystyle(-1)^{\textbf{b}[i+1]}\sum_{j=1}^{i}\sum_{r=1}^{2m}h_{r}\bm{1} \left(\nu_{r,i,\textbf{b}[i]}>0\right)\langle W^{(1)}_{r},\mathbf{E}(\textbf{b })[j]\rangle\frac{\partial\mathrm{softmax}\left(\mathbf{E}(\textbf{b})^{T}A \mathbf{E}(\textbf{b})\right)[j,i]}{\partial A}$ $\displaystyle=$ $\displaystyle(-1)^{\textbf{b}[i+1]+1}\sum_{j=1}^{i}\langle\kappa_{1,i,\textbf{ b}[i]},e_{j,\textbf{b}[j]}\rangle\frac{\partial\mathrm{softmax}\left(\mathbf{E }(\textbf{b})^{T}A\mathbf{E}(\textbf{b})\right)[j,i]}{\partial A}$ $$

We will use $Z$ to denote $\mathbf{E}(\textbf{b})^{T}A\mathbf{E}(\textbf{b})$, and calculate the derivative of th+e softmax function applied to $Z$ with respect to $A$.

$$ $\displaystyle\frac{\partial\mathrm{softmax}(Z)[j,i]}{\partial A}$ $\displaystyle=\frac{\partial}{\partial A}\left(\frac{e^{Z[j,i]}}{\sum_{p=1}^{i }e^{Z[p,i]}}\right)$ $\displaystyle=\frac{e^{Z[j,i]}\cdot\frac{\partial Z[j,i]}{\partial A}\cdot\sum _{p=1}^{i}e^{Z[p,i]}-e^{Z[j,i]}\cdot\sum_{p=1}^{i}e^{Z[p,i]}\cdot\frac{ \partial Z[p,i]}{\partial A}}{\left(\sum_{p=1}^{i}e^{Z[p,i]}\right)^{2}}$ $$

Given $Z[p,i]=\mathbf{E}(\textbf{b})[p]^{T}A\mathbf{E}(\textbf{b})[i]$, the derivative with respect to $A$ is:

$$ $\frac{\partial Z[p,i]}{\partial A}=\mathbf{E}(\textbf{b})[p]\cdot\mathbf{E}( \textbf{b})[i]^{T}$ $$

Substituting the derivatives into the quotient rule expression:

$$ $\frac{\partial\mathrm{softmax}(Z)[j,i]}{\partial A}=\frac{e^{Z[j,i]}\cdot \mathbf{E}(\textbf{b})[j]\mathbf{E}(\textbf{b})[i]^{T}\cdot\sum_{p=1}^{i}e^{Z[ p,i]}-e^{Z[j,i]}\cdot\sum_{p=1}^{i}e^{Z[p,i]}\cdot\mathbf{E}(\textbf{b})[p] \mathbf{E}(\textbf{b})[i]^{T}}{\left(\sum_{p=1}^{i}e^{Z[p,i]}\right)^{2}}$ $$

Thus, the derivative is:

$$ $\frac{\partial\mathrm{softmax}\left(\mathbf{E}(\textbf{b})^{T}A\mathbf{E}( \textbf{b})\right)[j,i]}{\partial A}=\mathrm{softmax}(Z)[j,i]\cdot\left( \mathbf{E}(\textbf{b})[j]-\sum_{p=1}^{i}\mathrm{softmax}(Z)[p,i]\mathbf{E}( \textbf{b})[p]\right)\mathbf{E}(\textbf{b})[i]^{T}$ $$

The proof is then complete.
∎

###### Lemma 21 .

After two training steps, the attention layer will show the following structure.
With probability $1-0.7\delta$, for all $i,b,i^{\prime},b_{2}$,

- 1.
If $j^{\prime}=S[i^{\prime}]$, then $\left\langle e_{j^{\prime},b_{1}^{\prime}},A^{(2)}e_{i^{\prime},b_{2}^{\prime}
}\right\rangle\geq\frac{\eta_{0}\eta_{1}}{256mn^{2}}$.
- 2.
If $j^{\prime}\neq S[i^{\prime}]$, then $|\left\langle e_{j^{\prime},b_{1}^{\prime}},A^{(2)}e_{i^{\prime},b_{2}^{\prime
}}\right\rangle|\leq\frac{\eta_{0}\eta_{1}}{512mn^{2}}$.

Further, events in Lemmas 18 and 17 hold.

###### Proof.

By Lemma 20,

$$ $\displaystyle\frac{\partial\mathrm{softmax}\left(\mathbf{E}(\textbf{b})^{T}A \mathbf{E}(\textbf{b})\right)[j,i]}{\partial A}=\frac{1}{i}\left(\mathbf{E}( \textbf{b})[j]-\sum_{p=1}^{i}\frac{1}{i}\mathbf{E}(\textbf{b})[p]\right) \mathbf{E}(\textbf{b})[i]^{T}$ $$

We can rewrite the gradient as

$$ $\displaystyle\frac{d\ell(\mathcal{T}(\textbf{b})[i],\textbf{b}[i+1])}{dA}=$ $\displaystyle-\frac{1}{i}(-1)^{\textbf{b}[i+1]}\left(\sum_{j=1}^{i}\langle \kappa_{i,\textbf{b}[i]},\mathbf{E}(\textbf{b})[j]\rangle\mathbf{E}(\textbf{b} )[j]\right)\mathbf{E}(\textbf{b})[i]^{T}$ $\displaystyle-\frac{1}{i^{2}}\sum_{j=1}^{i}\sum_{p=1}^{i}(-1)^{\textbf{b}[i+1] }\langle\kappa_{i,\textbf{b}[i]},\mathbf{E}(\textbf{b})[j]\rangle\mathbf{E}( \textbf{b})[p]\mathbf{E}(\textbf{b})[i]^{T}$ $$

If we define

$$ $\displaystyle\mu_{t,j,i,p,b_{1},b_{2},b_{3}}=\left(\sum_{s=1}^{B}(-1)^{\textbf {b}^{(t,s)}[i+1]}\bm{1}(\textbf{b}^{(t,s)}[j]=b_{1},\textbf{b}^{(t,s)}[i]=b_{2 },\textbf{b}^{(t,s)}[p]=b_{3})\right).$ $$

Summing over the second batch, we have that

$$ $\displaystyle-\frac{dL^{(1)}}{dA}=$ $\displaystyle\frac{1}{B}\sum_{i=1}^{n}\sum_{j=1}^{i}\sum_{b_{1}=0}^{1}\sum_{b_ {2}=0}^{1}\frac{1}{i}\alpha_{1,i,j,b_{1},b_{2}}\Delta_{1,i,j,b_{1},b_{2}}e_{j, b_{1}}e_{i,b_{2}}^{T}$ $\displaystyle-\frac{1}{B}\sum_{i=1}^{n}\sum_{j=1}^{i}\sum_{p=1}^{i}\sum_{b_{1} =0}^{1}\sum_{b_{2}=1}^{1}\sum_{b_{3}=0}^{1}\frac{1}{i^{2}}\mu_{1,j,i,m,b_{1},b _{2},b_{3}}\Delta_{1,i,j,b_{1},b_{2}}e_{k,b_{3}}e_{i,b_{2}}^{T}.$ $$

This implies $A^{(2)}$ is updated as,

$$ $\displaystyle A^{(2)}$ $\displaystyle=-\eta_{1}\frac{dL^{(1)}}{dA}$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i=n+1}^{n+k}\sum_{j=1}^{i}\sum_{b_{1}=0} ^{1}\sum_{b_{2}=0}^{1}\frac{1}{i}\alpha_{1,i,j,b_{1},b_{2}}\Delta_{1,i,j,b_{1} ,b_{2}}e_{j,b_{1}}e_{i,b_{2}}^{T}$ $\displaystyle-\frac{\eta_{1}}{B}\sum_{i=n+1}^{n+k}\sum_{j=1}^{i}\sum_{p=1}^{i} \sum_{b_{1}=0}^{1}\sum_{b_{2}=1}^{1}\sum_{b_{3}=0}^{1}\frac{1}{i^{2}}\mu_{1,j, i,p,b_{1},b_{2},b_{3}}\Delta_{1,i,j,b_{1},b_{2}}e_{p,b_{3}}e_{i,b_{2}}^{T}.$ $$

Recall that our goal is to calculate $e_{j^{\prime},b_{1}^{\prime}}^{T}A^{(2)}e_{i^{\prime},b_{2}^{\prime}}$. We can then calculate the contribution from each term by separating the calculation as follows:

$$ $\displaystyle T_{j}$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i\in[\max\{n+1,j\},n+k],b_{2}\in\{0,1\}, S[i]\neq j}\sum_{b_{1}=0}^{1}\frac{1}{i}\alpha_{1,i,j,b_{1},b_{2}}\Delta_{1,i, j,b_{1},b_{2}}\langle e_{j^{\prime},b_{1}^{\prime}},e_{j,b_{1}}\rangle\langle e _{i,b_{2}},e_{i^{\prime},b_{2}^{\prime}}\rangle.$ $\displaystyle R$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2 }=0}^{1}\frac{1}{i}\alpha_{1,i,S[i],b_{1},b_{2}}\Delta_{1,i,S[i],b_{1},b_{2}} \langle e_{j^{\prime},b_{1}^{\prime}},e_{S[i],b_{1}}\rangle\langle e_{i,b_{2}} ,e_{i^{\prime},b_{2}^{\prime}}\rangle.$ $\displaystyle U_{j,p}$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i=\max\{n+1,j,p\},S[i]\not\in\{j,p\}}^{n +k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\sum_{b_{3}=0}^{1}\frac{1}{i^{2}}\mu_{1 ,j,i,p,b_{1},b_{2},b_{3}}\Delta_{1,i,j,b_{1},b_{2}}\langle e_{j^{\prime},b_{1} ^{\prime}},e_{j,b_{1}}\rangle\langle e_{p,b_{3}},e_{i^{\prime},b_{2}^{\prime}}\rangle$ $\displaystyle V_{p}$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i=\max\{n+1,p\}}^{n+k}\sum_{b_{1}=0}^{1} \sum_{b_{2}=0}^{1}\frac{1}{i^{2}}\mu_{1,S[i],i,p,b_{1},b_{2},b_{3}}\Delta_{1,i ,S[i],b_{1},b_{2}}\langle e_{j^{\prime},b_{1}^{\prime}},e_{S[i],b_{1}}\rangle \langle e_{p,b_{3}},e_{i^{\prime},b_{2}^{\prime}}\rangle$ $\displaystyle W_{j}$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i=\max\{n+1,j\}}^{n+k}\sum_{b_{1}=0}^{1} \sum_{b_{2}=0}^{1}\sum_{b_{3}=0}^{1}\frac{1}{i^{2}}\mu_{1,j,i,S[i],b_{1},b_{2} ,b_{3}}\Delta_{1,i,j,b_{1},b_{2}}\langle e_{j^{\prime},b_{1}^{\prime}},e_{j,b_ {1}}\rangle\langle e_{p,b_{3}},e_{i^{\prime},b_{2}^{\prime}}\rangle.$ $\displaystyle Y$ $\displaystyle=\frac{\eta_{1}}{B}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2 }=0}^{1}\sum_{b_{3}=0}^{1}\frac{1}{i^{2}}\mu_{1,S[i],i,S[i],b_{1},b_{2},b_{3}} \Delta_{1,i,S[i],b_{1},b_{2}}\langle e_{j^{\prime},b_{1}^{\prime}},e_{S[i],b_{ 1}}\rangle\langle e_{S[i],b_{3}},e_{i^{\prime},b_{2}^{\prime}}\rangle.$ $$

Then we have that

$$ $\displaystyle e_{j^{\prime},b_{1}^{\prime}}^{T}A^{(2)}e_{i^{\prime},b_{2}^{ \prime}}$ $\displaystyle=\sum_{j=1}^{n+k}T_{j}+R-\sum_{j=1}^{n+k}\sum_{p=1}^{n+k}U_{j,p}- \sum_{j=1}^{n+k}W_{j}-\sum_{p=1}^{n+k}V_{p}+Y.$ $$

.By Lemma 31, with probability $1-\delta/100$,

$$ $\displaystyle\left|\langle e_{i,b_{2}},e_{i^{\prime},b_{2}^{\prime}}\rangle- \bm{1}(i=i^{\prime}\&b_{2}=b_{2}^{\prime})\right|\leq 2\sqrt{\frac{\log(100nd/ \delta)}{d}}.$ $$

We will also assume the event in Lemma 19 holds.

We will now discuss each term separately

- 1.
For $T_{j}$, by Lemma 6 and Azuma-Hoeffding bound, with probability $1-\delta/50$,
$\displaystyle\left|\sum_{i\in[\max\{n+1,j\},n+k],b_{2}\in\{0,1\},S[i]\neq j}
\sum_{b_{2}=0}^{1}\frac{1}{i}\alpha_{1,i,j,b_{1},b_{2}}\Delta_{1,i,j,b_{1},b_{
2}}\langle e_{j^{\prime},b_{1}^{\prime}},e_{j,b_{1}}\rangle\langle e_{i,b_{2}}
,e_{i^{\prime},b_{2}^{\prime}}\rangle\right|$
$\displaystyle\leq$
$\displaystyle 2\sum_{b_{1}=0}^{1}\sqrt{B\log(\frac{100n}{\delta})\sum_{i\in[
\max\{n+1,j\},n+k],b_{2}\in\{0,1\},S[i]\neq j}\frac{1}{i^{2}}\Delta_{1,1,i,j,b
_{1},b_{2}}^{2}\left(\langle e_{j^{\prime},b_{1}^{\prime}},e_{j,b_{1}}\rangle
\right)^{2}\left(\langle e_{i,b_{2}},e_{i^{\prime},b_{2}^{\prime}}\rangle
\right)^{2}}$
$\displaystyle\leq$
$\displaystyle\frac{2\sqrt{B\log(\frac{100n}{\delta})}}{n}\sup_{j\neq S[i]}|
\Delta_{1,i,j,b_{1},b_{2}}|\sum_{b_{1}=0}^{1}\left|\langle e_{j^{\prime},b_{1}
^{\prime}},e_{j,b_{1}}\rangle\right|\sqrt{\sum_{i\in[\max\{n+1,j\},n+k],b_{2}
\in\{0,1\},S[i]\neq j}\left(\langle e_{i,b_{2}},e_{i^{\prime},b_{2}^{\prime}}
\rangle\right)^{2}}$
Now by Lemma 19, we have that
$\displaystyle\sup_{j\neq S[i]}|\Delta_{1,i,j,b_{1},b_{2}}|\leq\frac{10\eta_{0}
}{mn}\frac{\sqrt{nk}}{\sqrt{Bd}}\log(300mnB/\delta).$
Further by Lemma 31,
$\displaystyle\sum_{i\in[\max\{n+1,j\},n+k],b_{2}\in\{0,1\},S[i]\neq j}\left(
\langle e_{i,b_{2}},e_{i^{\prime},b_{2}^{\prime}}\rangle\right)^{2}\leq$
$\displaystyle\frac{8k\log(100nd/\delta)}{d}+1\leq 4.$
$\displaystyle\left|\langle e_{j^{\prime},b_{1}^{\prime}},e_{j,b_{1}}\rangle
\right|\leq$
$\displaystyle\bm{1}(j=j^{\prime}\&b_{1}^{\prime}=b_{1})+2\sqrt{\frac{\log(100
nd/\delta)}{d}}.$
$\displaystyle|T_{j}|\leq\frac{40\eta_{1}\eta_{0}}{mn^{2}}\frac{\sqrt{nk}}{B
\sqrt{d}}\log^{1.5}(300mnB/\delta)\left(4\sqrt{\frac{\log(100nd/\delta)}{d}}+
\bm{1}(j=j^{\prime})\right).$
Summing over $j$, by Lemma 5, we have that the contribution is bounded by
$\displaystyle\sum_{j}|T_{j}|$
$\displaystyle\leq\frac{160\eta_{1}\eta_{0}}{mn^{2}}\log^{2}(300mnB/\delta)
\frac{n\sqrt{nk}}{d\sqrt{Bd}}\leq\frac{\eta_{1}\eta_{0}}{2000mn^{2}}.$
(3)
- 2.
For $R$, we will directly put everything to an upper bound. We will discuss three cases,
•
$\not\exists i^{\prime\prime},S[i^{\prime\prime}]=j$, by Lemma 19,
$\displaystyle|R|$
$\displaystyle\leq\frac{16k\eta_{1}\log(100nd/\delta)}{Bnd}\sup\alpha\Delta+
\frac{\eta_{1}}{B}\frac{1}{i^{\prime}}\alpha_{1,i^{\prime},S[i^{\prime}],b_{1}
,b_{2}^{\prime}}\Delta_{1,i^{\prime},S[i^{\prime}],b_{1},b_{2}^{\prime}}|
\langle e_{j^{\prime},b_{1}^{\prime}},e_{S[i],b_{1}}\rangle|$
$\displaystyle\leq\frac{16k\eta_{1}\log(100nd/\delta)}{Bnd}\frac{B}{2}\frac{
\eta_{0}}{4mn}+\frac{\eta_{1}}{Bn}B\frac{\eta_{0}}{4mn}\sqrt{\frac{\log(100nd/
\delta)}{d}}$
$\displaystyle=\frac{\eta_{0}\eta_{1}k\log(100nd/\delta)}{mn^{2}d}+\frac{\eta_{
0}\eta_{1}}{4mn^{2}}\sqrt{\frac{\log(100nd/\delta)}{d}}\leq\frac{\eta_{0}\eta_
{1}}{2000mn^{2}}.$
•
$\exists i^{\prime\prime},S[i^{\prime\prime}]=j,i^{\prime\prime}\neq i^{\prime}$, we can get a similar bound,
$\displaystyle|R|$
$\displaystyle\leq\frac{\eta_{0}\eta_{1}}{2000mn^{2}}.$
•
$S[i^{\prime}]=j^{\prime}$, in this case, we can show that,
$\displaystyle|R-\frac{\eta_{0}\eta_{1}}{8B^{2}mi^{\prime}}\bar{\alpha}_{i,S[i^
{\prime}],b_{1},b_{2}}^{2}|\leq\frac{\eta_{0}\eta_{1}}{2000mn^{2}}.$
This concludes that
$\displaystyle\left|R-\frac{\eta_{0}\eta_{1}}{8B^{2}mi^{\prime}}\bar{\alpha}_{i
,S[i^{\prime}],b_{1},b_{2}}^{2}\bm{1}(S[i^{\prime}]=j^{\prime})\right|\leq
\frac{\eta_{0}\eta_{1}}{2000mn^{2}}.$
(4)
- 3.
$U_{j,p}$ can ve bounded in the same way as $T_{j}$ and we have that,
$\displaystyle\sum_{j}\sum_{p}U_{j,p}\leq\frac{\eta_{1}\eta_{0}}{1000mn^{2}}.$
(5)
- 4.
For $V_{p},W_{j},Y$, we will directly put everything to an upper bound similar to the bound of $R$, we have that
$\displaystyle\sum_{p=1}^{n+k}|V_{p}|$
$\displaystyle\leq\frac{\eta_{1}\eta_{0}}{1000mn^{2}}.$
(6)
$\displaystyle\sum_{j=1}^{n+k}|W_{j}|$
$\displaystyle\leq\frac{\eta_{1}\eta_{0}}{1000mn^{2}}$
(7)
$\displaystyle|Y|$
$\displaystyle\leq\frac{\eta_{1}\eta_{0}}{1000mn^{2}}.$
(8)

Summing over Equations 3, 4, 5, 6, 7 and 8, we have that

$$ $\displaystyle\left|\left\langle e_{j^{\prime},b_{1}^{\prime}},A^{(2)}e_{i^{ \prime},b_{2}^{\prime}}\right\rangle-\frac{\eta_{0}\eta_{1}}{8Bmi^{\prime}} \bar{\alpha}_{i,S[i^{\prime}],b_{1},b_{2}}^{2}\bm{1}(S[i^{\prime}]=j^{\prime}) \right|\leq\frac{\eta_{1}\eta_{0}}{512mn^{2}}.$ $$

Further

$$ $\displaystyle\frac{\eta_{1}}{8B^{2}mi^{\prime}}\bar{\alpha}_{i,S[i^{\prime}],b _{1},b_{2}}^{2}\geq\frac{\eta_{0}\eta_{1}}{128mn^{2}}.$ $$

Hence, we can show that

- 1.
If $j^{\prime}=S[i^{\prime}]$, then $\left\langle e_{j^{\prime},b_{1}^{\prime}},A^{(2)}e_{i^{\prime},b_{2}^{\prime}
}\right\rangle\geq\frac{\eta_{0}\eta_{1}}{256mn^{2}}$.
- 2.
If $j^{\prime}\neq S[i^{\prime}]$, then $|\left\langle e_{j^{\prime},b_{1}^{\prime}},A^{(2)}e_{i^{\prime},b_{2}^{\prime
}}\right|\rangle\leq\frac{\eta_{0}\eta_{1}}{512mn^{2}}$.

The proof is complete.
∎

###### Lemma 22 .

Under the setting of Lemma 21, the attention output is approximately one-hot after the second step, with

$$ $\displaystyle\left|\mathcal{A}[A^{2}](\mathbf{E}(\textbf{b}))[i]-e_{S[i], \textbf{b}[S[i]]}\right|<1/n^{9}.$ $$

###### Proof.

This is a direct consequence of Lemma 21 with $\frac{\eta_{0}\eta_{1}}{256mn^{2}}>10\log n$ (Lemma 5).
∎

###### Lemma 16 .

###### Proof.

###### Lemma 17 .

###### Proof.

###### Lemma 18 .

###### Proof.

###### Lemma 19 .

###### Proof.

###### Lemma 20 .

###### Proof.

###### Lemma 21 .

###### Proof.

###### Lemma 22 .

###### Proof.

#### A.4.4 Third Step: Moving MLP in Linear Regime

###### Lemma 23 .

With probability $1-0.8\delta$, the gradient of the FFN layer on the third step can be written as,

$$ $\displaystyle\left|\left\langle\frac{d\mathcal{L}^{(2)}}{dW_{r,1:d}}\right \rangle\right|\leq$ $\displaystyle\frac{\sqrt{\log(300nB/\delta)}}{m\sqrt{B}}.$ $\displaystyle\frac{d\mathcal{L}^{(2)}}{dW_{r,d+1:2d}}=$ $\displaystyle h_{r}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\bm{1 }(\nu_{r,i,b_{2}}>0)(-1)^{b_{1}+b_{2}}e_{S[i],b_{1}}+O(\frac{1}{n^{8}}).$ $$

Further the event in Lemmas 21 and 22 hold.

###### Proof.

The proof is similar to Lemmas 14 and 12, switching the original output with the near one-hot output of the attention layer.
∎

###### Lemma 24 .

With probability $1-0.9\delta$, the attention output is almost unchanged after the third step, with

$$ $\displaystyle\left|\mathcal{A}[A^{3}](\mathbf{E}(\textbf{b}))[i]-e_{S[i], \textbf{b}[S[i]]}\right|<1/n^{9}.$ $$

Further, events in Lemma 23 hold.

###### Proof.

By Lemmas 20, 22 and 5, $A^{(3)}-A^{(2)}$ is of order $1/n^{8}$, this implies that the attention weight and output is almost unchanged.
∎

###### Lemma 25 .

With probability $1-0.91\delta$,

$$ $\displaystyle\left|\left\langle W_{r,d+1:2d}^{(2)},\mathcal{A}(A^{3})[\mathbf{ E}(\textbf{b})][i]\right\rangle+\mathrm{sign}(\nu_{r,i,\textbf{b}[S[i]]}h_{r}) \frac{2\epsilon}{3}\bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0)\right|<\frac{\epsilon}{300}.$ $$

Further, events in Lemma 24 hold.

###### Proof.

We will first calculate the projection of the gradient on $e_{S[i^{\prime}],b^{\prime}}$. When the event in Lemma 23 happens,

$$ $\displaystyle\left\langle\frac{d\mathcal{L}^{(2)}}{dW_{r,d+1:2d}},e_{S[i^{ \prime}],b^{\prime}}\right\rangle$ $\displaystyle=h_{r}\sum_{i=n+1}^{n+k}\sum_{b_{1}=0}^{1}\sum_{b_{2}=0}^{1}\bm{1 }(\nu_{r,i,b_{2}}>0)(-1)^{b_{1}+b_{2}}\langle e_{S[i^{\prime}],b^{\prime}},e_{ S[i],b_{1}}\rangle+O(\frac{1}{n^{8}}).$ $$

By Lemmas 30 and 5, with probability $1-0.01\delta$,

$$ $\displaystyle\left|\left\langle\frac{d\mathcal{L}^{(2)}}{dW_{r,d+1:2d}},e_{S[i ^{\prime}],b^{\prime}}\right\rangle-\mathrm{sign}(\nu_{r,i,b^{\prime}})h_{r} \bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0)\right|\leq\frac{1}{200m}$ $$

By Lemmas 17 and 5, we have that,

$$ $\displaystyle\left|\left\langle W_{r,d+1:2d}^{(2)},e_{S[i^{\prime}],b^{\prime} }\right\rangle+\mathrm{sign}(\nu_{r,i,b^{\prime}})h_{r}\eta_{2}\bm{1}(\nu_{r,i ,0}\nu_{r,i,1}<0)\right|<\frac{(\eta_{0}+\eta_{1})}{20mn}+\frac{\eta_{2}}{200m }<\frac{\eta_{2}}{100m}.$ $$

Combining with Lemma 24 and $\eta_{2}=\frac{4\epsilon m}{3}$, we have the result.
∎

###### Lemma 26 .

With probability $1-\delta$, for all $\textbf{b}\in\{0,1\}^{n+k}$, $i\in[n+1,n+k]$, we have that

$$ $\displaystyle\left|\mathcal{T}(\mathbf{E}(\textbf{b}))[i]-\frac{\epsilon(-1)^{ \textbf{b}[i+1]+1}}{3}\right|<\frac{\epsilon}{3}.$ $$

###### Proof.

By Lemma 23, the following holds,

$$ $\displaystyle\left\langle W_{r,1:d}^{(3)},\mathbf{E}(\textbf{b})[i]\right \rangle\mathrm{sign}(\nu_{r,i,\textbf{b}[i]})\in[5\epsilon/6,7\epsilon/6].$ $$

Combining with Lemma 25,

$$ $\displaystyle\left|\left\langle W_{r,d+1:2d}^{(2)},\mathcal{A}(A^{3})[\mathbf{ E}(\textbf{b})][i]\right\rangle+\mathrm{sign}(\nu_{r,i,\textbf{b}[S[i]]}h_{r}) \frac{2\epsilon}{3}\bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0)\right|<\frac{\epsilon}{300}.$ $$

Hence, we still have that

$$ $\displaystyle|\left\langle W_{r,d+1:2d}^{(3)},\mathcal{A}(A^{3})[\mathbf{E}( \textbf{b})][i]\right\rangle|<\frac{5\epsilon}{6}.$ $$

which implies that,

$$ $\displaystyle\bm{1}(\left\langle W_{r,1:d}^{(3)},\mathbf{E}(\textbf{b})[i] \right\rangle>0)=\bm{1}(\nu_{r,i,\textbf{b}[i]}>0).$ $$

Consider the output contribution of the attention part,

$$ $\displaystyle\Big{|}$ $\displaystyle\sum_{r=1}^{2m}h_{r}\left\langle W_{r,d+1:2d}^{(3)},\mathcal{A}(A ^{3})[\mathbf{E}(\textbf{b})][i]\right\rangle\bm{1}(\nu_{r,i,\textbf{b}[i]}>0)$ $\displaystyle+\sum_{r=1}^{2m}\frac{2h_{r}\epsilon}{3}\mathrm{sign}(\nu_{r,i, \textbf{b}[S[i]]}h_{r})\bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0,\nu_{r,i,\textbf{b}[i]} >0)\Big{|}<\frac{\epsilon}{300}.$ $$

The later term with $1-\delta/100$ satisfies,

$$ $\displaystyle\sum_{r=1}^{2m}\frac{2h_{r}\epsilon}{3}\mathrm{sign}(\nu_{r,i, \textbf{b}[S[i]]}h_{r})\bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0,\nu_{r,i,\textbf{b}[i]} >0)$ $\displaystyle=$ $\displaystyle\sum_{r=1}^{2m}\frac{\epsilon}{3m}\mathrm{sign}(\nu_{r,i,\textbf{ b}[S[i]]})\bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0,\nu_{r,i,\textbf{b}[i]}>0)$ $\displaystyle=$ $\displaystyle\sum_{r=1}^{2m}\frac{\epsilon}{3m}(-1)^{\textbf{b}[S[i]]+\textbf{ b}[i]}\bm{1}(\nu_{r,i,0}\nu_{r,i,1}<0)$ $\displaystyle=$ $\displaystyle\frac{\epsilon}{3}(-1)^{\textbf{b}[S[i]]+\textbf{b}[i]}+O(\frac{ \epsilon\log(100n/\delta)}{\sqrt{m}})$ $$

This shows that,

$$ $\displaystyle\left|\sum_{r=1}^{2m}h_{r}\left\langle W_{r,d+1:2d}^{(3)}, \mathcal{A}(A^{3})[\mathbf{E}(\textbf{b})][i]\right\rangle\bm{1}(\nu_{r,i, \textbf{b}[i]}>0)-\frac{\epsilon}{3}(-1)^{1+\textbf{b}[S[i]]+\textbf{b}[i]} \right|\leq\frac{\epsilon}{150}.$ $$

On the other hand, we have that

$$ $\displaystyle\left|\sum_{r=1}^{2m}h_{r}\left\langle W_{r,1:d}^{(3)},\mathbf{E} (\textbf{b})[i]\right\rangle\bm{1}(\nu_{r,i,\textbf{b}[i]}>0)\right|$ $\displaystyle\leq$ $\displaystyle\left|\sum_{r=1}^{2m}h_{r}\left\langle W_{r,1:d}^{(0)},\mathbf{E} (\textbf{b})[i]\right\rangle\bm{1}(\nu_{r,i,\textbf{b}[i]}>0)\right|+\left| \sum_{r=1}^{2m}h_{r}\left\langle W_{r,1:d}^{(3)}-W_{r,1:d}^{(0)},\mathbf{E}( \textbf{b})[i]\right\rangle\bm{1}(\nu_{r,i,\textbf{b}[i]}>0)\right|$ $\displaystyle\leq$ $\displaystyle\frac{\epsilon}{4}.$ $$

The first term is bounded due to standard concentration inequality over the axis of $r$. The second term is bounded by Lemmas 17 and 23. Combining the terms, we have that

$$ $\displaystyle\left|\mathcal{T}(\mathbf{E}(\textbf{b}))[i]-\frac{\epsilon(-1)^{ \textbf{b}[i]+\textbf{b}[S[i]]+1}}{3}\right|=\left|\mathcal{T}(\mathbf{E}( \textbf{b}))[i]-\frac{\epsilon(-1)^{\textbf{b}[i+1]+1}}{3}\right|<\frac{ \epsilon}{3}.$ $$

This concludes that the model is able to predict the correct output.
∎

###### Lemma 23 .

###### Proof.

###### Lemma 24 .

###### Proof.

###### Lemma 25 .

###### Proof.

###### Lemma 26 .

###### Proof.

### A.5 Extension to Hinge Loss

###### Lemma 27 .

The results in Lemmas 26 and 18 can be extended to hinge loss $\ell(\hat{y},y)=\max\{(-1)^{y}\hat{y}+1,0\}$ with the same probability.

###### Proof.

Our Lemmas 17 and 23 shows that the output of the FFN layer is bounded by $\epsilon/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.
∎

###### Lemma 27 .

###### Proof.

### A.6 Technical Lemma

###### Lemma 28 (Theorem 4.4.5 of [ 55 ] ) .

There exists universal constant $C$,let $A$ be an $m\times n$ random matrix whose entries $A_{ij}$ are independent, mean zero, sub-Gaussian random variables. Then, for any $t>0$, we have

$$ $\|A\|\leq CK\left(\sqrt{m}+\sqrt{n}+t\right)$ $$

with probability at least $1-2\exp(-t^{2})$. Here $K=\max_{i,j}\|A_{ij}\|_{\psi_{2}}$ is the maximum sub-Gaussian norm of $A_{ij}$.

###### Lemma 29 (Theorem 4.6.1 of [ 55 ] ) .

Let $A$ be an $m\times n$ matrix whose rows $A_{i}$ are independent, mean zero, sub-Gaussian isotropic random vectors in $\mathbb{R}^{n}$. Then, for any $t\geq 0$, we have

$$ $\sqrt{m}-CK^{2}(\sqrt{n}+t)\leq s_{n}(A)\leq s_{1}(A)\leq\sqrt{m}+CK^{2}(\sqrt {n}+t)$ $$

with probability at least $1-2\exp(-t^{2})$. Here, $K=\max_{i}\|A_{i}\|_{\psi_{2}}$.

Furthermore, a slightly stronger conclusion holds:

$$ $\left\|\frac{1}{m}A^{\top}A-I_{n}\right\|\leq K^{2}\max(\delta,\delta^{2}),$ $$

where $\delta=C\left(\sqrt{\frac{n}{m}}+\frac{t}{\sqrt{m}}\right).$

###### Lemma 30 .

Define $v=\sum_{i=1}^{n+k}\sum_{b=0}^{1}\lambda_{i,b}e_{i,b}$, then with probability $1-\delta$,

$$ $\left|\langle e_{i^{\prime},b^{\prime}},v\rangle-\lambda_{i^{\prime},b^{\prime }}\right|\leq\sqrt{\frac{2\log\left(\tfrac{2}{\delta}\right)}{d}\sum_{(i,b) \neq(i^{\prime},b^{\prime})}\lambda_{i,b}^{2}}.$ $$

###### Proof.

We aim to bound $|\langle e_{i^{\prime},b^{\prime}},v\rangle|$, where

$$ $v=\sum_{i=n+1}^{n+k}\sum_{b=0}^{1}\lambda_{i,b}e_{i,b}.$ $$

Note that each $e_{i,b}\in\mathbb{R}^{d}$ has entries that are independent random variables from $\left\{-\tfrac{1}{\sqrt{d}},\tfrac{1}{\sqrt{d}}\right\}$.

First, observe that

$$ $\langle e_{i^{\prime},b^{\prime}},v\rangle=\sum_{(i,b)}\lambda_{i,b}\langle e_ {i^{\prime},b^{\prime}},e_{i,b}\rangle=\lambda_{i^{\prime},b^{\prime}}+\sum_{( i,b)\neq(i^{\prime},b^{\prime})}\lambda_{i,b}\langle e_{i^{\prime},b^{\prime}} ,e_{i,b}\rangle.$ $$

Since $\langle e_{i^{\prime},b^{\prime}},e_{i^{\prime},b^{\prime}}\rangle=1$ and for $(i,b)\neq(i^{\prime},b^{\prime})$, the inner products $\langle e_{i^{\prime},b^{\prime}},e_{i,b}\rangle$ are sums of independent random variables with mean zero.

Define

$$ $X_{i,b}=\lambda_{i,b}\langle e_{i^{\prime},b^{\prime}},e_{i,b}\rangle,\quad \text{for }(i,b)\neq(i^{\prime},b^{\prime}).$ $$

Each $X_{i,b}$ is a sum of $d$ independent random variables bounded in $\left[-\tfrac{|\lambda_{i,b}|}{d},\tfrac{|\lambda_{i,b}|}{d}\right]$, because each component $e_{i^{\prime},b^{\prime},j}e_{i,b,j}$ is $\pm\tfrac{1}{d}$.

By the Azuma-Hoeffding inequality, for any $t>0$,

$$ $\Pr\left(\left|\sum_{(i,b)\neq(i^{\prime},b^{\prime})}X_{i,b}\right|\geq t \right)\leq 2\exp\left(-\frac{2dt^{2}}{\sum_{(i,b)\neq(i^{\prime},b^{\prime})} 4\lambda_{i,b}^{2}}\right)=2\exp\left(-\frac{dt^{2}}{2\sum_{(i,b)\neq(i^{ \prime},b^{\prime})}\lambda_{i,b}^{2}}\right).$ $$

Setting the right-hand side equal to $\delta$ and solving for $t$, we get

$$ $t=\sqrt{\frac{2\log\left(\tfrac{2}{\delta}\right)}{d}\sum_{(i,b)\neq(i^{\prime },b^{\prime})}\lambda_{i,b}^{2}}.$ $$

Therefore, with probability at least $1-\delta$,

$$ $\left|\langle e_{i^{\prime},b^{\prime}},v\rangle-\lambda_{i^{\prime},b^{\prime }}\right|\leq\sqrt{\frac{2\log\left(\tfrac{2}{\delta}\right)}{d}\sum_{(i,b) \neq(i^{\prime},b^{\prime})}\lambda_{i,b}^{2}}.$ $$

The proof is then complete.
∎

###### Lemma 31 .

With probability $1-\delta$, it holds that

$$ $\displaystyle\forall i,i^{\prime}\in[n+k],b,b^{\prime}\in\{0,1\},\|e_{i,b}^{T} e_{i^{\prime},b^{\prime}}-\bm{1}(i=i^{\prime}\&b=b^{\prime})\|_{2}\leq\sqrt{2 \frac{\log(8nd/\delta)}{d}}.$ $$

###### Proof.

This is a direct consequence combining Lemma 30 and union bound.
∎

###### Lemma 28 (Theorem 4.4.5 of [ 55 ] ) .

###### Lemma 29 (Theorem 4.6.1 of [ 55 ] ) .

###### Lemma 30 .

###### Proof.

###### Lemma 31 .

###### Proof.

## Appendix B Additional Experiment Results

In this section, we provide details of the experiment setup and present additional results. All training was conducted using PyTorch [] on NVIDIA RTX A10 GPUs.

### B.1 Experiment Details

The transformer architecture adopted in the experiment section is based on the GPT-2 model [] with a hidden size of $720$, an intermediate size of $3072$, and trainable position embeddings. For all experiments, we use Adam [] optimizer with random initialization, using hyperparameters $\beta_{1}=0.9,\beta_{2}=0.999$, a weight decay of 0 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.

##### Details of the experiments in Section 4.3

In this section, we analyze the normalized attention entropy of Qwen2-7B [] and Qwen2-Math-7B models [] on the GSM8K dataset [] with and without CoT prompting respectively. Both models consist of $28$ layers, each with $28$ attention heads. The normalized attention entropy is computed as the average across the GSM8K test set. 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.

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].”.

Figure: Figure 7: 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 (Section 4.3). Each bar represents the entropy of an attention head.
Refer to caption: x7.png

### B.2 Additional Results

In Figure 3 (Right), we present the attention pattern of a single-layer, single-head transformer trained on the $(n=20,k=40)$ parity problem with CoT data. In Figure 8, 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 an 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.

Figure: Figure 8: The attention pattern learned by a $2$-layer $2$-head transformer on $(n=20,k=6)$ parity problem with CoT.
Refer to caption: x8.png

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 (Figure 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 4, this problem is harder for the models to learn as the number of secret variables $k$ is larger. We examine various configurations: the number of layers and heads ranges from $1,2,3,4,6$, and $8$; the learning rate varies from $6\times 10^{-6}$, $8\times 10^{-8}$, to $1\times 10^{-4}$; and the training dataset size varies from $10^{4}$, $10^{5}$ to $10^{6}$, with corresponding epochs ranging from $1000$, $100$ to $10$. Across all configurations we examined, training without CoT fails to achieve non-trivial accuracy (Figure 9). 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.

Figure: Figure 9: Training with CoT on $(n=20,k=12)$ parity problem fails to achieve non-trivial accuracy under different configurations.
Refer to caption: x9.png

In Figure 5, a $4$-layer $4$-head transformer achieves perfect evaluation accuracy on $(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 a significant decrease in attention entropy, indicating the development of sparse attention, while entropy remains high when trained on a larger dataset. In Figure 10, we present more results across different architectures ($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.

Figure: Figure 10: Transformer trained on the $(n=20,k=6)$ parity problem without CoT with different sizes of training dataset and a fixed number of iterations.
Refer to caption: x10.png

Figure: (a) $2$-layer $3$-head transformer
Refer to caption: x11.png

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
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?

* 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.

* 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?

* 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?

minor typos: L247, "for any i" should be for any i, j and b vector?
soundness: 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.
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.
weaknesses: A couple of minor weaknesses (see also questions below):

* the lower bound is claimed to be about sample complexity, though it feels more like a computational lower bound. Perhaps this should be clarified.

* the initialization seems a bit non-standard, particularly W_{r,1:d}. Some more justification would be helpful.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: Please see the above concerns.
soundness: 3
strengths: 1. The motivation is good: understanding CoT is a significant direction to understand pre-training and finetuning LLMs, especially for reasoning tasks.
2. 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.
3. Both theory and experiments clearly provide a exponential gap between the sample complexity with and without CoT.
4. The observation of sparsity in Qwen models on GSM8K is interesting, especially the math model finetuned with CoT.
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.
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.
2. It would be better to include discussion regarding a recent related work [1].

[1] How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad. Abbe et al. June 2024.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: 1. Can your analysis be extended to few-shot CoT?
2. 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.
3. 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?
soundness: 2
strengths: The paper is well-written and easy to follow.
The experiments are overall solid to support the theory.
It 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.
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.
weaknesses: There are two major issues. 
1. 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.
2. 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.

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 3
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.
soundness: 3
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.
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.
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.

Besides, 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.

All in all, the rating reflects the fact that the ideas are interesting, but I believe polishing the paper would make it much better.

Examples 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):
- 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.
- line 296: “which only applies to one-pass training”, but Theorem 2 applies to multipass training?
- line 831: why $m>100$?
- 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.
- line 844: why $\delta/4$?
- lines 896: why $8$ instead of $10$?
- line 901: the $\mathcal{O}(1/n^7)$ needs to be explained.
- line 913: “by choosing $a_r$ and $b_r$ iteratively”: this needs to be explained.
- line 916: “the parameters can be encoded in $\mathcal{O}(\log n)$ precision”. Why?
- line 960: this is not a valid proof of the Lemma, the derivation of the inequalities is non-trivial from the Assumption.

Minor remarks (no impact on score):
- 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.
- line 77: $k$ is not introduced yet.
- line 191: the bounds of the sum should depend whether it is a CoT or non-CoT task.
- line 230: it should be $W_{r,1:m}$ instead of $W_{r,1:d}$.

Meta-review: This paper investigates how Chain-of-Thought (CoT) enhances the sample efficiency of Transformer models in reasoning tasks, focusing on the parity learning problem. The authors theoretically demonstrate an exponential gap in sample complexity between learning with and without CoT for the parity task. They show that CoT induces sparse sequential dependencies and leads to sparse, interpretable attention patterns. These theoretical findings are validated through experiments on both synthetic and real-world datasets.

The paper's primary strength lies in its novel theoretical analysis of Transformer learning dynamics, particularly challenging due to the architecture's non-linear nature. The rigorous proof of the exponential sample complexity gap between CoT and non-CoT learning for parity tasks is a significant contribution. The work also provides valuable insights into the connection between CoT and the induction of sparse attention patterns, supported by empirical validation on both synthetic and real-world datasets.

While the theoretical analysis is robust, some reviewers initially found parts of the proofs difficult to follow, though this was largely addressed in the author response. The theoretical setup uses some simplifying assumptions, such as a fixed secret for the parity task, which may limit generalizability. Additionally, while supportive, the real-world experimental results could be stronger to fully validate the theoretical claims.

Despite these minor weaknesses, the paper makes a significant theoretical contribution to understanding the mechanics of CoT in Transformer models. It provides a rigorous foundation for explaining the empirically observed benefits of CoT in LLMs, combining theoretical analysis with empirical validation. The authors have been responsive to reviewer concerns, improving the clarity and presentation of their work. These factors strongly support the decision to accept the paper.

Decision: Accept (Poster)