# Prediction Sets for High-Dimensional Mixture of Experts Models

Adel Javanmard, Simeng Shao, Jacob Bien
Data Sciences and Operations, University of Southern California  A. Javanmard
is partially supported by the Sloan Research Fellowship in mathematics, Adobe
Data Science Faculty Research Awards and the NSF CAREER Award DMS-1844481. J.
Bien was supported in part by NSF CAREER Award DMS-1653017.

###### Abstract

Large datasets make it possible to build predictive models that can capture
heterogenous relationships between the response variable and features. The
mixture of high-dimensional linear experts model posits that observations come
from a mixture of high-dimensional linear regression models, where the mixture
weights are themselves feature-dependent. In this paper, we show how to
construct valid prediction sets for an $\ell_{1}$-penalized mixture of experts
model in the high-dimensional setting. We make use of a debiasing procedure to
account for the bias induced by the penalization and propose a novel strategy
for combining intervals to form a prediction set with coverage guarantees in
the mixture setting. Synthetic examples and an application to the prediction
of critical temperatures of superconducting materials show our method to have
reliable practical performance.

##  1 Introduction

In traditional statistics, we imagine a universal relationship between
variables that holds across an entire population; observations not following
this relationship are dismissed as outliers. However, we know that reality is
more complex, with numerous subpopulations likely exhibiting distinct
behaviors. As datasets grow in size, we become better able to detect and
properly model this heterogeneity. The mixture of regressions model (Quandt &
Ramsey, 1978) is an important tool for extending linear regression to this
heterogeneity-aware setting. For a random response $y\in\mathbb{R}$ and a
random vector of predictors ${\boldsymbol{x}}\in\mathbb{R}^{p}$, we imagine a
latent subgroup membership $z\in\\{1,\ldots,K\\}$ that determines the
conditional distribution of $y$ given ${\boldsymbol{x}}$:

$\displaystyle y|{\boldsymbol{x}},z=k\sim
N({\boldsymbol{x}}^{T}{\boldsymbol{\beta}}_{k},\sigma_{k}^{2}).$ (1)

Making predictions with this model requires estimating for each subgroup a
coefficient vector ${\boldsymbol{\beta}}_{k}$, an error variance
$\sigma_{k}^{2}$, and a group membership probability
$\pi_{k}=\mathbb{P}(z=k)$. The mixture of experts model (MoE, Jordan & Jacobs
1994) is even more flexible, allowing these group membership probabilities to
depend on the predictors as well:

$\displaystyle
z|{\boldsymbol{x}}\sim\text{Multinomial}[\boldsymbol{\pi}({\boldsymbol{x}})]\qquad\text{with}\quad\pi_{k}({\boldsymbol{x}})=\frac{\exp{({\boldsymbol{x}}^{T}{\boldsymbol{\alpha}}_{k})}}{\sum_{\ell=1}^{K}\exp{({\boldsymbol{x}}^{T}{\boldsymbol{\alpha}}_{\ell})}}.$
(2)

This model is expressive enough to capture subpopulations that change in
prevalence depending on the conditions. For example, Hyun et al. (2020)
develop a high-dimensional mixture of experts approach to modeling
phytoplankton subpopulations as a function of environmental covariates in the
ocean. The (log) diameter of the phytoplankton cells within each specific
subpopulation are taken to be Gaussian with mean depending on environmental
covariates (expressed through nonzero values of
${\boldsymbol{\beta}}_{1},\ldots,{\boldsymbol{\beta}}_{K}$); however, the
prevalence of the different subpopulations also depends on these covariates
(expressed through nonzero values of
${\boldsymbol{\alpha}}_{1},\ldots,{\boldsymbol{\alpha}}_{K}$).

When making predictions, it is valuable to be able to quantify one’s level of
uncertainty. In this paper, we develop the machinery necessary to do so in the
context of the high-dimensional mixture of the experts model described above.
In particular, given a sample of $n$ observations from the model in (1)–(2), a
confidence level $q\in(0,1)$, and a new predictor vector of interest
$\boldsymbol{x}_{\textrm{new}}$, we show how to form a properly calibrated
prediction set $\Omega_{q}(\boldsymbol{x}_{\textrm{new}})$. That is, given a
new draw $(\boldsymbol{x}_{\textrm{new}},y_{\textrm{new}})$ from (1)–(2), we
have that

$\displaystyle\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|\boldsymbol{x}_{\textrm{new}})\geq
1-q.$ (3)

0.250.500.750.000.250.500.751.00$t$Probability01230.000.250.500.751.00$t$$y$01230.000.250.500.751.00$t$$y$0.9000.9250.9500.9751.0000.000.250.500.751.00$t$Cover
Rate0.81.01.20.000.250.500.751.00$t$Average LengthcolourGroup 1Group
2linetypeEstimatedTrue

Figure 1: (Top left) True (dashed) and estimated (solid) class probabilities
$\pi_{k}(t)$ and $\hat{\pi}_{k}(t)$. (Bottom left) Data set of $n=100$ points
$\\{(t,y_{t})$ (black dots) with true (dashed) and estimated (solid) mean
functions $\mu_{k}(t)$ and $\hat{\mu}_{k}(t)$. (Middle) Prediction set at
$95\%$ confidence level $\Omega_{0.05}(t)$ for $t\in[0,1]$ (shaded area).
(Right) Monte Carlo estimate of coverage and length of prediction sets
$\Omega_{0.05}(t)$ across $500$ training sets and $1000$ $y_{\textrm{new}}$
per training set.

As an illustration, consider the toy example shown in Figure 1. There are
$K=2$ subpopulations, one following a quadratic relation and the other being
constant (for details of the construction, see Section 4). The upper panel
shows that the mixture weight on the quadratic subpopulation decreases with
increasing $t$. This can also be seen by inspecting the scatterplot and noting
that for large $t$, most of the points are in the constant subpopulation. The
solid lines show our estimated mixture of experts model fit based on the
points in the scatterplot. Now suppose we are about to observe a new point
$y_{\textrm{new}}$ at, say, $t_{\text{new}}=0.5$. Can we form a set
$\Omega_{0.95}(t_{\text{new}})$ that is guaranteed to capture
$y_{\textrm{new}}$ at least 95% of the time? Given that we do not know which
subpopulation $y_{\textrm{new}}$ will be drawn from,
$\Omega_{0.95}(t_{\text{new}})$ will be a union of two intervals. The size and
location of these intervals will depend on our estimates of the subpopulation
means and variances (governed by $t_{\text{new}}$ and our estimates of
${\boldsymbol{\beta}}_{1}$, ${\boldsymbol{\beta}}_{2}$, $\sigma_{1}$, and
$\sigma_{2}$) as well as the estimated mixture weights (governed by
$t_{\text{new}}$ and our estimates of ${\boldsymbol{\alpha}}_{1}$ and
${\boldsymbol{\alpha}}_{2}$). The gray band in the middle panel shows our
constructed $\Omega_{0.95}(t_{\text{new}})$ as we vary $t_{\text{new}}$ from 0
to 1. When the means of the populations are far apart, the prediction set is a
union of two intervals, while the set becomes a single interval when the means
are close to each other. The larger a subpopulation’s mixture weight, the
wider that interval becomes. This “strategy” is reasonable because if it’s
unlikely that a point will fall in a certain subpopulation, one can get away
with less good coverage of that subpopulation (and thereby reduce the overall
size of the prediction set). The rightmost panel shows the result of a
simulation in which we generated 500 training sets, each time constructing
$\Omega_{0.95}(t)$ as a function of $t$; then, we generated 1000
$y_{\textrm{new}}$ at each $t$ and computed the coverage rate (averaging over
the $500\cdot 1000$ repetitions for each $t$). This verifies that our
procedure approximately attains the nominal 95% coverage. The bottom panel
shows the average size of the prediction set. Quite intuitively, the
prediction set is smallest when the two subpopulations are very close to each
other. One observes a bit of overcoverage when the subpopulations are close,
which makes sense since in this situation the interval from one subpopulation
can sometimes cover a point from the opposite subpopulation.

Constructing prediction sets in the context of a high-dimensional mixture of
experts model is challenging. In fact, making any sort of precise statements
about even the most simple mixture of regression models is nontrivial. For
example, much effort has gone into understanding the convergence and
estimation error of the expectation-maximization estimator (Dempster et al.,
1977) used in fitting such models (Yi et al. 2014; Balakrishnan et al. 2017;
Klusowski et al. 2019; Kwon et al. 2019; Kwon & Caramanis 2020) as well as
being able to test whether there are two groups versus one (Zhu & Zhang,
2004). Adding high dimensionality to the study of mixture of regression models
brings additional challenges. Städler et al. (2010) and Yi & Caramanis (2015)
proposed different $\ell_{1}$-regularized maximum likelihood estimators with
accompanying estimation error results. Wang et al. (2014) take this a step
further and develop a truncation-based high-dimensional estimator with both
estimation error results and the ability to construct confidence intervals for
low-dimensional components of the parameter vector. While their results hold
for general latent variable models, their application to mixture of regression
models is more of a proof of concept, with $K=2$, $\sigma_{1}=\sigma_{2}$
assumed known, $\pi_{1}=\pi_{2}=0.5$, and
${\boldsymbol{\beta}}_{1}=-{\boldsymbol{\beta}}_{2}$. Zhang et al. (2020)
provide inference for individual coefficients and differences of the form
${\boldsymbol{\beta}}_{1j}-{\boldsymbol{\beta}}_{2j}$ within the context of
this model using an $\ell_{1}$ penalty. They generalize to the case of unknown
mixture weight $\pi_{1}$,
${\boldsymbol{\beta}}_{1}\neq-{\boldsymbol{\beta}}_{2}$, and an unknown
covariance matrix for $X$ (which they take to be multivariate normal), but
they still assume $K=2$ and that the value $\sigma_{1}=\sigma_{2}$ is known.
While we are able to adopt in part a similar debiased approach, we will
highlight later why their technique, which works in the $K=2$ mixture of
regression setting (which involves a single unknown mixture parameter
$\pi_{1}$) does not easily generalize to our setting of a mixture of experts
model, in which mixture weights depend on the unknown parameter vectors
${\boldsymbol{\alpha}}_{1},\ldots,{\boldsymbol{\alpha}}_{K}\in\mathbb{R}^{p}$.

Furthermore, our interest in forming a prediction set (3) requires the ability
to make inferential statements about
$\boldsymbol{x}_{\textrm{new}}^{T}{\boldsymbol{\beta}}_{k}$, not just low-
dimensional components of ${\boldsymbol{\beta}}_{k}$. In this sense, Cai et
al. (2021) pursue a similar goal in performing inference for individualized
treatment effects
$\boldsymbol{x}_{\textrm{new}}^{T}({\boldsymbol{\beta}}_{1}-{\boldsymbol{\beta}}_{2})$;
however, unlike the mixture of regression setting, the group memberships are
known in their context.

To summarize, to the best of our knowledge, predictive inference in mixture of
expert models has not been addressed in the literature. Furthermore, we
address this problem for general $K$ and in the high-dimensional setting. We
make use of ideas from the debiased lasso literature. The debiasing approach
for constructing confidence intervals for coefficients has been widely used in
linear regression models in high-dimensional settings (Javanmard & Montanari
2014a, b; Zhang & Zhang 2014; Van de Geer et al. 2014; Javanmard & Montanari
2018). In recent years, there has been work on inference for general linear
functions (Cai & Guo 2017; Guo et al. 2021; Javanmard & Lee 2020; Zhu & Bradic
2018). In terms of debiasing in the non-mixture setting, Cai & Guo (2017);
Tripuraneni & Mackey (2019); Athey et al. (2018) proposed bias-corrected
estimators for a single linear regression model while, as we have noted above,
Cai et al. (2021) considers inference for
$\boldsymbol{x}_{\textrm{new}}^{T}({\boldsymbol{\beta}}_{1}-{\boldsymbol{\beta}}_{2})$
in the case of two observed (i.e., non-latent) groups.

Interest in predictive inference has led to an active area of work on
conformal prediction. These approaches are attractive for being distribution-
free and providing finite-sample coverage (see, e.g., Papadopoulos et al.
2002; Vovk et al. 2005; Lei et al. 2018; Romano, Patterson & Candes 2019).
They rely on very general ideas such as the exchangeability of draws from the
distribution. However, the coverage that these conformal methods attain is not
conditional on $\boldsymbol{x}_{\textrm{new}}$ as in (3) but rather holds
marginally over $\boldsymbol{x}_{\textrm{new}}$:

$\displaystyle\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}}))\geq
1-q.$

Indeed, it has been proven that to obtain finite length sets with conditional
coverage, one needs to make stronger assumptions (Vovk, 2012; Lei & Wasserman,
2014). In certain applications, one specifically desires coverage of the form
(3) and assuming the parametric form (1)–(2) can be a small price to pay. For
example, in the oceanographic example of mixture of experts (Hyun et al.,
2020), the mixture of Gaussians structure is visually well-supported.
Prediction sets with conditional coverage are desirable because we would like
to be able to say that for a given set of environmental conditions (e.g., at a
specific temperature and salinity level) our prediction set for the
phytoplankton diameters will have a 95% coverage guarantee. Marginal coverage
would mean that if we make predictions over many randomly sampled
environments, our coverage would average out to 95%. The latter means, for
example, that a procedure could be overconfident (i.e., undercovering) at high
temperatures and underconfident (i.e., overcovering) at low temperatures.
Setting conditional coverage (3) as the goal guards against this undesirable
property.

The rest of the paper is organized as follows. In Section 2, we describe our
approach to constructing prediction sets. This involves estimating parameters
using a penalized expectation-maximization approach (Section 2.1), then using
a debiasing technique on the coefficient vectors from each of the component
distributions (Section 2.2), and finally combining $K$ intervals into a
prediction set in a fashion that maintains proper coverage (Section 2.3). In
Section 3, we provide theoretical guarantees that establish the asymptotic
validity of our constructed prediction sets and provide insight into the
conditions under which we expect nominal coverage to hold. In Section 4 we
investigate the empirical performance of our prediction sets in a variety of
settings. Section 5 shows our sets and evaluates their performance empirically
in predicting the critical temperatures of superconducting materials. We
conclude this section with some notation.

Notation. Throughout the paper, we use $[p]$ for the set of integers
$1,...,p$. We use ${\boldsymbol{e}}_{i}$ to denote the $i$-th standard basis
vector. For a vector ${\boldsymbol{x}}\in\mathbb{R}^{p}$, we denote
$\left\|{\boldsymbol{x}}\right\|_{q}=\left(\sum_{j=1}^{p}|x_{j}|^{q}\right)^{1/q}$
for $q>0$ and $\|{\boldsymbol{x}}\|_{0}=|\textrm{supp}({\boldsymbol{x}})|$,
$\left\|{\boldsymbol{x}}\right\|_{\infty}=\max_{j\in[p]}|x_{j}|$. For a matrix
$\boldsymbol{A}\in\mathbb{R}^{n\times n}$, we use
$\left|\boldsymbol{A}\right|_{\infty}=\max_{1\leq i,j\leq n}|A_{i,j}|$. For
sequences $a_{n}$ and $b_{n}$, we use the notation $a_{n}\asymp b_{n}$ to
indicate that $a_{n}$ is bounded both above and below by $b_{n}$
asymptotically, i,e, for some constants $C_{0},C_{1}$ and for all $n\geq
n_{n_{0}}$ we have $C_{0}\leq|a_{n}/b_{n}|\leq C_{1}$. In addition, we write
$a_{n}=O_{p}(b_{n})$ if for any $\varepsilon>0$, there exists
$C_{\varepsilon}>0$ and large enough $n_{\varepsilon}$ such that
$\mathbb{P}(|a_{n}/b_{n}|>C_{\varepsilon})<\varepsilon$, for all $n\geq
n_{\varepsilon}$. we write $a_{n}=o_{p}(b_{n})$ if $a_{n}/b_{n}$ converges to
zero in probability, i.e.,
$\lim_{n\to\infty}\mathbb{P}\left(|a_{n}/b_{n}|\geq\varepsilon\right)=0,\forall\varepsilon>0$.
The notation $\stackrel{{\scriptstyle d}}{{\Rightarrow}}$ indicates
convergence in distribution.

We use the notation $\|\cdot\|_{\psi_{2}},\|\cdot\|_{\psi_{1}}$ to refer to
the sub-Gaussian and sub-exponential norms respectively. Specifically, for a
random variable $X$, we let

$\|X\|_{\psi_{1}}=\sup_{q\geq
1}\;q^{-1}(\mathbb{E}|X|^{q})^{1/q}\,,\quad\|X\|_{\psi_{2}}=\sup_{q\geq
1}\;q^{-1/2}(\mathbb{E}|X|^{q})^{1/q}\,.$

For a random vector ${\boldsymbol{x}}$, its sub-Gaussian and sub-exponential
norms are defined as

$\|{\boldsymbol{x}}\|_{\psi_{1}}=\sup_{\left\|\boldsymbol{u}\right\|_{2}\leq
1}\;\|\langle{\boldsymbol{x}},\boldsymbol{u}\rangle\|_{\psi_{2}}\,,\quad\|{\boldsymbol{x}}\|_{\psi_{1}}=\sup_{\left\|\boldsymbol{u}\right\|_{2}\leq
1}\;\|\langle{\boldsymbol{x}},\boldsymbol{u}\rangle\|_{\psi_{1}}\,.$

##  2 Methodology

We begin (in Section 2.1) with a review of a penalized maximum likelihood
procedure for the MoE model (1)–(2) such as is used in Hyun et al. (2020).
Sections 2.2 and 2.3 then introduce our proposed methodology for forming
prediction sets in this context.

###  2.1 Penalized EM-based Estimator

The expectation-maximization (EM) algorithm (Dempster et al. 1977) is a common
heuristic when faced with maximum-likelihood problems involving missing data,
especially in the form of latent variables. The algorithm operates in an
iterative fashion, alternating between the E (expectation) step and M
(maximization) step, while managing to increase the objective function.

Assume $n$ data points $\\{({\boldsymbol{x}}_{i},y_{i})\\}_{i=1}^{n}$ are
drawn independently from the MoE model (1)–(2). The EM algorithm aims at
maximizing the log-likelihood

$\displaystyle\ell({\boldsymbol{\theta}})=\frac{1}{n}\sum_{i=1}^{n}\log\left[\sum_{k=1}^{K}\pi_{k}({\boldsymbol{x}}_{i})\cdot\phi_{k}({\boldsymbol{x}}_{i},y_{i})\right],$
(4)

where
${\boldsymbol{\theta}}=\left(\\{{\boldsymbol{\beta}}_{k},{\boldsymbol{\alpha}}_{k},\sigma_{k}\\}_{k\in[K]}\right)\in\mathbb{R}^{(2p+1)K}$
represents the model parameters, $\phi_{k}({\boldsymbol{x}}_{i},y_{i})$ is
given by

$\displaystyle\phi_{k}({\boldsymbol{x}}_{i},y_{i})=\frac{1}{\sqrt{2\pi}\sigma_{k}}\exp\left(-\frac{(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k})^{2}}{2\sigma_{k}^{2}}\right),$
(5)

and $\pi_{k}({\boldsymbol{x}}_{i})$ is given in (2). The log-likelihood
$\ell({\boldsymbol{\theta}})$ is not concave even as a function of
${\boldsymbol{\beta}}_{k}$, treating ${\boldsymbol{\alpha}}_{k}$ and
$\sigma_{k}$ as fixed. This is due to the marginalizing over the latent
cluster memberships $z_{i}$. The EM algorithm instead employs a minorize-
maximize approach (Hunter & Lange, 2004) in which a minorizer to the log-
likelihood is repeatedly constructed and maximized. More specifically, given
some fixed $\widehat{{\boldsymbol{\theta}}}$ it maximizes a lower bound
function $Q({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})$ over
${\boldsymbol{\theta}}$ to make
$\ell({\boldsymbol{\theta}})-\ell(\widehat{{\boldsymbol{\theta}}})$ large. We
refer to Hastie et al. (2009) for a more detailed introduction to EM algorithm
and derivation of the function
$Q({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})$ and here only
provide the description of the EM algorithm for the MoE model.

Let $\gamma_{i,k}({\boldsymbol{\theta}})$ be the probability that $z_{i}=k$
conditioned on the observed variable $({\boldsymbol{x}}_{i},y_{i})$, i.e.,

$\gamma_{i,k}({\boldsymbol{\theta}}):=\mathbb{P}\left(z_{i}=k|{\boldsymbol{x}}_{i},y_{i}\right)=\frac{\pi_{k}({\boldsymbol{x}}_{i})\cdot\phi_{k}({\boldsymbol{x}}_{i},y_{i})}{\sum_{\ell=1}^{K}\pi_{\ell}({\boldsymbol{x}}_{i})\cdot\phi_{\ell}({\boldsymbol{x}}_{i},y_{i})}.$

The function $\gamma_{i,k}({\boldsymbol{\theta}})$ is sometimes referred to as
_responsibilities_ in the literature (see, e.g. Hyun et al. (2020)) because it
quantifies how “responsible” group $k$ is for point $i$.

Consider the function
$Q({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})$ defined as

$\displaystyle
Q({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})=-\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{K}\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})\left[\log\pi_{k}({\boldsymbol{x}}_{i})+\log\phi_{k}({\boldsymbol{x}}_{i},y_{i})\right]+\sum_{k=1}^{K}\lambda_{{\boldsymbol{\alpha}}}\|{\boldsymbol{\alpha}}_{k}\|_{1}+\sum_{k=1}^{K}\lambda_{{\boldsymbol{\beta}}}\|{\boldsymbol{\beta}}_{k}\|_{1},$
(6)

where the regularization terms
$\left\|{\boldsymbol{\beta}}_{k}\right\|_{1},\left\|{\boldsymbol{\alpha}}_{k}\right\|_{1}$
are added to enforce sparsity on the estimated parameters and allow for
applications in the high-dimensional sparse regime.

The EM algorithm iterates between estimating the conditional membership
probabilities $\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})$, and reducing
(6) by updating the parameters $\widehat{{\boldsymbol{\theta}}}$. The details
of the updates are given below:

  * •

E step: Estimating the conditional responsibility of membership based on the
latest parameter estimate.

$\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})=\frac{\frac{\hat{\pi}_{i,k}}{\widehat{\sigma}_{k}}\exp\Big{(}-\frac{(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k})^{2}}{2\widehat{\sigma}_{k}^{2}}\Big{)}}{\sum_{\ell=1}^{K}\frac{\hat{\pi}_{i,\ell}}{\widehat{\sigma}_{\ell}}\exp\Big{(}-\frac{(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{\ell})^{2}}{2\widehat{\sigma}_{\ell}^{2}}\Big{)}}.$

  * •

M step: Updating $\widehat{{\boldsymbol{\theta}}}$ to lower the objective
value in (6).

    * –

For each $k=1,\ldots,K$, update $\widehat{{\boldsymbol{\beta}}}_{k}$:

$\displaystyle\widehat{{\boldsymbol{\beta}}}_{k}=\underset{{\boldsymbol{\beta}}}{\arg\min}\sum_{i=1}^{n}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}^{2}_{k}}\frac{(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}})^{2}}{2n}+\lambda_{{\boldsymbol{\beta}}}\|{\boldsymbol{\beta}}\|_{1}.$
(7)

    * –

Update $\\{\widehat{{\boldsymbol{\alpha}}}_{k}\\}_{k=1}^{K}$:

$\displaystyle\\{\widehat{{\boldsymbol{\alpha}}}_{k}\\}=\underset{\\{{\boldsymbol{\alpha}}_{k}\\}}{\arg\min}-\frac{1}{n}\sum_{i=1}^{n}\Big{(}\sum_{k=1}^{K}\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}}){\boldsymbol{x}}_{i}^{T}{\boldsymbol{\alpha}}_{k}-\log\Big{(}\sum_{\ell=1}^{K}\exp{({\boldsymbol{x}}_{i}^{T}{\boldsymbol{\alpha}}_{\ell})}\Big{)}\Big{)}+\lambda_{{\boldsymbol{\alpha}}}\sum_{k=1}^{K}\|{\boldsymbol{\alpha}}_{k}\|_{1}.$
(8)

    * –

For each $k=1,\ldots,K$, update $\widehat{\sigma}_{k}$:

$\displaystyle\widehat{\sigma}^{2}_{k}=\frac{\sum_{i=1}^{n}\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k})^{2}}{\sum_{i=1}^{n}\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}.$
(9)

As noted in Hyun et al. (2020), this M-step update represents a decreasing in
the objective but not a minimization since
$({\boldsymbol{\beta}}_{k},\sigma_{k})$ would need to be jointly optimized.
The algorithm terminates when the improvement in (6) is below some threshold.

In practice, we use the R package flowmix (Hyun 2022) to carry out the EM
algorithm.

###  2.2 Debiased Prediction

From (1)–(2), we have that

$y_{\textrm{new}}|\boldsymbol{x}_{\textrm{new}}\sim
N(\boldsymbol{x}_{\textrm{new}}^{T}{\boldsymbol{\beta}}_{k},\sigma_{k}^{2})\text{
with probability }\pi_{k}(\boldsymbol{x}_{\textrm{new}}).$

Because of the penalization,
$\hat{\Gamma}_{k}:=\boldsymbol{x}_{\textrm{new}}^{T}\widehat{{\boldsymbol{\beta}}}_{k}$
is a biased estimate of
$\Gamma_{k}:=\boldsymbol{x}_{\textrm{new}}^{T}{\boldsymbol{\beta}}_{k}$, and
thus would give biased predictions for $y_{\textrm{new}}$ even if we knew that
$z_{\textrm{new}}=k$. We therefore propose a debaising procedure. Given an
arbitrary vector $\boldsymbol{u}_{k}$, we construct a debiased prediction

$\displaystyle\hat{\Gamma}^{d}_{k}:=\boldsymbol{x}_{\textrm{new}}^{T}\widehat{{\boldsymbol{\beta}}}_{k}+\boldsymbol{u}_{k}^{T}\frac{1}{n}\sum_{i}\frac{\partial\ell_{i}(\widehat{{\boldsymbol{\theta}}})}{\partial{\boldsymbol{\beta}}_{k}},$
(10)

where $\partial/\partial{\boldsymbol{\beta}}_{k}$ denotes the gradient with
respect to ${\boldsymbol{\beta}}_{k}$.

Our proposed choice for $\boldsymbol{u}_{k}$ depends on the estimated sample
Fisher information matrix, which is given by

$\displaystyle\tilde{\boldsymbol{I}}(\widehat{{\boldsymbol{\theta}}})=\frac{1}{n}\sum_{i=1}^{n}\nabla\ell_{i}(\widehat{{\boldsymbol{\theta}}})\nabla\ell_{i}^{T}(\widehat{{\boldsymbol{\theta}}}),$
(11)

where $\ell_{i}(\widehat{{\boldsymbol{\theta}}})$ are the summands in (4)
evaluated at
$\widehat{{\boldsymbol{\theta}}}=\left(\\{\widehat{{\boldsymbol{\beta}}}_{k},\widehat{{\boldsymbol{\alpha}}}_{k},\widehat{\sigma}_{k}\\}_{k\in[K]}\right)$,
and

$\nabla\ell_{i}=\Big{(}\Big{\\{}\frac{\partial\ell_{i}}{\partial{\boldsymbol{\beta}}_{k}}\Big{\\}}_{k=1}^{K},\Big{\\{}\frac{\partial\ell_{i}}{\partial{\boldsymbol{\alpha}}_{k}}\Big{\\}}_{k=1}^{K},\Big{\\{}\frac{\partial\ell_{i}}{\partial\sigma_{k}}\Big{\\}}_{k=1}^{K}\Big{)}^{T}\,.$

We would like a choice of $\boldsymbol{u}_{k}$ that will lead to a narrower
interval. This intuition (which is described in greater detail below) suggests
choosing $\boldsymbol{u}_{k}$ as the solution to the optimization problem

$\displaystyle\min_{\boldsymbol{u}}\
\boldsymbol{u}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}\qquad\qquad$
$\displaystyle\ \textrm{s.t.}\ \
\sup_{\omega\in\mathcal{C}}\left|\left\langle\omega,\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-\boldsymbol{x}_{\textrm{new}}\right\rangle\right|\leq\lambda_{k}\|\boldsymbol{x}_{\textrm{new}}\|_{2},\quad\left\|\boldsymbol{u}\right\|_{1}\leq
L\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2},$ (12)

where we define
$\tilde{{\boldsymbol{\Sigma}}}_{k}:=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}^{T}_{i}$,
$L$ is a sufficiently large constant, and $\lambda_{k}$ is a tuning parameter
(in Section 3, we discuss the proper rate of $\lambda_{k}$ and the choice of
$L$). In addition,
$\mathcal{C}=\\{{\boldsymbol{e}}_{1},...,{\boldsymbol{e}}_{p},\boldsymbol{x}_{\textrm{new}}/\|\boldsymbol{x}_{\textrm{new}}\|_{2}\\}$,
where ${\boldsymbol{e}}_{i}$ denotes the $i$-th standard Euclidean basis
vector. The matrix
$\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$ is the
estimated Fisher information matrix, constrained to
${\boldsymbol{\beta}}_{k}$, which is given by

$\displaystyle\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})=\frac{1}{n}\sum_{i=1}^{n}\left(\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})\frac{y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}\right)\left(\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})\frac{y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}\right)^{T}.$
(13)

The above characterization follows from the definition of
$\tilde{I}(\widehat{{\boldsymbol{\theta}}})$ given by (11) along with identity
(35), restricted to class $k$.

######  Remark.

Because $\tilde{{\boldsymbol{\Sigma}}}_{k}$ and
$\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$ depend
on both the response vector ${\bf y}$ and the covariate matrix $\mathbf{X}$,
the solution $\boldsymbol{u}_{k}$ to (12) is also dependent on them. This is
in contrast to the optimization problem proposed by Javanmard & Montanari
(2014a) for constructing the debiasing direction, as it only involved
covariates and hence conditional on the covariate matrix $\mathbf{X}$, the
debiasing direction was independent of the response vector ${\bf y}$.

To deal with the complications resulting from the dependence of
$\boldsymbol{u}_{k}$ on $(\mathbf{X},\mathbf{y})$, we do sample splitting.
Specifically, we split the data into two sets, $\mathcal{D}_{1}$ and
$\mathcal{D}_{2}$, where optimization (12) is solved on $\mathcal{D}_{1}$, and
$\widehat{{\boldsymbol{\beta}}}_{k},\hat{\Gamma}^{d}_{k}$ are calculated using
samples in $\mathcal{D}_{2}$, per (10).

The above approach to construct the direction $\boldsymbol{u}_{k}$ is
generalized from Javanmard & Montanari (2014a); Zhang & Zhang (2014); Cai &
Guo (2017) and aims to find a direction that minimizes the variance while
controlling the bias. However, the goal in Javanmard & Montanari (2014a);
Zhang & Zhang (2014); Cai & Guo (2017) is to establish inference for
coefficients, which differs from our goal of establishing inference for
prediction. Closer to our aim here, we follow the proposal of Cai et al.
(2021) to construct prediction intervals, under a linear regression model;
however, our setting of a MoE model is more complex and requires novel
methodology and analysis. In particular, our optimization problem differs from
Cai et al. (2021) in that we use a constrained Fisher information matrix,
$\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$, in the
objective of (12), and we allow the matrix $\tilde{{\boldsymbol{\Sigma}}}_{k}$
in the constraint of (12) to be different from the matrix in the objective,
while Cai et al. (2021) keeps them the same.

The first constraint in (12) can be decomposed as

$\displaystyle\left\|\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-\boldsymbol{x}_{\textrm{new}}\right\|_{\infty}$
$\displaystyle\leq\lambda_{k}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}$
(14)
$\displaystyle\left|\left\langle\boldsymbol{x}_{\textrm{new}},\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-\boldsymbol{x}_{\textrm{new}}\right\rangle\right|$
$\displaystyle\leq\lambda_{k}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}^{2}$
(15)

Similar to the general intuition behind the quadratic optimization of
Javanmard & Montanari (2014a), the objective value
$\boldsymbol{u}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}$
is related to the variance of $\hat{\Gamma}^{d}_{k}$, and the constraint (14)
relates to its bias. So the optimization is indeed aiming to minimize the
variance (and hence the length of prediction intervals which will be
constructed based on $\hat{\Gamma}^{d}_{k}$), while controlling the bias of
$\hat{\Gamma}^{d}_{k}$. That said, in the analysis we need to show that the
bias is dominated by the variance term and therefore need to establish a lower
bound on the variance. The constraint (15) is added for this step. It makes
the feasible set of the optimization problem smaller and makes it possible to
lower bound the optimal value of the objective (see Proposition A.3 for
technical arguments). This idea originates from Cai et al. (2021), which
introduced the “variance-enhancement projection direction”. While the general
intuition carries over to our current setting, characterizing the statistical
properties of $\hat{\Gamma}^{d}_{k}$ under the MoE requires a rather intricate
and technical analysis.

We denote the estimated variance for $\hat{\Gamma}^{d}_{k}$ as

$\displaystyle\widehat{V}_{k}=\frac{1}{n}\boldsymbol{u}^{T}_{k}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k},\quad
k=1,...,K,$ (16)

and the prediction variance estimate (conditional on $z_{\textrm{new}}=k$) as

$\displaystyle b^{2}_{k}:=\widehat{V}_{k}+\widehat{\sigma}^{2}_{k},\quad
k=1,...,K.$ (17)

###  2.3 Prediction Sets

Recall that our goal is to construct a $100(1-q)\%$ prediction set
$\Omega_{q}(\boldsymbol{x}_{\textrm{new}})$ satisfying (3). By the nature of
the mixture, we seek a prediction set of the form

$\Omega_{q}(\boldsymbol{x}_{\textrm{new}})=\cup_{k=1}^{K}\mathcal{L}_{k},$

where each $\mathcal{L}_{k}=[l_{k},u_{k}]$ is centered at a debiased estimator
$\hat{\Gamma}^{d}_{k}$. Regarding the length of
$\Omega_{q}(\boldsymbol{x}_{\textrm{new}})$ as our budget it is clear at an
intuitive level that we should spend more on the mixture components to which
$(\boldsymbol{x}_{\textrm{new}},y_{\textrm{new}})$ is more likely to belong,
i.e., to those groups $k$ with larger
$\pi_{k}(\boldsymbol{x}_{\textrm{new}})$. To this end, we form a probability
density function using a weighted mixture of Gaussian densities:

$\displaystyle
f(y)=\sum_{k=1}^{K}\frac{\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}})}{b_{k}}\cdot\phi\left(\frac{y-\hat{\Gamma}^{d}_{k}}{b_{k}}\right),$
(18)

where $\phi(\cdot)$ is the standard normal pdf. The particular form of this
density is justified in the proof of Theorem 3.2 given in Section B. We give a
schematic illustration of $f$ in Figure 2.

Figure 2: An illustration of the mixture density $f$, given by (18) for $K=2$
groups. The cutoff level (indicated by the red line) is the highest level such
that the region shaded in blue has area at least $1-q$.

We seek a set of intervals $[y^{-}_{1},y^{+}_{1}],...,[y^{-}_{K},y^{+}_{K}]$,
such that

$\displaystyle\sum_{k=1}^{K}\int_{y^{-}_{k}}^{y^{+}_{k}}f(y)dy=1-q,$ (19)

while minimizing $\sum_{k=1}^{K}|y^{+}_{k}-y^{-}_{k}|$. For this, we start
from a large cutoff (the horizontal line in the figure) and decrease that
until the area under $f$ and corresponding to $y$ with $f(y)$ above the cutoff
(the blue region in the figure) is $1-q$. To approximate the area we take a
discretization approach as outlined below.

Assume without loss of generality that
$\hat{\Gamma}_{1}^{d}\leq...\leq\hat{\Gamma}_{K}^{d}$. We start by considering
the interval

$\mathcal{Q}=[\hat{\Gamma}_{1}^{d}-b_{1}z_{1-q/2},\quad\hat{\Gamma}_{K}^{d}+b_{K}z_{1-q/2}],$

which we know has probability at least $1-q$. We then divide $\mathcal{Q}$
into segments of size $\delta$, denoted as
$\mathcal{Q}_{1},...,\mathcal{Q}_{\frac{|\mathcal{Q}|}{\delta}}$. The area
under the curve $f(y)$ confined to the segment $\mathcal{Q}_{i}$ is
approximately $\delta h_{i}$, where $h_{i}$ is the density $f$ evaluated at
the midpoint point of $\mathcal{Q}_{i}$. We next sort $h_{i}$’s corresponding
to each segment in decreasing order, i.e.,

$h_{(1)}\geq h_{(2)}\geq...\geq h_{(\frac{|\mathcal{Q}|}{\delta})},$

and find the smallest $N$ such that

$\displaystyle\delta\sum_{i=1}^{N}h_{(i)}\geq{1-q}.$ (20)

We return $\Omega_{q}(\boldsymbol{x}_{\textrm{new}}):=\cup_{i=1}^{N}\
\mathcal{Q}_{(i)}$ as the prediction set.

1

3

0: Confidence level $1-q$, discretization scale $\delta$, debiased estimate of
each center $\hat{\Gamma}^{d}_{k}$, prediction standard error estimate
$b_{k}$, membership probability estimate
$\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}}),k=1,...,K$.

20: Prediction set with $(1-q)$ coverage

1: Form a weighted mixture of Gaussian densities  $\displaystyle
f(y)=\sum_{k=1}^{K}\frac{\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}})}{b_{k}}\cdot\phi\left(\frac{y-\hat{\Gamma}^{d}_{k}}{b_{k}}\right).$

2: Choose a large enough interval $\mathcal{Q}$ so that the integral
$\int_{\mathcal{Q}}f(y)\geq 1-q$, i.e.
$\mathcal{Q}=[\hat{\Gamma}_{1}^{d}-b_{1}z_{1-q/2},\quad\hat{\Gamma}_{K}^{d}+b_{K}z_{1-q/2}]\,.$

3: Divide $\mathcal{Q}$ into segments of size $\delta$. Let $y_{i}$ be the
midpoint of $\mathcal{Q}_{i}$ and $h_{i}=f(y_{i})$.

4: Sort $h_{i}$’s in decreasing order,  $h_{(1)}\geq h_{(2)}\geq...\geq
h_{(\frac{|\mathcal{Q}|}{\delta})}.$

5: Find the smallest $N$ such that
$\displaystyle\delta\sum_{i=1}^{N}h_{(i)}\geq{1-q}.$

6: Return the union of the corresponding segments
$\displaystyle\Omega_{q}(\boldsymbol{x}_{\textrm{new}})=\cup_{i=1}^{N}\
\mathcal{Q}_{(i)}\,.$ (21)

Algorithm 1 Constructing prediction set
$\Omega_{q}(\boldsymbol{x}_{\textrm{new}})$ with $(1-q)$ coverage

##  3 Theoretical Guarantees

We a consider sequence of problems where the sample size $n\to\infty$ and
covariate dimension $p=p(n)\to\infty$, while the number of groups $K$ is
bounded, and we establish asymptotic validity of our prediction sets for the
MoE model (1)–(2). We first lay out several technical assumptions on the
estimation error of $\widehat{{\boldsymbol{\theta}}}$, the random covariate
vectors $\\{{\boldsymbol{x}}_{i}\\}_{i\in[n]}$, and the model parameters
${\boldsymbol{\theta}}$.

  * •

(A1) _Parameter estimation $\eta_{{\sf est}}$._ Suppose that

$\max_{k\in[K]}\left\|\widehat{{\boldsymbol{\theta}}}_{k}-{\boldsymbol{\theta}}_{k}\right\|_{1}=\max_{k\in[K]}\left(\left\|\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k}\right\|_{1}+\left\|\widehat{{\boldsymbol{\alpha}}}_{k}-{\boldsymbol{\alpha}}_{k}\right\|_{1}+|\widehat{\sigma}_{k}-\sigma_{k}|\right)=O_{p}(\eta_{{\sf
est}})\,,$

where $\eta_{{\sf est}}$ scales with $n,p$ and potentially other structure
associated with the parameters (e.g. sparsity levels). We assume that

$\displaystyle\sqrt{n}\log(np)\eta_{{\sf est}}^{2}=o(1)\,.$ (22)

  * •

(A2) _Distribution of features._ We have a positive-semidefinite matrix
${\boldsymbol{\Sigma}}\in\mathbb{R}^{p\times p}$ (or more precisely a sequence
of matrices of growing dimension) with bounded operator norm,
$\|{\boldsymbol{\Sigma}}\|_{{\rm op}}\leq C_{\Sigma}$, for a constant
$C_{\Sigma}$ as $p\to\infty$. Suppose that
${\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}$ are independent sub-
Gaussian vectors, with mean zero and sub-Gaussian norm
$\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}=O(1)$.

  * •

(A3) _Bounded noise and signal._ The noise variances $\sigma_{k}^{2}$ are
strictly positive and bounded constants for $k\in[K]$. We also assume that
$\max_{\ell,k\in[K]}\|{\boldsymbol{\beta}}_{k}-{\boldsymbol{\beta}}_{\ell}\|_{2}=O(1)$.

Condition (A1) assumes an $\ell_{1}$-consistency rate for the estimate
$\widehat{{\boldsymbol{\theta}}}$. Consistency presupposes identifiability,
which in the case of the $\\{{\boldsymbol{\alpha}}_{k}\\}$ may require
additional assumptions (since $\\{{\boldsymbol{\alpha}}_{k}+{\bf
c}:k\in[K]\\}$ corresponds to the same $\pi_{k}({\boldsymbol{x}})$ as
$\\{{\boldsymbol{\alpha}}_{k}:k\in[K]\\}$). Our theoretical results apply to
any estimator which satisfy condition (22). The proposed EM estimator in
Section 2.1 is just one specific choice. Instead of $\ell_{1}$-regularization,
one can follow other variants based on iterative truncation. For example, Wang
et al. (2014) analyzes a mixture of regression model with two groups and
proposes a truncated EM algorithm (with a gradient ascent implementation)
which achieves $\eta_{{\sf est}}=s\sqrt{\log(p)\log(n)/n}$, with $s$ the
sparsity level of model parameters. The work Zhang et al. (2020) derives a
similar $\ell_{1}$-consistency rate for the high-dimensional mixed linear
regression with two groups, for an iterative EM procedure which performs
$\ell_{1}$ regularization at each step. While the mixture of experts model is
more complicated we conjecture that a similar rate for $\eta_{{\sf est}}$
carries over to this setting. Under such conjecture, condition (22) simplifies
to

$\frac{s^{2}\log(np)\log(p)\log(n)}{\sqrt{n}}=o(1)\,.$

Condition (A2) is on the random covariate vectors ${\boldsymbol{x}}_{i}$ and
is a common assumption in high-dimensional statistical estimation; see e.g.
Bühlmann & Van De Geer (2011). Condition (A3) on the pairwise distances
$\left\|{\boldsymbol{\beta}}_{k}-{\boldsymbol{\beta}}_{\ell}\right\|_{2}$ and
noise variance $\sigma_{k}^{2}$ is to control the heterogeneity of data coming
from different groups.

Our first theorem is on asymptotic normality of the bias-corrected estimators
$\hat{\Gamma}^{d}_{k}$ defined in (10) and involves the matrix

${\boldsymbol{\Sigma}}_{k}=\mathbb{E}\left[\frac{\gamma_{1k}\left({\boldsymbol{\theta}}\right)}{\sigma_{k}^{2}}{\boldsymbol{x}}_{1}{\boldsymbol{x}}_{1}^{T}\right].$

######  Theorem 3.1.

Suppose that
$\frac{\left\|{\boldsymbol{\Sigma}}^{-1}_{k}\boldsymbol{x}_{\textrm{new}}\right\|_{1}}{\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}}=O(1)$
and $\log(p)=o(n^{1/4}/\sqrt{\log(n)})$. Let $\boldsymbol{u}_{k}$ be the
solution to the optimization problem (12) with constant
$L\geq\frac{\left\|{\boldsymbol{\Sigma}}^{-1}_{k}\boldsymbol{x}_{\textrm{new}}\right\|_{1}}{\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}}$
and $\lambda_{k}\asymp\eta_{{\sf est}}\log(np)+\sqrt{\log(p)/n}$. Under
Assumptions (A1), (A2), (A3), and Remark Remark, we have

$\displaystyle\frac{\sqrt{n}(\hat{\Gamma}_{k}^{d}-\Gamma_{k})}{\sqrt{\boldsymbol{u}^{T}_{k}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}\stackrel{{\scriptstyle
d}}{{\Rightarrow}}N(0,1),$ (23)

where $\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$,
given by (13), is the sample Fisher information matrix constrained to entries
corresponding to ${\boldsymbol{\beta}}_{k}$.

Now that we have established the asymptotic normality of our debiased
estimators, we are ready to prove that our prediction sets provide proper
asymptotic coverage.

######  Theorem 3.2.

Under the assumptions of Theorem 3.1, the prediction set
$\Omega_{q}(\boldsymbol{x}_{\textrm{new}})$ has asymptotically valid coverage.
Specifically, fix $\gamma>0$ arbitrarily small and in Algorithm 1 set the
discretization scale $\delta$ so that

$\delta\leq 11\sqrt{\gamma\min_{k\in[K]}(b_{k})/{\rm Len}(\mathcal{Q})}\,,$

with ${\rm Len}(\mathcal{Q})$ representing the length of interval
$\mathcal{Q}$. Then, for any $\boldsymbol{x}_{\textrm{new}}$ and its response
$y_{\textrm{new}}$ generated according to MoE, we have

$\lim_{n\to\infty}\mathbb{P}\left(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|\boldsymbol{x}_{\textrm{new}}\right)\geq
1-q-\gamma\,.$

Note that choosing $\gamma$ arbitrarily small we get a coverage arbitrarily
close to $1-q$. We refer to Section A for the proof of Theorems 3.1 and 3.2.

##  4 Numerical Study

In Section 4.1, we return to the low-dimensional example given in Section 1
and consider several variations to build greater understanding of the behavior
of our intervals. In Section 4.2, we assess the performance of our procedures
in a high-dimensional example.

###  4.1 A Low-Dimensional Example

Figure 1 shows a two-group example where the mean functions of the two groups
are

$\mu_{1}(t)=10(t-0.5)^{2}\qquad\text{and}\qquad\mu_{2}(t)=1,$

the error variances are $\sigma_{k}^{2}=0.15^{2}$, and the log odds of being
in the first group is given by

$\log\left[\frac{\pi_{1}(t)}{1-\pi_{1}(t)}\right]=\log(9)-2\log(9)\cdot t,$

for $n=100$ equally spaced values of $t$ ranging from 0 to 1. This means that
$\pi_{1}(t)$ decreases from 0.9 to 0.1. Writing this MoE model in the notation
of (1)–(2), we have ${\boldsymbol{x}}=(1,t,t^{2})^{T}$,
${\boldsymbol{\beta}}_{1}=(2.5,-10,10)^{T}$,
${\boldsymbol{\beta}}_{2}=(1,0,0)^{T}$, ${\boldsymbol{\alpha}}_{1}=(\log
9,-2\log 9,0)^{T}$, and ${\boldsymbol{\alpha}}_{2}=(0,0,0)^{T}$. This initial
setting is perfectly symmetric in the two groups other than the difference in
mean functions. For example, the first group at $t=0$ is as common as the
second group is at $t=1$. This symmetry is manifest in the width of the
intervals: In the middle panel, the width of the interval around the first
group at $t=0$ matches that of the second group at $t=1$. The symmetry is also
apparent in the third panel plots.

To build our intuition, we explore the effect of breaking this symmetry
between groups. For example, suppose that the error variances in the two
groups are not equal. In Figure 3, we take $\sigma_{1}=0.1$ and
$\sigma_{2}=0.2$. While the empirical coverage is maintained around the
nominal $0.95$ level, the average length of the prediction set
$\Omega_{0.05}(t)$ is no longer symmetric around $t=0.5$. This makes intuitive
sense since the class imbalance means that there is more uncertainty in
estimating ${\boldsymbol{\beta}}_{2}$ than ${\boldsymbol{\beta}}_{1}$; thus,
at $t=1$, where more of the prediction set is devoted to the second group
(since $\pi_{2}(1)>\pi_{1}(1)$), the average length of $\Omega_{0.5}(1)$ will
be larger compared to that of $\Omega_{0.5}(0)$.

0.250.500.750.000.250.500.751.00$t$Probability01230.000.250.500.751.00$t$$y$01230.000.250.500.751.00$t$$y$0.9000.9250.9500.9751.0000.000.250.500.751.00$t$Cover
Rate0.60.81.01.20.000.250.500.751.00$t$Average LengthcolourGroup 1Group
2linetypeEstimatedTrue

Figure 3: Unequal error variance case ($\sigma_{1}<\sigma_{2}$). See caption
of Figure 1 for description of panels.

Another way to break the symmetry would be by considering class imbalance. We
again assume $\sigma_{1}=\sigma_{2}=0.15$, but now suppose that overall about
60% of observations belong to the first group. In particular, we take instead
${\boldsymbol{\alpha}}_{1}=\log 10\cdot(1,-\log 4,0)^{T}$, so that
$\pi_{1}(t)$ ranges from about $0.91$ at $t=0$ to about $0.29$ at $t=1$.
Figure 4 shows the effect of this class imbalance. We see the same increasing
length as in the previous example despite the error variances being equal. In
this example, it is the class imbalance that leads to greater uncertainty in
estimating ${\boldsymbol{\beta}}_{2}$.

0.250.500.750.000.250.500.751.00$t$Probability01230.000.250.500.751.00$t$$y$01230.000.250.500.751.00$t$$y$0.9000.9250.9500.9751.0000.000.250.500.751.00$t$Cover
Rate0.60.81.01.20.000.250.500.751.00$t$Average LengthcolourGroup 1Group
2linetypeEstimatedTrue

Figure 4: Class imbalance case
($\mathbb{E}[\pi_{1}({\boldsymbol{x}})]>\mathbb{E}[\pi_{2}({\boldsymbol{x}})]$).
See caption of Figure 1 for description of panels.

###  4.2 High-Dimensional Case

We consider a high-dimensional case where $K=2$ and $p=501$ (including the
intercept). The oceanographic application in Hyun et al. (2020) has repeated
observations at multiple ${\boldsymbol{x}}_{t}$, and we mimic that setup with
$T=150$ feature vectors ${\boldsymbol{x}}_{t}$, and $n_{t}=5$ observations per
${\boldsymbol{x}}_{t}$ for a total of $n=750$ measurements. This is high-
dimensional since ${\boldsymbol{\theta}}$ has dimension $K(2p+1)=2002$. We
take $x_{t1}=1$ and $x_{tj}$ for $j>1$ to be independent standard Gaussians.
We generate $n_{t}$ responses $y_{it}$ for each ${\boldsymbol{x}}_{t}$
according to the MoE model (1)–(2) with $\sigma_{1}=\sigma_{2}=1$ and

$\displaystyle{\boldsymbol{\beta}}_{1}$ $\displaystyle=$ $\displaystyle($
$\displaystyle-2$ , $\displaystyle 4$ , $\displaystyle-2$ , $\displaystyle-4$
, $\displaystyle 6$ , $\displaystyle 2$ , $\displaystyle 0$ , $\displaystyle
0$ , $\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle 0$ ,
$\displaystyle{\bf 0}_{490}^{T}$ $\displaystyle)^{T},$
$\displaystyle{\boldsymbol{\beta}}_{2}$ $\displaystyle=$ $\displaystyle($
$\displaystyle 2$ , $\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle 0$
, $\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle 4$ ,
$\displaystyle-2$ , $\displaystyle-4$ , $\displaystyle 6$ , $\displaystyle 2$
, $\displaystyle{\bf 0}_{490}^{T}$ $\displaystyle)^{T},$

and

$\displaystyle{\boldsymbol{\alpha}}_{1}$ $\displaystyle=$ $\displaystyle($
$\displaystyle{\bf 0}_{491}$ , $\displaystyle 0$ , $\displaystyle 0$ ,
$\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle
0.7$ , $\displaystyle 0.7$ , $\displaystyle 0.7$ , $\displaystyle 0.7$ ,
$\displaystyle 0.7$ $\displaystyle)^{T},$
$\displaystyle{\boldsymbol{\alpha}}_{2}$ $\displaystyle=$ $\displaystyle($
$\displaystyle{\bf 0}_{491}$ , $\displaystyle-0.7$ , $\displaystyle-0.7$ ,
$\displaystyle-0.7$ , $\displaystyle-0.7$ , $\displaystyle-0.7$ ,
$\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle 0$ , $\displaystyle 0$
, $\displaystyle 0$ $\displaystyle)^{T}.$

We estimate the model using the penalized EM algorithm described in Section
2.1 as implemented in the flowmix R package (Hyun, 2022). We perform five-fold
cross validation to choose the parameters
$(\lambda_{{\boldsymbol{\alpha}}},\lambda_{{\boldsymbol{\beta}}})$ from a
$10\times 10$ logarithmically-spaced grid. While in Section 2.2 we remark that
sample splitting avoids the complications resulting from
$\boldsymbol{u}_{k}$’s dependence on $({\bf X},{\bf y})$, empirically we find
that coverage is attained even if we ignore this dependence. Therefore, in
this and all numerical results we do not use sample splitting for the
debiasing step.

To evaluate our method, we generate $100$ independent
$\boldsymbol{x}_{\textrm{new}}\in\mathbb{R}^{p}$. For each
$\boldsymbol{x}_{\textrm{new}}$, we compute
$\Omega_{0.05}(\boldsymbol{x}_{\textrm{new}})$ and record its length. We then
generate $100$ independent $y_{\text{new},i}$’s for each
$\boldsymbol{x}_{\textrm{new}}$ and record the proportion of
$y_{\text{new},i}$’s falling into the prediction set for that
$\boldsymbol{x}_{\textrm{new}}$.

We repeat the above procedure $500$ times, keeping the $100$
$\boldsymbol{x}_{\textrm{new}}$’s the same. For each
$\boldsymbol{x}_{\textrm{new}}$, we compute

  1. 1.

the average length of prediction sets (across the 500 runs), and

  2. 2.

the coverage probability (proportion across the 500 runs and 100
$y_{\text{new},i}$).

0.9000.9250.9500.9751.0000.000.250.500.751.00$\pi_{1}$Coverage Rate

56789100.000.250.500.751.00$\pi_{1}$Average Length

56789100102030DistanceAverage Length

Figure 5: High-dimensional simulation. Every point corresponds to a different
$\boldsymbol{x}_{\textrm{new}}\in\mathbb{R}^{p}$. The empirical coverage and
average length of the prediction set
$\Omega_{0.05}(\boldsymbol{x}_{\textrm{new}})$ is computed over 500 training
sets and 100 $y_{\text{new}}$ values for each $\boldsymbol{x}_{\textrm{new}}$.

In the left and middle panels of Figure 5, we plot these quantities as a
function of $\pi_{1}(\boldsymbol{x}_{\textrm{new}})$, the true probability of
being drawn from cluster 1 for each $\boldsymbol{x}_{\textrm{new}}$. As
desired, the coverage rate of our prediction sets meet the $95\%$ nominal
level, regardless of $\pi_{1}(\boldsymbol{x}_{\textrm{new}})$. We observe that
the prediction sets tend to be twice as long when
$\pi_{1}(\boldsymbol{x}_{\textrm{new}})\approx 0.5$ compared to at the
extremes. This is likely because at an extreme the less common group’s
interval can be very narrow without hurting coverage whereas when the two
groups are balanced, both intervals are needed.

In the right panel of Figure 5, we plot the average length of prediction
intervals against the distances between the two group means, i.e.,

$\left|\boldsymbol{x}_{\textrm{new}}^{T}({\boldsymbol{\beta}}_{1}-{\boldsymbol{\beta}}_{2})\right|.$

Recall that in the low-dimensional examples (Figures 1, 3, and 4) we observed
a marked decrease in the prediction set length when the means crossed each
other. A similar phenomenon is apparent here. There is a linear increasing
trend when
$\left|\boldsymbol{x}_{\textrm{new}}^{T}({\boldsymbol{\beta}}_{1}-{\boldsymbol{\beta}}_{2})\right|$
goes from 0 to 10. As suggested in the middle panel, the individual intervals
have average length around 5. Thus, this linear increase represents the two
overlapping intervals gradually being pulled apart. At a distance of 10, they
no longer overlap, which explains the leveling of this trend. The variability
in length seen for distances greater than 10 can be explained, for example, by
differing values of $\pi_{1}(\boldsymbol{x}_{\textrm{new}})$.

##  5 Superconductivity Data Application

We apply our method to the superconductivity data provided in Hamidieh (2018).
This dataset contains the critical temperature (in Kelvin) and a set of $p=81$
attributes for about 21,000 materials. The attributes used as predictors are
elemental property statistics and electronic structures of attributes. We
center and scale each predictor column and we take the response to be
$\log(1+\text{temperature})$. The log transform makes the data less skewed
right and adding 1 Kelvin to each temperature can be thought of as replacing
the log of extremely low temperatures (some are less than $1mK$) with 0.

We randomly split the observations into a training set of $n=200$ (used for
estimating model parameters, cross validation of
$\lambda_{{\boldsymbol{\alpha}}}$ and $\lambda_{{\boldsymbol{\beta}}})$, and
forming the prediction sets), a validation set of size 1000 (used to choose
$K$), and a test set of about 20,000 observation (for measuring the coverage
of our prediction sets). Table 1 shows the mean squared prediction error,
computed on the validation set, for $K$ ranging from 1 to 5. To make
predictions at a given ${\boldsymbol{x}}_{\text{validation}}$, we use

${\boldsymbol{x}}_{\text{validation}}^{T}\widehat{{\boldsymbol{\beta}}}_{\hat{k}({\boldsymbol{x}}_{\text{validation}})}\qquad\text{where}\qquad\hat{k}({\boldsymbol{x}}_{\text{validation}})=\arg\max_{k}\hat{\pi}_{k}({\boldsymbol{x}}_{\text{validation}})$

is the class with highest estimated probability. The prediction errors on the
validation set suggest that $K=2$ may be a suitable choice.

K |  Prediction error
---|---
1 |  8.953
2 |  0.867
3 |  0.978
4 |  1.318
5 |  0.869
Table 1: Prediction errors computed on a validation set for fitted models
where $K=1,2,3,4,5$.

For each observation $(\boldsymbol{x}_{\textrm{new}},y_{\textrm{new}})$ in the
test set, we form $\Omega_{0.05}(\boldsymbol{x}_{\textrm{new}})$ and note
whether $y_{\textrm{new}}\in\Omega_{0.05}(\boldsymbol{x}_{\textrm{new}})$.
Figure 6 displays the prediction sets for a random subset of 100 of the 20,000
intervals formed on the test set. We see that
$\Omega_{0.05}(\boldsymbol{x}_{\textrm{new}})$ is often a single interval,
although it also occasionally the union of two intervals. The overall coverage
on the test set is $97.1\%$. The average length of
$\Omega_{0.05}(\boldsymbol{x}_{\textrm{new}})$ (after being transformed back
from log-values) is around 42 Kelvin.

050255075100indexlog(1+temperature)

Figure 6: Plot of prediction sets for $100$ randomly sampled data points in
the test set. The black points are the $\log(1+\text{temperature})$ and the
bars correspond to the prediction sets for each observation. We use red when
there is only one interval and cyan when there are two intervals.

The $97.1\%$ coverage is averaged over all $\boldsymbol{x}_{\textrm{new}}$ and
yet our prediction sets are designed for conditional coverage in the sense of
(3). This stronger form of coverage implies that we can get coverage on
subsets of observations defined by $\boldsymbol{x}_{\textrm{new}}$, i.e.

$\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|\boldsymbol{x}_{\textrm{new}}\in\mathcal{S})\geq
1-q.$

The predictor that is most correlated with critical temperature is the
weighted standard deviation of thermal conductivity. We divide the range of
this variable into 5 equally-spaced sub-intervals, and divide test data points
into 5 subgroups accordingly. Table 2 confirms that our prediction sets meet
the nominal level within each subgroup.

Subgroup | Number of data | Coverage Rate
---|---|---
1 | 3098 | 95.9%
2 | 4039 | 97.2%
3 | 3960 | 97.1%
4 | 3999 | 95.6%
5 | 4000 | 94.8%
Table 2: Coverage rates of $95\%$ prediction sets conditional on subgroups
defined by the predictor “weighted standard deviation of thermal
conductivity.”

##  6 Conclusion

We have shown how to construct prediction sets for the high-dimensional
mixture of experts model. Mixture models are important for capturing the
heterogeneity that is present in many real-world situations. While in small
data samples it was common to dismiss deviations from the norm as outliers, in
large data sets it becomes possible to use models that can identify and model
these subgroups. While mixture of regression models allow for such
heterogeneity-aware predictive modeling, they assume that the relative sizes
of the subgroups are fixed. Importantly, mixture of experts models remove this
assumption and allow the prevalence of different subgroups to depend on the
features. This generalization is essential in many situations from ecology,
where the relative proportions of different subpopulations depends on
environmental covariates (Hyun et al., 2020), to politics, where the political
composition depends on demographic and geographic variables.

Our focus on conditional coverage can be crucial in certain applications. For
example, Romano, Barber, Sabatti & Candès (2019) emphasizes the importance of
ensuring that all subpopulations enjoy the same coverage guarantees and cast
this as a fairness issue when the subpopulations are defined based on a
protected attribute.

##  Appendix A Proof of Main Theorems

###  A.1 Proof of Theorem 3.1

Before we prove the asymptotic normality, we start by presenting a proposition
establishing that the optimization problem (12) is feasible. We postpone its
proof to Section C.1.

######  Proposition A.1.

(Feasibility) Under the assumptions of Theorem 3.1, there exists
$\boldsymbol{u}$ such that

$\sup_{\boldsymbol{\omega}\in\mathcal{C}}\left|\left\langle\boldsymbol{\omega},\tilde{\Sigma}_{k}\boldsymbol{u}-\boldsymbol{x}_{\textrm{new}}\right\rangle\right|\leq\lambda_{k}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\text{
and }\left\|\boldsymbol{u}\right\|_{1}\leq
L\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}$

is satisfied for $\lambda_{k}\asymp\eta_{{\sf est}}\log(np)+\sqrt{\log(p)/n}$,
$L\geq\frac{\left\|{\boldsymbol{\Sigma}}^{-1}_{k}\boldsymbol{x}_{\textrm{new}}\right\|_{1}}{\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}}$,
and
$\mathcal{C}=\left\\{\boldsymbol{e}_{1},\dots,\boldsymbol{e}_{p},\boldsymbol{x}_{\textrm{new}}/\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\right\\}$.

We next decompose the error of the bias-corrected estimator as follows:

$\displaystyle\hat{\Gamma}^{d}_{k}-\Gamma_{k}$
$\displaystyle=\boldsymbol{x}_{\textrm{new}}^{T}\widehat{{\boldsymbol{\beta}}}_{k}+\boldsymbol{u}_{k}^{T}\frac{1}{n}\sum_{i}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}){\boldsymbol{x}}_{i}-\boldsymbol{x}_{\textrm{new}}^{T}{\boldsymbol{\beta}}_{k}$
$\displaystyle=\left(\boldsymbol{x}_{\textrm{new}}-\frac{1}{n}\sum_{i}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}^{T}_{i}\boldsymbol{u}_{k}\right)^{T}(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k})+\boldsymbol{u}_{k}^{T}\frac{1}{n}\sum_{i}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}$
$\displaystyle=\left(\boldsymbol{x}_{\textrm{new}}-\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\right)^{T}(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k})+\boldsymbol{u}_{k}^{T}\frac{1}{n}\sum_{i}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}.$

Recall the estimated Fisher information matrix constrained to
${\boldsymbol{\beta}}_{k}$ is given by

$\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}^{2}_{k}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$

and define

$\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})=\mathbb{E}\left[\frac{\gamma^{2}_{1k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\frac{\left(y_{1}-{\boldsymbol{x}}_{1}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}{\sigma_{k}^{2}}{\boldsymbol{x}}_{1}{\boldsymbol{x}}_{1}^{T}\right].$

We then have

$\displaystyle\frac{\sqrt{n}(\hat{\Gamma}_{k}^{d}-\Gamma_{k})}{\sqrt{\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}$
$\displaystyle=\frac{\sqrt{n}\left(\boldsymbol{x}_{\textrm{new}}-\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\right)^{T}(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k})}{\sqrt{\boldsymbol{u}^{T}_{k}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}+\frac{\boldsymbol{u}_{k}^{T}\frac{1}{\sqrt{n}}\sum_{i}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}}{\sqrt{\boldsymbol{u}^{T}_{k}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}$
$\displaystyle=\underbrace{\frac{\sqrt{n}\left(\boldsymbol{x}_{\textrm{new}}-\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\right)^{T}(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k})}{\sqrt{\boldsymbol{u}^{T}_{k}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}}_{(\text{term
I})}+\underbrace{\frac{\boldsymbol{u}_{k}^{T}\frac{1}{\sqrt{n}}\sum_{i=1}^{n}\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma^{2}_{k}}(y_{i}-{\boldsymbol{x}}^{T}_{i}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}+(\text{term
II})}{\sqrt{\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}}_{(\text{term
III})},$ (24)

where

$\text{term
II}=\boldsymbol{u}_{k}^{T}\frac{1}{n}\sum_{i}\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}-\boldsymbol{u}_{k}^{T}\frac{1}{n}\sum_{i}\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\,.$

We will proceed by stating three propositions which will be used to control
terms I,II, and III. We defer the proof of these propositions to Section C.

The first proposition allows us to control the numerator of term I.

######  Proposition A.2.

Let $\boldsymbol{u}_{k}$ be the solution to optimization (12) with
$\lambda_{k}\asymp\eta_{{\sf est}}\log(np)+\sqrt{\log(p)/n}$ and
$L\geq\left\|{\boldsymbol{\Sigma}}^{-1}_{k}\boldsymbol{x}_{\textrm{new}}\right\|_{1}$.
Under Assumption (A1) we have

$\displaystyle\left|\left(\boldsymbol{x}_{\textrm{new}}-\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\right)^{T}(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k})\right|=O_{p}\left(\left(\eta_{{\sf
est}}\log(np)+\sqrt{\log(p)/n}\right)\eta_{{\sf
est}}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\right).$ (25)

The second proposition lower bounds the denominator of term I.

######  Proposition A.3.

Under the assumptions of Theorem 3.1, there exist constants $c,C>0$, such that

$\displaystyle\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}\geq
C\|\boldsymbol{x}_{\textrm{new}}\|^{2}_{2}\,.$ (26)

Combining Propositions A.2 and A.3 we obtain

$\text{term I}=O_{p}\left(\left(\eta_{{\sf
est}}\sqrt{n}\log(np)+\sqrt{\log(p)}\right)\eta_{{\sf est}}\right)=o_{p}(1),$

by Condition (A1), equation (22).

The next proposition controls term II.

######  Proposition A.4.

Under the assumptions of Theorem 3.1,

$\left\|\frac{1}{n}\sum_{i=1}^{n}\left(\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right)(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\right\|_{\infty}=o_{p}(1).$

Therefore,

$\displaystyle\text{term
II}\leq\left\|\frac{1}{n}\sum_{i=1}^{n}\left(\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right)(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\right\|_{\infty}\|\boldsymbol{u}\|_{1}=o_{p}(\|\boldsymbol{x}_{\textrm{new}}\|_{2})\,.$

The next proposition controls the difference of the sample Fisher information
at the estimated parameter
$\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$ and the
true Fisher information $\boldsymbol{I}^{\beta}_{k}({\boldsymbol{\theta}})$.

######  Proposition A.5.

Under the assumptions of Theorem 3.1, and for $\boldsymbol{u}_{k}$ the
solution of optimization (12), we have

$\displaystyle\left|\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}-\boldsymbol{u}_{k}^{T}\boldsymbol{I}^{\beta}_{k}({\boldsymbol{\theta}})\boldsymbol{u}_{k}\right|=o_{p}(\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2}).$
(27)

We next note that

$\displaystyle\nabla_{{\boldsymbol{\beta}}_{k}}\ell({\boldsymbol{\theta}})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma^{2}_{k}}(y_{i}-{\boldsymbol{x}}^{T}_{i}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\,,$

and by asymptotic normality of the score functions, see e.g. Van der Vaart
(2000), we have

$\frac{1}{\sqrt{n}}\sum_{i=1}^{n}\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma^{2}_{k}}(y_{i}-{\boldsymbol{x}}^{T}_{i}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\stackrel{{\scriptstyle
d}}{{\Rightarrow}}N(0,\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})).$

By sample splitting, $\boldsymbol{u}_{k}$ is independent of
$\\{({\boldsymbol{x}}_{i},y_{i})\\}$ and therefore is also independent of
$\nabla_{{\boldsymbol{\beta}}_{k}}\ell({\boldsymbol{\theta}})$. By invoking
Proposition A.5, and as an application of Slutsky’s Theorem, this implies that

$\displaystyle\frac{\boldsymbol{u}_{k}^{T}\frac{1}{\sqrt{n}}\sum_{i=1}^{n}\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma^{2}_{k}}(y_{i}-{\boldsymbol{x}}^{T}_{i}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}}{\sqrt{\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}}}\stackrel{{\scriptstyle
d}}{{\Rightarrow}}N(0,1)\,.$ (28)

Using the above distributional characterization, together with the bounds on
terms I, II, III in decomposition (24), we obtain the desired result.

##  Appendix B Proof of Theorem 3.2

Let $z_{\textrm{new}}$ denote the class of $\boldsymbol{x}_{\textrm{new}}$.
Conditioning on $z_{\textrm{new}}$ we have

$\displaystyle\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|\boldsymbol{x}_{\textrm{new}})$
$\displaystyle=\sum_{k=1}^{K}\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|z_{\textrm{new}}=k,\boldsymbol{x}_{\textrm{new}})\mathbb{P}(z_{\textrm{new}}=k|\boldsymbol{x}_{\textrm{new}})$
$\displaystyle=\sum_{k=1}^{K}\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|z_{\textrm{new}}=k,\boldsymbol{x}_{\textrm{new}})\pi_{k}(\boldsymbol{x}_{\textrm{new}})$
$\displaystyle=\sum_{k=1}^{K}\mathbb{P}(\Gamma_{k}+\varepsilon\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}})|\boldsymbol{x}_{\textrm{new}})\pi_{k}(\boldsymbol{x}_{\textrm{new}})\,,$
(29)

where we recall that
$\Gamma_{k}=\boldsymbol{x}_{\textrm{new}}^{T}{\boldsymbol{\beta}}_{k}$. Also
$\varepsilon\sim N(0,\sigma_{k}^{2})$ is independent of
$\boldsymbol{x}_{\textrm{new}}$.

As shown in the proof of Theorem 3.1 we have

$\frac{\Gamma_{k}-\hat{\Gamma}^{d}_{k}}{\widehat{V}_{k}^{1/2}}=W+\Delta\,,$

with $W\sim N(0,1)$ and $|\Delta|=o_{p}(1)$. Therefore,

$\displaystyle\frac{\Gamma_{k}+\varepsilon-\hat{\Gamma}^{d}_{k}}{b_{k}}$
$\displaystyle=\frac{\widehat{V}_{k}^{1/2}}{b_{k}}W+\frac{\sigma_{k}}{b_{k}}\varepsilon+\frac{\widehat{V}_{k}^{1/2}}{b_{k}}\Delta$
$\displaystyle=c_{k}W^{\prime}+\Delta^{\prime}\,,$

with $W^{\prime},\varepsilon\sim N(0,1)$ and
$|\Delta^{\prime}|<|\Delta|=o_{p}(1)$ and
$c_{k}^{2}:=\frac{\widehat{V}_{k}+\sigma_{k}^{2}}{b_{k}^{2}}$. Here we used
the fact that $W,\varepsilon$ are independent normal random variables.

Our procedure for constructing a prediction set returns a union of intervals:
$\Omega_{q}(\boldsymbol{x}_{\textrm{new}})=\cup_{k=1}^{K}\mathcal{L}_{k}(q)$
with $\mathcal{L}_{k}(q)=[y_{k}^{-},y_{k}^{+}]$. Since
$\mathcal{L}_{k}(q)\subseteq\Omega_{q}(\boldsymbol{x}_{\textrm{new}})$ we have

$\displaystyle\mathbb{P}(\Gamma_{k}+\varepsilon\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}}))$
$\displaystyle\geq\mathbb{P}(\Gamma_{k}+\varepsilon\in\mathcal{L}_{k}(q))$
$\displaystyle=\mathbb{P}\left(\frac{\Gamma_{k}+\varepsilon-\hat{\Gamma}^{d}_{k}}{b_{k}}\in\Big{[}\frac{y_{k}^{-}-\hat{\Gamma}^{d}_{k}}{b_{k}},\frac{y_{k}^{+}-\hat{\Gamma}^{d}_{k}}{b_{k}}\Big{]}\right)$
$\displaystyle=\mathbb{P}\left(c_{k}W^{\prime}+\Delta^{\prime}\in[l_{k},u_{k}]\right)\,,$
(30)

where we use the shorthands:

$l_{k}:=\frac{y_{k}^{-}-\hat{\Gamma}^{d}_{k}}{b_{k}},\quad
u_{k}:=\frac{y_{k}^{+}-\hat{\Gamma}^{d}_{k}}{b_{k}}\,.$

Fix $\gamma>0$ arbitrarily small. We write

$\displaystyle\mathbb{P}\left(c_{k}W^{\prime}+\Delta^{\prime}\in[l_{k},u_{k}]\right)$
$\displaystyle\geq\mathbb{P}\left(c_{k}W^{\prime}\in[l_{k}+\gamma,u_{k}-\gamma]\right)-\mathbb{P}(|\Delta^{\prime}|\geq\gamma)$
$\displaystyle\geq\mathbb{P}\left(W^{\prime}\in[(1+\gamma)(l_{k}+\gamma),(1-\gamma)(u_{k}-\gamma)]\right)-\mathbb{P}(|\Delta^{\prime}|\geq\gamma)-\mathbb{P}\Big{(}\Big{|}\frac{1}{c_{k}}-1\Big{|}\geq\gamma\Big{)}\,.$

By taking the limit $n\to\infty$ and using the fact that
$|\Delta^{\prime}|=o_{p}(1)$ and $|\widehat{\sigma}_{k}-\sigma_{k}|=o_{p}(1)$
(Condition (A1)), we get

$\displaystyle\lim_{n\to\infty}\mathbb{P}\left(c_{k}W^{\prime}+\Delta^{\prime}\in[l_{k},u_{k}]\right)$
$\displaystyle\geq\mathbb{P}\left(W^{\prime}\in[(1+\gamma)(l_{k}+\gamma),(1-\gamma)(u_{k}-\gamma)]\right)$
$\displaystyle=\Phi((1-\gamma)(u_{k}-\gamma))-\Phi((1+\gamma)(l_{k}+\gamma))\,.$

Since $\gamma>0$ was arbitrarily small and the left-hand side does not depend
on $\gamma$, by taking $\gamma\to 0$, we arrive at

$\displaystyle\lim_{n\to\infty}\mathbb{P}\left(c_{k}W^{\prime}+\Delta^{\prime}\in[l_{k},u_{k}]\right)\geq\Phi(u_{k})-\Phi(l_{k})\,.$
(31)

Using equations (30), (31) in (29) we obtain

$\displaystyle\lim_{n\to\infty}\mathbb{P}(y_{\textrm{new}}\in\Omega_{q}(\boldsymbol{x}_{\textrm{new}}))$
$\displaystyle\geq\sum_{k=1}^{K}(\Phi(u_{k})-\Phi(l_{k}))\pi_{k}(\boldsymbol{x}_{\textrm{new}})$
$\displaystyle=\sum_{k=1}^{K}\pi_{k}(\boldsymbol{x}_{\textrm{new}})\int_{l_{k}}^{u_{k}}\phi(t){\rm
d}t$
$\displaystyle=\sum_{k=1}^{K}\frac{\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}})}{b_{k}}\int_{y_{k}^{-}}^{y_{k}^{+}}\phi\left(\frac{y-\hat{\Gamma}_{k}^{d}}{b_{k}}\right){\rm
d}y$ $\displaystyle=\sum_{k=1}^{K}\int_{y_{k}^{-}}^{y_{k}^{+}}f(y){\rm d}y$
$\displaystyle=\delta\sum_{i=1}^{N}h_{(i)}-{\rm Err}$
$\displaystyle\geq(1-q)-{\rm Err}\,,$

where Err is the approximation error for replacing the integral with the
Riemann sum. To bound Err we need to upper bound the second derivative of
$f(y)$. Define the function $g_{a}(z)=\frac{1}{a}\phi(\frac{z}{a})$. We have

$|g^{\prime\prime}(z)|=\frac{1}{a^{3}}\Big{|}\frac{z^{2}}{a^{2}}-1\Big{|}\phi\Big{(}\frac{z}{a}\Big{)}\leq\frac{0.18}{a^{3}}\,.$

Therefore,

$\displaystyle|f^{\prime\prime}(y)|=\Big{|}\sum_{k=1}^{K}\frac{\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}})}{b_{k}}\phi\left(\frac{y-\hat{\Gamma}_{k}^{d}}{b_{k}}\right)\Big{|}\leq\sum_{k=1}^{K}\frac{0.18\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}})}{b_{k}}\leq\frac{0.18}{\min_{k}(b_{k})}\,,$

where we used the observation that
$\sum_{k=1}^{K}\hat{\pi}_{k}(\boldsymbol{x}_{\textrm{new}})=1$. Therefore, the
approximation error over an interval of size $N\delta$ is bounded as

$\displaystyle{\rm
Err}\leq\frac{0.18}{\min_{k}(b_{k})}\frac{(N\delta)^{3}}{24N^{2}}=0.0075\frac{N\delta^{3}}{\min_{k}(b_{k})}\,.$
(32)

We also note that the returned prediction set is a subset of the initial
interval $\mathcal{Q}$ of length ${\rm Len}(\mathcal{Q})$ which implies that
$N\delta\leq{\rm Len}(\mathcal{Q})$. Hence,

${\rm Err}\leq 0.0075{\rm
Len}(\mathcal{Q})\frac{\delta^{2}}{\min_{k}(b_{k})}\leq\gamma\,,$

by our choice of $\delta\leq 11\sqrt{\gamma\min_{k}(b_{k})/{\rm
Len}(\mathcal{Q})}$.

##  Appendix C Proof of Propositions

###  C.1 Proof of Proposition A.1

We prove the feasibility of the optimization problem by showing there exists
$\boldsymbol{u}$ such that

$\sup_{{\boldsymbol{w}}\in\mathcal{C}}\left|\left\langle{\boldsymbol{w}},\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-\boldsymbol{x}_{\textrm{new}}\right\rangle\right|\leq\lambda_{k}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\text{
and }\left\|\boldsymbol{u}\right\|_{1}\leq
L\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}$

is satisfied for $\lambda_{k}\asymp\eta_{{\sf est}}\log(np)+\sqrt{\log(p)/n}$
and
$\mathcal{C}=\left\\{{\boldsymbol{e}}_{1},\dots,{\boldsymbol{e}}_{p},\boldsymbol{x}_{\textrm{new}}/\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\right\\}$.

Recall the following quantities:

$\displaystyle\tilde{{\boldsymbol{\Sigma}}}_{k}$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$
$\displaystyle{\boldsymbol{\Sigma}}_{k}^{*}$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{ik}\left({\boldsymbol{\theta}}\right)}{\sigma_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$
$\displaystyle{\boldsymbol{\Sigma}}_{k}$
$\displaystyle=\mathbb{E}\left[\frac{\gamma_{1k}\left({\boldsymbol{\theta}}\right)}{\sigma_{k}^{2}}{\boldsymbol{x}}_{1}{\boldsymbol{x}}_{1}^{T}\right].$

Take
$\boldsymbol{u}={\boldsymbol{\Sigma}}^{-1}_{k}\boldsymbol{x}_{\textrm{new}}$.
By the assumption on $L$ stated in the statement of the proposition, we have
$\left\|\boldsymbol{u}\right\|_{1}\leq
L\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}$. To show that the other
constraint is satisfied we leverage Lemma D.6 for the set $\mathcal{C}$ which
implies that

$\displaystyle\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|(\tilde{{\boldsymbol{\Sigma}}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{\xi}\right\|_{\infty}=O_{p}\left(\eta_{{\sf
est}}\log(np)\right)\,,$
$\displaystyle\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|({\boldsymbol{\Sigma}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{\xi}\right\|_{\infty}=O_{p}\left(\sqrt{\frac{\log
p}{n}}\right).$

Combining the above two inequalities we get

$\displaystyle\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|(\tilde{{\boldsymbol{\Sigma}}}_{k}-{\boldsymbol{\Sigma}}_{k})\boldsymbol{\xi}\right\|_{\infty}=O_{p}\left(\eta_{{\sf
est}}\log(np)+\sqrt{\log(p)/n}\right)\,.$ (33)

For any ${\boldsymbol{w}}\in\mathcal{C}$,

$\displaystyle\langle{\boldsymbol{w}},\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-\boldsymbol{x}_{\textrm{new}}\rangle$
$\displaystyle=\langle{\boldsymbol{w}},\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-{\boldsymbol{\Sigma}}_{k}\boldsymbol{u}\rangle$
$\displaystyle\leq\left\|(\tilde{{\boldsymbol{\Sigma}}}_{k}-{\boldsymbol{\Sigma}}_{k}){\boldsymbol{w}}\right\|_{\infty}\left\|\boldsymbol{u}\right\|_{1}$
$\displaystyle\leq O_{p}\left(\eta_{{\sf
est}}\log(np)+\sqrt{\log(p)/n}\right)L\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}$
$\displaystyle\leq\lambda_{k}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\,,$

where in the second inequality we used (33). This completes the proof of the
feasibility claim.

###  C.2 Proof of Proposition A.2

As proved in Proposition A.1, for the choice of $\lambda_{k}\asymp\eta_{{\sf
est}}\log(np)+\sqrt{\log(p)/n}$ and
$L\geq\frac{\left\|{\boldsymbol{\Sigma}}^{-1}_{k}\boldsymbol{x}_{\textrm{new}}\right\|_{1}}{\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}}$,
optimization problem (12) is feasible. The claim follows readily from the
constraints of this optimization problem. Specifically,

$\displaystyle\left|\left(\boldsymbol{x}_{\textrm{new}}-\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\right)^{T}(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k})\right|$
$\displaystyle\leq\left\|\boldsymbol{x}_{\textrm{new}}-\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\right\|_{\infty}\left\|\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k}\right\|_{1}$
$\displaystyle\leq\lambda_{k}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\left\|\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k}\right\|_{1}$
$\displaystyle\leq O_{p}\left(\Big{(}\eta_{{\sf
est}}\log(np)+\sqrt{\log(p)/n}\Big{)}\eta_{{\sf
est}}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}\right).$

where the first inequality follows from Hölder’s inequality (duality of
$\ell_{1}-\ell_{\infty}$ norms).

###  C.3 Proof of Proposition A.3

We will use the proof strategy of Javanmard & Montanari (2014a, Lemma 3.1),
which was also used in Cai et al. (2021, Lemma 1) and modified to account for
the additional constraint (15). However, before doing that we need to deal
with the challenge that in optimization (12) the objective function is based
on $\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$,
while the constraints are in terms of $\tilde{{\boldsymbol{\Sigma}}}_{k}$. We
first relate
$\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}$
to
$\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}$.

We denote the first term in (6) by
$\tilde{Q}({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})$, i.e.,

$\displaystyle\tilde{Q}({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})=-\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{K}\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})\left[\log\pi_{k}({\boldsymbol{x}}_{i})+\log\phi_{k}({\boldsymbol{x}}_{i},y_{i})\right]\,.$
(34)

We restate a lemma from Wang et al. (2014) that allows us to connect the
derivatives of the function
$\tilde{Q}({\boldsymbol{\theta}}|\widehat{{\boldsymbol{\theta}}})$ with the
derivatives of the log-likelihood $\ell({\boldsymbol{\theta}})$.

######  Lemma C.1.

(Wang et al., 2014, Lemma 2.1) For the true parameter ${\boldsymbol{\theta}}$
and any $\tilde{{\boldsymbol{\theta}}}$, it holds that

$\displaystyle\nabla_{1}\tilde{Q}(\tilde{{\boldsymbol{\theta}}}|\tilde{{\boldsymbol{\theta}}})$
$\displaystyle=\nabla\ell(\tilde{{\boldsymbol{\theta}}})\,,$ (35)
$\displaystyle\mathbb{E}[\nabla^{2}_{1,1}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}})+\nabla^{2}_{1,2}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}})]$
$\displaystyle=-{\bf I}({\boldsymbol{\theta}})\,,$ (36)

where ${\bf
I}({\boldsymbol{\theta}})=-\mathbb{E}[\nabla^{2}\ell({\boldsymbol{\theta}})]$
is the Fisher information matrix and $\nabla_{1}$ and $\nabla_{2}$ denote
differentiation with respect to ${\boldsymbol{\beta}}_{k}$ in the first and
second arguments of $\tilde{Q}(\cdot|\cdot)$, respectively.

Following the same argument as in the proof of (Zhang et al., 2020, Lemma 3),
we have

$\displaystyle-\nabla_{1,2}^{2}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}})=-\frac{\partial^{2}}{\partial{\boldsymbol{\beta}}_{k}\partial{\boldsymbol{\beta}}_{k}^{\prime}}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}}^{\prime})|_{{\boldsymbol{\theta}}^{\prime}={\boldsymbol{\theta}}}\succeq
0\,.$ (37)

Combining (36) and (37) we get

$\displaystyle\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})$
$\displaystyle=-\mathbb{E}[\nabla^{2}_{1,1}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}})+\nabla^{2}_{1,2}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}})]$
$\displaystyle\succeq-\mathbb{E}[\nabla^{2}_{1,1}\tilde{Q}({\boldsymbol{\theta}}|{\boldsymbol{\theta}})]$
$\displaystyle=\mathbb{E}\left[\frac{\gamma_{1k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}{\boldsymbol{x}}_{1}{\boldsymbol{x}}_{1}^{T}\right]=:{\boldsymbol{\Sigma}}_{k}\,,$
(38)

where in the last step we used the notations defined in Lemma D.6.

We next write the following chain of terms:

$\displaystyle\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}=$
$\displaystyle\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}+\boldsymbol{u}_{k}^{T}(\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})-{\boldsymbol{\Sigma}}_{k})\boldsymbol{u}_{k}$
$\displaystyle+\boldsymbol{u}_{k}^{T}({\boldsymbol{\Sigma}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{u}_{k}+\boldsymbol{u}_{k}^{T}({\boldsymbol{\Sigma}}^{*}_{k}-\tilde{{\boldsymbol{\Sigma}}}_{k})\boldsymbol{u}_{k}$
$\displaystyle+\boldsymbol{u}_{k}^{T}(\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}}))\boldsymbol{u}_{k}$
$\displaystyle\geq\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}+\boldsymbol{u}_{k}^{T}(\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})-{\boldsymbol{\Sigma}}_{k})\boldsymbol{u}_{k}-O_{p}\left(\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2}\left\\{\log(pn)\sqrt{\frac{\log
p}{n}}+\eta_{{\sf est}}\sqrt{\log^{3}(pn)\log(nK)}\right\\}\right)$
$\displaystyle\geq\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}-o_{p}(\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2})\,,$
(39)

where in the first inequality we used Lemma D.6 and Proposition A.5, together
with the fact that $\left\|\boldsymbol{u}\right\|_{1}\leq
L\|\boldsymbol{x}_{\textrm{new}}\|_{2}$. The last step in (39) follows from
the condition $\sqrt{n}\log(np)\eta_{{\sf est}}^{2}=o(1)$ according to
Assumption (A1), and the assumption $\log(p)=o(n^{1/4}/\sqrt{\log(n)})$.

We next lower bound
$\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}$.
For any feasible solution $\boldsymbol{u}$ of optimization (12) we have:

$\displaystyle\boldsymbol{u}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}\geq\boldsymbol{u}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}+t((1-\lambda_{k})\|\boldsymbol{x}_{\textrm{new}}\|^{2}_{2}-\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}),$
(40)

for any $t>0$. The last inequality holds true because by the constraint of
optimization (12), we have

$\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2}-\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}\leq\left|\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}-\|\boldsymbol{x}_{\textrm{new}}\|^{2}_{2}\right|\leq\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2}\lambda_{k}.$

Minimizing over all feasible $\boldsymbol{u}$ gives

$\displaystyle\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\geq\min_{\boldsymbol{u}}\left\\{\boldsymbol{u}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}+t((1-\lambda_{k})\|\boldsymbol{x}_{\textrm{new}}\|^{2}_{2}-\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u})\right\\}\,.$

The minimizer $\boldsymbol{u}^{*}$ satisfies
$\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}^{*}=\frac{t}{2}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{x}_{\textrm{new}}$.
Substituting for $\boldsymbol{u}^{*}$, we obtain

$\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\geq-\frac{t^{2}}{4}\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{x}_{\textrm{new}}+t(1-\lambda_{k})\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}^{2}\,.$

Optimizing this bound over $t$, we get

$\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\geq\frac{(1-\lambda_{k})^{2}\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}^{4}}{\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{x}_{\textrm{new}}},$

with the optimal choice
$t^{*}=\frac{2(1-\lambda_{k})\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}^{2}}{\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{x}_{\textrm{new}}}>0$.

Similar to the proof of Lemma D.6, D.7, we have
$\left|\frac{\boldsymbol{x}_{\textrm{new}}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{x}_{\textrm{new}}}{\boldsymbol{x}_{\textrm{new}}^{T}{\boldsymbol{\Sigma}}_{k}\boldsymbol{x}_{\textrm{new}}}-1\right|=O_{p}(\sqrt{(\log
p)/n})$, and hence
$\boldsymbol{u}_{k}^{T}\tilde{{\boldsymbol{\Sigma}}}_{k}\boldsymbol{u}_{k}\geq
C\left\|\boldsymbol{x}_{\textrm{new}}\right\|_{2}^{2}$, which in conjunction
with equation (39) gives the desired result.

###  C.4 Proof of Proposition A.4

By the triangle inequality we have

$\displaystyle\left|\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|\leq\left|\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\sigma_{k}^{2}}\right|+\frac{|\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})-\gamma_{i,k}({\boldsymbol{\theta}})|}{\sigma_{k}^{2}}\,.$
(41)

The first term can be bounded as follows by recalling Condition (A1) on the
error term $|\widehat{\sigma}_{k}-\sigma_{k}|$:

$\left|\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\sigma_{k}^{2}}\right|=\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\sigma_{k}^{2}}\cdot\left|\frac{\sigma_{k}^{2}}{\widehat{\sigma}_{k}^{2}}-1\right|\leq\frac{1}{\sigma_{k}^{2}}\cdot\left|\frac{\sigma_{k}^{2}}{\widehat{\sigma}_{k}^{2}}-1\right|=O(|\sigma_{k}-\widehat{\sigma}_{k}|)=O_{p}(\eta_{{\sf
est}})\,,$

uniformly over all $i\in[n]$, $k\in[K]$. For the second term, by using Lemma
D.3 together with Assumption (A1) on the error term
$\left\|\widehat{{\boldsymbol{\theta}}}-{\boldsymbol{\theta}}\right\|_{1}$ we
get

$\sup_{i\in[n],k\in[K]}\frac{|\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})-\gamma_{i,k}({\boldsymbol{\theta}})|}{\sigma_{k}^{2}}=O_{p}(\sqrt{\log(np)\log(nK)}\eta_{{\sf
est}})\,.$

Combining the above two bounds into (41) we get that

$\displaystyle\sup_{i\in[n],k\in[K]}\left|\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|=O_{p}(\sqrt{\log(np)\log(nK)}\eta_{{\sf
est}})\,.$ (42)

Next, we observe that

$\displaystyle\left\|(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\right\|_{\infty}$
$\displaystyle\leq\max_{ik}|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}|\cdot\max_{i}\|x_{i}\|_{\infty}=O_{p}(\sqrt{\log(nK)}\cdot\sqrt{\log(np)})$

by Lemma D.2. Using this result combined with (42) gives the bound

$\displaystyle\left\|\frac{1}{n}\sum_{i=1}^{n}\left(\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right)(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\right\|_{\infty}$
$\displaystyle\leq\sup_{i\in[n],k\in[K]}\left|\frac{\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{i,k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|\cdot\frac{1}{n}\sum_{i=1}^{n}\left\|(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}\right\|_{\infty}$
$\displaystyle=O_{p}\Big{(}\log(np)\log(nK)\eta_{{\sf est}}\Big{)}$
$\displaystyle=o_{p}\Big{(}\sqrt{\log(np)}\cdot{\log(nK)}\cdot
n^{-1/4}\Big{)},$

where the last line follows from our assumption on the estimation rate
$\eta_{{\sf est}}$, cf. equation (22). Finally note that by invoking the
assumption $\log(p)=o(n^{1/4}/\sqrt{\log(n)})$, the last term is $o_{p}(1)$.

###  C.5 Proof of Proposition A.5

Note that for a matrix $\mathbf{A}$ and a vector $\boldsymbol{u}$ we have

$|\boldsymbol{u}^{T}\mathbf{A}\boldsymbol{u}|=|\sum_{ij}A_{ij}u_{j}u_{j}|\leq\left|\mathbf{A}\right|_{\infty}\sum_{ij}|u_{i}||u_{j}|=\left|\mathbf{A}\right|_{\infty}\left\|\boldsymbol{u}\right\|_{1}^{2}\,.$

From Lemma D.7,
$|\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})-\boldsymbol{I}^{\beta}_{k}({\boldsymbol{\theta}})|_{\infty}=O_{p}\left(\sqrt{\log^{3}(np)\log(nK)}\eta_{{\sf
est}}+\log(np)\sqrt{\frac{\log(p)}{n}}\right)$.

Given that $\left\|\boldsymbol{u}_{k}\right\|_{1}\leq
L\|\boldsymbol{x}_{\textrm{new}}\|_{2}$ for a constant $L$, and
$\sqrt{n}\log(np)\eta_{{\sf est}}^{2}=o(1)$ according to Assumption (A1) we
have

$\left|\boldsymbol{u}_{k}^{T}\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})\boldsymbol{u}_{k}-\boldsymbol{u}_{k}^{T}\boldsymbol{I}^{\beta}_{k}({\boldsymbol{\theta}})\boldsymbol{u}_{k}\right|=\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2}\cdot
O_{p}\left(\log(np)\sqrt{\log(nK)}n^{-1/4}+\log(np)\sqrt{\frac{\log(p)}{n}}\right)=o_{p}(\|\boldsymbol{x}_{\textrm{new}}\|_{2}^{2}),$

where the last step follows from the assumption
$\log(p)=o(n^{1/4}/\sqrt{\log(n)})$.

##  Appendix D Intermediate Lemmas and Proofs

This section summarizes several technical lemmas that were used in
establishing our theoretical results.

The first lemma is a classical maximal inequality for sub-Gaussian random
variables.

######  Lemma D.1.

(Maximal Inequality for Sub-Gaussians) Let
${\boldsymbol{x}}=(x_{1},\dots,x_{n})$ be a vector of zero-mean sub-Gaussian
random variables with variances $v_{1}^{2},\dots,v_{n}^{2}$, respectively.
Then, for any constant $c>0$ we have

$\sup_{i=1,\dots,n}\left|x_{i}\right|\leq v_{\max}\sqrt{2c\log n},$

with probability at least $1-2n^{1-c}$, where $v_{\max}=\max_{i\in[n]}v_{i}$.

The above result is obtained by using the tail bound of sub-Gaussian
variables, followed by a simple union bound, and therefore it does not require
the random variables to be independent.

Our next lemma is a simple corollary of the above maximal inequality.

######  Lemma D.2.

Define the following probability events:

$\displaystyle\mathcal{E}_{1}$
$\displaystyle:=\left\\{\max_{i\in[n]}\left\|{\boldsymbol{x}}_{i}\right\|_{\infty}\leq
C\sqrt{\log(np)}\right\\},$ $\displaystyle\mathcal{E}_{2}$
$\displaystyle:=\left\\{\max_{i\in[n],k\in[K]}\;|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}|\leq
C\sqrt{\log(nK)}\right\\}\,,$

and let $\mathcal{E}:=\mathcal{E}_{1}\cap\mathcal{E}_{2}$. Then, under
Assumptions (A2)–(A3) we have $\mathbb{P}(\mathcal{E})\geq 1-\frac{4}{n^{2}}$,
for large enough constant $C>0$.

###### Proof.

we write
$x_{ir}=\langle{\boldsymbol{\Sigma}}^{1/2}{\boldsymbol{e}}_{r},{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\rangle$.
Therefore,

$\displaystyle\|x_{ir}\|_{\psi_{2}}\leq\left\|{\boldsymbol{\Sigma}}^{1/2}{\boldsymbol{e}}_{r}\right\|_{2}\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}\leq
C_{\Sigma}^{1/2}\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}<\kappa\,,$
(43)

for all $r\in[p]$, $i\in[n]$ and some constant $\kappa<\infty$ by Assumption
(A2). By Lemma D.1, and for sufficiently large constant $C>0$, we get
$\mathbb{P}(\mathcal{E}_{1})\geq 1-\frac{2}{n^{2}}$.

To bound probability of $\mathcal{E}_{2}$, suppose that sample $i$ belongs to
group $\ell$, by which we can write

$y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}={\boldsymbol{x}}_{i}^{T}({\boldsymbol{\beta}}_{\ell}-{\boldsymbol{\beta}}_{k})+\varepsilon_{i}\,.$

We then have

$\displaystyle\|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\|_{\psi_{2}}$
$\displaystyle\leq\max_{\ell\in[K]}\|{\boldsymbol{x}}_{i}^{T}({\boldsymbol{\beta}}_{\ell}-{\boldsymbol{\beta}}_{k})+\varepsilon_{i}\|_{\psi_{2}}$
$\displaystyle\leq\max_{\ell\in[K]}\left(\|{\boldsymbol{x}}_{i}^{T}({\boldsymbol{\beta}}_{\ell}-{\boldsymbol{\beta}}_{k})\|_{\psi_{2}}+\|\varepsilon_{i}\|_{\psi_{2}}\right)$
$\displaystyle\leq\max_{\ell\in[K]}\left(\left\|{\boldsymbol{\Sigma}}^{1/2}({\boldsymbol{\beta}}_{\ell}-{\boldsymbol{\beta}}_{k})\right\|_{2}\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}+\sigma_{\ell}\right)$
$\displaystyle\leq\max_{\ell\in[K]}\left(C_{\Sigma}\left\|{\boldsymbol{\beta}}_{\ell}-{\boldsymbol{\beta}}_{k}\right\|_{2}\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}+\sigma_{\ell}\right)\,.$

Taking maximum over $i,k$ from both sides, we get

$\displaystyle\sup_{i\in[n],k\in[K]}\|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\|_{\psi_{2}}\leq
C_{\Sigma}\left(\max_{i\in[n]}\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}\right)\left(\max_{k,\ell\in[K]}\left\|{\boldsymbol{\beta}}_{\ell}-{\boldsymbol{\beta}}_{k}\right\|_{2}\right)+\max_{k\in[K]}\sigma_{k}\,.$
(44)

Recalling Assumptions (A2) and (A3), we get
$\sup_{i\in[n],k\in[K]}\|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\|_{\psi_{2}}\leq\kappa^{\prime}$,
for some constant $\kappa^{\prime}<\infty$. Therefore, by another application
of Lemma D.1, and for sufficiently large constant $C>0$, we get
$\mathbb{P}(\mathcal{E}_{2})\geq 1-\frac{2}{n^{2}}$.

Combining the two probability bounds we get $\mathbb{P}(\mathcal{E})\geq
1-\mathbb{P}(\mathcal{E}_{1})-\mathbb{P}(\mathcal{E}_{2})\geq
1-\frac{4}{n^{2}}$. ∎

While Assumption (A1) concerns the estimation error
$\left\|\widehat{{\boldsymbol{\theta}}}-{\boldsymbol{\theta}}\right\|_{1}$, in
our analysis we often need to control the perturbation of different functions
of ${\boldsymbol{\theta}}$. A useful step for these bounds is a control on the
Lipschitz factor of $\gamma_{i,k}({\boldsymbol{\theta}})$, which is the
subject of the next lemma.

######  Lemma D.3.

(Lipschitzness of $\gamma_{i,k}({\boldsymbol{\theta}})$) On event
$\mathcal{E}$, defined in Lemma D.2, the Lipschitz factor of
$\gamma_{i,k}({\boldsymbol{\theta}})$ with respect to the $\ell_{1}$ norm is
$O_{p}(\log(np))$, uniformly over all $i\in[n]$, $k\in[K]$. As a result,

$\sup_{i\in[n],k\in[K]}\frac{|\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})-\gamma_{i,k}({\boldsymbol{\theta}})|}{\left\|\widehat{{\boldsymbol{\theta}}}-{\boldsymbol{\theta}}\right\|_{1}}=O_{p}(\sqrt{\log(np)\log(nK)})\,.$

###### Proof.

To prove the claim, it suffices to show that

$\sup_{i\in[n],k\in[K]}\left(\left\|\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\beta}}_{k}}\right\|_{\infty},\left\|\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\alpha}}_{k}}\right\|_{\infty},\left|\frac{\partial\gamma_{i,k}}{\partial\sigma_{k}}\right|\right)=O_{p}(\sqrt{\log(np)\log(nK)}).$

Recall that
$\gamma_{i,k}({\boldsymbol{\theta}})=\frac{\pi_{k}({\boldsymbol{x}}_{i})\phi_{k}({\boldsymbol{x}}_{i},y_{i})}{\sum_{\ell=1}^{K}\pi_{\ell}({\boldsymbol{x}}_{i})\phi_{\ell}({\boldsymbol{x}}_{i},y_{i})}$,
with $\phi_{k}$ given by (5). A simple algebraic calculation shows that for a
function of form
$f_{k}(\boldsymbol{z})=\frac{c_{k}e^{z_{k}}}{\sum_{\ell}c_{\ell}e^{z_{\ell}}}$,
we have

$\displaystyle\frac{\partial}{\partial z_{k}}f_{k}=f_{k}(z)-f_{k}^{2}(z)\,.$
(45)

Applying this result, we obtain

$\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\beta}}_{k}}=(\gamma_{i,k}-\gamma^{2}_{i,k})\frac{1}{\sigma_{k}^{2}}(y_{i}-{\boldsymbol{x}}^{T}_{i}{\boldsymbol{\beta}}_{k}){\boldsymbol{x}}_{i}.$

Therefore, on the event $\mathcal{E}$ we have

$\displaystyle\left\|\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\beta}}_{k}}\right\|_{\infty}$
$\displaystyle\leq\frac{1}{\sigma_{k}^{2}}\left|\gamma_{i,k}-\gamma^{2}_{i,k}\right||y_{i}-{\boldsymbol{x}}^{T}_{i}{\boldsymbol{\beta}}_{k}|\cdot\left\|{\boldsymbol{x}}_{i}\right\|_{\infty}$
$\displaystyle\leq\frac{1}{4\sigma_{k}^{2}}C^{2}\sqrt{\log(np)\log(nK)}=O(\sqrt{\log(np)\log(nK)}),$

where we used the definition of event $\mathcal{E}$ and the fact that since
$0\leq\gamma_{i,k}\leq 1$, $\gamma_{i,k}(1-\gamma_{i,k})\leq\frac{1}{4}$.

Similarly, we can bound the partial derivative with respect to
${\boldsymbol{\alpha}}_{k}$. By another application of (45), we obtain

$\displaystyle\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\alpha}}_{k}}$
$\displaystyle=(\gamma_{i,k}-\gamma^{2}_{i,k}){\boldsymbol{x}}_{i}\,,$
$\displaystyle\left\|\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\alpha}}_{k}}\right\|_{\infty}$
$\displaystyle\leq|\gamma_{i,k}-\gamma^{2}_{i,k}|\;\left\|{\boldsymbol{x}}_{i}\right\|_{\infty}\leq\frac{1}{4}\left\|{\boldsymbol{x}}_{i}\right\|_{\infty}=O(\sqrt{\log(np)})\,,$
$\displaystyle\sup_{i\in[n],k\in[K]}\left\|\frac{\partial\gamma_{i,k}}{\partial{\boldsymbol{\alpha}}_{k}}\right\|_{\infty}$
$\displaystyle\leq\sup_{i}\|{\boldsymbol{x}}_{i}\|_{\infty}\leq
O(\sqrt{\log(np)})\,.$

Finally, we bound the partial derivative with respect to $\sigma_{k}$. By
another application of (45), we obtain

$\displaystyle\left|\frac{\partial\gamma_{i,k}}{\partial\sigma_{k}}\right|$
$\displaystyle=(\gamma_{i,k}-\gamma^{2}_{i,k})\frac{1}{\sigma_{k}^{3}}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k})^{2}=O(\log(nK))\,.$

This completes the proof.

∎

######  Lemma D.4.

(Mixture Estimation Error) On the event $\mathcal{E}$, defined in Lemma D.2,
we have

$\sup_{i\in[n],k\in[K]}\left|\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|=O_{p}\left(\sqrt{\log(np)\log(nK)}\eta_{{\sf
est}}\right).$

###### Proof.

We have that

$\sup_{i\in[n],k\in[K]}\left|\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|\leq\underbrace{\sup_{i,k}{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}\left|\frac{1}{\widehat{\sigma}_{k}^{2}}-\frac{1}{\sigma_{k}^{2}}\right|}_{(a)}+\underbrace{\sup_{i,k}\left|\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})-\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|}_{(b)}.$

For $(a)$, we have

$(a)\leq\sup_{k}\left|\frac{1}{\widehat{\sigma}_{k}^{2}}-\frac{1}{\sigma_{k}^{2}}\right|=\sup_{k}\frac{|\widehat{\sigma}_{k}-\sigma_{k}|\;|\widehat{\sigma}_{k}+\sigma_{k}|}{\sigma_{k}^{2}\widehat{\sigma}_{k}^{2}}=O_{p}\left(\eta_{{\sf
est}}\right)\,,$

by Assumption (A1) and (A3). For $(b)$, as shown in Lemma D.3 we have

$\sup_{i,k}\left|\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})-\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|\leq\frac{1}{\sigma_{k}^{2}}O_{p}(\sqrt{\log(np)\log(nK)}\|\widehat{{\boldsymbol{\theta}}}-{\boldsymbol{\theta}}\|_{1})=O_{p}\left(\sqrt{\log(np)\log(nK)}\eta_{{\sf
est}}\right),$

where in the last step we used Assumption (A1) and (A3). ∎

The next lemma is a concentration result on the covariate vectors
${\boldsymbol{x}}_{i}$ which will be used in our analysis.

######  Lemma D.5.

Let $\\{{\boldsymbol{x}}_{i}\\}_{i=1}^{n}$ satisfy Assumption (A2). Then for
any fixed unit vector $\boldsymbol{\xi}$ we have

$\mathbb{P}\left\\{\sup_{\ell\in[p]}\left|\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|-\mathbb{E}\left[\left|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|\right]\right|\geq
t\right\\}\leq
2p\exp\left(-c\min\left(\frac{nt^{2}}{C^{2}},\frac{nt}{C}\right)\right)\,,$

and

$\mathbb{E}\left[|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle|\right]\leq
C,$

for some positive constant $C$.

###### Proof.

To obtain the first result, we first note that
$\left|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|$
is a product of two sub-Gaussian random variables and thus is a sub-
exponential variable. Applying Theorem 2.8.1 of Vershynin (2018) to a fixed
$\ell\in[p]$, we get that for
$C\geq\max_{i\in[n]}\left\|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right\|_{\psi_{1}}$,

$\displaystyle\mathbb{P}\left\\{\left|\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|-\mathbb{E}\left[\left|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|\right]\right|\geq
t\right\\}\leq
2\exp\left(-c\min\left(\frac{nt^{2}}{C^{2}},\frac{nt}{C}\right)\right),$ (46)

where $\|\cdot\|_{\psi_{1}}$ is the sub-exponential norm of a random variable
and $c>0$ is an absolute constant.

To obtain our intended bound, we compute the sub-exponential norms
$\left\|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right\|_{\psi_{1}}$.
By Lemma 2.7.7 of Vershynin (2018), we have

$\displaystyle\|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\|_{\psi_{1}}$
$\displaystyle\leq\|x_{i,\ell}\|_{\psi_{2}}\|\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\|_{\psi_{2}}$
$\displaystyle=\|\langle{\boldsymbol{\Sigma}}^{1/2}{\boldsymbol{e}}_{\ell},{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\rangle\|_{\psi_{2}}\|\langle{\boldsymbol{\Sigma}}^{1/2}\boldsymbol{\xi},{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\rangle\|_{\psi_{2}}$
$\displaystyle\leq\left\|{\boldsymbol{\Sigma}}^{1/2}{\boldsymbol{e}}_{\ell}\right\|_{2}\left\|{\boldsymbol{\Sigma}}^{1/2}\boldsymbol{\xi}\right\|_{2}\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}^{2}$
$\displaystyle\leq C^{\prime}C_{\Sigma}=:C\,,$ (47)

where we used that $\|{\boldsymbol{\Sigma}}\|_{{\rm op}}\leq C_{\Sigma}$ and
$\|{\boldsymbol{\Sigma}}^{-1/2}{\boldsymbol{x}}_{i}\|_{\psi_{2}}\leq
C^{\prime}$, for some constants $C^{\prime},C_{\Sigma}$, per Assumption (A2).

To prove the second part of the lemma, we note that by definition

$\|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\|_{\psi_{1}}=\sup_{q\geq
1}q^{-1}\mathbb{E}[|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle|^{q}]^{1/q}\,,$

by which we have

$\mathbb{E}[|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle|]\leq\|x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\|_{\psi_{1}}\leq
C\,,$

for all $\ell\in[p]$. ∎

######  Lemma D.6.

(Covariance Matrix Estimation Error) Let

$\displaystyle\tilde{{\boldsymbol{\Sigma}}}_{k}$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$
$\displaystyle{\boldsymbol{\Sigma}}_{k}^{*}$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{ik}\left({\boldsymbol{\theta}}\right)}{\sigma_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$
$\displaystyle{\boldsymbol{\Sigma}}_{k}$
$\displaystyle=\mathbb{E}\left[\frac{\gamma_{1k}\left({\boldsymbol{\theta}}\right)}{\sigma_{k}^{2}}{\boldsymbol{x}}_{1}{\boldsymbol{x}}_{1}^{T}\right].$

Consider a set of unit-norm vectors
$\mathcal{C}=\\{\boldsymbol{\xi}_{1},\dotsc,\boldsymbol{\xi}_{p^{m}}\\}$, for
a fixed integer $m\geq 1$. Under Assumption (A2), and on the event
$\mathcal{E}$ defined in Lemma D.2, we have

$\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|(\tilde{{\boldsymbol{\Sigma}}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{\xi}\right\|_{\infty}=O_{p}\left(\eta_{{\sf
est}}\sqrt{\log(np)\log(nK)}\right)\,,$

and

$\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|({\boldsymbol{\Sigma}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{\xi}\right\|_{\infty}=O_{p}\left(\sqrt{\frac{\log
p}{n}}\right).$

###### Proof.

We have

$\displaystyle\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|(\tilde{{\boldsymbol{\Sigma}}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{\xi}\right\|_{\infty}$
$\displaystyle\leq\sup_{\ell\in[p],\boldsymbol{\xi}\in\mathcal{C}}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|\left|x_{i,\ell}\right|\left|\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|$
$\displaystyle\leq\sup_{i\in[n]}\left|\frac{\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\right|\cdot\sup_{\ell\in[p],\boldsymbol{\xi}\in\mathcal{C}}\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\right|\left|\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|.$

For the first component, we leverage Lemma D.4, which proves a
$O_{p}(\eta_{{\sf est}}\sqrt{\log(np)\log(nK)})$ bound. For the second
component, we use Lemma D.5 with $t=C\sqrt{\frac{2m\log(p)}{cn}}$ and union
bound over the set $\mathcal{C}$, which gives

$\sup_{\ell\in[p],\boldsymbol{\xi}\in\mathcal{C}}\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\right|\left|\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right|=O_{p}\left(1+\sqrt{\frac{\log
p}{n}}\right)=O_{p}(1).$

Putting them together, we have

$\displaystyle\sup_{\boldsymbol{\xi}\in\mathcal{C}}\left\|(\tilde{{\boldsymbol{\Sigma}}}_{k}-{\boldsymbol{\Sigma}}_{k}^{*})\boldsymbol{\xi}\right\|_{\infty}$
$\displaystyle=O_{p}\left(\eta_{{\sf est}}\sqrt{\log(np)\log(nK)}\right)\,.$

To prove the second result, we have

$\sup_{\ell\in[p],\boldsymbol{\xi}\in\mathcal{C}}\left\|({\boldsymbol{\Sigma}}_{k}^{*}-{\boldsymbol{\Sigma}}_{k})\boldsymbol{\xi}\right\|_{\infty}\leq\sup_{\ell\in[p],\boldsymbol{\xi}\in\mathcal{C}}\left|\frac{1}{n}\sum_{i=1}^{n}\left\\{\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle-\mathbb{E}\left[\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right]\right\\}\right|.$

As shown in (47),
$x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle$ have bounded
sub-exponential norm and since
$\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}$ is bounded, we
have that
$\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle$
has bounded sub-exponential norm. Therefore, by Theorem 2.8.1 of Vershynin
(2018) (similar to Lemma D.5) and a union bound over $\ell\in[p]$ and
$\boldsymbol{\xi}\in\mathcal{C}$, we obtain,

$\sup_{\ell\in[p],\boldsymbol{\xi}\in\mathcal{C}}\left|\frac{1}{n}\sum_{i=1}^{n}\left\\{\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle-\mathbb{E}\left[\frac{\gamma_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}x_{i,\ell}\langle{\boldsymbol{x}}_{i},\boldsymbol{\xi}\rangle\right]\right\\}\right|=O_{p}\left(\sqrt{\frac{\log
p}{n}}\right),$

which completes the proof of the second part. ∎

######  Lemma D.7.

(Fisher Information Estimation Error) Let

$\displaystyle\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}^{2}_{k}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}}{\widehat{\sigma}_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$
$\displaystyle\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma^{2}_{k}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}{\sigma_{k}^{2}}{\boldsymbol{x}}_{i}{\boldsymbol{x}}_{i}^{T},$
$\displaystyle\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})$
$\displaystyle=\mathbb{E}\left[\frac{\gamma^{2}_{1k}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\frac{\left(y_{1}-{\boldsymbol{x}}_{1}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}{\sigma_{k}^{2}}{\boldsymbol{x}}_{1}{\boldsymbol{x}}_{1}^{T}\right].$

Under Assumptions (A1), (A2) and (A3), We have

$\left|\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})-\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}=O_{p}\left(\sqrt{\log^{3}(np)\log(nK)}\eta_{{\sf
est}}\right),$

and

$\left|\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}=O_{p}\left(\log(np)\sqrt{\frac{\log(p)}{n}}\right).$

###### Proof.

We have

$\displaystyle\left|\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})-\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}$
$\displaystyle\leq\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}{\sigma_{k}^{2}}\right|\left|x_{i,\ell}\right|\left|x_{i,r}\right|$
$\displaystyle\leq\left(\sup_{i\in[n]}\left|\frac{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})}{\widehat{\sigma}_{k}^{2}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}}{\widehat{\sigma}_{k}^{2}}-\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}\frac{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}{\sigma_{k}^{2}}\right|\right)\left(\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\right|\left|x_{i,r}\right|\right)$
$\displaystyle\leq\left(\underbrace{\sup_{i\in[n]}\left|\frac{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})}{\sigma_{k}^{4}}\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}\right|\left|1-\frac{\sigma_{k}^{4}}{\widehat{\sigma}_{k}^{4}}\right|}_{(a)}+\underbrace{\sup_{i\in[n]}\frac{1}{{\sigma_{k}^{4}}}\left|{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}-\gamma^{2}_{ik}({\boldsymbol{\theta}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}\right|}_{(b)}\right)$
$\displaystyle\qquad\ \
\cdot\underbrace{\left(\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\right|\left|x_{i,r}\right|\right)}_{(c)}.$

On the event $\mathcal{E}$ we have that

$\displaystyle\sup_{i\in[n],k\in[K]}|y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}|$
$\displaystyle\leq\sup_{i\in[n],k\in[K]}|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}|+\sup_{i\in[n],k\in[K]}|{\boldsymbol{x}}_{i}^{T}({\boldsymbol{\beta}}_{k}-\widehat{{\boldsymbol{\beta}}}_{k})|$
$\displaystyle\leq
C\sqrt{\log(nK)}+(\sup_{i\in[n]}\left\|{\boldsymbol{x}}_{i}\right\|_{\infty})\left(\sup_{k\in[K]}\left\|{\boldsymbol{\beta}}_{k}-\widehat{{\boldsymbol{\beta}}}_{k}\right\|_{1}\right)$
$\displaystyle\leq C\left(\sqrt{\log(nK)}+\sqrt{\log(np)}\eta_{{\sf
est}}\right)$ $\displaystyle=O\left(\sqrt{\log(np)}\right)\,.$ (48)

Therefore part $(a)$ can be bounded as follows:

$\sup_{i\in[n]}\left|\frac{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}}{\sigma_{k}^{4}}\right|\left|1-\frac{\sigma_{k}^{4}}{\widehat{\sigma}_{k}^{4}}\right|\leq\frac{1}{\sigma_{k}^{4}}\left|1-\frac{\sigma_{k}^{4}}{\widehat{\sigma}_{k}^{4}}\right|\sup_{i\in[n]}\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}=O_{p}\left(\eta_{{\sf
est}}\log(np)\right),$

where we used that $\gamma_{i,k}(\widehat{{\boldsymbol{\theta}}})\leq 1$ and
$|\widehat{\sigma}_{k}-\sigma_{k}|=O_{p}(\eta_{{\sf est}})$ per Assumption
(A1).

For $(b)$, we write

$\displaystyle\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}-\gamma^{2}_{ik}({\boldsymbol{\theta}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}\right|$
$\displaystyle=\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}-\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}+\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}-\gamma^{2}_{ik}({\boldsymbol{\theta}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}\right|$
$\displaystyle\leq\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|(\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})-\gamma^{2}_{ik}({\boldsymbol{\theta}}))\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}\right|+\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|{\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)^{2}-\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}\right|$
$\displaystyle=\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|{\left(\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})-\gamma^{2}_{ik}({\boldsymbol{\theta}})\right)\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}\right|+\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|\gamma^{2}_{ik}(\widehat{{\boldsymbol{\theta}}})\left({\boldsymbol{x}}_{i}^{T}\left(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k}\right)\right)\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}+y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right)\right|$
$\displaystyle\leq\sup_{i\in[n]}\frac{2}{\sigma_{k}^{4}}|\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})-\gamma_{ik}({\boldsymbol{\theta}})|\left|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right|^{2}+\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left|{\boldsymbol{x}}_{i}^{T}\left(\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k}\right)\right|\left|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}+y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right|$
$\displaystyle\leq\sup_{i\in[n]}\frac{2}{\sigma_{k}^{4}}|\gamma_{ik}(\widehat{{\boldsymbol{\theta}}})-\gamma_{ik}({\boldsymbol{\theta}})|\left|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right|^{2}+\sup_{i\in[n]}\frac{1}{\sigma_{k}^{4}}\left\|{\boldsymbol{x}}_{i}\right\|_{\infty}\left\|\widehat{{\boldsymbol{\beta}}}_{k}-{\boldsymbol{\beta}}_{k}\right\|_{1}\left|y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}+y_{i}-{\boldsymbol{x}}_{i}^{T}\widehat{{\boldsymbol{\beta}}}_{k}\right|$
$\displaystyle=O(\log^{3/2}(np)\sqrt{\log(nK)}\eta_{{\sf
est}})+O(\sqrt{\log(nK)}\sqrt{\log(np)}\eta_{{\sf est}})$
$\displaystyle=O(\log^{3/2}(np)\sqrt{\log(nK)}\eta_{{\sf est}})\,,$

where in the penultimate step, we bounded the first term using Lemma D.3
together with Assumption (A1) and we bounded the second term using definition
of event $\mathcal{E}$ along with (48).

Finally, for $(c)$, we use Lemma D.5 with
$\boldsymbol{\xi}={\boldsymbol{e}}_{r}$ for $r\in[p]$ followed by a union
bound over $r$ to get

$\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|x_{i,\ell}\right|\left|x_{i,r}\right|=O_{p}(1).$

Putting $(a)$, $(b)$, and $(c)$ together, we have

$\left|\tilde{\boldsymbol{I}}_{k}^{\beta}(\widehat{{\boldsymbol{\theta}}})-\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}=O_{p}(\log^{3/2}(np)\sqrt{\log(nK)}\eta_{{\sf
est}}).$

To prove the second result, we have

$\displaystyle\mathbb{P}\left(\left|\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}\geq
t\right)$
$\displaystyle\leq\mathbb{P}\left(\left|\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}\geq
t;\mathcal{E}\right)+\mathbb{P}(\mathcal{E}^{c})$
$\displaystyle\leq\mathbb{P}\left(\left|\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}\geq
t;\mathcal{E}\right)+\frac{4}{n^{2}}\,,$ (49)

where we used the result of Lemma D.2 to bound $\mathbb{P}(\mathcal{E}^{c})$.
We next write

$\displaystyle\mathbb{P}\left(\left|\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}\geq
t;\mathcal{E}\right)$
$\displaystyle=\mathbb{P}\left(\sup_{\ell,r\in[p]}\left|\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}x_{i,\ell}x_{i,r}-\mathbb{E}\left[\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}{\left(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k}\right)^{2}}x_{i,\ell}x_{i,r}\right]\right|\geq
t;\mathcal{E}\right)$
$\displaystyle\leq\mathbb{P}\left(\Big{(}\sup_{i\in[n]}(y_{i}-{\boldsymbol{x}}_{i}^{T}{\boldsymbol{\beta}}_{k})^{2}\Big{)}\cdot\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}-\mathbb{E}\left[\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}\right]\right|\geq
t;\mathcal{E}\right)$
$\displaystyle\leq\mathbb{P}\left(\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}-\mathbb{E}\left[\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}\right]\right|\geq\frac{t}{C^{2}\log(np)};\mathcal{E}\right)$
$\displaystyle\leq\mathbb{P}\left(\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}-\mathbb{E}\left[\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}\right]\right|\geq\frac{t}{C^{2}\log(np)}\right)\,,$
(50)

where the second last step follows from definition of event $\mathcal{E}$.

Now by using equation (47) with $\boldsymbol{\xi}={\boldsymbol{e}}_{r}$, we
have that $x_{i,\ell}x_{i,r}$ is sub-exponential and since
$\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{2}}$ is bounded, we
get that
$\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}$
is a sub-exponential random variable. Thus, by Bernstein’s inequality (see
e.g. (Vershynin, 2018, Theorem 2.8.1)) and union bound over the $p^{2}$
choices of $\ell,r$, we obtain,

$\mathbb{P}\left(\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}-\mathbb{E}\left[\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}\right]\right|\geq\frac{t}{C^{2}\log(np)}\right)\leq
2p^{2}\exp\left(-c\frac{nt^{2}}{\log^{2}(np)}\right)\,.$

Therefore by choosing $t=2\sqrt{\frac{\log(p)}{cn}}\log(np)$ we obtain

$\displaystyle\mathbb{P}\left(\sup_{\ell,r\in[p]}\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}-\mathbb{E}\left[\frac{\gamma^{2}_{ik}({\boldsymbol{\theta}})}{\sigma_{k}^{4}}x_{i,\ell}x_{i,r}\right]\right|\geq\frac{2}{C^{2}}\sqrt{\frac{\log(p)}{cn}}\right)\leq
2p^{-2}\,.$ (51)

∎

Combining equations (49), (50) and (51) we arrive at

$\left|\tilde{\boldsymbol{I}}_{k}^{\beta}({\boldsymbol{\theta}})-\boldsymbol{I}_{k}^{\beta}({\boldsymbol{\theta}})\right|_{\infty}=O_{p}\left(\log(np)\sqrt{\frac{\log(p)}{n}}\right)\,,$

which completes the proof of the second claim.

## References

  * (1)
  * Athey et al. (2018) Athey, S., Imbens, G. W. & Wager, S. (2018), ‘Approximate residual balancing: debiased inference of average treatment effects in high dimensions’, Journal of the Royal Statistical Society: Series B (Statistical Methodology) 80(4), 597–623.
  * Balakrishnan et al. (2017) Balakrishnan, S., Wainwright, M. J. & Yu, B. (2017), ‘Statistical guarantees for the em algorithm: From population to sample-based analysis’, The Annals of Statistics 45(1), 77–120.
  * Bühlmann & Van De Geer (2011) Bühlmann, P. & Van De Geer, S. (2011), Statistics for high-dimensional data: methods, theory and applications, Springer Science & Business Media.
  * Cai & Guo (2017) Cai, T. T. & Guo, Z. (2017), ‘Confidence intervals for high-dimensional linear regression: Minimax rates and adaptivity’, The Annals of statistics 45(2), 615–646.
  * Cai et al. (2021) Cai, T., Tony Cai, T. & Guo, Z. (2021), ‘Optimal statistical inference for individualized treatment effects in high-dimensional models’, Journal of the Royal Statistical Society: Series B (Statistical Methodology) 83(4), 669–719.
  * Dempster et al. (1977) Dempster, A. P., Laird, N. M. & Rubin, D. B. (1977), ‘Maximum likelihood from incomplete data via the em algorithm’, Journal of the Royal Statistical Society: Series B (Methodological) 39(1), 1–22.
  * Guo et al. (2021) Guo, Z., Renaux, C., Bühlmann, P. & Cai, T. (2021), ‘Group inference in high dimensions with applications to hierarchical testing’, Electronic Journal of Statistics 15(2), 6633–6676.
  * Hamidieh (2018) Hamidieh, K. (2018), ‘A data-driven statistical model for predicting the critical temperature of a superconductor’, Computational Materials Science 154, 346–354.
  * Hastie et al. (2009) Hastie, T., Tibshirani, R., Friedman, J. H. & Friedman, J. H. (2009), The elements of statistical learning: data mining, inference, and prediction, Vol. 2, Springer.
  * Hunter & Lange (2004) Hunter, D. R. & Lange, K. (2004), ‘A tutorial on mm algorithms’, The American Statistician 58(1), 30–37.
  * Hyun (2022) Hyun, S. (2022), flowmix: Ocean Flow Cytometry Analysis.  R package version 0.0.0.9000.
  * Hyun et al. (2020) Hyun, S., Cape, M. R., Ribalet, F. & Bien, J. (2020), ‘Modeling cell populations measured by flow cytometry with covariates using sparse mixture of regressions’, arXiv preprint arXiv:2008.11251 .
  * Javanmard & Lee (2020) Javanmard, A. & Lee, J. D. (2020), ‘A flexible framework for hypothesis testing in high dimensions’, Journal of the Royal Statistical Society: Series B (Statistical Methodology) 82(3), 685–718.
  * Javanmard & Montanari (2014a) Javanmard, A. & Montanari, A. (2014a), ‘Confidence intervals and hypothesis testing for high-dimensional regression’, The Journal of Machine Learning Research 15(1), 2869–2909.
  * Javanmard & Montanari (2014b) Javanmard, A. & Montanari, A. (2014b), ‘Hypothesis testing in high-dimensional regression under the gaussian random design model: Asymptotic theory’, IEEE Transactions on Information Theory 60(10), 6522–6554.
  * Javanmard & Montanari (2018) Javanmard, A. & Montanari, A. (2018), ‘Debiasing the lasso: Optimal sample size for gaussian designs’, The Annals of Statistics 46(6A), 2593–2622.
  * Jordan & Jacobs (1994) Jordan, M. I. & Jacobs, R. A. (1994), ‘Hierarchical mixtures of experts and the em algorithm’, Neural computation 6(2), 181–214.
  * Klusowski et al. (2019) Klusowski, J. M., Yang, D. & Brinda, W. (2019), ‘Estimating the coefficients of a mixture of two linear regressions by expectation maximization’, IEEE Transactions on Information Theory 65(6), 3515–3524.
  * Kwon & Caramanis (2020) Kwon, J. & Caramanis, C. (2020), Em converges for a mixture of many linear regressions, in ‘International Conference on Artificial Intelligence and Statistics’, PMLR, pp. 1727–1736.
  * Kwon et al. (2019) Kwon, J., Qian, W., Caramanis, C., Chen, Y. & Davis, D. (2019), Global convergence of the em algorithm for mixtures of two component linear regression, in ‘Conference on Learning Theory’, PMLR, pp. 2055–2110.
  * Lei et al. (2018) Lei, J., G’Sell, M., Rinaldo, A., Tibshirani, R. J. & Wasserman, L. (2018), ‘Distribution-free predictive inference for regression’, Journal of the American Statistical Association 113(523), 1094–1111.
  * Lei & Wasserman (2014) Lei, J. & Wasserman, L. (2014), ‘Distribution-free prediction bands for non-parametric regression’, Journal of the Royal Statistical Society: Series B (Statistical Methodology) 76(1), 71–96.
  * Papadopoulos et al. (2002) Papadopoulos, H., Proedrou, K., Vovk, V. & Gammerman, A. (2002), Inductive confidence machines for regression, in ‘European Conference on Machine Learning’, Springer, pp. 345–356.
  * Quandt & Ramsey (1978) Quandt, R. E. & Ramsey, J. B. (1978), ‘Estimating mixtures of normal distributions and switching regressions’, Journal of the American statistical Association 73(364), 730–738.
  * Romano, Barber, Sabatti & Candès (2019) Romano, Y., Barber, R. F., Sabatti, C. & Candès, E. J. (2019), ‘With malice towards none: Assessing uncertainty via equalized coverage’, arXiv preprint arXiv:1908.05428 .
  * Romano, Patterson & Candes (2019) Romano, Y., Patterson, E. & Candes, E. (2019), ‘Conformalized quantile regression’, Advances in neural information processing systems 32.
  * Städler et al. (2010) Städler, N., Bühlmann, P. & van de Geer, S. A. (2010), ‘L1-penalization for mixture regression models’, TEST 19, 209–256.
  * Tripuraneni & Mackey (2019) Tripuraneni, N. & Mackey, L. (2019), ‘Debiasing linear prediction’, arXiv preprint arXiv:1908.02341 .
  * Van de Geer et al. (2014) Van de Geer, S., Bühlmann, P., Ritov, Y. & Dezeure, R. (2014), ‘On asymptotically optimal confidence regions and tests for high-dimensional models’, The Annals of Statistics 42(3), 1166–1202.
  * Van der Vaart (2000) Van der Vaart, A. W. (2000), Asymptotic statistics, Vol. 3, Cambridge university press.
  * Vershynin (2018) Vershynin, R. (2018), High-dimensional probability: An introduction with applications in data science, Vol. 47, Cambridge university press.
  * Vovk (2012) Vovk, V. (2012), Conditional validity of inductive conformal predictors, in ‘Asian conference on machine learning’, PMLR, pp. 475–490.
  * Vovk et al. (2005) Vovk, V., Gammerman, A. & Shafer, G. (2005), Algorithmic learning in a random world, Springer Science & Business Media.
  * Wang et al. (2014) Wang, Z., Gu, Q., Ning, Y. & Liu, H. (2014), ‘High dimensional expectation-maximization algorithm: Statistical optimization and asymptotic normality’, arXiv preprint arXiv:1412.8729 .
  * Yi & Caramanis (2015) Yi, X. & Caramanis, C. (2015), ‘Regularized em algorithms: A unified framework and statistical guarantees’, Advances in Neural Information Processing Systems 28.
  * Yi et al. (2014) Yi, X., Caramanis, C. & Sanghavi, S. (2014), Alternating minimization for mixed linear regression, in ‘International Conference on Machine Learning’, PMLR, pp. 613–621.
  * Zhang & Zhang (2014) Zhang, C.-H. & Zhang, S. S. (2014), ‘Confidence intervals for low dimensional parameters in high dimensional linear models’, Journal of the Royal Statistical Society: Series B (Statistical Methodology) 76(1), 217–242.
  * Zhang et al. (2020) Zhang, L., Ma, R., Cai, T. T. & Li, H. (2020), ‘Estimation, confidence intervals, and large-scale hypotheses testing for high-dimensional mixed linear regression’, arXiv preprint arXiv:2011.03598 .
  * Zhu & Zhang (2004) Zhu, H.-T. & Zhang, H. (2004), ‘Hypothesis testing in mixture regression models’, Journal of the Royal Statistical Society: Series B (Statistical Methodology) 66(1), 3–16.
  * Zhu & Bradic (2018) Zhu, Y. & Bradic, J. (2018), ‘Linear hypothesis testing in dense high-dimensional linear models’, Journal of the American Statistical Association 113(524), 1583–1600.