Mechanics of Next Token Prediction with Self-Attention
Yingcong Li Yixiao Huang M. Emrullah Ildiz Ankit Singh Rawat Samet Oymak
Transformer-based language models are trained on large datasets to predict
the next token given an input sequence. Despite this simple training objective,
they have led to revolutionary advances in natural language processing.
Underlying this success is the self-attention mechanism. In this work, we ask:
$\textit{What}$ $\textit{does}$ $\textit{a}$ $\textit{single}$
$\textit{self-attention}$ $\textit{layer}$ $\textit{learn}$ $\textit{from}$
$\textit{next-token}$ $\textit{prediction?}$ We show that training
self-attention with gradient descent learns an automaton which generates the
next token in two distinct steps: $\textbf{(1)}$ $\textbf{Hard}$
$\textbf{retrieval:}$ Given input sequence, self-attention precisely selects
the $\textit{high-priority}$ $\textit{input}$ $\textit{tokens}$ associated with
the last input token. $\textbf{(2)}$ $\textbf{Soft}$ $\textbf{composition:}$ It
then creates a convex combination of the high-priority tokens from which the
next token can be sampled. Under suitable conditions, we rigorously
characterize these mechanics through a directed graph over tokens extracted
from the training data. We prove that gradient descent implicitly discovers the
strongly-connected components (SCC) of this graph and self-attention learns to
retrieve the tokens that belong to the highest-priority SCC available in the
context window. Our theory relies on decomposing the model weights into a
directional component and a finite component that correspond to hard retrieval
and soft composition steps respectively. This also formalizes a related
implicit bias formula conjectured in [Tarzanagh et al. 2023]. We hope that
these findings shed light on how self-attention processes sequential data and
pave the path toward demystifying more complex architectures.