Découvrir les algorithmes de méso-optimisation dans les 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

Résumé

Les Transformers sont devenus le modèle dominant en apprentissage profond, mais la raison de leur performance supérieure reste mal comprise. Nous émettons ici l'hypothèse que la forte performance des Transformers découle d'un biais architectural en faveur de la méso-optimisation, un processus appris s'exécutant lors de la passe avant d'un modèle et consistant en deux étapes : (i) la construction d'un objectif d'apprentissage interne, et (ii) la recherche de sa solution correspondante via une optimisation. Pour tester cette hypothèse, nous avons rétro-conçu une série de Transformers autorégressifs entraînés sur des tâches simples de modélisation de séquences, révélant des algorithmes de méso-optimisation basés sur le gradient qui pilotent la génération des prédictions. De plus, nous montrons que l'algorithme d'optimisation appris lors de la passe avant peut être immédiatement réutilisé pour résoudre des tâches supervisées en few-shot, suggérant que la méso-optimisation pourrait sous-tendre les capacités d'apprentissage en contexte des grands modèles de langage. Enfin, nous proposons une nouvelle couche d'auto-attention, la méso-couche, qui résout explicitement et efficacement des problèmes d'optimisation spécifiés en contexte. Nous constatons que cette couche peut améliorer les performances dans des expériences synthétiques et préliminaires de modélisation du langage, renforçant notre hypothèse selon laquelle la méso-optimisation est une opération importante cachée dans les poids des Transformers entraînés.

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.