title: Scan and Snap: Understanding Training Dynamics and Token Composition in
    1-layer Transformer
summary: 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).
paper_authors._id: 64742ee8bf9b32c6bbcd84f2
paper_authors.name: Yuandong Tian
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-05-31T09:10:15.194Z
paper_authors._id: 6344cf73ee1504dbcd5bdfe7
paper_authors.avatarUrl: /avatars/6dd2bf1f9c5679e5c8c85d62c9836aac.svg
paper_authors.fullname: Yuandong Tian
paper_authors.type: user
paper_authors.user: tydsh
paper_authors._id: 64742ee8bf9b32c6bbcd84f3
paper_authors.name: Yiping Wang
paper_authors._id: 64742ee8bf9b32c6bbcd84f4
paper_authors.name: Beidi Chen
paper_authors._id: 64742ee8bf9b32c6bbcd84f5
paper_authors.name: Simon Du
paper_title: Scan and Snap: Understanding Training Dynamics and Token Composition in
    1-layer Transformer
paper_summary: 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).
paper_ai_keywords: Transformer architecture predictive loss representation emergence gradient training dynamics self-attention layer decoder layer SGD training dynamics next token prediction mathematical rigor black box inductive bias positional encoding long input sequence discriminative scanning algorithm uniform attention key tokens query token co-occurrence winner-takes-all phase transition learning rates scan and snap dynamics