Scan and Snap: Verständnis der Trainingsdynamik und Token-Zusammensetzung in 1-schichtigen Transformern

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

May 25, 2023

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

cs.AI

Zusammenfassung

Die Transformer-Architektur hat in zahlreichen Forschungsbereichen beeindruckende Leistungen gezeigt und ist zum Rückgrat vieler neuronaler Netzwerkmodelle geworden. Allerdings ist das Verständnis darüber, wie sie funktioniert, noch begrenzt. Insbesondere bleibt es ein Rätsel, wie sich die Repräsentation durch die Gradienten-Trainingsdynamik bei einem einfachen prädiktiven Verlust entwickelt. In diesem Artikel analysieren wir die SGD-Trainingsdynamik eines 1-Schicht-Transformers, bestehend aus einer Self-Attention-Schicht und einer Decoder-Schicht, für die Aufgabe der Vorhersage des nächsten Tokens auf mathematisch rigorose Weise. Wir öffnen die Blackbox des dynamischen Prozesses, wie die Self-Attention-Schicht Eingabe-Tokens kombiniert, und enthüllen die Natur des zugrunde liegenden induktiven Bias. Genauer gesagt beweisen wir unter den Annahmen (a) keine Positionskodierung, (b) lange Eingabesequenzen und (c) dass die Decoder-Schicht schneller lernt als die Self-Attention-Schicht, dass Self-Attention als ein diskriminierender Scan-Algorithmus agiert: Ausgehend von gleichmäßiger Aufmerksamkeit konzentriert sie sich schrittweise stärker auf bestimmte Schlüssel-Tokens für ein spezifisches, vorherzusagendes nächstes Token und schenkt weniger Aufmerksamkeit gemeinsamen Schlüssel-Tokens, die über verschiedene nächste Tokens hinweg auftreten. Unter den unterschiedlichen Tokens reduziert sie die Aufmerksamkeitsgewichte progressiv in der Reihenfolge von niedriger zu hoher Kookkurrenz zwischen dem Schlüssel- und dem Query-Token im Trainingsdatensatz. Interessanterweise führt dieser Prozess nicht zu einem Winner-Takes-All-Ergebnis, sondern verlangsamt sich aufgrund eines Phasenübergangs, der durch die Lernraten der beiden Schichten kontrollierbar ist, und hinterlässt (fast) feste Token-Kombinationen. Wir bestätigen diese \emph{Scan-and-Snap}-Dynamik anhand von synthetischen und realen Daten (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).