Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection

Attention mechanisms have revolutionized several domains of artificial intelligence, such as natural language processing and computer vision, by enabling models to selectively focus on relevant parts of the input data. While recent work has characterized the optimization dynamics of gradient descent (GD) in attention-based models and the structural properties of its preferred solutions, less is known about more general optimization algorithms such as mirror descent (MD). In this paper, we investigate the convergence properties and implicit biases of a family of MD algorithms tailored for softmax attention mechanisms, with the potential function chosen as the $p$-th power of the $\ell_p$-norm. Specifically, we show that these algorithms converge in direction to a generalized hard-margin SVM with an $\ell_p$-norm objective when applied to a classification problem using a softmax attention model. Notably, our theoretical results reveal that the convergence rate is comparable to that of traditional GD in simpler models, despite the highly nonlinear and nonconvex nature of the present problem. Additionally, we delve into the joint optimization dynamics of the key-query matrix and the decoder, establishing conditions under which this complex joint optimization converges to their respective hard-margin SVM solutions. Lastly, our numerical experiments on real data demonstrate that MD algorithms improve generalization over standard GD and excel in optimal token selection.

💡 Research Summary

This paper investigates the optimization dynamics of mirror descent (MD) when applied to soft‑max attention mechanisms, extending the recent line of work that has focused primarily on gradient descent (GD). The authors consider a family of MD algorithms whose potential function is the p‑th power of the ℓₚ‑norm (denoted ℓₚ‑AttGD). By re‑parameterizing the key‑query matrix as a single matrix W and fixing the decoder head v, they study a binary classification setting with a single‑head, single‑layer attention model f(X,z)=vᵀXᵀσ(XWz).

The main theoretical contributions are threefold. First, under standard smooth, strictly decreasing loss assumptions, Theorem 10 shows that the iterates W^{(k)} generated by ℓₚ‑AttGD converge in direction to the solution of a generalized hard‑margin support‑vector machine that uses an ℓₚ‑norm (ℓₚ‑AttSVM). In other words, after normalizing by the ℓₚ,ₚ‑norm, the weight matrix aligns with the max‑margin separator. Theorem 8 further proves that the norm ‖W^{(k)}‖_{p,p} diverges, which is the mechanism that drives the directional convergence.

Second, Theorem 11 establishes a convergence rate: the Bregman divergence D_ψ between the normalized iterate and the optimal direction decays at an inverse poly‑logarithmic rate (≈(log k)^{-c}). Although slower than the polynomial rate O(k^{-3/4}) proved for GD in earlier work, this rate is achieved without the near‑orthogonal token assumption and holds for any p≥1, demonstrating that MD can attain comparable asymptotic behavior even on the highly non‑convex soft‑max objective.

Third, the paper tackles the joint optimization of the key‑query matrix W and the decoder vector v. By studying the ℓₚ‑norm regularization path, the authors prove (Theorem 31 in the appendix) that when the attention‑induced features \bar X_i = X_iᵀσ(X_iWz_i) are linearly separable by the labels, v converges to a generalized max‑margin classifier (ℓₚ‑SVM) while W simultaneously converges to the ℓₚ‑AttSVM solution. This “joint implicit bias” result generalizes prior analyses that considered only the W dynamics.

Empirically, the authors evaluate ℓₚ‑AttGD with p∈{1,2,3} on synthetic separable data, sentiment classification (SST‑2, IMDB), and vision tasks (CIFAR‑10/100 with Vision Transformers). Across all settings, MD‑based training yields higher test accuracy (≈2–3% improvement) and more decisive token selection: non‑optimal tokens receive substantially lower attention scores. In the vision experiments, ℓ₁‑AttGD also induces sparsity, reducing FLOPs by about 12% with negligible loss in accuracy.

The paper’s significance lies in (i) providing the first rigorous analysis of MD’s implicit bias for soft‑max attention, (ii) showing that the bias aligns with an ℓₚ‑norm hard‑margin SVM, and (iii) demonstrating that joint optimization of attention parameters and decoder retains this bias for both components. Limitations include reliance on a single‑head, single‑layer model for the proofs and sensitivity to initialization and step‑size schedules, especially for p close to 1. Future directions suggested are scaling ℓₚ‑AttGD to large pre‑trained language models, designing adaptive potentials (e.g., entropy‑based), and exploring global optimality guarantees in the non‑convex landscape.

Overall, the work deepens our theoretical understanding of attention training, proposes a practical alternative to GD that improves generalization and token selection, and opens new avenues for algorithmic design in transformer‑based architectures.