SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Define for each  $i\in [n]$
$$
g _ {i} (Y, X) = \nabla L _ {i} (\hat {B} (Y, X) ^ {T} x _ {i}), \quad H _ {i} (Y, X) = \nabla^ {2} L _ {i} (\hat {B} (Y, X) ^ {T} x _ {i}),
$$
so that  $g_{i}(Y,X)\in \mathbb{R}^{K}$  and  $H_{i}(Y,X)\in \mathbb{R}^{K\times K}$ . Define
$$
G (Y, X) = \sum_ {i = 1} ^ {n} e _ {i} g _ {i} (Y, X) ^ {T},
$$
$$
V (Y, X) = \sum_ {i = 1} ^ {n} \left(H _ {i} (Y, X) - \left(H _ {i} (Y, X) \otimes x _ {i} ^ {T}\right) \left[ \sum_ {l = 1} ^ {n} H _ {l} (Y, X) \otimes \left(x _ {l} x _ {l} ^ {T}\right) \right] ^ {\dagger} \left(H _ {i} (Y, X) \otimes x _ {i}\right)\right),
$$
so that  $G(Y, X) \in \mathbb{R}^{n \times K}$  and  $V(Y, X) \in \mathbb{R}^{K \times K}$ . If the dependence on data  $(Y, X)$  is clear from context, we will simply write  $\hat{B}, g_i, H_i G$ , and  $V$ .
Theorem S5.1. Let Assumptions S5.1 and S5.2 be fulfilled. Let  $c_{*}, m_{*}, m^{*}, K$  be positive constants independent of  $n, p$ . Let  $U^{*} \subset \mathbb{R}^{p \times (K + 1)} \times \mathbb{R}^{n \times p}$  be an open set satisfying
(1) If  $\{(Y,X)\in U^{*}\}$ , then the minimizer  $\hat{B} (Y,X)$  in Assumption S5.2 exists,  $H_{i}\preceq I_{K}$  for each  $i\in [n]$ ,  $\frac{1}{n}\sum_{i = 1}^{n}H_{i}(Y,X)\otimes (x_{i}x_{i}^{T})\succeq c_{*}(I_{K}\otimes \Sigma)$  and  $m_{*}I_{K}\preceq \frac{1}{n} G(Y,X)^{T}G(Y,X)\preceq$ $m^{*}I_{K}$ .
(2) For any  $\{(Y,X),(Y,\tilde{X})\} \subset U^{*}$ ,  $\| G(Y,X) - G(Y,\tilde{X})\|_{F} \leq L\| (X - \tilde{X})\Sigma^{-1 / 2}\|_{F}$  holds for some positive constant  $L$ .
Then for any  $j \in [p]$  such that  $e_j^T B^* = \mathbf{0}_K^T$ , there exists a random variable  $\xi \in \mathbb{R}^K$  such that
$$
\mathbb {E} \left[ I \{(Y, X) \in U ^ {*} \} \right\rVert \frac {\left(G ^ {T} G\right) ^ {- 1 / 2} V \hat {B} ^ {T} e _ {j}}{\sqrt {\Omega_ {j j}}} - \xi \left\| ^ {2} \right] \leq \frac {C}{p - K},
$$
and  $\mathbb{P}(\| \xi \| ^2 >\chi_K^2 (\alpha))\leq \alpha$  for all  $\alpha \in (0,1)$ $C$  is a positive constant depending on  $(c_{*},m_{*},m^{*},K,L)$  only. If additionally  $\mathbb{P}((Y,X)\in U^{*})\to 1$  , then  $\xi$  in the previous display satisfies  $\xi \xrightarrow{\mathrm{d}} N(\mathbf{0},I_K)$  and
$$
\frac {\left(G ^ {T} G\right) ^ {- 1 / 2} V \hat {B} ^ {T} e _ {j}}{\sqrt {\Omega_ {j j}}} \xrightarrow {\mathrm {d}} N (\mathbf {0}, I _ {K}).
$$
The proof of Theorem S5.1 is given in next subsection.
S5.2 Proof of Theorem S5.1
In this subsection and next subsection, we will slightly abuse the notations  $A^{*}$  and  $\hat{A}$ , which have different definitions than the definitions in the main text.
Let  $\Sigma^{1/2}B^* = \sum_{k=1}^{K}s_ku_kv_k^T$  be the singular value decomposition of  $\Sigma^{1/2}B^*$ , where  $u_1, \ldots, u_K$  are the left singular vectors and  $v_1, \ldots, v_k$  the right singular vectors. If  $\Sigma^{1/2}B^*$  is of rank strictly less than  $K$ , we allow some  $s_k$  to be equal to 0 so that  $\Sigma^{1/2}B^* = \sum_{k=1}^{K}s_ku_kv_k^T$  still holds with orthonormal  $(u_1, \ldots, u_K)$  and orthonormal  $(v_1, \ldots, v_K)$ . We consider an orthogonal matrix  $\tilde{P} \in \mathbb{R}^{p \times p}$  such that
$$
\tilde {P} \tilde {P} ^ {T} = \tilde {P} ^ {T} \tilde {P} = I _ {p}, \quad \tilde {P} \frac {\sum^ {- 1 / 2} e _ {j}}{\| \sum^ {- 1 / 2} e _ {j} \|} = e _ {1}, \quad \tilde {P} u _ {k} = e _ {p - K + k}, \quad \forall k \in [ K ]. \tag {S5.1}
$$
Since  $e_j^T B^* = \mathbf{0}^T$  implies  $e_j^T \Sigma^{-1/2} u_k = 0$ , we can always find a matrix  $\tilde{P}$  satisfying (S5.1). From now on we fix this matrix  $\tilde{P}$  and consider the following change of variable,
$$
Z = X \Sigma^ {- 1 / 2} \tilde {P} ^ {T}, \quad A ^ {*} = \tilde {P} \Sigma^ {1 / 2} B ^ {*}. \tag {S5.2}
$$
It immediately follows that  $Z$  has i.i.d.  $N(0,1)$  entries and the first  $p - K$  rows of  $A^*$  are all zeros. Since the response  $y_i$  has the expression  $y_i = f(U_i, x_i^T B^*)$ ,  $Y$  is unchanged by the change of variable (S5.2) from  $ZA^* = XB^*$ . We now work on the multinomial logistic estimation with data  $(Y,Z)$  and the underlying coefficient matrix  $A^*$  in (S5.2). Parallel to the estimate  $\hat{B}$  of  $B^*$  in Assumption S5.2, we define the estimate of  $A^*$  using data  $(Y,Z)$  as
$$
\hat {A} (Y, Z) = \underset {A \in \mathbb {R} ^ {p \times K}} {\arg \min } \sum_ {i} L _ {i} \left(A ^ {T} z _ {i}\right),
$$
where  $z_{i} = Z^{T}e_{i}$  is the  $i$ -th row of  $Z$ . By construction, we have  $\hat{A} = \tilde{P}\Sigma^{1 / 2}\hat{B}$ , hence  $Z\hat{A} = X\hat{B}$  and  $e_1^T\hat{A} = e_j^T\hat{B} /\sqrt{\Omega_{jj}}$ . Furthermore, the quantities depending on  $(Y,X\hat{B})$  remain unchanged after the change of variable. In particular, the gradient and Hessian
$$
\nabla L _ {i} (\hat {B} ^ {T} x _ {i}) = \nabla L _ {i} (\hat {A} ^ {T} z _ {i}), \qquad \nabla^ {2} L _ {i} (\hat {B} ^ {T} x _ {i}) = \nabla^ {2} L _ {i} (\hat {A} ^ {T} z _ {i})
$$
are unchanged. It follows that the matrix  $G$  and  $V$  are unchanged. Therefore, we have
$$
\frac {e _ {j} ^ {T} \hat {B} V (G ^ {T} G) ^ {- 1 / 2}}{\sqrt {\Omega_ {j j}}} = e _ {1} ^ {T} \hat {A} V (G ^ {T} G) ^ {- 1 / 2}.
$$
In conclusion, with the change of variables (S5.2), we only need to prove Theorem S5.1 in the special case, where the design matrix  $X$  i.i.d.  $N(0,1)$  entries and the response  $Y$  is independent of the first  $p - K$  columns of  $X$ . To this end, we introduce the following Theorem S5.2, and the proof of Theorem S5.1 is a consequence of Theorem S5.2 as it proves the desired result for  $e_1^T \hat{A} V(G^T G)^{-1/2}$ .
Theorem S5.2. Let  $c_{*}, m_{*}, m^{*}, K$  be constants independent of  $n, p$ . Let  $Z \in \mathbb{R}^{n \times p}$  have i.i.d. rows  $(z_{1}, \ldots, z_{n})$  with  $z_{i} \sim N(0, I_{p})$ . Let  $y_{1}, \ldots, y_{n} \in \mathbb{R}^{(K + 1)}$  such that  $(y_{1}, \ldots, y_{n})$  is independent of the first  $p - K$  columns of  $Z$ . Consider twice continuously differentiable and strictly convex loss
functions  $(L_{i})_{i = 1,\dots ,n}$  with each  $L_{i}:\mathbb{R}^{K}\to \mathbb{R}$  depending on  $\mathbf{y}_i$  but not on  $z_{i}$  and define, provided that the minimizer admits a solution,
$$
\hat {A} (Y, Z) = \underset {A \in \mathbb {R} ^ {p \times K}} {\arg \min} \sum_ {i = 1} ^ {n} L _ {i} (A ^ {T} z _ {i}), g _ {i} (Y, Z) = \nabla L _ {i} (\hat {A} (Y, Z) ^ {T} z _ {i}), H _ {i} (Y, Z) = \nabla^ {2} L _ {i} (\hat {A} (Y, Z) ^ {T} z _ {i}),
$$
[ G(Y,Z) = \sum_{i=1}^{n} e_i g_i(Y,Z)^T \in \mathbb{R}^{n \times K}, \text{ and } V(Y,Z) = \sum_{i=1}^{n} \left( H_i - (H_i \otimes z_i^T) \left[ \sum_{l=1}^{n} H_l \otimes (z_l z_l^T) \right] \right)^{\dagger} (H_i \otimes z_i) \in \mathbb{R}^{K \times K}, \text{ where we dropped the dependence of } H_i \text{ on } (Y,Z) \text{ for simplicity.} ]
Let  $O \subset \mathbb{R}^{n \times (K + 1)} \times \mathbb{R}^{n \times p}$  be an open set satisfying
- If  $(Y,Z) \in O$ , then the minimizer  $\hat{A}(Y,Z)$  exists,  $H_{i} \preceq I_{K}$  for each  $i \in [n]$ ,  $c_{*}I_{pK} \preceq \frac{1}{n}\sum_{i=1}^{n} H_{i}(Y,Z) \otimes (z_{i}z_{i}^{T})$ , and  $m_{*}I_{K} \preceq \frac{1}{n}\sum_{i=1}^{n} G(Y,Z)^{T}G(Y,Z) \preceq m^{*}I_{K}$ .
- With the notation  $G(Y,Z) = \sum_{i=1}^{n} e_i g_i(Y,Z)^T$ , we have if two  $Z, \tilde{Z} \in \mathbb{R}^{n \times p}$  satisfy  $\{(Y,Z),(Y,\tilde{Z})\} \subset O$  then  $\| G(Y,Z) - G(Y,\tilde{Z}) \| \leq L \| Z - \tilde{Z} \|$ .
For  $e_1 \in \mathbb{R}^p$  the first canonical basis vector, there exists a random variable  $\xi \in \mathbb{R}^{K}$  such that
$$
\mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| (G ^ {T} G) ^ {- 1 / 2} V \hat {A} ^ {T} e _ {1} - \xi \right\| ^ {2} \right] \leq \frac {C}{p - K},
$$
and  $\mathbb{P}(\| \xi \| ^2 >\chi_K^2 (\alpha))\leq \alpha$  for all  $\alpha \in (0,1)$ ,  $C$  is a positive constant depending on  $(c_{*},m_{*},m^{*},K,L)$  only. If additionally  $\mathbb{P}((Y,Z)\in O)\to 1$ , then  $\xi$  in the previous display satisfies  $\xi\xrightarrow{\mathrm{d}}N(\mathbf{0},I_K)$  and
$$
e _ {1} ^ {T} \hat {A} V (G ^ {T} G) ^ {- 1 / 2} \xrightarrow {\mathrm {d}} N (\mathbf {0}, I _ {K}).
$$
The proof of Theorem S5.2 is presented in Section S5.3.
S5.3 Proof of Theorem S5.2
We first present a few useful lemmas, whose proofs are given at the end of this subsection.
Lemma S5.3 (Proof is given on page 24). Let  $z \sim N(0, \sigma^2 I_n)$  and  $F: \mathbb{R}^n \to \mathbb{R}^{n \times K}$  be weakly differentiable with  $\mathbb{E}\| F(z)\|_F^2 < \infty$ . Let  $\tilde{z}$  be an independent copy of  $z$ . Then
$$
\mathbb {E} \left[ \left\| z ^ {T} F (z) - \sigma^ {2} \sum_ {i = 1} ^ {n} \frac {\partial e _ {i} ^ {T} F (z)}{\partial z _ {i}} - z ^ {T} F (\tilde {z}) \right\| ^ {2} \right] \leq 3 \sigma^ {4} \mathbb {E} \sum_ {i = 1} ^ {n} \left\| \frac {\partial F (z)}{\partial z _ {i}} \right\| _ {F} ^ {2}.
$$
Lemma S5.4 (Proof is given on page 25). If  $G, \tilde{G} \in \mathbb{R}^{n \times K}$  satisfy  $m_* I_K \preceq \frac{1}{n} G^T G \preceq m^* I_K$  and  $m_* I_K \preceq \frac{1}{n} \tilde{G}^T \tilde{G} \preceq m^* I_K$  for some positive constants  $m_*$  and  $m^*$ . Then
$$
\left\| \left(G ^ {T} G\right) ^ {- 1 / 2} - \left(\tilde {G} ^ {T} \tilde {G}\right) ^ {- 1 / 2} \right\| _ {F} \leq L _ {1} n ^ {- 1} \| G - \tilde {G} \| _ {F},
$$
$$
\| G (G ^ {T} G) ^ {- 1 / 2} - \tilde {G} (\tilde {G} ^ {T} \tilde {G}) ^ {- 1 / 2} \| _ {F} \leq L _ {2} n ^ {- 1 / 2} \| G - \tilde {G} \| _ {F},
$$
where  $L_{1}, L_{2}$  are positive constants depending on  $(K, m_{*}, m^{*})$  only.
Lemma S5.5 (Proof is given on page 26). Let the assumptions in Theorem S5.2 be fulfilled. Let  $Y \in \mathbb{R}^{n \times (K + 1)}$  be fixed. If a minimizer  $\hat{A}(Y,Z)$  exists at  $Z$ , then  $Z \mapsto \hat{A}(Y,Z)$  exists and is differentiable in a neighborhood of  $Z$  with derivative
$$
\frac {\partial \operatorname {v e c} (\hat {A})}{\partial z _ {i j}} = - M [ g _ {i} \otimes e _ {j} + (H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}) ],
$$
$$
\frac {\partial g _ {l}}{\partial z _ {i j}} = - (H _ {l} \otimes z _ {l} ^ {T}) M [ g _ {i} \otimes e _ {j} + (H _ {i} \hat {A} ^ {T} e _ {j} \otimes z _ {i}) ] + I (l = i) H _ {l} \hat {A} ^ {T} e _ {j},
$$
where  $M = \left[\sum_{i=1}^{n} H_i \otimes (z_i z_i^T)\right]^{-1}$ . It immediately follows that
$$
\frac {\partial g _ {i}}{\partial z _ {i j}} = [ H _ {i} - (H _ {i} \otimes z _ {i} ^ {T}) M (H _ {i} \otimes z _ {i}) ] \hat {A} ^ {T} e _ {j} - (H _ {i} \otimes z _ {i} ^ {T}) M (g _ {i} \otimes e _ {j}).
$$
Corollary S5.6 (Proof is given on page 26). Under the same conditions of Lemma S5.5, for  $G = \sum_{i=1}^{n} e_i g_i^T$ , we have for each  $i \in [n], j \in [p]$
$$
\begin{array}{l} \sum_ {i = 1} ^ {n} \frac {\partial e _ {i} ^ {T} G (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i j}} \\ = e _ {j} ^ {T} \hat {A} V _ {i} ^ {T} (G ^ {T} G) ^ {- 1 / 2} + \sum_ {i = 1} ^ {n} \left[ - \left(g _ {i} ^ {T} \otimes e _ {j} ^ {T}\right) M \left(H _ {i} \otimes z _ {i}\right) \left(G ^ {T} G\right) ^ {- 1 / 2} + e _ {i} ^ {T} G \frac {\partial \left(G ^ {T} G\right) ^ {- 1 / 2}}{\partial z _ {i j}} \right]. \\ \end{array}
$$
Now we are ready to prove Theorem S5.2.
Proof of Theorem S5.2. Let  $h: O \to \mathbb{R}^{n \times K}$  be  $h(Y,Z) = G(Y,Z)(G(Y,Z)^T G(Y,Z))^{-1/2}$ . In most of this proof, we will omit the dependence  $(Y,Z)$  on  $h, \hat{A}, g_i, H_i, G, V$  to lighten notation. By Lemma S5.4, we know this  $h$  is  $LL_2n^{-1/2}$ -Lipschitz in the sense that  $\|h(Y,Z) - h(Y,\tilde{Z})\|_F \leq LL_2n^{-1/2}\|Z - \tilde{Z}\|_F$  for all  $\{(Y,Z),(Y,\tilde{Z})\} \subset O$ . By Kirschbraun theorem, there exists a function  $H: \mathbb{R}^{n \times (K+1)} \times \mathbb{R}^{n \times p} \to \mathbb{R}^{n \times K}$  (an extension of  $h$  from  $O$  to  $\mathbb{R}^{n \times (K+1)} \times \mathbb{R}^{n \times p}$ ) such that  $H(Y,Z) = h(Y,Z)$  for all  $(Y,Z) \in O$ ,  $\|H(Y,Z)\|_{op} \leq 1$  and  $\|H(Y,Z) - H(Y,\tilde{Z})\|_F \leq LL_2n^{-1/2}\|Z - \tilde{Z}\|_F$  for all  $\{(Y,Z),(Y,\tilde{Z})\} \subset \mathbb{R}^{n \times (K+1)} \times \mathbb{R}^{n \times p}$ .
For each  $j \in [p]$ , let  $\mathbf{z}_j = Z e_j$  be the  $j$ -th column of  $Z$  to distinguish it from the notation  $z_i$ , which means the  $i$ -th row of  $Z$ . Let  $\check{\mathbf{z}} \sim N(\mathbf{0}, I_n)$  be an independent copy of each columns of  $Z$ , and  $\check{Z}^j = Z(I_p - e_j e_j^T) + \check{z} e_j^T$ . That is,  $\check{Z}^j$  replaces the  $j$ -th column of  $Z$  by  $\check{\mathbf{z}}$ . By definition,  $\mathbf{z}_1 \perp \check{\mathbf{z}}$  and  $\mathbf{z}_1 \perp \check{Z}^1$ .
Let  $\xi = -[H(Y,\tilde{Z}^{1})]^{T}\mathbf{z}_{1}\in \mathbb{R}^{K}$ , then  $\| \xi \| ^2\leq \| \mathbf{z}_1\|$  since  $\| H(Y,\tilde{Z}^1)\|_{op}\leq 1$ . It follows that
$$
\mathbb {P} \big (\| \xi \| ^ {2} > \chi_ {K} ^ {2} (\alpha) \big) \leq \mathbb {P} \big (\| \mathsf {z} _ {1} \| ^ {2} > \chi_ {K} ^ {2} (\alpha) \big) = \alpha .
$$
Note that the first  $p - K$  columns of  $Z$  are exchangeable, because they are i.i.d. and independent of the response  $Y$ , we have for each  $\ell \in [p - K]$
$$
\begin{array}{l} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| (G ^ {T} G) ^ {- 1 / 2} V \hat {A} ^ {T} e _ {1} - \xi \right\| ^ {2} \right] \\ = \mathbb {E} \left[ I \{(Y, Z) \in O \} \| e _ {1} ^ {T} \hat {A} V (G ^ {T} G) ^ {- 1 / 2} + z _ {1} ^ {T} H (Y, \check {Z} ^ {1}) \| ^ {2} \right] \\ = \mathbb {E} \left[ I \{(Y, Z) \in O \} \| e _ {\ell} ^ {T} \hat {A} V (G ^ {T} G) ^ {- 1 / 2} + \mathbf {z} _ {\ell} ^ {T} H (Y, \check {Z} ^ {\ell}) \| ^ {2} \right], \\ \end{array}
$$
where the last line holds for any  $\ell \in [p - K]$  because  $(\mathbf{z}_1, e_1^T \hat{A}, \tilde{Z}) \stackrel{d}{=} (\mathbf{z}_\ell, e_\ell^T \hat{A}, \tilde{Z}^\ell)$ . Therefore,
$$
\begin{array}{l} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| e _ {1} ^ {T} \hat {A} V (G ^ {T} G) ^ {- 1 / 2} + z _ {1} ^ {T} H (Y, \tilde {Z} ^ {1}) \right\| ^ {2} \right] \\ = \frac {1}{p - K} \sum_ {\ell = 1} ^ {p - K} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| e _ {\ell} ^ {T} \hat {A} V (G ^ {T} G) ^ {- 1 / 2} + z _ {\ell} ^ {T} H (Y, \check {Z} ^ {\ell}) \right\| ^ {2} \right] \\ = \frac {1}{p - K} \sum_ {\ell = 1} ^ {p - K} \mathbb {E} \Big [ I \{(Y, Z) \in O \} \Big \| \sum_ {i} \frac {\partial e _ {i} ^ {T} G (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i \ell}} + z _ {\ell} ^ {T} H (Y, \check {Z} ^ {\ell}) - \mathrm {R e m} _ {\ell} \Big \| ^ {2} \Big ], \\ \end{array}
$$
where  $\mathrm{Rem}_{\ell} = \sum_{i=1}^{n}\left[-(g_i^T \otimes e_\ell^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\ell}}\right]$  from Corollary S5.6, and  $M = \left[\sum_{i=1}^{n}(H_i \otimes z_i z_i^T)\right]^{-1}$  from Lemma S5.5. Using  $(a + b)^2 \leq 2a^2 + 2b^2$ , the above display can be bounded by sum of two terms, denoted by  $(RHS)_1$  and  $(RHS)_2$ .
For the first term,
$$
(R H S) _ {1} = \frac {2}{p - K} \sum_ {\ell = 1} ^ {p - K} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| \sum_ {i} \frac {\partial e _ {i} ^ {T} h (Y , Z)}{\partial z _ {i \ell}} + z _ {\ell} ^ {T} H (Y, \check {Z} ^ {\ell}) \right\| ^ {2} \right].
$$
Let  $F(\mathbf{z}_{\ell}) = H(Y,Z(I - e_{\ell}e_{\ell}^{T}) + \mathbf{z}_{\ell}e_{\ell}^{T}) = H(Y,Z)$ , then  $F(\check{\mathbf{z}}) = H(Y,\check{Z}^{\ell})$ . Apply Lemma S5.3 to  $F(\mathbf{z}_{\ell})$  conditionally on  $Z(I - e_{\ell}e_{\ell}^{T})$ , we obtain
$$
\begin{array}{l} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| \sum_ {i} \frac {\partial e _ {i} ^ {T} h (Y , Z)}{\partial z _ {i \ell}} + z _ {\ell} ^ {T} F (\check {z}) \right\| ^ {2} \right] \\ = \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| \mathbf {z} _ {\ell} ^ {T} F (\mathbf {z} _ {\ell}) - \sum_ {i} \frac {\partial e _ {i} ^ {T} F (\mathbf {z} _ {\ell})}{\partial z _ {i \ell}} - \mathbf {z} _ {\ell} ^ {T} F (\check {\mathbf {z}}) \right\| ^ {2} \right] \\ \leq \mathbb {E} \left[ \left\| \mathbf {z} _ {\ell} ^ {T} F (\mathbf {z} _ {\ell}) - \sum_ {i} \frac {\partial e _ {i} ^ {T} F (\mathbf {z} _ {\ell})}{\partial z _ {i \ell}} - \mathbf {z} _ {\ell} ^ {T} F (\check {\mathbf {z}}) \right\| ^ {2} \right] \\ \leq 3 \sum_ {i} \mathbb {E} \left\| \frac {\partial F (\mathbf {z} _ {\ell})}{\partial z _ {i \ell}} \right\| _ {F} ^ {2} \\ = 3 \sum_ {i} \mathbb {E} \left\| \frac {\partial H (Y , Z)}{\partial z _ {i \ell}} \right\| _ {F} ^ {2}, \\ \end{array}
$$
where the first equality uses  $\mathbf{z}_{\ell}^{T}h(Y,Z) = 0$  from the KKT conditions  $Z^T G = 0$  and  $h(Y,Z) = G(G^T G)^{-1/2}$ . It follows that
$$
(R H S) _ {1} \leq \frac {6}{p - K} \mathbb {E} \left[ \sum_ {\ell = 1} ^ {p} \sum_ {i = 1} ^ {n} \left\| \frac {\partial H (Y , Z)}{\partial z _ {i \ell}} \right\| _ {F} ^ {2} \right].
$$
Note that the integrand in the last display is actually the squared Frobenius norm of the Jacobian of the mapping from  $\mathbb{R}^{n\times p}$  to  $\mathbb{R}^{n\times K}$ :  $Z\mapsto H(Y,Z)$ . This Jacobian is a matrix with  $nK$  rows and  $np$  columns, has rank at most  $nK$  and operator norm at most  $LL_{2}n^{-1 / 2}$  because  $Z\mapsto H(Y,Z)$  is  $LL_{2}n^{-1 / 2}$ -Lipschitz from Lemma S5.4. Using  $\| A\| _F^2\leq \mathrm{rank}(A)\| A\|_{op}^2$ , we obtain
$$
(R H S) _ {1} \leq 6 K (L L _ {2}) ^ {2} / (p - K).
$$
For the second term  $(RHS)_2 = \frac{2}{p - K}\sum_{\ell = 1}^{p - K}\mathbb{E}\big[I\{(Y,Z)\in O\} \| \mathrm{Rem}_\ell \| ^2\big]$ . By definition of  $\mathrm{Rem}_{\ell}$  and  $(a + b)^2\leq 2a^2 +2b^2$  , we obtain
$$
\begin{array}{l} (R H S) _ {2} \leq \frac {4}{p - K} \sum_ {\ell = 1} ^ {p - K} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| \sum_ {i = 1} ^ {n} \left(g _ {i} ^ {T} \otimes e _ {\ell} ^ {T}\right) M \left(H _ {i} \otimes z _ {i}\right) \left(G ^ {T} G\right) ^ {- 1 / 2} \right\| ^ {2} \right] (S5.3) \\ + \frac {4}{p - K} \sum_ {\ell = 1} ^ {p - K} \mathbb {E} \left[ I \{(Y, Z) \in O \} \left\| \sum_ {i = 1} ^ {n} e _ {i} ^ {T} G \frac {\partial \left(G ^ {T} G\right) ^ {- 1 / 2}}{\partial z _ {i \ell}} \right\| ^ {2} \right]. (S5.4) \\ \end{array}
$$
We next bound (S5.3) and (S5.4) one by one. For (S5.3), we focus on the norm without  $(G^T G)^{-1/2}$  which is  $\| \sum_{i=1}^{n} (g_i^T \otimes e_\ell^T) M(H_i \otimes z_i) \|$ . With  $\| a \| = \max_{u: \| u \| = 1} a^T u$  in mind, let us multiply to the right by a unit vector  $u \in \mathbb{R}^K$  and instead bound
$$
\sum_ {i = 1} ^ {n} \left(g _ {i} ^ {T} \otimes e _ {\ell} ^ {T}\right) M \left(H _ {i} \otimes z _ {i}\right) u = \operatorname {T r} \left[ \left(I _ {K} \otimes e _ {\ell} ^ {T}\right) M \sum_ {i} \left(H _ {i} u \otimes z _ {i}\right) g _ {i} ^ {T} \right] \leq K \| \left(I _ {K} \otimes e _ {\ell} ^ {T}\right) M \| _ {o p} \| \sum_ {i} \left(H _ {i} u \otimes z _ {i}\right) g _ {i} ^ {T} \| _ {o p}
$$
because the rank of the matrix inside the trace is at most  $K$  and  $\mathrm{Tr}[\cdot ]\leq K\| \cdot \|_{op}$  holds. Then
$$
\| \sum_ {i} (H _ {i} u \otimes z _ {i}) g _ {i} ^ {T} \| _ {o p} = \| (I _ {K} \otimes Z ^ {T}) \sum_ {i} (H _ {i} u \otimes e _ {i}) e _ {i} ^ {T} G \| _ {o p} \leq \| Z \| _ {o p} \| \sum_ {i} (H _ {i} u \otimes e _ {i} e _ {i} ^ {T}) \| _ {o p} \| G \| _ {o p}.
$$
Next,  $\| \sum_{i}(H_{i}u\otimes e_{i}e_{i}^{T})\|_{op} = \| \sum_{i}(H_{i}\otimes e_{i}e_{i}^{T})(u\otimes I_{n})\|_{op}\leq 1$  because  $H_{i}\preceq I_{K}$  and  $\| u\| = 1$  In summary, the norm in (S5.3) is bounded from above by
$$
K \| \left(G ^ {T} G\right) ^ {- 1 / 2} \| _ {o p} \| M \| _ {o p} \| Z \| _ {o p} \| G \| _ {o p}.
$$
To bound (S5.3), since in the event  $(Y,Z)\in O$ ,  $m_{*}I_{K}\preceq \frac{1}{n} G^{T}G\preceq m^{*}I_{K}$  and  $\frac{1}{n}\sum_{i = 1}^{n}(H_i\otimes z_i z_i^T)\succeq c_*I_{pK}$ , we have  $\| G\|_{op}\leq \sqrt{m^{*}n}$ , hence  $\| M\|_{op}\leq c_{*}^{-1}$ . Thus, the above display can be bounded by
$$
K (m _ {*} n) ^ {- 1 / 2} c _ {*} ^ {- 1} \| Z \| _ {o p} \sqrt {m ^ {*} n} = (m ^ {*} / m _ {*}) ^ {1 / 2} c _ {*} ^ {- 1} K \| Z \| _ {o p}.
$$
Since  $Z \in \mathbb{R}^{n \times p}$  has i.i.d.  $N(0,1)$  entries, [Davidson and Szarek, 2001, Theorem II.13] implies that  $\mathbb{E}\|Z\|_{op} \leq \sqrt{n} + \sqrt{p} \leq 2\sqrt{n}$ . Therefore,
$$
\left. \left(\mathrm {S} 5. 3\right) \leq C \left(c _ {*}, K, L\right) n ^ {- 1}. \right.
$$
Now we bound (S5.4). Since
$$
\begin{array}{l} \sum_ {\ell = 1} ^ {p - K} \left\| \sum_ {i = 1} ^ {n} e _ {i} ^ {T} G \frac {\partial (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i \ell}} \right\| ^ {2} \\ \leq \sum_ {j = 1} ^ {p} \left\| \sum_ {i = 1} ^ {n} e _ {i} ^ {T} G \frac {\partial \left(G ^ {T} G\right) ^ {- 1 / 2}}{\partial z _ {i j}} \right\| ^ {2} \\ = \sum_ {j = 1} ^ {p} \sum_ {k ^ {\prime} = 1} ^ {K} \left(\sum_ {i = 1} ^ {n} \sum_ {k = 1} ^ {K} e _ {i} ^ {T} G e _ {k} e _ {k} ^ {T} \frac {\partial (G ^ {T} G) ^ {- 1 / 2}}{\partial z _ {i j}} e _ {k ^ {\prime}}\right) ^ {2} \\ \leq \sum_ {j = 1} ^ {p} \sum_ {k ^ {\prime} = 1} ^ {K} \left[ \sum_ {i = 1} ^ {n} \sum_ {k = 1} ^ {K} \left(e _ {i} ^ {T} G e _ {k}\right) ^ {2} \sum_ {i = 1} ^ {n} \sum_ {k = 1} ^ {K} \left(e _ {k} ^ {T} \frac {\partial \left(G ^ {T} G\right) ^ {- 1 / 2}}{\partial z _ {i j}} e _ {k ^ {\prime}}\right) ^ {2} \right] \\ = \| G \| _ {F} ^ {2} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {p} \left\| \frac {\partial \left(G ^ {T} G\right) ^ {- 1 / 2}}{\partial z _ {i j}} \right\| ^ {2}. \\ \end{array}
$$
Using  $\| G\| _F^2\leq nK$ , and the mapping  $Z\mapsto (G^{T}G)^{-1 / 2}$  is  $LL_{1}n^{-1}$ -Lipschitz on  $O$  using Lemma S5.4, we conclude that
$$
(\mathrm {S} 5. 4) \leq 4 K ^ {3} L L _ {1} / (p - K).
$$
Combining the above bounds on  $(RHS)_1$  and  $(RHS)_2$ , we have
$$
\mathbb {E} \left[ I \{(Y, Z) \in O \} \| \left(G ^ {T} G\right) ^ {- 1 / 2} V \hat {A} ^ {T} e _ {1} - \xi \| ^ {2} \right] \leq \frac {C \left(c _ {*} , K , m _ {*} , m ^ {*} , L , L _ {1} , L _ {2}\right)}{p - K}, \tag {S5.5}
$$
where the constant depends on  $(c_{*},K,m_{*},m^{*},L)$  only because  $L_{1}$  and  $L_{2}$  are constants depending on  $(K,m_{*},m^{*})$  only.
If additionally  $\mathbb{P}((Y,Z)\in O)\to 1$ , we have  $\mathbb{P}((Y,\tilde{Z}^1)\in O)\to 1$  using  $(Y,Z)\stackrel {d}{=}(Y,\tilde{Z}^{1})$ . Therefore,
$$
\xi = - \left[ h \left(Y, \check {Z} ^ {1}\right) ^ {T} \mathbf {z} _ {1} \right] I \left(\left(Y, \check {Z} ^ {1}\right) \in O\right) - \left[ H \left(Y, \check {Z} ^ {1}\right) ^ {T} \mathbf {z} _ {1} \right] I \left(\left(Y, \check {Z} ^ {1}\right) \notin O\right) \xrightarrow {\mathrm {d}} N (\mathbf {0}, I _ {K}). \tag {S5.6}
$$
By (S5.5), we know  $(G^T G)^{-1/2} V \hat{A}^T e_1 - \xi \xrightarrow{\mathrm{d}} 0$  when  $\mathbb{P}((Y,Z) \in O) \to 1$ . Hence, we conclude
$$
(G ^ {T} G) ^ {- 1 / 2} V \hat {A} ^ {T} e _ {1} \xrightarrow {\mathrm {d}} N (\mathbf {0}, I _ {K}) \quad \text {a n d} \quad \| (G ^ {T} G) ^ {- 1 / 2} V \hat {A} ^ {T} e _ {1} \| ^ {2} \xrightarrow {\mathrm {d}} \chi_ {K} ^ {2}.
$$
We next prove Lemmas S5.3 to S5.5 and corollary S5.6.
Proof of Lemma S5.3. Let  $z_0 = (z^T, \tilde{z}^T)^T \in \mathbb{R}^{2n}$ , then  $z_0 \sim N(0, \sigma^2 I_{2n})$ . For each  $k \in [K]$ , let  $f^{(k)}: \mathbb{R}^{2n} \to \mathbb{R}^{2n}$  be
$$
f ^ {(k)} (z _ {0}) = \left( \begin{array}{c} [ F (z) - F (\tilde {z}) ] e _ {k} \\ 0 _ {n} \end{array} \right),
$$
so that  $z_0^T f^{(k)}(z_0) = z^T [F(z) - F(\tilde{z})]e_k$ , and  $\mathrm{div}f^{(k)}(z_0) = \sum_{i = 1}^{n}\frac{\partial e_i^T F(z)e_k}{\partial z_i}$ . Applying the second order Stein formula [Bellec and Zhang, 2021] (see also [Tan et al., 2022, Lemma F.1] for a
collection of such formulas) to  $f^{(k)}$  gives, with Jac denoting the Jacobian,