REVIEW 10 cited by
Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
read the original abstract
We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We establish the global convergence of gradient flow under suitable choices of initialization. In addition, we prove that an interesting "task allocation" phenomenon emerges during the gradient flow dynamics, where each attention head focuses on solving a single task of the multi-task model. Specifically, we prove that the gradient flow dynamics can be split into three phases -- a warm-up phase where the loss decreases rather slowly and the attention heads gradually build up their inclination towards individual tasks, an emergence phase where each head selects a single task and the loss rapidly decreases, and a convergence phase where the attention parameters converge to a limit. Furthermore, we prove the optimality of gradient flow in the sense that the limiting model learned by gradient flow is on par with the best possible multi-head softmax attention model up to a constant factor. Our analysis also delineates a strict separation in terms of the prediction accuracy of ICL between single-head and multi-head attention models. The key technique for our convergence analysis is to map the gradient flow dynamics in the parameter space to a set of ordinary differential equations in the spectral domain, where the relative magnitudes of the semi-singular values of the attention weights determines task allocation. To our best knowledge, our work provides the first convergence result for the multi-head softmax attention model.
Forward citations
Cited by 10 Pith papers
-
Specialization of softmax attention heads: insights from the high-dimensional single-location model
In a high-dimensional toy task, multi-head softmax attention first aligns all heads with the mean signal, then sequentially specializes to latent directions; the paper introduces Bayes-softmax, which attains the Bayes...
-
Learning Compositional Functions with Transformers from Easy-to-Hard Data
A transformer with O(log k) layers provably learns the k-fold permutation composition task in poly(N,k) samples with curriculum or mixed easy-to-hard data, despite an SQ lower bound requiring N^{Omega(k)} samples on h...
-
Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently
Under hand-designed masks and task-specific activations, RL fine-tuning learns a k-sparse Boolean reasoning chain in one gradient update while SFT learns it one CoT step per update.
-
How Can Mamba Learn In Context with Outliers and Generalize Provably?
A simplified one-layer Mamba provably learns in-context binary classification tolerating outlier fractions approaching 1, whereas a linear Transformer can only tolerate α < 1/2.
-
Unveiling the Mechanisms of Multi-Hop Reasoning in Transformers via Identity Bridge
Adding identity supervision on bridge tokens enables out-of-distribution two-hop reasoning in simple transformers, with a nuclear-norm theory explaining the benefit.
-
Train Once, Reuse Everywhere: Generalizable Implicit In-Context Learning by Routing Attention
ICR extracts shared attention directions from in-context learning and routes them at inference time, enabling zero-shot reuse across tasks.
-
Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression
A one-layer linear-attention transformer implementing noisy gradient descent gives provable bounds showing linear noise plus ensembling avoids label-noise overfitting, and majority voting beats greedy decoding in spar...
-
Transformers Meet In-Context Learning: A Universal Approximation Theory
A constructive theorem shows that transformers can perform in-context learning for any Barron-type function class by combining universal features with an emulated Lasso solver.
-
Federated In-Context Learning: Iterative Refinement for Improved Answer Quality
Fed-ICL iteratively refines QA answers via federated in-context learning with only label transmission, showing convergence on a linear attention model and gains on MMLU and TruthfulQA.
-
A Theoretical Study of (Hyper) Self-Attention through the Lens of Interactions: Representation, Training, Generalization
Single-layer linear self-attention can represent, train on, and length-generalize pairwise interaction functions under data-versatility and exact-realizability assumptions, and the paper introduces higher-order HyperA...
Discussion (0). Sign in to comment.