Escanear y Capturar: Comprender la Dinámica del Entrenamiento y la Composición de Tokens en Transformadores de 1 Capa

Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

May 25, 2023

Autores: Yuandong Tian, Yiping Wang, Beidi Chen, Simon Du

cs.AI

Resumen

La arquitectura Transformer ha demostrado un rendimiento impresionante en múltiples dominios de investigación y se ha convertido en la base de muchos modelos de redes neuronales. Sin embargo, existe una comprensión limitada sobre cómo funciona. En particular, con una simple pérdida predictiva, cómo emerge la representación a partir de la dinámica de entrenamiento por gradiente sigue siendo un misterio. En este artículo, para un Transformer de una capa con una capa de autoatención más una capa decodificadora, analizamos su dinámica de entrenamiento con SGD para la tarea de predicción del siguiente token de manera matemáticamente rigurosa. Abrimos la caja negra del proceso dinámico de cómo la capa de autoatención combina los tokens de entrada y revelamos la naturaleza del sesgo inductivo subyacente. Más específicamente, bajo los supuestos de (a) sin codificación posicional, (b) secuencia de entrada larga y (c) que la capa decodificadora aprende más rápido que la capa de autoatención, demostramos que la autoatención actúa como un algoritmo de escaneo discriminativo: partiendo de una atención uniforme, gradualmente presta más atención a tokens clave distintos para un siguiente token específico que se va a predecir, y presta menos atención a tokens clave comunes que aparecen en diferentes tokens siguientes. Entre los tokens distintos, reduce progresivamente los pesos de atención, siguiendo el orden de baja a alta co-ocurrencia entre el token clave y el token de consulta en el conjunto de entrenamiento. Curiosamente, este procedimiento no conduce a un comportamiento de "el ganador se lo lleva todo", sino que se desacelera debido a una transición de fase que es controlable por las tasas de aprendizaje de las dos capas, dejando una combinación de tokens (casi) fija. Verificamos esta dinámica de \emph{escaneo y ajuste} en datos sintéticos y del mundo real (WikiText).

English

Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss, how the representation emerges from the gradient training dynamics remains a mystery. In this paper, for 1-layer transformer with one self-attention layer plus one decoder layer, we analyze its SGD training dynamics for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a discriminative scanning algorithm: starting from uniform attention, it gradually attends more to distinct key tokens for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a phase transition that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this \emph{scan and snap} dynamics on synthetic and real-world data (WikiText).