Transformers Provably Solve Parity Efficiently with Chain of Thought
Authors: Juno Kim
Taiji Suzuki
Venue: ICLR 2025 Oral (2025)
Keywords: transformers
chain of thought
parity
self-consistency
Primary area: learning theory

Abstract: This work provides the first theoretical analysis of training transformers to solve complex problems by recursively generating intermediate states, analogous to fine-tuning for chain-of-thought (CoT) reasoning. We consider training a one-layer transformer to solve the fundamental $k$-parity problem, extending the work on RNNs by \citet{Wies23}. We establish three key results: (1) any finite-precision gradient-based algorithm, without intermediate supervision, requires substantial iterations to solve parity with finite samples. (2) In contrast, when intermediate parities are incorporated into the loss function, our model can learn parity in one gradient update when aided by \emph{teacher forcing}, where ground-truth labels of the reasoning chain are provided at each generation step. (3) Even without teacher forcing, where the model must generate CoT chains end-to-end, parity can be learned efficiently if augmented data is employed to internally verify the soundness of intermediate steps. Our findings, supported by numerical experiments, show that task decomposition and stepwise reasoning naturally arise from optimizing transformers with CoT; moreover, self-consistency checking can improve multi-step reasoning ability, aligning with empirical studies of CoT.

Published as a conference paper at ICLR 2025
TRANSFORMERS PROVABLY SOLVE PARITY EFFICIENTLY WITH CHAIN OF THOUGHT
Juno Kim
Taiji Suzuki!2
DDppartment of Mathematical Informatics, University of Tokyo Center for Advanced Intelligence Project, RIKEN "j junokim@berkeeey edu
ABSTRACT
This work provides the first theoretical analysis of training transformers to solve complex problems by recursively generating intermediate states, analogous to fine-tuning for bhain-of-thought (CoT) reasoning. We consider training one-layer transformer to solve the fundamental k-parity problem, extending the work on RNNs by Wies et al. (2023). We establish three key results: (1) any finie--preistie madient-bbsed algorithm, without intermediate supervision, requires substantial iterations to solve parity with finite samples. (2) In contrast, when intermediate parities incorporated into the loss function, our model can learn parity in gradient update when aided by teacher forcing, where ground-true labels of the reasoning chain are provided each generation step. (3) Even without teacher forcing, where the model must generate CoT chains end-to-end, parity can be learned efficiently if augmented data is employed to internally verify the soundness of intermediate steps. Our findings, supported by numerical experiments, show that task decomposition and stepwise reasoning naturally arise from optimizing transformers with CoT; moreover, self-consistency checking can improve multi- step reasoning ability, aligning with empirical studies of COT.
INTRODUCTION
I Large language models (LLMs) based on the transformer architecture (Vaswani et al., 2017) have achieved astounding success across a variety of natural language processing and machine learning tasks (see e.g. Wan et al., 2024; Minaee et al., 2024; Naveed et al., 2024; Zhao et al., 2024). However, they often struggle when tasked with solving complex reasoning problems, especially in a zero-shot setting without any form of intermediate guidance supervision (Geva et al., 2021; Rae et al., 2022; Arkoudas, 2023; Wang et al., 2024). These failures are particularly evident in tasks requiring multi-hop reasoning compounded logical steps (Sakarvadia et al., 2024).
promising approach to overcome these limitations is chain-of thtthutght (CoT) reasoning, where the model is prompted or fine-tuned solve complex tasks step-by-step by explicitly making intermediate reasoning steps to arrive at the desired answers (Wei et al., 2022; Kojima et al., 2022). Since its discovery, CoT reasoning has been shown to significantly enhance the problem-solving capabilities of LLMs while also increasing the interpretability and trustworthiness of the reasoning process, and has spawned numerous prompting techniques (Liu al., 2023; Qiao al., 2023) and applications for reasoning, mathematical
variety of downstream tasks including common-sente symbolic or multi-modal reasoning: see e.g. Zhang et al. (2023b): Yu et al. (2023); Chu al. (2024) for surveys on COT. In particular, besides being used as prompting method, directly training or fine-tuning models to generate CoT has also been shown to significantly improve multi-step reasoning performance (Nye et al., 2021; Wei et al., 2022; Zelikman et al., 2022; Lightman et al., 2024).
croblem-solving and
Despite these empirical successes, however, the theoretical understanding of the CoT mechanism and task decomposition in transformers is still limited. Existing works focus characterizing the expressivity of transformers equipped with COT. providing constructions which can solve certain complexity classes (Feng et al., 2023; Merrill & Sabharwal, 2023; 2024; Li et al., 2024b), studying the class of functions that can be learned in-context with CoT (Li et al., 2023; Bhattamishra et al., 2024), or analyzing the estimation error of multi-step models (Hu et al., 2024). Nevertheless, such

---

Published as a conference paper at ICLR 2025
approaches do not indicate how such capabilities might emerge when training transformers to generate reasoning chains. Li et al. (2024a) analyze the training dynamics of a one-layer transformer in an in-context learning setting and show that CoT ability may be acquired: however, they do not consider explicitly training with CoT chains, which is a more difficult problem since the objective depends on the recursive application of the transformer to itself.
In this paper, we seek to formalize the mechanism through which stepwise reasoning emerges in transformers optimized to generate CoT chains. We focus on the specific problem of bit subset parity (learning the parity of an unknown subset of k bits from a ddbit input), which is known to be impossible learn end-to-end with any inite-precissio Shwartz et al., 2017: Shamir, 2018). In contrast, Wies et al. (2023) have demonstrated that recurrent neural networks (RNNs) can solve parity efficiently when provided with intermediate supervision. We build on this direction to establish positive optimization guarantees for the transformer architecture. Our object of study a one-layer transformer incorporating a softmax attention layer, feedforward layer and positional encoding, that is recursively applied to its own output to generate sequence of intermediate parity computations to arrive at the desired output, analogous to CoT generation. Our contributions are summarized as follows.
gradient-based algorithm in polynomial steps (Shalev-
We extend the impossibility result for parity (Theorem 1), which was established only for population gradient descent, to the more realistic sanite-sample setting in Theorem 2. We prove that any iterative algorithm with access to approximate gradient oracle for the end-to-end empirical loss cannot solve a random target parity within a specific polynomial number of steps.
In contrast, we show that when the loss is summed over all intermediate states, by utilizing teacher forcing, form of process supervision wherein ground-truth intermediate steps provided during training," our model can learn any parity in a single gradient update (Theorem 5). This shows the benefits of training directly with CoT chains to acquire task decomposition ability.
We further consider training with CoT generated end-to-end without teacher forcing, ? and show that parity can still be learned in a logarithmic number of steps if augmented data is employed to check the validity of intermediate steps (Theorem 7), thereby mimicking self-consistency checks often used in CoT reasoning (Zelikman et al., 2022; Wang et al., 2023; Huang et al., 2023a).
We conduct numerical experiments supporting our findings (Section 4 and Appendix D).
Our results provide theoretical insights into how transformers can naturally and efficiently optimize to perform task decomposition, emphasizing the role of explicit intermediate supervision for complex tasks. Moreover, these findings corroborate recent empirical studies on CoT reasoning demonstrating improved performance through process supervision and internal validation of reasoning chains (Huang et al., 2023a; Tian et al., 2024; Lightman et al., 2024).
1.1 RELATED WORKS
Complexity of transformers. A line of work aims to understand the effectiveness of CoT from the perspective of complexity theory. Feng et al. (2023) show that autoregressive transformers of constant size can solve basic arithmetic tasks by recursively generating COT reasoning steps, which not possible when directly generating the solution; this separation arises because looping the generated outputs back to its inputs increases the 'effective depth' of the model. Works such as Chiang et al. (2023); Merrill & Sabharwal (2023) study the expressivity of fixed-precision transformer architectures in terms of classes of formal languages. Merrill & Sabharwal (2024); Li et al. (2024b) show that CoT reasoning enables recognizing wider language classes, and characterizes the increased expressivity depending on the length of the reasoning chain. Sanford al. (2024) studies the relation between transformers and massively parallel computation protocols, showing that logarithmic depth suffices to solve multi-hop induction tasks that cannot be efficiently solved by other sequence models.
'Teacher forcing or process supervision is a training procedure for recurrent models in which the model receives the ground truth output at time i as input at time t ++ 1 during training (Goodfellow et al., 2016, p.377). Many fine-tuning methods with ground-truth CoT chains implement teacher forcing. being more effective than output supervision with chains generated end-to-end (Deng al., 2023; Tian al., 2024; Lightman al., 2024). 'Teacher forcing can induce exposure bias where model is not robust to its own errors. In practice, partial (scheduled random) teacher forcing methods are used to overcome this issue (Bengio et al., 2015; Goyal et al., 2017; Mihaylova & Martins, 2019).
2

---

Published as a conference paper at ICLR 2025
Additionally, Li et al. (2023); Bhattamishra et al. (2024) study the class of functions that can be learned in context by transformers with CoT from the point of view of in-context learning.
Optimization and generalization of COT. Zhu al. (2024) study the 'reversal curse' via the training dynamics of a one-layer transformer and shows that the model fails to generalize from A A B C' to A C' as an argument for the necessity of explicit step-by-step reasoning. Hu et al. (2024) study CoTT prompting from a statistical estimation perspective by introducing a multi-step latent variable model for CoT and analyzing its approximation, generalization and prompting-bar based errors. Notably, Li et al. (2024a) study the training dynamics of a one-layer attention-only transformer model in an in-context learning setting and show that CoT generalization capability can be obtained. However, this does not address the possibility or benefits of training with CoT chains. Lightman et al. (2024) empirically study training L.MS with either process or outcome supervision, showing that the former significantly outperforms the latter when training to solve challenging reasoning tasks.
Parity and task decomposition. The difficulty of learning parity without task decomposition is established in Shalev-Shwarte et al. (2017); Shamir (2018). The work most relevant to our paper is Wies et al. (2023), which study task decomposition for parity with classical Elman RNNs. They show that by incorporating intermediate states into the loss function and utilizing teacher forcing, parity can be solved with polynomial iterations and embedding size. Our Theorem extends this positive result autoregressive transformers, rigorously establishing the benefits of CoT-based training.
2 PROBLEM SETUP
Notation. We write [n] :- 21, 2.... .n. for any integer n. Scalar operations apply componentwise
to vectors, e.g. for R" we write ((2) (#(z)), (b22) () =² z O 2 (z ... 2 and
|=| (121),-. ... (2.l))_ The 2-norm is always denoted by II- The multi-linees inner product or contraction ø1, ... 2, E R" for any r E N is denoted as (z1,-.. ,z, = >i= 21, 2p,5. In
particular, (21) = x 1,, and (21, z2) = 2 22.
2.1 THE PARITY PROBLEM
Let d 2 k > be integers and let P denote the set of size ls subsets of (1, ,d) equipped with the uniform distribution. In this paper, we study the k-parity problem for d-bit inputs æ = (zz) Unif({#1)
where the output y IIeep Ty is determined by the parity of unknown subset of bits p P. We abuse notation and identify the set of indices p with the corresponding parity mapping I + IhEp Zy. Given n samples (x', (²))E((). our goal is to predict the parity of any test input.
known that parity is fundamentally difficult the sense that it cannot be solved polynomial time by any finite-precision gradient-based algorithm, such as neural networks. More precisely, let {fe E }} be any differentiable (w.rt. 8) parametrized model with polynomially bounded gradients, fo())) (=)] O(poly(d)) and define the population loss F. E. [yy fo())) We presume access to an é-upproximate gradient oracle V for T., which takes any 0 E θ as query and returns vector V (() satisfying [|V Vī(0)]|2 S E, potentially in adversarial manner.
Then the following holds:
Theorem 1 (Wies et al. (2023), Theorem 4). Let Eu-1 be the zero-one loss. There exists an O(e-d/3) 4/3)- approximate oracle & such that' the output ((A) of any iterative algoritiu A which sequentially makes most O(poly (d)) queries to VL must satisfy
Ey [f0-
SecAy( Se()) 2 - O(e-9)
with probability at least 1 - O(e when the target parity p is uniformly sampled from P.
The intuition is that the set P of parity functions is exponentially large in the sense that all elements of P are pairwise orthogonal with respect to the data distribution. This implies that the variance of
'The original paper states that A can be any iterative gradient- algorithm which receives an li(e (d/3- approximation of the gradient at each step. However, to be more precise, the result is only valid for certain adversarial perturhation schemes.
3

---

Published as a conference paper at ICLR 2025
h =
2/3 =
h h= 2
221
222
h=
217
218
119
3/20
0
21
222
43
24
455
276
25
210
211
212
213
214
I15
216
Figure 1: A hierarchical decomposition of 8-parity problem for 16. Here .17 21 so that c 17] 1, c2(7] = 4, p[17] = 21 and h[17] 1.
each gradient call (0) with respect to the target parity p is exponentially small (Shalev-Shwartz al., 2017) and is drowned out by the noise from the adversarial oracle, so that information be gained on the target without exponentially many queries. See Section 3.1 for more details.
Task decomposition. As in Wies al. (2023), we assume k = 2" for an integer v for simplicity and decompose the problem into a hierarchy of 2-parity computations which can be efficiently learned sequential manner by our model. This is expressed as complete binary tree of height v and 2k nodes. The lowest level contains & nodes representing the bits 2 for m E [k]. The remaining nodes are labeled Id+1, -d+h-1 starting from the next lowest level and moving upwards, left
to right. The largest index level é for 0 S S v is denoted as de = Σド=1 2" do = Also, for each m d, the indices of the two child nodes of Im are denoted as C1 [m]] C2 [m] where 1 < Ci [m] < [m] m. In addition, the parent node index of xm denoted as p[mn] and the level or height of Im is denoted as h[m], so that di dnim]-1 < m S dním).
2.2 TRANSFORMER MODEL
C1
Cm
e
e2
ed
e
em
m
+ d
im
Bstk-1
R61
Recursive generation intermediate states.
(b) Filtered generation with data augmentation.
Figure 2: Illustration of the recursive data generation process by the transformer model. (a) Each token consists of one-hot positional encoding e, and parity data æj. The d input tokens (blue) fixed. The token &.., is generated at the (m dh step by computing attention scores based position, combining the previous tokens and applying the feedforward layer d. Æd+h-1 is returned as the model prediction. (b) For the no teacher forcing setup Section 3.3, data augmentation u,
implemented to check for self-consistency. If the augmented outputs from the previous generation (red) are uninformative, filter is applied zero out the subsequent output.
We study one-layer transformer architecture employing absolute positional encoding and single- head softmax attention layer followed by a shallow feedforward layer; skip connections are omitted for simplicity. See Figure 2 for visualization of our setup.
4

---

Published as a conference paper at ICLR 2025
Data encoding: Each input token z; =x))==1 for E [d] is the n-dimensional vector consisting of the jth bit of each sample æ'. We also add dummy tokens Id+1.".. .d+k-1 initially to O which will learn to sequentially generate the actual intermediate nodes. Each z; is concatenated with the one-hot positional encoding e; E R for j E [d [d+ k 1] to form the internal input
Pj - (x) e²)" E I to the attention layer.
Softmax attention layer: The attention layer is defined as in (I) in terms of key, query and value matrices K, Q. V. We fix the first columns of K, Q to zero so that the attention scores are determined by only the positional encodings. This ensures that the transformer focuses on learning which positions contribute to the parity at each step. K, Q are then reparametrized by single matrix W E R (d++=11) ")"; conversely, the value matrix is set to only preserve the z component, as follows.
K KJQ (0 Opxn
O Onx(daR-1)) W
V V= (Inxn (nx(d+x-1))-
(0(++R1))m
This type of reparametrization is common in the literature to make dynamical analysis tractable (Zhang al., 2023a; Huang al., 2023b; Mahankali al., 2023; Kim & Suzuki, 2024).
Feedforward layer: The feedforward layer realizes fixed link function & : 1, + [-1, 1]. applied elementwise and only to the æj component; the positional encodings are not affected. To exploit the decomposition of our task into 2-parities, we choose & such that ((0) 1, x = 1 so that sums are converted into parities, i.e. o( (()))) = ab for a, b E {+1}. Moreover, we require that ø' (0) = (1) O and assume Φ symmetric and sufficiently regular, so that we may expand @(t) = - at + (( and (t(t) 2ct + 0(
The transformer computes TF 1F 1d W) (#). -d+k-1) where the original data ē. = J. j E [d] remain unchanged and tokens 1d+1 - are computed as
im = ((2m,
2m
VD,- softmax(p] K (TQ))))
Σ j=1
(wm))
Ø=1
where the softmax scores 0 (wm) )
5m Σn=1 e 2ma.m Here, we have implicitly added the S d. Note that each 2m Fm will be
causal mask j -00 to the attention layer for contained in the cube [-1, 1jd as long as the input tokens are also contained in [-1,
Chain of thought. Consider repeatedly applying (F()) to its own output to generate a 'reasoning chain.' Since the input tokens are fixed, the token æ will be updated once and then always yield the same value afterwards. Next, since a depends on the input tokens and d will be updated twice before becoming fixed. Repeating this, the entire chain stops updating after at most k steps, yielding the output
TF (k-k-1)(2)) (x),.... Bd On, Oni W) = (æ11 (d+k-1)
where the intermediate predictions are recursively computed as t. = (E) Jg (2m))))) Finally, the top node is returned as the model prediction ÿ = 0atk-11
This process can be setting: instead of one-shot predicting y from x', the model starts by solving simpler subtasks and uses the information attack compound problems, learning to generate intermediate reasoning steps Id+1 + ... + Id+k-1 to finally arrive at the desired solution. Importantly, this process is not possible if the model is only trained on the one-shot data (a+ y y')xg(n] as we show in Section 3.1. Instead, we incorporate the prediction error for all intermediate states directly into our loss function (Lightman et al., 2024). We also consider shortening the reasoning chain by using a different causal mask in Section 3.3, which will result in improved control of error and faster convergence.
simplified version of CoT reasoning, albeit not in in-context learning
3
MAIN RESULTS
3.1 HARDNESS OF PARITY WITHOUT COT
Before analyzing our transformer model, first prove negative learning result in the absence of intermediate supervision that extends Theorem 1, which was stated with respect to the population objective T. and zero-one test loss fo fo-1, , finite samples and mean squared loss.
5

---

Published as a conference paper at ICLR 2025
R be any differentiable parametrized model and suppose select the target parity p uniformly at random from P. In the finite-sample setting, n i..d. samples (æ' "a",y")ee y")ee() are generated as x² ~ Unif((#1(4) y = p(æ') ) and we are given access to (approximate) gradients from the empirical loss
Let fo {=1}² -
1 2 folla
(0) 2m Σ w fo(x)))²
where Il is the empirical norm. important that the model fe is applied each x' separately and does not cross-reference between different samples, there exist more efficient parity-learning algorithms the data allowed be manipulated freely. For example, Gaussian elimination can solve parity with O(d) samples and (dd) iterations (Raz, 2018). Moreover, this implies that neural networks trained with stochastic gradient descent can also solve parity in polynomial time (Abbe Sandon, 2020). Instead, in setting the model is forced to learn from the averaged gradient signal and can only implicitly utilize the correlation between samples.
We show the following result for learning parities with samiles in Appendix A:
Theorem (hardness of finite-sample parity). Suppose k = (d).
(1) n n = (d()) and fo has polynomially bounded gradients, there exists e -9(()) -aprrxxiit gradient oracle V such that with probability 1 - e over random sampling, the output ((A) of any iterative (possibly randomized) algorithm which makes at most (poly(dd) queries to
v Ln has L2-loss lower bounded as
EpePae [(((x) - fOLA) 1e--0
(2) W (dd) and fell ), there exists an "(d-"))-approzinne gradient oracle A such that with probability 1 :id) over random sampling, the output (AA) of any iterative (possibly randomized) algorithm which makes at most ((d))) queries to V Ln has Lz-loss lower bounded, where 4v 4v2 + 2v3 + 2v4 as
P,p [(p(x)) [e(a)( 1 O
We remark that the bounds are asymptotically optimal since fa = is a valid estimator. Moreover, the expectation over p E P can be replaced by the corresponding "with high probability' statement.
A counter-intuitive aspect of the above result is that parity becomes potentially more difficult when the number of samples increases. Indeed, with exponential samples n = e(d) (1) we basically recover the statement of Theorem 1, while the guarantees for poly(d) (2) are also polynomial in d. This is because the difficulty of parity (Theorem 1) fundamentally depends on the following result:
et al. (2017), Theorem 1). Suppose x be a random variable in Rd. Let H. be a class of bounded real-valued functions on IR such that Ez (h(()))' = 0 for any distinct h, h E H and fe a differentiable parametric model with gradients bounded by E. fo & 2 Then for the loss (0) E. [((xx)) [o(x))
Proposition&
where is chosen uniformly
at random from H. the gradient variance is bounded
F(0)2
Var(0; H) = EAEH (1) ENEH F (1) <
H
Since all (3) O(d) parities in P are pairwise orthogonal with respect to the uniform distribution Unif((+1)(), it follows that the variance of V T. is exponentially small and the target signal can be drowned out by a correspondingly small noise from the oracle. However, this is not true for the empirical distribution which cannot distinguish all elements in P with only poly samples; the empirical correlation of two random parities will generally be (nn Therefore more careful decorrelation argument is needed, resulting in the weaker guarantees of Theorem 2(2). Another technical difference is that Theorem only considers the strong zero-one loss (more formally, their results can be seen to hold for any parity estimator PD(A) E P depending on the algorithm output), while prove the T.2 lower bound for any real-valued estimator fe(A)-

---

Published as a conference paper at ICLR 2025
3.2 COT WITH TEACHER FORCING
When training with teacher forcing, at each position d+1 S m S d the ground-truth labels of the preceding intermediate states æ1,--. 2m-1 1 are fed into the transformer input obtain the predictor o.m at the math position,
mm TF (1, 1, On Opt W) W)m.
The loss function then computes the squared error over all states,
d-k-1
L(W)= = 2n m=d+1
Σ
Zm
Im Imml?
(2)
Since each sequence of values Bd+1,i. - are generated depending only on the correspond-
ing sample x' and the parameter matrix W, this can be rewritten in terms of the augmented labels y' (1'+111" x2(+--1)
n
L(W)= 2n
Σ 1 (a';; W)Il (n'; W) Em.i. d+1s Sd
for a fixed mapping f {+1]d x (d+ + k-1 mirroring the setting of Theorem 2. Hence model does not cross-reference between samples; moreover, the gradient of fo is bounded as Lemma 4. For all W it holds uniformly that (";; w)ll S O(vd).
At inference time, test inputs T1, Zid are randomly generated and the prediction for ytest p(a1.... ... (Fa) is computed by iterating TF to generate all h reasoning steps without reference to ground-truth labels; Gtest (p((-1) (x),, ", Ba, On, On; W) W)d+k-1. Our positive learning result in this setting is as follows.
Theorem (CoT with teacher forcing). Suppose (2d2++))))))))) ) 0, d is sufficiently large and let V be any O(d ((d-2-e/8 -approiimte gradient oracle.* 4 Set initialization W()) and learning rate 17 ((d2+e/16) Then for any target parity p E P, it holds with probability cxp(-d) /d2 over random sampling that the one-step update ((1) W LL(() w.r.f. the objective (2)
with teacher forcing achieves loss
VBestt
Ytest ol - (8)
On the other hand, Theorem 2(2) shows that when ((d any iterative algorithm querying an Odd-2-e/8) (aapoiat))a oracle, with gradients bounded as in Lemma 4, requires more than ((d'/1) queries to attain a nontrivial (< 2)) loss. This establishes a strict separation between learning parities without intermediate supervision and our CoT transformer. The gap increases with more samples as increases; moreover, when eQ(d), we have much stronger separation by Theorem 2(1), where an exponential number of queries is required to learn
Sketch of proof. The result is shown by explicitly caleulating the gradient with respect to each weight Wjmm and extracting the gradient signal. As the softmax scores are uniform at initialization, the gradient can be expanded to obtain multilinear contraction or 'interaction' terms between the tokens 21, I m, one such example being
(m 2m, Sm)) n(m - 1)2 n
Σ (m Ra; mp) 0.,
In the above equation, if a are the two child nodes of m. the parity IaI = 1 will be trivial and (m Fo, xz) = n. On the other hand, for nontrivial parities the interaction strength will generally be O log d) due to sample concentration. For sufficiently large n, the trivial parities dominate, allowing us to extract the leading term. Performing these computations up to fourth order interaction terms, we show that the dominating signal of the gradient (dd=²) when j = ci [m], c2(m) and (dd-2-//8) otherwise. Hence the transformer learns increase only the weights at the relevant positions for each subtask, and is able to compute the desired 2-parity &m Rs o( ( [m] + @e.(mi)) As @c, (mj at each node during its forward pass. The full proof is provided in Appendix B.
In fact, we only require that each component of the gradient has error at most 0(d-2-e/8) for Theorems 5, which follows since the Tee is bounded above by T-2.
7

---

Published as a conference paper at ICLR 2025
3.3 COT WITHOUT TEACHER FORCING
In this section, we extend Theorem 5 to training transformer without teacher forcing, which is employed alongside teacher foreing in practice ensure robustness at test time (Bengio al., 2015; Goyal et al., 2017; Mihaylova & Martins, 2019). The main difficulty in this setting is that wrong answers propagate later generation steps, exponentially amplifying errors and drowning main gradient signals. Error accumulation is also central practical issue of CoT (Zhang & Parkes, 2023; Wang et al., 2023). To solve this issue, we make some modifications our transformer model.
the
First, we minimize the number of required reasoning steps by imposing a slightly stronger form of autoregressivity where each intermediate state x depends on all tokens æ J = 1. dn(m) up to the previous level, rather than the immediately preceding token. This can be expressed as the causal mask Wj. 6 -co for j dnim]-1 -1 or m S d; see Figure 3. This ensures that the model gradients polynomially bounded as in Theorem 2 and that emrors can propagate logarithmic rather than a linear number steps, and can be easily implemented as the indices de are known.
d
do
d1 d2
Figure 3: Causal mask for W with teacher forcing (left); without teacher forcing (right). The gray entries are set to -oc.
Second, we implement data augmentation technique where random d-bit strings ut ~ Unif((+1})) i [n'] are appended to the original dataset (x*" (e(n)- The resulting augmented tokens denoted as 2
(ar] u, E Ra**', U (wt) so that Pj ((at)) TT)TT (the notation is extended to > d), and the key, query and value matrices are appropriately enlarged. The ground truth labels as well as the intermediate states for the augmented data are unknown, so they are not included in the loss function. Nevertheless, unlabeled data can still suffice for self-consistency (Huang et al., 2023a); their purpose is to filter for 'faulty reasoning' in the following sense. If the weights are not sufficiently trained, the output of a node x; will consist of all nearly -Is and thus be uninformative for computing any parities. If the augmented tokens newly generated in the previous iteration of TF(-) (i.e. up to ud;-,) are uninformative, reasoning will be wrong. This is achieved by adding the following filter after the feedforward layer Φe
zero out its output on the basis that all subsequent
0
|u, + 1 lla < Eo for any de-2 < de-1:
+z+ E Ratn'
ne(2*)
x
otherwise.
Without teacher forcing, during training the entire reasoning chain is generated by iteratively applying TF to its output until convergence, which takes logg i rather than k 1 steps due to the imposed block autoregressivity. Henee TF(') ... . On W) (at, ... a a where the tokens # ...
" -d+k-1 are recursively generated per level as
daim) 1
- = [n/m] o ( (2
2+ =
Σ j=1
og (um)aa
The loss is computed against the ground-truth labels as in (2). As before, each sequence of generated states depends only on each sample æt and the augmented data U (u')sejn) so we may express
1 1 2n
L(W, U)
(x=;; (wr: W. U)l2, (x W, U) Im.i
(4)
for a fixed mapping fx so that the samples are again not cross-refereeccd. By considering the propagation of gradients up the chain, the gradient f can be shown to bounded as follows.
Lemma 6. For all x. W, U have / (æ;; W, U)] & O(d9) where g logg | aloc 1/2.
8

---

Published as a conference paper at ICLR 2025
The exact exponent 9 depends on the shape of Φ. Since ((0) = -1 and o(1) 1, it must hold that || > 2. Conversely, any such || may be achieved by taking & to be locally quadratic around 0, and smoothly joining the curve segments with straight lines of slope = + e). Furthermore, such a link function can be realized by a simple shallow feedforward layer using e.g. O(1) ReQU neurons. Hence 9 can be taken to be arbitrarily close to 1.5.
Finally, we implement a simple weight quantization method by rounding each entry of W to the nearest integer after every update; W(++1 = r[w()) 7/ WVWLLW( U)], where r IR + Z is the nearest-integer operator. Equivalently, the gradients themselves are quantized. Integer-based quantization methods are widely used in practice to accelerate training and reduce memory usage (Wu et al., 2020; Jacob et al., 2018), and have been successfully implemented in LLMSs to facilitate efficient fine-tuning (Dettmers et al., 2022; 2023). In our theoretical setting, quantization also allows us to simplify computations involving propagation of error.
In this setting, we obtain the following learning result.
Theorem 7 (CoT without teacher forcing). Suppose n = (dd for > O, n' = poly (d), is sufficiently large and let & be any O(d-2-e/8) 2 aapproiimate gradient oracle. Set W and " (dd++e/16). Then for any target parity p E P, it holds with probability 1 exp( exp(-d(<A1/2) over random sampling of (original and augmented) data that the sequence of updutes W(++1) ([(() 7)(L(W()) U)] w.r.t the objective (4) without teacher forcing achieves exponentially small loss Ueet ytest oloc I ex)( in log, k iterations.
This gives the same order of separation from Theorem 2(2) as in Section 3.2. Hence transformers can learn parities even without teacher forcing. if the consistency of the chain of reasoning is suitably controlled for. Moreover, our result shows that logarithmic time suffices to learn parity by exploiting the hierarchical decomposition in Figure This extends the circuit complexity result in Merrill Sabharwal (2024), which states that bounded-depth transformers with logarithmic number of CoT steps can express problems in log-space; Theorem guarantees that transformers of depth one can learn by gradient descent any such function in the exponentially large class P.
Sketch of proof. The idea is to inductively show that each 2-parity subtask xm at level will become solved at time = f. When < 6-22 2, Z.m cannot utilize its child nodes Ze, [m]] [e2[m]] since they will also not be optimized, so the weights do not change. At time - 1, its child nodes learn to output their parities with high precision, the objective is approximately equivalent to that of Theorem 5. Then the gradient signal will similarly concentrate on We, [m].m. [m/[m]] and Fm will become solved in the next step. It remains bound the gradients arising from the loss terms further down the chain xd+1 + ... + Id+k-1 (propagation of error), and verify that irrelevant weights wj. (p[i) # m) and already optimized weights do not change. The full proof provided in Appendix C.
NUMERICAL EXPERIMENTS
In this section, we present numerical experiments which support and complement our theoretical findings. Compared to the carefully calibrated step sizes and weight updates in Theorems 5 and 7, these experiments study a more realistic training scenario by taking relatively small learning rates and tracking the loss trajectories over longer period of training. We train one-layer transformers based on the architecture described in Section 2 to solve a random k-paritt problem with 64-bit inputs for k- = 8, 16, 32. Specifically, we implement and compare the following four models.
Direct: TF(-) is applied to itself A - 1 times to generate the reasoning chain end-to-end and the model prediction ý is directly compared the ground truth y with the prediction loss 2n. 11 y
CoT: (-)) is applied to itself to generate the reasoning chain end-to-end and the sequence of intermediate states is compared to the ground truth as in (2). Here, we also implement the causal mask in Figure 3 (right) so that only logg & iterations are needed, for additional stability.
CoT + teacher forcing: implements the model in Section 3.2 with teacher forcing.
CoT + self-consistency: implements the model in Section 3.3 with the causal mask in Figure (right) and data augmentation for consistency checks. Weight quantization is omitted.
Anyy polynomial order suffices for the number of augmented data samples.

---

Published as a conference paper at ICLR 2025
1.0
1.0
0.8
0.8
Mow Mm
CaT loss (K-322)
Target loss (x-32)
direct
0.6
CoT
CoT+forcing
0.4
Col Col+consiseeecy
0.4
0.0
0.0
50
100
150
200
250
300
300
50
100
150
200
250
300
350
Training
Training
Figure 4: CoT loss (left) and prediction loss (right) curves for the four models when d = 64, A: = 32, For the CoT cconsistency model, dashed lines indicate when the filters of each level are deactivated.
All models are optimized using full-batch gradient descent 100K 64-bit samples with single Tesla T4 GPU. The three CoT models are trained with the 'CoT loss' (2) scaled by k-ニ to match the prediction loss of the direct model. Figure 4 shows training curves for the CoT loss (left) and the prediction loss (right) over 350 epochs when k 32; results for all is and more details are provided in Appendix D.
We first note that the direct model (red) completely fails to learn the target, plateauing almost immediately. We observed that the weights become nearly uniform so that ý R O,, and the prediction error is stuck 0.5. This was improved by using multilayer transformer instead of repeated composition. The basic CoT model (yellow) is able to significantly decrease CoT loss but fails to fully solve the problem and eventually becomes unstable. Moreover, the prediction loss never improves beyond 0.5. Indeed, due to the hierarchical structure of parity, the model has no chance of making an informative prediction at the last level Ød+k-1 unless all preceding levels have been fully solved. In contrast, we verify that CoT with teacher forcing (blue) solves parity efficiently predicted in Section 3.2, even with a small learning rate. After a burn-in phase, the COT loss steadily decreases to nearly zero, at which point the prediction loss also decreases rapidly the final level is solved.
CoT with consisessst (green) is also able to solve parity efficiently as predicted. Furthermore, the corresponding CoT loss curve clearly exhibits multiple learning stages. In the beginning, the model is essentially optimizing only the first level as subsequent outputs are zeroed out. After a short burn-in phase, the weights are optimized so that the softmax scores concentrate nodes, at which point the CoT loss sharply decreases and the filters for the next level are deactivated, unlocking the next learning stage. This phased optimization repeats until all levels are fully solved and is crucial to arriving at the correct answer (in essence, teacher forcing is doing this at all levels simuitancously). Notably, similar behavior seems arise the basic CoT model well but fails due to accumulating error, further justifying the use of the filtering mechanism.
the relevant
These results confirm that training explicitly for CoT generation can improve performance on multi- step tasks, and that controlling error accumulation via teacher forcing or self-consistency is key to ensuring proper step-by by-step learning.
CONCLUSION
In this paper, by focusing on the k-parity problem, we provide an initial theoretical foundation for training transformers with CoT to perform stepwise reasoning. Our results show that gradient-based learning of parity requires significant iterations without intermediate supervision, but task decom- position using teacher forcing enables efficient learning in single gradient update. Furthermore, when transformers are trained to generate reasoning chains end-to-end, data augmentation and self- consistency checks enhance their ability to solve complex tasks. Our work takes the first steps towards understanding how CoT can be leveraged to improve multi-step reasoning capability of foundation models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
JK was partially supported by JST CREST (JPMJCR2015). TS was partially supported by JSPS KAKENHI (24K02905, 20H00576) and JST CREST (JPMICR2115).
REFERENCES
Emmanuel Abbe and Colin Sandon. On the universality of deep learning. In Advances in Neural Information Processing Systems, 2020.
Konstantine Arkoudas. GPT-4 Can't Reason. arXiv preprint arXiv:2308.025. 2023.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in Neural Information Processing Systems, 2015.
Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and LLMs by learning to learn discrete functions. In International Conference Learning Representations, 2024.
David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In International Conference on Machine Learning, 2023.
Zheng Chu, Jingchang Chen, Qianglong Chen, Weijiang Yu, Tao He, Haotian Wang, Weihua Peng, Ming Liu, Bing Qin, and Ting Liu. Navigate through enigmatic labyrinth: a survey chain of thought reasoning: advances, frontiers and future. In Association for Computational Linguistics, 2024.
Yuntian Deng, Kiran Prasad, Roland Fernandez, Paul Smolensky, Vishrav Chaudhary, and Stu- Shieber. Implicit chain of thought reasoning via knowledge distillation. arXiv preprint arXiv:2311.01460 2023.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. GPT3.int80): 8-bit matrix multiplication for transformers at scale. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: efficient finetuning
of quantized LLMs. In Advances in Neural Information Processing Systems, 2023.
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. In Advances in Neural Information Processing Systems, 2023.
Mor Geva, Daniel Khashabi, Elad Segal, Tushar Khot, Dan Roth, and Jonathan Berant. Did Aristotle Use a laptop? A question answering benchmark with implicit reasoning strategies. Transactions of the Association for Computational Linguistics, 9:346-361, 2021.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http:
//www. www..eeennniingooo.... org.
Kartik Goyal, Chris Dyer, and Taylor Berg- KirggKiikpatrick. Differentiable scheduled sampling for credit
assignment. In Association for Computational Linguistics, 2017.
Xinyang Hu, Fengzhuo Zhang, Siyu Chen, and Zhuoran Yang. Unveiling the statistical foundations of tfain-of-thought prompting methods. arXiv preprint arXiv:2408. 14511, 2024.
Jiaxin Huang, Shixiang Gu, Le Hou, Yuexin Wa, Xuezhi Wang. Hongkun Yu, and Jiawei Han. Large language models can self-improve. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 1051-1068, Singapore, December 2023a. Association for Computational Linguistics.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of Transformers. arXiv preprint arXiv:2310.0522 2023b.
11

---

Published as a conference paper at ICLR 2025
Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient anteger-aribbmeti--only inference. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: nonconvex mean-field dynamics on the attention landscape. In International Conference on Machine Learning. 2024.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In Advances in Neural Information Processing Systems, 2022.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear transformers acquire generalization-guarann CoT ability? In High-dimensional Learning Dynamics 2024: The Emergence of Structure and Reasoning, 2024a.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dis- secting chain-of oh-thotght: compositionality through in in-context filtering and learning. In Advances in Neural Information Processing Systems, 2023.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In International Conference on Learning Representations, 2024b.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. International Conference Learning Representations, 2024.
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhenghao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: a systematic survey of prompting methods in natural language processing. ACM Comput. Surv., 55(9), January 2023.
Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023.
William Merrill and Ashish Sabharwal. A logic for expressing log-precision transformers. In Advances in Neural Information Processing Systems, 2023.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In International Conference on Learning Representations, 2024.
Tsvetomila Mihaylova and André F. T. Martins. Scheduled sampling for transformers. Association for Computational Linguistics: Student Research Workshop, 2019.
Shervin Minace, Tomas Mikolov, Narjes Nikzad, Meysam Chenaghlu, Richard Socher. Xavier Amatriain, and Jianfeng Gao. Large language models: survey, 2024.
Humza Naveed, Asad Ullah Khan, Shi Qiu, Muhammad Saqib, Saeed Anwar, Muhammad Usman, Naveed Akhtar, Nick Barnes, and Ajmal Mian. comprehensive overview of large language models, 2024.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski. Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: scratchpads for intermediate computation with language models. arXiv preprint 122iv:21 14, 2021.
Shuofei Qiao, Yixin Ou, Ningyu Zhang, Xiang Chen, Yunzhi Yao, Shumin Deng, Chuanqi Tan, Fei Huang. and Huajun Chen. Reasoning with language model prompting: a survey. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3368-5393, Toronto, Canada, July 2023. Association for Computational Linguisties.
12

---

Published as a conference paper at ICLR 2025
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang. Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre. Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, Iason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: methods, analysis and insights from training Gopher. arXiv preprint arXiv:21 1446, 2022.
Ran Raz. Fast learning requires good memory: a time-space lower bound for parity learning. ACM, 66(1), 2018.
Mansi Sakarvadia, Aswathy Ajith, Arham Khan, Daniel Grzenda, Nathaniel Hudson, André Bauer, Kyle Chard, and Ian Foster. Memory injections: correcting multi-hop reasoning failures during inference in transformer-bas based language models. arXiv preprint wxXiv2230,005605 2024.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth. arXiv preprint (tt.::20020992666 2024.
Shai Shalev-Shwartz, Ohad Shamir, and Shaked Shammah. Failures of gradient-basee deep learning. In International Conference on Machine Learning, 2017.
Ohad Shamir. Distribution hardness of learning neural networks. Journal of Machine Learning Research, 19, August 2018.
Yijun Tian, Yikun Han, Xiusi Chen, Wei Wang, and Nitesh V. Chawla. TinyLLM: learning a small student from multiple large language models. arXiv preprint arXiv:2402.046 04616, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, Jiachen Liu, Zhongnan Qu, Shen Yan, Yi Zhu, Quanlu Zhang, Mosharaf Chowdhury, and Mi Zhang. Efficient large language models: survey. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. Survey Certification.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe V Le. Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain thought reasoning in language models. In International Conference on Learning Representations, 2023.
Zhiwei Wang, Yunji Wang, Zhongwang Zhang, Zhangchen Zhou, Hui Jin, Tianyang Hu, Jiacheng Sun, Zhenguo Li, Yaoyu Zhang, and Zhi-Qin John Xu. Towards understanding how transformer perform multi-step reasoning with matching operation. arXiv preprint arXiv:2405. 15302, 2024.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in sequence to sequence tasks. In International Conference on Learning Representations, 2023.
Hao Wu, Patrick Judd, Xiaojie Zhang, Mikhail Isaev, and Paulius Micikevicius. Integer quantization for deep learning inference: principles and empirical evaluation. arXiv preprint arXiv:2004.09602 2020.
13

---

Published as a conference paper at ICLR 2025
Zihan Yu, Liang He, Zhen Wu, Xinyu Dai, and Jiajun Chen. Towards better chain-of -thiig prompting strategies: a survey, 2023.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. STaR: bootstrapping reasoning with
reasoning. In Advances in Neural Information Processing Systems, 2022.
Hugh Zhang and David C. Parkes. Chain-of-thought reasoning a policy improvement operator, 2023.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained Transformers learn linear models in-context. arXiv preprint arXiv:2306.0992 2023a.
Zhuosheng Zhang, Yao Yao, Aston Zhang, Xiangru Tang, Xinbei Ma, Zhiwei He, Yiming Wang. Mark Gerstein, Rui Wang, Gongshen Liu, and Hai Zhao. Igniting language intelligence: the hitchhiker's guide from mauin-of-thour reasoning to language agents. arXiv preprint arXiv:2311.1197, 2023b.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang. Junjie Zhang. Zican Dong, Yifan Du, Chen Yang. Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang. Zikang Liu, Peiyu Liu, Jian- Yun Nie, and Ji-Rong Wen. A survey of large language models, 2024.
Hanlin Zhu, Baihe Huang. Shaolun Zhang, Michael Jordan, Jiantao Jiao, Yuandong Tian, and Stuart Russell. Towards a theoretical understanding of the 'reversal curse' via training dynamics. arXiv preprint 2005 206506699006 2024.
14

---

Published as a conference paper at ICLR 2025
APPENDIX
A PROOF OF THEOREM
Denote the empirical inner product on R by (f.g)n 1. Ei=1 ((x")g(x") and the corresponding norm as I (f.f)
We also write
1 Ln.p(0) 2 fellã
1
> (p(x') - fo(x)))2
2n i=1
to emphasize the dependency of Ln on Note that (d/h)* < (+) S (ed/k)* so that |P] = eld)
Bounding gradient variance. Consider the variance of the empirical gradient VLn.p w.r.t. the target parity p:
Var, (0; P) == Eper Lm. B L (0)]1
We proceed to evaluate the magnitude Var, (0; P). For p.p' E with # p' it holds that
n
(p.p')n 71
II JEP
II
II
J'Ep'
jEpAp'
Since HI 2j is i..d Unif(f#1)) for fixed p', by applying union bound over Hoeffding's inequality, follows for d / Add/n that
no² 2
dd) 1
2
2d
Pr sup (p. 2 < PP((PP - exp ppp'
c 2d <
Then with probability at least e (2()) random sampling, every off-diagonal component of the Gram matrix Gp == ((p, (") has magnitude at most 6, while the diagonal entries are equal 1. By the Gershgorin circle theorem, the maximum eigenvalue of Gp satisfies
Imax (Gp) 1| S (IP1-15
thus Amax (GP) < 2(1 V (P()). This implies that P constitutes partial frame for the empirical 2 norm with the corresponding frame upper bound. More specifically, for f IR - R, decompose
= f=Σ cp P fo. fu E (span +)+,
PEP
for coefficient sequence (cp)per. It follows that
|f > |lf follz-
Σ DDEE
(p.' =IIG
and
(f.p))
Σ PEP
||Gpc|22 (Gax
PEP
P'EP
Denoting = dim O we can therefore bound Var,,(B; P) as
Var, (8; P) = inf EpEP [IVLn, [u] PERO
(Folx)) - p(xx)) Vooaa) 1 Σ fola Vo(a')VVo(2>")))
< Eper
D
= EpeP
Σ
(Ve
V f
==1
PEP 1=
15

---

Published as a conference paper at ICLR 2025
D
mmx (Gp)
-
Vo, Selli
P
==1
4d
S
2
sup 0.
n
Now by Chebyshev's inequality, for any 0 it holds that
Var, (0; P)
Pr (IVV E VL m (0)1l > E
Constructing the oracle. As Shamir (2018), we define the e-approximate oracle V as
Tmn.p (0) = EpEP V (0)
VLn.p(0) Eprep S otherwise.
VL (f)]l S E,
(0)
By union bounding, we see that during steps the oracle always defaults to the mean gradient and does not reveal any information the true parity p. with probability at least
27 E²
4d
Pr(Q) 2 1
sup 0,2
fo()))?
PP
where Q C P denotes the corresponding subset of the hypothesis space. Note that the argument can be extended to any randomized algorithm and random initialization in a straightforward manner by lifting to the product probability space, and consider Q to be fixed. Then for any target parity p E Q. the output ((A) of the algorithm after steps does not depend on p, so the predictor f= So(A) also fixed.
Lower bounding the loss. We first remark that simpler proof can be given for the sup norm error, which enough to establish separation. Consider arbitrary p, p' P with p # p' and let E {+1]d be such that p(x)= F (xx) then
(p(x) - f()) lp'(z) f(z) 2 f(=) +1 f(x))
f(x) 2
Now let a Q + Q be any automorphism of Q with no fixed points. The Loo error can be bounded below by restricting to the noninformative set Q follows.
E EpEp
sup/p(z) [p(x) - fmaa)(a)s|
2 EpEP [IvEe) sup (p(x)(()))))))())) - f(x))
x
1
2 Pl sup (up)p(()> f(x)) + sup Ja p(x) f(xx))
1
2|Q Pr(Q).
P|
For mean squared error, we similarly restrict to Q so that
Epe P. [[v(æ)) -((()( > EyeP, [1ipe)) (p(x)) - f(=)))]] ((a))
Since the range of p is contained in [-1, 1]. the above loss will not increase when f is replaced by its clipped version f(x) (f(x) A 1) V (-1). Moreover, in Lemma (proved at the end of the section) show that (xpep(p(x)) < (id) holds with probability 1 e -Q(() over the sample space of æ. so that
Eper. [()) & A(d) ) [(()) + ((d) < 2(4)
and also
Eper, [()){))(())((=)
Eg EpeP. [x(x)f(x)] Epep, [())
< E (())
+ (1 - El [[(())]]
< e R(d) + (1 (1-Pr(Q))*
2Pr(Q) + Pr(Q),
Ex [f(x)?]
2
16

---

Published as a conference paper at ICLR 2025
Therefore may bound
EapP. (p((z) {
2 P.p (1E - -())))
= =Pr(Q) 2BpeP. 11 (pE0))()))(() + Pr(Q)) E [f(x)²
2 Pr(Q)) (1 - (1-Pr(Q)) (r))))2
Pr(Q) 2e 2(())
2 2 Pr(Q)
2e Q(d) -O(d)
4T
1
4d
2 1
sup 8.
2e" -2(4)
(5)
g²
P
n
where we have used the inequality 2 - The proof is completed by evaluating the following cases.
t) t) 2 1
- 2t, valid for t [0. 11-
(1) If n ((()) and T. || foll Oppoly(d)) the gradient variance is bounded as Var, (0; P) < (ar P)1/3, follows that Pr(Q)
R(d), By taking
((d)) and (5) yields the
lower bound
e -((()
(2) If n = I(d"), | fell (d ((d and T = O the gradient variance is bounded as Var, (0; P) - O( = O(d (dd22ma- and (5) yields the lower bound Odd"...
Lemma 8. If (d), it holds with probability at least 1 - e (l(d) over random sampling that
[ppe(p((x)] <e -((d)
Proof. Let denote the number of 1s in æ. By the Chernoff bound for the binomial distribution,
8d I
82d
Pr
m
2 1 2exp
2
6
for a constant 8 E 1) to be determined, we assume the above event throughout the proof. Moreover denoting the complement parity of = [4] p. it holds that p(æ) = 21 ... (a)) and |ppePlp( (p(=)))))) [p" so suffices to consider the case where 2k <
Without loss of generality, we may assume that æ (-1, ... -1,1, as (-1) pQ[m]] We bound the cardinality of the set P4 := {p E P p(x) = 1}. Each parity in P.
1) so that p(æ) is decided
can be determined by choosing 2j elements from [m]] and h 2j elements from [d] \ [m]. Denoting
by [E], the coefficient of operator of order j. we can evaluate
1m/2|
IP+1
Σ
m
d m
2ÿ
2j
j=0
[m/2]
1m/2.
m
m
[t]s-22( a)
Σ
[t]a[1 + 11)0-mmap m)d-ma2s
2jj
2j
j=0
[m/2]
(m)
1
[1 + 1) dm
Σ
22
Wall + 1) (1 +1) + (1 1)m))
((1+)))
2
1=0
1 1/
d
+ [4 12 (1 + st)d-2m 2m
1
2 k
1 (d)
sk
|k/2|
'mn) (d - 2mm
Σ (-1
+
2 k
2
{
k - 2j
1=0
17

---

Published as a conference paper at ICLR 2025
where m' A (d - m) and =1. It further follows that
P+ PP
1
1k/2
Σ
m
d
2m
2 (d//2)) ( [ód]
1/21
S P
2
P|
3
2j
1 - 2j
[k/2]
(
((1
(1a/2)
[ód]
d - [d/2] [od]
(K)
(1/21)
(1dd/2)
2
h - [k/2] 16d/2)
d (Ld/a))
d
d
k/4
= (()
4 [1/4]
-
K
Here, we have chosen 5 = ^ K 2d
(1) and used the inequality
(0) (01+ bbbbbubo
From this, conclude that
P
P+
PP+
E ppprpp(<))]]
se <(-( 2(d)
P
with probability eR((d)
B PROOF OF THEOREM 5
Proof of Lemma 4. For each d + 1 m < + k- - and 1 S j < m, the only component of f
depending on Wj.m is Im and
0f (z: W) @wi.m
22m
(2m))))
O1 g0,((um)
ao a0 (wm)
< |
Dij +
= -1o'lloo (1))((0mmmmm)) - oz((m))
Σ Da (wm)
axl
$ 1| (wm))((Wmmmm)))) (1-1 (1-))( + oy
Σ (2(wm)
afj
s 21|00 (w)
Hence it follows that
d Σ IVwfall 2
Σ
ay(wm) S <llo''1 1) O(d))
as desired.
We say that a parity Zn T. for < ji. .Jr S Sd+k is trivial it always equals or equivalently if its reduction to the independent bits .i, Wa cancel out mod 2. For example, the parity 117 in Figure I is trivial. Define Ir,m as the set of nontrivial index r-tules less than m:
Irom {(i, .. .r)) S j1, . S m 1, xi ... Ii # 1}.
In particular, 1,m = [m 1] since no single parity is trivial.
Lemma 9 (concentration of interaction terms). If each bit 2i, fo i [m], j E [d] i.i.d. generated from the uniform distribution on {=1}. for any p it holds with probability at least p that
(x æi S =
2 log
32d4
max
ISPS4 (Juen JeJ
n
7
p
18

---

Published as a conference paper at ICLR 2025
Proof. Each tuple (j1, .. fr) E Im computes specific nontrivial parity zy 2, for which the bits 2h ..."), = 1, are i.i.d Uni[((+1]) due to symmetry. By Hoeffding inequality we have that
Pr (1 ... +ل,رة ž A) - 2e 12
Moreover, y S (d + k 1 1) (2d 1)" that
1m +... + 1. 2 (2d 1) (2d 1)²² < (2d)*.
Therefore it follows by union bounding that
Pr max
K' , 22) S 332de -XXinn
1sr54. JojE
which implies the statement.
In particular, we take (( and p = exp(-dd/2) so that == O This will ensure that the informative gradient signals will dominate the irrelevant interaction terms.
We now proceed to the main proof of Theorem 5. The superscript (0) at initialization is omitted for simplicity. The loss can be written more explicitly as
L(W)
1
d+k-1
Σ
( - xm 12, 2m Σ
m
(um))
2n
m=d-1
==1
It is straightforward to verify for 1 S a < that
do (Wm)
(8ja - a (a(wm))aaaee = -0āa -o,(wm)))omwwp)
Fwl.m
and
22m
m-1
Σ (5ja - (wa (F(Wm)(I) km).
Dw..m
Then the gradient of with respect to each element Wjmm at initialization can be computed as
0L
(0)(2))
(W)
(6(2m) xmm iwijm
Jw..m
n
of( (2m)) (1(2m)) - x am & (2m); Bz 2m)
(6)
n
1
(30m: 202m, 2tj 2m)
n(m -1)
+
(-1n + c2?. 2c2m 2) - 2m)
(8)
n(m -
1
+ n(m 1)
(O(lēm 12), 2c2m It - -km)
1
n(m (0(2m) Em. 0(2mm11) 2m)
(10)
Computing interaction strengths. signal when j = c1 [m]], Cg[m], while the other terms can be bounded perturbations. Let é h [m] so that Im computes 2d -parity.
The term (7) will be shown to contain the dominating gradient
For term (7), substitute 2m m- E. va initialization to expand
1 (xmm 2m, D. - 2m)) n(m Σ (mm, Ba, xx)) 71.
Σ a.8
(m Ba, xg))
n(m 1)
19

---

Published as a conference paper at ICLR 2025
where the dummy indices a, B, ... taken to run over [m 1]. Let evaluate the third-order interaction terms (Imm, Ba Tz). If h [a] = f, Zm computes the parity of 2 independent bits from
Zd 50 2m 2 cannot be trivial, hence (m, a, () E 13.m and [ Za Tza) n. by Lemma 9. Similarly, hi]] = implies that (m, a, B) E T3.m- Suppose h[(x), h[3] I & - 1; unless the combined parity Ta will not contain enough independent bits to cancel
... ".....
h[a] = h[3] = L the 2 bits in Im again am,,,,3 13.m- Moreover if h[a] = h18]
Im=aata will be trivial if and only if {a, 3} = -ci| C2 (2(mm]). J. in which case (am, Lo. 2z) = n. Thus we have that
mmy Fa T3)) 2 n Σ
1
1
Σ
(Fm, Fa; 0z))
O(mm - 1))x).
n (m,a,8)E/s,,
Similarly, the contraction (mm: Da Fz) can be nontrivial only if plj] = m and only when is the other child node of mms so that
1)k)
pbj m, otherwise.
my Da. xz))
10(mm 1)m)
Since x O(d-1-e/4) and d < m S 2d we can therefore isolate the leading term of order (d as
1
Fms 2c2m Tj -2m)
n (m- 1)
2c
2c
(1 m) + +O(dk))= +
(2 + (dP)))
(m-
1)2
(m - 1)3
2c
1 (plsi=m O(d-2-e//4)
(m- 1)2
Next, for term (8), we expand
2c 1 (- 1n + cb 2căm: By 2m)
2c
20a
262
(2m xz}} +
(*
+
(s 23,)
(2
n
n
n2
n
7
The second-order terms can be computed
1 (2m xz) n2(m - 1)
(z)) + Σ
(D xz)
= m - +O(K),
π
a
1 (22) nmm - 1 Σ ( Da)
1
Σ axk
(D mp)
1 m
O(k).
We evaluate the fourth-order interaction terms by looking at when (,,8,7,0) & 14m can occur. Without loss of generality, suppose h[(x] < h[3] S h(7) < h88).
(i) If h[3]] < h[]] < h[o), the parities of Ta To, ZZ, I,, must combine without overlaps to cancel out Is, so it must hold that I., child of xy and Tas &g are the two children of the other child. This subtree is fully determined by the choice of the index 5 and one of its child nodes, so there
are most (d) trivial 4-tuples in this case.
(ii If h[3] = h[7] < hlo). it still must hold that h[7] = h[6] be children of æs; otherwise, the bits of as and the non-child node cannot be canceled out by the remaining nodes. Then either Fg = xyy or Fg = 2s and in both cases we see that
Morcover, both F3, Iyy must
22 cannot be trivial.
(ii))
If h[[]] < h[>] h[6]. J.j it must be that 7 = otherwise the bits of x.,.8 cannot be canceled out and B. so there (dd² trivial 4-tuples in this case. (iv) If h[6]] h/>] h66], it must again hold that two indices must be equal, and the remaining two indices must also be equal, so there are also ((d?) trivial 4-tuples.
by T. It follows that v.
Hence follows that
1
(2)h))
Σ 0,0,%
(Da, x.g Bays .28)
72
7(m-
1)4
20

---

Published as a conference paper at ICLR 2025
Σ n
Σ
(nxk)
n(m 1)a
n(m 1)a
(,,
[m - 1011 74.ml
|1.mml 1)4 (m-
O(k) O(d-² k.).
(m-1)4
Furthermore, suppose a, S, 8 are constrained to contain the index j. Then case (i) above counts O(1) nontrivial tuples, while case (i), while cases (iii).(iv) count at most (dd) tuples since there is only one free index to be determined. Hence we also have
1 (2m=z)) 23) n(mm- - 1)3
1
Σ 0,3,7
(Das 23,
((d) (m 1)3
+ 0(k) O(d-2 + +K).
n
Combining the above, we obtain that
2c
2c
(kk) O(d ((d22-214))
n(m 1) (-1n + cā 2c2m. Ij 2m)
+
(m 1)2
(m 1)²
m
For term (9), we note that <1=ml") = (2/)) = O(ndd= + nx) derived above. Then since each component of 2m D. 2m are contained in 1- 1), 22 2]. respectively, we have that
1 n(m - 1) (0(( (), 2c2m, I.j 2m) n(m - 1)
4c
O ( (+)) O(d Od--2-1/4) 2-1/4).
Finally for term (10), by the Cauchy-Schwarz inequality we have
1 () (3) 1mm, /
1/2
1/2
<
m,
22
m,
4
= (=m)
1/2
(zm) 1/2
1(nd-1111 .0(dd + nx:) Old-1-110)
and so we may bound
n(m - 1) (x(2m) - Im, (() x) - 2m) n(m 1)
O 0((12m13)) = O(d-2-1/8)
From (7)-(10) we conclude that
aL
2c
(W)
(m - 1)2 (pl]=m)) + O(d-2-e//8
@ws.m
and the same result applies the approximate gradient V "m.. L at initialization since the cutoff does not apply and each component of the noise bounded by O(d (d--8-)
Concentration of softmax scores. become
Taking = dctc/16 the updated weights W(1) = -1)
2ed2+e/16
(1)
(m - 1)2 [pb]=m) + O(d-r/10)
w j.m
In partieular, for each j + c1 [m], c2 [m] the softmax scores satisfy
( = JM j.m Eg w < (0(1) (1)) mmmli. S exp( cxp(-S((d<110))
As softmax scores must sum to 1, it holds that Te (w!!) (e (1) > exp( exp((2((ss/16) and moreover
[m] [c (uw)))
= C "mm... mm S (cxp(O(d=<110) - 0(d
so!1) (1)
[ea[m]] [ez(m][u0 (wo
21

---

Published as a conference paper at ICLR 2025
from the inequality e < 1 + (() for small / 0. By symmetry, (egim] (wm))) 1 1+ O1d)s/1)). By simple algebraic manipulation, we can conclude that
2 Old - Sei(m]] (w ) [m (w ( - 2 + Old
(1) That is, the updated attention layer 2m Σj oj um Jay has learned take the average of the two child nodes and ignore the remaining input tokens at each step.
(u(um
Evaluating the forward pass. Now to bound the updated prediction loss, we evaluate the error
|| 2m oo of each step of the forward pass for d + 1 S m S d + k 1. More precisely, define the increasing sequence
4 (1)
cm
max d<jsm
2j 00
Ed O.
Then
(1(1) [mm]
Teg |1.
2 ca(m)
S [m]m Cog[m]] S cm-1+
and for the intermediate values ((1) we have
2 (1)
+ x cgfmml
mm
(11) Tc + Feg[mm]
(1))
mm
ci
2
+ fm-
1
1
(2( + Ocl[m] (2llm)(200(1)
+
Dealm| (calm((((!11)
Em-1
2
2
2d expl - + 0 + cm-1
S 1/16 + 6m-1;
for some constant C1 > 0. Since behaves like quadratic near 0, +1, it follows that
|
Fm 00 (((()))) +b
Ccy [m]] + Fca[m]]
(C((iid'/16 + cm- 1)2
Em
for some constant C2 > 0 depending only on Φ). Then for sufficiently large d, by choosing C3 such that Cz(C1 + C3d-e/10) < Cz, for cm-1 S C/d follows that
<m S C2(Crd-e/16 + c/d= S Czd-e/8
thus = =(d--/88 inductively for all m. We conclude that | ylooo = =11 x 2d loc bounded for all inputs as (d-s//8)
C PROOF THEOREM 7
Proof of Lemma 6. For the iterative generation scheme (3), each Wy.m affects o,, well as all nodes &a on higher levels hía] > h[m] through 8..- We bound the contribution each term to the total gradient inductively with respect to the level. Define for each d m S d k- 1, S j I and ( < < v the quantity
Ii. Aw... m
Sjm,, = max a
We denote ||00 for brevity. Clearly Ej.m,e for l h hmm] and
Oim
Sj.m,hfm) =
< (or,(Wm))az, 2ml I (w)).
@w....
Moreover for any a with h[a] h him], we can bound by the chain rule
de
as Ow..m
0sg Dwcim
$ <10'(Ea)1
Σ
ao(wa)
S Krf.mmee-1
dajm)
1
22

---

Published as a conference paper at ICLR 2025
yielding the relation fj..e KGj.m,t-1. Iterating, obtain that Exm,e I 2nf-h(m]++ o Therefore we can bound the total gradient by the following.
d+k-1
de Σ
Σ
Vw
Σ
a=d+1
t= =1a=dd-1+1
de
d+k-1
m
Ot.
@wijm
a m m+ 1=1
d+h-1 m-1
(de de-1)55,m.e
m=d+1 1=1
d+k-1 m-1
Tr (Wm) )2
Σ
2" .4x
Ax:2 2r:22-2((m]++ 2h(m]+2
f=1
< 4
d4k-1
2"xx 2h(m]+2 +2 Σ
16
4k²
d+h-1
Σ
x
20(m]
R2
m=d+1
8=1
m=d-1
4k
4x
de
d 20=2)2244 20122+2
(2x²
)
22
R²
1=1
8=ml
4R²
(x² 2)(2x² 1) (2x3)) O logz x+1),
since 2" = k Odd).
We first provide a concentration bound for the augmented data, which we take to hold throughout the proof by conditioning on the high probability event.
Lemma 10 (concentration of augmented data). For n' poly(d), with probability 1 (((va) over random sampling of the augmented data u1. un it holds that [|u + 1 = 2 for all 1 S j I + -1 and
de
+
< (11
max 05550
0
de
Proof. The nodes Fdd ++,,". Tdd at each level compute independent parities, even though parities at different levels can be correlated. By Hoeffding's 's inequality and union bounding over coordinates, it follows that
de
2n'
2n' p
zz
2(de de-1) log
= 2
log
jmde
1+1
100
with probability at least p. Again union bounding, the above holds for all levels 0 S f S simultaneously with probability least
vp, so that
de
1
2n' log
2n' log
22((2+1) I
O(d 1/4-1/4)
d
d
P
p
j=1
C==0
for all if p = e v
In addition, the probability that u, = 1, for some j S d + k: is bounded
by 2d 2-n' Rid) otherwise, at least one entry is equal to 2.
Now to prove Theorem 7, we show by induction that with high probability, the weights can be written for constants Cj (1)
r[Chim]
h(ml S tt m,
(e)
w
j.m
x
j > d1jmi-1 m d,
(11)
otherwise.
Clearly (11) is satisfied when 0. Suppose (11) holds time t 1 for 1. 1.
23

---

Published as a conference paper at ICLR 2025
Evaluating the forward pass. We first evaluate the forward pass iteration of the transformer up to level h(m] I t; fixing 0 < C' < minggo Ce, holds that
(,(w!) ! -
(1) exp(wer m) + exp(w (ap] calm), ) + dhimm]1 S exp exp((CCd/1b)
when plj] A m and
- oel [m] (w!! Jeg (w0 - 2
2
For the augmented tokens, define the increasing per-level error sequence
(1) a; 2n =xx
6
6n
d<
We recursively bound as before up to this will simultaneously verify that the filter is not (0) + 1,, lloe 2 e, due Lemma 10.
applied for the first + levels since ||u.
For each state 2m (t) with h[m] l we have
(t)
xt 5 [m/]
+ c2
[m]
Σ oj (w + Oel(m) (w 2
+ (alm
+ 6e-1
2
< (2d - 2) exp( exp(-Cat<10 c
Since & behaves like quadratic near 0, +1, it follows that
S (2d 2) exp(eCdddp 1d6/16) tc--))
and we can inductively verify c S Scp(-cd=16) as well as || 2 lox S 2dexp(-Cdd/10 holds for all & I t for sufficiently large d.
On the other hand, for the forward pass for levels h[m] > & the softmax scores are uniform over dnímj-1 tokens; moreover, the filter will be applied to all tokens on level /. + 2 and higher. Indeed, the output of nodes on level h/m] = + reads
(()) i m
(02+(0))
2 +11)
1
(# + ... +
+11)) de
(o + .... + + O(ee), de
m
d,
1l I Od-11/1) 1/1) by Lemma 10. Then so that |mmm
(+()
(t)
i m + 1nll < C2lll S 1/2
so that 1/2 Eg. the filter zeroes out the output of each node on level 2. Then the intermediate states of nodes Vyn on level I + 2 read
+(t) di
de 21()))
+
d
d++1
which again activates the filter. Repeating this process for the remaining levels, we conclude that | ol < O(d-1/4) and so || + +1n+ < O(2-1/2) holds simultancously for all nodes h/m] (and all timesteps t for which (11) is valid).
Evaluating the updates.
Define 2!!)
Σe= ny that
de
Σ
211000 S <exp(-Cd/11))
00
de
j=l
24

---

Published as a conference paper at ICLR 2025
We proceed to evaluate the gradient of L at (11). For the weights Wymm with h[m² = + 1, by isolating the errors from the forward pass we have
0
1
(6(2m)
| 2m
2m 112
((=
mm
awy.. m
@wl.m
og ( (t)
(0(2/0) - Em. Ø' (3(3 2. =.0)
n
ndl ( (6(211))
Im (zx!!) ,, -x(")
( 4
+0
(1 + l| ac [o0) = 1l00
d2
(( 3. &l )z
- + ((exp(-Co' Cdt/16)).
-
nd,
Then the first term identical the initial gradient (6) analyzed the proof of Theorem 5 except for the differences in indices, and from the same computation we obtain the leading term:
1
2c,
[
xm =1m
dt (plil=m) + O(d 2-6/5)
Dw.,m
2n
which holds with probability 1 - exp( exp(-d=/2) if ((d under the same setting of Lemma 9,
For all other nodes on level + 1 or below, the output does not depend on the weight Wj.m so the gradient of the squared error with respect to Wj. is zero. Moreover, all nodes on level ( + 2 or above are zeroed out due to the filter and hence also has zero gradient. Then the oracle error absorbed into the second term and the update after time + 1 with learning rate fixed to = /2te/ 1/ noo reads
e+e/16
(1)
L(W()), U) 2cno
1 (pls|=m> + Odd
j.m
d him]-1
By choosing do such that none of the leading terms lands exactly on half-integer, we have that for sufficiently large d,
d2+e/16
(++1) w c1[m].m
w (1+1) 62 mmi.
=r
2cnoo
d2
w j.m
((+1) r(O(d=-/166) =
h[m]-1
if pli] # m. We verify that
d2te/16
CgO 2
S C. = 2cT
S 20no-
di_,
Also, the weights Wjmm such that h[m] 2 + 2 only affect the nodes that are zeroed out, so that the gradient is also bounded as Odd-</16) and wjm
(t+1) 0.
It remains to evaluate the gradient signal of weights wym with h(m] S I, which have already been updated in previous steps. Define the error
aæ
Ej.,, = max
15a5de
This is similar to the error control in the proof of Lemma 6, but we exploit the fact that parities up to
level t are solved to obtain a much tighter bound. Let us expand o'(t) = 2c (1-1) t) + O((1 - ))
near and 2c (-1 - t) + O(1 (1+ + t)²) near -1 for some positive constant c'. Recall
$ + (1) - = < 2deexp(Cdd//1)
hía] <
holds in the forward pass, so each component 20 is O(exp(-Cdr/1)))- to either of + It follows that 1 (x(( = O(exp(((xx///////
so we can bound
0a
Ejm =
< 2 22((0((2(()) (2(()) Mooo (w))) S O(oop((Caa///
Ow..m
25

---

Published as a conference paper at ICLR 2025
Moreover, for any level h[o]] L, h(mi < f the magnitude of the derivative of the output aa can be bounded as
de
<O(exp(-Cd//1 M..me-
@hwg..m
This implies that Ey,,,t = E,,m,r-1 on level h[o] = + holds that
hjmm, I O(exp(-Cdd/10)) Furthermore, for any
aal @w
0
| Do
(2)
$ (o0llosSs,mm. O(exp(-Cdd
8=1
Thus we have for all a with h[a] 11+1,
2n [
Va xaa1? 12)
((()
Bas
< 2
(exp(( -Cdd-
Dmej.m
Owi.m
@wij.
and the nodes on level + or above are zeroed out due to the filter. Hence the gradient signal is exponentially small and
j. L(W()) U) = O(dexp(-Ca</19) + O(d-2~e/8)
so that j.m = rlw j.m
O(d- = w This concludes the proof of (11).
Finally, after time the weights all levels have been updated, that repeating the analysis of the forward pass yields that
Dcest
Yresst 00
(t) d+ Datk
<e S exp(-CCds/10
as was to be shown.
D EXPERIMENTAL DETAILS
For the transformer architecture, the feedforward layer was fixed to the following piecewise quadratic link function:
4t2 8t - te 1,,,--11000 -055) 412 -1 4t + -t te 0.5, 1].
((t)
te -055, 0.5)
For CoT models, learning rates were fixed to n = 15, 50, 100 for l: 8, 32. For the direct model, the learning rate was scaled to 0.01n to ensure stability of training. For the self-consistency model, filtering was done through an equivalent weight-based filter, which checks any softmax score exceeds a threshold value, here set to 0.4. Moreover, we found that adding 10% fraction of the gradient from the prediction loss to that of the CoT loss resulted in more stable training.
Figure 5 shows CoT loss and prediction loss curves for k 8, 16, 32, extending Figure 4. The direct model fails to learn parity all cases, while CoT with teacher forcing always learns efficiently. For CoT with self-consistency, similar analysis as in Section 4 can be applied for A 8, 16 with two or three distinct learning stages. We also observe that the basic CoT model manages to fully solve the problem for l: 8 (two intermediate levels) but not for A: 16, 32 (three and four intermediate levels), indicating that assistance (teacher forcing or self-consistency checking) becomes necessary for more complex tasks.
26

---

Published as a conference paper at ICLR 2025
1.0
Target loss (o- (k=32)
CoT loss Ik- 32]
0.5
direct
0.0
Co1+forcing CoTsconsiseency
0.2
0.2
200
350
200
Training step
Training step
1.0
1.0
0.8
0.8
Prodicton loss (K=16)
COT loss (kx16)
0.6
W
0.4
0.4
0.0
0.0
30
100
150
200
250
30
100
150
200
250
Training
Training
0.8
0.8
Prodiction loss (k-88)
COT loss (1-5))
0.6
0.8
0.2
0.2
100
Training step
Training step
Figure 5: CoT loss (left) and prediction loss (right) curves for the four models when d = 64, k: = 32 (top), k = 16 (middle) and k = 8 (bottom). For the CoT+consistency model, dashed lines indicate when the filters each level are deactivated.
27


Reviews:
[{"review_id": "dAssu7up0z", "reviewer": "Reviewer_muTa", "summary": "The paper provides a theoretical foundations which explain the benefits of chain-of-though. The authors study a simple setup of k-parity problem with 1-layer Transformer and in this case provide a separation results for transformer trained without intermediate supervision and one trained with teacher forcing. They also provide positive results for the case where the transformer can generate chain-of-thought but is not explicitly supervised with the intermediate results. They show conditions which are required in this case (data augmentation and self-consistency checking).", "questions": "I did not came up with any useful and educated questions. I believe the paper is of high quality and will need to study it in more depth.", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I find the work to be of very high quality and very relevant to the current research in reasoning abilities of language model. The theoretical setup is in my opinion well chosen, very concise and easy to understand.  Even though this is a theoretical paper, the authors are well aware of the current research on the applied side and the questions studied reflect it. I believe it is a strong accept, but I'm giving accept because I'm not sure whether theoretical results of this kind should have spotlight talks at ICLR.", "weaknesses": "I was not able to judge how limiting is the setup described in the paper overall, but for example the special masking seems quite artificial but I understand that it is required for the theoretical results. It would be nice if the authors explicitly describe the limitations they are aware of.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lruHUjDbMd", "reviewer": "Reviewer_5ycE", "summary": "This paper provides a theoretical exploration of how transformers develop stepwise reasoning capabilities through recursive task decomposition, using the k-parity problem as a case study. Through formal mathematical proofs and asymptotic complexity analysis, the authors establish three key findings: without intermediate supervision, solving k-parity with finite-precision gradient-based algorithms requires exponentially many steps, extending prior impossibility results to finite-sample settings. In contrast, with teacher forcing—where ground-truth intermediate steps are provided—the model efficiently learns the task in a single gradient update, underscoring the value of explicit supervision for complex reasoning. Additionally, even without teacher forcing, the model achieves accurate reasoning by employing data augmentation to validate intermediate steps, implementing a form of self-consistency check that allows for logarithmic convergence. This research emphasizes theoretical analysis through mathematical proofs and complexity analysis, offering insights into how transformers can perform chain-of-thought reasoning and task decomposition, thereby aligning with recent empirical advances in complex reasoning tasks.", "questions": "1. For the k-parity problem, each layer’s intermediate state corresponds to a simple “2-parity” check, which is straightforward to define. In more complex tasks, however, defining such intermediate states may not be as intuitive. How could this approach be extended to real-world reasoning tasks where intermediate state decomposition is less clear?\n2. The paper uses data augmentation to validate intermediate states in the reasoning process. If applied to practical reasoning tasks, would this data augmentation method require significant adjustments to be effective?\n3. Given the proposed architecture, are there potential challenges that might arise in real-world training, particularly concerning stability and computational complexity?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper innovatively uses the k-parity problem to theoretically analyze how transformers develop stepwise reasoning capabilities.\n2. This paper introduces a novel hierarchical decomposition of the k-parity problem and designs a corresponding transformer architecture to handle this structure effectively.\n3. This paper proposes mechanisms like data augmentation and self-consistency checks, enabling transformers to perform CoT reasoning even in the absence of explicit intermediate supervision.\n4. The paper provides clear and logically rigorous analysis for the k-parity problem, offering strong, well-structured insights within this specific context.", "weaknesses": "1. The theoretical analysis is heavily focused on the k-parity problem, which, while illustrative, may not extend seamlessly to more complex or varied reasoning tasks. This limits the applicability of the findings to broader transformer applications.\n2. The paper lacks empirical experiments to support the theoretical conclusions, which could leave readers questioning the practical effectiveness of the proposed methods in real-world scenarios or on diverse datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DnteaumrHe", "reviewer": "Reviewer_anZP", "summary": "Authors showed that CoT here is beneficial to learn a very hard problem, i.e. bit subset parity. Authors trained a one-layer transformer with a fixed CoT and with a tree based approach to determine the subset and they showed that without CoT it's not possible to learn an algorithm to detect the right bit, _assuming_ data cannot be manipulated freely. They introduced a causal mask to allow less error propagation during CoT.", "questions": "Q1. Is there any particular reason why one-layer transformer has been chosen? Would the same results (or even stronger) apply to 2 layer transformer?\n\nQ2. Could we apply the same idea to less structured task which lack the hierarchical structure, hence the \"easier\" teaching forcing?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1. The study is very rigorous, sharing some lights on the importance of CoT on multi hop reasoning. \n\nS2. The introduction of the causal mask is indeed very simple to explain and very intuitive, but also very effective to limit the error compounding. I think this is the main contribution, given teaching forcing is not really useful on this problem.\n\nS3. This work extended previous related work on more realistic setting, offering a clearer picture of why CoT is indeed essential for the given problem.", "weaknesses": "W1. Main weakness is a lack of empirical validation. I would love to see just a trained model on various settings (like with and without teacher forcing, parity problem sizes with different $k$ and $d$ values) to show that without teaching forcing the model is indeed able to learn the task. \n\nW2. The claim on the conclusion about finetuning transformer to improve multi step reasoning seems really too strong. In particular, it would be nice to show that it's easy to produce CoT data for all reasoning tasks before claiming this. So I would rephrase that sentence or provide a more detailed discussion on why authors believe their results may extend to other reasoning tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)