SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We have provided a thorough optimization-theoretic characterization of the fundamental attention model  $f(\pmb{X}) = \pmb{v}^{\top}\pmb{X}^{\top}\mathbb{S}(\pmb{X}\pmb{W}\pmb{p})$  by formally connecting it to max-margin problems. We first established the convergence of gradient descent on  $\pmb{p}$  (or equivalently  $W$ ) in isolation. We also explored joint convergence of  $(\nu, p)$  via regularization path which revealed surprising implicit biases such as (10). These findings motivate several exciting avenues for future research. An immediate open problem is characterizing the (local) convergence of gradient descent for joint optimization of  $(\nu, p)$ . Another major direction is to extend similar analysis to study self-attention layer (4) or to allow for multiple tunable tokens (where  $\pmb{p}$  becomes a matrix). Either setting will enrich the problem by allowing the attention to discover multiple hyperplanes to separate tokens. While our convergence guarantees apply when tokens are separable, it would be interesting to characterize the non-separable geometry by leveraging results developed for logistic regression analysis [31, 22]. Ideas from such earlier results can also be useful for characterizing the non-asymptotic/transient dynamics of how gradient descent aligns with the max-margin direction. Overall, we believe that max-margin token selection is a fundamental characteristic of attention mechanism and the theory developed in this work lays the groundwork of these future extensions.
This work was supported by the NSF grants CCF-2046816 and CCF-2212426, Google Research Scholar award, and Army Research Office grant W911NF2110312.
The authors express their gratitude for the valuable feedback provided by the anonymous reviewers and Christos Thrampoulidis, which has significantly improved this paper.
[100] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022.
[101] Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. arXiv preprint arXiv:2305.15408, 2023.
[102] Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thought: A study on compositional in-context learning of mlps. arXiv preprint arXiv:2305.18869, 2023.
[103] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv:2305.16380, 2023.
[104] Tan Minh Nguyen, Tam Minh Nguyen, Nhat Ho, Andrea L Bertozzi, Richard Baraniuk, and Stanley Osher. A primal-dual framework for transformers and neural networks. In The Eleventh International Conference on Learning Representations, 2023.
[105] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308.16898, 2023.
Roadmap. The appendix is organized as follows: Section A provides basic facts about the training risk. Section B presents the proof of local and global gradient descent and regularized path for learning  $\pmb{p} \in \mathbb{R}^d$  with a fixed  $\pmb{\nu} \in \mathbb{R}^d$  choice. Section C provides the proof of regularized path applied to the general case of joint optimization of head  $\pmb{\nu}$  and attention weights  $\pmb{p}$  using a logistic loss function. Section D presents the regularized path applied to a more general model  $f(\pmb{X}) = \psi(\pmb{X}^\top \mathbb{S}(\pmb{X}\pmb{W}^\top \pmb{p}))$  with a nonlinear head  $\psi$ . Section E provides implementation details. Finally, Section F discusses additional related work on implicit bias and self-attention.
A Addendum to Section 1 18
A.1 Preliminaries on the Training Risk 18
A.2 Proof of Lemma 1 20
B Addendum to Section 2 21
B.1 Descent and Gradient Correlation Conditions 21
B.2 Proof of Theorem 1 29
B.3 Proof of Theorem 2 30
B.4 Proof of Theorem 3 30
B.5 Proof of Theorem 4 33
B.6 Proof of Lemma 2 39
C Addendum to Section 3 39
C.1 Proof of Theorem 5 39
C.2 Proof of Theorem 6 41
D Regularization Path of Attention with Nonlinear Head 43
D.1 Proof of Theorem 8 44
D.2 Application to Linearly-mixed Labels 45
E Implementation Details and Additional Experiments 46
F Addendum to Section 5 48
F.1 Related Work on Implicit Regularization 48
F.2 Related Work on Attention Mechanism 48
A Addendum to Section 1
A.1 Preliminaries on the Training Risk
By our assumption  $\psi : \mathbb{R}^d \to \mathbb{R}$  and  $\ell : \mathbb{R} \to \mathbb{R}$  are differentiable functions. Recall the objective
$$
\mathcal {L} (\boldsymbol {p}, \boldsymbol {W}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(Y _ {i} \cdot \psi \left(\boldsymbol {X} _ {i} ^ {\top} \mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right)\right)\right) \tag {11}
$$
with the generic prediction model  $\psi (X^{\top}\mathbb{S}(Kp))$  and  $K = XW^{\top}$
Here, we write down the gradients of  $\pmb{W}$  and  $\pmb{p}$  in (11) to highlight the connection. Set  $\pmb{q} \coloneqq \pmb{W}^{\top} \pmb{p}$ ,  $z\{X\} \coloneqq X^{\top} \mathbb{S}(Kp)$ , and  $a\{X\} \coloneqq Kp$ . Given  $\pmb{X}$  and using  $K = X W^{\top}$ , we have that
$$
\nabla_ {q} \psi (\boldsymbol {p}, \boldsymbol {W}) = \boldsymbol {X} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} \{X \}) \boldsymbol {X} \cdot \psi^ {\prime} (z \{X \}), \tag {12a}
$$
$$
\nabla_ {p} \psi (\boldsymbol {p}, \boldsymbol {W}) = \boldsymbol {W} \nabla_ {q} \psi (\boldsymbol {p}, \boldsymbol {W}), \tag {12b}
$$
$$
\nabla_ {\boldsymbol {W}} \psi (\boldsymbol {p}, \boldsymbol {W}) = \boldsymbol {p} \nabla_ {\boldsymbol {q}} ^ {\top} \psi (\boldsymbol {p}, \boldsymbol {W}), \tag {12c}
$$
$$
\mathbb {S} ^ {\prime} (\boldsymbol {a} \{\boldsymbol {X} \}) = \operatorname {d i a g} (\mathbb {S} (\boldsymbol {a} \{\boldsymbol {X} \})) - \mathbb {S} (\boldsymbol {a} \{\boldsymbol {X} \})) \mathbb {S} (\boldsymbol {a} \{\boldsymbol {X} \}) ^ {\top} \in \mathbb {R} ^ {T \times T}.
$$
Setting  $\psi (z) = \nu^{\top}z$  for linear head, we obtain
$$
\nabla_ {\boldsymbol {q}} \psi (\boldsymbol {p}, \boldsymbol {W}) = \boldsymbol {X} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} \{\boldsymbol {X} \}) \gamma , \tag {13a}
$$
$$
\nabla_ {p} \psi (\boldsymbol {p}, W) = W \nabla_ {\boldsymbol {q}} \psi (\boldsymbol {p}, W) = K ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} \{X \}) \gamma , \tag {13b}
$$
$$
\nabla_ {W} \psi (\boldsymbol {p}, W) = \boldsymbol {p} \nabla_ {\boldsymbol {q}} ^ {\top} \psi (\boldsymbol {p}, W) = \boldsymbol {p} \boldsymbol {v} ^ {\top} X ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} \{X \}) X. \tag {13c}
$$
Recalling (12b) and (12c), and defining  $\ell_i' := \ell'(Y_i \cdot \psi(z[X_i])) \in \mathbb{R}$ , we have that
$$
\nabla_ {\boldsymbol {p}} \mathcal {L} (\boldsymbol {p}, \boldsymbol {W}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot Y _ {i} \cdot \boldsymbol {W} \nabla_ {\boldsymbol {q}} \psi (\boldsymbol {p}, \boldsymbol {W}), \tag {14a}
$$
$$
\nabla_ {\boldsymbol {W}} \mathcal {L} (\boldsymbol {p}, \boldsymbol {W}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot Y _ {i} \cdot \boldsymbol {p} \nabla_ {\boldsymbol {q}} ^ {\top} \psi (\boldsymbol {p}, \boldsymbol {W}). \tag {14b}
$$
Setting  $\psi (z) = \nu^{\top}z$  for linear head and  $\gamma_{i} = Y_{i}\cdot X_{i}\nu$ , we obtain
$$
\nabla_ {\boldsymbol {p}} \mathcal {L} (\boldsymbol {p}, \boldsymbol {W}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} \left\{X _ {i} \right\}) \boldsymbol {\gamma} _ {i}, \tag {15a}
$$
$$
\nabla_ {\boldsymbol {W}} \mathcal {L} (\boldsymbol {p}, \boldsymbol {W}) = \boldsymbol {p} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \gamma_ {i} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} \left\{\boldsymbol {X} _ {i} \right\}) \boldsymbol {X} _ {i}\right). \tag {15b}
$$
Lemma 3 (Key Lemma) For any  $\pmb{p}, \pmb{q} \in \mathbb{R}^d$ , let  $\pmb{a} = \pmb{K}\pmb{q}$ ,  $s = \mathbb{S}(\pmb{K}\pmb{p})$ , and  $\gamma = X\nu$ . Set
$$
\Gamma = \sup  _ {t, \tau \in [ T ]} | \boldsymbol {\gamma} _ {t} - \boldsymbol {\gamma} _ {\tau} | \text {a n d} A = \sup  _ {t \in [ T ]} \| \boldsymbol {k} _ {t} \| \cdot \| \boldsymbol {q} \|.
$$
$$
\left| \boldsymbol {a} ^ {\top} \operatorname {d i a g} (\boldsymbol {s}) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} \boldsymbol {s} \boldsymbol {s} ^ {\top} \boldsymbol {\gamma} - \sum_ {t \geq 2} ^ {T} \left(\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}\right) \boldsymbol {s} _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right) \right| \leq 2 \Gamma A (1 - s _ {1}) ^ {2}.
$$
Proof. Set  $\bar{\gamma} = \sum_{t=1}^{T} \gamma_{t} s_{t}$ . We have
$$
\gamma_ {1} - \bar {\gamma} = \sum_ {t \geq 2} ^ {T} \left(\gamma_ {1} - \gamma_ {t}\right) s _ {t}, \text {a n d} | \gamma_ {1} - \bar {\gamma} | \leq \Gamma \left(1 - s _ {1}\right).
$$
$$
\begin{array}{l} \boldsymbol {a} ^ {\top} \operatorname {d i a g} (\boldsymbol {s}) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} s \boldsymbol {s} ^ {\top} \boldsymbol {\gamma} = \sum_ {t = 1} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {\gamma} _ {t} \boldsymbol {s} _ {t} - \sum_ {t = 1} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {s} _ {t} \sum_ {t = 1} ^ {T} \boldsymbol {\gamma} _ {t} \boldsymbol {s} _ {t} \\ = \boldsymbol {a} _ {1} s _ {1} \left(\boldsymbol {\gamma} _ {1} - \bar {\boldsymbol {\gamma}}\right) - \sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} s _ {t} (\bar {\boldsymbol {\gamma}} - \boldsymbol {\gamma} _ {t}). \tag {16} \\ \end{array}
$$
$$
\left| \sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {s} _ {t} (\bar {\gamma} - \boldsymbol {\gamma} _ {t}) - \sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {s} _ {t} (\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}) \right| \leq A \Gamma (1 - s _ {1}) ^ {2},
$$
$$
\begin{array}{l} \boldsymbol {a} ^ {\top} \operatorname {d i a g} (\boldsymbol {s}) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} s s ^ {\top} \boldsymbol {\gamma} = \boldsymbol {a} _ {1} s _ {1} (\boldsymbol {\gamma} _ {1} - \bar {\boldsymbol {\gamma}}) - \sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} s _ {t} (\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}) \pm A \Gamma (1 - s _ {1}) ^ {2} \\ = \boldsymbol {a} _ {1} s _ {1} \sum_ {t \geq 2} ^ {T} (\gamma_ {1} - \gamma_ {t}) s _ {t} - \sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} s _ {t} (\gamma_ {1} - \gamma_ {t}) \pm A \Gamma (1 - s _ {1}) ^ {2} \\ = \sum_ {t \geq 2} ^ {T} \left(\boldsymbol {a} _ {1} \boldsymbol {s} _ {1} - \boldsymbol {a} _ {t}\right) \boldsymbol {s} _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right) \pm A \Gamma \left(1 - \boldsymbol {s} _ {1}\right) ^ {2} \\ = \sum_ {t \geq 2} ^ {T} \left(\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}\right) \boldsymbol {s} _ {t} \left(\gamma_ {1} - \gamma_ {t}\right) \pm 2 A \Gamma \left(1 - \boldsymbol {s} _ {1}\right) ^ {2}. \\ \end{array}
$$
Here,  $\pm$  on the right hand side uses the fact that
$$
\left| \sum_ {t \geq 2} ^ {T} \left(\boldsymbol {a} _ {1} \boldsymbol {s} _ {1} - \boldsymbol {a} _ {1}\right) \boldsymbol {s} _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right) \right| \leq (1 - \boldsymbol {s} _ {1}) A \Gamma \sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t} = (1 - \boldsymbol {s} _ {1}) ^ {2} A \Gamma .
$$
Proof. Let us prove the result for a general step size sequence  $(\eta_{t})_{t\geq 0}$ . On the same training data  $(Y_{i},X_{i})_{i = 1}^{n}$ , recall the objectives  $\tilde{\mathcal{L}} (\pmb {p}) = \frac{1}{n}\sum_{i = 1}^{n}\ell (Y_{i}\cdot \psi (X_{i}^{\top}\mathbb{S}(X_{i}\pmb {p}))$  and  $\mathcal{L}(\boldsymbol {W}) = \frac{1}{n}\sum_{i = 1}^{n}\ell (Y_{i}\cdot \psi (X_{i}^{\top}\mathbb{S}(X_{i}\boldsymbol{W}^{\top}\boldsymbol {u})))$ . Suppose claim is true till iteration  $t$ . For iteration  $t + 1$ , using  $\boldsymbol {W}(t)^{\top}\boldsymbol {u} = \boldsymbol {p}(t)$  define and observe that
$$
\boldsymbol {S} _ {i} = \mathbb {S} ^ {\prime} \left(\boldsymbol {X} _ {i} \boldsymbol {W} (t) ^ {\top} \boldsymbol {u}\right) = \mathbb {S} ^ {\prime} \left(\boldsymbol {X} _ {i} \boldsymbol {p} (t)\right),
$$
$$
\boldsymbol {s} _ {i} = \mathbb {S} (\boldsymbol {X} _ {i} \boldsymbol {W} (t) ^ {\top} \boldsymbol {u}) = \mathbb {S} (\boldsymbol {X} _ {i} \boldsymbol {p} (t)),
$$
$$
z \left\{X _ {i} \right\} = X _ {i} ^ {\top} \mathbb {S} \left(X _ {i} \boldsymbol {p} (t)\right) = X _ {i} ^ {\top} \mathbb {S} \left(X _ {i} \boldsymbol {W} (t) ^ {\top} \boldsymbol {u}\right),
$$
for all  $i \in [n]$ .
Thus, using (14), we have that
$$
\nabla \tilde {\mathcal {L}} (\boldsymbol {p} (t)) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot Y _ {i} \cdot \boldsymbol {X} _ {i} ^ {\top} \boldsymbol {S} _ {i} \boldsymbol {X} _ {i} \cdot \psi^ {\prime} (\boldsymbol {z} \{\boldsymbol {X} _ {i} \}),
$$
$$
\nabla \mathcal {L} (\boldsymbol {W} (t)) = \boldsymbol {u} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot Y _ {i} \cdot \boldsymbol {X} _ {i} ^ {\top} \boldsymbol {S} _ {i} \boldsymbol {X} _ {i} \cdot \psi^ {\prime} (\boldsymbol {z} \{\boldsymbol {X} _ {i} \})\right) ^ {\top}.
$$
Consequently, we found that gradient is rank-1 with left singular space equal to  $\pmb{u}$ , i.e.,
$$
\nabla \mathcal {L} (\boldsymbol {W} (t)) = \boldsymbol {u} \nabla^ {\top} \tilde {\mathcal {L}} (\boldsymbol {p} (t)).
$$
Since  $W(t)$ 's left singular space is guaranteed to be in  $u$  (including  $W(0)$  by initialization), we only need to study the right singular vector. Using the induction till  $t$ , this yields
$$
\begin{array}{l} \boldsymbol {W} (t + 1) ^ {\top} \boldsymbol {u} = \boldsymbol {W} (t) ^ {\top} \boldsymbol {u} - \eta_ {t} \| \boldsymbol {u} \| ^ {- 2} \nabla^ {\top} \mathcal {L} (\boldsymbol {W} (t)) \boldsymbol {u} \\ = \boldsymbol {p} (t) - \eta_ {t} \| \boldsymbol {u} \| ^ {- 2} \boldsymbol {u} ^ {\top} \boldsymbol {u} \nabla \tilde {\mathcal {L}} (\boldsymbol {p} (t)) \\ = \boldsymbol {p} (t + 1). \\ \end{array}
$$
This concludes the induction.
2For simplicity, we use  $\pm$  on the right hand side to denote the upper and lower bounds.
B Addendum to Section 2
B.1 Descent and Gradient Correlation Conditions
The lemma below identifies conditions under which  $p^{mm\star}$  is a global descent direction for  $\mathcal{L}(p)$ .
Lemma 4 Suppose  $\ell(\cdot)$  is a strictly decreasing differentiate loss function and Assumption B holds. Then, for all  $\pmb{p} \in \mathbb{R}^d$ , the training loss (ERM) obeys  $\langle \nabla \mathcal{L}(\pmb{p}), \pmb{p}^{mm\star} \rangle < 0$ .
Proof. Set
$$
\gamma_ {i} = Y _ {i} \cdot X _ {i} v, \quad a _ {i} = K _ {i} p, \quad \bar {a} _ {i} = K _ {i} p ^ {m m \star}, \text {a n d} \ell_ {i} ^ {\prime} = \ell^ {\prime} \left(\gamma_ {i} ^ {\top} \mathbb {S} (K _ {i} p)\right). \tag {17}
$$
Let us recall the gradient evaluated at  $\pmb{p}$  which is given by
$$
\nabla \mathcal {L} (\boldsymbol {p}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} \left(\boldsymbol {a} _ {i}\right) \boldsymbol {\gamma} _ {i}. \tag {18}
$$
$$
\left\langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {p} ^ {m m \star} \right\rangle = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \left\langle \bar {\boldsymbol {a}} _ {i}, \mathbb {S} ^ {\prime} \left(\boldsymbol {a} _ {i}\right) \boldsymbol {\gamma} _ {i} \right\rangle . \tag {19}
$$
To proceed, we will prove that individual summands are all strictly negative. To show that, without losing generality, let us focus on the first input and drop the subscript  $i$  for cleaner notation. This yields
$$
\langle \bar {a}, \mathbb {S} ^ {\prime} (a) \gamma \rangle = \bar {a} ^ {\top} \operatorname {d i a g} (\mathbb {S} (a)) \gamma - \bar {a} ^ {\top} \mathbb {S} (a) \mathbb {S} (a) ^ {\top} \gamma . \tag {20}
$$
Without losing generality, assume optimal token is the first one and  $\gamma_{t}$  is a constant for all  $t\geq 2$
To proceed, we will prove the following: Suppose  $\gamma = \gamma_{t\geq 2}$  is constant,  $\gamma_1,\bar{a}_1$  are the largest indices of  $\gamma ,\bar{a}$ . Then, for any  $s$  obeying  $\sum_{t\in [T]}s_t = 1,s_t\geq 0$ , we have that  $\bar{\pmb{a}}^{\top}\mathrm{diag}(s)\pmb {\gamma} - \bar{\pmb{a}}^{\top}s s^{\top}\pmb {\gamma} > 0$ . To see this, we write
$$
\begin{array}{l} \bar {\boldsymbol {a}} ^ {\top} \operatorname {d i a g} (s) \boldsymbol {\gamma} - \bar {\boldsymbol {a}} ^ {\top} s s ^ {\top} \boldsymbol {\gamma} = \sum_ {t = 1} ^ {T} \bar {\boldsymbol {a}} _ {t} \boldsymbol {\gamma} _ {t} s _ {t} - \sum_ {t = 1} ^ {T} \bar {\boldsymbol {a}} _ {t} s _ {t} \sum_ {t = 1} ^ {T} \boldsymbol {\gamma} _ {t} s _ {t} \\ = \left(\bar {\boldsymbol {a}} _ {1} \boldsymbol {\gamma} _ {1} s _ {1} + \gamma \sum_ {t \geq 2} ^ {T} \bar {\boldsymbol {a}} _ {t} s _ {t}\right) - \left(\boldsymbol {\gamma} _ {1} s _ {1} + \gamma (1 - s _ {1})\right) \left(\bar {\boldsymbol {a}} _ {1} s _ {1} + \sum_ {t \geq 2} ^ {T} \bar {\boldsymbol {a}} _ {t} s _ {t}\right) \\ = \bar {a} _ {1} \left(\gamma_ {1} - \gamma\right) s _ {1} \left(1 - s _ {1}\right) + \left(\gamma - \left(\gamma_ {1} s _ {1} + \gamma \left(1 - s _ {1}\right)\right)\right) \sum_ {t \geq 2} ^ {T} \bar {a} _ {t} s _ {t} \tag {21} \\ = \bar {a} _ {1} (\gamma_ {1} - \gamma) s _ {1} (1 - s _ {1}) - (\gamma_ {1} - \gamma) s _ {1} \sum_ {t \geq 2} ^ {T} \bar {a} _ {t} s _ {t} \\ = (\gamma_ {1} - \gamma) (1 - s _ {1}) s _ {1} \left[ \bar {\boldsymbol {a}} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \bar {\boldsymbol {a}} _ {t} s _ {t}}{\sum_ {t \geq 2} ^ {T} s _ {t}} \right]. \\ \end{array}
$$
To proceed, let  $\gamma_{gap} = \gamma_1 - \gamma$  and  $a_{gap} = \bar{a}_1 - \max_{t\geq 2}\bar{a}_t$ . With these, we obtain
$$
\bar {a} ^ {\top} \operatorname {d i a g} (s) \gamma - \bar {a} ^ {\top} s s ^ {\top} \gamma \geq a _ {\text {g a p}} \gamma_ {\text {g a p}} s _ {1} \left(1 - s _ {1}\right). \tag {22}
$$
$$
\boldsymbol{a}_{gap}^{i}\geq \inf_{t\neq \operatorname{opt}_{i}}(\boldsymbol{k}_{i\mathrm{opt}_{i}} - \boldsymbol{k}_{it})^{\top}\boldsymbol{p}^{mm\star}\geq 1,
$$
$$
\gamma^ {i} _ {g a p} = \inf  _ {t \neq o p t _ {i}} \gamma_ {i o p t _ {i}} - \gamma_ {i t} > 0,
$$
$$
\mathbf {s} _ {i 1} \left(1 - \mathbf {s} _ {i 1}\right) > 0.
$$
On the other hand, by our assumption  $\ell_i^{\prime} < 0$ . Hence, infimum'ing (22) over all inputs, multiplying by  $\ell_i^\prime$  and using (19) give the desired result.
Lemma 5 (Gradient Correlation Conditions) Consider  $n = 1$  and let  $p^{mm} = p^{mm\star}$  be (ATT-SVM) solution separating  $\alpha = \mathrm{opt}$  from remaining tokens of input  $X$ . Suppose  $\ell(\cdot)$  is a strictly decreasing differentiate loss function and Assumption B holds. For any choice of  $\pi > 0$ , there exists  $R := R_{\pi}$  such that, for any  $p$  with  $\| p \| \geq R$ , we have
$$
\left\langle \nabla \mathcal {L} (\boldsymbol {p}), \frac {\boldsymbol {p}}{\| \boldsymbol {p} \|} \right\rangle \geq (1 + \pi) \left\langle \nabla \mathcal {L} (\boldsymbol {p}), \frac {\boldsymbol {p} ^ {m m}}{\| \boldsymbol {p} ^ {m m} \|} \right\rangle .
$$
Above, observe that as  $R\to \infty$  , we eventually get to set  $\pi = 0$
Proof. The proof is similar to Lemma 4 at a high-level. However, we also need to account for the impact of  $\pmb{p}$  besides  $\pmb{p}^{mm}$  in the gradient correlation. The main goal is showing that  $\pmb{p}^{mm}$  is the near-optimal descent direction, thus,  $\pmb{p}$  cannot significantly outperform it.
To proceed, let  $\bar{\pmb{p}} = \| \pmb{p}^{mm}\| \pmb {p} / \| \pmb {p}\|$ ,  $M = \sup_{t}\| \pmb{k}_{t}\|$ ,  $\Theta = 1 / \| p^{mm}\|$ ,  $s = \mathbb{S}(Kp)$ ,  $a = K\bar{p}$ ,  $\bar{a} = Kp^{mm}$ . Without losing generality assume  $\mathrm{opt} = 1$ . Set  $\gamma = \gamma_{t\geq 2}$ . Repeating the proof of Lemma 4 yields
$$
\left\langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {p} ^ {m m} \right\rangle = \ell^ {\prime} \cdot \left(\gamma_ {1} - \gamma\right) (1 - s _ {1}) s _ {1} \left[ \bar {\boldsymbol {a}} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \bar {\boldsymbol {a}} _ {t} s _ {t}}{\sum_ {t \geq 2} ^ {T} s _ {t}} \right],
$$
$$
\langle \nabla \mathcal {L} (\boldsymbol {p}), \bar {\boldsymbol {p}} \rangle = \ell^ {\prime} \cdot (\gamma_ {1} - \gamma) (1 - s _ {1}) s _ {1} \left[ \boldsymbol {a} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} s _ {t}}{\sum_ {t \geq 2} ^ {T} s _ {t}} \right].
$$
Given  $\pi$ , for sufficiently large  $R$ , we wish to show that
$$
\boldsymbol {a} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} \leq (1 + \pi) \cdot \left[ \bar {\boldsymbol {a}} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \bar {\boldsymbol {a}} _ {t} \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} \right]. \tag {23}
$$
We consider two scenarios.
Scenario 1:  $\| \bar{p} - p^{mm} \| \leq \epsilon \coloneqq \pi / (2M)$ . In this scenario, for any token, we find that
$$
\left| \boldsymbol {a} _ {t} - \bar {\boldsymbol {a}} _ {t} \right| = \left| \boldsymbol {k} _ {t} ^ {\top} (\bar {\boldsymbol {p}} - \boldsymbol {p} ^ {m m}) \right| \leq M \| \bar {\boldsymbol {p}} - \boldsymbol {p} ^ {m m} \| \leq M \epsilon .
$$
$$
\bar {\boldsymbol {a}} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \bar {\boldsymbol {a}} _ {t} \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} \geq \boldsymbol {a} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} - 2 M \epsilon = \boldsymbol {a} _ {1} - \frac {\sum_ {t \geq 2} ^ {T} \boldsymbol {a} _ {t} \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} - \pi .
$$
Also noticing  $\bar{\pmb{a}}_1 - \frac{\sum_{t\geq 2}^T\bar{a}_ts_t}{\sum_{t\geq 2}^Ts_t}\geq 1$  (thanks to  $\pmb{p}^{mm}$  satisfying  $\geq 1$  margin), this implies (23).
Scenario 2:  $\| \bar{p} - p^{mm} \| \geq \epsilon \coloneqq \pi / (2M)$ . In this scenario, for some  $\nu = \nu(\epsilon)$  and  $\tau \geq 2$ , we have that
$$
\bar {\boldsymbol {p}} ^ {\top} \left(\boldsymbol {k} _ {1} - \boldsymbol {k} _ {\tau}\right) = \boldsymbol {a} _ {1} - \boldsymbol {a} _ {\tau} \leq 1 - 2 v.
$$
Here  $\tau = \arg \max_{t\geq 2}\bar{\pmb{p}}^{\top}\pmb{k}_{t}$  denotes the nearest point to  $\pmb{k}_{1}$ . Recall that  $s = \mathbb{S}(\bar{R}\pmb{a})$  where  $\bar{R} = \| \pmb {p}\| /\| \pmb{p}^{mm}\|$ . To proceed, split the tokens into two groups: Let  $\mathcal{N}$  be the group of tokens obeying  $\bar{\pmb{p}}^{\top}(\pmb {k}_1 - \pmb {k}_t)\leq 1 - \nu$  for  $t\in \mathcal{N}$  and  $[T] - \mathcal{N}$  be the rest. Observe that
$$
\frac {\sum_ {t \in [ T ] - \mathcal {N}} \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} \leq \frac {\sum_ {t \in [ T ] - \mathcal {N}} \boldsymbol {s} _ {t}}{\boldsymbol {s} _ {\tau}} \leq T \frac {e ^ {\nu \bar {R}}}{e ^ {2 \nu \bar {R}}} = T e ^ {- \bar {R} \nu}.
$$
Set  $\bar{M} = M / \Theta$  and note that  $\| \pmb{a}_t\| \leq \| \pmb{p}^{mm}\| \cdot \| \pmb{k}_t\| \leq \bar{M}$ . Using  $\bar{\pmb{p}}^{\top}(\pmb{k}_1 - \pmb{k}_t) \leq 1 - \nu$  over  $t \in \mathcal{N}$  and plugging in the above bound, we obtain
$$
\begin{array}{l} \frac {\sum_ {t \geq 2} ^ {T} (\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}) \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} = \frac {\sum_ {t \in \mathcal {N}} (\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}) \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} + \frac {\sum_ {t \in [ T ] - \mathcal {N}} (\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}) \boldsymbol {s} _ {t}}{\sum_ {t \geq 2} ^ {T} \boldsymbol {s} _ {t}} \\ \leq 1 - \nu + 2 \bar {M} T e ^ {- \bar {R} \nu}. \\ \end{array}
$$
Using the fact that  $\bar{a}_1 - \frac{\sum_{t\geq 2}^T\bar{a}_ts_t}{\sum_{t\geq 2}^Ts_t}\geq 1$ , the above implies (23) with  $\pi^{\prime} = 2\bar{M} Te^{-\bar{R}\nu} - \nu$ . To proceed, choose  $R_{\pi} = \nu^{-1}\Theta^{-1}\log (2\bar{M} T / \pi)$  to ensure  $\pi^\prime \leq \pi$