On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
It provides incremental theoretical insights into in-context learning mechanisms for researchers in machine learning theory.
The paper tackles the theoretical understanding of transformers' in-context learning by analyzing training dynamics for Gaussian mixture classification, showing that a single-layer transformer converges linearly to an optimal model and that prediction error decreases with longer prompts.
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.