Gradient Descent and Attention Models: Challenges Posed by the Softmax Function

1. The setting is somewhat limited. In the simplified transformer model, only the tunable tokens are learned, while the attention matrix remains fixed. This approach may primarily address the challenges GD encounters in prompt-tuning tasks on transformers but does not extend to other tasks, such as language modeling. Therefore, I believe this simplified model does not fully capture the key features of transformers that make GD optimization challenging in practical applications.
2. In Assumption 2, it seems necessary to have $d \geq n \cdot N$ . Yet, in the underparameterized case where $d \ll N$ , your analysis still relies on Assumption 2. This appears contradictory, and as a result, the equation on line 345 may not hold as stated.
3. The proof of Theorem 2 appears inconsistent with its statement. In Eq. (47), you have $\lambda_{\text{max}} \geq L(\min z_i)^2$ and $\lambda_{\text{min}} \leq \mu(\max z_i)^2$, which suggests $\kappa \geq \frac{L(\min z_i)^2}{\mu(\max z_i)^2}$ . This would lower the bound, which seems to contradict your argument regarding $R$. Are there any potential typos?

1: This paper studies the convergence in an under-parameterized regime, but the analysis mainly focuses on the local convergence near the stationary point.

2: The paper's assumptions are relatively strict (assumption 2, PL conditions, and smooth conditions).

3: Some settings simplify the difficulty of analysis (assumption 1).

1. Unfortunately, after a very promising introduction and motivation on performance of GD vs Adam on transformers, I found the results somewhat underwhelming. I would hope that at least in this very simplified (not easy, but admittedly simplified with many assumptions to help analysis), the authors would present a setting where Adam>GD. This is not even given in the experiments. The main theoretical results are also somewhat limited. (For OP, it is rather well understood that sufficient overparameterization and appropriate initialization induces some form of PL with square loss. For UP, if I understand correct, the local analysis is essentially studying conditioning of a linearization of a quadratic approximation of GD updates of attention outputs; e.g. the quadratic approximation in (16) seems to be ‘hiding under the rug’ the fact that optimized parameters p_t need to diverge for z_t to end up being sparse.)

2. As mentioned above, some assumptions made are rather strong. In particular, the orthogonality assumption and the strong convexity assumptions. I ‘d prefer if the authors simply acknowledge that they need these for analysis and also provide some insights on where they are critical and how one could potentially relax them rather than attempting to justify them as reasonable (eg see Remark 1)

3. The experiments are rather limited. Even judging this as a theory paper, I d expect some more results to be presented (synthetic OK), e.g. (i) Adam performance for this synthetic model, (ii) estimation error (together with optimization error) since you are considering a student-teacher setup, (iii) another setting where sparsity is more explicitly imposed and desired at the output, (iv) other loss functions: does the same behavior hold for logistic loss, joint training of v and p

There are several missing references in the discussion on related work, e.g. [1-5]. The authors don’t cite [1], due to which some of the statements about the contribution seem misleading. [1] analyzes gradient-based methods with adaptive step size and shows global convergence and finite-time fast convergence rates for parameter convergence for a single-head self-attention model. In the Introduction, the authors state that prior works lack convergence rates, and in the Conclusion, they even mention analyzing adaptive step sizes as a direction for future work.

While the observation that using linear attention instead of softmax attention can speed up convergence is somewhat interesting, it is not very insightful given that [1] already shows that using adaptive step size provably speeds up convergence compared to gradient descent with fixed step size.

The analysis is done under very strong assumptions on the model. The authors consider prompt attention with a fixed linear decoder, instead of the commonly used self-attention. Since the decoder weights are fixed, the loss only decreases in a small range as the attention weights change and it does not converge to 0 in this setting. So, deriving a convergence rate in this setting is not very informative.

The experimental section is also very weak, the authors should include results with more complex models and datasets.

References:

[1] Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention, 2024.

[2] Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, Christos Thrampoulidis. On the Optimization and Generalization of Multi-head Attention, TMLR 2024.

[3] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, NeurIPS 2023.

[4] Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention, ICLR 2024.

[5] Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison H. Zhou. Implicit regularization of gradient flow on one-layer softmax attention, 2024.

Theoretical Limitations

(Assumption 1)

• Oversimplified Setup: The theoretical setup of prompt tuning might oversimplify the complex reasons why SGD struggles with Transformers. The softmax issue, while potentially significant, may not be the sole culprit.

• Ignoring Transformer Heterogeneity: Transformers are known for their heterogeneous curvature landscapes [1, 2]. Parameterizing query-key matrices as a single matrix can further alter this landscape [2].

I think it would be beneficial to compare how crucial each of these factors are for hindering the use of SGD with Transformers, and a discussion on these recent work would be helpful.

(Assumption 2)

• If what is being assumed is that for each sequence the token embeddings for distinct tokens are orthogonal, this could be just written as X X^\top = I. Why involve it together with key weights? And this is also used as such in lines 345-348, which would imply that this assumption is misleading, and it is beyond just the input data. Besides, I think this assumption is also quite drastic as there is no reason to assume that the tokens in a sentence are 'orthogonal'. This is also assumed for every sequence, if I understand correctly?

Methodological Limitations

• How so adaptive methods bypass this issue caused by Softmax? Can the authors comment how AdamW ends up combating this issue? Otherwise this would be an incomplete explanation.

• The focus here is on a local analysis, but the problems with SGD start appearing from the beginning of training language models with Transformers.

Empirical Limitations

• Experiments are quite limited: It is good to have experiments where one can precisely simulate the theoretical conditions. So that is fine. But it is important to still carry out experiments in scenarios which are a bit more closer to the practical setup. Right now the only other setting is ViT on MNIST, and thereby makes for a very weak empirical evaluation.

• Could the authors carry out an empirical study of R during the course of training an proper Transformer on a language modelling or translation task? This would allow us to see when this particular issue with softmax starts appearing. Also, it would be good to conduct this with both Adam and SGD to draw effective comparisons.

References:

[1] Zhang, et. al. (2024). Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.16793.

[2] Ormaniec, et. al. (2024). What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis. arXiv preprint arXiv:2410.10986.