Het blootleggen van mesa-optimalisatie-algoritmen in Transformers

Uncovering mesa-optimization algorithms in Transformers

September 11, 2023

Auteurs: Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymyrov, Razvan Pascanu, João Sacramento

cs.AI

Samenvatting

Transformers zijn het dominante model geworden in deep learning, maar de reden voor hun superieure prestaties is slecht begrepen. Hier stellen we de hypothese op dat de sterke prestaties van Transformers voortkomen uit een architectonische voorkeur voor mesa-optimalisatie, een geleerd proces dat plaatsvindt tijdens de forward pass van een model en bestaat uit de volgende twee stappen: (i) de constructie van een intern leerdoel, en (ii) de bijbehorende oplossing die wordt gevonden door optimalisatie. Om deze hypothese te testen, reverse-engineeren we een reeks autoregressieve Transformers die zijn getraind op eenvoudige sequentiemodelleringstaken, waarbij we onderliggende gradient-gebaseerde mesa-optimalisatie-algoritmen ontdekken die de voorspellingen sturen. Bovendien laten we zien dat het geleerde optimalisatie-algoritme in de forward pass direct kan worden hergebruikt om supervised few-shot taken op te lossen, wat suggereert dat mesa-optimalisatie ten grondslag zou kunnen liggen aan de in-context leer capaciteiten van grote taalmodellen. Tot slot introduceren we een nieuwe self-attention laag, de mesa-laag, die expliciet en efficiënt optimalisatieproblemen oplost die in de context zijn gespecificeerd. We ontdekken dat deze laag kan leiden tot verbeterde prestaties in synthetische en voorlopige taalmodelleringsexperimenten, wat onze hypothese ondersteunt dat mesa-optimalisatie een belangrijke operatie is die verborgen zit in de gewichten van getrainde Transformers.

English

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.