Scan e Snap: Comprendere le Dinamiche di Addestramento e la Composizione dei Token nei Trasformatori a 1 Strato

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

May 25, 2023

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

cs.AI

Abstract

L'architettura Transformer ha dimostrato prestazioni impressionanti in molteplici ambiti di ricerca ed è diventata il fulcro di numerosi modelli di reti neurali. Tuttavia, la comprensione del suo funzionamento è ancora limitata. In particolare, con una semplice funzione di perdita predittiva, rimane un mistero come la rappresentazione emerga dalle dinamiche di addestramento basate sui gradienti. In questo articolo, per un Transformer a un solo strato composto da un livello di self-attention e un livello di decodifica, analizziamo in modo matematicamente rigoroso le dinamiche di addestramento con SGD per il compito di predizione del token successivo. Apriamo la scatola nera del processo dinamico attraverso il quale il livello di self-attention combina i token di input e riveliamo la natura del bias induttivo sottostante. Più specificamente, assumendo (a) nessun encoding posizionale, (b) sequenze di input lunghe e (c) che il livello di decodifica apprenda più velocemente del livello di self-attention, dimostriamo che la self-attention agisce come un algoritmo di scansione discriminativo: partendo da un'attenzione uniforme, gradualmente si concentra maggiormente su token chiave distinti per un specifico token successivo da predire, e presta meno attenzione ai token chiave comuni che compaiono in diversi token successivi. Tra i token distinti, riduce progressivamente i pesi di attenzione, seguendo l'ordine dalla bassa alla alta co-occorrenza tra il token chiave e il token di query nel set di addestramento. Interessante è che questo procedimento non porta a una situazione di "winner-takes-all", ma rallenta a causa di una transizione di fase che è controllabile dai tassi di apprendimento dei due livelli, lasciando (quasi) fissa la combinazione di token. Verifichiamo questa dinamica di \emph{scansione e fissaggio} su dati sintetici e reali (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).