Descubriendo algoritmos de mesa-optimización en Transformers

Uncovering mesa-optimization algorithms in Transformers

September 11, 2023

Autores: 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

Resumen

Los Transformers se han convertido en el modelo dominante en el aprendizaje profundo, pero la razón de su rendimiento superior es poco comprendida. Aquí, planteamos la hipótesis de que el fuerte rendimiento de los Transformers surge de un sesgo arquitectónico hacia la mesa-optimización, un proceso aprendido que se ejecuta dentro del pase hacia adelante de un modelo y que consta de los siguientes dos pasos: (i) la construcción de un objetivo de aprendizaje interno, y (ii) su solución correspondiente encontrada mediante optimización. Para probar esta hipótesis, invertimos ingeniería en una serie de Transformers autorregresivos entrenados en tareas simples de modelado de secuencias, descubriendo algoritmos subyacentes de mesa-optimización basados en gradientes que impulsan la generación de predicciones. Además, demostramos que el algoritmo de optimización aprendido durante el pase hacia adelante puede reutilizarse inmediatamente para resolver tareas supervisadas de pocos ejemplos, lo que sugiere que la mesa-optimización podría ser la base de las capacidades de aprendizaje en contexto de los modelos de lenguaje grandes. Finalmente, proponemos una nueva capa de autoatención, la mesa-capa, que resuelve explícita y eficientemente problemas de optimización especificados en contexto. Encontramos que esta capa puede mejorar el rendimiento en experimentos sintéticos y preliminares de modelado de lenguaje, reforzando nuestra hipótesis de que la mesa-optimización es una operación importante oculta dentro de los pesos de los Transformers entrenados.

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.