On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures

Wei Shen, Ruida Zhou, Jing Yang, Cong Shen

TL;DR: Transformers trained via gradient descent can provably perform in-context classification of Gaussian mixtures.

Abstract: Although transformers have demonstrated impressive capabilities for in-context learning (ICL) in practice, theoretical understanding of the underlying mechanism that allows transformers to perform ICL is still in its infancy. This work aims to theoretically study the training dynamics of transformers for in-context classification tasks. We demonstrate that, for in-context classification of Gaussian mixtures under certain assumptions, a single-layer transformer trained via gradient descent converges to a globally optimal model at a linear rate. We further quantify the impact of the training and testing prompt lengths on the ICL inference error of the trained transformer. We show that when the lengths of training and testing prompts are sufficiently large, the prediction of the trained transformer approaches the ground truth distribution of the labels. Experimental results corroborate the theoretical findings.

Lay Summary: Transformers are powerful models known for their ability to learn new tasks by observing just a few examples and without updating their parameters—a capability known as in-context learning (ICL). While this ability performs well in practice, theoretical understanding of its underlying mechanisms is still in its infancy. This study takes a step toward addressing that gap by analyzing how transformers are trained to acquire the ability to perform in-context classification of different groups of data (Gaussian mixtures). We show that, under certain conditions, a simple one-layer transformer can successfully learn such tasks. We also provide a theoretical explanation of how using longer examples during both training and testing can help the model make more accurate predictions.

Primary Area: Theory->Learning Theory

Keywords: In-context learning, Transformer

Submission Number: 8550