Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ruiquan Huang, Yingbin Liang, Jing Yang

TL;DR

The paper studies how a one-layer transformer learns regular language recognition tasks, specifically even pairs and parity check, under gradient descent with CoT. It reveals a two-phase dynamic: Phase 1 aligns token and attention scores to create separable representations, and Phase 2 drives the linear layer along a max-margin direction while keeping the attention pattern stable, resulting in a sublinear loss decay $O(1/t)$. For parity check, the work shows two CoT-based approaches, including zero-shot truncation inference and CoT training with teacher forcing, both producing provable learning behavior and convergence to max-margin solutions. The results provide a theoretical framework for joint attention-linear training dynamics, highlight the essential role of CoT in parity reasoning, and are supported by experiments on synthetic data validating the predicted parameter growth and alignment patterns.

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

TL;DR

The paper studies how a one-layer transformer learns regular language recognition tasks, specifically even pairs and parity check, under gradient descent with CoT. It reveals a two-phase dynamic: Phase 1 aligns token and attention scores to create separable representations, and Phase 2 drives the linear layer along a max-margin direction while keeping the attention pattern stable, resulting in a sublinear loss decay . For parity check, the work shows two CoT-based approaches, including zero-shot truncation inference and CoT training with teacher forcing, both producing provable learning behavior and convergence to max-margin solutions. The results provide a theoretical framework for joint attention-linear training dynamics, highlight the essential role of CoT in parity reasoning, and are supported by experiments on synthetic data validating the predicted parameter growth and alignment patterns.

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 . Our experiments validate those theoretical results.
Paper Structure (19 sections, 18 theorems, 142 equations, 7 figures, 1 algorithm)

This paper contains 19 sections, 18 theorems, 142 equations, 7 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Formulation
  4. Even Pairs Problem
  5. Phase 1: Rapid Growth of Token and Attention Scores
  6. Phase 2: Margin Maximization and Implicit Bias
  7. Parity Check Problem
  8. Approach 1: Inference via Truncated CoT
  9. Approach 2: Training with CoT under Teacher Forcing
  10. Training Dynamics of Approach 2
  11. Experiments
  12. Conclusion
  13. Auxiliary Lemmas and Equations
  14. Gradients Calculation
  15. Proofs of Even Pairs Problem
  16. ...and 4 more sections

Key Result

Theorem 4.1

Let $-w$ denote the flip of token $w \in \{\texttt{a}, \texttt{b}\}$. Choose $\lambda=\Omega(L_{\max}^2), t_0 = O(1/(\eta L_{\max}) )$, and $\eta=O(\min\{1/L_{\max},1/\lambda^{2/3}\})$. Then, for all $t \leq t_0$, the parameters evolve as follows: (1) The dynamics of linear layer $u$ is governed by (2) The dynamics of attention layer $W$ is governed by the following inequalities. For any length $

Figures (7)

  • Figure 1: An illustration of CoT on input abb, where $L_0 = 3$.
  • 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 in length-3 sequences.
  • 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.
  • Figure 4: Training dynamics of learning Even pairs when $\lambda=10$
  • Figure 5: Training dynamics of learning Even pairs when $\lambda=18=L_{\max}^2/2$
  • ...and 2 more figures

Theorems & Definitions (28)

  • Theorem 4.1: Phase 1
  • Definition 4.2: Separable dataset
  • Proposition 4.3
  • Theorem 4.4: Phase 2
  • Theorem 4.5: Convergence of loss
  • Theorem 5.1: Phase 1
  • Theorem 5.2: Phase 2
  • Theorem 5.3
  • Lemma 1.1: gao2017properties
  • Lemma 2.1
  • ...and 18 more