Scan e Snap: Compreendendo a Dinâmica de Treinamento e Composição de Tokens em Transformers de 1 Camada

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

Resumo

A arquitetura Transformer tem demonstrado desempenho impressionante em múltiplos domínios de pesquisa e se tornou a base de muitos modelos de redes neurais. No entanto, há uma compreensão limitada sobre como ela funciona. Em particular, com uma simples função de perda preditiva, como a representação emerge da dinâmica de treinamento por gradiente permanece um mistério. Neste artigo, para um Transformer de uma camada composto por uma camada de autoatenção mais uma camada decodificadora, analisamos sua dinâmica de treinamento com SGD (Descida de Gradiente Estocástica) para a tarefa de previsão do próximo token de forma matematicamente rigorosa. Abrimos a caixa preta do processo dinâmico de como a camada de autoatenção combina os tokens de entrada e revelamos a natureza do viés indutivo subjacente. Mais especificamente, com as suposições de (a) ausência de codificação posicional, (b) sequência de entrada longa e (c) a camada decodificadora aprendendo mais rápido que a camada de autoatenção, provamos que a autoatenção age como um algoritmo de varredura discriminativo: partindo de uma atenção uniforme, ela gradualmente passa a atender mais a tokens-chave distintos para um próximo token específico a ser previsto, e presta menos atenção a tokens-chave comuns que ocorrem em diferentes próximos tokens. Entre os tokens distintos, ela progressivamente reduz os pesos de atenção, seguindo a ordem de baixa para alta co-ocorrência entre o token-chave e o token de consulta no conjunto de treinamento. Curiosamente, esse procedimento não leva a um cenário de "vencedor leva tudo", mas desacelera devido a uma transição de fase que é controlável pelas taxas de aprendizado das duas camadas, resultando em uma combinação de tokens (quase) fixa. Verificamos essa dinâmica de \emph{varredura e fixação} em dados sintéticos e do 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).