How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization

TL;DR

This paper investigates how transformers trained on in-context recall tasks learn, converge, and generalize, providing theoretical proofs of optimality, convergence rates, and out-of-distribution generalization, supported by empirical validation.

Contribution

It offers the first formal analysis of the training dynamics and generalization of transformers on in-context recall tasks, proving Bayes-optimality and convergence properties.

Findings

Transformers with linear, ReLU, or softmax attention are Bayes-optimal for recall tasks.

Expected loss converges linearly to Bayes risk during training.

Trained transformers can generalize out-of-distribution, but larger models may fail without proper parameterization.

Abstract

We study the approximation capabilities, convergence speeds and on-convergence behaviors of transformers trained on in-context recall tasks -- which requires to recognize the \emph{positional} association between a pair of tokens from in-context examples. Existing theoretical results only focus on the in-context reasoning behavior of transformers after being trained for the \emph{one} gradient descent step. It remains unclear what is the on-convergence behavior of transformers being trained by gradient descent and how fast the convergence rate is. In addition, the generalization of transformers in one-step in-context reasoning has not been formally investigated. This work addresses these gaps. We first show that a class of transformers with either linear, ReLU or softmax attentions, is provably Bayes-optimal for an in-context recall task. When being trained with gradient descent, we show via a finite-sample analysis that the expected loss converges at linear rate to the Bayes risks. Moreover, we show that the trained transformers exhibit out-of-distribution (OOD) generalization, i.e., generalizing to samples outside of the population distribution. Our theoretical findings are further supported by extensive empirical validations, showing that \emph{without} proper parameterization, models with larger expressive power surprisingly \emph{fail} to generalize OOD after being trained by gradient descent.