SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left[ \left(z ^ {T} F (z) e _ {k} - \sigma^ {2} \sum_ {i = 1} ^ {n} \frac {\partial e _ {i} ^ {T} F (z) e _ {k}}{\partial z _ {i}} - z ^ {T} F (\tilde {z}) e _ {k}\right) ^ {2} \right] \\ = \mathbb {E} \left[ \left(z _ {0} ^ {T} f ^ {(k)} (z _ {0}) - \sigma^ {2} \operatorname {d i v} f ^ {(k)} (z _ {0})\right) ^ {2} \right] \\ = \sigma^ {2} \mathbb {E} \| f ^ {(k)} (z _ {0}) \| ^ {2} + \sigma^ {4} \mathbb {E} \operatorname {T r} \left[ \left(\operatorname {J a c} f ^ {(k)} (z _ {0})\right) ^ {2} \right] \\ = \sigma^ {2} \mathbb {E} \| [ F (z) - F (\tilde {z}) ] e _ {k} \| ^ {2} + \sigma^ {4} \mathbb {E} \operatorname {T r} \left[ \left( \begin{array}{c c} \operatorname {J a c} [ F (z) e _ {k} ] & - \operatorname {J a c} [ F (\tilde {z}) e _ {k} ] \\ 0 _ {n \times n} & 0 _ {n \times n} \end{array} \right) ^ {2} \right] \\ = 2 \sigma^ {2} \mathbb {E} \| [ F (z) - \mathbb {E} F (z) ] e _ {k} \| ^ {2} + \sigma^ {4} \mathbb {E} \operatorname {T r} ((\operatorname {J a c} [ F (z) e _ {k} ]) ^ {2}) \\ \leq 3 \sigma^ {4} \mathbb {E} \| \operatorname {J a c} [ F (z) e _ {k} ] \| _ {F} ^ {2}, \\ \end{array}
$$
where the last inequality uses the Gaussian Poincaré inequality, and the Cauchy-Schwarz inequality  $\mathrm{Tr}(A^2) \leq \| A\|_F^2$ . Summing over  $k \in [K]$  gives the desired inequality.
Proof of Lemma S5.4. We first prove  $G \mapsto G^T G$  is Lipschitz by noting
$$
\begin{array}{l} \left\| G ^ {T} G - \tilde {G} ^ {T} \tilde {G} \right\| _ {o p} \\ = \| (G - \tilde {G}) ^ {T} G + \tilde {G} ^ {T} (G - \tilde {G}) \| _ {o p} \\ \leq \| G - \tilde {G} \| _ {o p} (\| G \| _ {o p} + \| \tilde {G} \| _ {o p}) \\ \leq 2 \sqrt {m ^ {*} n} \| G - \tilde {G} \| _ {\mathcal {O P}}. \\ \end{array}
$$
Then we show  $G^T G \mapsto (G^T G)^{-1}$  is Lipschitz. Let  $A = G^T G$  and  $\tilde{A} = \tilde{G}^T \tilde{G}$ , we have
$$
\begin{array}{l} \left\| A ^ {- 1} - \tilde {A} ^ {- 1} \right\| _ {o p} \\ = \left\| A ^ {- 1} (\tilde {A} - A) \tilde {A} ^ {- 1} \right\| _ {o p} \\ \leq \| A - \tilde {A} \| _ {o p} \| A ^ {- 1} \| _ {o p} \| \tilde {A} ^ {- 1} \| _ {o p} \\ \leq \left(m _ {*} n\right) ^ {- 2} \| A - \tilde {A} \| _ {o p}. \\ \end{array}
$$
We next prove  $(G^{T}G)^{-1}\mapsto (G^{T}G)^{-1 / 2}$  is Lipschitz. Let  $S = (G^T G)^{-1}$ ,  $S^{\prime} = (\tilde{G}^{T}\tilde{G})^{-1}$ , and if  $u$  with  $\| u\| = 1$  is the eigenvector of  $\sqrt{S} -\sqrt{\tilde{S}}$  with eigenvalue  $d$ , then
$$
\begin{array}{l} u ^ {T} (S - \tilde {S}) u = u ^ {T} (\sqrt {S} - \sqrt {\tilde {S}}) \sqrt {S} u + u ^ {T} \sqrt {\tilde {S}} (\sqrt {S} - \sqrt {\tilde {S}}) u \\ = d u ^ {T} \sqrt {S} u + d u ^ {T} \sqrt {\tilde {S}} u \\ = d u ^ {T} (\sqrt {S} + \sqrt {\tilde {S}}) u. \\ \end{array}
$$
As  $d$  can be chosen as  $\pm \| \sqrt{S} -\sqrt{\tilde{S}}\|_{op}$  (this argument is a special case of the Hemmen-Ando inequality [van Hemmen and Ando, 1980]), this implies
$$
\| \sqrt {S} - \sqrt {\tilde {S}} \| _ {o p} = \frac {| u ^ {T} (S - \tilde {S}) u |}{u ^ {T} (\sqrt {S} + \sqrt {\tilde {S}}) u} \leq \frac {\| S - \tilde {S} \| _ {o p}}{\lambda_ {\min} (\sqrt {S} + \sqrt {\tilde {S}})} \leq \frac {\| S - \tilde {S} \| _ {o p}}{2 / \sqrt {m ^ {*} n}}.
$$
Combining the above Lipschitz results, we have
$$
\| (G ^ {T} G) ^ {- 1 / 2} - (\tilde {G} ^ {T} \tilde {G}) ^ {- 1 / 2} \| _ {o p} \leq (m ^ {*} n) ^ {1 / 2} (m _ {*} n) ^ {- 2} (m ^ {*} n) ^ {1 / 2} \| G - \tilde {G} \| _ {o p} = \frac {m ^ {*}}{m _ {*} ^ {2}} n ^ {- 1} \| G - \tilde {G} \| _ {o p}.
$$
It immediately follows that
$$
\| (G ^ {T} G) ^ {- 1 / 2} - (\tilde {G} ^ {T} \tilde {G}) ^ {- 1 / 2} \| _ {F} \leq \sqrt {K} \frac {m ^ {*}}{m _ {*} ^ {2}} n ^ {- 1} \| G - \tilde {G} \| _ {F}.
$$
That is, the mapping  $G \mapsto (G^T G)^{-1/2}$  is  $L_1 n^{-1}$ -Lipschitz, where  $L_1 = \sqrt{K} m_*^{-2} m^*$ .
For the second statement, the result follows by
$$
\begin{array}{l} \left\| G \left(G ^ {T} G\right) ^ {- 1 / 2} - \tilde {G} \left(\tilde {G} ^ {T} \tilde {G}\right) ^ {- 1 / 2} \right\| _ {o p} \\ \leq \| G - \tilde {G} \| _ {o p} \| (G ^ {T} G) ^ {- 1 / 2} \| _ {o p} + \| \tilde {G} \| _ {o p} \| (G ^ {T} G) ^ {- 1 / 2} - (\tilde {G} ^ {T} \tilde {G}) ^ {- 1 / 2} \| _ {o p} \\ \leq \left\| G - \tilde {G} \right\| _ {o p} \left(m _ {*} n\right) ^ {- 1 / 2} + \left(m ^ {*} n\right) ^ {1 / 2} L _ {1} n ^ {- 1} \left\| G - \tilde {G} \right\| _ {o p} \\ \end{array}
$$
$$
\begin{array}{l} \left\| G \left(G ^ {T} G\right) ^ {- 1 / 2} - \tilde {G} \left(\tilde {G} ^ {T} \tilde {G}\right) ^ {- 1 / 2} \right\| _ {F} \\ \leq \sqrt {K} \left(m _ {*} ^ {- 1 / 2} + \left(m ^ {*}\right) ^ {1 / 2} L _ {1}\right) n ^ {- 1 / 2} \| G - \tilde {G} \| _ {F}, \\ \end{array}
$$
where  $L_{2} = \sqrt{K}\big(m_{*}^{-1 / 2} + (m^{*})^{1 / 2}L_{1}\big)$
Proof of Lemma S5.5. Recall the KKT conditions  $\sum_{l=1}^{n} z_l g_l^T = \mathbf{0}_{p \times K}$ . We look for the derivative with respect to  $z_{ij}$ . Denoting derivatives with a dot, we find by the chain rule and product rule
$$
\begin{array}{l} \dot {z} _ {l} = \frac {\partial z _ {l}}{\partial z _ {i j}} = I (l = i) e _ {j}, \\ \dot {g} _ {l} = \frac {\partial g _ {l}}{\partial z _ {i j}} = \frac {\partial g _ {l}}{\partial \hat {A} ^ {\top} z _ {l}} \frac {\partial \hat {A} ^ {\top} z _ {l}}{\partial z _ {i j}} = H _ {l} [ \dot {A} ^ {\top} z _ {l} + I (l = i) \hat {A} ^ {\top} e _ {j} ]. \\ \end{array}
$$
Thus, differentiating the KKT conditions w.r.t.  $z_{ij}$  by the product rule gives
$$
\sum_ {l = 1} ^ {n} \left[ I (l = i) e _ {j} g _ {l} ^ {T} + x _ {l} \left(\dot {A} ^ {\top} z _ {l} + I (l = i) \hat {A} ^ {\top} e _ {j}\right) ^ {T} H _ {l} \right] = 0.
$$
$$
e _ {j} g _ {i} ^ {T} + \sum_ {l = 1} ^ {n} z _ {l} z _ {l} ^ {T} \dot {A} H _ {l} + z _ {i} e _ {j} ^ {T} \hat {A} H _ {i} = 0.
$$
We then move the term involving  $\dot{A}$  to one side, and vectorize both sides,
$$
g _ {i} \otimes e _ {j} + \left(H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}\right) = - \sum_ {l = 1} ^ {n} \left(H _ {l} \otimes z _ {l} z _ {l} ^ {T}\right) \operatorname {v e c} (\dot {A}).
$$
With  $M = \left[\sum_{l=1}^{n}\left(H_{l} \otimes z_{l}z_{l}^{T}\right)\right]^{-1}$ , we obtain
$$
\operatorname {v e c} (\dot {A}) = - M \left[ g _ {i} \otimes e _ {j} + \left(H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}\right) \right].
$$
Hence, using  $\operatorname{vec}(H_l\dot{A}^\top z_l) = \operatorname{vec}(z_l^T\dot{A} H_l) = (H_l\otimes z_l^T)\operatorname{vec}(\dot{A})$  gives
$$
\begin{array}{l} \dot {g} _ {l} = \left(H _ {l} \otimes z _ {l} ^ {T}\right) \operatorname {v e c} (\dot {A}) + I (l = i) H _ {l} \hat {A} ^ {T} e _ {j} \\ = - \left(H _ {l} \otimes z _ {l} ^ {T}\right) M \left[ g _ {i} \otimes e _ {j} + \left(H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}\right) \right] + I (l = i) H _ {l} \hat {A} ^ {T} e _ {j}. \\ \end{array}
$$
$$
\begin{array}{l} \dot {g} _ {i} = - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left[ g _ {i} \otimes e _ {j} + \left(H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}\right) \right] + H _ {i} \hat {A} ^ {T} e _ {j} \\ = - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left(H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}\right) + H _ {i} \hat {A} ^ {T} e _ {j} - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left(g _ {i} \otimes e _ {j}\right) \\ = \left[ H _ {i} - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left(H _ {i} \otimes z _ {i}\right) \right] \hat {A} ^ {T} e _ {j} - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left(g _ {i} \otimes e _ {j}\right) \\ = V _ {i} \hat {A} ^ {T} e _ {j} - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left(g _ {i} \otimes e _ {j}\right), \\ \end{array}
$$
where  $V_{i} = [H_{i} - (H_{i}\otimes z_{i}^{T})M(H_{i}\otimes z_{i})]$
Proof of Corollary S5.6. For each  $i \in [n], j \in [p]$ , we have by the product rule
$$
\begin{array}{l} \frac {\partial e _ {i} ^ {T} G (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i j}} \\ = \frac {\partial g _ {i} ^ {T}}{\partial z _ {i j}} (G ^ {T} G) ^ {- 1 / 2} + e _ {i} ^ {T} G \frac {\partial (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i j}} \\ = \left[ V _ {i} \hat {A} ^ {T} e _ {j} - \left(H _ {i} \otimes z _ {i} ^ {T}\right) M \left(g _ {i} \otimes e _ {j}\right) \right] ^ {T} \left(G ^ {T} G\right) ^ {- 1 / 2} + e _ {i} ^ {T} G \frac {\partial \left(G ^ {T} G\right) ^ {- 1 / 2}}{\partial z _ {i j}} \\ = e _ {j} ^ {T} \hat {A} V _ {i} (G ^ {T} G) ^ {- 1 / 2} + \left[ - (g _ {i} ^ {T} \otimes e _ {j} ^ {T}) M (H _ {i} \otimes z _ {i}) (G ^ {T} G) ^ {- 1 / 2} + e _ {i} ^ {T} G \frac {\partial (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i j}} \right]. \\ \end{array}
$$
With  $V = \sum_{i=1}^{n} V_i$ , we further have
S6 Proof of Theorem S3.1
Recall that Theorem S5.1 holds for general loss function  $L_{i}:\mathbb{R}^{K}\to \mathbb{R}$  provided that conditions (1) and (2) in Theorem S5.1 hold. In this section, we consider the multinomial logistic loss function  $L_{i}$  defined in Section S3.1. To be specific,
$$
L _ {i} (u) = - \sum_ {k = 1} ^ {K + 1} \mathrm {y} _ {i k} e _ {k} ^ {T} Q u + \log \sum_ {k ^ {\prime} = 1} ^ {K + 1} \exp \left(e _ {k ^ {\prime}} ^ {T} Q u\right), \quad \forall u \in \mathbb {R} ^ {K}. \tag {S6.1}
$$
In order to apply Theorem S5.1, we need to verify that, when  $L_{i}$  in (S6.1) is used, the two conditions (1) and (2) in Theorem S5.1 hold. To this end, we present a few lemmas in the following two subsections, which will be useful for asserting the conditions (1) and (2) when we apply Theorem S5.1 to prove Theorem S3.1.
S6.1 Control of the singular values of the gradients and Hessians
Before stating the lemmas that assert the conditions in Theorem S5.1, define
$$
U = \left\{\left(Y, X\right) \in \mathbb {R} ^ {n \times (K + 1)} \times \mathbb {R} ^ {n \times p}: \hat {\mathsf {B}} \text {e x i s t s}, \| X \hat {\mathsf {B}} \left(I _ {K + 1} - \frac {\mathbf {1 1} ^ {T}}{K + 1}\right) \| _ {F} ^ {2} <   n \tau \right\},
$$
$$
U _ {y} = \left\{Y \in \mathbb {R} ^ {n \times (K + 1)}: \sum_ {i = 1} ^ {n} I \left(\mathrm {y} _ {i k} = 1\right) \geq \gamma n \text {f o r a l l} k \in [ K + 1 ] \right\}.
$$
Lemma S6.1 (deterministic result on gradient). Let  $L_{i}$  be defined as in (S6.1). Assume that either Assumption 2.2 or Assumption S1.1 holds. If  $Y \in U_{y}$ , for any  $M \in \mathbb{R}^{n \times K}$  such that  $\| MQ^T\|_F^2 \leq n\tau$ , we have
$$
m _ {*} I _ {K} \preceq n ^ {- 1} \sum_ {i = 1} ^ {n} \nabla L _ {i} (M ^ {T} e _ {i}) \nabla L _ {i} (M ^ {T} e _ {i}) ^ {T} \preceq K I _ {K},
$$
where  $m_*$  is a positive constant depending on  $(K, \gamma, \tau)$  only.
Proof of Lemma S6.1. Without loss of generality, let's assume that  $\gamma n$  is an integer. Otherwise, we can replace it with the greatest integer less than or equal to  $\gamma n$ , denoted as  $\lfloor \gamma n \rfloor$ .
If  $Y \in U_y$ , there exists at least  $\gamma n$  many disjoint index sets  $\{S_1, \ldots, S_{\gamma n}\}$  such that the following hold for each  $l \in [\gamma n]$ ,
$$
(i) S _ {l} \subset [ n ]; (i i) | S _ {l} | = K + 1; (i i i) \sum_ {i \in S _ {l}} y _ {i k} = 1, \quad \forall k \in [ K + 1 ].
$$
Since  $S_{l}$  are disjoint and  $\cup_{l = 1}^{\gamma^{n}}S_{l}\subset [n]$ , we have
$$
\sum_ {l = 1} ^ {\gamma n} \sum_ {i \in S _ {l}} \| Q M ^ {T} e _ {i} \| ^ {2} \leq \sum_ {i = 1} ^ {n} \| Q M ^ {T} e _ {i} \| ^ {2} = \| Q M ^ {T} \| _ {F} ^ {2} <   n \tau .
$$
It follows that at most  $\alpha n$  many of  $l \in \{1,2,\dots,\gamma n\}$  s.t.  $\sum_{i \in S_l} \| QM^T e_i \|^2 > \tau / \alpha$ , otherwise the previous display can not hold. In other words, there exists a subset  $L^* \subset \{1,2,\dots,\gamma n\}$  with  $|L^*| \geq (\gamma - \alpha)n$  s.t.  $\sum_{i \in S_l} \| QM^T e_i \|^2 \leq \tau / \alpha$  for all  $l \in L^*$ . Define the index set  $I = \cup_{l \in L^*} S_l$ , then  $|I| \geq (K + 1)n(\gamma - \alpha)$ , and  $\| QM^T e_i \|_\infty \leq \sqrt{\tau / \alpha}$  for all  $i \in I$ . Let us take  $\alpha = \gamma / 2$ , then  $|L^*| \geq \frac{\gamma}{2}n$  and  $|I| \geq \gamma (K + 1)n / 2$ . Recall that  $L_i(u) = L_i(Qu)$ , we have  $\nabla L_i(u) = Q^T\nabla L_i(Qu)$ . Thus,
$$
\nabla L _ {i} \left(M ^ {T} e _ {i}\right) = Q ^ {T} \nabla \mathsf {L} _ {i} \left(Q M ^ {T} e _ {i}\right) = Q ^ {T} \left(- \mathsf {y} _ {i} + \mathsf {p} _ {i}\right),
$$
where  $\mathfrak{p}_i\in \mathbb{R}^{K + 1}$  and its  $k$  -th entry satisfying
$$
\mathbf {p} _ {i k} = \frac {\exp \left(e _ {k} ^ {T} Q M ^ {T} e _ {i}\right)}{\sum_ {k ^ {\prime} = 1} ^ {K + 1} \exp \left(e _ {k ^ {\prime}} ^ {T} Q M ^ {T} e _ {i}\right)} \in [ c, 1 - c ], \tag {S6.2}
$$
for some constant  $c \in (0,1)$  depending on  $(\tau, \alpha, K)$  only. Therefore,
$$
\begin{array}{l} n ^ {- 1} \sum_ {i = 1} ^ {n} \nabla L _ {i} (M ^ {T} e _ {i}) \nabla L _ {i} (M ^ {T} e _ {i}) ^ {T} \\ = n ^ {- 1} Q ^ {T} \sum_ {i = 1} ^ {n} \left(\mathcal {Y} _ {i} - \mathcal {P} _ {i}\right) \left(\mathcal {Y} _ {i} - \mathcal {P} _ {i}\right) ^ {T} Q \\ \succeq n ^ {- 1} Q ^ {T} \sum_ {l = 1} ^ {\gamma^ {n}} \sum_ {i \in S _ {l}} \left(\mathcal {Y} _ {i} - \mathfrak {p} _ {i}\right) \left(\mathcal {Y} _ {i} - \mathfrak {p} _ {i}\right) ^ {T} Q \\ \succeq n ^ {- 1} Q ^ {T} \sum_ {l \in L ^ {*}} \sum_ {i \in S _ {l}} \left(\mathbf {y} _ {i} - \mathbf {p} _ {i}\right) \left(\mathbf {y} _ {i} - \mathbf {p} _ {i}\right) ^ {T} Q \\ := n ^ {- 1} Q ^ {T} \sum_ {l \in L ^ {*}} A _ {l} ^ {T} A _ {l} Q, \\ \end{array}
$$
where  $A_{l}\in \mathbb{R}^{(K + 1)\times (K + 1)}$  has  $K + 1$  rows  $\{\mathsf{y}_i - \mathsf{p}_i:i\in S_l\}$ . We further note that  $A_{l}$  is of the form  $(I_{K + 1} - \mathsf{P}_l)$  up to a rearrangement of the columns, where  $\mathsf{P}_l\in \mathbb{R}^{(K + 1)\times (K + 1)}$  is a stochastic matrix with entries of the form
$$
\frac {\exp (e _ {k} ^ {T} Q M ^ {T} e _ {i})}{\sum_ {k ^ {\prime} = 1} ^ {K + 1} \exp (e _ {k ^ {\prime}} ^ {T} Q M ^ {T} e _ {i})}, \qquad i \in S _ {l}, \quad k \in [ K + 1 ].
$$
By (S6.2), for each  $l \in L^*$ , the stochastic matrix  $\mathsf{P}_l$  is irreducible and aperiodic and  $\ker(I_{K+1} - \mathsf{P}_l)$  is the span of the all-ones vector  $\mathbf{1}$ . Therefore,
$$
(I _ {K + 1} - \mathsf {P} _ {l}) Q Q ^ {T} = (I _ {K + 1} - \mathsf {P} _ {l}) (I _ {K + 1} - \frac {\mathbf {1 1} ^ {T}}{K + 1}) = (I _ {K + 1} - \mathsf {P} _ {l}).
$$
$$
K = \operatorname {r a n k} (I _ {K + 1} - \mathsf {P} _ {l}) = \operatorname {r a n k} ((I _ {K + 1} - \mathsf {P} _ {l}) Q Q ^ {T}) \leq \operatorname {r a n k} ((I _ {K + 1} - \mathsf {P} _ {l}) Q) \leq K.
$$
We conclude that the rank of  $(I_{K + 1} - \mathsf{P}_l)Q$  is  $K$ .
If  $\mathcal{P}$  denotes the set of matrices  $\{\mathsf{P} \in \mathbb{R}^{(K + 1) \times (K + 1)} : \text{stochastic with entries in } [c, 1 - c]\}$ , and  $S^{K - 1} = \{a \in \mathbb{R}^K : \| a \| = 1\}$ . By compactness of  $\mathcal{P}$  and  $S^{K - 1}$ , we obtain
$$
\begin{array}{l} \frac {1}{n} \lambda_ {\min } \left(\sum_ {i = 1} ^ {n} \nabla L _ {i} \left(M ^ {T} e _ {i}\right) \nabla L _ {i} \left(M ^ {T} e _ {i}\right) ^ {T}\right) \\ \geq \frac {1}{n} \sum_ {l \in L ^ {*}} \lambda_ {\min } \left(Q ^ {T} A _ {l} ^ {T} A _ {l} Q\right) \\ \geq \frac {1}{n} \sum_ {l \in L ^ {*}} \min  _ {a \in S ^ {K - 1}} a ^ {T} Q ^ {T} A _ {l} ^ {T} A _ {l} Q a \\ \geq \frac {1}{n} \left| L ^ {*} \right| \min  _ {a \in S ^ {K - 1}, P \in \mathcal {P}} a ^ {T} Q ^ {T} \left(I _ {K + 1} - P\right) ^ {T} \left(I _ {K + 1} - P\right) Q a \\ \geq \frac {\gamma}{2} a _ {*} ^ {T} Q ^ {T} (I _ {K + 1} - \mathsf {P} _ {*}) ^ {T} (I _ {K + 1} - \mathsf {P} _ {*}) Q a _ {*} \\ \geq \frac {\gamma}{2} a _ {*} ^ {T} Q ^ {T} (I _ {K + 1} - \mathsf {P} _ {*}) ^ {T} (I _ {K + 1} - \mathsf {P} _ {*}) Q a _ {*} \\ := m _ {*}, \\ \end{array}
$$
where  $a_{*} \in S^{K-1}$ ,  $\mathsf{P}_{*} \in \mathcal{P}$ , and  $m_{*}$  is a positive constant depending on  $(K, \gamma, \tau)$  only. The first inequality above uses the property  $\lambda_{\min}(A + B) \geq \lambda_{\min}(A) + \lambda_{\min}(B)$ , where  $A$  and  $B$  are two positive semi-definite matrices.
In other words,  $\frac{1}{n}\sum_{i=1}^{n}\nabla L_i(M^T e_i)\nabla L_i(M^T e_i)^T \succeq m_*I_K$ . For the upper bound, since  $\|Q\|_{op} \leq 1$  by definition of  $Q$  and all the entries of  $(y_i - p_i)(y_i - p_i)^T$  are between  $-1$  and  $1$  if Assumption 2.2 or Assumption S1.1 holds, we have
$$
\| \sum_ {i = 1} ^ {n} \nabla L _ {i} \left(M ^ {T} e _ {i}\right) \nabla L _ {i} \left(M ^ {T} e _ {i}\right) ^ {T} \| _ {o p} = \| Q ^ {T} \sum_ {i = 1} ^ {n} \left(\mathbf {y} _ {i} - \mathbf {p} _ {i}\right) \left(\mathbf {y} _ {i} - \mathbf {p} _ {i}\right) ^ {T} Q \| _ {o p} \leq n K.
$$
Lemma S6.2 (deterministic result on Hessian). Let  $L_{i}$  be defined as in (S6.1). For all  $i\in [n]$ , we have  $\nabla^2 L_i(u)\preceq I_K$  for any  $u\in \mathbb{R}^K$  and
$$
\min_{u\in \mathbb{R}^{K},\| Qu\|_{\infty}\leq r}\nabla^{2}L_{i}(u)\succeq c_{*}I_{K},
$$
where  $c_*$  is a positive constant depending on  $(K, r)$  only.
Proof of Lemma S6.2. Recall that  $L_{i}(u) = \mathsf{L}_{i}(Qu)$ , we have
$$
\nabla^ {2} L _ {i} (u) = Q ^ {T} \nabla^ {2} \mathsf {L} _ {i} (Q u) Q,
$$
where  $\nabla^2\mathsf{L}_i(Qu) = \mathrm{diag}(\mathsf{p}_i) - \mathsf{p}_i\mathsf{p}_i^T$  and the  $k$ -th entry of  $\mathsf{p}_i \in \mathbb{R}^{K+1}$  is defined as  $\mathsf{p}_{ik} = \frac{\exp(e_k^T Q u)}{\sum_{k'=1}^{K+1} \exp(e_{k'}^T Q u)}$  for all  $i \in [n], k \in [K+1]$ . Thus,  $\mathsf{p}_{ik} \leq 1$  for all  $i \in [n], k \in [K+1]$ .
For any vector  $a \in S^{K-1}$ , we have
$$
\begin{array}{l} a ^ {T} \nabla^ {2} L _ {i} (u) a = a ^ {T} Q ^ {T} [ \operatorname {d i a g} (\mathfrak {p} _ {i}) - \mathfrak {p} _ {i} \mathfrak {p} _ {i} ^ {T} ] Q a \\ \leq a ^ {T} Q ^ {T} \operatorname {d i a g} \left(\mathfrak {p} _ {i}\right) Q a \\ \leq \left\| Q ^ {T} \operatorname {d i a g} (\mathfrak {p} _ {i}) Q \right\| _ {o p} \\ \leq 1, \\ \end{array}
$$
where the last inequality uses  $\| Q\|_{op}\leq 1$  and  $\mathfrak{p}_{ik}\leq 1$  for any  $k$ . Hence,  $\nabla^2 L_i(u)\preceq I_K$  for any  $i\in [n]$
Now we prove the lower bound. For any  $u \in \mathbb{R}^K$  such that  $\| Q u \|_{\infty} \leq r$ , we have
$$
\mathsf {p} _ {i k} = \frac {\exp (e _ {k} ^ {T} Q u)}{\sum_ {k ^ {\prime} = 1} ^ {K + 1} \exp (e _ {k ^ {\prime}} ^ {T} Q u)} \in [ c, 1 - c ]
$$
for some constant  $c$  depending on  $(K,r)$  only. For any vector  $a\in S^{K - 1}$ , let  $\eta = Qa\in \mathbb{R}^{K + 1}$ , then  $\mathbf{1}^T\eta = 0$  and
$$
\begin{array}{l} a ^ {T} \left[ \nabla^ {2} L _ {i} (u) \right] a = a ^ {T} Q ^ {T} \left[ \operatorname {d i a g} \left(\mathfrak {p} _ {i}\right) - \mathfrak {p} _ {i} \mathfrak {p} _ {i} ^ {T} \right] Q a \\ = \eta^ {T} \left[ \operatorname {d i a g} \left(\mathfrak {p} _ {i}\right) - \mathfrak {p} _ {i} \mathfrak {p} _ {i} ^ {T} \right] \eta \\ = \sum_ {k = 1} ^ {K + 1} p _ {i k} \eta_ {k} ^ {2} - \left(\sum_ {k = 1} ^ {K + 1} p _ {i k} \eta_ {k}\right) ^ {2} \\ > \sum_ {k} p _ {i k} \eta_ {k} ^ {2} - \sum_ {k} p _ {i k} \eta_ {k} ^ {2} \sum_ {k} p _ {i k} \\ = \sum_ {k} p _ {i k} \eta_ {k} ^ {2} (1 - \sum_ {k} p _ {i k}) \\ = 0, \\ \end{array}
$$
where the last equality uses  $\sum_{k=1}^{K+1} \mathsf{p}_{ik} = 1$ , and the inequality follows by  $(\sum_{k} \mathsf{p}_{ik} \eta_{k})^{2} = (\sum_{k} \sqrt{\mathsf{p}_{ik}} \sqrt{\mathsf{p}_{ik}} \eta_{k})^{2} \leq \sum_{k} \mathsf{p}_{ik} \sum_{k} \mathsf{p}_{ik} \eta_{k}^{2}$  using the Cauchy-Schwarz inequality, and here “ $=$ ” holds if and only if  $\sqrt{\mathsf{p}_{ik}} \propto \sqrt{\mathsf{p}_{ik}} \eta_{k}$  for each  $k$ , which is not true since  $\mathsf{p}_{ik} \in [c, 1-c]$  and  $\mathbf{1}^{T} \eta = 0$ .
Let  $\mathcal{H} = \{Q^T (\mathrm{diag}(\mathfrak{p}) - \mathfrak{pp}^T)Q:\mathfrak{p}\in [c,1 - c]^{K + 1}\}$ , then  $\mathcal{H}$  is compact and
$$
\min  _ {u \in \mathbb {R} ^ {K}, \| Q u \| _ {\infty} \leq r} \lambda_ {\min } (\nabla^ {2} L _ {i} (u)) \geq \min  _ {a \in S ^ {K - 1}, H \in \mathcal {H}} a ^ {T} H a = a _ {*} ^ {T} H _ {*} a _ {*} > 0
$$
for some  $a_{*}\in S^{K - 1}$  and  $H_{*}\in \mathcal{H}$ . Therefore,
$$
\min  _ {u \in \mathbb {R} ^ {K}, \| Q u \| _ {\infty} \leq r} \nabla^ {2} L _ {i} (u)) \succeq c _ {*} I _ {K}
$$
S6.2 Lipschitz conditions
We first restate the definitions of following sets,
$$
U = \left\{\left(Y, X\right) \in \mathbb {R} ^ {n \times (K + 1)} \times \mathbb {R} ^ {n \times p}: \hat {\mathrm {B}} \text {e x i s t s}, \| X \hat {\mathrm {B}} \left(I _ {K + 1} - \frac {\mathbf {1 1} ^ {T}}{K + 1}\right) \| _ {F} ^ {2} <   n \tau \right\},
$$
$$
U _ {y} = \left\{Y \in \mathbb {R} ^ {n \times (K + 1)}: \sum_ {i = 1} ^ {n} I \left(y _ {i k} = 1\right) \geq \gamma n \text {f o r a l l} k \in [ K + 1 ] \right\}
$$
Lemma S6.3. Assume  $p / n \leq \delta^{-1} < 1 - \alpha$  for some  $\alpha \in (0,1)$  and  $\delta > 1$ . Let  $\mathcal{I} = \{ I \subset [n] : |I| = \lceil n(1 - \alpha) \rceil \}$  and  $P_{I} = \sum_{i \in I} e_{i} e_{i}^{T}$ . Define
$$
U _ {x} = \left\{X \in \mathbb {R} ^ {n \times p}: \min  _ {I \in \mathcal {I}} \lambda_ {\min } \left(\frac {\Sigma^ {- 1 / 2} X ^ {T} P _ {I} X \Sigma^ {- 1 / 2}}{n}\right) \geq \phi_ {*} ^ {2}, \frac {\| X \Sigma^ {- 1 / 2} \| _ {o p}}{\sqrt {n}} \leq \phi^ {*} \right\} \tag {S6.3}
$$
for some positive constants  $\phi_{*},\phi^{*}$ , which depend on  $(\delta ,\alpha)$  only. Let  $U^{*} = \{(Y,X)\in U:Y\in$ $U_y,X\in U_x\}$ . Then under Assumptions 2.1, 2.3 and 2.4, and if either Assumption 2.2 or Assumption S1.1 holds, we have
(i)  $\mathbb{P}((Y,X)\in U^{*})\to 1$  as  $n,p\to \infty$
(ii) Let  $G$  be defined in Assumption S5.2. If  $\{(Y,X),(Y,\tilde{X})\} \subset U^{*}$ , we have
$$
\| G (Y, X) - G (Y, \tilde {X}) \| _ {F} \leq L \| (X - \tilde {X}) \Sigma^ {- 1 / 2} \| _ {F},
$$
where  $L$  is a positive constant depending on  $(K,\gamma ,\tau ,\alpha)$  only.
Proof of Lemma S6.3. We first prove statement (i). Under Assumption 2.1, [Bellec, 2022, Lemma 7.7] implies
$$
\mathbb {P} \left(\min  _ {I \in \mathcal {I}} \lambda_ {\min } \left(\frac {\Sigma^ {- 1 / 2} X ^ {T} P _ {I} X \Sigma^ {- 1 / 2}}{n}\right) \geq \phi_ {*} ^ {2}\right)\rightarrow 1
$$
for some positive constant  $\phi_*$  depending on  $(\delta, \alpha)$  only. Furthermore, [Davidson and Szarek, 2001, Theorem II.13] implies
$$
\mathbb {P} \big (\frac {\| X \Sigma^ {- 1 / 2} \| _ {o p}}{\sqrt {n}} \leq \phi^ {*} \big) \to 1
$$
for some positive constant  $\phi^*$  depending on  $\delta$  only. Therefore,  $\mathbb{P}(X\in U_x)\to 1$ . Under Assumption 2.3, we have  $\mathbb{P}(Y\in U_y)\to 1$ . Under Assumption 2.4, we have  $\mathbb{P}((Y,X)\in U)\to 1$ . In conclusion, under Assumptions 2.1, 2.3 and 2.4, we have  $\mathbb{P}((Y,X)\in U^{*})\to 1$  as  $n,p\to \infty$ .
Now we prove the statement (ii). For a fixed  $Y$ , let  $(Y, X), (Y, \tilde{X}) \in U^{*}$ ,  $\hat{B}, \tilde{B}$  be their corresponding minimizers of (S3.3), and  $G, \tilde{G}$  be their corresponding gradient matrices. We first provide some useful results derived from the KKT conditions. From the KKT conditions  $X^T G = \tilde{X}^T \tilde{G} = 0$ , we have
$$
\begin{array}{l} \langle X \hat {B} - \tilde {X} \tilde {B}, G - \tilde {G} \rangle = \langle \hat {B} - \tilde {B}, \tilde {X} ^ {T} \tilde {G} - X ^ {T} G \rangle + \langle X \hat {B} - \tilde {X} \tilde {B}, G - \tilde {G} \rangle \\ = - \langle (X - \tilde {X}) (\hat {B} - \tilde {B}), G \rangle + \langle (X - \tilde {X}) \hat {B}, G - \tilde {G} \rangle . \\ \end{array}
$$
Since  $\| \nabla^2 L_i(u)\|_{op}\leq 1$  for any  $u\in \mathbb{R}^{K}$  from Lemma S6.2,  $\nabla L_{i}(\cdot)$  is 1-Lipschitz. Thus,
$$
\begin{array}{l} \langle X \hat {B} - \tilde {X} \tilde {B}, G - \tilde {G} \rangle = \sum_ {i = 1} ^ {n} \langle \hat {B} ^ {T} x _ {i} - \tilde {B} ^ {T} \tilde {x} _ {i}, \nabla L _ {i} (\hat {B} ^ {T} x _ {i}) - \nabla L _ {i} (\tilde {B} ^ {T} \tilde {x} _ {i}) \rangle \\ \geq \sum_ {i = 1} ^ {n} \left\langle \nabla L _ {i} \left(\hat {B} ^ {T} x _ {i}\right) - \nabla L _ {i} \left(\tilde {B} ^ {T} \tilde {x} _ {i}\right), \nabla L _ {i} \left(\hat {B} ^ {T} x _ {i}\right) - \nabla L _ {i} \left(\tilde {B} ^ {T} \tilde {x} _ {i}\right) \right\rangle \\ = \| G - \tilde {G} \| _ {F} ^ {2}. \\ \end{array}
$$