Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers

• URL: http://arxiv.org/abs/2409.10559v1
• Date: Mon, 9 Sep 2024 18:10:26 GMT
• Title: Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers
• Authors: Siyu Chen, Heejune Sheen, Tianhao Wang, Zhuoran Yang,
• Abstract summary: We study how a two-attention-layer transformer is trained to perform ICL on $n$-gram Markov chain data. We prove that the gradient flow with respect to a cross-entropy ICL loss converges to a limiting model.
• Score: 54.20763128054692
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: In-context learning (ICL) is a cornerstone of large language model (LLM) functionality, yet its theoretical foundations remain elusive due to the complexity of transformer architectures. In particular, most existing work only theoretically explains how the attention mechanism facilitates ICL under certain data models. It remains unclear how the other building blocks of the transformer contribute to ICL. To address this question, we study how a two-attention-layer transformer is trained to perform ICL on $n$-gram Markov chain data, where each token in the Markov chain statistically depends on the previous $n$ tokens. We analyze a sophisticated transformer model featuring relative positional embedding, multi-head softmax attention, and a feed-forward layer with normalization. We prove that the gradient flow with respect to a cross-entropy ICL loss converges to a limiting model that performs a generalized version of the induction head mechanism with a learned feature, resulting from the congruous contribution of all the building blocks. In the limiting model, the first attention layer acts as a $\mathit{copier}$, copying past tokens within a given window to each position, and the feed-forward network with normalization acts as a $\mathit{selector}$ that generates a feature vector by only looking at informationally relevant parents from the window. Finally, the second attention layer is a $\mathit{classifier}$ that compares these features with the feature at the output position, and uses the resulting similarity scores to generate the desired output. Our theory is further validated by experiments.

Related papers

• Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers [14.59741397670484]
  We consider a deep multi-head self-attention network, that is closely related to Transformers yet analytically tractable. We develop a statistical mechanics theory of Bayesian learning in this model. Experiments confirm our findings on both synthetic and real-world sequence classification tasks.
  arXiv  Detail & Related papers  (2024-05-24T20:34:18Z)
• Toward a Theory of Tokenization in LLMs [26.516041872337887]
  We study tokenization from a theoretical point of view by studying the behavior of transformers on simple data generating processes. We show that even the simplest unigram models (over tokens) learnt by transformers are able to model the probability of sequences drawn from $ktextth$-order Markov sources near optimally.
  arXiv  Detail & Related papers  (2024-04-12T09:01:14Z)
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations [98.7450564309923]
  This paper takes initial steps on understanding in-context learning (ICL) in more complex scenarios, by studying learning with representations. We construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.
  arXiv  Detail & Related papers  (2023-10-16T17:40:49Z)
• Transformers as Support Vector Machines [54.642793677472724]
  We establish a formal equivalence between the optimization geometry of self-attention and a hard-margin SVM problem. We characterize the implicit bias of 1-layer transformers optimized with gradient descent. We believe these findings inspire the interpretation of transformers as a hierarchy of SVMs that separates and selects optimal tokens.
  arXiv  Detail & Related papers  (2023-08-31T17:57:50Z)
• Combiner: Full Attention Transformer with Sparse Computation Cost [142.10203598824964]
  We propose Combiner, which provides full attention capability in each attention head while maintaining low computation complexity. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention. An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach.
  arXiv  Detail & Related papers  (2021-07-12T22:43:11Z)
• $O(n)$ Connections are Expressive Enough: Universal Approximability of Sparse Transformers [71.31712741938837]
  We show that sparse Transformers with only $O(n)$ connections per attention layer can approximate the same function class as the dense model with $n2$ connections. We also present experiments comparing different patterns/levels of sparsity on standard NLP tasks.
  arXiv  Detail & Related papers  (2020-06-08T18:30:12Z)
• Supervised Quantile Normalization for Low-rank Matrix Approximation [50.445371939523305]
  We learn the parameters of quantile normalization operators that can operate row-wise on the values of $X$ and/or of its factorization $UV$ to improve the quality of the low-rank representation of $X$ itself. We demonstrate the applicability of these techniques on synthetic and genomics datasets.
  arXiv  Detail & Related papers  (2020-02-08T21:06:02Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.