Transformers Provably Learn Two-Mixture of Linear Classification via Gradient Flow
Keywords: transformer, training dynamics, gradient flow, mixture of linear classification
TL;DR: We study the gradient flow dynamics of training transformer on two-mixture of linear classification.
Abstract: Understanding how transformers learn and utilize hidden connections between tokens is crucial to understand the behavior of large language models.
To understand this mechanism, we consider the task of two-mixture of linear classification which possesses a hidden correspondence structure among tokens, and study the training dynamics of a symmetric two-headed transformer with ReLU neurons.
Motivated by the stage-wise learning phenomenon in our experiments, we design and theoretically analyze a three-stage training algorithm, which can effectively characterize the actual gradient descent dynamics when we simultaneously train the neuron weights and the softmax attention.
The first stage is a neuron learning stage, where the neurons align with the underlying signals.
The second stage is a attention feature learning stage, where we analyze the feature learning process of how the attention learns to utilize the relationship between the tokens to solve certain hard samples.
In the meantime, the attention features evolve from a nearly non-separable state (at the initialization) to a well-separated state.
The third stage is a convergence stage, where the population loss is driven towards zero.
The key technique in our analysis of softmax attention is to identify a critical sub-system inside a large dynamical system and bound the growth of the non-linear sub-system by a linear system.
Finally, we discuss the setting with more than two mixtures.
We empirically show the difficulty of generalizing our analysis of the gradient flow dynamics to the case even when the number of mixtures equals three, although the transformer can still successfully learn such distribution.
On the other hand, we show by construction that there exists a transformer that can solve mixture of linear classification given any arbitrary number of mixtures.
Primary Area: learning theory
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 7208
Loading