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

Authors: Aaron Alvarado Kristanto Julistiono, Davoud Ataee Tarzanagh, Navid Azizan
Primary area: optimization
Keywords: Attention Mechanism, Mirror Descent, Implicit Regularization, Transformers

Abstract:
Attention mechanisms have revolutionized numerous domains of artificial intelligence, including natural language processing and computer vision, by enabling models to selectively focus on relevant parts of the input data. Building on recent results characterizing the optimization dynamics of gradient descent (GD) and the structural properties of its preferred solutions in attention-based models, this paper explores the convergence properties and implicit bias of a family of mirror descent (MD) algorithms designed for softmax attention mechanisms, with the potential function chosen as the $p$-th power of the $\ell_p$-norm. Specifically, we show the directional convergence of these algorithms to a generalized hard-margin SVM with an $\ell_p$-norm objective when applied to a classification problem using a one-layer softmax attention model. Our theoretical results demonstrate that these algorithms not only converge directionally to the generalized max-margin solutions but also do so at a rate 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.

Decision: reject
Reviewer scores: 6, 3, 3, 6

Reviews and discussion:
<thread>
creation date:
1730692415287

thread_type:
official_review

<thread_object>
creator: reviewer_7bsg

<content>
summary:
This paper examines how mirror descent (MD) algorithms approach convergence and exhibit bias when optimizing attention mechanisms in softmax attention, with potential functions using $\ell_p$-norms as a basis for analysis. Key theoretical findings include demonstrating convergence towards hard margin Support Vector Machine (SVM) solutions and achieving convergence rates similar to gradient descent, in more straightforward models. The research builds upon findings related to descent by expanding the scope to encompass a wider range of optimization algorithms and shedding light on properties related to token selection.

strengths:
1. The paper has a strong theoretical foundation, it provides rigorous mathematical analysis and proofs for the convergence properties of mirror descent in attention optimization, extending previous work on gradient descent to a more general framework.
2. This paper provides a novel algorithmic insight, the introduction of $\ell_p$-AttGD generalizes both $\ell_p$-GD and attention GD, offering new perspectives on attention optimization and token selection.
3. This work provides a complete theoretical treatment of the optimization dynamics by examining both fixed-head and joint optimization scenarios.

weaknesses:
1. The empirical evaluation is limited, the paper includes experiments, and they are primarily focused on synthetic data and a single real-world dataset (Stanford Large Movie Review Dataset). More diverse real-world applications would strengthen the practical implications.
2. Theoretical results are highly dependent on assumption, Theorem 2, Theorem 3, and Theorem 4, rely on specific assumptions about initialization and step sizes, which may limit their practical applicability.
3. The paper does not thoroughly address the computational overhead of implementing mirror descent compared to standard gradient descent.

questions:
Please refer to the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730731695736

thread_type:
official_review

<thread_object>
creator: reviewer_vf65

<content>
summary:
This paper investigates the optimization dynamics of mirror descent (MD) for training attention mechanisms, specifically focusing on $l_p$-AttGD. The authors claim that $l_p$-AttGD converges directionally to a generalized hard-margin SVM with an $l_p$ norm objective when applied to binary classification with a single-layer attention model. Some experiments on synthetic and real data are presented.

strengths:
The paper attempts to provide a theoretical analysis of mirror descent for attention training, extending prior work focused on gradient descent. It derives convergence results to a generalized hard-margin SVM and establishes convergence rates. The use of $l_p$ norms offers a degree of generality in the theoretical analysis.

weaknesses:
1. The core idea of connecting attention optimization to SVM-like objectives is not new and has been explored in prior work, notably in "A Primal-Dual Framework for Transformers and Neural Networks" by Nguyen et al. and related papers. These prior works establish the fundamental link between attention and SVMs, including the optimization perspective. While this paper extends the analysis to mirror descent, the incremental contribution feels minimal and lacks motivation. Other core ideas of analyzing the implicit bias of GD/MD algorithms for softmax attention is already present in the cited works (Tarzanagh et al., 2023, 2024; Vasudeva et al., 2024a; Sheen et al., 2024) without fundamentally changing the nature of the problem or leading to significantly different insights.

2. The synthetic experiments are too simplistic and lack the complexity needed to represent realistic attention training scenarios. The real-data experiments, while showing some improvements in generalization, are insufficient to support the claims. The demonstration of improved token selection is based on a handful of examples (Figure 8) without any rigorous quantitative evaluation. Comparisons with GD do not include other commonly used optimization algorithms like Adam, making it impossible to judge the relative merits of MD for attention training.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730759711875

thread_type:
official_review

<thread_object>
creator: reviewer_kvfq

<content>
summary:
This paper proposes to study the dynamics of a simplified one-layer attention network. The implicit bias of gradient descent for such a problem is already known, this paper proposes to study what happens when the training algorithm is mirror descent with p-norms. The paper extends previous results known for gradient descent to this setting, and then demonstrate that using mirror descent to train a transformer on a sentiment analysis dataset yields better generalization.

strengths:
- The paper is well written 
- It is quite easy to understand
- The fact that MD achieves sparser weights and leads to better generalization is interesting.

weaknesses:
- The motivation is quite unclear. “A broader understanding of general descent algorithms, including the mirror descent (MD) family and their token selection properties, is essential.” why? In practice, nobody trains attention layers with mirror descent. The observation that it works better than gradient descent is not very strong in my opinion, because in practice transformers cannot be trained with gradient descent either. In the experiments, it would be worthwhile to compare the proposed methods to the widely used adam algorithm. Also, providing train losses values for the real data experiment would be enlightening.
- There are many unclear things about global and optimal tokens (def.3). First, the fact that there is a feasible solution is very important, precise conditions for it to happen must be put in the paper. The authors simply state that “under mild overparameterization d ≥ max{T − 1, n}, the problem is almost always feasible (Tarzanagh et al., 2023, Theorem 1).” These conditions must be put in the paper. Also, it is unclear whether the statement is “for all $\alpha_i$, the problem is feasible” or “the problem is feasible for at least one set of $\alpha_i$.
- I cannot understand why Theorem 1 does not require those conditions: clearly, $W^{opt}_{mm}$ must exist for it to hold, and there is nothing in the assumptions guaranteeing it.
- It is never properly stated in the paper if “globally optimal tokens” are also “locally optimal”. It seems not to be the case, since globally optimal tokens depend on $v$ while locally optimal tokens do not. This also causes issues in thm 1: globally optimal tokens do not seem to be always feasible, so W^{opt}_{mm} is not always defined, while the lp-AttRp always has a solution. For instance, it could very well have a global solution of fixed norm.

questions:
- Can the authors clarify the questions above regarding feasibility? 
- Can the authors compare the training curves with that of adam?
- The model sizes do not seem to be specified. The architectural details should be put in the appendix


**Minor remarks and typos**

- L68: ERM refers to an equation much later in the manuscript
- L71: W(k) has not bee introduced thus far
- L113: the softmax acts here on a matrix, but later in (2) it acts on vectors.
- L139: (Blair, 1985)
- L165: (Tarzanagh et al. (2024; 2023)).  
- L165: the sense of nearby is unclear here.
- L169: it would be good to specify to which set the indices belong. Generally, it would be clearer to indicate in which space lives each newly introduced variable.
- L176: “the αi component has a significantly higher probability” : this statement should also involve the norm of W, right? Because if all values in the pre-softmax vector are high, then the probability vector will be almost constant.
- Eq 4b: typo W -> W’
- L280: $exp(2)$ could be $\exp(2)$
- In eq.5, it would be good to recall that $D_\psi$ depends on $p$.
- All figures should be in pdf format
- “In the Appendix G, we show that W and v generated by ℓp-JointRP converge to their respective…” it would strengthen that section to put it in the main text.
- Fig4: a log scale would make things clearer, especially highlighting the practical convergence rates.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731028303725

thread_type:
official_review

<thread_object>
creator: reviewer_hzdv

<content>
summary:
The paper introduces a novel approach to optimizing attention mechanisms using Mirror Descent (MD), specifically focusing on a generalized max-margin token selection strategy for softmax attention models. The authors propose a family of MD algorithms, termed $\ell_{p}$-AttGD, which generalize traditional gradient descent by using the $p$-th power of the $\ell_p$-norm as the potential function. The main contributions include proving the convergence of these algorithms to generalized max-margin Support Vector Machine (SVM) solutions for optimizing the attention mechanism, both for fixed and jointly optimized parameters (key-query matrix and decoder).

strengths:
1. Theoretical Contributions: The paper provides a solid theoretical foundation for understanding the convergence properties and implicit bias of MD in attention models. The extension to $\ell_p$-norm objectives adds flexibility in modeling and opens up new avenues for optimizing attention mechanisms.

2. Generalization of Attention Optimization: The approach generalizes previous work on attention models by using MD with a broad class of potential functions, allowing a deeper exploration of the optimization landscape beyond vanilla gradient descent.

weaknesses:
1. As far as I know, mirror descent is not a popular optimization algorithm for training deep learning models. I agree that a simplified model (e.g., the one layer model considered in this paper and previous work) could provide valuable insights on understanding transformers, but it is not clear what is the role/implication the $\ell_p$-norm for deep learning. If possible, it would helpful if the authors could highlight a few practical mirror descent-based optimizers in the revision.

2. For Line431, for the $\ell_{1,1}$/$\ell_3$ optimizer, the authors provide code for implementing the $\ell_{1,1}$/$\ell_3$ optimizer for optimizing the transformer. In term of efficiency, would the $\ell_{1,1}$/$\ell_3$ optimizer be as efficient as popular optimizers like AdamW [LH2019]? Could the authors provide comparison with AdamW in Table 1?

[LH2019] Decoupled Weight Decay Regularization. Ilya Loshchilov, Frank Hutter.

questions:
1. What does the comparison token $z_i$ mean in practice (Line 55)? For example, for a real-world application problem, how to set the comparison token $z_i$ given the input $X$ and label $y$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732747761834

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their valuable feedback, which greatly improved our manuscript. Below, we summarize the main contributions **(C1–C3)**, the reviewers' key points **(P1–P4)**, and our actions **(A1–A4)**. Revisions and **new experiments** are highlighted in blue in the paper.

**C1. Implicit Bias Analysis with Mirror Descent (MD):** We provide a comprehensive theoretical analysis of MD for softmax attention mechanisms, introducing the $\ell_p$-AttGD framework. This approach demonstrates directional convergence to generalized hard-margin SVM solutions with $\ell_p$-norm objectives and extends prior work limited to gradient descent.
 Reviewer kVFQ noted: "The fact that MD achieves sparser weights and leads to better generalization is interesting." This highlights the significance of analyzing MD's implicit bias and token selection properties.

**C2. Different Optimizers and Their Token Selection:**  Our motivation is to demonstrate that **different optimizers have different token selection properties**, which in turn leads to different generalization performances.  This observation, missing in the attention optimization literature (Tarzanagh et al., 2023, 2024; Vasudeva et al., 2024a; Sheen et al., 2024), is validated by analyzing $\ell_p$-AttGD convergence to $\ell_p$-SVM under different $\ell_p$-norm objectives within the nonconvex softmax framework. Reviewer HzdV stated: " The paper provides a solid theoretical foundation for understanding the convergence properties and implicit bias of MD in attention models. The extension to $\ell_p$-norm objectives adds flexibility in modeling and opens up new avenues for optimizing attention mechanisms.” 

**C3. Empirical Insights on Generalization:** Through experiments on synthetic and real-world datasets (e.g., the Stanford Large Movie Review Dataset), we demonstrate that MD algorithms improve token selection and generalization compared to GD. Furthermore, $\ell_{1.1}$-AttGD exhibits superior sparsity and focus, offering improved token selection capabilities. Reviewer 7Bsg mentioned: "The fact that the real-data experiments show some improvements in generalization is noteworthy," validating the practical value of our empirical contributions.

**P1. Practical Relevance and Role of MD:** Several reviewers questioned the practicality of MD, emphasizing its limited usage in deep learning models (Reviewers HzdV and kVFQ). Reviewer HzdV requested examples of MD-based optimizers, while Reviewer kVFQ emphasized the need for clearer motivation for exploring MD in attention training.

**P2. Comparison to Adam and Architecture Details:** Reviewers requested comparisons between MD and popular optimizers like AdamW (Reviewers HzdV and vf65). Reviewer kVFQ also suggested including model architecture details in the appendix for reproducibility.

**P3. Feasibility and Assumptions of $\ell_p$-AttSVM:** Reviewers kVFQ and vf65 raised concerns about feasibility conditions in the $\ell_p$-AttSVM problem and requested explicit clarification of assumptions related to initialization and step size.

**P4. Broader Experimental Scope and Evaluation:** Reviewers vf65 and 7Bsg noted that the synthetic experiments are overly simplistic and called for more diverse datasets to support the claims. Additionally, quantitative evaluations of token selection and training loss values were suggested for real-data experiments.

**A1. Clarified Motivation and Practical Relevance of MD:** We revised the introduction to emphasize the novel insights provided by MD, including its token selection properties and implicit bias. Further, examples of MD-based optimizers, such as those used in convex optimization and reinforcement learning, have been added.

    Further details are provided in the response to Reviewers HzdV and kVFQ.

**A2. Comparisons with Adam and Architecture Details:** We added new experiments to compare the test accuracy and weight distribution of Adam and $\ell_{1.1}$-MD on training a Vision Transformer model to learn the CIFAR-10 dataset, to show a competitive performance and a superior explainability of MD, as shown **in the newly added Figures 11 and 12 in Appendix H5**.

    Further details are provided in the response to Reviewers kVFQ, vf65, and HzdV.

**A3. Feasibility and Assumptions in $\ell_p$-AttSVM:** We explicitly stated the feasibility conditions and assumptions (e.g., overparameterization $d \geq \max(T-1, n)$) in the main text. These conditions were also discussed in the context of Theorem 1 and other results.

    Further details are provided in the response to Reviewers kVFQ and vf65.

**A4. Broader Experiments and Quantitative Evaluation:** To strengthen the empirical results, we conducted additional experiments on tasks like sentiment classification with more diverse datasets. Quantitative metrics for token selection accuracy and training loss values have been incorporated.

    Further details are provided in the response to Reviewers vf65 and 7Bsg.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733292853302

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We appreciate the reviewers' feedback, which has greatly improved our paper.

We understand that the response period has concluded and sincerely hope our responses have addressed all concerns. If so, we kindly request a re-evaluation of our work.

Best,

Authors
</content>
</thread_object>

</thread>


Metareview: