Search papers, labs, and topics across Lattice.
This paper presents a Bayesian theory of feature learning in attention mechanisms, specifically analyzing how the copy subcircuit in induction heads emerges during training. By deriving a closed-form posterior over the attention matrix and reducing it to a low-dimensional order parameter space, the authors identify a first-order phase transition in softmax attention as training data increases, contrasting it with the second-order phase transition observed in linear attention. The findings provide a theoretical framework that explains the abrupt emergence of specific attention patterns, which is crucial for understanding in-context learning in transformers.
A first-order phase transition in softmax attention reveals how copy mechanisms emerge abruptly during training, challenging existing notions of gradual feature learning.
Attention is the key mechanism underlying in-context learning in transformers, and attention patterns have been observed empirically to emerge abruptly during training. We present a Bayesian theory of feature learning in attention; we then focus on how the copy subcircuit in the first layer of an induction head is learned by analyzing a single-layer softmax attention network trained on a copy task. We derive a closed-form posterior over the attention matrix and reduce it to a low-dimensional order parameter space. This reduction reveals a phase transition in the amount of training data, which we verify using both Bayesian sampling and standard training with Adam. We contrast our results with linear attention and find that softmax attention exhibits a \emph{first-order phase transition} while in linear attention an initial \emph{second-order phase transition} is followed by a smooth, continuous evolution toward the structured attention pattern (\emph{crossover}). Our work provides a first-principles theoretical account of the abrupt emergence of the copy subcircuit, reminiscent of the one observed in training large language models.
