Infinite Limits of Multi-head Transformer Dynamics

2024/05/24 by Blake Bordelon, Hamza Tahir Chaudhry, Bordelon, Blake +3 · 1 voice · 17 citations
Engineering · #Computer science #Control Systems in Engineering #Control and Stability of Dynamical Systems #Electrical engineering #Engineering #Geology #Geomorphology #Head (geology) #Physics and Engineering Research Articles #Transformer #Voltage

paper · pdf · doi:10.48550/arxiv.2405.15712

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2024/05/24 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/05

Abstract

In this work, we analyze various scaling limits of the training dynamics of transformer models in the feature learning regime. We identify the set of parameterizations that admit well-defined infinite width and depth limits, allowing the attention layers to update throughout training--a relevant notion of feature learning in these models. We then use tools from dynamical mean field theory (DMFT) to analyze various infinite limits (infinite key/query dimension, infinite heads, and infinite depth) which have different statistical descriptions depending on which infinite limit is taken and how attention layers are scaled. We provide numerical evidence of convergence to the limits and discuss how the parameterization qualitatively influences learned features.

Cited by

• Dropout and Random Gradient Masking Are Asymptotically Equivalent in Large ResNets
• There Will Be a Scientific Theory of Deep Learning
• Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
• Training Dynamics of In-Context Learning in Linear Attention
• Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales
• Precise Dynamics of Diagonal Linear Networks: A Unifying Analysis by Dynamical Mean-Field Theory
• Gaussian Equivalence for Self-Attention: Asymptotic Spectral Analysis of Attention Matrix
• Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time
• Theoretical Foundations of Representation Learning using Unlabeled Data: Statistics and Optimization
• The Importance of Being Lazy: Scaling Limits of Continual Learning
• Revisiting Transformers with Insights from Image Filtering and Boosting
• On the Convergence of Gradient Descent on Learning Transformers with Residual Connections
• Bayes optimal learning of attention-indexed models
• Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs
• Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation
• Scaling Diffusion Transformers Efficiently via μP
• Don't be lazy: CompleteP enables compute-efficient deep transformers

Discussions

• Come by at Neurips to hear Hamza present about properties of various feature learning infinite parameter limits of transformer models. Poster in Hall A-C #4804 at 11 AM PST Paper arxiv.org/abs/2405. [bsky, 17 points, 0 comments]

Related

• Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape
• In-Context Convergence of Transformers
• The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit
• How Transformers Get Rich: Approximation and Dynamics Analysis
• The Asymptotic Behavior of Attention in Transformers