SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {(\mu - \epsilon)}{1 - \mu} \frac {\rho (t)}{\| \nabla \mathcal {L} (\boldsymbol {p} (t)) \| ^ {2}} \geq \frac {\mu - \epsilon}{1 - \mu} \frac {1}{1 - \epsilon} \frac {c}{C} \frac {\Theta}{\bar {A}} \frac {1}{\bar {A} C T} e ^ {R _ {\mu} ^ {0} \Theta / 2} \\ \geq \frac {\mu}{2 (1 - \mu) \left(1 - \frac {\mu}{2}\right)} \frac {c}{C} \frac {\Theta}{\bar {A}} \frac {1}{\bar {A} C T} e ^ {R _ {\mu} ^ {0} \Theta / 2} \geq \eta . \tag {56} \\ \end{array}
$$
Here, the second inequality uses our choice of  $\epsilon \in (0,\mu /2)$  (see Step 2), and the first inequality is obtained from Lemma 7 since
$$
\frac {\rho (t)}{\| \nabla \mathcal {L} (\boldsymbol {p} (t)) \|} = - \frac {1}{1 - \epsilon} \left\langle \frac {\nabla \mathcal {L} (\boldsymbol {p} (t))}{\| \nabla \mathcal {L} (\boldsymbol {p} (t)) \|}, \frac {\boldsymbol {p} ^ {m m}}{\| \boldsymbol {p} ^ {m m} \|} \right\rangle \geq \frac {1}{1 - \epsilon} \frac {c}{C} \frac {\Theta}{\bar {A}},
$$
$$
\frac {1}{\| \nabla \mathcal {L} (\boldsymbol {p} (t)) \|} \geq \frac {1}{\bar {A} C _ {n} ^ {\frac {1}{n}} \sum_ {i = 1} ^ {n} \left(1 - s _ {i \alpha_ {i}}\right)} \geq \frac {1}{\bar {A} C T e ^ {- R _ {\mu} ^ {0} \Theta / 2}}
$$
for some data dependent constants  $c$  and  $C$ ,  $\bar{A} = \max_{i\in [n],t,\tau \in [T]}\| \pmb{k}_{it} - \pmb{k}_{it}\|$ , and  $\Theta = 1 / \| \pmb{p}^{mm}\|$ .
Next, we will demonstrate that the choice of  $\eta$  in (56) does indeed meet our step size condition as stated in the theorem, i.e.,  $\eta \leq 1 / L_p$ . Recall that  $1 / (1 + \pi) = 1 - \epsilon$ , which implies that  $\pi = \epsilon / (1 - \epsilon)$ . Combining this with (46), we obtain:
$$
\begin{array}{l} R _ {\pi} \geq \frac {\operatorname* {m a x} (2 , \delta^ {- 1})}{\Theta} \log \left(\frac {C _ {0} T \Gamma A}{\pi \gamma_ {\mathrm {m i n}} ^ {\mathrm {g a p}}}\right), \quad \text {w h e r e} \quad C _ {0} \geq 6 4 \pi , \\ \Rightarrow R _ {\epsilon} \geq \frac {\operatorname* {m a x} (2 , \delta^ {- 1})}{\Theta} \log \left(\frac {(1 - \epsilon) C _ {0} T \Gamma A}{\epsilon \gamma_ {\mathrm {m i n}} ^ {\mathrm {g a p}}}\right), \quad \text {w h e r e} \quad C _ {0} \geq 6 4 \frac {\epsilon}{1 - \epsilon}. \\ \end{array}
$$
On the other hand, at the initialization, we have  $\epsilon = \mu /2$  which implies that
$$
R _ {\mu} ^ {0} \geq \frac {\operatorname* {m a x} (2 , \delta^ {- 1})}{\Theta} \log \left(\frac {(2 - \mu) C _ {0} T \Gamma A}{\mu \gamma_ {\min } ^ {\text {g a p}}}\right), \quad \text {w h e r e} \quad C _ {0} \geq 6 4 \frac {\mu}{(2 - \mu)}. \tag {57}
$$
In the following, we will determine a lower bound on  $C_0$  such that our step size condition in Theorem 3's statement, i.e.,  $\eta \leq 1 / L_p$ , is satisfied. Note that for the choice of  $\eta$  in (56) to meet the condition  $\eta \leq 1 / L_p$ , the following condition must hold:
$$
\frac {1}{L _ {p}} \leq \frac {\mu}{(2 - \mu)} \frac {1}{C _ {2} T} e ^ {R _ {\mu} ^ {0} \Theta / 2} \Rightarrow R _ {\mu} ^ {0} \geq \frac {2}{\Theta} \log \left(\frac {1}{L _ {p}} \frac {(2 - \mu)}{\mu} C _ {2} T\right), \tag {58}
$$
where  $C_2 = (1 - \mu)\frac{\bar{A}^2C^2}{\Theta c}$
This together with (57) implies that for sufficiently large
$$
R _ {\mu} ^ {0} \geq \frac {\max  (2 , \delta^ {- 1})}{\Theta} \log \left(\frac {(2 - \mu) C _ {3} T}{\mu}\right), \quad \text {w h e r e} \quad C _ {3} = \frac {C _ {0} \Gamma A}{\gamma_ {\min } ^ {\text {g a p}}} \vee \frac {C _ {2}}{L _ {p}},
$$
the step size bound in (56) ensures that  $\eta \leq 1 / L_p$  guarantees (55). Hence,  $\pmb{p}(t + 1)$  remains within the cone, i.e.,  $\pmb{p}(t + 1) \in C_{\mu, R_{\epsilon}}(\pmb{p}^{mm})$ .
Step 4: The correlation of  $\pmb{p}(t)$  and  $\pmb{p}^{mm}$  increases over  $t$ . The remainder is similar to the proof of Theorem 2. From Step 3, we have that all iterates remain within the initial conic set i.e.  $\pmb{p}(t) \in C_{\mu,R_\mu^0}(\pmb{p}^{mm})$  for all  $t \geq 0$ . Note that it follows from Lemma 7 that  $\langle \nabla \mathcal{L}(\pmb{p}), \pmb{p}^{mm} / \| \pmb{p}^{mm} \| \rangle < 0$ , for any finite  $\pmb{p} \in C_{\mu,R_\mu^0}(\pmb{p}^{mm})$ . Hence, there are no finite critical points  $\pmb{p} \in C_{\mu,R_\mu^0}(\pmb{p}^{mm})$ , for which  $\nabla \mathcal{L}(\pmb{p}) = 0$ . Now, based on Lemma 6, which guarantees that  $\nabla \mathcal{L}(\pmb{p}(t)) \to 0$ , this implies that  $\| \pmb{p}(t) \| \to \infty$ . Consequently, for any choice of  $\epsilon \in (0,\mu/2)$  there is a time  $t_\epsilon$  such that, for all  $t \geq t_\epsilon$ ,  $\pmb{p}(t) \in C_{\mu,R_\epsilon}(\pmb{p}^{mm})$ . Once within  $C_{\mu,R_\epsilon}(\pmb{p}^{mm})$ , following similar steps in (51) and (52), for any  $t \geq t_\epsilon$ ,
$$
\left\langle \frac {\boldsymbol {p} (t)}{\| \boldsymbol {p} (t) \|}, \frac {\boldsymbol {p} ^ {m m}}{\| \boldsymbol {p} ^ {m m} \|} \right\rangle \geq 1 - \epsilon + \frac {C _ {2} (\epsilon , \eta)}{\| \boldsymbol {p} (t) \|}, \quad \boldsymbol {p} (t) \in C _ {\mu , R _ {\epsilon}} (\boldsymbol {p} ^ {m m}),
$$
for some finite constant  $C_2(\epsilon, \eta)$ . Consequently,
$$
\lim  _ {t \rightarrow \infty} \inf  _ {\epsilon} \left\langle \frac {\boldsymbol {p} (t)}{\| \boldsymbol {p} (t) \|}, \frac {\boldsymbol {p} ^ {m m}}{\| \boldsymbol {p} ^ {m m} \|} \right\rangle \geq 1 - \epsilon , \text {w h e r e} \boldsymbol {p} (t) \in C _ {\mu , R _ {\epsilon}} (\boldsymbol {p} ^ {m m}).
$$
Since the choice of  $\epsilon \in (0,\mu /2)$  is arbitrary, we obtain  $\pmb {p}(t) / \| \pmb {p}(t)\| \to \pmb{p}^{mm} / \| \pmb{p}^{mm}\|$ .
B.5.1 Supporting Lemma
We present a lemma that will aid in simplifying our analysis. We begin with a definition.
Definition 3 (Selected-tokens, Neighbors, Margins, and Neighbor-optimality of a direction) Let  $\mathbf{q} \in \mathbb{R}^d - \{\mathbf{0}\}$  and  $(Y_i, K_i, X_i)_{i=1}^n$  be our dataset. We define the (possibly non-unique) selected-tokens of  $\mathbf{q}$  as follows:
$$
\alpha_ {i} \in \arg \max  _ {t \in [ T ]} \boldsymbol {k} _ {i t} ^ {\top} \boldsymbol {q}. \tag {59}
$$
Next, we define the margin and directional-neighbors for  $\mathbf{q}$  as the minimum margin tokens to the selected-tokens, i.e.,
$$
\Gamma_ {\boldsymbol {q}} = \min  _ {i \in [ n ], t \neq \alpha_ {i}} \left(\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}\right) ^ {\top} \boldsymbol {q}, \tag {60}
$$
$$
\mathcal {M} _ {\boldsymbol {q}} = \left\{(i, t) \mid \left(\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}\right) ^ {\top} \boldsymbol {q} = \Gamma_ {\boldsymbol {q}} \right\}. \tag {61}
$$
Finally, we say that  $\mathbf{q}$  is neighbor-optimal if the scores of its directional-neighbors are strictly less than the corresponding selected-token. Concretely, for all  $(i,t) \in \mathcal{M}_q$ , we require that
$$
\boldsymbol {\gamma} _ {i t} = Y _ {i} \cdot \boldsymbol {x} _ {i t} ^ {\top} \boldsymbol {v} <   \boldsymbol {\gamma} _ {i \alpha_ {i}} = Y _ {i} \cdot \boldsymbol {x} _ {i \alpha_ {i}} ^ {\top} \boldsymbol {v}.
$$
Lemma 8 (When does one direction dominate another?) Suppose  $\pmb{q}, \pmb{p} \in \mathbb{R}^d$  be two unit Euclidean norm vectors with identical selected tokens. Specifically, for each  $i \in [n]$ , there exists unique  $\alpha_i \in [T]$  such that  $\alpha_i = \arg \max_{t \in [T]} \pmb{k}_{it}^\top \pmb{q} = \arg \max_{t \in [T]} \pmb{k}_{it}^\top \pmb{p}$ . Suppose directional margins obey  $\Gamma_q < \Gamma_p$  and set  $\delta_\Gamma = \Gamma_p - \Gamma_q$ .
- Suppose  $\mathbf{q}$  and  $\mathbf{p}$  are both neighbor-optimal. Then, for some  $R(\delta_{\Gamma})$  and all  $R > R(\delta_{\Gamma})$ , we have that  $\mathcal{L}(R \cdot \mathbf{p}) < \mathcal{L}(R \cdot \mathbf{q})$ .
- Suppose  $\mathbf{q}$  has a unique directional-neighbor and is not neighbor-optimal (i.e. this neighbor has higher score). Let  $\delta_{\mathbf{q}}$  be the margin difference between unique directional-neighbor and the second-most minimum-margin neighbor (i.e. the one after the unique one, see (65)) of  $\mathbf{q}$ . Then, for some  $R(\delta_{\Gamma} \wedge \delta_{\mathbf{q}})$  and all  $R > R(\delta_{\Gamma} \wedge \delta_{\mathbf{q}})$ , we have that  $\mathcal{L}(R \cdot \mathbf{q}) < \mathcal{L}(R \cdot \mathbf{p})$ .
Proof. We prove these two statements in order. First define the directional risk baseline induced by letting  $R \to \infty$  and purely selecting the tokens  $\alpha = (\alpha_{i})_{i=1}^{n}$ . This is given by
$$
\mathcal {L} _ {\star} := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(Y _ {i} \cdot v ^ {\top} x _ {i \alpha_ {i}}\right).
$$
$^3$ If  $\alpha_{i}$  is unique for all  $i\in [n]$ , let us call it, unique selected tokens.
We evaluate  $\pmb{q}, \pmb{p}$  with respect to  $\mathcal{L}_{\star}$ . To proceed, let  $s_i = \mathbb{S}(R\pmb{K}_i\pmb{q})$ . Define  $\Gamma_{\pmb{q}}^{it} = \pmb{k}_{i\alpha_i}^{\top}\pmb{q} - \pmb{k}_{it}^{\top}\pmb{q}$ . Note that, the smallest value for  $t \neq \alpha_i$  is achieved for  $\Gamma_{\pmb{q}}$ . For sufficiently large  $R \gtrsim O(\log (T) / \Gamma_{\pmb{q}})$ , observe that, for  $t \neq \alpha_i$
$$
e ^ {- R \Gamma_ {q} ^ {i t}} \geq s _ {i t} = \frac {e ^ {R k _ {i t} ^ {\top} q}}{\sum_ {t \in [ T ]} e ^ {R k _ {i t} ^ {\top} q}} \geq 0. 5 e ^ {- R \Gamma_ {q} ^ {i t}}. \tag {62}
$$
Recalling the score definition and let  $M_{+}, M_{-}$  be the upper and lower bounds on  $-\ell'$  over its bounded domain that scores fall on, respectively. Note that, for some intermediate  $M_{+} \geq M_{i} \geq M_{-}$  values, we have
$$
\begin{array}{l} \mathcal {L} (R \boldsymbol {q}) - \mathcal {L} _ {\star} = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (\sum_ {t \in [ T ]} \boldsymbol {s} _ {i t} \boldsymbol {\gamma} _ {i t}) - \ell (\boldsymbol {\gamma} _ {i \alpha_ {i}}) \\ = \frac {1}{n} \sum_ {i = 1} ^ {n} M _ {i} \sum_ {t \neq \alpha_ {i}} s _ {i t} \left(\gamma_ {i \alpha_ {i}} - \gamma_ {i t}\right). \\ \end{array}
$$
Now, using (62) for a refreshed  $M_{+} \geq M_{it} \geq 0.5M_{-}$  values, we can write
$$
\mathcal {L} (R \boldsymbol {q}) - \mathcal {L} _ {\star} = \frac {1}{n} \sum_ {i = 1} ^ {n} \sum_ {t \neq \alpha_ {i}} M _ {i t} e ^ {- R \Gamma_ {q} ^ {- i t}} \left(\boldsymbol {\gamma} _ {i \alpha_ {i}} - \boldsymbol {\gamma} _ {i t}\right). \tag {63}
$$
The same bound also applies to  $\pmb{p}$  with some  $M_{it}'$ , multipliers
$$
\mathcal {L} (R p) - \mathcal {L} _ {\star} = \frac {1}{n} \sum_ {i = 1} ^ {n} \sum_ {t \neq \alpha_ {i}} M _ {i t} ^ {\prime} e ^ {- R \Gamma_ {p} ^ {- i t}} \left(\boldsymbol {\gamma} _ {i \alpha_ {i}} - \boldsymbol {\gamma} _ {i t}\right). \tag {64}
$$
We can now proceed with the proof.
Case 1:  $q$  and  $p$  are both neighbor-optimal. This means that  $\gamma_{i\alpha_i} - \gamma_{it} > 0$  for all  $i \in [n], t \neq \alpha_i$ . Let  $K_+ > K_- > 0$  be upper and lower bounds on  $\gamma_{i\alpha_i} - \gamma_{it}$  values. We can now upper bound the right hand side of (63) via
$$
M _ {+} K _ {+} T e ^ {- R \Gamma_ {q}} \geq \mathcal {L} (R q) - \mathcal {L} _ {\star} \geq \frac {1}{2 n} M _ {-} K _ {-} e ^ {- R \Gamma_ {q}}.
$$
Consequently,  $\mathcal{L}(R\pmb {q}) > \mathcal{L}(R\pmb {p})$  as soon as  $\frac{1}{2n} M_{-}K_{-}e^{-R\Gamma_{q}} > M_{+}K_{+}Te^{-R\Gamma_{p}}$ . Since  $M_{+},K_{+},n,T$  are global constants, this happens under the stated condition on the margin gap  $\Gamma_p - \Gamma_q$ .
Case 2:  $q$  has a unique directional-neighbor and is not neighbor-optimal. In this scenario,  $\mathcal{L}(Rq) - \mathcal{L}_{\star}$  is actually negative for large  $R$ . To proceed, define the maximum score difference  $K_{+} = \sup_{i,t\neq \alpha_{i}}|\gamma_{i\alpha_{i}} - \gamma_{it}|$ . Also let  $(j,\beta)$  be the unique directional neighbor achieving the minimum margin  $\Gamma_q$ . Then,  $\delta_q$  – the margin difference between unique directional-neighbor and the second minimum-margin neighbor (i.e. the one after the unique one) of  $q$  – is defined as
$$
\delta_ {q} = \min  _ {i \in [ n ], t \neq \alpha_ {i}, (i, t) \neq (j, \beta)} \Gamma_ {q} ^ {i t} - \Gamma_ {q}. \tag {65}
$$
To proceed, we can write
$$
\mathcal {L} (R p) - \mathcal {L} _ {\star} \geq - M _ {+} K _ {+} T e ^ {- R \Gamma_ {p}}.
$$
On the other hand, setting  $\kappa = \gamma_{j\beta} - \gamma_{j\alpha_j} > 0$  we can bound
$$
\begin{array}{l} \mathcal {L} (R \boldsymbol {q}) - \mathcal {L} _ {\star} = - \frac {1}{n} M _ {j \beta} e ^ {- R \Gamma_ {q}} K + \frac {1}{n} \sum_ {i \in [ n ], t \neq \alpha_ {i}, (i, t) \neq (j, \beta)} M _ {i t} e ^ {- R \Gamma_ {q} ^ {i t}} (\boldsymbol {\gamma} _ {i \alpha_ {i}} - \boldsymbol {\gamma} _ {i t}) \\ \leq - \frac {1}{n} M _ {-} e ^ {- R \Gamma_ {q} K} + M _ {+} K _ {+} T e ^ {- R \left(\Gamma_ {q} + \delta_ {q}\right)}. \\ \end{array}
$$
Consequently, we have found that  $\mathcal{L}(Rp) > \mathcal{L}(Rq)$  as soon as
$$
\frac {1}{n} M _ {-} e ^ {- R \Gamma_ {q} K} \geq M _ {+} K _ {+} T \left(e ^ {- R \left(\Gamma_ {q} + \delta_ {q}\right)} + e ^ {- R \Gamma_ {p}}\right)
$$
This happens when  $R \gtrsim \frac{1}{\delta_q \wedge (\Gamma_p - \Gamma_q)}$  (up to logarithmic terms) establishing the desired statement.
B.5.2 Proof of Theorem 4
Define the locally-optimal unit directions
$$
\mathcal {P} ^ {m m} = \left\{\frac {\boldsymbol {p} ^ {m m} (\alpha)}{\| \boldsymbol {p} ^ {m m} (\alpha) \|} \mid \alpha \text {i s a l o c a l l y - o p t i m a l s e t o f i n d i c e s} \right\}.
$$
The theorem below shows that cone-restricted regularization paths can only directionally converge to an element of this set.
Theorem 7 (Non-LOMM Regularization Paths Fail) Fix a unit Euclidean norm vector  $\pmb{q} \in \mathbb{R}^d$  such that  $\pmb{q} \notin \mathcal{P}^{mm}$ . Assume that the token scores are distinct (i.e.,  $\gamma_{it} \neq \gamma_{i\tau}$  for  $t \neq \tau$ ) and the key embeddings  $\pmb{k}_{it}$  are in general position. Specifically, we require the following conditions to hold:
- When  $m = d$ , all matrices  $\bar{\pmb{K}} \in \mathbb{R}^{m \times d}$  where each row of  $\bar{\pmb{K}}$  has the form  $\pmb{k}_{it} - \pmb{k}_{i\alpha_i}$  for a unique  $(i, \alpha_i, t \neq \alpha_i)$  tuple, are full-rank.
- When  $m = d + 1$ , the vector of all ones is not in the range space of any such  $\tilde{K}$  matrix.
Fix arbitrary  $\epsilon > 0, R_0 > 0$ . Define the local regularization path of  $\pmb{q}$  as its  $(\epsilon, R_0)$ -conic neighborhood:
$$
\bar {\boldsymbol {p}} (R) = \underset {\boldsymbol {p} \in C _ {\epsilon , R _ {0}} (\boldsymbol {q}), \| \boldsymbol {p} \| \leq R} {\arg \min } \mathcal {L} (\boldsymbol {p}), \quad \text {w h e r e} \quad C _ {\epsilon , R _ {0}} (\boldsymbol {q}) = \operatorname {c o n e} _ {\epsilon} (\boldsymbol {q}) \cap \left\{\boldsymbol {p} \in \mathbb {R} ^ {d} \mid \| \boldsymbol {p} \| \geq R _ {0} \right\}. \tag {66}
$$
Then, either  $\lim_{R\to \infty}\| \bar{\pmb{p}} (R)\| <  \infty$  or  $\lim_{R\to \infty}\bar{\pmb{p}} (R) / \| \bar{\pmb{p}} (R)\| \neq \pmb {q}$ . In both scenarios  $\lim_{R\to \infty}\bar{\pmb{p}} (R) / R\neq \pmb{q}$ .
Proof. We will prove the result by dividing the problem into distinct cases. In each case, we will construct an alternative direction that achieves a strictly better objective than some  $\delta = \delta(\epsilon) > 0$  neighborhood of  $q$ , thereby demonstrating the suboptimality of the  $q$  direction. Let's define the  $\delta$  neighborhood as follows:
$$
\mathcal {N} _ {\delta} = \left\{\boldsymbol {p} \mid \left\| \frac {\boldsymbol {p}}{\| \boldsymbol {p} \|} - \boldsymbol {q} \right\| \leq \delta \quad \text {a n d} \quad \| \boldsymbol {p} \| \geq R _ {0} \right\}. \tag {67}
$$
Now, let's recall a few more definitions based on Definition 3. First, the tokens selected by  $\pmb{q}$  are given by (59). To proceed, let's initially consider the scenario where  $\alpha_{i}$  is unique for all  $i \in [n]$ , meaning that as we let  $c \to \infty$ ,  $c \cdot q$  will choose a single token per input. Later, we will revisit the setting when arg max is not a singleton, and  $\pmb{q}$  is allowed to select multiple tokens.
Additionally, it's important to note that  $\| \bar{p}(R) \|$  is non-decreasing by definition. Suppose it has a finite upper bound  $\| \bar{p}(R) \| \leq M$  for all  $R < \infty$ . In that scenario, we have  $\lim_{R \to \infty} \frac{\bar{p}(R)}{R} = 0 \neq q$ .
- (A)  $q$  selects a single token per input: Given that the indices  $\alpha = (\alpha_{i})_{i=1}^{n}$  defined in (59) are uniquely determined, we can conclude that the  $q$  direction eventually selects tokens  $k_{i\alpha_{i}}$ . Recall the definition of the margin  $\Gamma_{q}$  from (60) and the set of directional neighbors, which is defined as the indices that achieve  $\Gamma_{q}$ , as shown in (61). Let us refer to  $q$  as neighbor-optimal if  $\gamma_{it} < \gamma_{i\alpha_{i}}$  for all  $(i,t) \in \mathcal{M}_{q}$ .
We will consider two cases for this scenario: when  $\mathbf{q}$  is neighbor-optimal and when  $\mathbf{q}$  is not neighbor-optimal.
$\diamond$  (A1)  $q$  is neighbor-optimal. In this case, we will argue that max-margin direction  $\bar{p}^{mm} := p^{mm}(\alpha) / \| p^{mm}(\alpha)\|$  can be used to construct a strictly better objective than  $q$ . Note that  $p^{mm}(\alpha)$  exists because  $q$  is already a viable separating direction for tokens  $\alpha$ . Specifically, consider the direction  $q' = \frac{q + \epsilon \bar{p}^{mm}}{\|q + \epsilon \bar{p}^{mm}\|}$ . Observe that,  $q'$  lies within  $\mathrm{cone}_{2\epsilon}(q)$ ,<sup>5</sup>
$$
\boldsymbol {q} ^ {\top} \boldsymbol {q} ^ {\prime} \geq \frac {1 - \epsilon}{1 + \epsilon} \geq 1 - 2 \epsilon .
$$
We now argue that, there exists  $\delta = \delta_{\epsilon} > 0$  such that for all  $R > R_{\epsilon}$
$$
\min  _ {R _ {\epsilon} \leq r \leq R} \mathcal {L} (r \cdot q ^ {\prime}) <   \min  _ {p \in \mathcal {N} _ {\delta}, R _ {\epsilon} \leq \| p \| \leq R} \mathcal {L} (p). \tag {68}
$$
This requirement holds for general data because it is guaranteed by adding arbitrarily small independent gaussian perturbations to keys  $\pmb{k}_{it}$ .
As a result, let us prove the result for  $\epsilon \gets 2\epsilon$  without losing generality.
To prove this, we study the margin  $\Gamma_{q'}$  induced by  $q'$  and the maximum margin  $\Gamma_{\delta}$  induced within  $p \in \mathcal{N}_{\delta}$ . Concretely, we will show that  $\Gamma_{q'} > \Gamma_{\delta}$  and directly apply the first statement of Lemma 8 to conclude with (68).
Let  $\Gamma = 1 / \| \pmb{p}^{mm}(\alpha)\|$  be the margin induced by  $\bar{\pmb{p}}^{mm}$ . Note that  $\Gamma >\Gamma_q$  by the optimality of  $p^{mm}(\alpha)$  and the fact that  $\pmb {q}\neq \pmb{p}^{mm}(\pmb {\alpha})$ . Consequently, we can lower and upper bound the margins via
$$
\begin{array}{l} \Gamma_ {q ^ {\prime}} = \min  _ {i \in [ n ]} \min  _ {t \neq \alpha_ {i}} (\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}) ^ {\top} \boldsymbol {q} ^ {\prime} \geq \frac {\Gamma_ {q} + \epsilon \Gamma}{1 + \epsilon} \geq \Gamma_ {q} + \frac {\epsilon}{2} (\Gamma - \Gamma_ {q}), \\ \Gamma_ {\delta} = \max  _ {\boldsymbol {p} \in \mathcal {N} _ {\delta}} \min  _ {i \in [ n ]} \min  _ {t \neq \alpha_ {i}} (\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}) ^ {\top} \boldsymbol {p} / \| \boldsymbol {p} \| \\ \leq \max  _ {\| \boldsymbol {r} \| \leq 1} \min  _ {i \in [ n ]} \min  _ {t \neq \alpha_ {i}} (\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}) ^ {\top} (\boldsymbol {q} + \delta \boldsymbol {r}) \\ \leq \Gamma_ {q} + M \delta , \\ \end{array}
$$
where  $M = \max_{i,t,\tau} \| \pmb{k}_{it} - \pmb{k}_{i\tau} \|$ .
Consequently, setting  $\delta = \frac{\epsilon}{4M} (\Gamma -\Gamma_q)$ , we find that
$$
\Gamma_ {q ^ {\prime}} \geq \Gamma_ {\delta} + \frac {\epsilon}{4} (\Gamma - \Gamma_ {q}).
$$
Equipped with this inequality, we apply the first statement of Lemma 8 which concludes that for some  $R_{\epsilon} = R(\frac{\epsilon}{4} (\Gamma - \Gamma_q))$  and all  $R > R_{\epsilon}$ , (68) holds. This in turn implies that, within  $C_{\epsilon}$ , the optimal solution is
- either upper bounded by  $R_{\epsilon}$  in  $\ell_2$  norm (i.e.  $\lim_{R \to \infty} \| \bar{p}(R) \| < \infty$ ) or
- at least  $\delta = \delta(\epsilon) > 0$  away from  $\pmb{q}$  after  $\ell_2$ -normalization i.e.  $\| \frac{\bar{p}(R)}{\|\bar{p}(R)\|} - \pmb{q} \| \geq \delta$ .
In either scenario, we have proven that  $\lim_{R\to \infty}\frac{\bar{p}(R)}{R}\neq q$
$\diamond$  (A2)  $q$  is not neighbor-optimal. In this scenario, we will prove that  $\bar{p}(R)$  is finite to obtain  $\lim_{R \to \infty} \bar{p}(R)/R = 0 \neq q$ . To start, assume that conic neighborhood  $\epsilon$  of  $q$  is small enough so that selected-tokens  $\alpha$  remain unchanged within  $C_{\epsilon}$ . This is without generality because if directional convergence fails in a small neighborhood of  $q$ , it will fail in the larger neighborhood as well. Secondly, if  $\lim_{R \to \infty} \| \bar{p}(R) \| \to \infty$  and  $\bar{p}(R) \in C_{\epsilon}$ , since softmax will eventually perfectly select  $\alpha$  (i.e. assigning probability 1 on token indices  $(i, \alpha_i)$ ), we would have
$$
\lim  _ {R \rightarrow \infty} \mathcal {L} (\bar {\boldsymbol {p}} (R)) = \mathcal {L} _ {\star} = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (\boldsymbol {\gamma} _ {i \alpha_ {i}}).
$$
Note that, this is simply by selection of  $\alpha$  and regardless of  $\bar{p}(R)$  directionally converges to  $q$ . This means that, if there exists a finite  $p \in C_{\epsilon}$  such that  $\mathcal{L}(p) < \mathcal{L}_{\star}$  (i.e. outperforming the training loss of  $\| \bar{p}(R) \| \to \infty$ ), then  $\| \bar{p}(R) \| < \infty$ . This would conclude the proof.
Thus, we will simply find such a  $\pmb{p}$  obeying  $\mathcal{L}(\pmb{p}) < \mathcal{L}_{\star}$ . To this aim, we first prove the following lemma.
Lemma 9 Given a fixed unit Euclidean norm vector  $\pmb{p}$ , if all directional neighbors of  $\pmb{p}$  consistently have higher scores for their associated selected tokens, i.e.,  $\gamma_{i\alpha_i} < \gamma_{i\beta}$  for all  $(i,\alpha_i)$  and directional neighbor  $(i,\beta)$ , then there exists  $\bar{R}$  such that for all  $R > \bar{R}$ ,
$$
\mathcal {L} (R \cdot p) <   \mathcal {L} _ {\star} = \lim  _ {R \rightarrow \infty} \mathcal {L} (R \cdot p). \tag {69}
$$
Proof. Define the maximum score difference  $K_{+} = \sup_{i\in [n],t\neq \alpha_{i}}|\gamma_{i\alpha_{i}} - \gamma_{it}|$ . Also let  $\mathcal{M}_p$  be the set of directional neighbors achieving the minimum margin  $\Gamma_p$ ; see (61). Define  $\Gamma^{it} = k_{i\alpha_i}^\top p - k_{it}^\top p$ . Define  $\delta_p$  to be the margin difference between the directional-neighbors and the second-most minimum-margin neighbors defined as
$$
\delta_ {p} = \min  _ {i \in [ n ], t \neq \alpha_ {i}, (i, t) \notin \mathcal {M} _ {p}} \Gamma_ {p} ^ {i t} - \Gamma_ {p}. \tag {70}
$$
Here, we apply this lemma to compare  $q'$  against all  $p \in \mathcal{N}_{\delta}$ . We can do this uniform comparison because the  $R$  requirement in Lemma 8 only depends on the margin difference and global problem variables and not the particular choice of  $p \in \mathcal{N}_{\delta}$ .
To proceed, setting  $\kappa = \min_{(j,\beta)\in \mathcal{M}_p}\gamma_{j\beta} - \gamma_{j\alpha_j} > 0$  and using (63), we can bound
$$
\begin{array}{l} \mathcal {L} (R p) - \mathcal {L} _ {\star} \leq - \frac {1}{n} \sum_ {(j, \beta) \in \mathcal {M} _ {p}} M _ {j \beta} e ^ {- R \Gamma_ {p} K} + \frac {1}{n} \sum_ {i \in [ n ], t \neq \alpha_ {i}, (i, t) \notin \mathcal {M} _ {p}} M _ {i t} e ^ {- R \Gamma_ {p} ^ {i t}} (\boldsymbol {\gamma} _ {i \alpha_ {i}} - \boldsymbol {\gamma} _ {i t}) \\ \leq - \frac {1}{n} M _ {-} e ^ {- R \Gamma_ {p} K} + M _ {+} K _ {+} T e ^ {- R \left(\Gamma_ {p} + \delta_ {p}\right)}. \\ \end{array}
$$
Consequently, we have found that  $\mathcal{L}(Rp) < \mathcal{L}_{\star}$  as soon as
$$
\frac {1}{n} M _ {-} e ^ {- R \Gamma_ {q}} \kappa \geq M _ {+} K _ {+} T e ^ {- R (\Gamma_ {q} + \delta_ {q})}.
$$
This happens when  $R \gtrsim 1 / \delta_q$  (up to logarithmic terms) establishing the desired statement.