Scan and Snap : Comprendre la dynamique d'apprentissage et la composition des tokens dans un Transformer à une couche

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

May 25, 2023

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

cs.AI

Résumé

L'architecture Transformer a démontré des performances impressionnantes dans de multiples domaines de recherche et est devenue la pierre angulaire de nombreux modèles de réseaux neuronaux. Cependant, la compréhension de son fonctionnement reste limitée. En particulier, avec une simple fonction de perte prédictive, la manière dont la représentation émerge des dynamiques d'entraînement par gradient demeure un mystère. Dans cet article, pour un Transformer à une couche composé d'une couche d'auto-attention et d'une couche de décodage, nous analysons de manière mathématiquement rigoureuse ses dynamiques d'entraînement par descente de gradient stochastique (SGD) pour la tâche de prédiction du token suivant. Nous ouvrons la boîte noire du processus dynamique par lequel la couche d'auto-attention combine les tokens d'entrée, et révélons la nature du biais inductif sous-jacent. Plus précisément, en supposant (a) l'absence d'encodage positionnel, (b) des séquences d'entrée longues, et (c) que la couche de décodage apprend plus rapidement que la couche d'auto-attention, nous prouvons que l'auto-attention agit comme un algorithme de balayage discriminatif : partant d'une attention uniforme, elle se concentre progressivement davantage sur les tokens clés distincts pour un token spécifique à prédire, et accorde moins d'attention aux tokens clés communs qui apparaissent dans différents tokens suivants. Parmi les tokens distincts, elle réduit progressivement les poids d'attention, suivant l'ordre de faible à forte co-occurrence entre le token clé et le token de requête dans l'ensemble d'entraînement. Fait intéressant, ce processus ne conduit pas à une situation de type "le gagnant prend tout", mais ralentit en raison d'une transition de phase contrôlable par les taux d'apprentissage des deux couches, laissant une combinaison de tokens (presque) fixe. Nous vérifions cette dynamique de \emph{balayage et fixation} sur des données synthétiques et réelles (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).