Scan en Snap: Inzicht in Trainingsdynamiek en Tokencompositie in 1-laagse Transformers

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

Samenvatting

De Transformer-architectuur heeft indrukwekkende prestaties getoond in meerdere onderzoeksdomeinen en is de ruggengraat geworden van veel neurale netwerkmodellen. Er is echter beperkt begrip over hoe deze werkt. Met name blijft het een raadsel hoe de representatie ontstaat uit de gradienttrainingsdynamiek met een eenvoudig voorspellend verlies. In dit artikel analyseren we, voor een 1-laags Transformer met één self-attention-laag plus één decoder-laag, de SGD-trainingsdynamiek voor de taak van volgende tokenvoorspelling op een wiskundig rigoureuze manier. We openen de black box van het dynamische proces van hoe de self-attention-laag invoertokens combineert, en onthullen de aard van de onderliggende inductieve bias. Meer specifiek, onder de aannames (a) geen positionele codering, (b) lange invoerreeks, en (c) de decoder-laag leert sneller dan de self-attention-laag, bewijzen we dat self-attention fungeert als een discriminerend scan-algoritme: beginnend met uniforme aandacht, richt het zich geleidelijk meer op onderscheidende sleuteltokens voor een specifieke volgende token die moet worden voorspeld, en besteedt het minder aandacht aan gemeenschappelijke sleuteltokens die voorkomen bij verschillende volgende tokens. Onder onderscheidende tokens laat het progressief de aandachtgewichten dalen, volgens de volgorde van lage naar hoge co-voorkomens tussen de sleutel- en de querytoken in de trainingsset. Interessant genoeg leidt dit proces niet tot een winner-takes-all-situatie, maar vertraagt het door een faseovergang die beheersbaar is door de leerpercentages van de twee lagen, waardoor (bijna) vaste tokencombinaties overblijven. We verifiëren deze \emph{scan en snap}-dynamiek op synthetische en real-world data (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).