How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias
Authors: Ruiquan Huang
Yingbin Liang
Jing Yang
Venue: ICML 2025 poster (2025)
Keywords: Transformers
Training dynamics
Implicit bias
language recognition
Primary area: theory->learning_theory
TL;DR: We analyze how a one-layer transformer learns to solve two regular language tasks: even pairs and parity check. The model can learn to solve even pairs directly and parity check through CoT.

Abstract: Language recognition tasks are fundamental in natural language processing (NLP) and have been widely used to benchmark the performance of large language models (LLMs). These tasks also play a crucial role in explaining the working mechanisms of transformers. In this work, we focus on two representative tasks in the category of regular language recognition, known as 'even pairs' and 'parity check', the aim of which is to determine whether the occurrences of certain subsequences in a given sequence are even. Our goal is to explore how a one-layer transformer, consisting of an attention layer followed by a linear layer, learns to solve these tasks by theoretically analyzing its training dynamics under gradient descent. 
While even pairs can be solved directly by a one-layer transformer, parity check need to be solved by integrating Chain-of-Thought (CoT), either into the inference stage of a transformer well-trained for the even pairs task, or into the training of a one-layer transformer. For both problems, our analysis shows that the joint training of attention and linear layers exhibits two distinct phases. In the first phase, the attention layer grows rapidly, mapping data sequences into separable vectors. In the second phase, the attention layer becomes stable, while the linear layer grows logarithmically and approaches in direction to a max-margin hyperplane that correctly separates the attention layer outputs into positive and negative samples, and the loss decreases at a rate of $O(1/t)$. Our experiments validate those theoretical results.

How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias
Ruiquan Huang'
Yingbin Liang? Jing Yang?
Abstract
1. Introduction
Language recognition tasks fundamental in natural language processing (NLP) and have been widely used to benchmark the performance of large language models (LLMs). These tasks also play a crucial role in explaining the working mech- anisms of transformers. In this work, we focus on two representative tasks in the category of regular language recognition, known as 'even pairs' and "parity check', the aim of which is to determine whether the occurrences of certain subsequences in a given sequence are even. Our goal is to ex- plore how one-layer transformer, consisting of an attention layer followed by a linear layer, learns to solve these tasks by theoretically analyzing its training dynamies under gradient descent. While even pairs can be solved directly by one-layer transformer, parity check need to be solved by integrating Chain-of- ohought (CoT), either into the inference stage a transformer well-trained for the even pairs task, or into the training of a one-layer transformer. For both problems, our analysis shows that the joint training of attention and linear layers exhibits two distinct phases. In the first phase, the attention layer grows rapidly, mapping data sequences into separable vectors. In the second phase, the attention layer becomes stable, while the linear layer grows logarithmi- cally and approaches in direction to a max-margin hyperplane that correctly separates the attention layer outputs into positive and negative samples, and the loss decreases at rate of O(/()). Our experiments validate those theoretical results.
Transformers (Vaswani al., 2017) have become founda- tional in modern machine learning, revolutionizing natural language processing (NLP) tasks such as language model- ing (Devlin et al., 2018), translation (Wang et al., 2019), and text generation (Radford et al., 2019). Among these, language recognition tasks are fundamental to NLP and are widely used to benchmark the empirical performance of large language models (LLMs) (Bhattamishra al., 2020; Deletang al., 2023). Beyond their practical applications, these tasks hold significant potential for uncovering the un- derlying working mechanism of transformers. A growing body of research has explored the expressiveness and learn- ability of transformers in these settings (Strobl et al., 2024; Hahn & Rofin, 2024; Chiang & Cholak, 2022; Merrill Sabharwal, 2023; Hahn, 2020). Despite this, there has been little effort to understand transformers' training dynamics in language recognition tasks.
In this work, we take the first step towards bridging this gap by focusing on two fundamental pattern recognition tasks in formal language recognition, known as 'even pairs' and 'parity check' problems, and explore how transformers can be trained learn these tasks from theoretical perspective. Specifically, the objective of the 'even pairs' problem is to determine whether the total number of specific subsequences in a binary sequence is even, and the objective of the 'parity check' problem is to determine whether the total occurrence of a single pattern is even. These tasks are particularly compelling for studying transformers because they require the model recognize parity constraints and capture global dependencies across long sequences, which are essential for real-world applications such syntax parsing and error detection in communication systems.
For the two problems of our interest, the even pairs prob- lem has not been studied before theoretically. The parity check problem has recently been studied Kim Suzuki (2024b); Wen et al. (2024), which characterized the train- ing dynamics of CoT for learning parity. However, Kim & Suzuki (2024b) analyzed the training of an attention layer only, leaving more general characterization of joint train- ing of feed-forward and attention layers yet to be studied. Wen et al. (2024) analyzed three iteration steps in training
Department of Electrical Engincering, The Pennsylvania State University. State College, PA 16801, USA 'Department of Electri- cal and Computer Engineering. The Ohio State University, Colum- bus, OH 40210, USA Department of Electrical and Computer Engineering, University of Virginia, Charlottesville, VA 22904, USA. Correspondence to: Jing Yang <yangjing@ Sraanging@virginia.coe virginia.edu>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Training Dynamics of Transformers
without establishing the convergence of the entire training process. Our goal is to develop more general training dynamics characterization, including joint training of the attention and linear feed-forward layers, the convergence rate of the loss functions, and the implicit bias of the train- ing parameters. Furthermore, while both the 'even pairs' and 'parity check' problems are classification tasks, existing theoretical studies on transformers in classification settings (Li al., 2023; Tarzanagh et al., 2023b;a; Vasudeva al., 2024; Deora et al., 2023; Yang et al., 2024a; Magen et al., 2024; Jiang et al., 2024; Sakamoto & Sato, 2024) have pri- marily focused on cases where class distinctions are based on identifiable features. In contrast, these language recog- nition tasks will pose unique challenges, which require the transformer to leverage its attention mechanism to uncover intricate dependencies inherent in data sequences. By ex- ploring these tasks, our work will offer new insights into the fundamental mechanism of transformers.
ing) using truncated CoT inference. Such surprising result is based the intricate connection between the even pairs problem and the parity check problem. For the second ap- proach, it trains a one-layer transformer with CoT under teacher forcing, where we further include the training loss of even pairs to stabilize the training process. We show that with a two-phase training process, similarly to that of the even pairs problem, gradient descent provably renders a one-layer transformer that can solve parity check via COT.
Third, we introduce a novel analytical technique for study- ing joint training of attention and linear layers. Specifically, we employ higher-order Taylor expansions to precisely an- alyze the coupling between gradients of two layers and its impact on parameter updates in Phase 1. Then we incorpo- rate implicit bias principles to further characterize the train- ing dynamics in Phase 2. Since all parameters are actively updated, we must bound the perturbations in the attention layer and analyze their effects on the linear layer. This is achieved by carefully designing the scaling factor in the attention mechanism, which not only stabilizes training but also underscores its critical role in transformer architectures.
In this work, we investigate how one-layer transformer, consisting of an attention layer followed by a linear layer, learns to perform the 'even pairs' and 'parity check" tasks. We will theoretically analyze the model dynamics during the training process of gradient descent, and examine how transformer parameters will be guided to converge to lution with implicit bias. Here, we will jointly analyze the training process of the attention layer and linear layer, which will be significantly different from most existing analysis of the training dynamics of transformers for classification problems, where joint training is not studied (Huang et al., 2024a; Tarzanagh et al., 2023b; Kim & Suzuki, 2024b; Li et al., 2024b).
Finally, we conduct experiments to validate our theoreti- cal findings, demonstrating consistent parameter growth, alignment behavior, and loss convergence.
2. Related Work
Due to the recent extensive theoretical studies of transform- ers from various perspectives, the following summary will mainly focus on the training dynamics characterization for transformers, which is highly relevant to this paper.
Our major contributions are four-fold:
Learning regular language recognition problems via transformers. Regular language recognition tasks are fun- damental to NLP and are widely used to benchmark the em- pirical performance of large language models (LLMs) (Bhat- tamishra et al., 2020; Deletang et al., 2023). Thcoreti- cal understanding of transformers for solving these tasks are mainly focusing on the expressiveness and learnabil- ity Strobl et al. (2024); Hahn & Rofin (2024); Chiang & Cholak (2022); Merrill & Sabharwal (2023); Hahn (2020). Among these studies, several negative results highlight the limitations of transformers in learning problems such as par- ity checking. Notably, Merrill & Sabharwal (2023) demon- strated that chain-of-thought (CoT) reasoning significantly enhances the expressive power of transformers. We refer readers to the comprehensive survey by Strobl et al. (2024) for a detailed discussion on the expressiveness and learn- ability of transformers. Regarding the study on the training transformers, Kim & Suzuki (2024b); Wen et al. (2024) showed that it is impossible to successfully learn the parity check problem by applying transformer once. They further developed CoT method and showed that attention model
First, for the even pairs problem, we identify two distinct learning phases. In Phase 1, both linear and attention layers grow rapidly, inducing separable outputs of the attention layer. In Phase 2, the attention layer remains almost un- changed, while the dynamics of the linear layer is governed by an implicit bias, which converges in direction to the max-margin hyperplane that correctly separates the atten- tion layer's s outputs into positive and negative samples. We also show that the loss function decays to the global min- imum sublinearly in time. To the best of our knowledge, this is the first theoretical study on the training dynamics of transformers for the even pairs problem.
Second, we innovatively leverage the insights from the even pairs problem and Chain-of- -hhhught (CoT) to solve the par- ity check problem through two different approaches. In the first approach, we introduce truncated CoT into the infer- ence stage of a trained transformer. We demonstrate that a transformer, well-trained on the even pairs problem but with- out CoT training, can successfully solve the parity check problem in zero-shot manner (without any additional train-
2

---

Training Dynamics of Transformers
with CoT can be trained to provably learn parity check. Dif- ferently from Kim & Suzuki (2024b), our work analyzed the softmax attention model joindy trained with a linear layer with CoT for learning parity check. Our CoT training design is also different from that in Kim & Suzuki (2024b). Further, Wen et al. (2024) studied the sample complexity of training a one-layer transformer for learning complex parity problems. They analyzed three steps in the training procedure without establishing the convergence of the entire training process, which is one of our focuses here.
Yang et al. (2024b), next-token prediction (NTP) in Tian et al. (2023a:b); Li et al. (2024b); Huang et al. (2024a); Thrampoulidis (2024), unsupervised learning in Huang et al. (2024b), regression problem in Roix-Adsera et al. (2023), etc. Those studies for the classification problem and CoT training have been discussed above.
Implicit bias. Our analysis of the convergence guarantee develops implicit bias of gradient descent for transformers. Such characterization has been previously established in Soudry et al. (2018) 8); Nacson et al. (2019); Ji & Telgarsky (2021); Ji et al. (2021) for gradient busscnttbbasee optimiza- tion and in Phuong & Lampert (2020); Frei et al. (2022); Kou et al. (2024) for training ReLU/Leaky-Rer UeUU networks on orthogonal data. More studies along this line can be found in comprehensive survey Vardi (2023). Most rele- vant to our study are the recent works Huang et al. (2024a) and Tarzanagh et al. (2023b;a); Sheen et al. (2024), which established implicit bias for training transformers for next- token prediction and classification problems, respectively. Differently from those work on transformers, our study here focuses on the even pairs and parity check problems, which have unique structures not captured in those work.
Training dynamies of transformers with CoT. CoT is a powerful technique that enables transformers to solve more complex tasks by breaking down problem-solvingg into intermediate reasoning steps (Wei et al., 2022; Kojima et al., 2022). Recently, training dynamics of transformers with CoT has been studied in Kim & Suzuki (2024b); Wen al. (2024) for the parity problems (as discussed above) and in Li et al. (2024a) for in-context supervised learning.
Training dynamics of transformers for classification problems. A recent active line of research has focused on studying the training dynamics of transformers for clas- sification problems. Tarzanagh et al. (2023b;a) showed that the training dynamics of an attention layer for classifi- cation problem is equivalent to support vector machine problem. Further, Vasudeva et al. (2024) established the convergence rate for the training of a one-layer transformer for a classification problem. Li et al. (2023) characterized the training dynamics of vision transformers and provide converging upper bound on the generalization error. Deora et al. (2023) studied the training and generalization error under the neural tangent kernel (NTK) regime. Yang et al. (2024a) characterized the training dynamics of gradient flow for a word co-occurrence recognition problem. Magen et al. (2024); Jiang et al. (2024); Sakamoto & Sato (2024) studied the benign overfitting of transformers in learning classifica- tion tasks. Although the problems of our interest here (i.c., even pairs and parity check) generally fall into the classifica- tion problem, these language recognition tasks pose unique challenges that require transformer to leverage its attention mechanism to uncover intricate dependencies inherent in data sequences, which have not been addressed in previous studies of the conventional classification problems.
3. Problem Formulation
Notations. All vectors considered in this paper are column vectors. For a matrix W, |[W]] represents its Frobenious norm. For a vector c, we use [o]; to denote the i-th coor- dinate of v. We use {eiheedd to denote the canonical basis of Rd. We use Φ(v) to denote the softmax function, i.c., [o(v)]; = exp(v:)/ Σ,, exp(e, v). which can be applied to any vector with arbitrary dimension. The inner product (A, B) of two matrices or vectors A, B equals Trace Trace(AB (ABT) ). For a set X, we use X to denote the Carte- sian product of n copies of X, and 2 = 2
In this work, we consider pattern recognition tasks in the context of formal language recognition, which challenges the ability of machine learning models such as transformers to recognize patterns over long sequences.
First, introduce the general pattern recognition task in binary sequence. The set of all binary sequences is denoted by {a, b) Skao, where Lmax is the maximum length of the sequence. The pattern of interest is a set of sequences P = (p,,..., . C (a, S.max Given a sequence X and a pattern P, let Np be the number of total matching subsequences in X, where a subsequence of X matches if it equals some p; E P. The pattern recognition task is to determine whether Np satisfies a predefined condition, such as whether Np is even.
Training dynamics of transformers for other problems. Due to the rapidly increasing studies in this area, we include only some example papers in each of the following topics. In order to understand the working mechanism of transform- ers, training dynamics has been intensively investigated for various machine learning problems, for example, in- context learning problems in Ahn et al. (2024); Mahankali al. (2023); Zhang et al. (2023); Huang et al. (2023); Cui et al. (2024); Cheng et al. (2023); Kim & Suzuki (2024a); Nichani et al. (2024); Chen et al. (2024); Chen & Li (2024);
In the following, we describe two representative tasks that are particularly interesting in the study of regular language recognition due to their simple formulation and inherent
3

---

Training Dynamics of Transformers
learning challenges. More details about formal language recognition (which includes regular and non-regular lan- guage recognition) can be found in Deletang et al. (2023).
label I means the input sequence contains even pairs, and 1 otherwise. Similarly, for the parity check task, positive predicted label 1 means the input sequence contains even number of the pattern of interest, and -1 otherwise.
Even pairs. The pattern of interest is P = {ab, baj. For example, in the sequence X = aabba, there is one ab and one ba, resulting in total of two, which is even.
Learning objective. We adopt the logistic loss for these binary classification tasks. We denote training dataset, where It denotes the set of all length- se- quences. An individual training data consists of sequence X(") = [x] ........" 2L_ and a label yn E {1, -1]. where n is the index of the data. With slight abuse of notation, we also use the notation n E Iy to indicate that X((*) has length L. Then, the loss function can be expressed as:
-.y-nyy Ir as the
= Parity check. The pattern of interest is P= {b}. In other words, this task is simply to determine whether the number of bs in a binary string X E b) is even or odd.
Although even pairs may appear to involve more complex pattern recognition than parity check, it can be shown that this task is equivalent to determining whether the first and last tokens of the sequence are equal, which can be solved in (1) time (Deletang et al., 2023). However, parity check usually requires O(L) time complexity given an L-length sequence since we need to check every token at least once.
[(u,
= -Eter 1 Enes, log (1 + exp (-ynToo To TEl
The use of 1/14g| in the loss function ensures that sequences of different lengths contribute equally to training.
In this work, our aim is to apply transformer models to solve these problems while leveraging these tasks to understand the underlying working mechanisms of transformer models.
Our goal is to train the transformer to minimize the loss function, i.e., to solve the problem ming-(1,W) [(u, W). We adopt 2-phase gradient descent (GD) to mínimize the loss function L(u, W) as follows. We adopt zero initializa- tion, i.e., Gu = 0. Then at early steps t tu (to be specified later), we update 0 = (u, W) as follows:
Embedding strategy. We employ the following embedding strategy for each input binary sequence. For a token a at position É, its embedding is given by E? E2c-1, and for a token b at position L, its embedding is given by E = e2c. Such an embedding ensures that token embeddings are orthogonal to each other, which is a widely adopted condition for transformers.
u++1 = 4 V W411 = We - V Ce,
where is the learning rate, and La is the abbreviation of C(u,, Wt).
One-layer transformer. We consider a one-layer trans- former, denoted To Ledxx R. which takes sequence of token embeddings input and outputs scaler value, and 0 includes all trainable parameters of the trans- former. Specifically, let X = (II...., x1.] denote the input sequence, where xe E IRd for each I l < L and L is the length of the sequence. Let the value, key and query matrices be denoted by W., We, Wg e Rdxd Then, the attention layer is given by W,, E/ Kepr, where Pe = [Ø(X W Wgzz//Ale and À E IR is a scaling pa- rameter. Hence, the one-layer transformer has the form Tx) W. E/_ Zefet, where Wa Rixd de- notes the linear feed-forward layer. For simplicity, we reparameterize W.W, as u and W as W, as com- monly taken in Huang et al. (2023); Tian et al. (2023a); Li et al. (2024b). Hence, the transformer is reformulated as
After step to. we update I as follows:
Me+1 = -4 V/
W411 W.
We remark that such learning rate schedule can be viewed as an approximation of GD with decaying learning rate or Adam (Kingma, 2014). To be more specific, Adam updates the parameter in the form of 0z-1 = Or - rpm// Vv, + where m, ~ 9t. Ue R g,, e is a positive constant, and ge is the gradient. Thus, during the early training steps, Adam's update behaves closely to the high learning rate regime in GD. In the subsequent training steps when the gradient is relatively small, Adam behaves similar to vanilla GD due to the constant e in the denominator / U/g + E.
(X) = u Σte Zeffe, where Sr = [Ø(X" T and all trainable parameters are captured in = (u, W). We will call W and respectively as attention and linear layer parameters.
W ///////) A)]e.
4. Even Pairs Problem
In this section, we characterize the training dynamics of a one-layer transformer for learning the even pairs task. For simplicity, we choose IT. = b)
When the transformer is used to for the regular language recognition tasks, for given input X, it will take the sign (i.e., 1 or -1) of the transformer output (X) as the pre- dicted label. For the even pairs task, a positive predicted
Key challenges. The key challenge in this analysis arises from the joint training of the linear and attention layers. The intertwined updates of these layers create a coupled stochas- tie process, complicating the analysis of parameter evolution.
4

---

Training Dynamics of Transformers
Furthermore, since every token contributes both positive and negative samples, this leads to gradient cancellation during analysis, making the analysis more difficult.
the first token score. This is because sequences of length L 1 (which always have positive labels) dominate the early training, and create an initial bias for the first token. (b) The first two equations of part (2) indicate that the attention weight of the first token increases in positive samples (where last token by, and first token ET share the same value w), and is suppressed in negative samples. In other words, the attention layer allocates more weights on non-leading tokens (with E 2 2) in negative samples. Consequently, the transformer output of the negative samples relies more on the token scores at non-leading positions. In order to minimize the loss, those token scores become increasingly negative, as shown in the second equation of part (1). (c) The last equation of part (1) indicates that the token score of the second token decreases faster than other non-leading tokens, as appears more frequently across samples. This rapid token score decrease drives the attention layer to allocate more attention weight on it over other non-leading tokens to reduce the loss, as shown in the last equation in part (2).
To address the above challenges, we employ higher-order Taylor expansions to precisely analyze the coupling between gradients of two layers and its impact on parameter updates. We next present our results for these two phases and explain the insights that these results imply. Before we proceed, we introduce the following concepts.
Token score. We define the token score of Eig,, w E {a, b) as (ue, by",, which quantifies the alignment between the token embedding and the linear layer.
Attention score. For a given attention layer W,, we define the (raw) attention score of E as (EY, W, W, Ey. ), w, w/ E (a. For token zg, its attention weight is given by PP = Waz))e which is proportional to the exponen- tial of its attention score. Importantly, it is the differences between the attention scores of different tokens that govern the attention weight distribution.
Due to the fact that attention weight is determined by the differences between attention scores, Theorem 4.1 1 suggests that the end of Phase the attention layer focuses on the first token (with şp (n.to) > I/L. for
Relationship with the transformer output. Consequently, the transformer output T(X) for an input X is weighted average of the token scores of ze, with the attention weights be serving as the weighting factors.
in pos- itive samples, and on the second token (with sp ((m.to)) > L for length-L samples) in negative samples. Hence, the atten- tion layer maps data samples to satisfy a separable property, To be more specific, we first introduce the definition of separable data.
4.1. Phase 1: Rapid Growth of Token and Attention Scores
In Phase 1 of training, the linear and attention layers exhibit mutually reinforcing dynamics, with the featuring dynamics captured in the following theorem.
Definition 4.2 (Separable dataset). A dataset of d- dimensional vectors and their labels {(v(") Bn))y separable if there exists u E Rd succ that
Theorem 4.1 (Phase 1). Let ww denote the flip of token
(u, (n))) > 00 vigngN.
w E {a. b). Choose x= = a. 1o = m() and 1j = O(min(1/Imama Lmmxx 1/)2/31). Then, for all t S to.
Intuitively speaking, a dataset is separable indicates that there exists a hyperplane that can correctly separate the positive and negative samples into two half-spaces.
the parameters evolve as follows:
(I) The dynamics of linear layer is governed by the fol- lowing inequalities.
At the end of Phase 1, the attention layer Wt, maps data samples to E/_ 2 (n)
'Pt (n.fo) It can be shown that if the linear layer has parameter u = Ej + Ej - Eg - Eg. then, the predicted label of the data sample, which is the sign of (n.to) matches with the ground-truth label 'Pi
(we; ETr = @nt)), Vw,
(ue; ETr) = (172tt),
vl 2 2, Yi.
(we, ET - E S - Ve 2 3, Vw.
u Ele ae (1) "pe
x (m.to)) > O for all n. Thus, we In. i.c.. (u, 3r E/ It
(2) The dynamics of attention layer W is governed by the following inequalities. For any lenghh L S Lmaax we have
have the following proposition.
Proposition 4.3. Let v(n) = = I1 4i (inda) with label Yn- Then, at the end of phase 1, the dataset {(vlm), yw) } is separable by u = F2 + F.p- Fng FEg.
(n))
(ET - Ey', ,,EE)) 2 2)722). (ET'' - ",,) 2 2))2t),
VE 2 2, Vw, w'
V8 2 2, Vw, w/
(Ex'- EF"" W,EY) Z 1()tt),
AE Z 3, Vw, w/, w"
We note that at the end of Phase 1, the linear layer is trained to be uy which may not necessarily separate the attention layer's 's outputs. In fact, the continual training into Phase 2 will further update the linear layer, so that the attention layer's outputs can be separated by it.
Theorem 4.1 characterizes the following featuring dynamics of the linear and attention layers in Phase 1. (a) The first equation of part (1) indicates that there is a rapid growth of
5

---

Training Dynamics of Transformers
4.2. Phase 2: Margin Maximization and Implicit Bias
been shown in Pérez al. (2021) that it is impossible to recognize parity by applying transformer once. However, it has recently been shown in Kim & Suzuki (2024b) that chain-of siiin-of-the of-thought (CoT) can serve as an advanced approach to solving such a task.
In this phase, the transformer shifts focus from rapid feature alignment to margin maximization, driven by the implicit bias of gradient descent.
Since the outputs of the attention layer at the end of Phase become separable, we can define the max-margin solution of the corresponding separating hyperplane as follows.
In this section, we provide two new approaches to solv- ing the parity check problem, by integrating CoT with the solution for the even pairs problem studied in Section 4. The first approach solves parity check by applying the one- layer transformer well-trained for even pairs via a truncated CoT-type inference, which does not require any additional training. The second approach trains a one-layer transformer with CoT under teacher forcing. where GD provably renders transformer that solves the parity check problem.
"gr= = arg min m
s.t. No yu Ets xb oe slabto)) 21, T.
The following theorem characterizes the training dynamics of Phase 2, which shows that the linear layer will converge to the above max-margin solution in direction.
5.1. Approach 1: Inference via Truncated CoT
Theorem 4.4 (Phase 2). There exists a constant t2 = (1(1) and T = O(x2/3/31))) 0 }, such that for L2 S L1 < T, we have ||uel| 2 (2oog t), and
Inspired by the 2-state machine, we show that by simply tak- ing inference via truncated CoT, the one-layer transformer well-trained for the even pairs problem can solve parity check efficiently without additional training.
Me |Tuell
BEr > 2 21-2 2 u2pll
1
)
To formalize this, we first outline the method to solve parity check through the lens of a 2-state finite automaton. Recall that the parity check problem is to determine whether the number of bs in a binary string X E ba, even or odd. Given an X = w/ W2.......0 w',, the automaton initializes its state $1 to wy and updates 9t for & > 2 sequentially as follows. At each step 2 2, the state transits to $1+1 if 8t = W4+1, or 81+1 = b if 8t * Wt+1- For example, for X = abb, the state transitions are = + 82 b (as 81 + w2 b) + 83 = a (as $2 = w3 = b). The character a of the last state indicates that the sequence takes even parity (we equate with token a and 1 with token b). It is worth noting that the core step of parity check involves comparing two characters each time, which is also performed in even pairs except that the characters to be checked are fixed to be the first and last tokens.
Moreover, we have w|,- 1i S O(1.
Theorem 4.4 characterizes the following dynamics of the linear and attention layers in Phase 2. (a) The norm ||u, of the linear layer continues to grow logarithmically to increase the classification margin. (b) The updates of the attention layer We is negligible due to the scaling factor X, as indi- cated by I|W, W.,1 || I O1)). Hence, attention patterns at the end of Phase persist, i.e., in both positive and negative samples, the attention scores of the first and second tokens still dominate, respectively. (c) The linear layer u enters regime governed by implicit bias, which converges to the max-margin solution for separating the attention layer's out- puts. This dynamics is also observed in an empirical work (Merrill et al., 2021).
Theorem 4.5 (Convergence of loss). For t S T, we have L. 0 Lmox ll
Inspired by this observation, we propose the following in- ference method via truncated CoT to solve parity check.
nvt
1/2
Inference via truncated COT. Recall that the even pairs problem is equivalent to labeling whether the first and the last tokens are the same. Hence, the transformer trained for even pairs can provide correct labels for such task during the inference. We thus propose truncated CoT that leverages the label predicted by the trained transformer for even pairs iteratively to obtain the answer for parity check. Such an inference process runs as follows, with the pseu- docode provided in Algorithm 1. Given a binary sequence
Theorem 4.5 indicates that the loss converges to O
X373
Therefore, as long as the scaling factor À = 2 (022)) ). the loss can achieve arbitrarily small value e. The full proof in this section can be found in Appendix B.
In summary, the trained transformer utilizes its attention to decide if two tokens are equal and the linear layer increases the classification margin and enable fast loss decay.
X = W1 W' each iteration E L-....,-111 1) of CoT, it performs the following steps: (1) check whether the first and the last tokens of X are equal by applying the one-layer transformer trained for even pairs; (2) append the predicted
5. Parity Check Problem
The parity check problem is generally considered to be more difficult than the even pairs problem. For instance, it has
6

---

Training Dynamics of Transformers
Algorithm 1 Truncated CoTT
use the set It of length- sequences (where I, Lo + 1), and for each sequence X TL, set its label y to be 1 if the last token matches the token at position -- Lo + 1 t, and -1 otherwise. The superscript 'p' in yn indicates that the label constructed for parity check. Thus, the train- ing CoT will use data sequences with length where Lo S L S 20g and the total loss given by
Input: Binary sequence X w1 ... WL for t E {1, .t....... ., L 1 do
Predict y/2 (X)
Let WL++
if yt = 1 and WL+t = if 1.
Update X = W1+1 WL+
end for Output 91-1-
LCor(u, W) = = 1 (iz] EnEI, log (1 + cxp (-
label y = a or y = b to the end of the sequence X (we equate with token a and -1 with token b); (3) remove the first token in X to maintain the sequence length. Then after L- iterations, the final prediction provides the parity of the original input sequence.
where we adopt the one-layer transformer described in Sec- tion with the dimension 2Lu- 2.
Furthermore,
observe in experiments that if we directly train transformers over the above CoT loss Lcor (u, W), the gradient vanishes. Interestingly, initializ- ing transformer by that trained for even pairs helps to avoid such a case. Motivated by such an observation, we introduce the even pairs loss to regularize the training process. this end, we include data sequences with length L Lo for the even pairs loss, and label those sequences by their even pairs labels. Namely, for any sequence X (m), where E l7 with L < Lo, set the label yf to be if the last token matches the first token, and -1 otherwise. The superscript "e" in vi indicates that the label is constructed for even pairs. These data sequences provide a regularization loss given by
5.2. Approach 2: Training with CoT under Teacher Forcing
In this section, train one-layer transformer with the full version of Chain-of-Thought Thought (CoT) to solve the parity check problem. Unlike the truncated CoT described Section 5.1, we keep the original sequence and only append the predicted label to the end the sequence at each iteration.
To explain how we train one-layer transformer execute CoT learning of parity, first describe how the parity of sequence can be obtained via CoT step by step.
For given sequence X = W1 ... WLo with length Lo. we generate CoT inputs X Xla and their correspond- ing labels to learn the parity as follows. First, let X X. For each t {l..... Lo 1), take Xt and compare its to- kens w WLo+t11 and wy. If they are the same, let WLo+t a; otherwise, let WLo+ b. Then set LLo+ to be the la- bel of X and append wh to the end of X to obtain X++1 = W1 ... WLo+t for the next step of COT. Finally, the label W25.0 is the parity of X (See Figure 1).
CReg(u, W) = -to-l 1 [T1] Ener, log (1 + exp ( (-V,To(X(
The role of the regularization loss to initially guide the linear layer to rapidly increase along the direction of solving even pairs problems, which will also initialize the parame- ters learn parity check in stable way. We note that such regularization is also equivalent to data mixing technique.
Hence, the total training loss for parity check is given by
2-step CoT
C = Lcor (u, W) + CReg W).
We minimize the above loss function by gradient descent as described in Section 3 for training the parity check problem. In the following, we choose I = b) for simplicity.
input
b
output
5.3. Training Dynamics of Approach 2
Figure 1. An illustration of CoT input abb, where 3.
The training process of CoT under teacher forcing can also be divided into two training phases. Below, we present the theoretical characterization of those two phases.
Training design. To train transformer learn each step 1 IS < Lo of CoT, our dataset should be labeled to compare the token WLa+t-1 (last token of the input se- quence of the current CoT step) and w,, and such a label is used to train (supervise) the transformer to conduct the f-th step of CoT correctly. Thus, to trann the 1-th step of CoT, we
Theorem 5.1 (Phase 1). Let -w denote the flip of token E {a, b). Choose X = l(LL2 and to = 0(1 ((1l((/ )).
Then, for all t & ty. the parameters evolve as follows:
(1) The dynamics of linear layer governed by the fol-
7

---

Training Dynamics of Transformers
lowing inequalities.
All Set Diff
Token Ser
Loss
(we; FT) = (nt)), vw;
(ue; ET) = (() vl 2 2, Ve,
300 600 900 Step
0
300 600 900
300 600 900 Step
(ut. ET Eiii) - ve 2 3, Vw.
Step
Figure 2. Results of one-layer transformer on even pairs. From the left to the right: (1) Loss decay over training. (2) Token scores at first three positions. (3) Attention scores length-3 sequences.
(2) The dynamics of the attention layer W is governed by the following inequalities. For any length I. < Lo, we have
(EY - E, W.EEY) 2 (17(t). (EE' - -ET", ,,, 2 (),
AA 2 2, Vw, w',
ws
Alta Scr Diff
x CD Wed
Token Scr
2
3A 2 2, Vw, w', Ae 2 3, Vue, w', w".
3
(EY' - EYY", ,,EEE) 2 ())
300 600 900 Step
300 600 900 Step
300 600 900 Step
For length 2 Lo, let Eo = Lo + 1, and we have (F%O Fy' W,FF)) 2 21121))
ve + la, Vwu, w', ve * lo, Vw, w', "Y",,,ETT W.ET) 2 1)1 V f 1, lo Vw, w', w".
Figure 3. Results of one-layer transformer on parity check. From the left to the right: (1) Loss decay over training. (2) Token scores at first three positions. (3) Attention scores in length-5 sequences.
(EE Ero" W,ET) 2 2(12(1) (ET' -ET",
have ||u,|| 2 Blog t). and
Since the loss function Equation (1) is regularized by the loss of even pairs, the linear layer exhibits the same dynamics as the even pairs problem in Theorem 4.1. The key difference between CoT training in Theorem 5.1 and even pairs training in Theorem 4.1 lies in attention dynamics on sequences with length L, > Lo, which are labeled for CoT training of parity check. In particular, the last three inequalities Theorem 5.1 suggests that it is the token at position L Lo + 1 that differs most from other tokens.
Mi Twell
1-11 [illbbrr
Moreover, we have || - Weoll S (1).
The above theorem also implies that the total loss converges sublinearly, which further implies that both the CoT and regularization losses enjoy the same decay rate.
Theorem 5.3. For t S T. we have L Pariity
At the end of Phase 1, the outputs of the attention layer are also separable. Namely, there exists linear classifier that provides correct labels for all CoT steps. i.c., labels all training sequences with length Lo < L < 2Lg correctly. As a by-product, such a linear classifier also provides correct even pairs labels for the sequences with L < Lo. For those separable data sequences, we define thee max- max-margin solution for the separating hyperplane as
O ( Lemaa 11u
(Amazllltzz
6. Experiments
In this section, we provide experiments on synthetic datasets to verify our thcoretical findings. Specifically, we choose Lmax = 6, Lo = 4. Then, we train u and W by gradient descent with step size 7 0.1. We choose to = 100, and X = 2. As observed in Figures 2 and 3, the first plot in each figure shows the rapid decay of the loss to the global minimum. The second plot shows the dynamics of token scores at the first three positions, where the first token score grows (blue curve) and the second and third token scores de- crease (green and orange curves) during training. The third plot illustrates dynamics of the attention weights, where the first token receives more attention in positive samples (blue curve) and less attention in negative samples (orange curve). These plots validate our theoretical findings on dy- namics of token scores and attention scores characterized in Theorem 4.1 and Theorem 5.1. Furthermore, the ver- tical blue dashed line in these plots indicates the end of first phase. The vertical orange dashed line in these plots
"Ccor arg min ||u/].
s.t. (u. Ela xe pp (m.to) 21, I,
((X(m))
(n)
where ((X( denotes the length of (m), Note that ucor slightly abuses notation the dataset also includes se- quences with lengths L < Lo for even pairs.
The following theorem shows that the training enters Phase 2 we continue to update the parameters by gradient descent, during which the attention layer has negligible change, but the linear layer converges to the max-man -argin solution u'or-
Theorem 5.2 (Phase 2). There exists a constant t/2 = 2(1) such that for 1/ S St5 5 T, we
and T =
8

---

Training Dynamics of Transformers
indicates l2 600 in Theorem 4.4 and Theorem 5.2, where |u,ll starts to grow logarithmically (second plot of both figures).
Boix-Adsera, E., Litwwin E., Abbe, E., Bengio, S., and Susskind, J. Transformers learn through gradual rank increase. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Additional Experiments. In Appendix D. we conduct more experiments on different configurations of scaling parameter and the two-phase learning dynamics.
Chen, S. and Li, Y. Provably learning multi-head attention layer. arXiv preprint arXiv:2402.0444 04004, 2024.
All experiments are conducted on PC equipped with an i5- 15-12400F processor and 16GB of memory.
Chen, S., Sheen, H., Wang. T., and Yang, Z. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. /9442, 2024.
7. Conclusion
Cheng, X., Chen, Y., and Sra, S. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06522. 2023.
In this work, we provide a theoretical characterizing of two training phases to uncover how a one-layer transformer can be trained to solve two regular language recognition problems: even pairs and parity check. In order to charac- terize the joint training of attention and linear layers, we employ higher-order Taylor expansions to precisely ana- lyze the coupling between gradients of two layers and its impact on parameter updates. Our results not only offer deeper insights into the training behavior of transformers but also highlight the critical role of CoT in solving par- ity problems. Experimental validation further supports our theoretical findings, confirming key aspects of parameter evolution and convergence. The analysis tools developed in this work can be useful for future understanding the implicit biases and training dynamics of transformers in structured learning tasks.
Chiang, D. and Cholak, P. Overcoming theoretical limita- tion of self-attention. arXiv preprint 1221v 2022.
Cui, Y., Ren, J., He, P., Tang, J., and Xing, Y. Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401. 177-66 2024.
Deletang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, E.,, Cundy, C., Hutter, M., Legg, S., Veness, J., et al. Neural networks and the chomsky hi- erarchy. In The Eleventh International Conference on Learning Representations, 2023.
Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310.126800 20233.
Acknowledgements
The work of Y. Liang was supported in part by the U.S. Na- tional Science Foundation under the grants ECCS-21 13860 and DMS-213414 supported in part by the U.S. National Science Foundation under the grants CN5-1956276 and CNS-2114542.
Devlin, J., Chang. M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. arXiv preprint arXiv:1810.044 2018.
The work of R. Huang and J. Yang was
Frei, S., Vardi, G., Bartlett, P. L., Srebro, N., and Hu, W. Implicit bias in leaky relu networks trained on high- dimensional data. arXiv preprint urXiv:2210.07082. 2022.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Gao, B. and Pavel, L. On the properties of the softmax func- tion with application in game theory and reinforcement learning. arXiv preprint arXiv: arXiv:1704.00805 2017.
References
Hahn, M. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 1561171, 2020.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Hahn, M. and Rofin, M. Why are sensitive functions hard for transformers? arXiv preprint arXi::2442.00990 2024.
Huang, R., Liang, Y., and Yang, J. Non-asymptotic conver- gence of training transformers for next-token prediction. In Proc. Advances in Neural Information Processing Sys- tems (NeurIPS), 2024a.
Bhattamishra, S., Ahuja, K., and Goyal, N. On the abil- ity and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009, //244, 2020.

---

Training Dynamics of Transformers
Huang, Y., Cheng, Y., and Liang, Y. In-context convergence of transformers. arXiv preprint arXiv:23 2023.
Magen, R., Shang, S., Xu, Z., Frei, S., Hu, W., and Vardi, G. Benign overfitting in single-head attention. arXiv preprint arXiv:2410.07746, 2024.
Huang, Y., Wen, Z., Chi, Y., and Liang. How transform- ers learn diverse attention correlations in masked vision pretraining, arXiv preprint arXiv:2403.02233, 2024b.
Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent is provably the optimal in-context leamer with one layer of linear self-attention. arXiv preprint arXiv:2307.0352 2023.
Ji, Z. and Telgarsky, M. Characterizing the implicit bias via primal-dual analysis. Algorithmic Learning Theory, pp. 772-804. PMLR, 2021.
Merrill, W. and Sabharwal, A. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.0792. 2023.
Ji, Z., Srebro, N., and Telgarsky, M. Fast margin maximiza- tion via dual acceleration. International Conference on Machine Learning, pp. 186004869 PMLR, 2021.
Merrill, W. Ramanujan, V., Goldberg, Y., Schwartz, R., and Smith, N. A. Effects of parameter norm growth during transformer training: Inductive bias from gradi- ent descent. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 2021.
Jiang, J., Huang, W., Zhang, M., Suzuki, T., and Nie, L.. Unveil benign overfitting for transformer in vision: Train- ing dynamics, convergence, and generalization. arXiv preprint arXiv:2409. 19345, 2024.
Karpathy, A. nanogpt.
htrp:://qithin com/ tirppthhynannoce 2023. Accessed: 2025-04-28.
Nacson, M. S., Lee, J., Gunasekar, S., Savarese, H. P., Srebro, N., and Soudry, D. Convergence of gradient descent on separable data. In The 22nd International Conference Artificial Intelligence and Statistics, pp.
Kim, J. and Suzuki, T. Transformers learn nonlinear fea- tures in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258, 2024a.
3420-3428. PMLR, 2019.
Nichani, E., Damian, A., and Lee, J. D. How transform- ers learn causal structure with gradient descent. arXiv preprint arXiv:2402. 14735, 2024.
Kim, J. and Suzuki, T. Transformers provably solve par- ity efficiently with chain of thought. arXiv preprint arXiv:2410.0863: 2024b.
Pérez, J., Barceló, P., and Marinkovic, J. Attention is turing- complete. Journal of Machine Learning Research, 22 (75):1-35, 2021.
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint
2014.
Phuong, M. and Lampert, C. H. The inductive bias of relu
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. Ad- vances in neural information processing systems, 35: 22199-22213, 2022.
networks on orthogonally separable data. In International Conference on Learning Representations, 2020.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
Kou, Y. Chen, Z., and Gu, Q Implicit bias of gradient descent for two-layer relu and leaky relu networks on nearly-orthogonal data. Advances in Neural Information Processing Systems, 36, 2024.
Sakamoto, K. and Sato, I. Benign not- net-benign overfitting in token selection of attention mechanism. arXiv preprint arXiv:2409. 17625, 2024.
Li, H., Wang, M., Liu, S., and Chen, P.-Y. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint vrXiv:2302.06015. 2023.
Sheen, H., Chen, S., Wang. T., and Zhou, H. H. Implicit regularization of gradient flow on one-layer softmax at- tention. arXiv preprint urXiv:2433.080 2024.
Li, H., Wang, M., Lu. S., Cui, X., and Chen, P.-Y. Training nonlinear transformers for chain-of-thought inference: A theoretical generalization analysis. arXiv preprint arXiv:2410.02167, 2024a.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. The implicit bias of gradient descent on sep- arable data. Journal of Machine Learning Research, 19 (70):1-57, 2018.
Li, Y., Huang, Y., Idiz, M. E., Rawat, A. S., and Oymak, S. Mechanics of next token prediction with self- attention. In International Conference on Artificial Intelligence and Statistics, pp. 663, PMLR, 2024b.
Strobl, L., Merrill, W., Weiss, G., Chiang, D., and Angluin, D. What formal languages can transformers express? a survey. Transactions of the Association for Computa- tional Linguistics, 12:543-566 561, 2024.
10

---

Training Dynamics of Transformers
Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. Transformers support vector machines. arXiv preprint arXiv:2308. 16898, 2023a.
Yang. T., Huang, Y., Liang, Y., and Chi, Y. In-context learning with representations: Contextual generalization of trained transformers. In Proc. Advances in Neural Information Processing Systems (NeurIPS), 2024b.
Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. Max- margin token selection in attention mechanism. In Thirty- seventh Conference on Neural Information Processing Systems, 2023b.
Zhang, R., Frei, S., and Bartlett, P. L. Trained trans- formers learn linear models in-context, arXiv preprint arXiv:2306.0990 2023.
Thrampoulidis, C. Implicit bias of next-token prediction, 2024.
Tian, Y., Wang, Y., Chen, B., and Du, S. S. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems, 36:71911-71947, 2023a.
Tian, Y., Wang, Y., Zhang, Z., Chen, B., and Du, S. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00555. 2023b.
Vardi, G. On the implicit bias in deep-learnig algorithms. Communications of the ACM, 66(6):86-93, 2023.
Vasudeva, B., Deora, P., and Thrampoulidis, C. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738, 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Wang. Q. Li, B., Xiao, T., Zhu, J., Li, C., Wong, D. F., and Chao, L. S. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1810-1822, 2019.
Wei, J., Wang, X., Schuurmans, D.,, Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-though prompting elicits reasoning in large language models. Advances in neural information processing systems, 15:24824-24887 2022.
Wen, K., Zhang, H., Lin, H., and Zhang, J. From sparse dependence to sparse attention: Unveiling how chain-of- thought enhances transformer sample efficiency. arXiv preprint arXiv:2410.05459 2024.
Yang. H., Kailkhura, B., Wang, Z., and Liang, Y. Train- ing dynamics of transformers to recognize word co- occurrence via gradient flow analysis. In Proc. Advances in Neural Information Processing Systems (NeurlPS), 2024a.

---

Training Dynamics of Transformers
A. Auxiliary Lemmas and Equations
Lemma A.1 (Gao & Pavel (2017)). The softmax function with scaling factor X is Lipsschit continuous. Mahematically, we have
|lØ(x/A) - -(y/A)I - x -1 v/1.
By Lemma A.1 and (n.t) exp ( W,a(")//a) the subsequence proofs.
we have the following inequalities, which will be itensively used in
Ex exp (w))) ( //a)"//a))
10t Pp ( S
(n.t)
W Wi
A
Let ((x) log(1 + e )) be the logistic loss. We also frequently use the following Taylor expansion about (2).
(
1
J(x)' - J(x'| s 1r x'1.
s xx
(2)
+ cxx(xx
2
4
A.1. Gradients Calculation
Recall that the loss at time step /
I
Lemax
L.
J
Ur
Σ (0) (mnd)
pp
2L
REIL
For simplicity, we use the following notation. For each sample n and time step t, we define
J(me)) 1 exp(y,,
where Tx(X(()) = u E, x
(n))
p(
(n.t)
Therefore, the gradients can written as follows. The gradients at time are
Lmox VuL.
Σ
Knn Σ
(n) (n.1) xe Pi
NEIL L=1
Lenax
V
2L
Un Σ
ut x1
fit
(ait)
ze
(n.t) (m) -
-
oe
Iz (
(T(
L=2
nELL
Note that V. < L Liaxx and V Cill S on each token embeddings.
|luel Lmmxx/ /X. We will also frequently the following projection of gradients
Σ
1
Σ
(m.t)
(m.))
2
L>l
nEl,.
(m)) ["))=w
(V.
1
Σ
(u, x(()) -T(X(1))
21
o( )=(()) (1)
=
(ET", (v
1
Σ
[u x1(1) - T(X(m)))
21
nEtz. ((2 (,))) =
(EY, (Vw E.)
1
Σ
(mit) (u, x -(((()))
2
n (*) (=(
)=w
where ((r)") (n)) the character of token embedding xt (a) For example, c(F w.
12

---

Training Dynamics of Transformers
B. Proofs of Even Pairs Problem
In this section, we provide the full proof of training dynamics of transformers on even pairs problem. Before proceeding to analyzing phase and phase 2, we first show that the token scores only depend on the position, which helps us reduce the complexity of the subsequent proof.
Lemma B.1. During the entire training process, for any {{a, D), we have
(we, ETr) (ue, 17 V 1.
For attention scores, have the following equalities.
(ET. W.EY) (Ei ".) LET.W,ER (E,", W,EE). 4e 2 2.
Proof. Note that results are valid for i = O. Assume the results hold at time I, we aim to prove the results hold for / + 1.
suffices to prove the following equalities.
(Vah,, = ET) (V 1E
(3)
(EY, (V (V) (ET", (VwL.)ET), 34 2,
(4)
(V,L)) = (ET", (V
(5)
We first show that Equation (3) is true.
For 2 have
(V WV,,, E2) Σ
Σ
ImnniWnip/mm (m.))
L2f L2l nel )=w
L2 RET, (n )=w exp(ynT((X("))) Un y
Σ
Σ
(n.))
For any In satisfying (() w. let nt' be the sample that only replace with w at the 8-th position. Then, due to the induction hypothesis, we have vp! vp (n'm)) 2) and Tx(x(()) (x(""))) Since f 2 2, changing one token at the position does not change the label, we have yn = Yn' Therefore, we have
(VuLe, ET) = ( (V F ") 2
For we have,
(VT,,
Σ Σ
Unmss (1,1)
121 nET1, n( (n) )=
Σ
Σ
In gn (m,t) (n
1 ex
121
Now, for any E Tr. satisfying c(x(")) w, let n' be the sample that flips the first and the last token at the same time. Then, due to the induction hypothesis, we have ot (mit) = (n'.t)
yn you and Tr (x(x())) = Tr(XE'), Therefore, we have
(V.C, = Lv, ET>)
We conclude that Equation (3) is true.
13

---

Training Dynamics of Transformers
Then, we show that Equation (4) is true. For É 2 2, have
(V
1
Σ
(n.s) ((rit)
u (TTT(a)) (x)
( Tx()
2L
netn Cl0z )=(0=
1
1+ (vnTr(X(m))) Unipi (mit) (u. x2 - T(()
2
)
For any n E I satistyfing ((x)")) c(xl")) = w, let n' be the sample that only replace with w at the é-th position. Then, due the induction hypothesis, have yp (mm.)) Pi (n,,)) and (x() Tx(,),)) Since 2 2, changing one token the position l does not change the label, have yo yn- Therefore, we have
E (V WWLL)E)) = =( (E (VwL.)EE, 2
We conclude that Equation (4) is true.
Finally, we show that Equation (5) is frue.
Note that
(VWL.)
E
[1 ( (uTaf(i)-TTa - (x(()
2
nele, ((a, "(a)") )=(()))))) (m))
E
p(mi)) (m) (
2
nEle,0(a) [( + exp (xr((()
Now, for any In satisfying ((x)")) c(z
(n) w, let n' be the sample that flips the first and the last token the Then, due to the induction hypothesis, we have o (n.1t) (n' .t)
same time, i.e., (()(")) c(r("))
se
Yn yur and
(1) ) Tx(x - (n)). Therefore, we have
(u ETT (V FT
We conclude that Equation (5) is true.
Therefore, the proof is complete by induction.
Due to above lemma, for each length L, we only need to analyze two types of sequence, i.c., the one with positive label and the one with negative label. We use x(() = [EY,, E ..... ET to represent the sequence with positive labels, and x( [ EY,..... E7] to represent the sequence with negative labels.
B.1. Phase
In this section, we characterize the training dynamics in phase 1. In general, we prove the results by induction. First, we characterize the initialization dynamics.
Lemma B.2 (Initialization). the beginning (t 2,3), for the linear layer, we have
(ug, ET = ((tt), Ye,
(w2, ETT - AA 2 2, Ve,
(w2, LT ET" S -0(17 v4 3, Vw.
For the attention layer, we have
112
EE - E F
> Ω
Vie
L
14

---

Training Dynamics of Transformers
11²
(ET - E2 WgEE - -
Ve,
L
11²
(ET - Et W3b 2 Ω
VP
3, Vw.
L
Proof. Since (n.0) 2. and pp (m,0)
for 12 E 17 (the length is L), we have
(ui, DTr r7( vV. ET)
Lmax
Σ
2L
L=1
nEl
)=w
(n)
=
where the last equality is due the cancellation between positive and negative samples whose length is greater than 1. Due the same reason, for any e 2 2, have
(u1, E1 = (7(-7vVo VuLoo Etr)
I_ax
Σ
(-in (2.0)
2L
L=2
nE nEI2,0(0r )
0
Regarding the attention, for any token EY", we have
(E W.E = T7 (- E.,
1
Σ
(n.0) (uTzz² Un (n) x To((((n)))
2
yur ((x()"))=
0,
where the last inequality is due to the fact that uo = 0.
In summary, time step 1, only the token score the first position increases, and all other token scores remain 0, and the attention scores are all 0, resulting Pi (n.1) 1 for n E Iz. Note that, we also have
L.
1
1
KKn. = + exp( IT
n E I2; -Kk. = 1 + expl 3)
E 17.
Next, characterize the token scores and attention scores at time step 2.
{u2, ETT)
(ui, ETT
VuLi,
Linos
Σ
(1mpp))) (n,0)
+
4
2t
1.m2
)=)) >=
Linax
7
1
+
+
4 2
exp(11/)
2
1
expl 11
+ expl
n1
L
1=2
1T.
47.
Thus,
(12, E
2n
<
1
+
Lmxx
Σ
17
17
4
2
+ exp(n/4)
2
2L 2L
1
15

---

Training Dynamics of Transformers
3r
8
where the last inequality is due to the Lipchitz continuity of f(x) = 1/(1+=*) (with Lipchitz constant 1) and 3IL=11L1 S
Similarly, for 22 2, we have
(U2, Et = 17( Vn Et)
Lmax
Σ
(n.0) (n
2L
L=tt
ne of
mmax
1
1
2
exp(
1 exp(
T
L=
42
Lenax
1/2
4L2
r/²
4
In addition, regarding the difference of token scores, we have
Lenax
Imax
(u2, Ex
E7
Σ
1
!!
+ exp( (.)
+ exp(ppp2
-11))
1 e +t)
1 exp(p (
L=t
8-1
2
1 expl 41
+ exp(
11
T.
L=2
172
16
where / 2 3.
Therefore, already prove that
(u2, ETT = (())), Vw,
(w2, ET") = -)))??) VE 2 2, Vw,
(w2, E E?") S - v4 2 3, Vwn
Next, we analyze the attention score.
A (E
1 Σ
[(JJa, (n,1)
(uu, Elir) -2)) ))
21.
Σ
(- (-J( (n,1) (!ues Fizt) ((11,
2
4L
(( (
2( ((u
-un, E())
8
1 1 exp
11
L
47.
+
1
1
1
(ui, E
17
8
+ exp(
L
4L
1I
(0) 1
3n2²
17
+
(u1, E
7L
1t exp
2
4
8L 2L
1
7
+ 18 1 exppp - 2
4L
16

---

Training Dynamics of Transformers
!
52
where (a) is due to the Lipchitz continuity of [(x) = 1111 + e²), and (b) is due to (ui, Egg) S /66
Similarly, for negative samples, since, we only flip the label Vn. from 1 to -1, we directly have
2 ( F2 ( < - (7)
Following the same argument and algebra, we have
A (ET , (-V 2 Ω (7.)
A (FT - Fa V < - (7)
which implies that
172
(E W3 F1 > Ω
I
172
-E E", W3E $ -2
L
Since ET) for all L 2 2, we have
{EY EF, V 0.
Finally, we aim to show that at time step = 3, the attention layer also distinguishes between non-leading tokens. This can be done by noting that (u2, EY EE7) < (2) Specifically, we have
1 E. ( V E2)
17
Σ
(-Imm,)) (m.22) (u22 - (x(
2L
nETL, )a(a m
Σ
(n.2) ((.2)(12)))) ((u2, ETT) - T2 (2()
Eu)
2L
mEl2,0(0) (1)) (0)
)=(x)()))))))))))))))))))))))))) )=
(a)
11 (- 1(+.)) P2 (+.2)
Ku2, Bbr - (uz, 1.7")]
7 (-)( )42 (-)2 [(u2 ET) - (u2, ETM
2 0(1/LL))
where (a) is due to the fact that ee (n.2) are equal for any 8 2 2, and the last inequality follows from that (wz, Fige)
(u2, bTr) < s and (-J(++2) - (- < -Q(n/L) L).
Thus, the proof is complete.
Then we prove Theorem 4.1 through induction. Note that the statement in the following is essentially the same as that in Theorem 4.1 due to Lemma B.1.
Theorem B.3 (Restatement of Theorem 4.1), Choose X = D(Dhaaw) = mmxxm))max) and =
1/12/3J). we have
(1) The dynamics of linear layer is governed by the following inequalities.
(w,, ET")) = ((r)),
Yw E (a, b},
17

---

Training Dynamics of Transformers
(ur: EEr) - Vt 2, Vw E (a, (ug: E2g EEg) S -Q((pP)) Vw {a,
(2) The dynamics of attention layer W is governed by the following inequalities. For any length L S Lmax have
{ Et WT 2 / V4 22
EEz - E ve wi WW 2 007 LI, 22
(1+
7
(F FT" F. 20 Ω
AP V223. 3,
XL3
AL
Proof. By Lemma B.2, know that the results hold for f 3. the following, suppose the results are true for We will intensively use the following fact that characterize the norm of attention parameter W,. Since
Liaa
V Call 2
Σ 2
(m.2)
L-2
NEIL
t=1
I 0 (Imax))))
we have
1- WWel $
Σ M C. S (Lmaxx)) 2
We first show that the dynamics of each token score (u,, ET) is true.
For the first token, we have
Limax
(- u-V,L )
Σ
(-Ins))Waa) (n,t)
2
1=1
n(l
)= ==
+ cxp (1ue, E1T))
Imax
Σ
(1010)))1111 (0.0)
+
2L
L=2
nETL,C(a(m) ymw
Lenax
Σ
(-Jom.)) + J(m.os) ynspl (m.0)
+
2E
L=2
n(( (n))
Lenax
Σ
(-Imne)) (p(( - -pp (m)))
2L.
1.m2
1
1 exp(((u;
2
Lmox
Linox 1
+
Σ
(-J(r.)) + J(m.o)) ynsl (n.0)
Σ
Σ
(-J(mnn))
2
1
1
nET
where (a) follows from the fact that J/n,0; = 2. and (p()))) = =
Then, by the Lipchitz continuity (Lemma A.)), we have
18

---

Training Dynamics of Transformers
Hence,
1
1
Lmax
1
Lmax
W/ll
(. E
4
S
2
(ue: ET
2
2
+ exp(unT((X(n))
L
A
L-2
nEJi, (n) j=w
L=2
(4)
(u,, ET |
+
Lenax
Σ
1
(4
1
( + Ie (((x(()))))
L I Wel
2
I
A
1=
n(a) )=w
1n
Temax
Lmax WWll
Kute BTT
max 1885L
Keee, 1.7711
2L
L=2
L 2t
W7
(we, 1T + log Lmax) +
A
$ O(nt),
where (a) is due to Equation (2), (b) follows from nyt 0()), and the last inequality follows from À = (( and induction hypothesis.
Similarly, for 2 2, we have
(-V Et)
Σ
Σ
(-mme))ess (2.1)
2
L=tt
Linoe
Σ
Σ
(-imas) (mn1)
Σ (-(me)))
2
L=t
mEEL ae (
netE )(u
Imaa
Σ
(-J(+ (ttt) 1-00)-))p--00---1111100100-001 --N((-0))--))
Lmax
1
1
(x((t()
1 1
_x(((( O(1 max(ue; pE)13)) (5-10)
4
2 4
L=t
(6)
Lmax
1
Σ
L
+ +/ (1)
-1) (ue, Ei} x5-
16
T.
(0)
Tmmax Σ
1
nt
1
Ω
T
T.
L=8
P=
Ω
(nt),
where (a) is due to Equation (2), (b) follows from Tpl = 0(1), and (c) is due to the induction hypothesis. In addition, for any l 2 3. we have
t-1 1 Σ
(-V Ey - ET1)) =
4
1=
Timax
(( (t0)) - -/(t00)- (--(00))) (-.00 --16--))660-----00)))))
+
L=t
(n()
Lmax
(0) + (15t (t.t)
Σ
4
Lme
19

---

Training Dynamics of Transformers
1 + 4
Tx(x))+ + +0(1()) (15- - -1)
(6)
Wall
-1(()))
+
A
L=
S {((()),
where (a) follows from the same argument of the previous analysis on EV",, and Equation (2), (b)) is due the fact that T₃ (X( 2 0, spf".t)
(n.f) 2 O, and the last inequality follows from À = S2(
Therefore,
(UL41, FT)
71( 1)
Σn
1 V-VCs, FT) 7
4
s 0 (1)8))
Similarly, for 2 2,
(11-1, 1b = Σ 17 V VuL.. ET)
x=2
S -1((78)).
In addition,
(ur41, D2
ET Σ 77 (- C.C. E3 ET")
2=2
- 217
We conclude that the dynamics of un in Phase is true.
Then, we show that the dynamics of W, in Phase is true. In the following, we fix the length L, and focus on the L-th column of W,,
We denote
A EE EY W,EY)
B(0 (ET ET", W.EY)
(1) (EY EY, W. W.FY),
which are three quantities we aim to analyze. Due to the gradient update in phase 1, We have
A
(t+1)
(
(ET - E (-)
1
Σ
(-(m.)))11))) (ma))
(T) -T(x(()))
20
m
7
Σ (-Imn)))nnpp (mit) (lu, Ete) T(X(1)))
2
nEII, (()))) ==
71 - (-((t,)))
(pittt) (lu,, ET') - -(((())) -((tt))))
(t.l)
-
8
20

---

Training Dynamics of Transformers
+ 7
J (()))))) (+:t)) (w,,
T(x(+)) -((x(())) + 8 ( (1)) (wee -.a E,) - T(X(-1))
8
17 1 1
(T((x((t) + +
(202 (1/ttt))))0 (ui, 1iTr) - 442 (+.1) {ue, -i²²) - (2pp - 5/2 (22x1))((x((t))
8 2 7
1
- + -((((((-)))))
1
(55-1) (Kue Ey'i + |T(((())))
+ 7
8
2 nt
1 11 - no
1 (Hmipe + 1
nz
n't
2 0
T
where (a) is due to Equation (2), and (b) follows from the induction hypothesis. Thus,
1/22
A (
(1(51)) - 1 + A = 52
I.
2m2
Similarly, for any 8 2 2, we have
B (1+11)
B
FT" ( (-V FF)
1
Σ
(- )m (n) (lui, Ext) - T(x(m)))
2L
mETZ.0(e =C0))
+ 1 (1),,)) (1-))
(u,, ( - T(X(-))
"
(ue, (-)(-a) (2p) (-1) ETr) 0p: -00-00 (ue, Etii) - -
- ((-())((X(--
+ "
( 1(1)))98
(+(0) (wu, - T(x((t))
8
2
1
1 _ nn (71P2 + 1
1/12
T.
T."
T.76)
1/(
Ω
L
where (a) is due to Equation (2), and the induction hypothesis. Thus,
7222
B 2 Ω
(1+1) m
T
Finally, we show that the attention score the second position surpasses those of other no-leading tokens. We have
A
(1+1)
All = "MMEET
Er
(-V
28
17
Σ
(-Tme)))n (m.t) (wu. - T(x(m))
2
((e))))
1
Σ
(-J( (-Jmno))nnpp (m)))) (lu, Exy) - T(X(NO))
2
nElL,C(a (2) )-
(m)) >=
30
"
(J((+))) (+.1)
(ue, F2r) - 4 (ut Fx (yp (1))t 2 (
(+11)
(62
" ((- (-J(-o) (-t)
() (u E7Ph) (42 ( (t) 25 ())x)))(--))
(
(ue, Enat) - up
8
"
[ (t(1) (-u,, E - pi
(++)) ( Us, E ) + (12 (4,4) - 2p )
8
21

---

Training Dynamics of Transformers
77
1(-1)) 822
Ur E2) Pi (1) (-u,, ki + (p2 -1 /p -0 IT(-)
8
(
1+ +
exp(A12 (1)////(/-11//0))))))// We: E2) (, E + (cxp(A2) /2) 1)T(X-))
7
A), eep((A2/ /AX W1 E3y - (-- Et (exp( 1)T(X))
8
7
1-)0-00 -.) (exp(A2) A) -1 1) (T(X -((x)))
=
8
(-)(- (1)
(-J(+,0) )
8
(-u,, ur, P27
(-ue; F ) (exp( (A/2 /A) 1)(( 1), E3 T(X+))
7
1)) (-(0) (exp(A))) /2) (>) - ( -T((tt)
14/1-1)
2 775
17t
(exp(A)) (A) -
nt
n0 ( (exp(-19//1) -1) L L
A 17
99 Ω
1/ 1
LmaxT)/
A
L
1
LA
X
where (a) follows from the definition of softmax, (b) is due to the induction hypothesis, and (c) due to À [2(Tmaxx). By noting that ( 2 Lf,/L, we have
((1+
22
28 2
AL3
AL
We conclude that the dynamics of W, in Phase is true.
B.2. Phase
In this section, we analyze the training dynamics during Phase 2. Roughly speaking, both the linear layer's parameter vector u, and the attention layer's parameters increase in norm over time. However, due to the scaling factor A, the linear layer dominates the loss reduction, contributing more significantly to optimization progress than the attention layer. In the following, we leverage implicit bias theory demonstrate that the growth of ||u,]] induces sublinear convergence rate for the loss, governed by (1/tt). On the other hand, the attention layer does not change significantly.
First, we show that at the end of phase ( (l=lo), the attention layer make data samples separable.
(n) (n.ta) Then, at the end of phase 1, the dataset Proposition B.4 (Restatement of Proposition 4.3). Let u(m) = Σ/ De 4i ((v(%) }n)) is separable.
Proof. Let u Ewe(arb) (EY ET.") Then, for any positive sequence, where yn 1, we have
(u, Σ ze pt p(ntto)) = p/
(m.to) - (p(m
(n,ta)
8=1
(m)to) (exp ((E)) -E,, W,ET) 1)
= =02
> 0,
where the last inequality is due to Theorem 4.1. Similarly, for any negative samples, we have
L
(u, Vn Σ
(n) (ntt))
= -6p( (m.to) + 122
(n,to)
x pe
(0) (mnta)
(exp ((F
FT W.FF)) - -1)
22

---

Training Dynamics of Transformers
0,
where the last inequality is due to Theorem 4.1.
Thus, The proof is complete. Parameter setup. Before we present the technical lemmas, we first introduce two parameters. Let
A2/3
A2/3
T min
/3llugrll 22/3
22/3
Co 2 log 4 4Lmax +
We remark that since À (4 max). T is well defined. We first provide the property T. Note that for any 4, we have
|| - 1 - ueal + ||uelll
S
VaCell
=
S Tmmme
Similarly,
IW, - Wiol S
Σ
W Lill
s=to
Ug Lmax
17 L D t?
A
Therefore, for any / < T, we have
IIterlllv Wi, < and | lull
We S 2
Next, we provide the key lemma in phase 2, which characterizes the alignment of the gradient - V. and the logarithm growth of [[ue|].
Lemma B.5, Let upp be the solution of the following problem.
uEP arg min [|u]],
u, Un Σ
(n,ta) (n) Pi It
21, 12,
Then, for all to < T, we have
Wt V. Lt,
1+
Collugrll |uell
1)
ugr -Vaae,
||elll
|
and
|uell 3| log(t to 1)
3 3
9 91
log
"
Proof.
First, note that for any < we have the following lower bound of Co-
log 4 Jmmax + 2|]ux|]||We - W
Co 2
X 1 [l[ W
23

---

Training Dynamics of Transformers
We aim to show that L
ugr
Ilwlll |
+ Co) W:)
< C(ue, We).
Since ut is not the optimal solution of the problem
(nito (n)
arg min ||u]|,
u, Un
pe
21,
we have at least one sample n of length L such that
Wt
(n.ta)
(n)
|
Un
Pe
Ft
|ugpl
which implies there are at least 2 samples satisfies the same inequality. Hence, for this n, we have
Yn (X(mm))
(n,1) (n)
Ue yn
1=1
(n,to) (m) Ft
( (m) (1p
(n.1)
Pt
+
88
1=1
||uell
W
Wioll
+
||eelll
|ugell
A
where the inequality due to the Cauchy's ' inequality and the Lipchitz continuity of softmax function (Lemma A.1).
Thus,
Lemax
C(u, Wi)
Σ
Σ log
1 + exp exp
(- ((((()))))
2
L=1
nElL
||uell
Wwel/|WW - Wcol )) A | MllllWWe Weo
2
log
+ exp
2
|
1
|uel
4
exp
"Brl
where the last inequality follows from log(1 + x)) On the other hand, we have
when E (0,
L
L ugP, Un Σ
I
Σ
(m,) (n)
(n,to) (m)
(m))
+ "Epp Un. (vplre (m))))
p/ (m)t0) yar
"pP Un
It
1=1
W
We,ll
21 1110 11
A
where the inequality is due to the definition of upp and Cauchy's inequality. Thus,
|eulll |upell
(- (-(1- 1 |ugr ll| W
||uell
)
L upp
<1EE
Co) W.
< I Tamax log
eep
1 + exp
Co
1 Iu'rll
||u6||
|uell| -Weoll exp
|lue W
Wooll
< Linax exp
-llo log 4.max
| ||
| ||
|uull|WW, Wrol
exp
24

---

Training Dynamics of Transformers
Therefore, we conluede that (Ep
+ Co) W:)
- C(ur, W.).
Since C(u, Wi) is convex respect to u for any We, by the first order optimality, we have
(uule uer
( (
|1EEpll
||uell
+ Co
u
0
By rearranging, we conclude that
Lt
Ut [eel
1
ColluEp!! [[ue]]
uEP (Vaaf
|
Next, show that the norm of un grows at least logarithmically. On one hand, we have
Leaax
X
VuLs:
EEP ||ugl||
Σ (-(mm.))
UHo
Σ 1=1
of (m,)) afmmeo) x
2
L=1
nElL
Lmux
I[, Wioll
Σ
2L
1
+
Txp(unn
A
L=1
nEI.
(0) 1
I W
2 + exp
Iuell |
[uelll - W/oll/a)
||wEPl
+
3 U + 2 exp
|uu.l
where (a) follows from the non-optimality of u. and (b) is due to that IIW, - Wi. S बाा
Thus,
1
|usl| 2 311 ur s'=oo 1 + 2 exp
||ut0
Let f4 maxxxft/Vs S ||u,|| S kj. We next show that t has an upper bound.
We have for any E tk.
à: 2 || 2
to) 3 u + 2 expp(
7
| ||
Thus,
31luErr
k
t. to
(1: [|/utol)
1
2 exx
Tueelll
17
9 91|uEp'1 112
k
|
+
| exp
17
2k
k
3 31luuppl
+1
exp
|Iwgell
14 44
5k
exp
7
where the last inequality is due to x + 1 < et Hence, & has upper bound:
14 14|
5k
t - to 5
exp
11
3 1 gpl
25

---

Training Dynamics of Transformers
Therefore, by the definition of 1g. have
|1t++1|
3 3l log
7 |4| uppl
(1x to)
Equivalently, we have
3 3/ugpl
3 31
14 14 log
|uel
log(t to
5
5
1
We conclude that the norm of u, grows at least logarithmically. The proof is complete.
In the following, show that there exists threshold 12, such that after time step 22, the transformer correctly label the data, which eventually leads to loss decay.
Lemma B.6 (Formal version of Theorem 4.4). There exists l2 such that for any 2 12 have
2
Ut
upP
|
|ugrl
2
6/
and
5 well
yo (x(())
2
(6)
u Rgrll
Proof. The proof consists of two parts that leverage Lemma B.5 in a different way. In the first part, we aim to find an interval (/2 la] such that the result holds. In the second part, aim to prove that for any 2 la, the result holds.
Let
141luEr|1
5
f1 to +1
exp
7
3A
Then, we have
|uel 2
where A will be determined later.
We aim to find an interval [t2, ta] such that Equation (6) true.
By Lemma B.5, for t 2 t1, we have
Co Collurrl
wgr
V Lt,
< +
- -uuLi
l
|ulll
|u'pl
uEP
< (1 AColluEpll
vuLe,
|
which implies
WEP
1
Ue
U1+1 Ut
|
1 AColluE
Fpll
U+41 Ue:
|
1
|
well"
|ue+
||ur/|1
+ ACo ugrl
2 21|utel ||
1
17² L A
[
luell
+ ACol|uupp
26

---

Training Dynamics of Transformers
By telescoping from 4 to i 1 and rearranging, have
Ut ||ell
uEp
(1- |/t,, ll |
s7² 1. max AlL (1)
||us,
2
||
+ &Collugpl
2 21
|
Now choose l2 such that
9 E
5 Me, 11
t2 = to + 1
exp
3 u u
which gives us |luell
Mey II
Thus,
Ut
EP
2
17
122
max
A(1 - 11)
|/4 ||
|uell
|
1 AC ACollu uelll
2 1/1191 ||
114
(a)
AC Eulll
1
21[ue,
[1/[|[ce]]
Liaxxxx A ~
t2)/22
1 + AC || Epll
AC [Epel
2 1
2A
L A3 L
12)
I + ACul "ppll
where (a) follows from | 2 ||u,, / 2 1/42, and the last inequality is due to that || |1/l|u/ll S A.
By choosing = o inequalities hold.
Tمlي and noting that S T
22/3
we conclude that for / E (/2, T2, the following
nlemax"
Ut
MEP
Ut
VEP
2 2
[gp
|
|
2 2/ ||
+ 4. + L2 A² - 41)/2
ACol ||werll
A
2C 2ColluEpl!
+ 4 172 L2 max A t1)/2
61lurll
1
6 61
||ue||
where the last inequality is due to that |/ 2 1/A Hence,
(x((())))) =
Ue, yn
4 xt
(n (n)
t=
(niso) (n)
| [, Weall
Us, Un
A
1=1
|
||ue||
||
Vpp: Un
Σ x0
(nnto) (n)
We - ||uell
|
UEP
yn
L
(m,ta) (n)
xt
|/4
W|
W
||uce]
||uell
21
Uta
[u llllwW Wooll
|
|
1 |l
A
13
|eell
|!ulll
-1
Eell
61lueep
27

---

Training Dynamics of Transformers
5 Ue
6 6ll
where (a) follows from Cauchy's inequality and the Lipchitz continuity of softmax function (Lemma A.I), and (b) is the due that / < T.
The proof is complete.
Theorem B.7 (Restatement of Theorem 4.5). Recall that ' = =0
12/3 (2///
For any T, we have
Limax /ppll
L, =
n/t
Proof. By Lemmas B.5 and B.6, we have proved that for any T 2 b2 we have
5 well
yn (xT((
Vn.
6/ u "EPll
and
31
nit to 1) log
|u,|| 2
5
14 14
Thus,
Lmax
L1 _ 2T. log
11exp((yyTT(
L=1
Vl- - 10 1 141lurr!
1
S Immax log
Lmmax|| [upps
S
nyt
12/3 n/mmax
Finally, since T =
we note that
1/2 1/.5 1Ep 112
Cr 50
11/3
C. Proofs of Parity Check Problem
In this section, we provide the proof of training dynamics of transformers parity check with COT. Our strategy is similar to that used in Appendix B. Since once the data after the attention layer is separable, the analysis of the dynamics would be the same. Hence, we omit the phase analysis in parity check and focus on showing that at the end phase 1, the attention layer makes the data separable.
We first show that the token scores only depend the position, which helps reduce the complexity of the subsequent proof.
Lemma C.1. During the entire training process, for any E {a, D}, we have
(ue: ET = (ue, B) VL 1.
For attention scores, when L Lo, have the following equalities.
(E W.FLE) (FIT" W.FFT")
(ET W.ET) = =EE (BE"")))EE".EEE" W.ET), 2 2.
28

---

Training Dynamics of Transformers
When L Lg let to L Lo + 1, then we have the following equalities.
E W.EY) = (ET," ".) WEE) - ,WEF",,EE) vl to.
Proof. We only check the last two equalities when L 2 Lo. since the others can proved similarly to Lemma B.I.
Note that the results are valid for / = 0. Assume the results hold at time I, we aim to prove the results hold for + 1. It suffices to prove the following equalities.
EE (-V = -B,,", (- (E," (- (W = ( (-V
(7)
Vl +
(8)
We first show that Equation (8) is true.
For + to L Lo + L, we have
(v) 1
Σ
(u, Innn (1.1)
(x) T(()((1))
2L
ne netz c(r(")) --())"))))))------
1
exp (x Unipl (mitt) (u Tx()
-1
2L
veT1.00 0( (1) )((() (1)))))))
For any n E In satistyfing C(2)")) = )(x()))))))))))) w, let n' be the sample that only replace w with - at the é-th position. =
Then, due the induction hypothesis, we have w/thtt) Pi (n',t) and (x(X(") Tx(""")). Since > 2, changing one token the position é does not change the label, have yo ya- Therefore, we have
(EY ()) = (E," (VwL))E), 42 2.
We conclude that Equation (8) is true.
Then, we show that Equation (7) is true.
Note that
EN (VwL.))E))
1
Σ
(m (1 (x())
2
n( (m) - (m)) -
1
(n.0) (m) ( - (x(x
mETz,0(x!!)) (m)) - (()r + exp (V(Ti()(())
Now, for any n E I satisfying ((x!")) = c(z(")) = w, let n' be the sample that flips the fo-th and the last tokens at the same time, i.e., (nn) (xx("²))
(m.11) (n' ,t)
Then, due to the induction hypothesis, we have spt 4p
Un Yn' and
(n))) -xx( Therefore, have
(EE (VwL))E)) = (E7"", (V
We conclude that Equation (7) is true.
Therefore, the proof is complete by induction.
29

---

Training Dynamics of Transformers
Due above lemma, for each length L, we only need to analyze two types of sequence, i.e., the one with positive label and the one with negative label. We also use xL (+) = , E ..... EL s] to represent the sequence with positive labels, and xf_ = - LT, ...., E.] to represent the sequence with negative labels, similar to Appendix B.
Next, we characterize the initialization dynamics. Lemma C.2 (Initializationn. At the beginning (t 2,3), for the linear layer, we have
(w2, ET ) = (n))), Vi,
(u2, E - - 34 2 2, Vw,
(w2, Eg EX" S ( ve 2 3, Vin.
For the atlention layer, when L S Lo we have
112 I
(E - E2 2 Ω
Vw,
112 I
E,T - Do - W3 W3EL 2 -Ω
Yw,
1/2
(Er Et W3 W3 2 !
44 2 3, Vw.
T.
When > Lo, let Eo = Lo + we have
(
722
Bio w - Et, WgEL) 2 Ω
44 + lo, aw,
EEY - EY W3 W3 2 Ω
ve + lo Yw.
n
(ET ET W3 W3E 2 Ω
Vl + lo, V..
T.
Proof. Since (m.0) 1
2: and Pe
(n,0) for n E T (the length is L), we have
(u1, F = 17( VV. El
Lmax
Σ
(-I(m.0))8n1pp1)))) (m.0)
2L
L-1
nET2,0(ar
)=
(n)
Jm
!
where the last equality is due to the cancellation between positive and negative samples whose length is greater than 1. Due to the same reason, for any 2 2, we have
(u1, F7
( VuLo. Fi())
Imax
Σ
(7,0)
2L
1=2
nET2,0(a (n)
)=
Regarding the attention, for any token Ev", we have
(DV', W.EY) = n){ (-VwL:)
1
Σ
(m.)
(NT x. To((((1))
2L
nEI (m) yur (.))
30

---

Training Dynamics of Transformers
0
where the last inequality is due to the fact that Wo 0.
In summary, similar to what happens in even pairs problem, at time step 1, only the token score at the first position increases, and all other token scores remain 0, and the attention scores are all resulting iPe (n.1) for n E Iz. Note that, we also have
1
1 T.eEi I2: -K,,1) + exp( 11
Km.11 =
E L.
exp II
Next, we characterize the token scores and attention scores at time step 2.
(u2, ETT)
(w1, F
V
FT")
Lmn
= 7
+
Σ
Σ
(- (7.0)
4
2
1.=2
)=)
)=w
Lmax
" + "
4 2 1 exp(n/4)
2
1+
expx
47
11
1 1
exp(-
az
21
L
Thus,
Lmxx
21
1
11 "1
(12, ETr
< 2
+
4
+ exp(n/4)
2
2L L=
2L
37/
8
where the last inequality is due to the Lipchitz continuity of f(x) 1/(1+=*) (with Lipchitz constant 1) and Σ 1/L2 S
1.
Similarly, for 22 2, we have
(u2, Ein) = 77( V V.Li, ET)
Lmax
Σ
(-ii (m,0)
n
2
L=
ne (
Lanax
11
1
1
2
1
exp(1))
1 -xxp(- 11.)
L
L=i
Linax
n
4L2
T.=
r/2
412
In addition, regarding the difference of token scores, we have
Lmnox
Lmnx 1)
(u2 Ex
E
1 1 exp((
+ eexpp--3 11)
exexp(
+ exp( - 21
T.
L=2
L=f
6-1
Σ
1
1 exp 9
+ expl
L
31

---

Training Dynamics of Transformers
?72
16"
where 8 2 3. Therefore, we already prove that
(u2, 1 (nt), Vw,
(w2, Ee()) = - 44 2 2, Vw;
(u2, E2 E S (² vt 3, Ve.
Next, we analyze the attention score. For any length L < Lo, the proof follows the same steps in Lemma B.2. Here, we just present the results.
A ( ET (-VwL)) 2 2 (7)
A ( E , (- S -2 (7)
In addition,
A AET - E2. (-v 2 (")
A A - - ET (-V s -9 (7)
which implies that for any L S Lo. we have
1/2
EET - Ex Wy E1) 2 Ω
T.
112
(FT - FV. W3 F. < Ω
L
Next, show that the initial dynamics of attention scores in sequence with length I La-
Recall that to = L - Lo + 1. We have
2 (
1
Σ
(-Hms)) (n.1))
(lw1, 111 )
=
21
nETZ
1
Σ
(- (m)n1)) (n.1) ( (u,, E
7
2
AL
nETz,0(z(
(m)) ==
1
-2(u1, E + (21, E +
8
1 + exp(
11
L
4L
1
(u,,, E
11
8 + expl aT.)
T.
4T.
(0)
Ω T
where (a) is due to the Lipchitz continuity of f(x) = 1/(1 + e and (un, ET) > 1(1)
Similarly,
X { - EV (- W- < -22 (m)
Since (u,,Ey) 0 for all L > 2, we have
X FE Th - F EW, (-) =
32

---

Training Dynamics of Transformers
Finally, we aim show that time step the attention layer also distinguishes between first and other tokens. This can be done by noting that (uz, EV"
E 2 ((()) for & + to. More importantly, (p(t12) > x f Specifically, we have
n Ey. (- V L2) FE
"
Σ
(-J 1m. (m (lu2, EY> Tx(
2
nEIL.0(u)
E
(- (-J(2)))Mu (n.2) (w22, Fitt T2(x((N))
2
(
((+))) (4)) ((+12) (u2. ETTT) - :p (+.2) (u2, ETr))
1 (-J(- |0| (++2) (u2, ET Pi (+
fug, Es
8
2 2(nn'111
where (a) is due to the fact that Pi p/m..) are equal for any & 2 2 and & * fo. and the last inequality follows from that
(u2, bT") (u2, E 2 (2))).
Thus, the proof is complete.
For the rest of the proof, the steps follow the same as in Appendix B. Specifically, by induction, we have Theorem C.3 (Phase 1). [Restatement of Theorem 5.1] Let ww denote the flip of token E {a, Choose À (2 and to = ((1/((xxx))) Then, for all t to the parameters evolve follows:
(I) The dynamics of linear layer u is governed by the following inequalities.
(ue, ET) = ((t)), Vw,
(u,, ET" = (t? AE 2, Vi.
(u,, ET ET" < -1017
YE 2 3.
(2) The dynamics of the attention layer W is governed by the following inequalities. For any length L < Lo, we have
(EY - Ey' W,EE) 2 2(77t)),
A4 2 2, Vw, w',
(E ET", W,ET) 2 0(72t), (Ey' -EV", ,EY,,,EY" 2 ((1)t),
V 2, Vw, w',
AA 2 3, Yu, w', w".
For length I, 2 Lop let lo = Lo + 1, and we have
(ET%" " EX', W.ET) 2 2(128),
vl + to Vw, w', vl + to, Vw, w', vl 1, Co. Vw, u/ w".
E ET" W,ET) 2 0(12t),
(ET' ET" W,ET) 2 1t1"t),
(n) (m.ta) yn)} is separable. This fact enable us to Therefore, at the end of phase 1, i.e., S to, the dataset (E Ft 9i performa the same implicit bias analysis as in Appendix B. Thus, we conclude that similar theorems Theorem 5.2 and Theorem 5.3 hold true.
D. Additional Experiments
Imapet of the scaling parameter A. We first select two additional À configurations for training on the Even Pairs task to demonstrate that the training dynamics remain consistent with those reported in the main paper.
Two-phase dynamics. As shown in Figure 6, the two-phase phenomenon naturally emerges even when a fixed step size used throughout the entire gradient descent training process. Notably, this behavior also appears in real-world datasets: Figure demonstrates that the two-phase learning dynamies persist when training with NanoGPT (Karpathy, 2023).
33

---

Training Dynamics of Transformers
UMA Iamodp=30
Attention Score Difference
Token Score
Loss
0.5
-10
2003
1000
3000
4000
3000
3000
3000
3000
4000
5000
Step
Step
Step
Figure 4. Training dynamics of learning Even pairs when À = 10
Attention Score Difference
47-ppWEE -Fpile sftwee
4.0
Token Score
30
Loss
2.0
1000
2000
3000
4000
5000
Step
Step
Step
Figure 5. Training dynamics of learning Even pairs when 18 Lmaxx
Loa, Lambdlon,, consant
Attention Score Difference
7.5
Loss
lt +f/wef - k 1//1/
Token
0.5
-108
3000
3000
5000
Step
Step
Step
Figure Training dynamics of training transformer on even pairs using vanilla GD (constant Icarning rate) and À
Netraa 17'4
MLP Porameter North
6i7a
Attention Parameter Norm
Aprisc
Training Loss
Validation Loos
1000
3508
400
1000
1850
8003
1250
Iter
Rer
Figure 7. Results of nanoGPT, trained on 'shakespeare' dataset. Configuration: block-size=64, batch-size=12, n-layer=4, n-head=4, n-embd=128
34


Reviews:
[{"review_id": "xBZ3PJGJyj", "reviewer": "Reviewer_TZ3J", "summary": "This paper focuses on two typical regular languages: even pairs and parity check. The authors show that one-layer transformers can learn even pairs directly without CoT. For parity check, it is shown that one-layer transformers can learn it with CoT and with a small amount of data mixing with even pairs.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments. Please note that all our new experiment results (i.e., the figures we quote below) can be accessed via the link https://anonymous.4open.science/r/icml2025-4BCC/Figures%20for%20ICML.pdf\n\nQ1: It is not clear how the results can be extended to deep transformers\n\nA1: We thank the reviewer for the thoughtful comment. We acknowledge that there are several settings that make our results seem better than empirical study. E.g., we construct an orthonormal embedding vectors, whereas Deletang et al. (2023) needs to learn the embedding functions and may lead to additional errors. However, our results still offer several key insights for deep transformers as discussed below.\n\n**(i) Layer roles:** Our theoretical analysis uncovers distinct roles played by different components of the transformer during training. Specifically, the attention layer $W$ learns to capture token-level semantics—e.g. token equality—while the linear layer $u$ encodes positional information. These functional separations are expected to persist in deeper architectures, where different layers may specialize to recognize different information such as content or positional information. **(ii) From simple to complex tasks via CoT**: We theoretically justify how transformers trained on simple tasks can generalize to complex tasks using CoT (Algorithm 1). This helps explain the success of CoT in scaling model reasoning capabilities without additional training. **(iii) Stabilizing training with simple-task supervision:** We demonstrate that incorporating simple tasks data (e.g. even pairs) improve the initial training of complex tasks (e.g. parity check) using CoT. This finding provides practical guidance for constructing datasets in deep transformers—including simpler examples can stabilize training process.\n\nQ2: The proof may critically depend on some seemingly weak assumptions: (1) zero initialization; (2) two-phase LR schedule; (3) merging $W_u W_v$ as $u$ and $W_k W_q$ as $W$\n\nA2: We would like to point out that these assumptions can be relaxed or justified as follows. \n\n(1) can be relaxed to random initialization such as Gaussian initialization. Then concentration inequalities can ensure that initial parameters are relatively small with high probability. Then, similar techniques can be applied to analyze the training dynamics of transformers to show that GD updates will push parameters on the right track. (2) serves as a simplified approximation of the decaying learning rate in Adam, which takes high learning rate in early training and then takes smaller learning rate in later training. See lines 182-191 (right column) for more detailed explanation. \n(3) can be removed by separately training $W_k, W_q$, $W_v$ and $u$. Our analysis can be extended to include analyzing the changes of more terms such as $W_k W_k^\\top$, $W_q^\\top W_q$, and $\\langle x_\\ell W_k , W_q x_L\\rangle$ during the training. We expect that the core idea remains similar. \n\nQ3: The analysis may not be really for the joint training of both layers.\n\nA3: Our analysis is indeed joint training. While the learning rates of the two layers are different in the first phase, our analysis does not de-couple their changes into two-time scales. \nSpecifically, on pages 16-21, we explicitly analyze the effect of gradients of $u$ and $W$ on parameters update at *each time step* and take into consideration how one parameter affects another. \n\nQ4: I feel the analysis is too ad-hoc to the two regular languages studied in this paper.\n\nA4: Thanks for the insightful question. While our analysis focuses on two specific regular languages, the core techniques, token dependencies, and the role of CoT, are generalizable via the following unified view of regular language tasks.\n\nEvery regular language can be recognized by a finite-state automaton, which operates by updating a state $s$ based on an input symbol $w$ via a transition function $\\delta(s,w)$. This abstraction aligns well with our framework, where even pairs compare the last token (state) with another token in the sequence, appending the result at the end, while CoT iteratively applies this process to determine the final answer. To generalize our approach to arbitrary regular languages, the key is extending from simple token-equality comparisons to a more general transition function $\\delta(\\cdot,\\cdot)$. This can be done by find an attention map that can approximate a general transition function $\\delta$. The training dynamics analysis will leverage some key techniques that we develop here such as handling the coupling of layers and implicit bias.\n\nQ5: It would be better to phrase the method as data mixing rather than adding a regularization.\n\nA5: Thanks for the suggestion! We will make the change or add discussions.\n\n------\nThank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.", "claims_and_evidence": "The theoretical claims are supported by clear theorems and proofs.", "methods_and_evaluation_criteria": "/", "theoretical_claims": "Due to the large amount of theoretical papers I have been assigned to review, I cannot verify every proof in full detail. The theorem statements look correct to me.", "experimental_designs_or_analyses": "The experiments on one-layer transformers are presented. The experimental designs are sound and valid.", "relation_to_broader_scientific_literature": "It is crucial to understand whether transformers can learn formal languages defined according to Chomsky's hierarchy, and this has been studied in many prior works empirically or theoretically. This paper contributes to this line of research by proving that one-layer transformers can learn even pairs and parity check, the two most basic regular languages, under a neat theoretical setting.", "essential_references_not_discussed": "/", "other_strengths_and_weaknesses": "Strengths:\n1. The theoretical setting is very neat and clean.\n2. Many prior theoretical works train layer weights of transformers separately, which is not very natural. This paper directly analyzes the joint training of both layers.\n\nWeaknesses:\n1. It is not clear how the results, or even how this line of research, can be extended to deep transformers. In fact, Deletang et al. (2023) (cited in the paper) showed that deep transformers are not able to learn even pairs (acc is not 100%). It might be the case that the setting studied in this paper is just a delicate setting that happens to enable transformers to learn these two regular languages. If this is the case, it would be important to discuss how generalizable the results are and whether there is something to learn from this paper for deep transformers.\n2. Although I really appreciate that the authors made the setting clean and neat, I do feel that the proof may critically depend on some seemingly weak assumptions: (1) zero initialization; (2) two-phase LR schedule; (3) merging $W_u W_v$ as $u$ and $W_kW_q$ as W.\n3. The authors claim that one of their main technical contributions is to analyze the joint training of both layers. However, in the first phase of the two-phase LR schedule, the learning rate of the two layers is decoupled, which makes me suspect that the analysis may implicitly assume that one layer is trained much faster than the other one. In this case, the analysis is not really for the joint training of both layers.\n4. I feel the analysis is too ad-hoc to the two regular languages studied in this paper. It is not clear how easily their analysis can be extended to other regular languages.", "other_comments_or_suggestions": "In the second approach to parity check, it would be better to phrase the method of adding even pairs as data mixing rather than adding a regularization loss, since data mixing is a more intuitive term to describe this process. I agree that its effect is similar to regularization, though.", "questions_for_authors": "1. I would like to see the authors' thoughts on Weakness 1. I would like to raise my score if the authors could point out some useful insight that can be extended to deep transformers.\n2. I would like to see the authors' thoughts on Weaknesses 2, 3, 4. I would like to raise my score if the authors could actually show that the proof is flexible enough to hold with big variations of the setting.", "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": "FQtvRMVav7", "reviewer": "Reviewer_btPg", "summary": "The authors theoretically study language recognition tasks with transformer. Formally, they study the training dynamics of transformers trained with gradient descent on the parity check and even pairs problems. Considering a single-layer simplified transformer, they first show that the even pairs problem can be solved directly before showing how Chain-of-Thought was needed for the harder parity check problems, either at inference with a model pre-trained on even pairs or during the training itself. For both tasks, the authors identify two distinct phases of learning where the attention grows in the first, mapping data to separable vectors. Then, once the attention module is stable, the linear layer grows and implements ma ax-margin decision boundary to separate the attention outputs between positive and negative samples. They provide a convergence rate for the training loss and empirically validate their findings on synthetic data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments. Please note that all our new experiment results (i.e., the figures we quote below) can be accessed via the link https://anonymous.4open.science/r/icml2025-4BCC/Figures%20for%20ICML.pdf\n\nQ1: It seems that the embedding strategy does not depend on the token value.\n\nA1: We would like to clarify that the embedding strategy does depend on both the token value and the position. To see this, suppose the token is at the $\\ell$-th position in the sequence. Then if the token is symbol 'a', embedding vector $E$ has 1 at the $\\ell$-th **odd** index, i.e., $E[2\\ell-1]=1$; otherwise if the token value is 'b', the embedding vector has 1 at the $\\ell$-th **even** index, i.e., $E[2\\ell]=1$. Clearly, embedding also depends on the token value. \n\nQ2: In experiments, authors take $L_{\\max}=6$ and $\\lambda=2$, which does not match $\\lambda=\\Omega(L_{\\max}^2)$ in Theorem 4.1. \n\nA2: We provide new experimental results (see Figures 1 and 2) under $\\lambda=L_{\\max}^2$, which exhibit similar training phases as our original experiments and validate our theory. We further note that $\\lambda=\\Omega(L_{\\max}^2)$ of Theorem 4.1 is a sufficient condition to guarantee the convergence of the training theoretically, but may not be necessary. In practice, as  demonstrated in our original experiment, a much smaller $\\lambda=2$ can also lead to desired convergence even with a faster rate. \n\nQ3: Theorem 4.1 chooses $\\lambda=\\Omega(L_{\\max}^2)$ instead of the common $\\sqrt{d}$ in the vanilla transformer. Will this lead experimentally to shrinking the logits range and outputting almost uniform attention values?\n\nA3: As discussed in A2, $\\lambda = \\Omega(L_{\\max}^2)$ achieves the same desirable output (non-uniform attention) as smaller $\\lambda$, although the training becomes slower. We also provide new experiments on deeper and larger models (see Figure 3), and observe that choosing a large scaling factor may not always degrade the performance drastically. Several reasons can possibly explain this. (a) Deeper models have layer normalization and adaptive learning rates, which might compensate for smaller gradients. (b) The dynamics of feed-forward layers can help activate softmax to learn so that attention is not uniform. (c) The attention weights will learn to increase its value to mitigate the high scaling factor (as shown in the Figure 3). We will include those experiments and discussion on the choice of $\\lambda$ in the paper.\n\nQ4: I am intrigued by the gradient descent in two phases considered in the analysis. Why did the authors choose this framework instead of a more classical GD, SGD, or more adapted Adam (noting that adaptive optimizers are better for transformer and attention-based models [4])?\n\nA4: Our two-phase learning-rate schedule serves as a simplified approximation of the decaying learning rate in Adam, which takes high learning rate in early training and then takes smaller learning rate in later training. See lines 182-191 (right column) for more detailed explanation. \n\nNevertheless, even for classical GD with fixed stepsize throughout the training process, our new experiments (Figure 4) show that the training still exhibits the similar two-phase process as characterized in our theorem. \n\nQ5: How did the author select the value $t_0$?\n\nA5: In Theorem 4.1, we provide a parameter setting that guarantees a desirable training output, where $t_0 = 1/(\\eta L_{\\max})$, and $\\eta = O(\\min\\{1/L_{\\max}, 1/\\lambda^{2/3}\\})$. For example, if we choose $\\eta=1/\\lambda$ and $\\lambda=L_{\\max}^2$, making the learning rate at the first phase $O(1)$, then $t_0=O(L_{\\max})$. This choice suggests how to set the duration of the warm-up phase in practice, which can be chosen based on the length of input sequences. \n\nThank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.", "claims_and_evidence": "The theoretical claims are well supported by clear and detailed proofs, and the authors also provide experimental validation of their theory in a controlled setting.", "methods_and_evaluation_criteria": "The main contributions of this work are theoretical, and the authors provided detailed proofs and brief experimental computation to validate the theory in a synthetic setting. I believe the methods and evaluation criteria make sense for the problem.", "theoretical_claims": "The theoretical findings are very detailed and greatly explain the training dynamics. The proofs are well thought, detailed, and clear. In my opinion, these are the main strengths of this work.", "experimental_designs_or_analyses": "The experimental setup is well explained and remains in a controlled setting. I believe the current submission would be strengthened with additional visualization to illustrate the different training dynamics and phases identified by the authors.", "relation_to_broader_scientific_literature": "I find that related work and prior works are well introduced and compared. The submission's contributions are interesting and are part of a growing interest in the literature on the theoretical understanding of transformers. The novelty seems to come mostly from the slightly more general model considered and the study of language recognition tasks with truncated CoT framework. The max-margin solution seems connected to [1] that showed how gradient-descent on softmax model converges in direction to a max-margin solution that separates\nlocally-optimal tokens from non-optimal ones. \n\n[1] Tarzanagh et al. Max-Margin Token Selection in Attention Mechanism, NeurIPS 2023", "essential_references_not_discussed": "To the best of my knowledge, there are no essential references not discussed in the current submission.", "other_strengths_and_weaknesses": "**Strengths**\n\n- The paper is very well-written and motivated\n- The problem tackled is of great interest\n- I appreciate that the model considered is (slightly) more general than prior works [1], which leads to a more involved analysis\n- The findings are interesting and well explained to grasp their impact on our understanding of the training dynamics\n- The proofs of clear and elegant\n\n**Weakness**\n\nI list below what I think are weaknesses, but I would be happy to be corrected if I misunderstood some important aspects of the authors' contributions.\n\n- While the authors' approach is indeed more general thatn [1] because of the linear layer, it is still less general that [2] since the linear layer is incorporated the value matrix $W_V$ in the optimization which amounts to simply consider an attention module and not a transformer: typically feed-forward blocks have non-linearities as it was considered in [2].\n- I find some part of the setting oversimplified compared to more practical settings and some of these simplifications are not well-motivated in my opinion (see questions). \n- The current submission lacks additional visualization or empirical validation since the considered setting seems a little bit ad-hoc and remains simple. It would be interesting to see how the figures vary with other hyperparameters for the controlled setting.\n\nOverall, I find the paper interesting and the analysis well conducted. I believe this is valuable work that can be improved with more visualization and variety in the controlled setting to illustrate the theoretical findings better. This is the reason for my current score, but I remain open to modifying my score, provided the authors clarify the points mentioned in the weaknesses section.\n\n*References*\n\n[1] Kim, J. and Suzuki, T. Transformers provably solve parity efficiently with chain of thought. ICLR 2025\n\n[2] Wen et al. From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency, arxiv 2024\n\n**Update after rebutal**: increased score from 3 to 4.", "other_comments_or_suggestions": "I list below some potential typos:\n\n- l. 121, second column: \"Frobenious\" --> \"Frobenius\"\n- l 400, first column: \"thee\" --> \"the\"", "questions_for_authors": "1) It seems that the embedding strategy described from l.171 only depends on the position in the sequence and not on the token value. Could the author clarify this point? If this is the case, then I believe it is not very realistic since for the tasks considered in the submission, the value of the token matters a lot (and it is, of course, the case in more practical settings).\n\n2) In Theorem 41, the authors propose to choose $\\lambda = \\Omega(L_{max}^2)$ instead of the common $\\sqrt{d}$ in the vanilla transformer (motivated to avoid small gradients when the magnitude of logits increases). However, the order of $L_{max}$ and $d$ in common transformers is typically $\\approx 10^2$ (time series forecasting [1], ViT [2], BERT [3], etc.). It means replacing a scaling of order $10^1$ by a scaling of order $10^4$. I believe this will lead experimentally to shrinking the logits range and outputting almost uniform attention values. Could the authors elaborate on that?\n\n3) Related to the point above, I notice that in their experiments, the authors take $L_{max} = 6$, which seems quite small, given again the type of sequences processed by transformers in more practical settings, and $\\lambda=2$. This does not seem to match the predicted $\\lambda=\\Omega(L_{max}^2)$ of Theorem 4.1. Maybe I misunderstood the setting, but could the authors clarify that point?\n\n4) I am intrigued by the gradient descent in two phases considered in the analysis. Why did the authors choose this framework instead of a more classical GD, SGD, or more adapted Adam (noting that adaptive optimizers are better for transformer and attention-based models [4])? \n\n5) Related to the point above, how did the author select the value $t_0$? It seems that the optimization can be heavily impacted because of that, especially given that the parameter $\\lambda$ is used in front of the learning rate. Could the authors elaborate on that?\n\n*References*\n\n[1] Ilbert et al. SAMformer: Unlocking the Potential of Transformers in Time Series Forecasting with Sharpness-Aware Minimization and Channel-Wise Attention, ICML 2024\n\n[1] Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR 2021\n\n[3] Devlin et al. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, NAACL 2019\n\n[4] Zhang et al. Why are adaptive methods good for attention models? NeurIPS 2020", "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": "vtbGDrqtdp", "reviewer": "Reviewer_yZje", "summary": "This paper presents a detailed theoretical analysis of how a one-layer transformer learns two sequence recognition/classification tasks: even pairs and parity check. The analysis decomposes the factors driving the attention weights and token score, with a discussion of the training dynamics in detail (e.g. attention weight reliance on different tokens, linear layer performing a max-margin solution).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments. Please note that all our new experiment results (i.e., the figures we quote below) can be accessed via the link  https://anonymous.4open.science/r/icml2025-4BCC/Figures%20for%20ICML.pdf\n\nQ: Add more diagrams, or interleave the theoretical analyses with the empirical results. \n\nA: Thanks for the suggestion. We will add a diagram to illustrate what each training phase of the even pairs problem learn, and will add a diagram to illustrate how the CoT-based inference approach solves the parity check problem. We will also present the experiments of the learning phases for even pairs and parity check problems right after theorems for their training dynamics. \n\nQ: It was unclear to me if the two-phase phenomena is entirely emergent or tied to the two-phase gradient descent schedule, which from Figure 3 I take to be the former (emergent).\n\nA: Yes. As demonstrated by our additional experiments (Figure 4), the two-phase phenomena arises naturally even under a fixed stepsize for the entire gradient descent training process. \n\nThank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.", "claims_and_evidence": "They seem fine overall.", "methods_and_evaluation_criteria": "The two tasks seem appropriate and standard from the mentioned related work.", "theoretical_claims": "I found sections 4 and 5 a bit hard to follow, see notes below.", "experimental_designs_or_analyses": "The experiment in Section 6 makes sense.", "relation_to_broader_scientific_literature": "This work extends prior theoretical work on the learning dynamics of transformers on a more representative architecture and on multiple tasks.", "essential_references_not_discussed": "Not to my knowledge", "other_strengths_and_weaknesses": "This work is a bit outside of my expertise so I'll comment from a general perspective. The biggest concern I have is the accessibility of this work to a broader audience. While I get the gist of ideas, I find the paper a bit hard to follow.", "other_comments_or_suggestions": "Related to the accessibility point above, I think it would be great if the authors could add more diagrams, or interleave the theoretical analyses with the empirical results. These can help a more general reader understand the work and in turn interest a broader audience.", "questions_for_authors": "It was unclear to me if the two-phase phenomena is entirely emergent or tied to the two-phase gradient descent schedule, which from Figure 3 I take to be the former (emergent).", "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": "Sbtsbl7ff4", "reviewer": "Reviewer_fgjp", "summary": "The work shows a theoretical analysis studying how a single layer of a Transformer (more precisely, an attention layer with a linear layer on top) learns to solve \"even pairs\" and \"parity check\" - two regular language recognition tasks. The authors begin by analyzing the even pairs task, showing that the Transformer trained in two phases learns to solve the even pairs problem, and discuss the training dynamics. Next, the authors discuss how parity check can be solved with Chain-of-Thought (CoT), either by inference with a trained Transformer or by training with CoT data. The authors verify their results experimentally.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments. Please note that all our new experiment results (i.e., the figures we quote below) can be accessed via the link https://anonymous.4open.science/r/icml2025-4BCC/Figures%20for%20ICML.pdf\n\nQ1: Main takeaway about transformers: (i) How analysis of parity learning with CoT differs from previous results; (ii) How learning \"even pairs\" help with learning of parity check.\n\nA1: **(i)** There are two key differences of our analysis from previous theoretical analysis [1,2] of parity check. *First*, [1,2] analyze only the first one or three gradient descent steps, not the entire training process. They do not prove the convergence of the training algorithm as the number of iterations goes large. In contrast, we characterize the dynamics of the entire training process, with the loss converging to global minimum (nearly zero) (Theorems 5.1–5.3), which is a comprehensive analysis of training process. *Second*, prior studies requires all input data sequence to have the same length, whereas our approach allows inputs to have  different lengths -- a more realistic setting. Such a more general setting is much more challenging to analyze, because tokens appearing in different sequence lengths influence each other’s gradient updates. Our analysis explicitly captures these dependencies and provides a fine-grained analysis of individual token values, making our technique more generalizable to broader learning scenarios. \n\n**(ii)** To connect the two problems, the learning output of \"even pairs\" serves as the first step in CoT towards solving parity check. Hence, using \"even pairs\" as a regularizer in parity check loss can provide an initial momentum to start the training of parity check. The reviewer can refer to lines 330–347 right column to see more detailed explanation in terms of how parameter updates. \n\n**Further insights:** (a) Layer roles: Our analysis characterizes distinct roles of feed-forward and attention layers via their joint training dynamics. Specifically, the attention layer $W$ learns to capture token-level semantics—e.g. token equality—while the linear layer $u$ encodes positional information. These functional separations are expected to persist in deeper architectures, where different layers may specialize to recognize different information such as content or positional information. (b) Training process exhibits different phases, where the first phase is parameter alignment and the second phase enables growing parameter norm and loss fast decay. Such dynamics is also observed in real-world problems (see Figure 3). (c) Our developed techniques for analyzing joint training of linear and attention layers can be useful for studying more general transformers.\n\n[1] Kim, J. and Suzuki, T. Transformers provably solve parity efficiently with chain of thought. ICLR 2025\n\n[2] Wen et al. From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency, arXiv 2024\n\nQ2: (i) It is unclear whether Algorithm 1 is interesting on its own. (ii) Why adding the regularization is needed, and how does the \"even pairs\" regularization improves training.\n\n\nA2: **(i):** Our Approach 1 in Section 5.1 is new and of independent interest. Existing works studying CoT typically require training with CoT data [1,2]. In contrast, our Approach 1 shows that transformers trained on even pairs can solve parity check by simply using CoT at **inference time** without additional training. \n\n**(ii):** The output of \"even pairs” algorithm serves as the first step in CoT towards solving parity check. Hence, using \"even pairs\" as regularization provides initial momentum to the training process for parity check. Also note that such a regularization approach is equivalent to changing the data distribution, i.e., mixing the even pairs data and parity data together and training the transformer model. \n\nQ3: Does the two-phase learning happen with a standard learning-rate schedule?\n\nA3: Yes. For vanilla GD with fixed stepsize throughout the training process, our new experiments (Figure 4) show that the training still exhibits the similar two-phase process as characterized in our theorem.  \n\nQ4: Connecting theoretical conclusions to real-world problems through additional experiments\n\nA4: We provide new experiments in Figure 3 on real-world dataset 'shakespeare' with deeper and realistic transformers (nanoGPT). It can be observed that the training  with vanilla Adam optimizer also exhibits a two-phase learning curve, i.e., the parameter norm grows fast in phase 1 and slows down in phase 2. \n\n**Generalization:** Our analysis framework can be extended to study general regular language problems. More details are provided in A4 of our response to Reviewer TZ3J.\n\nThank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.", "claims_and_evidence": "I think that the theoretical analysis of Transformers trained on the tasks studied in the paper is novel, and the results are interesting in the context of understanding training dynamics and capabilities of Transformers for solving language learning tasks. However, there are some flaws in the paper that I believe the authors need to address:\n- My main concern is that, beyond showing nice theoretical analysis and some study of training dynamics, it is unclear what is the main takeaway from the paper. What do these results teach us about what Transformers can do, or about how they learn, that we didn't know before? One conclusion that is addressed in the discussion is the importance of CoT for learning complex tasks like parity check. However, this seems to overlap with results shown in previous works, and it is unclear to me how the analysis of parity learning with CoT shown in this work differs from previous results of similar flavor. Another interesting result is that learning \"even pairs\" can help with learning parity check, but this is not shown directly and not discussed. I believe that stating clearly what is the main novel conclusion from the theoretical analysis and what it informs us about learning with Transformers can greatly improve the paper.\n- Related to the above, I believe that the relation between learning \"even pairs\" and \"parity check\" is interesting, but it is not studied in a setting that captures learning with language models. Specifically, the result for using transformers trained on even pairs to compute parity during inference seems particularly synthetic, and it is unclear whether this is just a way to introduce the next result in the section, or whether this results is interesting on its own. For the analysis of the second approach, adding the regularization with respect to the even pairs appears synthetic and it is unclear why it is needed. Can a similar result be shown by just changing the data mixture (e.g., training on a mixture of \"even pairs\" and \"parity check\" task, and testing on \"parity check\" with CoT)? How does the \"even pairs\" regularization improves training, given other results showing that parity learning is possible just with CoT?\n- The authors emphasize that learning happens in two phases, but it seems that the phases arise from the change of hyper-parameters at some point in training. While this is reasonable for the theoretical analysis, I believe that the authors should discuss whether the two-phase learning should happen with a standard learning-rate schedule. In particular, I think that experimental results in a standard setting (standard, maybe even deep, Transformer, standard learning rate schedule etc.) would be helpful.\n- Additionally, I think that connecting the theoretical conclusions to real-world problems (even somewhat synthetic problems) through additional experiments could be a good way to make this work more appealing to a broader audience.\n\nTo summarize, I believe that the theoretical setting and results are interesting, but the bottom-line conclusion is unclear, and some of the results are perhaps too synthetic.", "methods_and_evaluation_criteria": "See above.", "theoretical_claims": "See above.", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "See above.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See above.", "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:

Meta-review: This paper examines the training dynamics of transformers for the even pair and parity check problems. All reviewers acknowledge the theoretical contributions of this paper, including the joint training dynamic analysis and the clean theoretical setup for regular language recognition problems. The major concerns raised by the reviewers focus on the simplicity of the theoretical setting and its generalization to more complex scenarios, such as multi-layer transformers and other challenging tasks.

I agree that the contribution is somewhat limited given its idealized setting and simplified theoretical analysis. However, as a theoretical work, this paper is self-contained and explores new problems with novel techniques, which meets the acceptance criteria. Therefore, I recommend acceptance.
Decision: Accept (poster)