Training attention before output layer lags in transformers on long inputs
Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning
Machine Learning
Summary
This paper looks at how transformers learn to process very long inputs, like sentences or code of unlimited length. The authors found that when training a simple transformer to recognize certain structured languages, the output layer learns important features first, and only then does the attention mechanism improve. This is surprising because it was thought attention should be learned first. They also show that this delayed learning of attention helps in other language and learning tasks.
What this means in practice
- •For language model engineers: Develop training schedules that first update output layers before attention for better learning of language patterns in long inputs.
- •For machine learning practitioners: Improve regression and in-context learning tasks by adopting phased updates that delay attention training until after output layer tuning.
Authors
Yaomengxi Han, Debarghya Ghoshdastidar
Abstract
The success of transformers is attributed to several factors, including their capacity to express formal languages over variable-length inputs, and ability of the attention mechanism to focus on relevant input features. However, there is limited understanding of the training dynamics of shallow transformers, specifically how to decouple the learning of the attention and the output (feedforward) layer. A natural extension of the theory of one-step gradient feature learning in feedforward networks suggests that one could first learn the attention weights with a single gradient step, and then train the output layer. We show that this fails for a shallow transformer trained to recognize a star-free language from unbounded-length input. We prove that the transformer can be trained to recognize the language with a three-phase gradient descent: the output layer is updated first, followed by an attention update, and the output layer is fine-tuned in the third phase. In contrast, the language cannot be learned by training only the attention or output layer, or by a single attention-layer update followed by finetuning the output layer. We empirically show that ``attention learning lagging output layer learning" is beneficial in other language recognition, regression, and in-context learning tasks as well.