How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias

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

Q1: 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.

A1: (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.

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

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.

[1] Kim, J. and Suzuki, T. Transformers provably solve parity efficiently with chain of thought. ICLR 2025

[2] Wen et al. From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency, arXiv 2024

Q2: (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.

A2: (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.

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

Q3: Does the two-phase learning happen with a standard learning-rate schedule?

A3: 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.

Q4: Connecting theoretical conclusions to real-world problems through additional experiments

A4: 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.

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.

Thank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.

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

Q: Add more diagrams, or interleave the theoretical analyses with the empirical results.

A: 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.

Q: 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).

A: 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.

Thank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.

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

Q1: It seems that the embedding strategy does not depend on the token value.

A1: 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.

Q2: In experiments, authors take $L_{\max}=6$ and $\lambda=2$, which does not match $\lambda=\Omega(L_{\max}^2)$ in Theorem 4.1.

A2: 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.

Q3: 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?

A3: 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.

Q4: 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])?

A4: 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.

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

Q5: How did the author select the value $t_0$?

A5: 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.

Thank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.

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

Q1: It is not clear how the results can be extended to deep transformers

A1: 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.

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

Q2: 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$

A2: We would like to point out that these assumptions can be relaxed or justified as follows.

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

Q3: The analysis may not be really for the joint training of both layers.

A3: 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. Specifically, 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.

Q4: I feel the analysis is too ad-hoc to the two regular languages studied in this paper.

A4: 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.

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

Q5: It would be better to phrase the method as data mixing rather than adding a regularization.

A5: Thanks for the suggestion! We will make the change or add discussions.

Thank you again for your insightful comments. We hope our responses addressed your concerns and would greatly appreciate your kind consideration in increasing your score.