Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent

Transformers have demonstrated remarkable in-context learning (ICL) capabilities. The strong ICL performance of transformers is commonly believed to arise from their ability to implicitly execute certain algorithms on the context, thereby enhancing prediction and generation. In this work, we investigate how transformers with softmax attention perform in-context learning on linear classification data. We first construct a class of multi-layer transformers that can perform in-context logistic regression, with each layer exactly performing one step of normalized gradient descent on an in-context loss. Then, we show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model. Training convergence guarantees of the self-attention layer and out-of-distribution generalization guarantees of the looped model are provided. Our results advance the theoretical understanding of ICL mechanism by showcasing how softmax transformers can effectively act as in-context learners.

Large language models can often learn from examples placed directly in the prompt, without updating their parameters. This ability is called in-context learning, but we still do not fully understand what kind of learning process happens inside the model. We study this question in a simplified classification setting where the underlying rule is mathematically clear. We show that a transformer with softmax attention can behave like an iterative learner: it uses the examples in the prompt to gradually improve an internal classifier. More specifically, this behavior matches a normalized form of gradient descent, a standard method for improving models during training. We also show when this mechanism can be learned from data and when it can generalize to new input distributions. These results give a precise explanation of one way transformers can perform in-context learning.