SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Case 1:  $h_R / R$  does not converge. Under this scenario there exists  $\delta, \gamma = \gamma(\delta) > 0$  such that we can find arbitrarily large  $R$  with  $\| h_R / R - \Xi p^{\text{relax}} \| \geq \delta$ . This implies that margin induced by  $h_R / R$  is at most  $\Xi (1 - \gamma)$  over the support vectors  $S$  (from strong convexity of (10)). The reason is that,  $h_R$  satisfies  $h_R^\top (k_{i\alpha_i} - k_{it}) \geq 0$  for all  $t \neq \alpha_i$  by construction as  $h_R \in Q$ . Thus, a constraint over the support vectors have to be violated (when normalized to the same  $\ell_2$  norm as  $\| p^{\text{relax}} \| = 1 / \Xi$ ).
As usual, we will construct a solution strictly superior to  $\pmb{h}_R$  and contradicts with its optimality.
Construction of competitor: Rather than using  $p^{relax}$  direction, we will choose a slightly deviating direction that ensures the selection of the correct tokens over non-supports  $\bar{S}$ . Specifically, consider the solution of (10) where we tighten the non-support constraints by arbitrarily small  $\epsilon > 0$ .
$$
\boldsymbol {p} ^ {\epsilon - r l x} = \underset {p} {\arg \min } \| \boldsymbol {p} \| \quad \text {s u c h t h a t} \quad \boldsymbol {p} ^ {\top} \left(\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}\right) \geq \left\{ \begin{array}{l l} 1 & \text {f o r a l l} \quad t \neq \alpha_ {i}, i \in S \\ \epsilon & \text {f o r a l l} \quad t \neq \alpha_ {i}, i \in \bar {S} \end{array} . \right. \tag {78}
$$
Let  $\pmb{p}^{mm}$  be the solution of (ATT-SVM) with  $\alpha = (\alpha_{i})_{i=1}^{n}$  (which was assumed to be separable). Observe that  $\pmb{p}_{\epsilon}^{mm} = \epsilon \pmb{p}^{mm} + (1 - \epsilon) \pmb{p}^{relax}$  satisfies the constraints of (78). Additionally,  $\pmb{p}_{\epsilon}^{mm}$  would achieve a margin of  $\frac{1}{(1 - \epsilon) / \Xi + \epsilon / \Delta} = \frac{\Delta \Xi}{\Delta + \epsilon(\Xi - \Delta)}$  where  $\Delta = 1 / \| \pmb{p}^{mm} \|$ . Using optimality of  $\pmb{p}^{\epsilon - rlx}$ , this implies that the reduced margin  $\Xi_{\epsilon} = 1 / \| \pmb{p}^{\epsilon - rlx} \|$  (by enforcing  $\epsilon$  over non-support) over the support vectors is a Lipschitz function of  $\epsilon$ . That is  $\Xi_{\epsilon} \geq \Xi - \epsilon M$  for some  $M \geq 0$ . To proceed, choose an  $\epsilon > 0$  such that, it is strictly superior to margin induced by  $h_{R}$ , that is,
$$
\Xi_ {\epsilon} \geq \Xi \left(1 - \frac {\gamma}{2}\right).
$$
To proceed, set  $\tilde{p}^{\epsilon - rlx} = R\Xi_{\epsilon}p^{\epsilon - rlx}$ . Let us recall the following notation from the proof of Theorem 5:  $s_i^p = \mathbb{S}(K_i p)$  and  $q_i^p = 1 - s_{i\alpha_i}$ . Set  $\hat{q}_{\mathrm{max}} = \max_{i\in S} q_i^{h_R}$  to be worst non-optimality of  $h_R$  over support set. Similarly, define  $q_{\mathrm{max}}^{\star} = \max_{i\in S} q_i^{\bar{p}^{\epsilon - rlx}}$  to be the same for  $\tilde{p}^{\epsilon - rlx}$ . Repeating the identical arguments to (71), (72), (73), and using the fact that  $p^{\epsilon - rlx}$  achieves a margin  $\Xi(1 - \frac{\gamma}{2}) \leq \Xi_{\epsilon} \leq \Xi$ , we end up with the lines
$$
\log \left(\hat {q} _ {\max }\right) \geq - (1 - \gamma) \Xi R - \log T, \tag {79a}
$$
$$
- \Xi R - \log T \leq \log \left(q _ {\max } ^ {\star}\right) \leq - \Xi (1 - 0. 5 \gamma) R + \log T. \tag {79b}
$$
In what follows, we will prove that  $\tilde{p}^{\epsilon -rlx}$  achieves a strictly smaller logistic loss contradicting with the optimality of  $p_R$  (whenever  $\| h_R / R - \Xi p^{relax}\| \geq \delta$ ).
Upper bounding logistic loss. Let us now upper bound the logistic loss of  $(\tilde{\nu}^{mm},\tilde{p}^{\epsilon -rlx})$  where  $\tilde{\nu}^{mm} = r\Gamma \tilde{\nu}^{mm}$  with  $\nu^{mm}$  being the solution of (SVM) with  $\boldsymbol {r}_i\gets \boldsymbol{x}_{i\alpha_i}$  and  $\Gamma = 1 / \| \boldsymbol{v}^{mm}\|$ . Set  $\boldsymbol {r}_i = X_i^\top \mathbb{S}(K_i\tilde{\boldsymbol{p}}^{\epsilon -rlx})$ . Set  $\nu = \min_{i\in \bar{S}}Y_{i}\cdot x_{i\alpha_{i}}^{\top}\nu^{mm} - 1$  to be the additional margin buffer that non-support vectors have access to. Also set  $M = \sup_{i\in [n],t,\tau \in [T]}\| x_{it} - x_{i\tau}\|$ . Observe that we can write
$$
\boldsymbol {x} _ {i \alpha_ {i}} - \boldsymbol {r} _ {i} = \sum_ {t \neq \alpha_ {i}} \boldsymbol {s} _ {i t} (\boldsymbol {x} _ {i \alpha_ {i}} - \boldsymbol {x} _ {i t}) \Rightarrow \| \boldsymbol {x} _ {i \alpha_ {i}} - \boldsymbol {r} _ {i} \| \leq q _ {i} M.
$$
Non-supports achieve strong label-margin: Using above and (78) for all  $i \in \bar{S}$  and  $t \neq \alpha_i$ , we have that  $s_{it} \leq e^{-\epsilon \Xi_{\epsilon} R} s_{i\alpha_i} \leq e^{-\epsilon \Xi(1 - \gamma/2)R} s_{i\alpha_i}$ . Consequently, whenever  $R \geq \bar{R}_0 := (\epsilon \Xi(1 - \gamma/2))^{-1} \log(\frac{T M}{\Gamma \nu})$ ,
$$
q _ {i} ^ {\bar {p} ^ {\epsilon - r l x}} \leq \frac {\sum_ {t \neq \alpha_ {i}} \boldsymbol {s} _ {i t}}{\boldsymbol {s} _ {i \alpha_ {i}}} \leq T e ^ {- \epsilon \Xi (1 - \gamma / 2) R} \leq \frac {\Gamma \nu}{M}.
$$
This implies that, on  $i\in \bar{S}$
$$
Y _ {i} \cdot \boldsymbol {r} _ {i} ^ {\top} \boldsymbol {v} ^ {m m} \geq 1 + v + Y _ {i} \cdot \left(\boldsymbol {r} _ {i} - \boldsymbol {x} _ {i \alpha_ {i}}\right) ^ {\top} \boldsymbol {v} ^ {m m} \geq 1 + v - q _ {i} M \| \boldsymbol {v} ^ {m m} \| \geq 1. \tag {80}
$$
In words: Above a fixed  $\bar{R}_0$  that only depends on  $\gamma = \gamma (\delta)$ , features  $\pmb{r}_i$  induced by all non-support indices  $i\in \bar{S}$  achieve margin at least 1. What remains is analyzing the margin shrinkage over the support vectors as in Theorem 5.
Controlling support margin and combining bounds: Over  $S$ , suppose  $\pmb{v}^{mm}$  satisfies the SVM constraints on  $\pmb{r}_i$  with  $Y_i \cdot \pmb{r}_i^\top \pmb{v}^{mm} \geq 1 - \epsilon_i / \Gamma$ . Consequently, setting  $\epsilon_{\max} = \sup_{i \in [n]} \epsilon_i$ ,  $\pmb{v}^{mm}$  achieves a label-margin of  $\Gamma - \epsilon_{\max}$  on the dataset  $(Y_i, r_i)_{i \in [n]}$ . Next, we recall the fact (79b) that worst-case perturbation is  $\epsilon_{\max} \leq M \exp(-\Xi(1 - 0.5\gamma)R + \log T) = MT \exp(-\Xi(1 - 0.5\gamma)R)$ . With this and (80), we upper bound the logistic loss of  $(\tilde{\nu}^{mm}, \tilde{\pmb{p}}^{\epsilon - r|x})$  as follows.
$$
\begin{array}{l} \mathcal {L} \left(\tilde {\boldsymbol {v}} ^ {m m}, \tilde {\boldsymbol {p}} ^ {m m}\right) \leq \max  _ {i \in [ n ]} \log \left(1 + \exp \left(- Y _ {i} \boldsymbol {r} _ {i} ^ {\top} \tilde {\boldsymbol {v}} ^ {m m}\right)\right). \\ \leq \max  _ {i \in [ n ]} \exp (- Y _ {i} \boldsymbol {r} _ {i} ^ {\top} \tilde {\boldsymbol {v}} ^ {m m}) \\ \leq \exp (- r \Gamma + r \epsilon_ {\max }) \\ \leq e ^ {r M T \exp (- \Xi (1 - 0. 5 \gamma) R)} e ^ {- r \Gamma}. \tag {81} \\ \end{array}
$$
Conversely, we obtain a lower bound for  $(\nu_r, h_R)$ . Set  $r_i = X_i^\top \mathbb{S}(K_i h_R)$ . Recall the lower bound (79a) over the support vector set  $S$ . Combining this with our Assumption C over the support vectors
of (SVM) implies that, solving (SVM) on  $(Y_{i},r_{i})_{i\in [n]}$  achieves at most  $\Gamma -\nu e^{-(1 - \gamma)\Xi R} / T$  margin. Consequently, we have
$$
\begin{array}{l} \mathcal {L} \left(\boldsymbol {v} _ {r}, \boldsymbol {h} _ {R}\right) \geq \frac {1}{n} \max  _ {i \in [ n ]} \log \left(1 + \exp \left(- Y _ {i} \boldsymbol {r} _ {i} ^ {\top} \boldsymbol {v} _ {r}\right)\right) \\ \geq \frac {1}{2 n} \max  _ {i \in [ n ]} \exp (- Y _ {i} \boldsymbol {r} _ {i} ^ {\top} \boldsymbol {v} _ {r}) \wedge \log 2 \\ \geq \frac {1}{2 n} \exp (- r (\Gamma - v e ^ {- (1 - \gamma) \Xi R} / T)) \wedge \log 2 \\ \geq \frac {1}{2 n} e ^ {r (\nu / T) \exp (- (1 - \gamma) \Xi R)} e ^ {- r \Gamma} \wedge \log 2. \tag {82} \\ \end{array}
$$
$$
(\nu / T) e ^ {- (1 - \gamma) \Xi R} \geq M T e ^ {- \Xi (1 - 0. 5 \gamma) R} \Longleftrightarrow R \geq R _ {0} := \frac {2}{\gamma \Xi} \log \left(\frac {M T ^ {2}}{\nu}\right).
$$
Thus, we obtain the desired contradiction since  $\tilde{\pmb{p}}^{\epsilon - rlx}$  is a strictly better solution compared to  $\pmb{p}_R = \pmb{h}_R$  (once  $R$  is sufficiently large).
- Case 2:  $\nu_r / r$  does not converge. This is the simpler scenario: There exists  $\delta > 0$  such that we can find arbitrarily large  $r$  obeying  $\| \nu_r / r - \nu^{mm} / \| \nu^{mm} \| \| \geq \delta$ . First, note that, due to the strong convexity of (SVM), for some  $\gamma := \gamma(\delta) > 0$ ,  $\nu_r$  achieves a margin of at most  $(\Gamma - \gamma)r$  on the dataset  $(Y_i, x_{i1})_{i \in [n]}$ . By theorem's condition, we are provided that  $\mathbb{S}(K_i p_R)_{\alpha_i} \to 1$ . This immediately implies that, for any choice of  $\epsilon = \gamma/3 > 0$ , above some sufficiently large  $(r_0, R_0)$ , we have that  $\| x_i^{p_R} - r_i \| \leq \epsilon$ . Following (81), this implies that, choosing  $\tilde{\nu}^{mm} = r\nu^{mm} / \| \nu^{mm} \|$  achieves a logistic loss of at most  $e^{r\gamma/3} e^{-r\Gamma}$ . Again using  $\| x_i^{p_R} - r_i \| \leq \epsilon$ , for sufficiently large  $(r, R)$  we have that
$$
\begin{array}{l} \min  _ {i \in [ n ]} Y _ {i} \boldsymbol {v} _ {r} ^ {\top} \boldsymbol {r} _ {i} \leq \min  _ {i \in [ n ]} Y _ {i} \boldsymbol {v} _ {r} ^ {\top} \boldsymbol {x} _ {i 1} + \sup  _ {i \in [ n ]} | \boldsymbol {v} _ {r} ^ {\top} (\boldsymbol {r} _ {i} - \boldsymbol {x} _ {i 1}) | \\ \leq (\Gamma - \gamma) r + \epsilon r \\ \leq (\Gamma - 2 \gamma / 3) r. \\ \end{array}
$$
This in turn implies that logistic loss is lower bounded by (following (82)),
$$
\mathcal {L} \left(\boldsymbol {v} _ {r}, \boldsymbol {p} _ {R}\right) \geq \frac {1}{2 n} e ^ {2 \gamma r / 3} e ^ {- r \Gamma} \wedge \log 2.
$$
This dominates the above upper bound  $e^{r\gamma /3}e^{-r\Gamma}$  of  $\tilde{\nu}^{mm}$  whenever  $\frac{1}{2n} e^{\gamma r / 3} > 1 \iff r > \frac{3}{\gamma}\log (2n)$ , (that is, when  $r$  is sufficiently large), again concluding the proof.
D Regularization Path of Attention with Nonlinear Head
So far our discussion has focused on the attention model with linear head. However, the conceptual ideas on optimal token selection via margin maximization also extends to a general nonlinear model under mild assumptions. The aim of this section is showcasing this generalization. Specifically, we consider the prediction model  $f(\boldsymbol{X}) = \psi(\boldsymbol{X}^{\top} \mathbb{S}(\boldsymbol{K}\boldsymbol{p}))$  where  $\psi(\cdot): \mathbb{R}^{d} \to \mathbb{R}$  generalizes the linear head  $\nu$  of our attention model. For instance, following exposition in Section 1.1,  $\psi(\cdot)$  can represent a multilayer transformer with  $\boldsymbol{p}$  being a tunable prompt at the input layer. Recall that  $(X_{i}, K_{i}, Y_{i})_{i=1}^{n}$  is the dataset of the input-key-label tuples. We consider the training risk
$$
\mathcal {L} (\boldsymbol {p}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(Y _ {i}, \psi \left(\boldsymbol {X} _ {i} ^ {\top} s _ {i} ^ {p}\right)\right), \quad \text {w h e r e} \quad s _ {i} ^ {p} = \mathbb {S} \left(\boldsymbol {K} _ {i} \boldsymbol {p}\right) \in \mathbb {R} ^ {T}. \tag {83}
$$
The challenge with nonlinear  $\psi(\cdot)$  is that, we lack a clear score function (Def. 1) unlike the previous sections. The assumption below introduces a generic condition that splits the tokens of each  $X_{i}$  into an optimal set  $\mathcal{O}_i$  and non-optimal set  $\bar{\mathcal{O}}_i = [T] - \mathcal{O}_i$ . In words, non-optimal tokens are those that strictly increase the training risk  $\mathcal{L}(\pmb{p})$  if they are not fully suppressed by attention probabilities  $s_i^p$ .
Assumption D (Mixing non-optimal tokens hurt) There exists sets  $(O_i)_i = [T]$  as follows. Let  $q_i^p = \sum_{t\in \bar{O}_i}s_{it}^p$  be the sum of softmax similarities over the non-optimal set for  $p$ . Set  $q_{\mathrm{max}}^p = \max_{i\in [n]}q_i^p$ . For any  $\Delta >0$ , there exists  $\rho < 0$  such that:
$$
\text {F o r a l l} \boldsymbol {p}, \boldsymbol {p} ^ {\prime} \in \mathbb {R} ^ {d}, \text {i f} \log \left(q _ {\max } ^ {p}\right) \leq (1 + \Delta) \log \left(q _ {\max } ^ {p ^ {\prime}}\right) \wedge \rho , \text {t h e n} \mathcal {L} (\boldsymbol {p}) <   \mathcal {L} (\boldsymbol {p} ^ {\prime}).
$$
This assumption is titled mixing hurts because the attention output  $X_{i}^{\top}s_{i}^{p}$  is mixing the tokens of  $X_{i}$  and our condition is that, to achieve optimal risk, this mixture should not contain any non-optimal tokens. In particular, we require that, a model  $p$  that contains exponentially less non-optimality (quantified via  $\log (q_{\mathrm{max}})$ ) compared to  $p'$  is strictly preferable. As we discuss in the supplementary material, Theorem 1 is in fact a concrete instance (with linear head  $\nu$ ) satisfying this condition.
Before stating our generic theorem, we need to introduce the max-margin separator towards which regularization path of attention will converge. This is a slightly general version of Section 2's (ATT-SVM) problem where we allow for a set of optimal tokens  $O_{i}$  for each input.
$$
\boldsymbol {p} ^ {m m} = \arg \min  _ {\boldsymbol {p}} \| \boldsymbol {p} \| \quad \text {s u b j e c t t o} \quad \max  _ {\alpha \in O _ {i}} \min  _ {\beta \in \bar {O} _ {i}} \boldsymbol {p} ^ {\top} (\boldsymbol {k} _ {i \alpha} - \boldsymbol {k} _ {i \beta}) \geq 1, \quad \text {f o r a l l} \quad i \in [ n ]. \quad (\text {A T T - S V M} ^ {*})
$$
Unlike (ATT-SVM), this problem is not necessarily convex when the optimal set  $O_{i}$  is not a singleton. To see this, imagine  $n = d = 1$  and  $T = 3$ : Set the two optimal tokens as  $k_{1} = 1$  and  $k_{2} = -1$  and the non-optimal token as  $k_{3} = 0$ . The solution set of (ATT-SVM') is  $p^{mm} \in \{-1, 1\}$  whereas their convex combination  $p = 0$  violates the constraints. To proceed, our final result establishes the convergence of regularization path to the solution set of (ATT-SVM') under Assumption D.
Theorem 8 Let  $\mathcal{G}^{mm}$  be the set of global minima of (ATT-SVM'). Suppose its margin  $\Xi := 1 / \| \pmb{p}^{mm} \| > 0$  and Assumption  $D$  holds. Let  $\text{dist}(\cdot, \cdot)$  denote the  $\ell_2$ -distance between a vector and a set. Following (83), define  $\bar{\pmb{p}}(R) = \arg \min_{\| \pmb{p} \| \leq R} \mathcal{L}(\pmb{p})$ . We have that  $\lim_{R \to \infty} \text{dist}\left(\frac{\bar{p}(R)}{\Xi R}, \mathcal{G}^{mm}\right) = 0$ .
We note that Theorem 1 is a corollary of this result where  $O_{i}$ 's and  $\mathcal{G}^{mm}$  are singleton. Based on this result, with multiple optimal tokens, Theorem 1 would gracefully generalize to solve (ATT-SVM').
Proof. The key idea is showing that, thanks to the exponential tail of softmax-attention, (harmful) contribution of the non-optimal token with the minimum margin can dominate the contribution of all other tokens as  $R \to \infty$ . This high-level approach is similar to earlier works on implicit bias of gradient descent with logistic loss [31, 22].
Pick  $\pmb{p}^{mm} \in \mathcal{G}^{mm}$  and set  $\pmb{p}_R^\star = R\frac{\pmb{p}^{mm}}{\|\pmb{p}^{mm}\|}$ . This will be the baseline model that  $\pmb{p}_R$  has to compete against. Also let  $\bar{\pmb{p}}_R = \frac{\pmb{p}_R}{\Xi R}$ . Now suppose  $\text{dist}(\bar{\pmb{p}}_R, \mathcal{G}^{mm}) \to 0$  as  $R \to \infty$ . Then, there exists  $\delta > 0$  such that, we can always find arbitrarily large  $R$  obeying  $\text{dist}(\bar{\pmb{p}}_R, \mathcal{G}^{mm}) \geq \delta$ .
Since  $\bar{\pmb{p}}_R$  is  $\delta > 0$  bounded away from  $\mathcal{G}^{mm}$ , and  $\| \bar{\pmb{p}}_R\| = \| \pmb{p}^{mm}\|$ ,  $\bar{\pmb{p}}_R$  strictly violates at least one of the inequality constraints in (ATT-SVM'). Otherwise, we would have  $\bar{\pmb{p}}_R \in \mathcal{G}^{mm}$ . Without losing generality, suppose  $\bar{\pmb{p}}_R$  violates the first margin constraint, that is, for some  $\gamma := \gamma(\delta) > 0$ ,  $\max_{\alpha \in \mathcal{O}_1} \min_{\beta \in \bar{\mathcal{O}}_1} \bar{\pmb{p}}_R^{\top} (\pmb{k}_{1\alpha} - \pmb{k}_{1\beta}) \leq 1 - \gamma$ . Now, we will argue that this will lead to a contradiction as  $R \to \infty$  since we will show that  $\mathcal{L}(\pmb{p}_R^\star) < \mathcal{L}(\pmb{p}_R)$  for sufficiently large  $R$ .
First, let us control  $\mathcal{L}(\pmb{p}_R^\star)$ . We study  $s_i^\star = \mathbb{S}(\pmb{K}_i\pmb{p}_R^\star)$  and let  $\alpha_{i}\in \mathcal{O}_{i}$  be the index  $\alpha$  in (ATT-SVM') for which  $\mathrm{margin}_i = \max_{\alpha \in O_i}\min_{\beta \in \bar{\mathcal{O}}_i}(\pmb{k}_{i\alpha} - \pmb{k}_{i\beta})^\top \pmb{p}^{mm}\geq 1$  is attained. Then, we bound the non-optimality amount  $q_{i}^{\star}$  of  $\pmb{p}_R^\star$  as
$$
q _ {i} ^ {\star} = \frac {\sum_ {t \in \bar {\mathcal {O}} _ {i}} \exp \left(\boldsymbol {k} _ {i t} ^ {\top} \boldsymbol {p} _ {R} ^ {\star}\right)}{\sum_ {t \in [ T ]} \exp \left(\boldsymbol {k} _ {i t} ^ {\top} \boldsymbol {p} _ {R} ^ {\star}\right)} \leq \frac {\sum_ {t \in \bar {\mathcal {O}} _ {i}} \exp \left(\boldsymbol {k} _ {i t} ^ {\top} \boldsymbol {p} _ {R} ^ {\star}\right)}{\exp \left(\boldsymbol {k} _ {i \alpha_ {i}} ^ {\top} \boldsymbol {p} _ {R} ^ {\star}\right)} \leq T \exp (- \Xi R).
$$
Thus,  $q_{\mathrm{max}}^{\star} = \max_{i\in [n]}q_{i}^{\star}\leq T\exp (-\Xi R)$ . Secondly, we wish to control  $\mathcal{L}(\pmb {p}_R)$  by lower bounding the non-optimality in  $\pmb{p}_R$ . Focusing on the first margin constraint, let  $\alpha \in O_1$  be the index in (ATT-SVM') for which  $\mathrm{min}_{1}\leq 1 - \gamma$  is attained. Denoting the amount of non-optimality of the first input as  $\hat{q}_1$ , we find
$$
\hat {q} _ {1} = \frac {\sum_ {t \in \bar {O} _ {1}} \exp \left(\boldsymbol {k} _ {1 t} ^ {\top} \boldsymbol {p} _ {R}\right)}{\sum_ {t \in [ T ]} \exp \left(\boldsymbol {k} _ {1 t} ^ {\top} \boldsymbol {p} _ {R}\right)} \geq \frac {1}{T} \frac {\sum_ {t \in \bar {O} _ {1}} \exp \left(\boldsymbol {k} _ {1 t} ^ {\top} \boldsymbol {p} _ {R}\right)}{\exp \left(\boldsymbol {k} _ {1 \alpha} ^ {\top} \boldsymbol {p} _ {R}\right)} \geq T ^ {- 1} \exp (- \Xi R (1 - \gamma)).
$$
<sup>7</sup>Here, we assumed margin is non-negative i.e.  $\pmb{k}_{1\alpha}^{\top}\pmb{p}_R \geq \sup_{t \in \bar{O}_1} \pmb{k}_{1t}^{\top}\pmb{p}_R$ . Otherwise,  $\sup_{t \in [T]} \pmb{k}_{1t}^{\top}\pmb{p}_R$  is attained in  $\bar{O}_1$  which implies  $\hat{q}_1 \geq T^{-1}$ . Thus, we can still use the identical inequality (84) with the choice  $\gamma = 1$ .
We similarly have  $q_{\mathrm{max}}^{\star} \geq T^{-1}\exp(-\Xi R)$ . In conclusion, for  $\pmb{p}_R, \pmb{p}_R^\star$ , denoting maximum nonoptimality by  $\hat{q}_{\mathrm{max}} \geq \hat{q}_1$  and  $q_{\mathrm{max}}^{\star}$ , we respectively obtained
$$
\begin{array}{l} \log (\hat {q} _ {\max }) \geq - (1 - \gamma) (\Xi R) - \log T, \\ - (\Xi R) - \log T \leq \log \left(q _ {\max } ^ {\star}\right) \leq - (\Xi R) + \log T. \tag {84} \\ \end{array}
$$
The above inequalities satisfy Assumption D as follows where  $\pmb{p} \gets \pmb{p}_R^\star$  and  $\pmb{p}' \gets \pmb{p}_R$ : Set  $R_0 = 3\gamma^{-1}\Xi^{-1}\log T$  so that  $\log T = \frac{\gamma\Xi R_0}{3}$ . Secondly, set  $\rho_0 = -\Xi R_0 - \log T$ . This way,  $\rho_0 \geq \log (q_{\max}^\star)$  implies  $R \geq R_0$  and  $\log T \leq \frac{\gamma\Xi R}{3}$ . Using the latter inequality, we bound the  $\log T$  terms to obtain
-  $\log (\hat{q}_{\mathrm{max}}) \geq -(1 - 2\gamma / 3)(\Xi R)$ , and
-  $\log (q_{\max}^{\star}) \leq -(1 - \gamma / 3)(\Xi R)$ .
To proceed, we pick  $1 + \Delta = \frac{1 - \gamma / 3}{1 - 2\gamma / 3}$  implying  $\Delta := \frac{\gamma}{3 - 2\gamma}$ . Finally, for this  $\Delta$ , there exists  $\rho(\Delta)$  which we need to ensure  $\log (\hat{q}_{\max}) \leq \rho(\Delta)$ . This can be guaranteed by picking sufficiently large  $R$  that ensures  $\log (q_{\max}^{\star}) \leq -(1 - \gamma / 3)(\Xi R) \leq \rho(\Delta)$  to satisfy all conditions of Assumption D. Since such large  $R$  exists by initial assumption  $\text{dist}(\bar{p}_R, \mathcal{G}^{mm}) \to 0$ , Assumption D in turn implies that  $\mathcal{L}(p_R^{\star}) < \mathcal{L}(p_R)$  contradicting with the optimality of  $p_R$  in (83).
D.2 Application to Linearly-mixed Labels
The following example shows that if non-optimal tokens result in reduced score (in terms of the alignment of prediction and label), Assumption D holds. The high-level idea behind this lemma is that, if the optimal risk is achieved by setting  $q_{\max}^p = 0$ , then, Assumption D will hold.
Lemma 14 (Linear label mixing) Recall  $q_{i}^{p} = \sum_{t\in \bar{O}_{t}}s_{it}^{p}$  from Assumption D. Suppose  $Y_{i}\in \{-1,1\}$  and
$$
Y _ {i} \cdot \psi \left(\boldsymbol {X} _ {i} ^ {\top} \boldsymbol {s} _ {i} ^ {\boldsymbol {p}}\right) = \nu_ {i} \left(1 - q _ {i} ^ {\boldsymbol {p}}\right) + Z _ {i},
$$
for some  $(\nu_{i})_{i = 1}^{n} > 0$ . Here  $Z_{i} = Z_{i}(\pmb{p})$  is the contribution of non-optimal tokens to prediction. For some  $C, \epsilon > 0$  and for all  $\pmb{p} \in \mathbb{R}^d$ , assume
$$
- C q _ {\max } ^ {p} \leq Z _ {i} \leq (1 - \epsilon) v _ {i} q _ {i} ^ {p}. \tag {85}
$$
Then, Assumption  $D$  holds for  $\mathcal{L}(\pmb{p}) = \frac{1}{n}\sum_{i=1}^{n}\ell(Y_i\cdot\psi(\pmb{X}_i^\top\pmb{s}_i^{\pmb{p}}))$  when  $\ell(\cdot)$  is a strictly decreasing loss function with continuous derivative.
Proof. Recall the assumption  $Y_{i} \cdot \psi (X_{i}^{\top}s_{i}^{p}) = \nu_{i}(1 - q_{i}^{p}) + Z_{i}$  with  $Z_{i}$  obeying (85). Let us also write the loss function
$$
\mathcal {L} (\boldsymbol {p}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(v _ {i} \left(1 - s _ {i} ^ {\boldsymbol {p}}\right) + Z _ {i}\right).
$$
Define  $q_{\mathrm{max}}^p = \sup_{i\leq [n]}q_i^p$ . Let  $M$  be the maximum absolute value of score over tokens. Let
$$
B = \max  _ {| x | \leq M} - \ell^ {\prime} (x) \geq A = \min  _ {| x | \leq M} - \ell^ {\prime} (x) > 0.
$$
Through Taylor's Theorem (integral remainder), we have that
$$
B (q _ {i} ^ {p} \nu_ {i} - Z _ {i}) \geq \ell (\nu_ {i} (1 - q _ {i} ^ {p}) + Z _ {i}) - \ell (\nu_ {i}) \geq A (q _ {i} ^ {p} \nu_ {i} - Z _ {i}) \geq \epsilon A \nu_ {i} q _ {i} ^ {p}.
$$
Set  $\mathcal{L}_{\star} = \frac{1}{n}\sum_{i=1}^{n}\ell(\nu_i)$ . Set  $C_+ = B(C + \max_{i\in [n]}\nu_i)$  and  $C_- = n^{-1}A\epsilon \min_{i\in [n]}\nu_i$ . This also implies
$$
\begin{array}{l} C _ {+} q _ {\max } ^ {p} \geq \frac {1}{n} \sum_ {i \in [ n ]} B \left(q _ {i} ^ {p} v _ {i} - Z _ {i}\right) \geq \mathcal {L} (\boldsymbol {p}) - \mathcal {L} _ {\star} \geq \frac {1}{n} \sum_ {i \in [ n ]} A \left(q _ {i} ^ {p} v _ {i} - Z _ {i}\right) \\ \geq \frac {1}{n} \sum_ {i \in [ n ]} \epsilon A v _ {i} q _ {i} ^ {p} \geq C _ {-} q _ {\max } ^ {p}. \\ \end{array}
$$
Thus, to prove  $\mathcal{L}(\pmb{p}^{\prime}) > \mathcal{L}(\pmb{p})$ , we simply need to establish the stronger statement  $C_{-}q_{\mathrm{max}}^{p^{\prime}} > C_{+}q_{\mathrm{max}}^{p}$ .
Going back to the condition of Assumption D, any  $\log (q_{\mathrm{max}}^p) \leq (1 + \Delta)\log (q_{\mathrm{max}}^{p'})$  obeys  $q_{\mathrm{max}}^p \leq (q_{\mathrm{max}}^{p'})^{1 + \Delta}$  i.e.  $q_{\mathrm{max}}^{p'} \geq (q_{\mathrm{max}}^p)^{(1 + \Delta)^{-1}}$ . Following above, we wish to ensure  $q_{\mathrm{max}}^{p'} > \Theta q_{\mathrm{max}}^p$  for such  $(p, p')$  pairs where  $\Theta = C_+ / C_- > 1$ . This is guaranteed by
$$
(q _ {\max } ^ {p}) ^ {(1 + \Delta) ^ {- 1} - 1} > \Theta \iff \frac {\Delta}{1 + \Delta} \log (q _ {\max } ^ {p}) <   - \log (\Theta).
$$
The above is satisfied by choosing a  $\rho (\Delta)\coloneqq -2(1 + \Delta^{-1})\log (\Theta)$  in Assumption D. Thus, all  $p,p^{\prime}$  with  $\log (q_{\mathrm{max}}^p)\leq \rho = \rho (\Delta)$  satisfies the condition of Assumption D finishing the proof.
In this section, we provide implementation details and additional experiments.