a

#  Multinomial Logistic Regression: Asymptotic Normality on Null Covariates in
High-Dimensions

Kai Tan  Department of Statistics, Rutgers University, Piscataway, NJ 08854,
USA<EMAIL_ADDRESS>and  Pierre C. Bellec  Department of Statistics,
Rutgers University, Piscataway, NJ 08854, USA<EMAIL_ADDRESS>
###### Abstract.

This paper investigates the asymptotic distribution of the maximum-likelihood
estimate (MLE) in multinomial logistic models in the high-dimensional regime
where dimension and sample size are of the same order. While classical large-
sample theory provides asymptotic normality of the MLE under certain
conditions, such classical results are expected to fail in high-dimensions as
documented for the binary logistic case in the seminal work of Sur and Candès
(2019). We address this issue in classification problems with 3 or more
classes, by developing asymptotic normality and asymptotic chi-square results
for the multinomial logistic MLE (also known as cross-entropy minimizer) on
null covariates. Our theory leads to a new methodology to test the
significance of a given feature. Extensive simulation studies on synthetic
data corroborate these asymptotic results and confirm the validity of proposed
p-values for testing the significance of a given feature.

##  1\. Introduction

Multinomial logistic modeling has become a cornerstone of classification
problems in machine learning, as witnessed by the omnipresence of both the
cross-entropy loss (multinomial logistic loss) and the softmax function
(gradient of the multinomial logistic loss) in both applied and theoretical
machine learning. We refer to Cramer (2002) for an account of the history and
early developments of logistic modeling.

Throughout, we consider a classification problem with $K+1$ possible labels
where $K$ is a fixed constant. This paper tackles asymptotic distributions of
multinomial logistic estimates (or cross-entropy minimizers) in generalized
linear models with moderately high-dimensions, where sample size $n$ and
dimension $p$ have the same order, for instance $n,p\to+\infty$ simultaneously
while the ratio $p/n$ converges to a finite constant. Throughout the paper,
let $[n]=\\{1,2,\ldots,n\\}$ for all $n\in\mathbb{N}$, and
$I\\{\text{statement}\\}$ be the 0-1 valued indicator function, equal to 1 if
statement is true and 0 otherwise (e.g., $I\\{y_{i}=1\\}$ in the next
paragraph equals 1 if $y_{i}=1$ holds and 0 otherwise).

#### The case of binary logistic regression.

Let $\rho(t)=\log(1+e^{t})$ be the logistic loss and
$\rho^{\prime}(t)=1/(1+e^{-t})$ be its derivative, often referred to as the
sigmoid function. In the current moderately-high dimensional regime where
$n,p\to+\infty$ with $p/n\to\kappa>0$ for some constant $\kappa$, recent works
(Candès and Sur, 2020; Sur and Candès, 2019; Zhao et al., 2022) provide a
detailed theoretical understanding of the behavior of the logistic Maximum
Likelihood Estimate (MLE) in binary logistic regression models. Observing
independent observations $(x_{i},y_{i})_{i\in[n]}$ from a logistic model
defined as $\mathbb{P}(y_{i}=1|x_{i})=\rho^{\prime}(x_{i}^{T}\beta)$ where
$x_{i}\sim N(\bm{0},n^{-1}I_{p})$, and
$\lim_{n\to\infty}\|\beta\|^{2}/n=\gamma^{2}$ for a constant $\gamma$ for the
limiting squared norm of the unknown regression vector $\beta$. These works
prove that the behavior of the MLE
$\hat{\beta}=\operatorname*{arg\,min}_{b\in\mathbb{R}^{p}}\sum_{i=1}^{n}\rho(x_{i}^{T}b)-I\\{y_{i}=1\\}x_{i}^{T}b$
is summarized by the solution $(\alpha_{*},\sigma*,\lambda_{*})$ of the system
of three equations

(1.1)
$\begin{cases}\sigma^{2}&=\frac{1}{\kappa^{2}}\mathbb{E}[2\rho^{\prime}(\gamma
Z_{1})(\lambda\rho^{\prime}({\mathrm{prox}}_{\lambda\rho}(-\alpha\gamma
Z_{1}+\sqrt{\kappa}\sigma Z_{2})))^{2}]\\\ 0&=\mathbb{E}[\rho^{\prime}(\gamma
Z_{1})\lambda\rho^{\prime}({\mathrm{prox}}_{\lambda\rho}(-\alpha\gamma
Z_{1}+\sqrt{\kappa}\sigma Z_{2}))]\\\
1-\kappa&=\mathbb{E}[{2\rho^{\prime}(\gamma
Z_{1})}\big{/}\bigl{(}{1+\lambda\rho^{\prime\prime}({\mathrm{prox}}_{\lambda\rho}(-\alpha\gamma
Z_{1}+\sqrt{\kappa}\sigma Z_{2}))}\bigr{)}]\end{cases},$

where $(Z_{1},Z_{2})$ are i.i.d. $N(0,1)$ random variables and the proximal
operator is defined as
${\mathrm{prox}}_{\lambda\rho}(z)=\operatorname*{arg\,min}_{t\in\mathbb{R}}\bigl{\\{}\lambda\rho(t)+(t-z)^{2}/2\bigr{\\}}$.
The system (1.1) characterize, among others, the following behavior of the MLE
$\hat{\beta}$: for almost any $(\gamma,\kappa)$, the system admits a solution
if and only if $\hat{\beta}$ exists with probability approaching one and in
this case, $\|\hat{\beta}\|^{2}/n$ and $\|\hat{\beta}-\beta\|^{2}/n$ both have
finite limits that may be expressed as simple functions of
$(\alpha_{*},\sigma_{*},\lambda_{*})$, and for any feature $j\in[p]$ such that
$\beta_{j}=0$ (i.e., $j$ is a null covariate), the $j$-th coordinate of the
MLE satisfies

$\hat{\beta}_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(0,\sigma_{*}^{2}).$

The proofs in Sur and Candès (2019) are based on approximate message passing
(AMP) techniques; we refer to Berthier et al. (2020); Feng et al. (2022);
Gerbelot and Berthier (2021) and the references therein for recent surveys and
general results. More recently, Zhao et al. (2022) extended the result of Sur
and Candès (2019) from isotropic design to Gaussian covariates with an
arbitrary covariance structure: if now $x_{i}\sim N(\bm{0},\Sigma)$ for some
positive definite $\Sigma$ and
$\lim_{n,p\to+\infty}\beta^{T}\Sigma\beta=\kappa$, null covariates $j\in[p]$
(in the sense that $y_{i}$ is independent of $x_{ij}$ given
$(x_{ik})_{k\in[p]\setminus\\{j\\}}$) of the MLE satisfy

(1.2)
$(n/\Omega_{jj})^{1/2}\hat{\beta}_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(0,\sigma_{*}^{2}),$

where $\sigma_{*}$ is the same solution of (1.1) and $\Omega=\Sigma^{-1}$.
Zhao et al. (2022) also obtained asymptotic normality results for non-null
covariates, that is, features $j\in[p]$ such that $\beta_{j}\neq 0$. The
previous displays can be used to test the null hypothesis $H_{0}:y_{i}$ is
independent of $x_{ij}$ given $(x_{ik})_{k\in[p]\setminus\\{j\\}}$ and develop
the corresponding p-values if $\sigma_{*}$ is known; in this binary logistic
regression model the ProbeFrontier (Sur and Candès, 2019) and SLOE Yadlowsky
et al. (2021) give means to estimate the solutions
$(\alpha_{*},\sigma_{*},\lambda_{*})$ of system (1.1) without the knowledge of
$\gamma$. Mai et al. (2019) studied the performance of Ridge regularized
binary logistic regression in mixture models. Salehi et al. (2019) extended
Sur and Candès (2019) to separable penalty functions. Bellec (2022) derived
asymptotic normality results similar to (1.2) in single-index models including
binary logistic regression without resorting to the system (1.1), showing that
for a null covariate $j\in[p]$ in the unregularized case that

(1.3)
$(n/\Omega_{jj})^{1/2}(\hat{v}/\hat{r})\hat{\beta}_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(0,{1})$

where
$\hat{v}=\frac{1}{n}\sum_{i=1}^{n}\rho^{\prime\prime}(x_{i}^{T}\hat{\beta})-\rho^{\prime\prime}(x_{i}^{T}\hat{\beta})^{2}x_{i}^{T}[\sum_{l=1}^{n}x_{l}\rho^{\prime\prime}(x_{l}^{T}\hat{\beta})x_{l}^{T}]^{-1}x_{i}$
is scalar and so is
$\hat{r}^{2}=\frac{1}{n}\sum_{i=1}^{n}(I\\{y_{i}=1\\}-\rho^{\prime}(x_{i}^{T}\hat{\beta}))^{2}$.
In summary, in this high dimensional binary logistic model,

  1. (i)

The phase transition from Candès and Sur (2020) splits the $(\gamma,\kappa)$
plane into two connected components: in one component the MLE does not exist
with high probability, in the other component the MLE exists and
$\|\Sigma^{1/2}\hat{\beta}\|^{2}$ is bounded with high probability
(boundedness is a consequence of the fact that
$\|\Sigma^{1/2}\hat{\beta}\|^{2}$ or $\|\Sigma^{1/2}(\hat{\beta}-\beta)\|^{2}$
admit finite limits);

  2. (ii)

In the component of the $(\gamma,\kappa)$ plane where the MLE exists, for any
null covariate $j\in[p]$, the asymptotic normality results (1.2)-(1.3) holds.

#### Multiclass classification.

The goal of this paper is to develop a theory for the asymptotic normality of
the multinomial logistic regression MLE (or cross-entropy minimizer) on null
covariates when the number of classes, $K+1$, is greater than 2 and $n,p$ are
of the same order. In other words, we aim to generalize results such as (1.2)
or (1.3) for three or more classes. Classification datasets with 3 or more
classes are ubiquitous in machine learning (MNIST, CIFAR to name a few), which
calls for such multiclass generalizations. In Gaussian mixtures and logistic
models, Thrampoulidis et al. (2020) derived characterizations of the
performance of of least-squares and class-averaging estimators, excluding
cross-entropy minimizers or minimizers of non-linear losses. Loureiro et al.
(2021) extended Sur and Candès (2019); Zhao et al. (2022); Salehi et al.
(2019) to multiclass classification problems in a Gaussian mixture model, and
obtained the fixed-point equations that characterize the performance and
empirical distribution of the minimizer of the cross-entropy loss plus a
convex regularizer. In the same vein as Loureiro et al. (2021), Cornacchia et
al. (2022) studied the limiting fixed-point equations in a multiclass teacher-
student learning model where labels are generated by a noiseless channel with
response $\operatorname*{arg\,min}_{k\in\\{1,...,K\\}}x_{i}^{T}\beta_{k}$
where $\beta_{k}\in\mathbb{R}^{p}$ is unknown for each class $k$. These two
aforementioned works assume a multiclass Gaussian mixture model, which is
different than the normality assumption for $x_{i}$ used in the present paper.
More importantly, these results cannot be readily used for the purpose testing
significant covariates (cf. (1.10) below) since solving the fixed-point
equations require the knowledge of several unknown parameters, including the
limiting spectrum of the mixture covariances and empirical distributions of
the mixture means (cf. for instance Corollary 3 in Loureiro et al. (2021)). In
the following sections, we fill this gap with a new methodology to test the
significance of covariates. This is made possible by developing new asymptotic
normality results for cross-entropy minimizers that generalize (1.3), without
relying on the low-dimensional fixed-point equations.

#### Notation.

Throughout, $I_{p}\in\mathbb{R}^{p\times p}$ is the identity matrix, for a
matrix $A\in\mathbb{R}^{m\times n}$, $A^{T}$ denotes the transpose of $A$,
$A^{\dagger}$ denotes the Moore-Penrose inverse of $A$. If $A$ is psd,
$A^{1/2}$ denotes the unique symmetric square root, i.e., the unique positive
semi-definite matrix such that $(A^{1/2})^{2}=A$. The symbol $\otimes$ denotes
the Kronecker product of matrices. Given two matrices $A\in\mathbb{R}^{n\times
k},B\in\mathbb{R}^{n\times q}$ with the same number or rows,
$(A,B)\in\mathbb{R}^{n\times(k+q)}$ is the matrix obtained by stacking the
columns of $A$ and $B$ horizontally. If $v\in\mathbb{R}^{n}$ is a column
vector with dimension equal to the number of rows in $A$, we construct
$(A,v)\in\mathbb{R}^{n\times(k+1)}$ similarly. We use $\bm{0}_{n}$ and
$\bm{1}_{n}$ to denote the all-zeros vector and all-ones vector in
$\mathbb{R}^{n}$, respectively; we do not bold vectors and matrices other than
$\bm{0}_{n}$ and $\bm{1}_{n}$. We may omit the subscript giving the dimension
if clear from context; e.g., in $I_{K+1}-\frac{\bm{1}\bm{1}^{T}}{K+1}$ the
vector $\bm{1}$ is in $\mathbb{R}^{K+1}$. The Kronecker product between two
matrices is denoted by $\otimes$ and $\operatorname{vec}(M)\in\mathbb{R}^{nd}$
is the vectorization operator applied to a matrix $M\in\mathbb{R}^{n\times
d}$. For an integer $K\geq 2$ and $\alpha\in(0,1)$, the quantile
$\chi^{2}_{K}(\alpha)$ is the unique real number satisfying
$\mathbb{P}(W>\chi^{2}_{K}(\alpha))=\alpha$ where $W$ has a chi-square
distribution with $K$ degrees of freedom. The symbols
${\smash{\xrightarrow{~{}\mathrm{d}~{}}}}$ and
${\smash{\xrightarrow{~{}\mathrm{p}~{}}}}$ denote convergence in distribution
and in probability.

Throughout, _classical asymptotic regime_ refers to the scenario where the
feature dimension $p$ is fixed and the sample size $n$ goes to infinity. In
contrast, the term _high-dimensional regime_ refers to the situation where $n$
and $p$ both tend to infinity with the ratio $p/n$ converging to a limit
smaller than 1.

###  1.1. Multinomial logistic regression

Consider a multinomial logistic regression model with $K+1$ classes. We have
$n$ i.i.d. data samples $\\{(x_{i},\mathsf{y}_{i})\\}_{i=1}^{n}$, where
$x_{i}\in\mathbb{R}^{p}$ is the feature vector and
$\mathsf{y}_{i}=(\mathsf{y}_{i1},...,\mathsf{y}_{i(K+1)})^{T}\in\mathbb{R}^{K+1}$
is the response. Each response $\mathsf{y}_{i}$ is the one-hot encoding of a
single label, i.e., $\mathsf{y}_{i}\in\\{0,1\\}^{K+1}$ with
$\sum_{k=1}^{K+1}\mathsf{y}_{ik}=1$ such that $\mathsf{y}_{ik}=1$ if and only
if the label for $i$-th observation is $k$. A commonly used generative model
for $\mathsf{y}_{i}$ is the multinomial regression model, namely

(1.4)
$\mathbb{P}(\mathsf{y}_{ik}=1|x_{i})=\frac{\exp(x_{i}^{T}\mathsf{B}^{*}e_{k})}{\sum_{k^{\prime}=1}^{K+1}\exp(x_{i}^{T}\mathsf{B}^{*}e_{k^{\prime}})},\quad
k\in\\{1,2,\ldots,K+1\\}$

where $\mathsf{B}^{*}\in\mathbb{R}^{p\times(K+1)}$ is an unknown logistic
model parameter and
$e_{k}\in\mathbb{R}^{K+1},e_{k^{\prime}}\in\mathbb{R}^{K+1}$ are the $k$-th
and $k^{\prime}$-th canonical basis vectors. The MLE for $\mathsf{B}^{*}$ in
the model (1.4) is any solution that minimizes the cross-entropy loss,

(1.5)
$\textstyle\hat{\mathsf{B}}\in\operatorname*{arg\,min}_{\mathsf{B}\in\mathbb{R}^{p\times(K+1)}}\sum_{i=1}^{n}\mathsf{L}_{i}(\mathsf{B}^{T}x_{i}),$

where $\mathsf{L}_{i}:\mathbb{R}^{K+1}\to\mathbb{R}$ is defined as
$\mathsf{L}_{i}(\mathsf{u})=-\sum_{k=1}^{K+1}\mathsf{y}_{ik}\mathsf{u}_{k}+\log\sum_{k^{\prime}=1}^{K+1}\exp(\mathsf{u}_{k^{\prime}}).$
If the solution set in (1.5) is non-empty, we define for each observation
$i\in[n]$ the vector of predicted probabilities
$\hat{\mathsf{p}}_{i}=(\hat{\mathsf{p}}_{i1},...,\hat{\mathsf{p}}_{i(K+1)})^{T}$
with

(1.6) $\hat{\mathsf{p}}_{ik}\stackrel{{\scriptstyle\text{\tiny
def}}}{{\coloneqq}}\mathbb{P}(\hat{\mathsf{y}}_{ik}=1)=\frac{\exp(x_{i}^{T}\hat{\mathsf{B}}e_{k})}{\sum_{k^{\prime}=1}^{K+1}\exp(x_{i}^{T}\hat{\mathsf{B}}e_{k^{\prime}})}\qquad\qquad\text{
for each }k\in\\{1,...,K+1\\}.$

Our results will utilize the gradient and Hessian of $\mathsf{L}_{i}$
evaluated at $\hat{\mathsf{B}}^{T}x_{i}$, denoted by

(1.7) $\mathsf{g}_{i}\stackrel{{\scriptstyle\text{\tiny
def}}}{{\coloneqq}}\nabla\mathsf{L}_{i}(\hat{\mathsf{B}}^{T}x_{i})=-\mathsf{y}_{i}+\hat{\mathsf{p}}_{i},\qquad\mathsf{H}_{i}\stackrel{{\scriptstyle\text{\tiny
def}}}{{\coloneqq}}\nabla^{2}\mathsf{L}_{i}(\hat{\mathsf{B}}^{T}x_{i})=\operatorname{diag}(\hat{\mathsf{p}}_{i})-\hat{\mathsf{p}}_{i}\hat{\mathsf{p}}_{i}^{T}.$

The quantities
$(\hat{\mathsf{B}},\hat{\mathsf{p}}_{i},\mathsf{g}_{i},\mathsf{H}_{i})$ can be
readily computed from the data $\\{(x_{i},\mathsf{y}_{i})\\}_{i=1}^{n}$. To be
specific, the MLE $\hat{\mathsf{B}}$ in (1.5) can be obtained by invoking a
multinomial regression solver (e.g., `sklearn.linear_model.LogisticRegression`
from Pedregosa et al. (2011)), and the quantities
$\hat{\mathsf{p}}_{i},\mathsf{g}_{i},\mathsf{H}_{i}$ can be further computed
from eqs. 1.6 and 1.7 by a few matrix multiplications and application of the
softmax function.

#### Log-odds model and reference class.

The matrix $\mathsf{B}^{*}$ in (1.4) is not identifiable since the conditional
distribution of $\mathsf{y}_{i}|x_{i}$ in the model (1.4) remains unchanged if
we replace columns of $\mathsf{B}^{*}$ by $\mathsf{B}^{*}-b\bm{1}_{K+1}^{T}$
for any $b\in\mathbb{R}^{p}$. In order to obtain an identifiable model, a
classical and natural remedy is to model the log-odds, here with the class
$K+1$ as the reference class:

(1.8)
$\log\frac{\mathbb{P}(\mathsf{y}_{ik}=1|x_{i})}{\mathbb{P}(\mathsf{y}_{i(K+1)}=1|x_{i})}=x_{i}^{T}A^{*}e_{k},\qquad\forall
k\in[K]$

where $e_{k}$ is the $k$-th canonical basis vector of $\mathbb{R}^{K}$, and
$A^{*}\in\mathbb{R}^{p\times K}$ is the unknown parameter. The matrix
$A^{*}\in\mathbb{R}^{p\times K}$ in log-odds model (1.8) is related to
$\mathsf{B}^{*}\in\mathbb{R}^{p\times(K+1)}$ in the model (1.4) by
$A^{*}=\mathsf{B}^{*}(I_{K},-\bm{1}_{K})^{T}$. This log-odds model has two
benefits: First it is identifiable since the unknown matrix $A^{*}$ is
uniquely defined. Second, the matrix $A^{*}$ lends itself well to
interpretation as its $k$-th column represents the contrast coefficient
between class $k$ and the reference class $K+1$.

The MLE $\hat{A}$ of $A^{*}$ in (1.8) is
$\hat{A}=\operatorname*{arg\,min}_{A\in\mathbb{R}^{p\times
K}}\sum_{i=1}^{n}\mathsf{L}_{i}((A,\bm{0}_{p})^{T}x_{i})$. If the solution set
in (1.5) is non-empty, $\hat{A}$ is related to any solution $\hat{\mathsf{B}}$
in (1.5) by $\hat{A}=\hat{\mathsf{B}}(I_{K},-\bm{1}_{K})^{T}$. Equivalently,

(1.9) $\hat{A}_{jk}=\hat{\mathsf{B}}_{jk}-\hat{\mathsf{B}}_{j(K+1)}$

for each $j\in[p]$ and $k\in[K]$.

If there are three classes (i.e. $K+1=3$), this parametrization allows us to
draw scatter plots of realizations of
$\sqrt{n}e_{j}^{T}\hat{A}=(\sqrt{n}\hat{A}_{j1},\sqrt{n}\hat{A}_{j,2})$ as in
Figure 1.

###  1.2. Hypothesis testing for the $j$-th feature and classical asymptotic
normality for MLE

#### Hypothesis testing for the $j$-th feature.

Our goal is to develop a methodology to test the significance of the $j$-th
feature. Specifically, for a desired confidence level $(1-\alpha)\in(0,1)$
(say, $1-\alpha=0.95$) and a given feature $j\in[p]$ of interest, our goal is
to test

(1.10) $H_{0}:\mathsf{y}_{i}\text{ is conditionally independent of
}x_{ij}\text{ given }(x_{ij^{\prime}})_{j^{\prime}\in[p]\setminus\\{j\\}}.$

Namely, we want to test whether the $j$-th variable is independent from the
response given all other explanatory variables
$(x_{ij^{\prime}},j^{\prime}\in[p]\setminus\\{j\\})$. Assuming normally
distributed $x_{i}$ and a multinomial model as in (1.4) or (1.8), it is
equivalent to test

(1.11) $H_{0}:e_{j}^{T}A^{*}=\bm{0}_{K}^{T}\qquad\text{ versus }\qquad
H_{1}:e_{j}^{T}A^{*}\neq\bm{0}_{K}^{T},$

where $e_{j}\in\mathbb{R}^{p}$ is the $j$-th canonical basis vector.

If the MLE $\hat{\mathsf{B}}$ in (1.5) exists in the sense that the solution
set in (1.5) is nonempty, the conjecture that rejecting $H_{0}$ when
$e_{j}^{T}\hat{\mathsf{B}}$ is far from $\bm{0}_{K+1}$ is a reasonable
starting point. The important question, then, is to determine a quantitative
statement for the informal “far from $\bm{0}_{K+1}$”, similarly to (1.2) or
(1.3) in binary logistic regression.

#### Classical theory with $p$ fixed.

If $p$ is fixed and $n\to\infty$ in model (1.8), classical maximum likelihood
theory (Van der Vaart, 1998, Chapter 5) provides the asymptotic distribution
of the MLE $\hat{A}$, which can be further used to test (1.11).

Briefly, if $x$ has the same distribution as any $x_{i}$, the MLE $\hat{A}$ in
the multinomial logistic model is asymptotically normal with

$\sqrt{n}(\operatorname{vec}(\hat{A})-\operatorname{vec}(A^{*})){\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},\mathcal{I}^{-1})\quad\text{
where
}\quad\mathcal{I}=\mathbb{E}[(xx^{T})\otimes(\operatorname{diag}(\pi^{*})-\pi^{*}\pi^{*}{}^{T})]$

is the Fisher information matrix evaluated at the true parameter $A^{*}$,
vec$(\cdot)$ is the usual vectorization operator, and
$\pi^{*}\in\mathbb{R}^{K}$ has random entries

$\pi^{*}_{k}=\exp(x^{T}A^{*}e_{k})\big{/}({1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}A^{*}e_{k^{\prime}})})$

for each $k\in[K]$. In particular, under
$H_{0}:e_{j}^{T}A^{*}=\bm{0}_{K}^{T}$,

(1.12)
$\sqrt{n}\hat{A}{}^{T}e_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},S_{j})$

where $S_{j}=(e_{j}^{T}\otimes I_{K})\mathcal{I}^{-1}(e_{j}\otimes
I_{K})=e_{j}^{T}({\mathrm{cov}}(x))^{-1}e_{j}[\mathbb{E}~{}(\text{diag}(\pi^{*})-\pi^{*}\pi^{*}{}^{T})]^{-1}$.
When (1.12) holds, by the delta method we also have
$\sqrt{n}S_{j}^{-1/2}\hat{A}^{T}e_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},I_{K})$
and

(1.13)
$n\|S_{j}^{-1/2}\hat{A}^{T}e_{j}\|^{2}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}\chi^{2}_{K}.$

where the limiting distribution is chi-square with $K$ degrees of freedom.
This further suggests the size $\alpha$ test that rejects $H_{0}$ when
$T_{n}^{j}(X,Y)>\chi^{2}_{K}(\alpha)$, where
$T_{n}^{j}(X,Y)=n\|S_{j}^{-1/2}\hat{A}^{T}e_{j}\|^{2}$ is the test statistic.
If (1.13) holds, this test is guaranteed to have a type I error converging to
$\alpha$. The p-value of this test is given by

(1.14) $\textstyle\int_{T_{n}^{j}(X,Y)}^{+\infty}f_{\chi^{2}_{K}}(t)dt,$

where $f_{\chi^{2}_{K}}(\cdot)$ is the density of the chi-square distribution
with $K$ degrees of freedom.

As discussed in the introduction, Sur and Candès (2019) showed that in binary
logistic regression, classical normality results for the MLE such as (1.12)
fail in the high-dimensional regime because the variance in (1.12)
underestimates the variability of the MLE even for null covariates; see also
the discussion surrounding (1.2). Our goal is to develop, for classification
problems with $K+1\geq 3$ classes, a theory that correctly characterize the
asymptotic distribution of $\hat{A}^{T}e_{j}$ for a null covariate $j\in[p]$
in the high-dimensional regime.

(a) $(n,p)=(2000,600)$

(b) $(n,p)=(3500,1000)$

(c) $(n,p)=(5000,1500)$

Figure 1. Scatter plot of pairs $(\sqrt{n}\hat{A}_{j1},\sqrt{n}\hat{A}_{j2})$
with $K=2$ over 1000 repetitions. The blue ellipsoid is the boundary of the
$95\%$ confidence set for $\sqrt{n}\hat{A}^{T}e_{j}$ under $H_{0}$ from the
classical MLE theory (1.12)-(1.13) based on the Fisher information, the dashed
red ellipsoids are the boundaries of the $95\%$ confidence set for
$\sqrt{n}\hat{A}^{T}e_{j}$ under $H_{0}$ from this paper (cf. (2.3) below).
Each of the 1000 repetition gives a slightly different dashed ellipsoid. The
solid red ellipsoid is the average of these 1000 dashed ellipsoids. Each row
of $X$ is i.i.d. sampled from $N(\bm{0},\Sigma)$ with
$\Sigma=(0.5^{|i-j|})_{p\times p}$. The first $\lceil p/4\rceil$ rows of
$A^{*}$ are i.i.d. sampled from $N(\bm{0},I_{K})$ while other rows are set to
zeros. We further normalize $A^{*}$ such that $A^{*}{}^{T}\Sigma A^{*}=I_{K}$.
The last coordinate $j=p$ is used as the null coordinate.

We present first some motivating simulations that demonstrate the failure of
classical normal approximation (1.12) in finite samples. These simulations are
conducted for various configurations of $(n,p)$ with $K+1=3$ classes. We fix
the true parameter $A^{*}$ and obtain 1000 realizations of
$(\hat{A}_{j1},\hat{A}_{j2})$ by independently resampling the data
$\\{(x_{i},\mathsf{y}_{i})\\}_{i=1}^{n}$ 1000 times. If the result (1.12)
holds, then
$\mathbb{P}(\sqrt{n}\hat{A}^{T}e_{j}\in\mathcal{C}_{\alpha}^{j})\to 1-\alpha$,
where
$\mathcal{C}_{\alpha}^{j}=\\{u\in\mathbb{R}^{K}:\|S_{j}^{-1/2}u\|\leq\chi^{2}_{K}(\alpha)\\}$.
Figure 1 displays scatter plots of $\sqrt{n}(\hat{A}_{j1},\hat{A}_{j2})$ along
with the boundary of 95% confidence set $\mathcal{C}_{\alpha}^{j}$ with
$\alpha=0.05$. We observe that, across the three different configurations of
$(n,p)$, the 95% confidence sets from our theory (Theorem 2.2 presented in
next section) cover around 95% of the realizations, while the set
$\mathcal{C}_{\alpha}^{j}$ from classical theory only covers approximately
$30\%$ of the points, which is significantly lower than the desired coverage
rate of $95\%$. Intuitively and by analogy with results in binary
classification (Sur and Candès, 2019), this is because the classical theory
(1.12) underestimates the variation of the MLE in the high-dimensional regime.
Motivated by this failure of classical MLE theory and the results in binary
classification (Sur and Candès, 2019, among others), the goal of this paper is
to develop a theory for multinomial logistic regression that achieves the
following objectives:

  * •

Establish asymptotic normality of the multinomial MLE $\hat{A}^{T}e_{j}$ for
null covariates as $n,p\to+\infty$ simultaneously with a finite limit for
$n/p$.

  * •

Develop a valid methodology for hypothesis testing of (1.10) in this regime,
i.e., testing for the presence of an effect of a feature $j\in[p]$ on the
multiclass response.

The contribution of this paper is two-fold: (i) For a null covariate
$j\in[p]$, we establish asymptotic normality results for $\hat{A}^{T}e_{j}$
that are valid in the high-dimensional regime where $n$ and $p$ have the same
order; (ii) we propose a user-friendly test for assessing the significance of
a feature in multiclass classification problems.

##  2\. Main result: asymptotic normality of $\hat{\mathsf{B}}{}^{T}e_{j}$ and
$\hat{A}^{T}e_{j}$ on null covariates

In this section, we present the main theoretical results of our work and
discuss their significance. We work under the following assumptions.

######  Assumption 2.1.

For constants $\delta>1$, assume that $n,p\to\infty$ with
$p/n\leq\delta^{-1}$, and that the design matrix $X\in\mathbb{R}^{n\times p}$
has $n$ i.i.d. rows $(x_{i})_{i\in[n]}\sim N(\bm{0},\Sigma)$ for some
invertible $\Sigma\in\mathbb{R}^{p\times p}$. The observations
$(x_{i},\mathsf{y}_{i})_{i\in[n]}$ are i.i.d. and each $\mathsf{y}_{i}$ is of
the form $\mathsf{y}_{i}=f(U_{i},x_{i}^{T}\mathsf{B}^{*})$ for some
deterministic function $f$, deterministic matrix
$\mathsf{B}^{*}\in\mathbb{R}^{p\times(K+1)}$ such that
$\mathsf{B}^{*}\bm{1}_{K+1}=\bm{0}_{p}$, and latent random variable $U_{i}$
independent of $x_{i}$.

######  Assumption 2.2 (One-hot encoding).

The response matrix $Y$ is in $\mathbb{R}^{n\times(K+1)}$. Its $i$-th row
$\mathsf{y}_{i}$ is a one-hot encoded vector, that is, valued in
$\\{0,1\\}^{K+1}$ with $\sum_{k=1}^{K+1}\mathsf{y}_{ik}=1$ for each $i\in[n]$.

The model $\mathsf{y}_{i}=f(U_{i},x_{i}^{T}\mathsf{B}^{*})$ for some
deterministic $f$ and $\mathsf{B}^{*}$ and latent random variable $U_{i}$ in
Assumption 2.1 is more general than a specific generative model such as the
multinomial logistic conditional probabilities in (1.4), as broad choices for
$f$ are allowed. In words, the model
$\mathsf{y}_{i}=f(U_{i},x_{i}^{T}\mathsf{B}^{*})$ with
$\mathsf{B}^{*}\bm{1}_{K+1}=0_{p}$ means that $\mathsf{y}_{i}$ only depends on
$x_{i}$ through a $K$ dimensional projection of $x_{i}$ (the projection on the
row-space of $\mathsf{B}^{*}$). The assumption $p/n\leq\delta^{-1}$ is more
general than assuming a fixed limit for the ratio $p/n$; this allows us to
cover low-dimensional settings satisfying $p/n\to 0$ as well.

The following assumption requires the labels to be “balanced”: we observe each
class at least $\gamma n$ times for some constant $\gamma>0$. If
$(\mathsf{y}_{i})_{i\in[n]}$ are i.i.d. as in Assumption 2.1 with distribution
independent of $n,p$, by the law of large numbers this assumption is
equivalent to $\min_{k\in[K+1]}\mathbb{P}(\mathsf{y}_{ik}=1)>0$.

######  Assumption 2.3.

There exits a constant $\gamma\in(0,\tfrac{1}{K+1}]$, such that for each
$k\in[K+1]$, with probability approaching one at least $\gamma n$ observations
$i\in[n]$ are such that $\mathsf{y}_{ik}=1$. In other words,
$\mathbb{P}(\sum_{i=1}^{n}I(\mathsf{y}_{ik}=1)\geq\gamma n)\to 1$ for each
$k\in[K+1]$.

As discussed in item list (i) on page 2, in binary logistic regression, Candès
and Sur (2020); Sur and Candès (2019) show that the plane
$(\frac{p}{n},\|\Sigma^{1/2}\beta^{*}\|)$ is split by a smooth curve into two
connected open components: in one component the MLE does not exist with high
probability, while in the other component, with high probability the MLE
exists and is bounded in the sense that
$\|\Sigma^{1/2}\hat{\beta}\|^{2}<\tau^{\prime}$ or equivalently
$\frac{1}{n}\|X\hat{\beta}\|^{2}<\tau$ for constants $\tau,\tau^{\prime}$
independent of $n,p$. The next assumption requires the typical situation of
the latter component, in the current multiclass setting: $\hat{\mathsf{B}}$ in
(1.5) exists in the sense that the minimization problem has solutions, and at
least one solution is bounded.

######  Assumption 2.4.

Assume $\mathbb{P}(\hat{\mathsf{B}}\text{~{}exists}\text{ and
}\|X\hat{\mathsf{B}}(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})\|_{F}^{2}\leq
n\tau)\to 1$ as $n,p\to+\infty$ for some large enough constant $\tau$.

Note that the validity of Assumption 2.4 can be assessed using the data at
hand; if a multinomial regression solver (e.g.
`sklearn.linear_model.LogisticRegression`) converges and
$\tfrac{1}{n}\|X\hat{\mathsf{B}}(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})\|_{F}^{2}$
is no larger than a predetermined large constant $\tau$, then we know
Assumption 2.4 holds. Otherwise the algorithm does not converge or produces an
unbounded estimate: we know Assumption 2.4 fails to hold and we need collect
more data.

Our first main result, Theorem 2.1, provides the asymptotic distribution of
$\hat{\mathsf{B}}^{T}e_{j}$ where $j\in[p]$ is a null covariate, where
$\hat{\mathsf{B}}$ is any minimizer $\hat{\mathsf{B}}$ of (1.5). Throughout,
we denote by $\Omega$ the precision matrix defined as $\Omega=\Sigma^{-1}$.

######  Theorem 2.1.

Let Assumptions 2.1, 2.2, 2.3 and 2.4 be fulfilled. Then for any $j\in[p]$
such that $H_{0}$ in (1.10) holds, and any minimizer $\hat{\mathsf{B}}$ of
(1.5), we have

(2.1)
$\underbrace{\vphantom{\sum_{i=1}^{n}}\sqrt{\frac{n}{\Omega_{jj}}}}_{\text{scalar}}\Bigl{(}\underbrace{\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}(\mathsf{y}_{i}-\hat{\mathsf{p}}_{i})(\mathsf{y}_{i}-\hat{\mathsf{p}}_{i})^{T}\Bigr{)}^{1/2}}_{\text{square
root pseudo-inverse
}\mathbb{R}^{(K+1)\times(K+1)}}\Bigr{)}^{\dagger}\underbrace{\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}\Bigr{)}}_{\mathbb{R}^{(K+1)\times(K+1)}}\underbrace{\vphantom{\sum_{i=1}^{n}}\hat{\mathsf{B}}^{T}e_{j}}_{\mathbb{R}^{K+1}}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N\Bigl{(}\bm{0},\underbrace{\vphantom{\sum_{i=1}^{n}}I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1}}_{\text{cov.
}\mathbb{R}^{(K+1)\times(K+1)}}\Bigr{)},$

where $\mathsf{V}_{i}=\mathsf{H}_{i}-(\mathsf{H}_{i}\otimes
x_{i}^{T})[\sum_{l=1}^{n}\mathsf{H}_{l}\otimes(x_{l}x_{l}^{T})]^{\dagger}(\mathsf{H}_{i}\otimes
x_{i})$.

The proof of Theorem 2.1 is given in Supplementary Section S3. Theorem 2.1
establishes that under $H_{0}$, $\hat{\mathsf{B}}^{T}e_{j}$ converges to a
singular multivariate Gaussian distribution in $\mathbb{R}^{K+1}$. In (2.1),
the two matrices
$\frac{1}{n}\sum_{i=1}^{n}(\mathsf{y}_{i}-\hat{\mathsf{p}}_{i})(\mathsf{y}_{i}-\hat{\mathsf{p}}_{i})^{T}$
and $\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}$ are symmetric with kernel being
the linear span of $\bm{1}_{K+1}$, and similarly, if a solution exists, we may
replace $\hat{\mathsf{B}}$ by
$\hat{\mathsf{B}}(I_{K+1}-\frac{\bm{1}\bm{1}^{T}}{K+1})$ which is also
solution in (1.5). In this case, all matrix-matrix and matrix-vector
multiplications, matrix square root and pseudo-inverse in (2.1) happen with
row-space and column space contained in the orthogonal component of
$\bm{1}_{K+1}$, so that the limiting Gaussian distribution in
$\mathbb{R}^{K+1}$ is also supported on this $K$-dimensional subspace.

Since the distribution of the left-hand side of (2.1) is asymptotically
pivotal for all null covariates $j\in[p]$, Theorem 2.1 opens the door of
statistical inference for multinomial logistic regression in high-dimensional
settings. By construction, the multinomial logistic estimate
$\hat{A}\in\mathbb{R}^{p\times K}$ in (1.9) ensures $(\hat{A},\bm{0}_{p})$ is
a minimizer of (1.5). Therefore, we can deduce the following theorem from
Theorem 2.1.

######  Theorem 2.2.

Define the matrix $R=(I_{K},\bm{0}_{K})^{T}\in\mathbb{R}^{(K+1)\times K}$
using block matrix notation. Let Assumptions 2.1, 2.2, 2.3 and 2.4 be
fulfilled. For $\hat{A}$ in (1.9) and any $j\in[p]$ such that $H_{0}$ in
(1.10) holds,

(2.2)
$\underbrace{\vphantom{\sum_{i=1}^{n}}\Bigl{(}I_{K}+\frac{\bm{1}_{K}\bm{1}_{K}^{T}}{\sqrt{K+1}+1}\Bigr{)}R^{T}}_{\text{matrix
}\mathbb{R}^{K\times(K+1)}}\underbrace{\vphantom{\sum_{i=1}^{n}}\sqrt{\frac{n}{\Omega_{jj}}}}_{\text{scalar}}\Bigl{(}\underbrace{\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{g}_{i}\mathsf{g}_{i}^{T}\Bigr{)}^{1/2}}_{\text{matrix
}\mathbb{R}^{(K+1)\times(K+1)}}\Bigr{)}^{\dagger}\underbrace{\vphantom{\sum_{i=1}^{n}}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}R\Bigr{)}}_{\mathbb{R}^{(K+1)\times
K}}\underbrace{\vphantom{\sum_{i=1}^{n}}\hat{A}^{T}e_{j}}_{\mathbb{R}^{K}}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0}_{K},I_{K})$

where $\mathsf{g}_{i}$ is defined in (1.7) and $\mathsf{V}_{i}$ is defined in
Theorem 2.1. Furthermore, for the same $j\in[p]$,

(2.3) $\mathcal{T}_{n}^{j}(X,Y)\stackrel{{\scriptstyle\text{\tiny
def}}}{{\coloneqq}}\frac{n}{\Omega_{jj}}\Bigl{\|}\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{g}_{i}\mathsf{g}_{i}^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}\Bigr{)}R\hat{A}^{T}e_{j}\Bigr{\|}^{2}~{}\text{
satisfies
}~{}\mathcal{T}_{n}^{j}(X,Y){\smash{\xrightarrow{~{}\mathrm{d}~{}}}}\chi_{K}^{2}.$

Theorem 2.2 is proved in Supplementary Section S4. To the best of our
knowledge, Theorem 2.2 is the first result that characterizes the distribution
of null MLE coordinate $\hat{A}^{T}e_{j}$ in high-dimensional multinomial
logistic regression with 3 or more classes. It is worth mentioning that the
quantities $(\mathsf{g}_{i},\mathsf{V}_{i},\hat{A})$ used in Theorem 2.2 can
be readily computed from the data $(X,Y)$. Therefore, Theorem 2.2 lets us test
the significance of a specific feature: for testing $H_{0}$, this theorem
suggests the test statistic $\mathcal{T}_{n}^{j}(X,Y)$ in (2.3) and the
rejection region $\mathcal{E}_{\alpha}^{j}\stackrel{{\scriptstyle\text{\tiny
def}}}{{\coloneqq}}\bigl{\\{}(X,Y):\mathcal{T}_{n}^{j}(X,Y)\geq\chi^{2}_{K}(\alpha)\bigr{\\}}.$
Under the null hypothesis $H_{0}$ in (1.10), Theorem 2.2 guarantees
$\mathbb{P}\big{(}(X,Y)\in\mathcal{E}_{\alpha}^{j}\big{)}\to\alpha.$ In other
words, the test that rejects $H_{0}$ if $(X,Y)\in\mathcal{E}_{\alpha}^{j}$ has
type I error converging to $\alpha$. The p-value of this test is

(2.4)
$\textstyle\text{p-value}=\int_{\mathcal{T}_{n}^{j}(X,Y)}^{+\infty}f_{\chi^{2}_{K}}(t)dt,$

where $f_{\chi^{2}_{K}}(\cdot)$ is the density of the chi-square distribution
with $K$ degrees of freedom.

Unknown $\Omega_{jj}=e_{j}^{T}\Sigma^{-1}e_{j}$. If $\Sigma$ is unknown, we
describe a consistent estimate of the quantity $\Omega_{jj}$ appearing in
(2.1), (2.2), and (2.3). Under the Gaussian Assumption 2.1, the quantity
$\Omega_{jj}$ is the reciprocal of the conditional variance
$\operatorname{Var}(x_{ij}|x_{i,-j})$, which is also the noise variance in the
linear model of regressing $Xe_{j}$ onto $X_{-j}$ (the submatrix of $X$
excluding the $j$-th column). According to standard results in linear models,
we have
$\Omega_{jj}\|[I_{n}-X_{-j}(X_{-j}^{T}X_{-j})^{-1}X_{-j}^{T}]Xe_{j}\|^{2}\sim\chi^{2}_{n-p+1}.$
Since $\chi^{2}_{n-p+1}/(n-p+1)\to 1$ almost surely by the strong law of large
numbers,

(2.5)
$\hat{\Omega}_{jj}=(n-p+1)\big{/}\|[I_{n}-X_{-j}(X_{-j}^{T}X_{-j})^{-1}X_{-j}^{T}]Xe_{j}\|^{2}$

is a consistent estimator of $\Omega_{jj}$. Therefore, the previous asymptotic
results in Theorems 2.1 and 2.2 still hold by Slutsky’s theorem if we replace
$\Omega_{jj}$ by the estimate $\hat{\Omega}_{jj}$ in (2.5).

##  3\. Numerical experiments

This section presents simulations to examine finite sample properties of the
above results and methods.

Simulation settings.

We set $p=1000$ and consider different combinations of $(n,K)$. The covariance
matrix $\Sigma$ is specified to be the correlation matrix of an AR(1) model
with parameter $\rho=0.5$, that is, $\Sigma=(0.5^{|i-j|})_{p\times p}$. We
generate the regression coefficients $A^{*}\in\mathbb{R}^{p\times K}$ once and
for all as follows: sample $A_{0}\in\mathbb{R}^{p\times K}$ with first $\lceil
p/4\rceil$ rows being i.i.d. $N(\bm{0},I_{K})$, and set the remaining rows to
0. We then scale the coefficients by defining $A^{*}=A_{0}(A_{0}^{T}\Sigma
A_{0})^{-1/2}$ so that $A^{*}{}^{T}\Sigma A^{*}=I_{K}$. With this
construction, the $p$-th variable is always a null covariate , and we use this
null coordinate $j=p$ to demonstrate the effectiveness of our theoretical
results presented in Theorem 2.2 and the suggested test for testing $H_{0}$ as
described in (1.10). Using the above settings, we generate the design matrix
$X\in\mathbb{R}^{n\times p}$ from $N(0,\Sigma)$, and then simulate the labels
from a multinomial logistic model as given in (1.8), using the coefficients
$A^{*}\in\mathbb{R}^{p\times K}$. For each simulation setting, we perform
5,000 repetitions.

Assessment of $\chi^{2}$ approximations. To assess the $\chi^{2}$
approximation (2.3) from this paper and that of the classical theory (1.13),
we compute the two $\chi^{2}_{K}$ test statistics for each sample
$(x_{i},y_{i})_{i=1}^{n}$. Figure 2 shows the empirical quantiles of the two
statistics versus the $\chi^{2}_{K}$ distribution quantiles. The results
demonstrate that the quantiles (in blue) from our high-dimensional theory
closely match the 45-degree line (in red), whereas the quantiles (in orange)
from the classical theory significantly deviate from the 45-degree line. These
findings highlight the accuracy of our proposed $\chi^{2}$ approximation (2.3)
over the classical result (1.13) when $p$ is not sufficiently small compared
to $n$.

(a) $(n,K)=(4000,2)$

(b) $(n,K)=(5000,3)$

(c) $(n,K)=(6000,4)$

Figure 2. Q-Q plots of the test statistic in the left-hand side of (1.13) (in
orange) and in the left-hand side of (2.3) (in blue) for different $(n,K)$ and
$p=1000$.

(a) $(n,K)=(4000,2)$

(b) $(n,K)=(5000,3)$

(c) $(n,K)=(6000,4)$

Figure 3. Histogram for p-values of the classical test (1.14) (in orange) and
of the proposed test (2.4) (in blue) under $H_{0}$ in simulated data with
different $(n,K)$ and $p=1000$.

Uniformity of null p-values. Recall that the p-value from the classical test
(1.13) is given by (1.14), while the p-value from this paper taking into
account high-dimensionality is given by (2.4). Figure 3 displays the
histograms of these two sets of p-values out of 5000 repetitions. The results
in Figure 3 show that the p-values obtained from the classical test deviate
significantly from the uniform distribution, with a severe inflation in the
lower tail. This indicates that the classical test tends to produce large type
I errors due to the excess of p-values close to 0. In contrast, the p-values
proposed in this paper exhibit a uniform distribution, further confirming the
effectiveness and applicability of the theory in Theorem 2.2 for controlling
type I error when testing for null covariates with (1.10).

Unknown $\Omega_{jj}$. In the situation where the covariance matrix $\Sigma$
is unknown, we can estimate the diagonal element
$\Omega_{jj}=e_{j}^{T}\Sigma^{-1}e_{j}$ by ${\smash{\hat{\Omega}_{jj}}}$
defined in (2.5). To evaluate the accuracy of the normal and chi-square
approximations and the associated test with $\Omega_{jj}$ replaced by
${\smash{\hat{\Omega}_{jj}}}$, we conduct simulations similar to those in
Figures 2 and 3, but we replace $\Omega_{jj}$ with its estimate
${\smash{\hat{\Omega}_{jj}}}$. The results are presented in Figure S1. The
plots are visually indistinguishable from the plots using $\Omega_{jj}$. These
confirm that the chi-square approximation and the associated test using
${\smash{\hat{\Omega}_{jj}}}$ are accurate.

Non-Gaussian covariates and unknown $\Omega_{jj}$. Although our theory assumes
Gaussian covariates, we expect that the same results hold for other
distributions with sufficiently light tails. To illustrate this point, we
consider the following two types of non-Gaussian covariates: (i) The design
matrix $X$ has i.i.d. Rademacher entries, i.e., $\mathbb{P}(x_{ij}=\pm
1)=\frac{1}{2}$, (ii) Each $x_{ij}$ takes on values 0, 1 and 2 with
respectively probabilities $a_{j}^{2},2a_{j}(1-a_{j})$, and $(1-a_{j})^{2}$,
where $a_{j}$ varies in $[0.25,0.75]$. Each columns of $X$ are then centered
and normalized to have $0$ mean and unit variance. This generation of non-
Gaussian covariates is adopted from single-nucleotide poly-morphisms (SNPs)
example in Sur and Candès (2019). For these two types of non-Gaussian
covariates, we further rescale the feature vectors to ensure that $x_{i}$ has
the same covariance as in the Gaussian case at the beginning of Section 3,
that is $\Sigma=(0.5^{|i-j|})_{p\times p}$. We present the Q-Q plots in Figure
S2 using the same settings as in Figure S1, with the only difference being
that the covariates in Figure S2 are non-Gaussian distributed. The Q-Q plots
of $\mathcal{T}_{n}^{j}(X,Y)$ in (2.3) plotted in Figure S2 still closely
match the diagonal line. These empirical successes suggest that the normal and
$\chi^{2}_{K}$ approximations (2.1)-(2.3) apply to a wider range of covariate
distributions beyond normally distributed data.

##  4\. Discussion and future work

Multinomial logistic regression estimates and their p-values are ubiquitous
throughout the sciences for analyzing the significance of explanatory
variables on multiclass responses. Following the seminal work of Sur and
Candès (2019) in binary logistic regression, this paper develops the first
valid tests and p-values for multinomial logistic estimates when $p$ and $n$
are of the same order. For 3 or more classes, this methodology and the
corresponding asymptotic normality results in Theorems 2.1 and 2.2 are novel
and provide new understanding of multinomial logistic estimates (also known as
cross-entropy minimizers) in high-dimensions. We expect similar asymptotic
normality and chi-square results to be within reach for loss functions
different than the cross-entropy or a different model for the response
$\mathsf{y}_{i}$; for instance Section S1 provides an extension to the
$q$-repeated measurements model, where $q$ responses are observed for each
feature vector $x_{i}$.

Let us point a few follow-up research directions that we leave open for future
work. A first open problem regards extensions of our methodology to confidence
sets for $e_{j}^{T}\mathsf{B}^{*}$ when $H_{0}$ in (1.10) is violated for the
$j$-th covariate. This would require more stringent assumptions on the
generative model than Assumption 2.1 as $\mathsf{B}^{*}$ there is not
identifiable (e.g., modification of both $\mathsf{B}^{*}$ and $f(\cdot,\cdot)$
in Assumption 2.1 is possible without changing $\mathsf{y}_{i}$). A second
open problem is to relate this paper’s theory to the fixed-point equations and
limiting Gaussian model obtained in multiclass models, e.g., Loureiro et al.
(2021). While it may be straightforward to obtain the limit of
$\frac{1}{n}\sum_{i=1}^{n}\mathsf{g}_{i}\mathsf{g}_{i}^{T}$ and of the
empirical distribution of the rows of ${\smash{\hat{A}}}$ in this context
(e.g., using Corollary 3 in Loureiro et al. (2021)), the relationship between
the fixed-point equations and the matrix
$\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}$ appearing in (2.1) is unclear and
not explained by typical results from this literature. A third open problem is
to characterize the exact phase transition below which the multinomial
logistic MLE exists and is bounded with high-probability (Assumption 2.4);
while this is settled for two classes (Candès and Sur, 2020) and preliminary
results are available for 3 or more classes (Loureiro et al., 2021; Kini and
Thrampoulidis, 2021), a complete understanding of this phase transition is
currently lacking. A last interesting open problem is to prove that our theory
extend to non-Gaussian data, as observed in simulations. This challenging
problem is often referred to as “universality” and has received intense
attention recently (Montanari and Saeed, 2022; Gerace et al., 2022; Pesce et
al., 2023; Dandi et al., 2023), showing that in several settings of interest
(although none exactly the one considered here), the asymptotic behavior of
the minimizers is unchanged if the distribution of the covariates is modified
from normal to another distribution with the same covariance.

## References

  * Bellec [2022] Pierre C Bellec.  Observable adjustments in single-index models for regularized m-estimators.  _arXiv preprint arXiv:2204.06990_ , 2022.
  * Bellec and Zhang [2021] Pierre C. Bellec and Cun-Hui Zhang.  Second-order Stein: SURE for SURE and other applications in high-dimensional inference.  _Ann. Statist._ , 49(4):1864–1903, 2021.  ISSN 0090-5364.  URL .
  * Berthier et al. [2020] Raphael Berthier, Andrea Montanari, and Phan-Minh Nguyen.  State evolution for approximate message passing with non-separable functions.  _Information and Inference: A Journal of the IMA_ , 9(1):33–79, 2020.
  * Candès and Sur [2020] Emmanuel J. Candès and Pragya Sur.  The phase transition for the existence of the maximum likelihood estimate in high-dimensional logistic regression.  _The Annals of Statistics_ , 48(1):27 – 42, 2020\.  doi: 10.1214/18-AOS1789.  URL .
  * Cornacchia et al. [2022] Elisabetta Cornacchia, Francesca Mignacco, Rodrigo Veiga, Cédric Gerbelot, Bruno Loureiro, and Lenka Zdeborová.  Learning curves for the multi-class teacher-student perceptron.  _arXiv preprint arXiv:2203.12094_ , 2022.
  * Cramer [2002] Jan Salomon Cramer.  The origins of logistic regression.  _Tinbergen Institute Working Paper_ , 2002.
  * Dandi et al. [2023] Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová.  Universality laws for gaussian mixtures in generalized linear models.  _arXiv preprint arXiv:2302.08933_ , 2023.
  * Davidson and Szarek [2001] Kenneth R Davidson and Stanislaw J Szarek.  Local operator theory, random matrices and banach spaces.  _Handbook of the geometry of Banach spaces_ , 1(317-366):131, 2001.
  * Feng et al. [2022] Oliver Y Feng, Ramji Venkataramanan, Cynthia Rush, Richard J Samworth, et al.  A unifying tutorial on approximate message passing.  _Foundations and Trends® in Machine Learning_ , 15(4):335–536, 2022.
  * Gerace et al. [2022] Federica Gerace, Florent Krzakala, Bruno Loureiro, Ludovic Stephan, and Lenka Zdeborová.  Gaussian universality of linear classifiers with random labels in high-dimension.  _arXiv preprint arXiv:2205.13303_ , 2022.
  * Gerbelot and Berthier [2021] Cédric Gerbelot and Raphaël Berthier.  Graph-based approximate message passing iterations.  _arXiv preprint arXiv:2109.11905_ , 2021.
  * Kini and Thrampoulidis [2021] Ganesh Ramachandra Kini and Christos Thrampoulidis.  Phase transitions for one-vs-one and one-vs-all linear separability in multiclass gaussian mixtures.  In _ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 4020–4024. IEEE, 2021.
  * Loureiro et al. [2021] Bruno Loureiro, Gabriele Sicuro, Cédric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborová.  Learning gaussian mixtures with generalized linear models: Precise asymptotics in high-dimensions.  _Advances in Neural Information Processing Systems_ , 34:10144–10157, 2021.
  * Mai et al. [2019] Xiaoyi Mai, Zhenyu Liao, and Romain Couillet.  A large scale analysis of logistic regression: Asymptotic performance and new insights.  In _ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 3357–3361. IEEE, 2019.
  * Montanari and Saeed [2022] Andrea Montanari and Basil N Saeed.  Universality of empirical risk minimization.  In _Conference on Learning Theory_ , pages 4310–4312. PMLR, 2022\.
  * Pedregosa et al. [2011] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay.  Scikit-learn: Machine learning in Python.  _Journal of Machine Learning Research_ , 12:2825–2830, 2011\.
  * Pesce et al. [2023] Luca Pesce, Florent Krzakala, Bruno Loureiro, and Ludovic Stephan.  Are gaussian data all you need? extents and limits of universality in high-dimensional generalized linear estimation.  _arXiv preprint arXiv:2302.08923_ , 2023.
  * Salehi et al. [2019] Fariborz Salehi, Ehsan Abbasi, and Babak Hassibi.  The impact of regularization on high-dimensional logistic regression.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Sur and Candès [2019] Pragya Sur and Emmanuel J. Candès.  A modern maximum-likelihood theory for high-dimensional logistic regression.  _Proc. Natl. Acad. Sci. USA_ , 116(29):14516–14525, 2019.  ISSN 0027-8424.  doi: 10.1073/pnas.1810420116.  URL .
  * Thrampoulidis et al. [2020] Christos Thrampoulidis, Samet Oymak, and Mahdi Soltanolkotabi.  Theoretical insights into multiclass classification: A high-dimensional asymptotic view.  _Advances in Neural Information Processing Systems_ , 33:8907–8920, 2020.
  * Van der Vaart [1998] Aad W Van der Vaart.  _Asymptotic statistics_ , volume 3.  Cambridge university press, 1998.
  * van Hemmen and Ando [1980] J Leo van Hemmen and Tsuneya Ando.  An inequality for trace ideals.  _Communications in Mathematical Physics_ , 76:143–148, 1980\.
  * Yadlowsky et al. [2021] Steve Yadlowsky, Taedong Yun, Cory Y McLean, and Alexander D’Amour.  Sloe: A faster method for statistical inference in high-dimensional logistic regression.  _Advances in Neural Information Processing Systems_ , 34:29517–29528, 2021.
  * Zhao et al. [2022] Qian Zhao, Pragya Sur, and Emmanuel J Candes.  The asymptotic distribution of the mle in high-dimensional logistic models: Arbitrary covariance.  _Bernoulli_ , 28(3):1835–1861, 2022.
  * Zhu and Hastie [2004] Ji Zhu and Trevor Hastie.  Classification of gene microarrays by penalized logistic regression.  _Biostatistics_ , 5(3):427–443, 2004.

Supplementary Material of “Multinomial Logistic Regression: Asymptotic
Normality on Null Covariates in High-Dimensions”
Let us define some standard notation that will be used in the rest of this
supplement. For a vector $v\in\mathbb{R}^{n}$, let
$\|v\|_{\infty}=\max_{i\in[n]}|v_{i}|$ denote the infinity norm of vector $v$.
If $A$ is symmetric, we define $\lambda_{\min}(A)$ and $\|A\|_{op}$ as the
minimal and maximal eigenvalues of $A$, respectively. For two symmetric
matrices $A,B$ of the same size, we write $A\preceq B$ if and only if $B-A$ is
positive semi-definite.

## Diagram: Organization of the proofs

The following diagram summarizes the different theorems and lemmas, and the
relationships between them.

Theorem 2.2 Asymptotic normality for $\hat{A}^{T}e_{j}$ on null covariates,
where $\hat{A}\in\mathbb{R}^{p\times K}$ is the multinomial logistic MLE with
class $K+1$ fixed as the reference class (see (1.9)). Theorem 2.1 Asymptotic
normality for $\hat{\mathsf{B}}^{T}e_{j}$ on null covariates, where
$\hat{\mathsf{B}}\in\mathbb{R}^{p\times(K+1)}$ is the multinomial logistic MLE
in (1.5). Theorem S3.1 Asymptotic normality for $\hat{B}^{T}e_{j}$ on null
covariates, where $\hat{B}\in\mathbb{R}^{p\times K}$ is the multinomial
logistic MLE using the parameter space from Section S3.1. The proof uses that
the conditions in Theorem S5.1 on the loss function are satisfied by the
cross-entropy. $K$-dimensional orthogonal parametrization defined by the
matrix the $Q$ Section S3.1 defines the matrix $Q\in\mathbb{R}^{(K+1)\times
K}$ and discusses a convenient parametrization of the model isometric to the
subspace orthogonal to $\bm{1}_{K+1}$. Control of $\mathsf{g}_{i}$ and
$\mathsf{H}_{i}$ for the cross-entropy loss Lemmas S6.2 and S6.1 give
deterministic arguments to control the gradients and Hessians of the cross-
entropy loss. Lemma S6.4 controls the Hessian of the cross-entropy loss at the
minimizer, in a specific high-probability event. Lemma S6.3 defines this high-
probability event. Theorem S5.1 Asymptotic normality on null covariates for
general loss functions, $\Sigma\neq I_{p}$. Deduced from Theorem S5.2 by
rotational invariance. Theorem S5.2 Asymptotic normality on null covariates
for general loss functions, $\Sigma=I_{p}$. Lemma S5.3 Normal and $\chi^{2}$
approximations for random variables defined as a differentiable function of
standard normal vectors. Lemma S5.5 Lemma S5.5 computes the derivatives of the
minimizer with respect to $X$, used in the proof of Theorem S5.2.

##  S1. Extension: $q$ repeated measurements

Let integer $q\geq 1$ be a constant independent of $n,p$. Our results readily
extend if $q$ labels are observed for each observed feature vector $x_{i}$,
and the corresponding $q$ one-hot encoded vectors are averaged into
$\mathsf{y}_{i}\in\\{0,\frac{1}{q},\frac{2}{q},...,1\\}^{K+1}$. Concretely,
for each observation $i\in[n]$, $q$ i.i.d. labels $(Y_{i}^{m})_{m\in[q]}$ are
observed with each $Y_{i}^{m}\in\\{0,1\\}^{K+1}$ one-hot encoded and
$\mathsf{y}_{ik}=\frac{1}{q}\sum_{m=1}^{q}Y_{ik}^{m}$, for instance in a
repeated multinomial regression model with $\mathbb{P}(Y_{ik}^{m}=1|x_{i})$
equal to right-hand side of (1.4). In this case where $(Y_{i}^{m})_{m\in[q]}$
are i.i.d., Assumption 2.3 is satisfied by the law of large numbers if
$\min_{k\in[K+1]}\mathbb{P}(Y_{ik}^{m}=1)>0$ since $q$ is constant. For this
$q$ repeated measurements model, the negative log-likelihood function of a
parameter $\mathsf{B}\in\mathbb{R}^{p\times(K+1)}$ is

$\displaystyle-\sum_{i=1}^{n}\sum_{m=1}^{q}\sum_{k=1}^{K+1}Y_{ik}^{m}\Bigl{[}x_{i}^{T}\mathsf{B}e_{k}-\log\sum_{k^{\prime}=1}^{K+1}\exp(x_{i}^{T}\mathsf{B}e_{k^{\prime}})\Bigr{]}$
$\displaystyle=~{}$ $\displaystyle
q\sum_{i=1}^{n}\sum_{k=1}^{K+1}\mathsf{y}_{ik}\Bigl{[}-x_{i}^{T}\mathsf{B}e_{k}+\log\sum_{k^{\prime}=1}^{K+1}\exp(x_{i}^{T}\mathsf{B}e_{k^{\prime}})\Bigr{]}$
$\displaystyle=~{}$ $\displaystyle
q\sum_{i=1}^{n}\Bigl{[}\sum_{k=1}^{K+1}-\mathsf{y}_{ik}x_{i}^{T}\mathsf{B}e_{k}+\log\sum_{k^{\prime}=1}^{K+1}\exp(x_{i}^{T}\mathsf{B}e_{k^{\prime}})\Bigr{]}$
$\displaystyle=~{}$ $\displaystyle
q\sum_{i=1}^{n}\mathsf{L}_{i}(\mathsf{B}^{T}x_{i}),$

where the first equality uses
$\mathsf{y}_{ik}=\frac{1}{q}\sum_{m=1}^{q}Y_{ik}^{m}$, the second equality
uses $\sum_{k=1}^{K+1}\mathsf{y}_{ik}=1$ under the following Assumption S1.1,
and the last equality uses the definition of $\mathsf{L}_{i}$ after (1.5).

######  Assumption S1.1.

For all $i\in[n],$ the response $\mathsf{y}_{i}$ is in
$\\{0,1/q,2/q,...,1\\}^{K+1}$ with $\sum_{k=1}^{K+1}\mathsf{y}_{ik}=1$.

In such repeated measurements model, we replace Assumption 2.2 with Assumption
S1.1 under which the following Theorem S1.1 holds.

######  Theorem S1.1.

Let $q\geq 2$ be constant. Let Assumptions 2.1, S1.1, 2.3 and 2.4 be
fulfilled. For any $j\in[p]$ such that $H_{0}$ in (1.10) holds, we have the
convergence in distribution (2.1), (2.2) and (2.3).

###### Proof of Theorem S1.1.

Under the assumptions in Theorem S1.1, the MLE $\hat{\mathsf{B}}$ for this $q$
repeated measurements model is the minimizer of the optimization problem

$\hat{\mathsf{B}}\in\operatorname*{arg\,min}_{\mathsf{B}\in\mathbb{R}^{p\times(K+1)}}\sum_{i=1}^{n}\mathsf{L}_{i}(\mathsf{B}^{T}x_{i})$

as in (1.5). Similar to the non-repeated model, the MLE $\hat{A}$ for the
identifiable log-odds model can be expressed as

$\hat{A}=\operatorname*{arg\,min}_{A\in\mathbb{R}^{p\times
K}}\sum_{i=1}^{n}\mathsf{L}_{i}((A,\bm{0}_{p})^{T}x_{i}).$

The only difference between this $q$ repeated measurements model and the non-
repeated model considered in the main text is that the response
$\mathsf{y}_{ik}$ for this $q$ repeated measurements model is now valued in
$\\{0,1/q,2/q,...,1\\}$. Because the proofs of Theorems 2.1 and 2.2 do not
require the value of $\mathsf{y}_{ik}$ to be $\\{0,1\\}$-valued. Theorem S1.1
can be proved by the same arguments used in the proof of Theorems 2.1 and 2.2.
∎

##  S2. Implementation details and additional figures

The pivotal quantities in our main results Theorems 2.1 and 2.2 involve only
observable quantities that can be computed from the data
$(x_{i},\mathsf{y}_{i})_{i\in[n]}$. In this section we provide an efficient
way of computing the matrix $\mathsf{V}_{i}$ appearing in Theorems 2.1 and
2.2.

#### Fast computation of $\mathsf{V}_{i}$.

Recall the definition of $\mathsf{V}_{i}$ in Theorem 2.1,

$\mathsf{V}_{i}=\mathsf{H}_{i}-(\mathsf{H}_{i}\otimes
x_{i}^{T})\Bigl{[}\sum_{l=1}^{n}\mathsf{H}_{l}\otimes(x_{l}x_{l}^{T})\Bigr{]}^{\dagger}(\mathsf{H}_{i}\otimes
x_{i}).$

The majority of computational cost in calculating $\mathsf{V}_{i}$ lies in the
step of calculating its second term

$(\mathsf{H}_{i}\otimes
x_{i}^{T})\Bigl{[}\sum_{l=1}^{n}\mathsf{H}_{l}\otimes(x_{l}x_{l}^{T})\Bigr{]}^{\dagger}(\mathsf{H}_{i}\otimes
x_{i}^{T}).$

Here we provide an efficient way to compute this term using the Woodbury
matrix identity. Since $\mathsf{H}_{i}\bm{1}_{K+1}=\bm{0}_{K+1}$, we have
$\ker(\mathsf{H}_{i}\otimes(x_{i}x_{i}^{T}))$ is the span of
$\\{\bm{1}_{K+1}\otimes e_{j}:j\in[p]\\}$, where $\bm{1}_{K+1}$ is the all-
ones vector in $\mathbb{R}^{K+1}$. Therefore, the second term in
$\mathsf{V}_{i}$ can be rewritten as

$\displaystyle(\mathsf{H}_{i}\otimes
x_{i}^{T})\Bigl{[}\sum_{l=1}^{n}\mathsf{H}_{l}\otimes(x_{l}x_{l}^{T})\Bigr{]}^{\dagger}(\mathsf{H}_{i}\otimes
x_{i})$ $\displaystyle=~{}$ $\displaystyle(\mathsf{H}_{i}\otimes
x_{i}^{T})\Bigl{[}\sum_{l=1}^{n}\mathsf{H}_{l}\otimes(x_{l}x_{l}^{T})-\sum_{j=1}^{p}(\bm{1}\otimes
e_{j})(\bm{1}\otimes e_{j})^{T}\Bigr{]}^{-1}(\mathsf{H}_{i}\otimes x_{i}).$

We now apply the Woodbury matrix identity to compute the matrix inversion in
the above display. Recall
$\mathsf{H}_{i}=\operatorname{diag}(\hat{\mathsf{p}}_{i})-\hat{\mathsf{p}}_{i}\hat{\mathsf{p}}_{i}^{T}$,
we have

$\sum_{i=1}^{n}\mathsf{H}_{i}\otimes(x_{i}x_{i}^{T})=\sum_{k=1}^{K+1}(e_{k}e_{k}^{T})\otimes(\sum_{i=1}^{n}\hat{\mathsf{p}}_{ik}x_{i}x_{i}^{T})-\sum_{i=1}^{n}(\hat{\mathsf{p}}_{i}\otimes
x_{i})(\hat{\mathsf{p}}_{i}\otimes x_{i})^{T}.$

Let
$A=\sum_{k=1}^{K+1}(e_{k}e_{k}^{T})\otimes(\sum_{i=1}^{n}\hat{\mathsf{p}}_{ik}x_{i}x_{i}^{T})$,
and $U\in\mathbb{R}^{p(K+1)\times(n+p)}$ with the first $n$ columns being
$(\hat{\mathsf{p}}_{i}\otimes x_{i})_{i\in[n]}$ and the following $p$ columns
$(\bm{1}\otimes e_{j})_{j\in[p]}$. Then the term we want to invert is
$A-UU^{T}$, where $A$ is a block diagonal matrix and can be inverted by
inverting each block separately. By the Woodbury matrix identity, we have

$\displaystyle(A-UU^{T})^{-1}$
$\displaystyle=A^{-1}-A^{-1}U(-I_{n+p}+U^{T}A^{-1}U)^{-1}U^{T}A^{-1}.$

The gain of using the above formula is significant for large $K$: instead of
inverting the $p(K+1)\times p(K+1)$ matrix
$\sum_{l=1}^{n}\mathsf{H}_{l}\otimes(x_{l}x_{l}^{T})$ in the left-hand side,
the right-hand side only requires to invert a block diagonal matrix $A$ and a
$(n+p)\times(n+p)$ matrix $-I_{n+p}+U^{T}A^{-1}U$.

(a) $(n,K)=(4000,2)$

(b) $(n,K)=(5000,3)$

(c) $(n,K)=(6000,4)$

Figure S1. The upper row: Q-Q plots of the test statistics from (2.3) (in
blue) and (1.13) (in orange) for different $(n,K)$ and $p=1000$ using
$\hat{\Omega}_{jj}$. The lower row: histograms of p-values from classical test
and our test for different $(n,K)$ and $p=1000$ using $\hat{\Omega}_{jj}$.

(a) $(n,K)=(4000,2)$

(b) $(n,K)=(5000,3)$

(c) $(n,K)=(6000,4)$

Figure S2.  Q-Q plots of the test statistics from (2.3) (in blue) and (1.13)
(in orange) for different $(n,K)$ and $p=1000$ using $\hat{\Omega}_{jj}$. The
upper row: covariates are sampled from Rademacher distribution. The lower row:
covariates are sampled from distribution of SNPs.

(a) $q=1$

(b) $q=2$

(c) $q=4$

Figure S3. Scatter plot of pairs $(\sqrt{n}\hat{A}_{j1},\sqrt{n}\hat{A}_{j2})$
with the same data generating process as in Figure 1 (a) except using
different $q$.

##  S3. Proof of Theorem 2.1

Before proving Theorem 2.1, we present another parametrization of the
multinomial logistic regression model. The asymptotic theory of MLE for this
new parametrized multinomial logistic model will be used to prove Theorem 2.1.

###  S3.1. Another parametrization of multinomial logistic regression

Recall the symbol “$\in$" is used in (1.5) to emphasize that the minimizer
$\hat{\mathsf{B}}$ in (1.5) is not unique: if $\hat{\mathsf{B}}$ is a
minimizer of (1.5) then $\hat{\mathsf{B}}-b\bm{1}_{K+1}^{T}$ is also a
minimizer of (1.5), for any $b\in\mathbb{R}^{p}$ and the all-ones vector
$\bm{1}_{K+1}$ in $\mathbb{R}^{K+1}$.

Besides the log-odds model (1.8), here we consider another identifiable
parametrization of multinomial logistic regression, whose unknown parameter,
denoted by $B^{*}$, is in $\mathbb{R}^{p\times K}$.

#### Orthogonal complement.

To obtain an identifiable multinomial logistic regression model from (1.4), we
consider the symmetric constraint
$\mathsf{B}^{*}\bm{1}=\sum_{k=1}^{K+1}\mathsf{B}^{*}e_{k}=\bm{0}$ as in [Zhu
and Hastie, 2004], thus
$\mathsf{B}^{*}=\mathsf{B}^{*}(I_{K+1}-\frac{\bm{1}\bm{1}^{T}}{K+1})$, where
$\bm{1}$ is the all-ones vector in $\mathbb{R}^{K+1}$. Let
$Q\in\mathbb{R}^{(K+1)\times K}$ be any matrix such that

(S3.1) $I_{K+1}-\tfrac{1}{K+1}\bm{1}\bm{1}^{T}=QQ^{T},\qquad Q^{T}Q=I_{K}.$

We fix one choice of $Q$ satisfying (S3.1) throughout this supplement. Let
$B^{*}=\mathsf{B}^{*}Q$, then $\mathsf{B}^{*}=B^{*}Q^{T}$ and the model (1.4)
can be parameterized using $B^{*}$ as

(S3.2)
$\mathbb{P}(\mathsf{y}_{ik}=1|x_{i})=\frac{\exp(x_{i}^{T}B^{*}Qe_{k})}{\sum_{k^{\prime}=1}^{K+1}\exp(x_{i}^{T}B^{*}Qe_{k^{\prime}})},\quad
k\in\\{1,2,\ldots,K+1\\}.$

The multinomial logistic MLE of $B^{*}$ in (S3.2) is given by

(S3.3) $\textstyle\hat{B}=\operatorname*{arg\,min}_{B\in\mathbb{R}^{p\times
K}}\sum_{i=1}^{n}L_{i}(B^{T}x_{i}),$

where $L_{i}:\mathbb{R}^{K}\to\mathbb{R}$ is defined by
$L_{i}(u)=\mathsf{L}_{i}(Qu)\text{ for all }u\in\mathbb{R}^{K}.$ By this
construction, we have $\hat{B}=\hat{\mathsf{B}}Q$ for any minimizer
$\hat{\mathsf{B}}$ of (1.5). Furthermore, by the chain rule using the
expressions (1.7), the gradient and Hessian of $L_{i}$ evaluated at
$\hat{B}^{T}x_{i}$ are

(S3.4) $g_{i}:=\nabla L_{i}(\hat{B}^{T}x_{i})=Q^{T}\mathsf{g}_{i},\qquad
H_{i}:=\nabla^{2}L_{i}(\hat{B}^{T}x_{i})=Q^{T}\mathsf{H}_{i}Q.$

Throughout, we use serif upright letters to denote quantities defined on the
unidentifiable parameter space $\mathbb{R}^{p\times(K+1)}$:

$\mathsf{B}^{*},\hat{\mathsf{B}}\in\mathbb{R}^{p\times(K+1)},\quad\mathsf{L}_{i}:\mathbb{R}^{K+1}\to\mathbb{R},\qquad\hat{\mathsf{p}}_{i},\mathsf{y}_{i},\mathsf{g}_{i}\in\mathbb{R}^{K+1},\qquad\mathsf{H}_{i}\in\mathbb{R}^{(K+1)\times(K+1)}$

and the normal italic font to denote analogous quantities for the identifiable
parameter space $\mathbb{R}^{p\times K}$:

$B^{*},\hat{B}\in\mathbb{R}^{p\times K},\qquad
L_{i}:\mathbb{R}^{K}\to\mathbb{R},\qquad g_{i}\in\mathbb{R}^{K},\qquad
H_{i}\in\mathbb{R}^{K\times K}.$

Theorem S3.1 provides the asymptotic normality and the chi-square
approximation of null MLE coordinates in high-dimensions where $n,p\to\infty$
with the ratio $n/p$ converging to a finite limit.

######  Theorem S3.1 (Proof is given on page S6).

Let Assumptions 2.1, 2.3 and 2.4 be fulfilled. Assume that either Assumption
2.2 or Assumption S1.1 holds. Then for any $j\in[p]$ such that $H_{0}$ in
(1.10) holds,

(S3.5)
$\sqrt{n}\Omega_{jj}^{-1/2}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}g_{i}g_{i}^{T}\Bigr{)}^{-1/2}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}V_{i}\Bigr{)}\hat{B}^{T}e_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(0,I_{K}),$

where $V_{i}=H_{i}-(H_{i}\otimes
x_{i}^{T})[\sum_{l=1}^{n}H_{l}\otimes(x_{l}x_{l}^{T})]^{-1}(H_{i}\otimes
x_{i})$.

A direct consequence of (S3.5) is the $\chi^{2}$ result,

(S3.6)
$\|\sqrt{n}\Omega_{jj}^{-1/2}\bigl{(}\frac{1}{n}\sum_{i=1}^{n}g_{i}g_{i}^{T}\bigr{)}^{-1/2}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}V_{i}\Bigr{)}\hat{B}^{T}e_{j}\|^{2}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}\chi^{2}_{K}.$

The proof of Theorem S3.1 is deferred to Section S5 and Section S6. In the
next subsection, we prove Theorem 2.1 using Theorem S3.1.

###  S3.2. Proof of Theorem 2.1

We restate Theorem 2.1 for convenience. See 2.1

The proof of Theorem 2.1 is a consequence of Theorem S3.1. To begin with, we
state the following useful lemma.

######  Lemma S3.2.

For $\mathsf{V}_{i}$ and $V_{i}$ defined in Theorems S3.1 and 2.1, we have
$V_{i}=Q^{T}\mathsf{V}_{i}Q$.

###### Proof of Lemma S3.2.

Since $H_{i}=Q^{T}\mathsf{H}_{i}Q$, we have

$\displaystyle V_{i}=$ $\displaystyle H_{i}-(H_{i}\otimes
x_{i}^{T})\Bigl{[}\sum_{i=1}^{n}H_{i}\otimes(x_{i}x_{i}^{T})\Bigr{]}^{-1}(H_{i}\otimes
x_{i})$ $\displaystyle=$ $\displaystyle H_{i}-[(Q^{T}\mathsf{H}_{i}Q)\otimes
x_{i}^{T}]\Bigl{[}\sum_{i=1}^{n}(Q^{T}\mathsf{H}_{i}Q)\otimes(x_{i}x_{i}^{T})\Bigr{]}^{-1}[(Q^{T}\mathsf{H}_{i}Q)\otimes
x_{i}]$ $\displaystyle=$ $\displaystyle H_{i}-Q^{T}(\mathsf{H}_{i}\otimes
x_{i}^{T})(Q\otimes I_{p})\Bigl{[}(Q^{T}\otimes
I_{p})[\sum_{i=1}^{n}\mathsf{H}_{i}\otimes(x_{i}x_{i}^{T})](Q\otimes
I_{p})\Bigr{]}^{-1}(Q^{T}\otimes I_{p})(\mathsf{H}_{i}\otimes x_{i})Q$
$\displaystyle=$ $\displaystyle
Q^{T}\mathsf{H}_{i}Q-Q^{T}(\mathsf{H}_{i}\otimes
x_{i}^{T})\Bigl{[}\sum_{i=1}^{n}(\mathsf{H}_{i}\otimes
x_{i}x_{i}^{T})\Bigr{]}^{\dagger}(\mathsf{H}_{i}\otimes x_{i})Q$
$\displaystyle=$ $\displaystyle Q^{T}\mathsf{V}_{i}Q,$

where the penultimate equality is proved as follows.

Let $A=Q\otimes I_{p}$ and
$\mathsf{D}=\sum_{i=1}^{n}\mathsf{H}_{i}\otimes(x_{i}x_{i}^{T})$ only in the
remaining of this proof. It remains to prove

(S3.7) $A[A^{T}\mathsf{D}A]^{-1}A^{T}=\mathsf{D}^{\dagger}.$

Since $\mathsf{H}_{i}\bm{1}=0$, we have $\mathsf{D}(\bm{1}\otimes I_{p})=0$.
Since $Q^{T}\bm{1}=0$ by definition of $Q$, we have $A^{T}(\bm{1}\otimes
I_{p})=0$. If we write the eigen-decomposition of $\mathsf{D}$ as
$\mathsf{D}=\sum_{i=1}^{pK}\lambda_{i}u_{i}u_{i}^{T}$, then
$u_{i}^{T}(\bm{1}\otimes I_{p})=0$. Hence, with $v_{i}=A^{T}u_{i}$,

$\displaystyle A^{T}\mathsf{D}A=~{}$
$\displaystyle\sum_{i=1}^{pK}\lambda_{i}v_{i}v_{i}^{T}.$

Since
$v_{i}^{T}v_{i^{\prime}}=u_{i}^{T}AA^{T}u_{i^{\prime}}=u_{i}^{T}[(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})\otimes
I_{p}]u_{i^{\prime}}=u_{i}^{T}u_{i^{\prime}}=I(i=i^{\prime})$, we have

$\displaystyle
A[A^{T}\mathsf{D}A]^{-1}A^{T}=A\bigl{(}\sum_{i=1}^{pK}\lambda_{i}^{-1}v_{i}v_{i}^{T}\bigr{)}A^{T}=\sum_{i=1}^{pK}\lambda_{i}^{-1}u_{i}u_{i}^{T}=\mathsf{D}^{\dagger},$

where the second equality uses $Av_{i}=AA^{T}u_{i}=u_{i}$. The proof of (S3.7)
is complete. ∎

Now we are ready to prove that Theorem 2.1 is a consequence of Theorem S3.1.

###### Proof of Theorem 2.1.

By definition of $\mathsf{g}_{i}$ and $\mathsf{V}_{i}$, we have
$\bm{1}^{T}\mathsf{g}_{i}=0$ and $\bm{1}^{T}\mathsf{V}_{i}=\bm{0}^{T}$. Thus,
we have $QQ^{T}\mathsf{g}_{i}=\mathsf{g}_{i}$ and
$QQ^{T}\mathsf{V}_{i}=\mathsf{V}_{i}$. Therefore, we can rewrite the left-hand
side of (2.1) (without $\sqrt{n}\Omega_{jj}^{-1/2}$) as

$\displaystyle\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{g}_{i}\mathsf{g}_{i}^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}\Bigr{)}\hat{\mathsf{B}}^{T}e_{j}$
$\displaystyle=~{}$
$\displaystyle\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}QQ^{T}\mathsf{g}_{i}\mathsf{g}_{i}^{T}QQ^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}QQ^{T}\mathsf{V}_{i}QQ^{T}\Bigr{)}\hat{\mathsf{B}}^{T}e_{j}$
$\displaystyle=~{}$
$\displaystyle\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}Qg_{i}g_{i}^{T}Q^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}QV_{i}Q^{T}\Bigr{)}\hat{\mathsf{B}}^{T}e_{j}$
$\displaystyle=~{}$ $\displaystyle
Q\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}g_{i}g_{i}^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}Q^{T}Q\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}V_{i}Q^{T}\Bigr{)}\hat{\mathsf{B}}^{T}e_{j}$
$\displaystyle=~{}$ $\displaystyle
Q\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}g_{i}g_{i}^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}V_{i}\Bigr{)}\hat{B}^{T}e_{j},$

where the first equality uses $QQ^{T}\mathsf{g}_{i}=\mathsf{g}_{i}$ and
$QQ^{T}\mathsf{V}_{i}=\mathsf{V}_{i}$, the second equality uses
$g_{i}=Q^{T}\mathsf{g}_{i}$ and $V_{i}=Q^{T}\mathsf{V}_{i}Q$ from Lemma S3.2,
the third equality follows from the same argument of (S3.7), and the last
equality uses $Q^{T}Q=I_{K}$ and $\hat{B}=\hat{\mathsf{B}}Q$.

Therefore, Theorem S3.1 implies that the limiting covariance for the left-hand
side of (2.1) is $QQ^{T}=I_{K}-\tfrac{\bm{1}\bm{1}^{T}}{K+1}$. This completes
the proof. ∎

##  S4. Proof of Theorem 2.2

We restate Theorem 2.2 for convenience. See 2.2

The proof is a direct consequence of Theorem 2.1.

###### Proof of Theorem 2.2.

By definition of $\hat{A}$ in (1.9), we have
$\hat{A}=\hat{\mathsf{B}}(I_{K},-\bm{1}_{K})^{T}$ and

$\hat{A}~{}(I_{K},\bm{0}_{K})=\hat{\mathsf{B}}(I_{K},-\bm{1}_{K})^{T}(I_{K},\bm{0}_{K})=\hat{\mathsf{B}}(I_{K+1}-e_{K+1}\bm{1}^{T})=\hat{\mathsf{B}}-(\hat{\mathsf{B}}e_{K+1})\bm{1}^{T},$

which is of the form $\hat{\mathsf{B}}-b\bm{1}^{T}$ with
$b=\hat{\mathsf{B}}e_{K+1}$. Therefore, $\hat{A}~{}(I_{K},\bm{0}_{K})$ is also
a solution of (1.5). Taking $\hat{\mathsf{B}}$ in Theorem 2.1 to be
$\hat{A}~{}(I_{K},\bm{0}_{K})=\hat{A}R^{T}$ gives the desired $\chi^{2}$
result (2.3) and

(S4.1)
$\sqrt{n}\Omega_{jj}^{-1/2}\Bigl{(}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{g}_{i}\mathsf{g}_{i}^{T}\Bigr{)}^{1/2}\Bigr{)}^{\dagger}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}\mathsf{V}_{i}\Bigr{)}R\hat{A}^{T}e_{j}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N\Bigl{(}0,I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1}\Bigr{)}.$

Multiplying $(R^{T}(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})R)^{-1/2}R^{T}$ to
the left of the above display gives the desired normality result (2.2) by
observing
$(R^{T}(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})R)^{-1/2}=(I_{K}+\frac{\bm{1}_{K}\bm{1}_{K}^{T}}{\sqrt{K+1}+1}).$
This completes the proof. ∎

##  S5. Preliminary results for proving Theorem S3.1

###  S5.1. Results for general loss functions

In this subsection, we will work under the following assumptions with a
general convex loss function. Later in Section S6, we will apply the general
results of this subsection to the multinomial logistic loss discussed in
Section S3.1.

######  Assumption S5.1.

Suppose we have data $(Y,X)$, where $Y\in\mathbb{R}^{n\times(K+1)}$ with rows
$(\mathsf{y}_{1},...,\mathsf{y}_{n})$, and $X\in\mathbb{R}^{n\times p}$ has
i.i.d. rows $(x_{1},...,x_{n})$ with $x_{i}\sim N(\bm{0},\Sigma)$ and
invertible $\Sigma$. The observations $(\mathsf{y}_{i},x_{i})_{i\in[n]}$ are
i.i.d. and $\mathsf{y}_{i}$ has the form
$\mathsf{y}_{i}=f(U_{i},x_{i}^{T}B^{*})$ for some deterministic function $f$,
deterministic $B^{*}\in\mathbb{R}^{p\times K}$, and latent random variable
$U_{i}$ independent of $x_{i}$. Assume $p/n\leq\delta^{-1}<1$.

######  Assumption S5.2.

Given data $(Y,X)$, consider twice continuously differentiable and strictly
convex loss functions $(L_{i})_{i\in[n]}$ with each
$L_{i}:\mathbb{R}^{K}\to\mathbb{R}$ depending on $\mathsf{y}_{i}$ but not on
$x_{i}$.

Provided that the following minimization problem admits a solution, define

$\displaystyle\hat{B}(Y,X)=\operatorname*{arg\,min}_{B\in\mathbb{R}^{p\times
K}}\sum_{i=1}^{n}L_{i}(B^{T}x_{i}).$

Define for each $i\in[n]$,

$\displaystyle 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

$\displaystyle G(Y,X)$ $\displaystyle=\sum_{i=1}^{n}e_{i}g_{i}(Y,X)^{T},$
$\displaystyle V(Y,X)$
$\displaystyle=\sum_{i=1}^{n}\Bigl{(}H_{i}(Y,X)-(H_{i}(Y,X)\otimes
x_{i}^{T})\Bigl{[}\sum_{l=1}^{n}H_{l}(Y,X)\otimes(x_{l}x_{l}^{T})\Bigr{]}^{\dagger}(H_{i}(Y,X)\otimes
x_{i})\Bigr{)},$

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.2 and S5.1 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^{*}=\bm{0}^{T}_{K}$, there exists
a random variable $\xi\in\mathbb{R}^{K}$ such that

$\mathbb{E}\bigl{[}I\\{(Y,X)\in
U^{*}\\}\bigl{\|}\tfrac{(G^{T}G)^{-1/2}V\hat{B}^{T}e_{j}}{\sqrt{\Omega_{jj}}}-\xi\bigr{\|}^{2}\bigr{]}\leq\tfrac{C}{p-K},$

and $\mathbb{P}(\|\xi\|^{2}>\chi^{2}_{K}(\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{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},I_{K})$ and

$\frac{(G^{T}G)^{-1/2}V\hat{B}^{T}e_{j}}{\sqrt{\Omega_{jj}}}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{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_{k}u_{k}v_{k}^{T}$ be the singular
value decomposition of $\Sigma^{1/2}B^{*}$, where $u_{1},...,u_{K}$ are the
left singular vectors and $v_{1},...,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_{k}u_{k}v_{k}^{T}$
still holds with orthonormal $(u_{1},...,u_{K})$ and orthonormal
$(v_{1},...,v_{K})$. We consider an orthogonal matrix
$\tilde{P}\in\mathbb{R}^{p\times p}$ such that

(S5.1)
$\tilde{P}\tilde{P}^{T}=\tilde{P}^{T}\tilde{P}=I_{p},\quad\tilde{P}\frac{\Sigma^{-1/2}e_{j}}{\|\Sigma^{-1/2}e_{j}\|}=e_{1},\quad\tilde{P}u_{k}=e_{p-K+k},\quad\forall
k\in[K].$

Since $e_{j}^{T}B^{*}=\bm{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,

(S5.2) $Z=X\Sigma^{-1/2}\tilde{P}^{T},\qquad
A^{*}=\tilde{P}\Sigma^{1/2}B^{*}.$

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 $\mathsf{y}_{i}$ has
the expression $\mathsf{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

$\displaystyle\hat{A}(Y,Z)=\operatorname*{arg\,min}_{A\in\mathbb{R}^{p\times
K}}\sum_{i}L_{i}(A^{T}z_{i}),$

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

$\displaystyle\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

$\displaystyle\frac{e_{j}^{T}\hat{B}V(G^{T}G)^{-1/2}}{\sqrt{\Omega_{jj}}}=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},...,z_{n})$ with
$z_{i}\sim N(\bm{0},I_{p})$. Let
$\mathsf{y}_{1},...,\mathsf{y}_{n}\in\mathbb{R}^{(K+1)}$ such that
$(\mathsf{y}_{1},...,\mathsf{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,...,n}$ with each
$L_{i}:\mathbb{R}^{K}\to\mathbb{R}$ depending on $\mathsf{y}_{i}$ but not on
$z_{i}$ and define, provided that the minimizer admits a solution,

$\displaystyle\hat{A}(Y,Z)=\operatorname*{arg\,min}_{A\in\mathbb{R}^{p\times
K}}\sum_{i=1}^{n}L_{i}(A^{T}z_{i}),\quad g_{i}(Y,Z)=\nabla
L_{i}(\hat{A}(Y,Z)^{T}z_{i}),\quad
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}$, and
$V(Y,Z)=\sum_{i=1}^{n}\bigl{(}H_{i}-(H_{i}\otimes
z_{i}^{T})\bigl{[}\sum_{l=1}^{n}H_{l}\otimes(z_{l}z_{l}^{T})\bigr{]}^{\dagger}(H_{i}\otimes
z_{i})\bigr{)}\in\mathbb{R}^{K\times K}$, where we dropped the dependence of
$H_{i}$ on $(Y,Z)$ 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}\bigl{[}I\\{(Y,Z)\in
O\\}\bigl{\|}(G^{T}G)^{-1/2}V\hat{A}^{T}e_{1}-\xi\bigr{\|}^{2}\bigr{]}\leq\tfrac{C}{p-K},$

and $\mathbb{P}(\|\xi\|^{2}>\chi^{2}_{K}(\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{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},I_{K})$ and

$e_{1}^{T}\hat{A}V(G^{T}G)^{-1/2}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{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 section S5.3).

Let $z\sim N(\bm{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}\Bigl{[}\Bigl{\|}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})\Bigr{\|}^{2}\Bigr{]}\leq
3\sigma^{4}\mathbb{E}\sum_{i=1}^{n}\Big{\|}\frac{\partial F(z)}{\partial
z_{i}}\Big{\|}_{F}^{2}.$

######  Lemma S5.4 (Proof is given on section S5.3).

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

$\displaystyle\|(G^{T}G)^{-1/2}-(\tilde{G}^{T}\tilde{G})^{-1/2}\|_{F}$
$\displaystyle\leq L_{1}n^{-1}\|G-\tilde{G}\|_{F},$
$\displaystyle\|G(G^{T}G)^{-1/2}-\tilde{G}(\tilde{G}^{T}\tilde{G})^{-1/2}\|_{F}$
$\displaystyle\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 section S5.3).

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

$\displaystyle\pdv{\operatorname{vec}(\hat{A})}{z_{ij}}=-M[g_{i}\otimes
e_{j}+(H_{i}\hat{A}^{T}e_{j}\otimes z_{i})],$
$\displaystyle\pdv{g_{l}}{z_{ij}}=-(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=[\sum_{i=1}^{n}H_{i}\otimes(z_{i}z_{i}^{T})]^{-1}$. It immediately
follows that

$\pdv{g_{i}}{z_{ij}}=[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 section S5.3).

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]$,

$\displaystyle\sum_{i=1}^{n}\pdv{e_{i}^{T}G(G^{T}G)^{-1/2}}{z_{ij}}$
$\displaystyle=~{}$ $\displaystyle
e_{j}^{T}\hat{A}V_{i}^{T}(G^{T}G)^{-1/2}+\sum_{i=1}^{n}\Bigl{[}-(g_{i}^{T}\otimes
e_{j}^{T})M(H_{i}\otimes
z_{i})(G^{T}G)^{-1/2}+e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{ij}}\Bigr{]}.$

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_{2}n^{-1/2}$-Lipschitz in the sense that
$\|h(Y,Z)-h(Y,\tilde{Z})\|_{F}\leq LL_{2}n^{-1/2}\|Z-\tilde{Z}\|_{F}$ for all
$\\{(Y,Z),(Y,\tilde{Z})\\}\subset O$. By Kirszbraun 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_{2}n^{-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 $\mathsf{z}_{j}=Ze_{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{\mathsf{z}}\sim N(\bm{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{\mathsf{z}}e_{j}^{T}$. That is,
$\check{Z}^{j}$ replaces the $j$-th column of $Z$ by $\check{\mathsf{z}}$. By
definition, $\mathsf{z}_{1}\perp\check{\mathsf{z}}$ and
$\mathsf{z}_{1}\perp\check{Z}^{1}$.

Let $\xi=-[H(Y,\check{Z}^{1})]^{T}\mathsf{z}_{1}\in\mathbb{R}^{K}$, then
$\|\xi\|^{2}\leq\|\mathsf{z}_{1}\|$ since $\|H(Y,\check{Z}^{1})\|_{op}\leq 1$.
It follows that

$\mathbb{P}(\|\xi\|^{2}>\chi^{2}_{K}(\alpha))\leq\mathbb{P}(\|\mathsf{z}_{1}\|^{2}>\chi^{2}_{K}(\alpha))=\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]$,

$\displaystyle\mathbb{E}\bigl{[}I\\{(Y,Z)\in
O\\}\bigl{\|}(G^{T}G)^{-1/2}V\hat{A}^{T}e_{1}-\xi\bigr{\|}^{2}\bigr{]}$
$\displaystyle=~{}$ $\displaystyle\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\|e_{1}^{T}\hat{A}V(G^{T}G)^{-1/2}+\mathsf{z}_{1}^{T}H(Y,\check{Z}^{1})\|^{2}\Bigr{]}$
$\displaystyle=~{}$ $\displaystyle\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\|e_{\ell}^{T}\hat{A}V(G^{T}G)^{-1/2}+\mathsf{z}_{\ell}^{T}H(Y,\check{Z}^{\ell})\|^{2}\Bigr{]},$

where the last line holds for any $\ell\in[p-K]$ because
$(\mathsf{z}_{1},e_{1}^{T}\hat{A},\check{Z})\stackrel{{\scriptstyle
d}}{{=}}(\mathsf{z}_{\ell},e_{\ell}^{T}\hat{A},\check{Z}^{\ell})$. Therefore,

$\displaystyle\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Bigl{\|}e_{1}^{T}\hat{A}V(G^{T}G)^{-1/2}+\mathsf{z}_{1}^{T}H(Y,\check{Z}^{1})\Bigr{\|}^{2}\Bigr{]}$
$\displaystyle=~{}$
$\displaystyle\frac{1}{p-K}\sum_{\ell=1}^{p-K}\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Bigl{\|}e_{\ell}^{T}\hat{A}V(G^{T}G)^{-1/2}+\mathsf{z}_{\ell}^{T}H(Y,\check{Z}^{\ell})\Bigr{\|}^{2}\Bigr{]}$
$\displaystyle=~{}$
$\displaystyle\frac{1}{p-K}\sum_{\ell=1}^{p-K}\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Bigl{\|}\sum_{i}\pdv{e_{i}^{T}G(G^{T}G)^{-1/2}}{z_{i\ell}}+\mathsf{z}_{\ell}^{T}H(Y,\check{Z}^{\ell})-\operatorname{Rem}_{\ell}\Bigr{\|}^{2}\Bigr{]},$

where $\operatorname{Rem}_{\ell}=\sum_{i=1}^{n}\bigl{[}-(g_{i}^{T}\otimes
e_{\ell}^{T})M(H_{i}\otimes
z_{i})(G^{T}G)^{-1/2}+e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{i\ell}}\bigr{]}$ from
Corollary S5.6, and $M=[\sum_{i=1}^{n}(H_{i}\otimes z_{i}z_{i}^{T})]^{-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,

$\displaystyle(RHS)_{1}=\frac{2}{p-K}\sum_{\ell=1}^{p-K}\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Bigl{\|}\sum_{i}\pdv{e_{i}^{T}h(Y,Z)}{z_{i\ell}}+\mathsf{z}_{\ell}^{T}H(Y,\check{Z}^{\ell})\Bigr{\|}^{2}\Bigr{]}.$

Let
$F(\mathsf{z}_{\ell})=H(Y,Z(I-e_{\ell}e_{\ell}^{T})+\mathsf{z}_{\ell}e_{\ell}^{T})=H(Y,Z)$,
then $F(\check{\mathsf{z}})=H(Y,\check{Z}^{\ell})$. Apply Lemma S5.3 to
$F(\mathsf{z}_{\ell})$ conditionally on $Z(I-e_{\ell}e_{\ell}^{T})$, we obtain

$\displaystyle\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Big{\|}\sum_{i}\pdv{e_{i}^{T}h(Y,Z)}{z_{i\ell}}+\mathsf{z}_{\ell}^{T}F(\check{\mathsf{z}})\Big{\|}^{2}\Bigr{]}$
$\displaystyle=~{}$ $\displaystyle\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Big{\|}\mathsf{z}_{\ell}^{T}F(\mathsf{z}_{\ell})-\sum_{i}\pdv{e_{i}^{T}F(\mathsf{z}_{\ell})}{z_{i\ell}}-\mathsf{z}_{\ell}^{T}F(\check{\mathsf{z}})\Big{\|}^{2}\Bigr{]}$
$\displaystyle\leq~{}$
$\displaystyle\mathbb{E}\Bigl{[}\Big{\|}\mathsf{z}_{\ell}^{T}F(\mathsf{z}_{\ell})-\sum_{i}\pdv{e_{i}^{T}F(\mathsf{z}_{\ell})}{z_{i\ell}}-\mathsf{z}_{\ell}^{T}F(\check{\mathsf{z}})\Big{\|}^{2}\Bigr{]}$
$\displaystyle\leq~{}$ $\displaystyle
3\sum_{i}\mathbb{E}\Big{\|}\pdv{F(\mathsf{z}_{\ell})}{z_{i\ell}}\Big{\|}_{F}^{2}$
$\displaystyle=~{}$ $\displaystyle
3\sum_{i}\mathbb{E}\Big{\|}\pdv{H(Y,Z)}{z_{i\ell}}\Big{\|}_{F}^{2},$

where the first equality uses $\mathsf{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

$\displaystyle(RHS)_{1}\leq~{}$
$\displaystyle\frac{6}{p-K}\mathbb{E}\Bigl{[}\sum_{\ell=1}^{p}\sum_{i=1}^{n}\Big{\|}\pdv{H(Y,Z)}{z_{i\ell}}\Big{\|}_{F}^{2}\Bigr{]}.$

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\operatorname{rank}(A)\|A\|_{op}^{2}$, we
obtain

$(RHS)_{1}\leq 6K(LL_{2})^{2}/(p-K).$

For the second term
$(RHS)_{2}=\frac{2}{p-K}\sum_{\ell=1}^{p-K}\mathbb{E}\bigl{[}I\\{(Y,Z)\in
O\\}\|\operatorname{Rem}_{\ell}\|^{2}\bigr{]}$. By definition of
$\operatorname{Rem}_{\ell}$ and $(a+b)^{2}\leq 2a^{2}+2b^{2}$, we obtain

(S5.3) $\displaystyle(RHS)_{2}\leq$
$\displaystyle\frac{4}{p-K}\sum_{\ell=1}^{p-K}\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Big{\|}\sum_{i=1}^{n}(g_{i}^{T}\otimes e_{\ell}^{T})M(H_{i}\otimes
z_{i})(G^{T}G)^{-1/2}\Big{\|}^{2}\Bigr{]}$ (S5.4)
$\displaystyle+\frac{4}{p-K}\sum_{\ell=1}^{p-K}\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\Bigl{\|}\sum_{i=1}^{n}e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{i\ell}}\Bigr{\|}^{2}\Bigr{]}.$

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}(g_{i}^{T}\otimes e_{\ell}^{T})M(H_{i}\otimes
z_{i})u=\operatorname{Tr}\Bigl{[}(I_{K}\otimes
e_{\ell}^{T})M\sum_{i}(H_{i}u\otimes z_{i})g_{i}^{T}\Bigr{]}\leq
K\|(I_{K}\otimes e_{\ell}^{T})M\|_{op}\|\sum_{i}(H_{i}u\otimes
z_{i})g_{i}^{T}\|_{op}$

because the rank of the matrix inside the trace is at most $K$ and
$\operatorname{Tr}[\cdot]\leq K\|\cdot\|_{op}$ holds. Then

$\|\sum_{i}(H_{i}u\otimes z_{i})g_{i}^{T}\|_{op}=\|(I_{K}\otimes
Z^{T})\sum_{i}(H_{i}u\otimes
e_{i})e_{i}^{T}G\|_{op}\leq\|Z\|_{op}\|\sum_{i}(H_{i}u\otimes
e_{i}e_{i}^{T})\|_{op}\|G\|_{op}.$

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\|(G^{T}G)^{-1/2}\|_{op}\|M\|_{op}\|Z\|_{op}\|G\|_{op}.$

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\|_{op}\sqrt{m^{*}n}=(m^{*}/m_{*})^{1/2}c_{*}^{-1}K\|Z\|_{op}.$

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,

$\eqref{eq:rem1}\leq C(c_{*},K,L)n^{-1}.$

Now we bound (S5.4). Since

$\displaystyle\sum_{\ell=1}^{p-K}\Bigl{\|}\sum_{i=1}^{n}e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{i\ell}}\Bigr{\|}^{2}$
$\displaystyle\leq~{}$
$\displaystyle\sum_{j=1}^{p}\Bigl{\|}\sum_{i=1}^{n}e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{ij}}\Bigr{\|}^{2}$
$\displaystyle=~{}$
$\displaystyle\sum_{j=1}^{p}\sum_{k^{\prime}=1}^{K}\Bigl{(}\sum_{i=1}^{n}\sum_{k=1}^{K}e_{i}^{T}Ge_{k}e_{k}^{T}\pdv{(G^{T}G)^{-1/2}}{z_{ij}}e_{k^{\prime}}\Bigr{)}^{2}$
$\displaystyle\leq~{}$
$\displaystyle\sum_{j=1}^{p}\sum_{k^{\prime}=1}^{K}\Bigl{[}\sum_{i=1}^{n}\sum_{k=1}^{K}(e_{i}^{T}Ge_{k})^{2}\sum_{i=1}^{n}\sum_{k=1}^{K}\Bigl{(}e_{k}^{T}\pdv{(G^{T}G)^{-1/2}}{z_{ij}}e_{k^{\prime}}\Bigr{)}^{2}\Bigr{]}$
$\displaystyle=~{}$
$\displaystyle\|G\|_{F}^{2}\sum_{i=1}^{n}\sum_{j=1}^{p}\Bigr{\|}\pdv{(G^{T}G)^{-1/2}}{z_{ij}}\Bigr{\|}^{2}.$

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

$\eqref{eq:rem2}\leq 4K^{3}LL_{1}/(p-K).$

Combining the above bounds on $(RHS)_{1}$ and $(RHS)_{2}$, we have

(S5.5) $\mathbb{E}\Bigl{[}I\\{(Y,Z)\in
O\\}\|(G^{T}G)^{-1/2}V\hat{A}^{T}e_{1}-\xi\|^{2}\Bigr{]}\leq\frac{C(c_{*},K,m_{*},m^{*},L,L_{1},L_{2})}{p-K},$

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,\check{Z}^{1})\in O)\to 1$ using $(Y,Z)\stackrel{{\scriptstyle
d}}{{=}}(Y,\check{Z}^{1})$. Therefore,

(S5.6) $\xi=-[h(Y,\check{Z}^{1})^{T}\mathsf{z}_{1}]I((Y,\check{Z}^{1})\in
O)-[H(Y,\check{Z}^{1})^{T}\mathsf{z}_{1}]I((Y,\check{Z}^{1})\notin
O){\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},I_{K}).$

By (S5.5), we know
$(G^{T}G)^{-1/2}V\hat{A}^{T}e_{1}-\xi{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}0$
when $\mathbb{P}((Y,Z)\in O)\to 1$. Hence, we conclude

$\displaystyle(G^{T}G)^{-1/2}V\hat{A}^{T}e_{1}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},I_{K})\quad\text{
and
}\quad\|(G^{T}G)^{-1/2}V\hat{A}^{T}e_{1}\|^{2}{\smash{\xrightarrow{~{}\mathrm{d}~{}}}}\chi^{2}_{K}.$

∎

We next prove Lemmas S5.3, S5.4, S5.5 and 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(\bm{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})=\begin{pmatrix}[F(z)-F(\tilde{z})]e_{k}\\\
0_{n}\end{pmatrix},$

so that $z_{0}^{T}f^{(k)}(z_{0})=z^{T}[F(z)-F(\tilde{z})]e_{k}$, and
$\operatorname{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] to $f^{(k)}$ gives, with $\operatorname{Jac}$
denoting the Jacobian,

$\displaystyle\mathbb{E}\Bigl{[}\Bigl{(}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}\Bigr{)}^{2}\Bigr{]}$ $\displaystyle=~{}$
$\displaystyle\mathbb{E}\Bigl{[}\Bigl{(}z_{0}^{T}f^{(k)}(z_{0})-\sigma^{2}\operatorname{div}f^{(k)}(z_{0})\Bigr{)}^{2}\Bigr{]}$
$\displaystyle=~{}$
$\displaystyle\sigma^{2}\mathbb{E}\|f^{(k)}(z_{0})\|^{2}+\sigma^{4}\mathbb{E}\operatorname{Tr}[(\operatorname{Jac}f^{(k)}(z_{0}))^{2}]$
$\displaystyle=~{}$
$\displaystyle\sigma^{2}\mathbb{E}\|[F(z)-F(\tilde{z})]e_{k}\|^{2}+\sigma^{4}\mathbb{E}\operatorname{Tr}\Bigl{[}\begin{pmatrix}\operatorname{Jac}[F(z)e_{k}]&-\operatorname{Jac}[F(\tilde{z})e_{k}]\\\
0_{n\times n}&0_{n\times n}\end{pmatrix}^{2}\Bigr{]}$ $\displaystyle=~{}$
$\displaystyle
2\sigma^{2}\mathbb{E}\|[F(z)-\mathbb{E}F(z)]e_{k}\|^{2}+\sigma^{4}\mathbb{E}\operatorname{Tr}((\operatorname{Jac}[F(z)e_{k}])^{2})$
$\displaystyle\leq~{}$ $\displaystyle
3\sigma^{4}\mathbb{E}\|\operatorname{Jac}[F(z)e_{k}]\|_{F}^{2},$

where the last inequality uses the Gaussian Poincaré inequality, and the
Cauchy-Schwarz inequality $\operatorname{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

$\displaystyle\|G^{T}G-\tilde{G}^{T}\tilde{G}\|_{op}$ $\displaystyle=~{}$
$\displaystyle\|(G-\tilde{G})^{T}G+\tilde{G}^{T}(G-\tilde{G})\|_{op}$
$\displaystyle\leq~{}$
$\displaystyle\|G-\tilde{G}\|_{op}(\|G\|_{op}+\|\tilde{G}\|_{op})$
$\displaystyle\leq~{}$ $\displaystyle 2\sqrt{m^{*}n}\|G-\tilde{G}\|_{op}.$

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

$\displaystyle\|A^{-1}-\tilde{A}^{-1}\|_{op}$ $\displaystyle=~{}$
$\displaystyle\|A^{-1}(\tilde{A}-A)\tilde{A}^{-1}\|_{op}$
$\displaystyle\leq~{}$
$\displaystyle\|A-\tilde{A}\|_{op}\|A^{-1}\|_{op}\|\tilde{A}^{-1}\|_{op}$
$\displaystyle\leq~{}$ $\displaystyle(m_{*}n)^{-2}\|A-\tilde{A}\|_{op}.$

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

$\displaystyle u^{T}(S-\tilde{S})u$
$\displaystyle=u^{T}(\sqrt{S}-\sqrt{\tilde{S}})\sqrt{S}u+u^{T}\sqrt{\tilde{S}}(\sqrt{S}-\sqrt{\tilde{S}})u$
$\displaystyle=du^{T}\sqrt{S}u+du^{T}\sqrt{\tilde{S}}u$
$\displaystyle=du^{T}(\sqrt{S}+\sqrt{\tilde{S}})u.$

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

$\displaystyle\|\sqrt{S}-\sqrt{\tilde{S}}\|_{op}=\frac{|u^{T}(S-\tilde{S})u|}{u^{T}(\sqrt{S}+\sqrt{\tilde{S}})u}\leq\frac{\|S-\tilde{S}\|_{op}}{\lambda_{\min}(\sqrt{S}+\sqrt{\tilde{S}})}\leq\frac{\|S-\tilde{S}\|_{op}}{2/\sqrt{m^{*}n}}.$

Combining the above Lipschitz results, we have

$\displaystyle\|(G^{T}G)^{-1/2}-(\tilde{G}^{T}\tilde{G})^{-1/2}\|_{op}\leq(m^{*}n)^{1/2}(m_{*}n)^{-2}(m^{*}n)^{1/2}\|G-\tilde{G}\|_{op}=\frac{m^{*}}{m_{*}^{2}}n^{-1}\|G-\tilde{G}\|_{op}.$

It immediately follows that

$\displaystyle\|(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

$\displaystyle\|G(G^{T}G)^{-1/2}-\tilde{G}(\tilde{G}^{T}\tilde{G})^{-1/2}\|_{op}$
$\displaystyle\leq~{}$
$\displaystyle\|G-\tilde{G}\|_{op}\|(G^{T}G)^{-1/2}\|_{op}+\|\tilde{G}\|_{op}\|(G^{T}G)^{-1/2}-(\tilde{G}^{T}\tilde{G})^{-1/2}\|_{op}$
$\displaystyle\leq~{}$
$\displaystyle\|G-\tilde{G}\|_{op}(m_{*}n)^{-1/2}+(m^{*}n)^{1/2}L_{1}n^{-1}\|G-\tilde{G}\|_{op}$

Hence,

$\displaystyle\|G(G^{T}G)^{-1/2}-\tilde{G}(\tilde{G}^{T}\tilde{G})^{-1/2}\|_{F}$
$\displaystyle\leq~{}$
$\displaystyle\sqrt{K}\bigl{(}m_{*}^{-1/2}+(m^{*})^{1/2}L_{1}\bigr{)}n^{-1/2}\|G-\tilde{G}\|_{F},$

where $L_{2}=\sqrt{K}\bigl{(}m_{*}^{-1/2}+(m^{*})^{1/2}L_{1}\bigr{)}$. ∎

###### Proof of Lemma S5.5.

Recall the KKT conditions $\sum_{l=1}^{n}z_{l}g_{l}^{T}=\bm{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

$\displaystyle\dot{z}_{l}$ $\displaystyle=\frac{\partial z_{l}}{\partial
z_{ij}}=I(l=i)e_{j},$ $\displaystyle\dot{g}_{l}$ $\displaystyle=\frac{\partial
g_{l}}{\partial z_{ij}}=\frac{\partial
g_{l}}{\partial\hat{A}^{\top}z_{l}}\frac{\partial\hat{A}^{\top}z_{l}}{\partial
z_{ij}}=H_{l}[\dot{A}^{\top}z_{l}+I(l=i)\hat{A}^{\top}e_{j}].$

Thus, differentiating the KKT conditions w.r.t. $z_{ij}$ by the product rule
gives

$\sum_{l=1}^{n}\bigl{[}I(l=i)e_{j}g_{l}^{T}+x_{l}\bigl{(}\dot{A}^{\top}z_{l}+I(l=i)\hat{A}^{\top}e_{j}\bigr{)}^{T}H_{l}\bigr{]}=0.$

That is,

$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}+(H_{i}\hat{A}^{T}e_{j}\otimes
z_{i})=-\sum_{l=1}^{n}(H_{l}\otimes
z_{l}z_{l}^{T})\operatorname{vec}(\dot{A}).$

With $M=[\sum_{l=1}^{n}(H_{l}\otimes z_{l}z_{l}^{T})]^{-1}$, we obtain

$\operatorname{vec}(\dot{A})=-M[g_{i}\otimes
e_{j}+(H_{i}\hat{A}^{T}e_{j}\otimes z_{i})].$

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

$\displaystyle\dot{g}_{l}$ $\displaystyle=(H_{l}\otimes
z_{l}^{T})\operatorname{vec}(\dot{A})+I(l=i)H_{l}\hat{A}^{T}e_{j}$
$\displaystyle=-(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}.$

Thus,

$\displaystyle\dot{g}_{i}$ $\displaystyle=-(H_{i}\otimes
z_{i}^{T})M[g_{i}\otimes e_{j}+(H_{i}\hat{A}^{T}e_{j}\otimes
z_{i})]+H_{i}\hat{A}^{T}e_{j}$ $\displaystyle=-(H_{i}\otimes
z_{i}^{T})M(H_{i}\hat{A}^{T}e_{j}\otimes
z_{i})+H_{i}\hat{A}^{T}e_{j}-(H_{i}\otimes z_{i}^{T})M(g_{i}\otimes e_{j})$
$\displaystyle=[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})$
$\displaystyle=V_{i}\hat{A}^{T}e_{j}-(H_{i}\otimes z_{i}^{T})M(g_{i}\otimes
e_{j}),$

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

$\displaystyle\pdv{e_{i}^{T}G(G^{T}G)^{-1/2}}{z_{ij}}$ $\displaystyle=~{}$
$\displaystyle\pdv{g_{i}^{T}}{z_{ij}}(G^{T}G)^{-1/2}+e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{ij}}$
$\displaystyle=~{}$ $\displaystyle[V_{i}\hat{A}^{T}e_{j}-(H_{i}\otimes
z_{i}^{T})M(g_{i}\otimes
e_{j})]^{T}(G^{T}G)^{-1/2}+e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{ij}}$
$\displaystyle=~{}$ $\displaystyle
e_{j}^{T}\hat{A}V_{i}(G^{T}G)^{-1/2}+\Bigl{[}-(g_{i}^{T}\otimes
e_{j}^{T})M(H_{i}\otimes
z_{i})(G^{T}G)^{-1/2}+e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{ij}}\Bigr{]}.$

With $V=\sum_{i=1}^{n}V_{i}$, we further have

$\displaystyle\sum_{i=1}^{n}\pdv{e_{i}^{T}G(G^{T}G)^{-1/2}}{z_{ij}}$
$\displaystyle=~{}$ $\displaystyle
e_{j}^{T}\hat{A}V_{i}^{T}(G^{T}G)^{-1/2}+\sum_{i=1}^{n}\Bigl{[}-(g_{i}^{T}\otimes
e_{j}^{T})M(H_{i}\otimes
z_{i})(G^{T}G)^{-1/2}+e_{i}^{T}G\pdv{(G^{T}G)^{-1/2}}{z_{ij}}\Bigr{]}.$

∎

##  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,

(S6.1)
$L_{i}(u)=-\sum_{k=1}^{K+1}\mathsf{y}_{ik}e_{k}^{T}Qu+\log\sum_{k^{\prime}=1}^{K+1}\exp(e_{k^{\prime}}^{T}Qu),\qquad\forall
u\in\mathbb{R}^{K}.$

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

$\displaystyle U$
$\displaystyle=\bigl{\\{}(Y,X)\in\mathbb{R}^{n\times(K+1)}\times\mathbb{R}^{n\times
p}:\hat{\mathsf{B}}\text{
exists},\|X\hat{\mathsf{B}}(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})\|_{F}^{2}<n\tau\bigr{\\}},$
$\displaystyle U_{y}$
$\displaystyle=\Bigl{\\{}Y\in\mathbb{R}^{n\times(K+1)}:\sum_{i=1}^{n}I(\mathsf{y}_{ik}=1)\geq\gamma
n\text{ for all }k\in[K+1]\Bigr{\\}}.$

######  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 KI_{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];(ii)~{}|S_{l}|=K+1;(iii)~{}\sum_{i\in
S_{l}}\mathsf{y}_{ik}=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}}\|QM^{T}e_{i}\|^{2}\leq\sum_{i=1}^{n}\|QM^{T}e_{i}\|^{2}=\|QM^{T}\|_{F}^{2}<n\tau.$

It follows that at most $\alpha n$ many of $l\in\\{1,2,...,\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,...,\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)=\mathsf{L}_{i}(Qu)$, we have
$\nabla L_{i}(u)=Q^{T}\nabla\mathsf{L}_{i}(Qu)$. Thus,

$\displaystyle\nabla
L_{i}(M^{T}e_{i})=Q^{T}\nabla\mathsf{L}_{i}(QM^{T}e_{i})=Q^{T}(-\mathsf{y}_{i}+\mathsf{p}_{i}),$

where $\mathsf{p}_{i}\in\mathbb{R}^{K+1}$ and its $k$-th entry satisfying

(S6.2)
$\mathsf{p}_{ik}=\frac{\exp(e_{k}^{T}QM^{T}e_{i})}{\sum_{k^{\prime}=1}^{K+1}\exp(e_{k^{\prime}}^{T}QM^{T}e_{i})}\in[c,1-c],$

for some constant $c\in(0,1)$ depending on $(\tau,\alpha,K)$ only. Therefore,

$\displaystyle n^{-1}\sum_{i=1}^{n}\nabla L_{i}(M^{T}e_{i})\nabla
L_{i}(M^{T}e_{i})^{T}$ $\displaystyle=~{}$ $\displaystyle
n^{-1}Q^{T}\sum_{i=1}^{n}(\mathsf{y}_{i}-\mathsf{p}_{i})(\mathsf{y}_{i}-\mathsf{p}_{i})^{T}Q$
$\displaystyle\succeq~{}$ $\displaystyle n^{-1}Q^{T}\sum_{l=1}^{\gamma
n}\sum_{i\in
S_{l}}(\mathsf{y}_{i}-\mathsf{p}_{i})(\mathsf{y}_{i}-\mathsf{p}_{i})^{T}Q$
$\displaystyle\succeq~{}$ $\displaystyle n^{-1}Q^{T}\sum_{l\in
L^{*}}\sum_{i\in
S_{l}}(\mathsf{y}_{i}-\mathsf{p}_{i})(\mathsf{y}_{i}-\mathsf{p}_{i})^{T}Q$
$\displaystyle:=~{}$ $\displaystyle n^{-1}Q^{T}\sum_{l\in
L^{*}}A_{l}^{T}A_{l}Q,$

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}QM^{T}e_{i})}{\sum_{k^{\prime}=1}^{K+1}\exp(e_{k^{\prime}}^{T}QM^{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 $\bm{1}$. Therefore,

$(I_{K+1}-\mathsf{P}_{l})QQ^{T}=(I_{K+1}-\mathsf{P}_{l})(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})=(I_{K+1}-\mathsf{P}_{l}).$

It follows that

$K=\operatorname{rank}(I_{K+1}-\mathsf{P}_{l})=\operatorname{rank}((I_{K+1}-\mathsf{P}_{l})QQ^{T})\leq\operatorname{rank}((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

$\displaystyle\frac{1}{n}\lambda_{\min}(\sum_{i=1}^{n}\nabla
L_{i}(M^{T}e_{i})\nabla L_{i}(M^{T}e_{i})^{T})$ $\displaystyle~{}\geq$
$\displaystyle\frac{1}{n}\sum_{l\in
L^{*}}\lambda_{\min}(Q^{T}A_{l}^{T}A_{l}Q)$ $\displaystyle~{}\geq$
$\displaystyle\frac{1}{n}\sum_{l\in L^{*}}\min_{a\in
S^{K-1}}a^{T}Q^{T}A_{l}^{T}A_{l}Qa$ $\displaystyle~{}\geq$
$\displaystyle\frac{1}{n}|L^{*}|\min_{a\in
S^{K-1},\mathsf{P}\in\mathcal{P}}a^{T}Q^{T}(I_{K+1}-\mathsf{P})^{T}(I_{K+1}-\mathsf{P})Qa$
$\displaystyle~{}\geq$
$\displaystyle\frac{\gamma}{2}a_{*}^{T}Q^{T}(I_{K+1}-\mathsf{P}_{*})^{T}(I_{K+1}-\mathsf{P}_{*})Qa_{*}$
$\displaystyle~{}\geq$
$\displaystyle\frac{\gamma}{2}a_{*}^{T}Q^{T}(I_{K+1}-\mathsf{P}_{*})^{T}(I_{K+1}-\mathsf{P}_{*})Qa_{*}$
$\displaystyle~{}:=$ $\displaystyle m_{*},$

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
$(\mathsf{y}_{i}-\mathsf{p}_{i})(\mathsf{y}_{i}-\mathsf{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}(M^{T}e_{i})\nabla
L_{i}(M^{T}e_{i})^{T}\|_{op}=\|Q^{T}\sum_{i=1}^{n}(\mathsf{y}_{i}-\mathsf{p}_{i})(\mathsf{y}_{i}-\mathsf{p}_{i})^{T}Q\|_{op}\leq
nK.$

∎

######  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

$\displaystyle\nabla^{2}L_{i}(u)=Q^{T}\nabla^{2}\mathsf{L}_{i}(Qu)Q,$

where
$\nabla^{2}\mathsf{L}_{i}(Qu)=\operatorname{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}Qu)}{\sum_{k^{\prime}=1}^{K+1}\exp(e_{k^{\prime}}^{T}Qu)}$
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

$\displaystyle a^{T}\nabla^{2}L_{i}(u)a=$ $\displaystyle
a^{T}Q^{T}[\operatorname{diag}(\mathsf{p}_{i})-\mathsf{p}_{i}\mathsf{p}_{i}^{T}]Qa$
$\displaystyle\leq$ $\displaystyle
a^{T}Q^{T}\operatorname{diag}(\mathsf{p}_{i})Qa$ $\displaystyle\leq$
$\displaystyle\|Q^{T}\operatorname{diag}(\mathsf{p}_{i})Q\|_{op}$
$\displaystyle\leq$ $\displaystyle 1,$

where the last inequality uses $\|Q\|_{op}\leq 1$ and $\mathsf{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
$\|Qu\|_{\infty}\leq r$, we have

$\mathsf{p}_{ik}=\frac{\exp(e_{k}^{T}Qu)}{\sum_{k^{\prime}=1}^{K+1}\exp(e_{k^{\prime}}^{T}Qu)}\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 $\bm{1}^{T}\eta=0$ and

$\displaystyle a^{T}\bigl{[}\nabla^{2}L_{i}(u)\bigr{]}a=$ $\displaystyle
a^{T}Q^{T}[\operatorname{diag}(\mathsf{p}_{i})-\mathsf{p}_{i}\mathsf{p}_{i}^{T}]Qa$
$\displaystyle=$
$\displaystyle\eta^{T}[\operatorname{diag}(\mathsf{p}_{i})-\mathsf{p}_{i}\mathsf{p}_{i}^{T}]\eta$
$\displaystyle=$
$\displaystyle\sum_{k=1}^{K+1}\mathsf{p}_{ik}\eta_{k}^{2}-(\sum_{k=1}^{K+1}\mathsf{p}_{ik}\eta_{k})^{2}$
$\displaystyle>$
$\displaystyle\sum_{k}\mathsf{p}_{ik}\eta_{k}^{2}-\sum_{k}\mathsf{p}_{ik}\eta_{k}^{2}\sum_{k}\mathsf{p}_{ik}$
$\displaystyle=$
$\displaystyle\sum_{k}\mathsf{p}_{ik}\eta_{k}^{2}(1-\sum_{k}\mathsf{p}_{ik})$
$\displaystyle=$ $\displaystyle 0,$

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 $\bm{1}^{T}\eta=0$.

Let
$\mathcal{H}=\\{Q^{T}(\operatorname{diag}(\mathsf{p})-\mathsf{p}\mathsf{p}^{T})Q:\mathsf{p}\in[c,1-c]^{K+1}\\}$,
then $\mathcal{H}$ is compact and

$\displaystyle\min_{u\in\mathbb{R}^{K},\|Qu\|_{\infty}\leq
r}\lambda_{\min}(\nabla^{2}L_{i}(u))\geq\min_{a\in
S^{K-1},H\in\mathcal{H}}a^{T}Ha=a_{*}^{T}H_{*}a_{*}>0$

for some $a_{*}\in S^{K-1}$ and $H_{*}\in\mathcal{H}$. Therefore,

$\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. ∎

###  S6.2. Lipschitz conditions

We first restate the definitions of following sets,

$\displaystyle U$
$\displaystyle=\bigl{\\{}(Y,X)\in\mathbb{R}^{n\times(K+1)}\times\mathbb{R}^{n\times
p}:\hat{\mathsf{B}}\text{
exists},\|X\hat{\mathsf{B}}(I_{K+1}-\tfrac{\bm{1}\bm{1}^{T}}{K+1})\|_{F}^{2}<n\tau\bigr{\\}},$
$\displaystyle U_{y}$
$\displaystyle=\Bigl{\\{}Y\in\mathbb{R}^{n\times(K+1)}:\sum_{i=1}^{n}I(\mathsf{y}_{ik}=1)\geq\gamma
n\text{ for all }k\in[K+1]\Bigr{\\}}$

######  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

(S6.3) $\displaystyle U_{x}=\bigl{\\{}X\in\mathbb{R}^{n\times
p}:\min_{I\in\mathcal{I}}\lambda_{\min}(\tfrac{\Sigma^{-1/2}X^{T}P_{I}X\Sigma^{-1/2}}{{n}})\geq\phi_{*}^{2},\tfrac{\|X\Sigma^{-1/2}\|_{op}}{\sqrt{n}}\leq\phi^{*}\bigr{\\}}$

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

  1. (i)

$\mathbb{P}((Y,X)\in U^{*})\to 1$ as $n,p\to\infty$.

  2. (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}\bigl{(}\min_{I\in\mathcal{I}}\lambda_{\min}(\tfrac{\Sigma^{-1/2}X^{T}P_{I}X\Sigma^{-1/2}}{{n}})\geq\phi_{*}^{2}\bigr{)}\to
1$

for some positive constant $\phi_{*}$ depending on $(\delta,\alpha)$ only.
Furthermore, [Davidson and Szarek, 2001, Theorem II.13] implies

$\mathbb{P}(\tfrac{\|X\Sigma^{-1/2}\|_{op}}{\sqrt{n}}\leq\phi^{*})\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

$\displaystyle\langle X\hat{B}-\tilde{X}\tilde{B},G-\tilde{G}\rangle=$
$\displaystyle\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$ $\displaystyle=$
$\displaystyle-\langle(X-\tilde{X})(\hat{B}-\tilde{B}),G\rangle+\langle(X-\tilde{X})\hat{B},G-\tilde{G}\rangle.$

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,

$\displaystyle\langle X\hat{B}-\tilde{X}\tilde{B},G-\tilde{G}\rangle=$
$\displaystyle\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$
$\displaystyle\geq$ $\displaystyle\sum_{i=1}^{n}\langle\nabla
L_{i}(\hat{B}^{T}x_{i})-\nabla L_{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$
$\displaystyle=$ $\displaystyle\|G-\tilde{G}\|_{F}^{2}.$

If $(Y,X),(Y,\tilde{X})\in U$, we have
$\|X\hat{B}Q^{T}\|_{F}^{2}+\|\tilde{X}\tilde{B}Q^{T}\|_{F}^{2}\leq 2n\tau$.
That is,

$\sum_{i=1}^{n}\bigl{(}\|Q\hat{B}^{T}x_{i}\|^{2}+\|Q\tilde{B}^{T}\tilde{x}_{i}\|^{2}\bigr{)}\leq
2n\tau.$

Define the index set

$I=\\{i\in[n]:\|Q\hat{B}^{T}x_{i}\|^{2}+\|Q\tilde{B}^{T}\tilde{x}_{i}\|^{2}\leq\tfrac{2\tau}{\alpha}\\},$

then we have $|I|\geq(1-\alpha)n$ by Markov’s inequality. Thus, for all $i\in
I$, we have
$\|Q\hat{B}^{T}x_{i}\|_{\infty}\vee\|Q\tilde{B}^{T}\tilde{x}_{i}\|_{\infty}\leq\sqrt{\tfrac{2\tau}{\alpha}}$.

Applying Lemma S6.2 with $r=\sqrt{\tfrac{2\tau}{\alpha}}$ gives

$\min_{\|Qu\|_{\infty}\leq\sqrt{\tfrac{2\tau}{\alpha}}}\nabla^{2}L_{i}(u)\succeq
c_{*}I_{K},$

where $c_{*}$ is a constant depending on $(K,\tau,\alpha)$. Therefore,

$\displaystyle c_{*}\|P_{I}(X\hat{B}-\tilde{X}\tilde{B})\|_{F}^{2}=$
$\displaystyle c_{*}\sum_{i\in
I}\|\hat{B}^{T}x_{i}-\tilde{B}^{T}\tilde{x}_{i}\|^{2}$ $\displaystyle\leq$
$\displaystyle\sum_{i\in
I}\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$
$\displaystyle\leq$
$\displaystyle\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$
$\displaystyle=$ $\displaystyle\langle
X\hat{B}-\tilde{X}\tilde{B},G-\tilde{G}\rangle$ $\displaystyle=$
$\displaystyle-\langle(X-\tilde{X})(\hat{B}-\tilde{B}),G\rangle+\langle(X-\tilde{X})\hat{B},G-\tilde{G}\rangle.$

We next bound the first line from below by expanding the squares,

$\displaystyle\|P_{I}(X\hat{B}-\tilde{X}\tilde{B})\|_{F}^{2}$
$\displaystyle=\|P_{I}\tilde{X}(\hat{B}-\tilde{B})+P_{I}(X-\tilde{X})\hat{B}\|_{F}^{2}$
$\displaystyle\geq\|P_{I}\tilde{X}(\hat{B}-\tilde{B})\|_{F}^{2}+2\langle
P_{I}\tilde{X}(\hat{B}-\tilde{B}),P_{I}(X-\tilde{X})\hat{B}\rangle$
$\displaystyle\geq
n\phi_{*}^{2}\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}^{2}+2\langle\tilde{X}(\hat{B}-\tilde{B}),P_{I}(X-\tilde{X})\hat{B}\rangle,$

where in the last inequality we use the constant $\phi^{*}$ in (S6.3).
Therefore, we obtain

$\displaystyle c_{*}\phi_{*}^{2}n\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}^{2}$
$\displaystyle\leq$
$\displaystyle-\langle(X-\tilde{X})(\hat{B}-\tilde{B}),G\rangle+\langle(X-\tilde{X})\hat{B},G-\tilde{G}\rangle-2c_{*}\langle\tilde{X}(\hat{B}-\tilde{B}),P_{I}(X-\tilde{X})\hat{B}\rangle.$

Together with the inequality that $\|G-\tilde{G}\|_{F}^{2}\leq\langle
X\hat{B}-\tilde{X}\tilde{B},G-\tilde{G}\rangle,$ we obtain

$\displaystyle
c_{*}\phi_{*}^{2}n\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}^{2}+\|G-\tilde{G}\|_{F}^{2}$
$\displaystyle\leq$
$\displaystyle-2\langle(X-\tilde{X})(\hat{B}-\tilde{B}),G\rangle+2\langle(X-\tilde{X})\hat{B},G-\tilde{G}\rangle-2c_{*}\langle\tilde{X}(\hat{B}-\tilde{B}),P_{I}(X-\tilde{X})\hat{B}\rangle$
$\displaystyle\leq$
$\displaystyle(4+2c_{*}\phi^{*})\|(X-\tilde{X})\Sigma^{-1/2}\|_{op}\bigl{(}\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}\vee\frac{\|G-\tilde{G}\|_{F}}{\sqrt{n}}\bigr{)}\bigl{(}\|\Sigma^{1/2}\hat{B}\|_{F}\vee\frac{\|G\|_{op}}{\sqrt{n}}\bigr{)}\sqrt{n},$

where we bound
$\langle\tilde{X}(\hat{B}-\tilde{B}),P_{I}(X-\tilde{X})\hat{B}\rangle$ by
definition of $\phi^{*}$,

$\displaystyle\langle\tilde{X}(\hat{B}-\tilde{B}),P_{I}(X-\tilde{X})\hat{B}\rangle$
$\displaystyle=$
$\displaystyle\langle\Sigma^{1/2}(\hat{B}-\tilde{B}),\Sigma^{-1/2}\tilde{X}^{T}P_{I}(X-\tilde{X})\hat{B}\rangle$
$\displaystyle\leq$
$\displaystyle\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}\|P_{I}\tilde{X}\Sigma^{-1/2}\|_{op}\|(X-\tilde{X})\Sigma^{-1/2}\|_{op}\|\Sigma^{1/2}\hat{B}\|_{F}$
$\displaystyle\leq$
$\displaystyle\sqrt{n}\phi^{*}\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}\|(X-\tilde{X})\Sigma^{-1/2}\|_{op}\|\Sigma^{1/2}\hat{B}\|_{F}.$

Now we derive a bound of the form
$\|\Sigma^{1/2}\hat{B}\|_{F}\lesssim\|G\|_{F}/\sqrt{n}$. To this end, since
$\phi_{*}\|\Sigma^{1/2}\hat{B}\|_{F}\leq\|P_{I}X\hat{B}\|_{F}/\sqrt{n}\leq\|X\hat{B}\|_{F}/\sqrt{n}=\|X\hat{B}Q^{T}\|_{F}/\sqrt{n}\leq\sqrt{\tau}$.

Applying Lemma S6.1 to $M=X\hat{B}$, we have
$\frac{1}{n}\sum_{i=1}^{n}g_{i}g_{i}^{T}\succeq m_{*}I_{K}$. Therefore,

$\displaystyle\frac{1}{n}\|G\|_{F}^{2}=\frac{1}{n}\sum_{i=1}^{n}\|g_{i}\|^{2}=\frac{1}{n}\sum_{i=1}^{n}\operatorname{Tr}(g_{i}g_{i}^{T})\geq
Km_{*}.$

This implies that

$\displaystyle\phi_{*}^{2}\|\Sigma^{1/2}\hat{B}\|_{F}^{2}\leq\tau\leq\frac{\tau}{Km_{*}(I)}\|G\|_{F}^{2}/n.$

In conclusion, if $\\{(Y,X),(Y,\tilde{X})\\}\subset U^{*}$ then

$\displaystyle\sqrt{n}\|\Sigma^{1/2}(\hat{B}-\tilde{B})\|_{F}+\|G-\tilde{G}\|_{F}\leq
Cn^{-1/2}\|(X-\tilde{X})\Sigma^{-1/2}\|_{op}\|G\|_{F}\leq
CK\|(X-\tilde{X})\Sigma^{-1/2}\|_{op},$

where $C$ is a constant depending on $(K,\gamma,\tau,\alpha)$ only. Note that
$\|G\|_{F}\leq\sqrt{nK}$ since all entries of $G$ are in $[-1,1]$. ∎

######  Lemma S6.4.

If $p/n\leq\delta^{-1}<(1-\alpha)$ for some $\alpha\in(0,1)$ and $\delta>1$.
If $(Y,X)\in U$ and $X\in U_{x}$, where $U_{x}$ is defined in Lemma S6.3, we
have

$\frac{1}{n}\sum_{i=1}^{n}H_{i}\otimes(x_{i}x_{i}^{T})\succeq
c_{1}(I_{K}\otimes\Sigma),$

where $c_{1}$ is a positive constant depending on $(K,\tau,\alpha,\phi_{*})$
only.

###### Proof of Lemma S6.4.

If $(Y,X)\in U$, we have $\|X\hat{B}Q^{T}\|_{F}^{2}\leq n\tau$. Define the
index set

$I=\\{i\in[n]:\|Q\hat{B}^{T}x_{i}\|\leq\tfrac{\tau}{\alpha}\\},$

then we have $|I|\geq(1-\alpha)n$ by Markov’s inequality. Therefore, for any
$i\in I$, $\|Q\hat{B}^{T}x_{i}\|_{\infty}\leq\tfrac{\tau}{\alpha}$. Applying
Lemma S6.2 with $u=\hat{B}^{T}x_{i}$ and $r=\tfrac{\tau}{\alpha}$, we have for
any $i\in I$, $H_{i}=\nabla^{2}L_{i}(\hat{B}^{T}x_{i})\succeq c_{*}I_{K}$ for
some positive constant $c_{*}$ depending on $(K,\tau,\alpha)$ only. Therefore,
if $(Y,X)\in U$ and $X\in U_{x}$, we have

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}H_{i}\otimes(x_{i}x_{i}^{T})\succeq~{}$
$\displaystyle\frac{1}{n}\sum_{i\in I}H_{i}\otimes(x_{i}x_{i}^{T})$
$\displaystyle\succeq~{}$ $\displaystyle c_{*}\frac{1}{n}\sum_{i\in
I}I_{K}\otimes(x_{i}x_{i}^{T})$ $\displaystyle=~{}$ $\displaystyle
c_{*}(I_{K}\otimes\tfrac{X^{T}P_{I}X}{n})$ $\displaystyle\succeq~{}$
$\displaystyle c_{*}\phi_{*}(I_{K}\otimes\Sigma)$ $\displaystyle=~{}$
$\displaystyle c_{1}(I_{K}\otimes\Sigma),$

where $P_{I}=\sum_{i\in I}e_{i}e_{i}^{T}$ and $c_{1}$ is a positive constant
depending on $(K,\tau,\alpha,\phi_{*})$ only. ∎

###  S6.3. Proof of Theorem S3.1

The proof of Theorem S3.1 is a direct consequence of Theorem S5.1 by noting

$\sqrt{n}\Omega_{jj}^{-1/2}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}g_{i}g_{i}^{T}\Bigr{)}^{-1/2}\Bigl{(}\frac{1}{n}\sum_{i=1}^{n}V_{i}\Bigr{)}\hat{B}^{T}e_{j}=\Omega_{jj}^{-1/2}(G^{T}G)^{-1/2}V\hat{B}^{T}e_{j},$

which is a consequence of the identities $G=\sum_{i=1}^{n}e_{i}g_{i}^{T}$ and
$V=\sum_{i=1}^{n}V_{i}$.

It thus remains to verify the conditions (1) and (2) in Theorem S5.1 from the
assumptions in Theorem S3.1.

Applying Lemma S6.3 with $\alpha$ chosen as $1-\delta^{-1/2}$, we have for
$\\{(Y,X),(Y,\tilde{X})\\}\subset U^{*}$,

$\|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,\delta)$ only.

Apply Lemma S6.4 with the same $\alpha=1-\delta^{-1/2}$, we have for $(Y,X)\in
U^{*}$,

$\frac{1}{n}\sum_{i=1}^{n}H_{i}\otimes(x_{i}x_{i}^{T})\succeq
c_{*}(I_{K}\otimes\Sigma),$

where $c_{*}$ is a positive constant depending on $(K,\tau,\delta)$ only.

Applying Lemma S6.1 with $M=X\hat{B}$, we have for $(Y,X)\in U^{*}$,

$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 KI_{K},$

where $m_{*}$ is a positive constant depending on $(K,\gamma,\tau)$ only.
Therefore, the conditions (1) and (2) in Theorem S5.1 hold when the
multinomial logistic loss is used. This completes the proof of Theorem S3.1.

##  S7. Other proof

###  S7.1. Proof of Equation 1.12 (Classical asymptotic theory with fixed $p$)

Here we provide a derivation of the asymptotic distribution of MLE under
classical setting, where $p$ is fixed and $n$ tends to infinity.

We first calculate the Fisher information matrix of the multinomial logistic
log-odds model (1.8) with covariate $x\sim N(\bm{0},\Sigma)$ and response
$\mathsf{y}\in\\{0,1\\}^{K+1}$ one-hot encoded satisfying
$\sum_{k=1}^{K+1}\mathsf{y}_{k}=1$. Note that the model (1.8) can be rewritten
as

$\displaystyle\mathbb{P}(\mathsf{y}_{k}=1|x)$
$\displaystyle=\frac{\exp(x^{T}A^{*}e_{k})}{1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}A^{*}e_{k^{\prime}})},\quad\forall
k\in\\{1,...,K\\}$ $\displaystyle\mathbb{P}(\mathsf{y}_{K+1}=1|x)$
$\displaystyle=\frac{1}{1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}A^{*}e_{k^{\prime}})}.$

The likelihood function of a parameter $A\in\mathbb{R}^{p\times K}$ is

$\displaystyle L(A)$
$\displaystyle=\prod_{k=1}^{K}\Bigl{[}\frac{\exp(x^{T}Ae_{k})}{1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}Ae_{k^{\prime}})}\Bigr{]}^{\mathsf{y}_{k}}\Bigl{[}\frac{1}{1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}Ae_{k^{\prime}})}\Bigr{]}^{\mathsf{y}_{K+1}}$
$\displaystyle=\prod_{k=1}^{K}[\exp(x^{T}Ae_{k})]^{\mathsf{y}_{k}}\frac{1}{1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}Ae_{k^{\prime}})},$

where we used $\sum_{k=1}^{K+1}\mathsf{y}_{k}=1$. Thus, the log-likelihood
function is

$\displaystyle\ell(A)$
$\displaystyle=\sum_{k=1}^{K}\mathsf{y}_{k}x^{T}Ae_{k}-\log\bigl{[}1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}Ae_{k^{\prime}})\bigr{]}.$

It is more convenient to calculate the Fisher information matrix on the vector
space $\mathbb{R}^{pK}$ instead of the matrix space $\mathbb{R}^{p\times K}$.
To this end, let $\theta=\operatorname{vec}(A^{T})$, then
$x^{T}Ae_{k}=e_{k}^{T}A^{T}x=(x^{T}\otimes
e_{k}^{T})\operatorname{vec}(A^{T})=(x^{T}\otimes e_{k}^{T})\theta$, and the
log-likelihood function parameterized by $\theta$ is

$\ell(\theta)=\sum_{k=1}^{K}y_{k}(x^{T}\otimes
e_{k}^{T})\theta-\log[1+\sum_{k^{\prime}=1}^{K}\exp((x^{T}\otimes
e_{k^{\prime}}^{T})\theta)].$

By multivariate calculus, we obtain the Fisher information matrix evaluated at
$\theta^{*}=\operatorname{vec}(A^{*}{}^{T})$,

$\displaystyle\mathcal{I}(\theta^{*})$
$\displaystyle=-\mathbb{E}\Bigl{[}\frac{\partial}{\partial\theta}\frac{\partial\ell(\theta)}{\partial\theta^{T}}\Bigr{]}\Big{|}_{\theta=\theta^{*}}$
$\displaystyle=\mathbb{E}[(xx^{T})\otimes(\operatorname{diag}(\pi^{*})-\pi^{*}\pi^{*}{}^{T})],$

where $\pi^{*}\in\mathbb{R}^{K}$ with $k$-th entry
$\pi_{k}^{*}=\frac{\exp(x^{T}A^{*}e_{k})}{1+\sum_{k^{\prime}=1}^{K}\exp(x^{T}A^{*}e_{k^{\prime}})}$
for each $k\in[K]$.

From classical maximum likelihood theory, for instance [Van der Vaart, 1998,
Chapter 5], we have

$\displaystyle\sqrt{n}(\hat{\theta}-\theta^{*}){\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},\mathcal{I}_{\theta^{*}}^{-1}),$

where $\hat{\theta}=\operatorname{vec}(\hat{A}^{T})$ and $\hat{A}$ is the MLE
of $A^{*}$. Furthermore, if the $j$-th covariate is independent of the
response, we know $e_{j}^{T}A^{*}=\bm{0}^{T}$, then

$\displaystyle\sqrt{n}\hat{A}^{T}e_{j}=\sqrt{n}(\hat{A}^{T}e_{j}-A^{*}{}^{T}e_{j})=\sqrt{n}(e_{j}^{T}\otimes
I_{K})(\hat{\theta}-\theta^{*}){\smash{\xrightarrow{~{}\mathrm{d}~{}}}}N(\bm{0},S_{j}),$

where

$\displaystyle S_{j}$ $\displaystyle=(e_{j}^{T}\otimes
I_{K})\mathcal{I}_{\theta^{*}}^{-1}(e_{j}\otimes I_{K})$
$\displaystyle=e_{j}^{T}{\mathrm{cov}}(x)^{-1}e_{j}[\mathbb{E}(\operatorname{diag}(\pi^{*})-\pi^{*}\pi^{*}{}^{T})]^{-1}$

holds by the independence between the $j$-th covariate and the response under
$H_{0}$ in (1.10). This completes the proof of Equation 1.12.