[2309.05858] Uncovering mesa-optimization algorithms in Transformers

Abstract:Transformers have become the dominant model in deep learning, but the reason for their superior performance is poorly understood. Here, we hypothesize that the strong performance of Transformers stems from an architectural bias towards mesa-optimization, a learned process running within the forward pass of a model consisting of the following two steps: (i) the construction of an internal learning objective, and (ii) its corresponding solution found through optimization. To test this hypothesis, we reverse-engineer a series of autoregressive Transformers trained on simple sequence modeling tasks, uncovering underlying gradient-based mesa-optimization algorithms driving the generation of predictions. Moreover, we show that the learned forward-pass optimization algorithm can be immediately repurposed to solve supervised few-shot tasks, suggesting that mesa-optimization might underlie the in-context learning capabilities of large language models. Finally, we propose a novel self-attention layer, the mesa-layer, that explicitly and efficiently solves optimization problems specified in context. We find that this layer can lead to improved performance in synthetic and preliminary language modeling experiments, adding weight to our hypothesis that mesa-optimization is an important operation hidden within the weights of trained Transformers.

Authors:Johannes von Oswald, Maximilian Schlegel, Alexander Meulemans, Seijin Kobayashi, Eyvind Niklasson, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymyrov, Razvan Pascanu, João Sacramento View PDF HTML (experimental) 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…

saved by

• Aleksi Maunu

related reading

• [2212.07677] Transformers learn in-context by gradient descentarxiv.org
• In-context Learning and Induction Headstransformer-circuits.pub
• A Mathematical Framework for Transformer Circuitstransformer-circuits.pub
• Transformer Circuits Threadtransformer-circuits.pub
• Transformers from Scratche2eml.school
• NL.pdfabehrouz.github.io
• The Annotated Transformernlp.seas.harvard.edu
• What Can Transformers Learn In-Context?arxiv.org
• [2402.01258] Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscapearxiv.org
• [1906.01820] Risks from Learned Optimization in Advanced Machine Learning Systemsarxiv.org
• Thinking like Transformersrush.github.io
• Learning without training: The implicit dynamics of in-context learningarxiv.org