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

TL;DR

This paper provides a theoretical analysis of how single-layer transformers trained with gradient descent can effectively perform in-context classification of Gaussian mixtures, with insights into prompt length effects and convergence rates.

Contribution

It offers the first theoretical convergence analysis of transformers for in-context classification of Gaussian mixtures, including prompt length impact.

Findings

Transformer trained via gradient descent converges to a global optimum at a linear rate.

Longer training and testing prompts improve inference accuracy.

Experimental results support the theoretical predictions.

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.