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

We thank the reviewer for acknowledging this paper as having a nice and original idea. We especially appreciate the reviewer’s detailed suggestions. Following the reviewer’s suggestion, we have polished our writing and proofs.

Before delving into the detailed edits, we would like to point out that our proof of Theorem 3 is technically challenging to simplify because we are analyzing finite-batch dynamics with random embedding. While our proof for Theorem 3 can be greatly simplified without these two constraints, both constraints are necessary for us to derive a sample complexity upper bound under the regime of Theorem 2.

To improve the readability of our proof of Theorem 3, we now format the proof to first define and bound necessary auxiliary statistics in Section A.4.1. We also include a table listing out the definitions and orders of these auxiliary statistics. We then analyze the trajectory of weight step by step in Section A.4.2 to A.4.4, where each section corresponds to a different step. Finally, we conclude our proof in Section A.4.5 and prove our technical lemmas in Section A.4.5.

We would like to address the reviewers’ concerns below. We will start by addressing high-level concerns and list detailed edits to the end. We are happy to make further edits.

W1.1: Unclear Role of Initialization Scheme

The initialization we used is mainly used to guarantee the activation pattern of the MLP (i.e. whether the $r$-th neuron is activated by the ReLU activation) is solely determined by the position $i$, the boolean value $\\mathbf b$ i $$ and the randomness of v\_{r,i,\\mathbf b i } since \\langle e\_{i,\\mathbf b i },W\_{r,1:d}\\rangle \\approx v\_{r,i,b} with high probability. As a result, for a fixed $(i,\\mathbf b$ i $)$, the FFN can be viewed as a linear function, which consists of the neurons with v\_{r,i,\\mathbf b i }\>0. At phase 1, this linear function will develop a stronger correlation with embedding e\_{(S i $,\\mathbf b$ j )}. We have updated our manuscript to state this clearly in the main text.

W1.2: Confusing details of Theorem 1

In our theoretical analysis, we assume that $k \\in$ n / \log^5 (n / \delta), n / \log^4(n/\delta) $$. Hence the construction in Theorem 1 will fall in the scope of Theorem 2. We thank the reviewer for pointing out this missing assumption and have added it in the updated version.

W2: The real-world experiments are not convincing enough.

For the real-world experiments, we would like to clarify and highlight the following points in response: As shown in Figure 7 in the appendix, the attention entropy patterns across all layers of the Transformer model consistently reflect the trend discussed in the paper: 1. The attention entropy is lower when processing With CoT data compared to No CoT data (18.0% reduction on average across all layers), indicating that real-world CoT data indeed exhibits a sparser structure; 2. On the same With CoT data, Qwen2-Math-7B model demonstrates lower attention entropy compared to Qwen2-7B model (17.4% reduction on average across all layers), suggesting that fine-tuning on CoT data promotes the development of sparser attention patterns. We argue that these reductions (18.0% and 17.4%) are statistically significant, indicating the sparse structure inherent in CoT data as well as the development of sparser attention from training on CoT data. )

Q1: Why do we use random embedding instead of orthogonal ones?

We agree that using random embedding adds complexity to our analysis. However, this complexity is necessary because when we use orthogonal embeddings, the number of parameters in the model will not fall into the regime of Theorem 2 and our lower bound will no longer hold. Hence, if we use orthogonal embeddings, we could not show theoretically there is an exponential sample complexity gap in this case.

We thank the reviewer for acknowledging the paper's novelty and clarity and address their questions below.

Why do we study the fixed secret parity setup?

In Weakness 2, the reviewer asked whether the location of the secret is fixed during training and inference time. This is indeed correct. We would like to restate the motivation for our study on the parity problem with a fixed secret key.

1. We aim to study how Chain of Thought (CoT) improves sample complexity in reasoning.
   1. Hence, we choose learning a binary secret serves as a minimal example of learning an algorithm
   2. We both lower bound the sample complexity without CoT and upper bound the sample complexity with CoT, showing that there is an exponential sample complexity gap.
   3. Given our lower bound, we believe that the simplicity of the fixed-secret learning task is indeed a strength rather than a limitation. Our lower bound shows that without CoT, Transformers will struggle to learn even simple algorithms like parity, let alone other more complex algorithms.
2. We intentionally omitted a few-shot examples to focus on studying the role of CoT independently.
   1. It has been shown empirically that models can learn to perform CoT without any prompting or in-context examples.
   2. As our primary goal is to understand the role of CoT instead of in-context examples, we deliberately choose not to include few-shot examples in our analysis.
   3. This is also consistent with previous works $1,2$ which study CoT in the setting without in-context examples from a representation theory perspective.
3. Our results already shed insight into more realistic setups such as real-world math problems.
   1. While we studied CoT in the synthetic parity task, our results showing that training on CoT data leads to sparser attention is validated on real-world math datasets.

$1$ Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

$2$ Chain of Thought Empowers Transformers to Solve Inherently Serial Problems

W1: It is difficult to find proof for Theorem 3

In the submitted version, we first outline our proof between line 961 and line 971 and then prove the supporting lemmas for the proof of Theorem 3 in the rest of the appendix. Following the reviewer’s suggestion, we updated our draft and added an explicit pointer to the proof in section A.5.

W2: We use fixed secret for training and inference

The location being fixed during training and inference is indeed correct and we have restated our motivation in the response above. We agree with the reviewer that extending this work to a meta-learning setting—where the model would need to learn an algorithm for the parity problem when all queries and keys are provided in the context—would be an exciting direction. However, this presents significantly greater challenges for theoretical analysis, so we leave it for future exploration.

We thank the reviewers for recognizing the clear motivation of our work and the interesting results it provides.

W1: Is the permutation of indices in CoT sequence random or fixed during training?

The permutation $S$ is fixed during training. The reviewer's comment is entirely correct: If the permutation is not fixed, then it is impossible for any model to achieve perfect accuracy. We have provided further clarification in the revised version.

W2: Include discussion regarding a recent work.

We thank the reviewer for introducing this work and have added it to our related work section. In this work, the authors propose the concept of "globality degree", which shares a similar high-level intuition to the "sparsity dependence" we explore. The paper conjectures that a distribution is weakly learnable for transformers if and only if its globality is constant, and that scratchpads can enhance learning by breaking down the globality of the problem. These conjectures are mainly supported by extensive experiments, but lack theoretical proof. They also discuss the parity problem as a special example.

While the overall message aligns with our own, there are key differences. First, we theoretically establish a separation between the sample complexity lower bound without CoT and the upper bound with CoT for the parity problem, whereas [1] offers only an empirical conjecture. Second, we demonstrate that CoT data induces sparse attention supported by theoretically analyzing the training dynamics as well as empirical experiments, an aspect not emphasized in [1].

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

W2 & Q2: Justification for Initialization; Why the initialization on $W_{r,1:d}$ in is helpful?

Why the initialization on $W_{r,1:d}$ in is helpful? The initialization we used on the FFN weights $W_{r,1:d}$ is $W_{r,1:d}=\sum_{i=n+1}^{n+k}\sum_{b=0}^1 v_{r,i,b}e_{i,b}$, where $v_{r,i,b}\sim U \\{ -\epsilon,\epsilon \\}$. This initialization ensures that the activation pattern of the FFN (i.e. whether the $r$-th neuron is activated by the ReLU activation) is solely determined by the position $i$, the boolean value $\mathbf b[i]$ and the randomness of $v_{r,i,\mathbf b[i]}$ since $\langle e_{i,\mathbf b[i]},W_{r,1:d}\rangle \approx v_{r,i,b}$ with high probability due to near-orthogonal embedding. As a result, for a fixed $(i,\mathbf b[i])$, the FFN can be viewed as a linear function, which consists of the neurons with $v_{r,i,\mathbf b[i]}>0$. At phase 1, this linear function will develop a stronger correlation with embedding $e_{(S[i],\mathbf b[j])}$.

The explanation is mainly provided in the proof sketch of Phase 1. We've provided further explanation in the revised version.

Justification for the initialization: Our analysis relies on the embeddings $e_{i,b}$ being near-orthogonal to each other. If we assume the embeddings $e_{i,b}$ are exactly orthogonal, such as one-hot vectors, then the initialization of $W_{r,1:d}$ involves randomly sampling each coordinate from $\{-\epsilon, \epsilon\}$, which is a standard approach. To accommodate the more general case where $e_{i,b}$ are not strictly orthogonal, we rotate to a non-orthogonal basis defined by $\{e_{i,b}\}$, which simplifies our analysis.

We appreciate the reviewer’s recognition of our work's novelty and significance. Below, we clarify and address your questions.

W1 & Q1: It should be clarified that the lower bound is about sample complexity. Could you comment on how the established lower bound compares to other known lower bounds for parities?

As stated in Theorem 2, the lower bound we establish is indeed about sample complexity. Our result builds on memory-sample lower bounds from classical online learning communities [1] to show that $2^{\Omega (k)}$ samples are required to learn parity when the number of parameters is constrained to $o(nk)$, regardless of how much computation is put. There are indeed other lower bounds for learning parity in the literature, such as those under statistical query framework [2] and in Abbe and Sandon 2020 [3]. However, these lower bounds apply to setups distinct from ours. Below, we provide a detailed comparison of our results with these prior works.

Comparison with the statistical query model[2]: It is well-known that $\Omega(n^k)$ queries are necessary in the SQ framework to learn parity [2]. In this framework, algorithms access an oracle that gives estimates of the expected value of some query function (for example, population loss gradient) over the data distribution. However, this framework does not apply to the mini-batch SGD training considered in our work, where gradients are averaged over small batches rather than the full data distribution. Our choice to focus on mini-batch SGD is motivated by two reasons: (1) Our primary goal is to analyze the sample complexity of learning parity with and without CoT, where population GD is not applicable. (2) mini-batch SGD is more aligned with real-world practice.

Comparison with Abbe and Sandon 2020 [3]: [3] establishes two lower bounds for learning parity with one-pass SGD in different setups. The high-level idea is that under additional gradient noise or memory constraints, the output of a neural net trained on a parity function is statistically indistinguishable from the output of the neural net trained on junk data labels. However, these results rely on setups that diverge significantly from ours:

• SGD with additional randomness [3, Theorem 6]: This result demonstrates that under inverse-polynomial Gaussian noise, polynomial-size neural networks fail to learn parity within polynomial steps of SGD. By contrast, our work assumes an accurate SGD oracle free of such gradient noise.
• SGD with memory constraints: [3, Theorem 5 & Remark 6]: This result shows that when only $o(n / \log n)$ weights are updated per training step, polynomial-size networks fail to learn parity in polynomial steps of SGD. However, we consider standard SGD training where $\Omega(n)$ parameters are updated simultaneously in each training step.

Furthermore, our lower bound applies to constant-pass training, while theirs only apply to one-pass SGD.

To summarize, our lower bound focuses on constant-pass mini-batch SGD training without additional noise when the number of parameters is limited ($o(nk)$), but all parameters are updated simultaneously in each training step, which is different from the above setups. We have revised our related work section to include a discussion of these differences

[1] Xin Lyu, Avishay Tal, Hongxun Wu, and Junzhao Yang. Tight time-space lower bounds for constantpass learning, 2023.
[2]Kearns. Efficient noise-tolerant learning from statistical queries, 1998. [3]Emmanuel Abbe and Colin Sandon. Poly-time universality and limitations of deep learning, 2020.