Analysis and control of mean-field self-attention dynamics in Transformers

Investigate the dynamics of the McKean–Vlasov partial differential equation ∂_s μ_s + div(μ_s A_{ω_s}(μ_s)) = 0 that models non-causal Transformer self-attention in the mean-field limit, and establish theoretical guarantees for optimizing the parameter path θ = (Q_s, K_s, V_s)_{s ∈ [0,1]} via gradient descent viewed as a PDE control problem.

Background

The paper derives a mean-field representation of attention leading to a McKean–Vlasov-type PDE for the token distribution, distinct from the Wasserstein gradient flows analyzed for two-layer MLP training.

Although some asymptotic behaviors (e.g., clustering) are known in specific cases, the authors state that a deeper understanding of the PDE’s evolutions and of parameter optimization via gradient descent remains open, naturally framing it as a PDE control problem.