Uncovering mesa-optimization algorithms in Transformers

TL;DR

This paper investigates how Transformers trained on sequence prediction develop an internal optimization process, explaining in-context learning as a consequence of their training objective and revealing potential for designing new optimization-based layers.

Contribution

It uncovers the emergence of a subsidiary optimization algorithm within Transformers, linking in-context learning to gradient-based processes during training.

Findings

Transformers develop an internal optimization process during training.

This process explains in-context learning capabilities.

The findings suggest new directions for designing optimization-based Transformer layers.

Abstract

Some autoregressive models exhibit in-context learning capabilities: being able to learn as an input sequence is processed, without undergoing any parameter changes, and without being explicitly trained to do so. The origins of this phenomenon are still poorly understood. Here we analyze a series of Transformer models trained to perform synthetic sequence prediction tasks, and discover that standard next-token prediction error minimization gives rise to a subsidiary learning algorithm that adjusts the model as new inputs are revealed. We show that this process corresponds to gradient-based optimization of a principled objective function, which leads to strong generalization performance on unseen sequences. Our findings explain in-context learning as a product of autoregressive loss minimization and inform the design of new optimization-based Transformer layers.