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

TL;DR

This paper provides a theoretical analysis of how one-layer transformers learn regular language recognition tasks, revealing training phases, implicit biases, and the role of Chain-of-Thought in solving parity checks.

Contribution

It offers a novel theoretical framework analyzing training dynamics and implicit bias of transformers on regular language tasks, including the integration of Chain-of-Thought.

Findings

Attention layer rapidly separates data in early training

Linear layer approaches max-margin hyperplane logarithmically

Training loss decreases at a rate of O(1/t)

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.