Scan and Snap: Understanding Training Dynamics and Token Composition in
  1-layer Transformer
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).
64742ee8bf9b32c6bbcd84f2
Yuandong Tian
admin_assigned
2023-05-31T09:10:15.194Z
6344cf73ee1504dbcd5bdfe7
/avatars/6dd2bf1f9c5679e5c8c85d62c9836aac.svg
Yuandong Tian
user
tydsh
64742ee8bf9b32c6bbcd84f3
Yiping Wang
64742ee8bf9b32c6bbcd84f4
Beidi Chen
64742ee8bf9b32c6bbcd84f5
Simon Du
Scan and Snap: Understanding Training Dynamics and Token Composition in
  1-layer Transformer
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).
The analysis of SGD training dynamics in a one-layer transformer reveals that self-attention operates as a discriminative scanning algorithm, focusing on distinct tokens for prediction and demonstrating a phase transition controlled by layer learning rates.
Transformer architecture self-attention layer decoder layer SGD training dynamics inductive bias positional encoding long input sequence scan and snap phase transition