SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The following lemma states the descent property of gradient descent for  $\mathcal{L}(\pmb{p})$  under Assumption A. It is important to note that although the infimum of the optimization problem is  $\mathcal{L}^*$ , it is not achieved at any finite  $\pmb{p}$ . Additionally, there are no finite critical points  $\pmb{p}$ .
Lemma 6 Under Assumption  $A$ , the function  $\mathcal{L}(\pmb{p})$  is  $L_{p}$ -smooth, where
$$
L _ {p} := \frac {1}{n} \sum_ {i = 1} ^ {n} \left(M _ {0} \| \boldsymbol {v} \| ^ {2} \| \boldsymbol {W} \| ^ {2} \| \boldsymbol {X} _ {i} \| ^ {4} + 3 M _ {1} \| \boldsymbol {v} \| \| \boldsymbol {W} \| ^ {2} \| \boldsymbol {X} _ {i} \| ^ {3}\right). \tag {24}
$$
Furthermore, if  $\eta \leq 1 / L_p$ , then, for any initialization  $\pmb{p}(0)$ , with the GD sequence  $\pmb{p}(t + 1) = \pmb{p}(t) - \eta \nabla \mathcal{L}(\pmb{p}(t))$ , we have
$$
\mathcal {L} (\boldsymbol {p} (t + 1)) - \mathcal {L} (\boldsymbol {p} (t)) \leq - \frac {\eta}{2} \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2}, \tag {25}
$$
for all  $t \geq 0$ . This implies that
$$
\sum_ {t = 0} ^ {\infty} \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2} <   \infty , \text {a n d} \lim  _ {t \rightarrow \infty} \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2} = 0. \tag {26}
$$
Proof. Recall that we defined  $\gamma_{i} = Y_{i}\cdot X_{i}\nu$  and  $a_{i} = K_{i}p$ . The gradient of  $\mathcal{L}(\pmb {p})$  is given by
$$
\nabla \mathcal {L} (\boldsymbol {p}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell^ {\prime} \left(\boldsymbol {\gamma} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})\right) \cdot \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {a} _ {i}) \boldsymbol {\gamma} _ {i}.
$$
Note that for any  $\pmb{p} \in \mathbb{R}^d$ , the Jacobian of  $\mathbb{S}(\pmb{K}_i\pmb{p})$  is given by
$$
\frac {\partial \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})}{\partial \boldsymbol {p}} = \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \boldsymbol {p}) \boldsymbol {K} _ {i} = \left(\operatorname {d i a g} (\mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})) - \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p}) \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p}) ^ {\top}\right) \boldsymbol {K} _ {i}. \tag {27}
$$
The Jacobian (27) together with the definition of the softmax function  $\mathbb{S}(\cdot)$  implies that  $\| \partial \mathbb{S}(\pmb {K}_i\pmb {p}) / \partial \pmb {p}\| \leq \| \pmb {K}_i\|$ . Hence, for any  $\pmb {p},\dot{\pmb{p}}\in \mathbb{R}^d$ , we have
$$
\| \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p}) - \mathbb {S} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}}) \| \leq \| \boldsymbol {K} _ {i} \| \| \boldsymbol {p} - \dot {\boldsymbol {p}} \|, \tag {28a}
$$
$$
\begin{array}{l} \left\| \mathbb {S} ^ {\prime} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) - \mathbb {S} ^ {\prime} \left(\boldsymbol {K} _ {i} \boldsymbol {\dot {p}}\right) \right\| \leq \left\| \operatorname {d i a g} \left(\mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right)\right) - \operatorname {d i a g} \left(\mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {\dot {p}}\right)\right) \right\| \\ + \left\| \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p}) \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p}) ^ {\top} - \mathbb {S} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}}) \mathbb {S} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}}) ^ {\top} \right\| \\ \leq 3 \| \boldsymbol {K} _ {i} \| \| \boldsymbol {p} - \dot {\boldsymbol {p}} \|. \tag {28b} \\ \end{array}
$$
Here, the last inequality uses the fact that  $|ab - cd| \leq |d||a - c| + |a||b - d|$ .
Next, for any  $\pmb{p},\dot{\pmb{p}}\in \mathbb{R}^d$  , we have
$$
\begin{array}{l} \| \nabla \mathcal {L} (\boldsymbol {p}) - \nabla \mathcal {L} (\dot {\boldsymbol {p}}) \| \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\| \ell^ {\prime} \left(\boldsymbol {\gamma} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})\right) \cdot \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \boldsymbol {p}) \boldsymbol {\gamma} _ {i} - \ell^ {\prime} \left(\boldsymbol {\gamma} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}})\right) \cdot \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}}) \boldsymbol {\gamma} _ {i} \right\| \\ \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\| \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}}) \boldsymbol {\gamma} _ {i} \right\| \left| \ell^ {\prime} \left(\boldsymbol {\gamma} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})\right) - \ell^ {\prime} \left(\boldsymbol {\gamma} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \dot {\boldsymbol {p}})\right) \right| \\ + \frac {1}{n} \sum_ {i = 1} ^ {n} \left| \ell^ {\prime} \left(\boldsymbol {\gamma} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})\right) \right| \left\| \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) \boldsymbol {\gamma} _ {i} - \boldsymbol {K} _ {i} ^ {\top} \mathbb {S} ^ {\prime} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) \boldsymbol {\gamma} _ {i} \right\| \\ \leq \frac {1}{n} \sum_ {i = 1} ^ {n} M _ {0} \| \boldsymbol {\gamma} _ {i} \| ^ {2} \| \boldsymbol {K} _ {i} \| \| \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p}) - \mathbb {S} (\boldsymbol {K} _ {i} \hat {\boldsymbol {p}}) \| \\ + \frac {1}{n} \sum_ {i = 1} ^ {n} M _ {1} \| \boldsymbol {\gamma} _ {i} \| \| \boldsymbol {K} _ {i} \| \| \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \boldsymbol {p}) - \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \boldsymbol {\hat {p}}) \|, \tag {29} \\ \end{array}
$$
where the second inequality follows from the fact that  $|ab - cd| \leq |d||a - c| + |a||b - d|$  and the third inequality uses Assumption A.
Substituting (28a) and (28b) into (29), we get
$$
\begin{array}{l} \| \nabla \mathcal {L} (\boldsymbol {p}) - \nabla \mathcal {L} (\dot {\boldsymbol {p}}) \| \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left(M _ {0} \| \boldsymbol {\gamma} _ {i} \| ^ {2} \| \boldsymbol {K} _ {i} \| ^ {2} + 3 M _ {1} \| \boldsymbol {K} _ {i} \| ^ {2} \| \boldsymbol {\gamma} _ {i} \|\right) \| \boldsymbol {p} - \dot {\boldsymbol {p}} \| \\ \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left(M _ {0} \| \boldsymbol {v} \| ^ {2} \| \boldsymbol {W} \| ^ {2} \| \boldsymbol {X} _ {i} \| ^ {4} + 3 M _ {1} \| \boldsymbol {v} \| \| \boldsymbol {W} \| ^ {2} \| \boldsymbol {X} _ {i} \| ^ {3}\right) \| \boldsymbol {p} - \dot {\boldsymbol {p}} \| \\ \leq L _ {p} \| \boldsymbol {p} - \dot {\boldsymbol {p}} \|, \\ \end{array}
$$
where  $L_{p}$  is defined in (24).
The remaining proof follows standard gradient descent analysis (see e.g. [22, Lemma 10]). Since  $\mathcal{L}(\pmb{p})$  is  $L_{p}$ -smooth, we get
$$
\begin{array}{l} \mathcal {L} (\boldsymbol {p} (t + 1)) \leq \mathcal {L} (\boldsymbol {p} (t)) + \nabla \mathcal {L} (\boldsymbol {p} (t)) ^ {\top} (\boldsymbol {p} (t + 1) - \boldsymbol {p} (t)) + \frac {L _ {p}}{2} \| \boldsymbol {p} (t + 1) - \boldsymbol {p} (t) \| ^ {2} \\ = \mathcal {L} (\boldsymbol {p} (t)) - \eta \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2} + \frac {L _ {p} \eta^ {2}}{2} \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2} \\ = \mathcal {L} (\boldsymbol {p} (t)) - \eta \left(1 - \frac {L _ {p} \eta}{2}\right) \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2} \\ \leq \mathcal {L} (\boldsymbol {p} (t)) - \frac {\eta}{2} \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2}, \\ \end{array}
$$
where the last inequality follows from our assumption on the stepsize.
The above inequality implies that
$$
\sum_ {t = 0} ^ {\infty} \| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2} \leq \frac {2}{\eta} (\mathcal {L} (\boldsymbol {p} (0)) - \mathcal {L} ^ {*}), \tag {30}
$$
where the right hand side is upper bounded by a finite constant. This is because, by Assumption A,  $\mathcal{L}(\pmb{p}(0)) < \infty$  and  $\mathcal{L}^* \leq \mathcal{L}(\pmb{p}(t))$ , where  $\mathcal{L}^*$  denotes the minimum objective.
Finally, (30) yields the expression (26).
In the following lemma, we demonstrate the existence of parameters  $\mu = \mu (\alpha) > 0$  and  $R_{\mu} > 0$  such that when  $R_{\mu}$  is sufficiently large, there are no stationary points within  $C_{\mu ,R_\mu}(\pmb{p}^{mm})$ . Additionally, we provide the local gradient correlation condition.
Lemma 7 (Local Gradient Condition) Suppose Assumption A on the loss function  $\ell$  holds. Let  $\alpha = (\alpha_{i})_{i=1}^{n}$  be indices of locally-optimal tokens per Definition 2.
L1. There exists a positive scalar  $\mu = \mu (\alpha) > 0$  such that for sufficiently large  $\bar{R}_{\mu}$ , no stationary point exists within  $C_{\mu ,\bar{R}_\mu}(\pmb{p}^{mm})$ , where  $C_{\mu ,\bar{R}_\mu}$  is defined in (8).
L2. For all  $\pmb{q}, \pmb{p} \in \mathrm{cone}_{\mu}(\pmb{p}^{mm})$  with  $\| \pmb{q} \| = \| \pmb{p}^{mm} \|$  and  $\| \pmb{p} \| \geq \bar{R}_{\mu}$  with same  $\bar{R}_{\mu}$  choice as (L1.), there exist dataset dependent constants  $C, c > 0$  such that
$$
C \cdot \frac {1}{n} \sum_ {i \in [ n ]} \left\{1 - \mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) _ {\alpha_ {i}} \right\} \geq - \langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {q} \rangle \geq c \cdot \frac {1}{n} \sum_ {i \in [ n ]} \left\{1 - \mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) _ {\alpha_ {i}} \right\} > 0, \tag {31a}
$$
$$
\left\| \nabla \mathcal {L} (\boldsymbol {p}) \right\| \leq \bar {A} C \cdot \frac {1}{n} \sum_ {i \in [ n ]} \left\{1 - \mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) _ {\alpha_ {i}} \right\} \leq \bar {A} C T e ^ {- \bar {R} _ {\mu} \Theta / 2}. \tag {31b}
$$
$$
- \left\langle \frac {\boldsymbol {q}}{\| \boldsymbol {q} \|}, \frac {\nabla \mathcal {L} (\boldsymbol {p})}{\| \nabla \mathcal {L} (\boldsymbol {p}) \|} \right\rangle \geq \frac {c \Theta}{C \bar {A}} > 0, \tag {31c}
$$
Here,  $\bar{A} = \max_{i\in [n],t,\tau \in [T]}\| \pmb {k}_{it} - \pmb {k}_{i\tau}\|$  and  $\Theta = 1 / \| \pmb{p}^{mm}\|$
L3. For any  $\pi > 0$ , there exists  $R_{\pi}$  such that  $R_{\pi} \geq \bar{R}_{\mu}$  and all  $\pmb{p} \in C_{\mu, R_{\pi}}(\pmb{p}^{mm})$  obeys
Proof. Let  $p^{mm} = p^{mm}(\alpha)$  be the solution of (ATT-SVM). Recall
$$
\mathcal {C} _ {\mu , \bar {R} _ {\mu}} (\boldsymbol {p} ^ {m m}) = \operatorname {c o n e} _ {\mu} (\boldsymbol {p} ^ {m m}) \bigcap \left\{\boldsymbol {p} \mid \| \boldsymbol {p} \| \geq \bar {R} _ {\mu} \right\}.
$$
Let  $(\mathcal{T}_i)_{i=1}^n$  be the sets of all SVM-neighbors per Definition 2. Let  $\bar{\mathcal{T}}_i = [T] - \mathcal{T}_i - \{\alpha_i\}$  be the set of non-SVM-neighbor tokens,  $i \in [n]$ . Let
$$
\Theta = 1 / \| \boldsymbol {p} ^ {m m} \|,
$$
$$
\delta = 0.5\min_{i\in [n]}\min_{t\in \mathcal{T}_{i},\tau \in \bar{\mathcal{T}}_{i}}(\boldsymbol{k}_{it} - \boldsymbol{k}_{i\tau})^{\top}\boldsymbol{p}^{mm},
$$
$$
A = \max  _ {i \in [ n ], t \in [ T ]} \| \boldsymbol {k} _ {i t} \| / \Theta , \tag {32}
$$
$$
\mu \leq \mu (\delta) = \frac {1}{8} \left(\frac {\min (0 . 5 , \delta)}{A}\right) ^ {2}.
$$
When  $\bar{\mathcal{T}}_i = \emptyset$  for all  $i\in [n]$  (i.e. globally-optimal indices), we set  $\delta = \infty$  as all non-neighbor related terms will disappear. Since  $\pmb{p}^{mm}$  is the max-margin model ensuring  $(\pmb{k}_{i\alpha_i} - \pmb{k}_{it})^\top \pmb{p}^{mm}\geq 1$  for all  $i\in [n]$ , the following inequalities hold for all  $\pmb {q}\in \mathrm{cone}_{\mu}(\pmb{p}^{mm})$ ,  $\| \pmb {q}\| = \| \pmb{p}^{mm}\|$  and all  $i\in [n],t\in \mathcal{T}_i,\tau \in \mathcal{T}_i$ :
$$
\left(\boldsymbol {k} _ {i t} - \boldsymbol {k} _ {i \tau}\right) ^ {\top} \boldsymbol {q} \geq \delta > 0,
$$
$$
\left(\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i \tau}\right) ^ {\top} \boldsymbol {q} \geq 1 + \delta , \tag {33}
$$
$$
3 / 2 \geq (\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}) ^ {\top} \boldsymbol {q} \geq 1 / 2.
$$
Here, we used  $\| \pmb{q} - \pmb{p}^{mm}\| ^2 /\| \pmb{p}^{mm}\| ^2\leq 2\mu$  which implies  $\| \pmb {q} - \pmb{p}^{mm}\| \leq \sqrt{2\mu} /\Theta$
L1. and L2.. Now that the choice of local cone is determined, we need to prove the main claims. We will lower bound  $-\pmb{q}^{\top}\nabla \mathcal{L}(\pmb{p})$  and establish its strict positivity for  $\| \pmb {p}\| \geq R$ , where  $R = \bar{R}_{\mu}$ . This will show that there is no stationary point as a by product.
Consider any  $\pmb{q} \in \mathbb{R}^d$  satisfying  $\| \pmb{q} \| = \| \pmb{p}^{mm} \|$ . To proceed, we write the gradient correlation following (18) and (21)
$$
\left\langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {q} \right\rangle = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \left\langle \boldsymbol {a} _ {i}, \mathbb {S} ^ {\prime} \left(\boldsymbol {a} _ {i} ^ {\prime}\right) \gamma_ {i} \right\rangle , \tag {34}
$$
where we denoted  $\ell_i^\prime = \ell^\prime (Y_i\cdot \pmb{v}^\top \pmb {X}_i^\top \mathbb{S}(\pmb {K}_i\pmb {p}))$ $a_{i} = K_{i}\pmb {q},a_{i}^{\prime} = K_{i}\pmb {p},s_{i} = \mathbb{S}(K_{i}\pmb {p})$
Using (33), for all  $t \in \mathcal{T}_i$ ,  $\tau \in \tilde{\mathcal{T}}_i$ , for all  $\pmb{p} \in C_{\mu,R}(\pmb{p}^{mm})$ , we have that
$$
\boldsymbol {a} _ {i \alpha_ {i}} ^ {\prime} - \boldsymbol {a} _ {i \tau} ^ {\prime} \geq R \Theta (1 + \delta), \quad \text {a n d} \quad \boldsymbol {a} _ {i t} ^ {\prime} - \boldsymbol {a} _ {i \tau} ^ {\prime} \geq R \Theta \delta .
$$
Consequently, we can bound the softmax probabilities  $s_i = \mathbb{S}(K_i p)$  as follows: For all  $i \in [n]$ ,
$$
S _ {i} := \sum_ {\tau \in \mathcal {T} _ {i}} \boldsymbol {s} _ {i \tau} \leq 1 - \boldsymbol {s} _ {i \alpha_ {i}} = \sum_ {\tau \neq \alpha_ {i}} \boldsymbol {s} _ {i \tau} \leq T e ^ {- R \Theta / 2} \boldsymbol {s} _ {i \alpha_ {i}} \leq T e ^ {- R \Theta / 2},
$$
$$
Q _ {i} := \sum_ {\tau \in \mathcal {T} _ {i}} s _ {i \tau} \leq T e ^ {- R \Theta \delta} s _ {i t _ {i}} \leq T e ^ {- R \Theta \delta} S _ {i}, \forall t _ {i} \in \mathcal {T} _ {i}. \tag {35}
$$
Recall scores  $\pmb{\gamma}_{it} = Y_{i}\cdot \pmb{v}^{\top}\pmb{x}_{it}$ . Define the score gaps over neighbors:
$$
\gamma_ {i} ^ {\text {g a p}} = \gamma_ {i \alpha_ {i}} - \max  _ {t \in \mathcal {T} _ {i}} \gamma_ {i t}, \quad \text {a n d} \quad \bar {\gamma} _ {i} ^ {\text {g a p}} = \gamma_ {i \alpha_ {i}} - \min  _ {t \in \mathcal {T} _ {i}} \gamma_ {i t}.
$$
It follows from (32) that
$$
A = \max  _ {i \in [ n ], t \in [ T ]} \| \boldsymbol {k} _ {i t} \| / \Theta \geq \max  _ {i \in [ n ], t \in [ T ]} \| \boldsymbol {a} _ {i t} \| = \| \boldsymbol {k} _ {i t} \boldsymbol {q} \|.
$$
Define the  $\alpha$ -dependent global scalar  $\Gamma = \sup_{i\in [n],t,\tau \in [T]}|\gamma_{it} - \gamma_{i\tau}|$ . Let us focus on a fixed datapoint  $i\in [n]$ , assume (without losing generality)  $\alpha_{i} = 1$ , and drop subscripts  $i$ , that is,  $\alpha \coloneqq \alpha_{i} = 1$ ,  $X\coloneqq X_{i}$ ,  $Y\coloneqq Y_{i}$ ,  $K\coloneqq K_{i}$ ,  $a^{\prime} = Kp$ ,  $a = Kq$ ,  $s = \mathbb{S}(Kp)$ ,  $\gamma = Y\cdot X\nu$ ,  $\gamma^{gap}\coloneqq \gamma_i^{gap}$ ,  $\bar{\gamma}^{gap}\coloneqq \bar{\gamma}_i^{gap}$ ,  $Q\coloneqq Q_{i}$ , and  $S\coloneqq S_{i}$ . Directly applying Lemma 3, we obtain
$$
\left| \boldsymbol {a} ^ {\top} \operatorname {d i a g} (s) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} s s ^ {\top} \boldsymbol {\gamma} - \sum_ {t \geq 2} ^ {T} (\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}) s _ {t} (\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}) \right| \leq 2 \Gamma A (1 - s _ {1}) ^ {2}.
$$
To proceed, let us decouple the non-neighbors within  $\sum_{t\geq 2}^{T}(\pmb{a}_{1} - \pmb{a}_{t})s_{t}(\pmb{\gamma}_{1} - \pmb{\gamma}_{t})$  via
$$
\left| \sum_ {t \in \widetilde {\mathcal {T}}} \left(\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}\right) \boldsymbol {s} _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right) \right| \leq 2 Q \Gamma A.
$$
Aggregating these, we found
$$
\left| \boldsymbol {a} ^ {\top} \operatorname {d i a g} (s) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} s s ^ {\top} \boldsymbol {\gamma} - \sum_ {t \in \mathcal {T} _ {i}} \left(\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}\right) s _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right) \right| \leq 2 \Gamma A \left((1 - s _ {1}) ^ {2} + Q\right). \tag {36}
$$
To proceed, let us upper/lower bound the gradient correlation. We use two bounds depending on  $\pmb{q} \in \mathrm{cone}_{\mu}(\pmb{p}^{mm})$  (Case 1) or general  $\pmb{q} \in \mathbb{R}^d$  (Case 2).
- Case 1:  $\pmb{q} \in \mathrm{cone}_{\mu}(\pmb{p}^{mm})$ . Since  $1.5 \geq \pmb{a}_1 - \pmb{a}_t \geq 0.5$  following (33), we find
$$
1. 5 \cdot S \cdot \bar {\gamma} ^ {g a p} \geq \sum_ {t \in \mathcal {T} _ {i}} \left(\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}\right) s _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right) \geq 0. 5 \cdot S \cdot \gamma^ {g a p}.
$$
Next we claim that  $S$  dominates  $((1 - s_1)^2 + Q)$  for large  $R$ . Specifically, we wish for
$$
S \cdot \gamma^ {\text {g a p}} / 4 \geq 4 \Gamma A \max  \left((1 - s _ {1}) ^ {2}, Q\right) \Longleftrightarrow S \geq 1 6 \frac {\Gamma A}{\gamma^ {\text {g a p}}} \max  \left((1 - s _ {1}) ^ {2}, Q\right). \tag {37}
$$
Now choose  $R \geq \delta^{-1} \log(T) / \Theta$  to ensure  $Q \leq S$  since  $Q \leq Te^{-R\Theta \delta}S$ . Consequently
$$
(1 - s _ {1}) ^ {2} = (Q + S) ^ {2} \leq 4 S ^ {2} \leq 4 S T e ^ {- R \Theta / 2}.
$$
Combining these, what we wish is ensured by guaranteeing
$$
S \geq 1 6 \frac {\Gamma A}{\gamma^ {\text {g a p}}} \max  \left(4 S T e ^ {- R \Theta / 2}, T e ^ {- R \Theta \delta} S\right). \tag {38}
$$
This in turn is ensured for all inputs  $i \in [n]$  by choosing
$$
R = \frac {\operatorname* {m a x} (2 , \delta^ {- 1})}{\Theta} \log \left(\frac {6 4 T \Gamma A}{\gamma_ {\min } ^ {\text {g a p}}}\right), \tag {39}
$$
where  $\gamma_{\mathrm{min}}^{\mathrm{gap}} = \min_{i\in [n]}\gamma_i^{\mathrm{gap}}$  is the global scalar which is the worst case score gap over all inputs. With the above choice of  $R$  we guaranteed
$$
2 (1 - s _ {1}) \cdot \bar {\gamma} ^ {g a p} \geq 2 \cdot S \cdot \bar {\gamma} ^ {g a p} \geq \boldsymbol {a} ^ {\top} \operatorname {d i a g} (s) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} s s ^ {\top} \boldsymbol {\gamma} \geq \frac {S \cdot \gamma^ {g a p}}{4} \geq \frac {(1 - s _ {1}) \gamma^ {g a p}}{8}.
$$
Since this holds over all inputs, going back to the gradient correlation (34) and averaging above over all inputs  $i \in [n]$  and plugging back the indices  $i$ , we obtain the advertised bound by setting  $q_{i} = 1 - s_{i\alpha_{i}}$  (where we set  $\alpha_{i} = 1$  above without losing generality)
$$
\frac {2}{n} \sum_ {i \in [ n ]} - \ell_ {i} ^ {\prime} \cdot q _ {i} \cdot \tilde {\gamma} _ {i} ^ {\text {g a p}} \geq - \langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {q} \rangle \geq \frac {1}{8 n} \sum_ {i \in [ n ]} - \ell_ {i} ^ {\prime} \cdot q _ {i} \cdot \gamma_ {i} ^ {\text {g a p}}. \tag {40}
$$
Let  $-\ell_{\min / \max}^{\prime}$  be the min/max values negative loss derivative admits over the ball  $[-B, B]$  for  $B = \| \pmb{v} \| \cdot \max_{i,t} \| \pmb{x}_{it} \|$  and note that  $\max_{i \in [n]} \bar{\gamma}_i^{gap} > 0$  and  $\min_{i \in [n]} \gamma_i^{gap} > 0$  are dataset-dependent constants. Then, we declare the constants  $C = -2\ell_{\max}^{\prime} \cdot \max_{i \in [n]} \bar{\gamma}_i^{gap} > 0$ ,  $c = -(1/8)\ell_{\min}^{\prime} \cdot \min_{i \in [n]} \gamma_i^{gap} > 0$  to obtain the bound
$$
\frac {C}{n} \sum_ {i \in [ n ]} q _ {i} \geq - \langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {q} \rangle \geq \frac {c}{n} \sum_ {i \in [ n ]} q _ {i}, \tag {41}
$$
which is the desired statement in (31a).
- Case 2:  $\pmb{q} \in \mathbb{R}^d$  and  $\| \pmb{q} \| = \| \pmb{p}^{mm} \|$ . Define  $\bar{A} = \max_{i \in [n], t, \tau \in [T]} \| \pmb{k}_{it} - \pmb{k}_{i\tau} \|$ . For any  $\| \pmb{q} \| = \| \pmb{p}^{mm} \|$ , we use the fact that
$$
\left\| \boldsymbol {a} _ {1} - \boldsymbol {a} _ {t} \right\| \leq \left\| \boldsymbol {k} _ {1} - \boldsymbol {k} _ {t} \right\| \cdot \left\| \boldsymbol {q} \right\| \leq \frac {\bar {A}}{\Theta}.
$$
Note that by definition  $\frac{\bar{A}}{\Theta} \geq 1$ . To proceed, we can upper bound
$$
\frac {\bar {A}}{\Theta} \cdot S \cdot \bar {\gamma} ^ {\text {g a p}} \geq \sum_ {t \in \mathcal {T}} \left(\boldsymbol {a} _ {1} - \boldsymbol {a} _ {t}\right) \boldsymbol {s} _ {t} \left(\boldsymbol {\gamma} _ {1} - \boldsymbol {\gamma} _ {t}\right). \tag {42}
$$
By choosing the same  $R$  as in (39) to ensure  $S$  dominates  $((1 - s_1)^2 +Q)$  and since  $\frac{\bar{A}}{\Theta}\geq 1$ , we guaranteed
$$
\frac {2 \bar {A}}{\Theta} \cdot S \cdot \bar {\gamma} ^ {\text {g a p}} \geq \boldsymbol {a} ^ {\top} \operatorname {d i a g} (s) \boldsymbol {\gamma} - \boldsymbol {a} ^ {\top} s s ^ {\top} \boldsymbol {\gamma}.
$$
Going back to the gradient correlation (34) and averaging above over all inputs  $i \in [n]$ , with the same definition of  $C > 0$ , we obtain
$$
\frac {\bar {A} C}{\Theta n} \sum_ {i \in [ n ]} q _ {i} \geq - \langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {q} \rangle . \tag {43}
$$
To proceed, since (43) holds for any  $\pmb{q} \in \mathbb{R}^d$  and  $\| \pmb{q} \| = \| \pmb{p}^{mm} \|$ , we observe that when choosing  $\pmb{q} = \frac{\|\pmb{p}^{mm}\|}{\|\nabla \mathcal{L}(\pmb{p})\|} \cdot \nabla \mathcal{L}(\pmb{p})$ , this implies that
$$
\langle \nabla \mathcal {L} (\boldsymbol {p}), \boldsymbol {q} \rangle = \| \nabla \mathcal {L} (\boldsymbol {p}) \| \cdot \| \boldsymbol {p} ^ {m m} \| \leq \frac {\bar {A} C}{\Theta n} \sum_ {i \in [ n ]} q _ {i}.
$$
Simplifying  $\Theta = 1 / \| \pmb{p}^{mm}\|$  on both sides yields (31b). Incorporating (35) in the bound above provides the exponential upper bound that decay with  $R$ .
Combining this with (41), we obtain that for all  $\pmb{q}, \pmb{p} \in \mathrm{cone}_{\mu}(\pmb{p}^{mm})$  and  $\| \pmb{q} \| \geq \bar{R}_{\mu}$
$$
- \left\langle \frac {\boldsymbol {q}}{\| \boldsymbol {q} \|}, \frac {\nabla \mathcal {L} (\boldsymbol {p})}{\| \nabla \mathcal {L} (\boldsymbol {p}) \|} \right\rangle \geq \frac {c \Theta}{C \bar {A}}.
$$
This gives the desired result in (31c).
L3.: Establishing gradient correlation.
Our final goal is establishing gradient comparison between  $\pmb{p}, \pmb{p}^{mm}$  for the same choice of  $\mu > 0$  provided in (32). Define  $\bar{\pmb{p}} = \| \pmb{p}^{mm} \| \pmb{p} / \| \pmb{p} \|$  to be the normalized vector. Set notations  $\pmb{a}_i = K_i \bar{\pmb{p}}$ ,  $\bar{\pmb{a}}_i = K_i \pmb{p}^{mm}$ , and  $\gamma_i = Y_i \cdot X_i \nu$ .
To establish the result, using (34), we will prove that, for any  $\pi > 0$ , there is sufficiently large  $R = R_{\pi}$  such that for any  $\pmb{p} \in C_{\mu,R}(\pmb{p}^{mm})$ :
$$
\begin{array}{l} \left\langle - \nabla \mathcal {L} (\boldsymbol {p}), \frac {\boldsymbol {p}}{\| \boldsymbol {p} \|} \right\rangle = - \frac {1}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \left\langle \boldsymbol {a} _ {i}, \mathbb {S} ^ {\prime} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) \boldsymbol {\gamma} _ {i} \right\rangle \\ \leq - \frac {1 + \pi}{n} \sum_ {i = 1} ^ {n} \ell_ {i} ^ {\prime} \cdot \left\langle \bar {\boldsymbol {a}} _ {i}, \mathbb {S} ^ {\prime} (\boldsymbol {K} _ {i} \boldsymbol {p}) \boldsymbol {\gamma} _ {i} \right\rangle = (1 + \pi) \left\langle - \nabla \mathcal {L} (\boldsymbol {p}), \frac {\boldsymbol {p} ^ {m m}}{\| \boldsymbol {p} ^ {m m} \|} \right\rangle . \tag {44} \\ \end{array}
$$
Following (36), for all  $i \in [n]$ , for all  $\pmb{q} \in \mathrm{cone}_{\mu}(\pmb{p}^{mm})$  with  $\| \pmb{q} \| = \| \pmb{p}^{mm} \|$ ,  $\pmb{a}' = \pmb{K}\pmb{q}$  and  $s = \mathbb{S}(\pmb{K}\pmb{p})$ , we have found