Research map / Optimization / Implicit bias and generalization
High-Dimensional Learning Dynamics of Attention-Indexed Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes online SGD on attention-indexed models in a high-dimensional limit, showing a finite-dimensional population loss but an infinite moment hierarchy for dynamics, and that tied vs untied attention factorizations change whether weak recovery occurs on d^2 log d samples.
Problem
Training dynamics of attention with extensive-rank attention matrices in the high-dimensional regime are poorly understood; it is unclear how attention parameterization (direct, tied S=WW^T, untied S=UV^T) affects feature learning and escape from uninformative initializations.
Approach
Attention-indexed model framework with extensive-rank quadratic forms: population loss reduces to a finite set of trace order parameters, while online SGD for tied/untied factorizations yields an infinite moment hierarchy (exponentially accurate truncations); tied PSD factorization breaks symmetry automatically, while untied dynamics have a fast mean phase and slow overlap phase whose selected sta
Key results
- Tied attention S=WW^T achieves weak recovery in Theta(d^2 log d) samples under a mild non-degeneracy condition, while direct optimization over S can remain trapped on this scale.
- Untied attention S=UV^T exhibits fast-slow dynamics: a weak-recovery dichotomy where recovery on the d^2 log d scale occurs iff the fast-selected state breaks the relevant symmetry.
- Finite-degree moment truncations approximate the infinite hierarchy with error decaying exponentially in truncation degree M (Corollary on truncation), verified numerically at d=400 for L=2 softmax MSE.
Surprising findings
- Direct optimization over the attention matrix S can stay on an uninformative manifold on the d^2 log d sample scale, while the tied factorization escapes it automatically.
Limitations
- Detailed dynamics only for tied and untied parameterizations, not full Q/K/V dynamics of deep transformers.
- Weak recovery on longer timescales governed by information exponents is not analyzed.
- Analysis restricted to online SGD; ERM and empirical gradient flow not covered.
Experimental setup
Metrics: overlap weak recovery sample complexity
Baselines: direct optimization over S
Builds on: information exponent framework sequence-indexed models multi-indexed models
Closest papers by meaning
- Bayes optimal learning of attention-indexed models (NeurIPS 2025 · similarity 0.90)
- Scaling Limits of Long-Context Transformers (NeurIPS 2026 · similarity 0.87)
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.87)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.87)
- Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers (NeurIPS 2024 · similarity 0.87)
- Training Dynamics of In-Context Learning in Linear Attention (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.