Descobrindo algoritmos de mesa-otimização em 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

Resumo

Os Transformers tornaram-se o modelo dominante no aprendizado profundo, mas a razão para seu desempenho superior é pouco compreendida. Aqui, hipotetizamos que o forte desempenho dos Transformers decorre de um viés arquitetônico em direção à mesa-otimização, um processo aprendido que ocorre durante a passagem direta de um modelo e consiste nos seguintes dois passos: (i) a construção de um objetivo de aprendizado interno, e (ii) sua solução correspondente encontrada por meio de otimização. Para testar essa hipótese, realizamos uma engenharia reversa de uma série de Transformers autorregressivos treinados em tarefas simples de modelagem de sequências, revelando algoritmos subjacentes de mesa-otimização baseados em gradiente que impulsionam a geração de previsões. Além disso, mostramos que o algoritmo de otimização aprendido durante a passagem direta pode ser imediatamente reutilizado para resolver tarefas supervisionadas de poucos exemplos, sugerindo que a mesa-otimização pode estar na base das capacidades de aprendizado em contexto de grandes modelos de linguagem. Por fim, propomos uma nova camada de autoatenção, a mesa-camada, que resolve explicitamente e de forma eficiente problemas de otimização especificados em contexto. Descobrimos que essa camada pode levar a melhorias de desempenho em experimentos sintéticos e preliminares de modelagem de linguagem, reforçando nossa hipótese de que a mesa-otimização é uma operação importante oculta nos pesos dos Transformers treinados.

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.