High-dimensional Analysis of Knowledge Distillation: Weak-to-Strong Generalization and Scaling Laws

## Abstract

Abstract A growing number of machine learning scenarios rely on knowledge distillation where one uses the output of a surrogate model as labels to supervise the training of a target model. In this work, we provide a sharp characterization of this process for ridgeless, high-dimensional regression, under two settings: (i) model shift, where the surrogate model is arbitrary, and (ii) distribution shift, where the surrogate model is the solution of empirical risk minimization with out-of-distribution data. In both cases, we characterize the precise risk of the target model through non-asymptotic bounds in terms of sample size and data distribution under mild conditions. As a consequence, we identify the form of the optimal surrogate model, which reveals the benefits and limitations of discarding weak features in a data-dependent fashion. In the context of weak-to-strong (W2S) generalization, this has the interpretation that (i) W2S training, with the surrogate as the weak model, can provably outperform training with strong labels under the same data budget, but (ii) it is unable to improve the data scaling law. We validate our results on numerical experiments both on ridgeless regression and on neural network architectures.

## 1 Introduction

The increasing number and diversity of machine learning models has motivated the development of techniques that leverage the output of one model to train a different one – a process known as knowledge distillation (Hinton et al., ). Variations of this approach include generating synthetic data from powerful language models (Wang et al., ; Gunasekar et al., ; Abdin et al., ), weak-to-strong generalization to obtain stronger models under weak supervision (Burns et al., ), and filtering/curating ML datasets via a smaller model to train a larger model (Fang et al., ; Lin et al., ). The diversity of these applications motivates a deeper understanding of the statistical properties and limits of the distillation process.

In this work, we focus on the scenario where a target/student model is trained on the labels of a surrogate/teacher model. Let ${\cal{D}}_{t}$ denote the target distribution, $(\bm{x}_{i},y_{i})_{i=1}^{n}$ be sampled i.i.d. from this distribution, and $p$ denote the dimension of the features $\bm{x}_{i}$. Given a surrogate model $s$, we create the synthetic labels $y_{i}^{s}=s(\bm{x}_{i})$ and obtain the target model by minimizing the empirical risk, i.e.,

$$ $\displaystyle\hat{f}=\arg\min_{f\in\mathcal{F}}\frac{1}{n}\sum_{i=1}^{n}\ell(y _{i}^{s},f(\bm{x}_{i})),$ (1) $$

where $\mathcal{F}$ denotes the hypothesis class and $\ell$ the loss function. This procedure motivates a few fundamental questions regarding (1): *(i)* What is the excess risk of $\hat{f}$ compared to that of the minimizer of the population risk $f_{\star}=\arg\min_{f\in\mathcal{F}}\operatorname{\mathbb{E}}_{(\bm{x},y)\sim{
\cal{D}}_{t}}[\ell(y,f(\bm{x}))]$? *(ii)* What is the optimal surrogate model $s$ that minimizes such excess risk? *(iii)* Can the optimal $s$ strictly outperform using true labels $y_{i}$ or setting $s=f_{\star}$ in (1)? Furthermore, in practice, $s$ itself is the outcome of an empirical risk minimization (ERM) procedure. Specifically, let ${\cal{D}}_{s}$ be the surrogate distribution, $(\tilde{\bm{x}}_{i},\tilde{y}_{i})_{i=1}^{m}$ be sampled i.i.d. from ${\cal{D}}_{s}$ (having feature dimension $p$), and assume

$$ $\displaystyle s=\arg\min_{f\in\mathcal{S}}\frac{1}{m}\sum_{i=1}^{m}\ell(\tilde {y}_{i},f(\tilde{\bm{x}}_{i})),$ (2) $$

where $\mathcal{S}$ denotes the surrogate hypothesis class.
Thus, as a final (and more challenging question), one may ask: *(iv)* How do the sample sizes $n,m$ and the data distributions ${\cal{D}}_{t},{\cal{D}}_{s}$ affect the performance of $\hat{f}$?

Figure: (a) Ground-truth and surrogate model weights
Refer to caption: x1.png

Main contributions. We address these questions in the context of high-dimensional ridgeless regression, where $\mathcal{F},\mathcal{S}$ are the class of linear models and $\ell$ is the quadratic loss. We provide a sharp non-asymptotic characterization of the test risk of $\hat{f}$ in two core settings:

- 1.
The surrogate $s$ is provided and we solve $\hat{f}$ via (1), which addresses the first three questions above;
- 2.
$\hat{f}$ is obtained via two stages of ERM, i.e., (2) followed by (1), which addresses the last question.

We focus on the
regime where the sample sizes $n,m$ and the feature dimension $p$ are all proportional and also allow for a distribution shift between ${\cal{D}}_{t}$ and ${\cal{D}}_{s}$, which correspond to the two ERM stages. Our theoretical guarantees in Section 3 precisely characterize the regression coefficients ${\bm{\beta}}_{s}$ of the optimal
surrogate model in terms of
the corresponding coefficients ${\bm{\beta}}_{\star}$ of the population risk minimizer $f_{\star}$ and the feature covariance. This is depicted in Figure 1(a) where ${\bm{\beta}}_{\star}$ and ${\bm{\beta}}_{s}$ are displayed by blue and green curves, respectively. This unveils a remarkable phenomenology in the process of knowledge distillation, that can be described as follows. Define the per-feature ‘gain’ of the optimal surrogate as $\texttt{gain}={\bm{\beta}}_{s}/{\bm{\beta}}_{\star}$ where the division is entrywise, which corresponds to the ratio between green and blue curves. We show that the gain vector is entirely controlled by the *covariance statistics*, denoted by $(\zeta_{i})_{i=1}^{p}$, that summarize the role of feature covariance and finite sample size $n$ in the test risk. There is a well-defined transition point, $\zeta_{i}=1-\Omega$ in Fig. 1(a), where the gain passes from strict amplification ($\texttt{gain}_{i}>1$) to strict shrinkage ($\texttt{gain}_{i}<1$) as we move from principal eigendirections to tail. Our theory also clarifies when we are better off discarding the weak features. The yellow curve shows the optimal surrogate when ${\bm{\beta}}_{s}$ is restricted to be a 0-1 mask on the entries of ${\bm{\beta}}_{\star}$ so that the surrogate has the direct interpretation of feature pruning. We show that beyond the transition point $\zeta_{i}^{2}=1-\Omega$, truncating the weak features that lie on the tail of the spectrum is strictly beneficial to distillation. This masked surrogate model can be viewed as a weak supervisor as it contains strictly fewer features compared to the target, revealing a success mechanism for weak-to-strong supervision.

Notably, as the sample size $n$ decreases, the optimal surrogate provably becomes sparser (transition points shift to the left) under a power-law decay covariance model. In Section 4, under this power-law model, we also quantify the performance gain that arises from the optimal surrogate and show that while the surrogate can strictly improve the test risk, it does not alter the exponent of the scaling law. This is depicted in Figure 1(b) where the surrogate risks are smaller but behave similarly to the ground-truth. Finally, in Section 5, we study the more intricate problem of two-stage ERM ((2) followed by (1)) and establish a non-asymptotic risk characterization that precisely captures the influences of the sample sizes $m,n$ and of the surrogate/target covariance matrices.

### 1.1 Related work

Our work relates to the topics of high-dimensional learning, distribution shift, scaling laws, and distillation.

High-dimensional risk characterization. There is a large body of literature dedicated to the study of linear regression in over- and under-parameterized regimes. Via random matrix theory tools, one can precisely study the test risk and various associated phenomena, such as benign overfitting (Bartlett et al., ) or double descent (Belkin et al., ).
Specifically, asymptotic and non-asymptotic risk characterizations for the minimum $\ell_{2}$-norm interpolator (i.e., ridgeless regression with $p\geq n$) have been provided in a recent line of work (Hastie et al., ; Cheng & Montanari, ; Han & Xu, ; Wu & Xu, ; Richards et al., ).

While the standard ERM formulation is well studied, the analysis of the distillation problem necessitates a fine-grained characterization of the ERM process. A series of papers (Chang et al., ; Montanari et al., ; Han & Xu, ) utilize Gaussian process theory (Thrampoulidis et al., ) to characterize the distribution of ridgeless estimators. Our theory builds on these to precisely characterize the distillation performance by *(i)* accounting for model and covariate shift and *(ii)* tracking the distribution across the two-stage problem where (2) is followed by (1). Our setting strictly subsumes the problem of characterizing the test risk under distribution shift, which relates to the recent papers (Patil et al., ; Song et al., ; Yang et al., ; Mallinar et al., ). A distinguishing feature of our work is that we precisely characterize the optimal surrogate model that minimizes the downstream target risk, as highlighted in Figure 1.

Distillation and weak-to-strong generalization. Mobahi et al. () provide a theoretical analysis of self-distillation, whereas Nagarajan et al. () study knowledge distillation in a teacher-student setting. A related problem is self-training which relies on progressively generating pseudo-labels for unlabeled data (Frei et al., ; Oymak & Gulcu, ; Wei et al., ). While these works consider low-dimensional settings or provide loose bounds, our study provides a sharp analysis of ridgeless over-parameterized regression. Closer to us, Jain et al. () investigate the benefit of *surrogate data* by employing both real and surrogate data in a single step of ERM, but the analysis is limited to isotropic covariance. Kolossov et al. () consider the problem of surrogate-based data selection. Finally, Charikar et al. (); Lang et al. () aim to demystify weak-to-strong generalization by formalizing the intuition that W2S generalization occurs when the strong
model avoids fitting the mistakes of the weak teacher. In contrast, our theory reveals that the strong student can benignly overfit the weak teacher, and in fact, a carefully crafted weak teacher provably outperforms strong labels, see again Figure 1.

Scaling laws. The dependence of the performance on the available statistical and computational resources is often empirically well-captured by a power-law (Hestness et al., ; Kaplan et al., ). This experimental evidence has led to a flurry of theoretical work aimed at characterizing the emergence of scaling laws, mostly focusing on linear regression (Spigler et al., ; Simon et al., ; Bahri et al., ; Paquette et al., ; Bordelon et al., ; Lin et al., ; Maloney et al., ). Bordelon et al. () analyze a random feature model trained with gradient descent via dynamical mean field theory. Jain et al. () consider scaling laws with surrogate data, whereas Sorscher et al. () study the benefits of data pruning.

## 2 Problem setup

Notation. Let $[p]$ denote the set $\{1,\cdots,p\}$ for an integer $p\geq 1$.
We use lower-case and upper-case bold letters (e.g., $\bm{x},{\bm{X}}$) to represent vectors and matrices, respectively; $x_{i}$ denotes the $i$-th entry of the vector $\bm{x}$, ${\bm{X}}^{\dagger}$ the pseudo-inverse of the matrix ${\bm{X}}$, and $\texttt{tr}\left({\bm{X}}\right)$ the trace of ${\bm{X}}$.

We consider a two-stage linear learning problem. In the first stage, pairs of labels and input features from the distribution ${\cal{D}}_{s}$ are used to produce an estimate of the ground-truth parameter, which is then used to generate labels along with input features from a different distribution ${\cal{D}}_{t}$. The second stage uses these generated labels to obtain the final estimate of the ground-truth parameter. The models trained in the first and second stages are referred to as surrogate and target models, respectively.

Stage 1: Surrogate model. We consider a data distribution $(\tilde{\bm{x}},\tilde{y})\sim{\cal{D}}_{s}$ following the linear model $\tilde{y}=\tilde{\bm{x}}^{\top}{\bm{\beta}_{\star}}+{\tilde{z}}$, where ${\bm{\beta}_{\star}}\in\mathbb{R}^{p}$, $\tilde{\bm{x}}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{s})$ and ${\tilde{z}}\sim\mathcal{N}(0,{\sigma}_{s}^{2})$ is independent of $\tilde{\bm{x}}$. Let $\{(\tilde{\bm{x}}_{i},\tilde{y}_{i})_{i=1}^{m}\}$ be the dataset for the surrogate model drawn i.i.d. from ${\cal{D}}_{s}$. We analyze both under- and over-parametrized settings: in the former, we estimate ${\bm{\beta}_{\star}}$ by minimizing the quadratic loss; in the latter, we estimate ${\bm{\beta}_{\star}}$ as the minimum norm interpolator. As a result, the estimator of the surrogate model can be written as follows:

$$ $\displaystyle{\bm{\beta}}^{s}=\text{Est}(\tilde{\bm{X}},\tilde{\bm{y}}):= \begin{cases}\arg\min_{\bm{\beta}}\|{\tilde{\bm{y}}-\tilde{\bm{X}}{\bm{\beta}} }\|_{2}^{2},&\quad\text{if $m\geq p$},\\ \arg\min_{\bm{\beta}}\{\|{{\bm{\beta}}}\|_{2}:\tilde{\bm{X}}{\bm{\beta}}= \tilde{\bm{y}}\},&\quad\text{if $m<p$},\end{cases}$ (3) $$

where $\tilde{\bm{X}}=[\tilde{\bm{x}}_{1}^{\top},\dots,\tilde{\bm{x}}_{m}^{\top}]^{
\top}\in\mathbb{R}^{m\times p}$ and $\tilde{\bm{y}}=[y_{1},\dots,y_{m}]^{\top}\in\mathbb{R}^{m}$.

Stage 2: Target model. Given ${\bm{\beta}}^{s}\in\mathbb{R}^{p}$, we consider another data distribution $(\bm{x},y^{s})\sim{\cal{D}}_{t}({\bm{\beta}}^{s})$ following the linear model $y^{s}=\bm{x}^{\top}{\bm{\beta}}^{s}+z$, where $\bm{x}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{t})$ and $z\sim\mathcal{N}(0,{\sigma}_{t}^{2})$. Let $\{(\bm{x}_{i},y^{s}_{i})_{i=1}^{n}\}$ be the dataset for the target model drawn i.i.d. from ${\cal{D}}_{t}({\bm{\beta}}^{s})$. As for the surrogate model, the estimator for the target model is defined as

$$ $\displaystyle{\bm{\beta}}^{s2t}=\text{Est}({\bm{X}},\bm{y}^{s}),$ (4) $$

where ${\bm{X}}=[\bm{x}_{1}^{\top},\ldots,\bm{x}_{n}^{\top}]^{\top}\in\mathbb{R}^{n
\times p}$ and $\bm{y}^{s}=[y^{s}_{1},\ldots,y^{s}_{n}]^{\top}\in\mathbb{R}^{n}$. Our analysis will generally apply to an arbitrary ${\bm{\beta}}^{s}$ choice and will not require it to be the outcome of (3).
Finally, we define the excess (population) risk for a given estimator $\hat{\bm{\beta}}\in\mathbb{R}^{p}$ as

$$ $\displaystyle\mathcal{R}(\hat{\bm{\beta}}):=\operatorname{\mathbb{E}}_{(\bm{x} ,y)\sim{\cal{D}}_{t}({\bm{\beta}_{\star}})}[(y-\bm{x}^{\top}\hat{\bm{\beta}})^ {2}]-{\sigma}_{t}^{2}=\|{\bm{\Sigma}_{t}^{1/2}(\hat{\bm{\beta}}-{\bm{\beta}_{ \star}})}\|_{2}^{2}.$ (5) $$

Throughout the paper, we compare the surrogate-to-target model with two different reference models.

Reference 1: Standard target model. We study the generalization performance of ${\bm{\beta}}^{s2t}$ with respect to the standard target model, which has access to the ground-truth parameter through labeling. Specifically, consider the dataset $\{(\bm{x}_{i},y_{i})_{i=1}^{n}\}$ drawn i.i.d. from ${\cal{D}}_{t}({\bm{\beta}_{\star}})$; then, the estimation is

$$ $\displaystyle{\bm{\beta}}^{t}:=\text{Est}({\bm{X}},\bm{y}),$ (6) $$

where ${\bm{X}}=[\bm{x}_{1}^{\top},\ldots,\bm{x}_{n}^{\top}]^{\top}\in\mathbb{R}^{n
\times p}$ and $\bm{y}=[y_{1},\ldots,y_{n}]^{\top}\in\mathbb{R}^{n}$.
We compare the excess risks of the surrogate-to-target model $\mathcal{R}({\bm{\beta}}^{s2t})$ with that of the standard target model $\mathcal{R}({\bm{\beta}}^{t})$.

Reference 2: Covariance shift model (Mallinar et al., ; Patil et al., ). Given ${\bm{\beta}_{\star}}\in\mathbb{R}^{p}$, let $\{(\bm{x}_{i},y_{i})_{i=1}^{n}\}$ be a dataset drawn i.i.d. from ${\cal{D}}_{s}^{\text{cs}}$, where $\bm{x}_{i}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{s})$, $z_{i}\sim\mathcal{N}(0,{\sigma}_{t}^{2})$, and $y_{i}=\bm{x}_{t}^{\top}{\bm{\beta}_{\star}}+z_{i}$; then, using the same notation ${\bm{X}}=[\bm{x}_{1}^{\top},\ldots,\bm{x}_{n}^{\top}]^{\top}\in\mathbb{R}^{n
\times p}$ and $\bm{y}=[y_{1},\ldots,y_{n}]^{\top}\in\mathbb{R}^{n}$,
the estimation is $\hat{\bm{\beta}}^{\text{cs}}:=\text{Est}({\bm{X}},\bm{y})$. The test risk of $\hat{\bm{\beta}}^{\text{cs}}$ is calculated under covariance shift. Let $(\bm{x},y)\sim{\cal{D}}_{t}$ be a distribution such that $\bm{x}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{t})$, $z\sim\mathcal{N}(0,{\sigma}_{t}^{2})$, and $y=\bm{x}^{\top}{\bm{\beta}_{\star}}+z$. Then, the excess transfer risk is

$$ $\displaystyle\mathcal{R}(\hat{\bm{\beta}}^{\text{cs}}):=\operatorname{\mathbb{ E}}_{(\bm{x},y)\sim{\cal{D}}_{t}}[(y-\bm{x}^{\top}\hat{\bm{\beta}})^{2}]-{ \sigma}_{t}^{2}.$ $$

We discuss the equivalence between the surrogate-to-target model and the covariance shift model in Section 3.

## 3 Analysis for model shift

We start by examining the behavior of the surrogate-to-target model when there is a model shift ${\bm{\beta}}^{s}\neq{\bm{\beta}_{\star}}$. First, we provide a non-asymptotic bound on the risk conditioned on ${\bm{\beta}}^{s}$, and then we optimize this quantity with respect to ${\bm{\beta}}^{s}$ (finding also a closed-form expression of the corresponding optimal value of ${\bm{\beta}}^{s}$). In addition, we build a connection between our surrogate-to-target model and knowledge distillation (Hinton et al., ), as well as weak-to-strong generalization (Burns et al., ).

We note that when ${\bm{\beta}}^{s}={\bm{\beta}_{\star}}$, the calculations in this section simplify to the non-asymptotic risk characterization of the minimum $\ell_{2}$-norm interpolator, recently studied by Hastie et al. (); Cheng & Montanari (); Han & Xu (). When ${\bm{\beta}}^{s}\neq{\bm{\beta}_{\star}}$, the problem is instead equivalent to covariance shift in transfer learning, as formalized by the following observation whose proof is deferred to Appendix A.

###### Observation 1 .

The model shift in the surrogate-to-target model is equivalent to the covariance shift model (Mallinar et al., ). Formally, given any ${\bm{\beta}_{\star}}\in\mathbb{R}^{p}$ and any jointly diagonalizable covariance matrices $\bm{\Sigma}_{s},\bm{\Sigma}_{t}\in\mathbb{R}^{p\times p}$, there exists a unique ${\bm{\beta}}^{s}\in\mathbb{R}^{p}$ such that the risk of the surrogate-to-target problem $\mathcal{R}({\bm{\beta}}^{s2t})$ with $({\bm{\beta}_{\star}},{\bm{\beta}}^{s},\bm{\Sigma}_{t})$ is equivalent to the risk of the covariance shift model $\mathcal{R}^{\text{cs}}(\hat{\bm{\beta}})$ with $({\bm{\beta}_{\star}},\bm{\Sigma}_{s},\bm{\Sigma}_{t})$.

The joint diagonalizability of $\bm{\Sigma}_{s}$ and $\bm{\Sigma}_{t}$ is also required by Mallinar et al. () (see their Assumption 2.1), where upper and lower bounds for the bias and variance are provided by adapting results from (Bartlett et al., ; Tsigler & Bartlett, ).
In contrast, our approach utilizes the non-asymptotic characterization of the $\ell_{2}$-norm interpolator by Han & Xu (). This allows us to directly characterize the non-asymptotic risk (instead of giving upper and lower bounds, as in (Mallinar et al., )) and, thus, to obtain the optimal surrogate parameter ${\bm{\beta}}^{s}$.
We begin by defining the asymptotic risk of the surrogate-to-target model, given the surrogate parameter ${\bm{\beta}}^{s}$, in the proportional regime where $p,n\xrightarrow{}\infty$ and the ratio $\kappa_{t}=p/n>1$ is kept fixed.

###### Definition 1 .

Let $\kappa_{t}=p/n>1$ and $\tau_{t}\in\mathbb{R}$ be the unique solution of the following equation

$$ $\displaystyle\kappa_{t}^{-1}=\frac{1}{p}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau _{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}\right).$ (7) $$

Let ${\bm{\theta}}_{1}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}$ and ${\bm{\theta}}_{2}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2
}\frac{\bm{g}_{t}}{\sqrt{p}}$ where $\bm{g}_{t}\sim\mathcal{N}(\bm{0},{\bm{I}}_{p})$. Then, the asymptotic risk estimate is defined as

$$ $\begin{split}\bar{\mathcal{R}}^{s2t}_{\kappa_{t},{\sigma}_{t}}(\bm{\Sigma}_{t} ,{\bm{\beta}_{\star}},{\bm{\beta}}^{s})&:=({\bm{\beta}}^{s}-{\bm{\beta}_{\star }})^{\top}{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\beta} }^{s}-{\bm{\beta}_{\star}})+\gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{ \mathbb{E}}_{\bm{g}_{t}}[{\bm{\theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_ {2}]\\ &+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t }({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{\star}}^{\top} ({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\beta }}^{s}-{\bm{\beta}_{\star}}),\end{split}$ (8) $$

where the function $\gamma_{t}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ as

$$ $\gamma_{t}^{2}({\bm{\beta}}^{s})=\kappa_{t}\left({\sigma}_{t}^{2}+\bar{ \mathcal{R}}^{s2t}_{\kappa_{t},{\sigma}_{t}}(\bm{\Sigma}_{t},{\bm{\beta}}^{s}, {\bm{\beta}}^{s})\right).$ (9) $$

We now state our non-asymptotic characterization of the risk. Its proof is deferred to Appendix A.

###### Theorem 1 .

Suppose that, for some constant $M_{t}>1$, we have $1/M_{t}\leq\kappa_{t},{\sigma}_{t}^{2}\leq M_{t}$ and $\left\|{\bm{\Sigma}_{t}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{t}^{-1}}
\right\|_{\text{op}}\leq M_{t}$. Recall from (5) that $\mathcal{R}({\bm{\beta}}^{s2t})$ represents the risk of the surrogate-to-target model given ${\bm{\beta}}^{s}$. Then, there exists a constant $C=C(M_{t})$ such that, for any $\varepsilon\in(0,1/2]$, the following holds with $R+1<M_{t}$:

$$ $\displaystyle\sup_{{\bm{\beta}_{\star}},{\bm{\beta}}^{s}\in{\bm{B}}_{p}(R)} \operatorname{\mathbb{P}}(\left|\mathcal{R}({\bm{\beta}}^{s2t})-\bar{\mathcal{ R}}^{s2t}_{\kappa_{t},{\sigma}_{t}}(\bm{\Sigma}_{t},{\bm{\beta}_{\star}},{\bm{ \beta}}^{s})\right|\geq\varepsilon)\leq Cpe^{-p\varepsilon^{4}/C}.$ (10) $$

The precise non-asymptotic characterization in (10) under model shift enables us to derive the optimum surrogate parameter ${\bm{\beta}}^{s*}$, which is visualized as the green curve in Figure 1.

###### Proposition 1 .

Let $\Omega=\frac{\texttt{tr}\left(\bm{\Sigma}_{t}^{2}(\bm{\Sigma}_{t}+\tau_{t}{\bm
{I}})^{-2}\right)}{n}$. The optimal surrogate ${\bm{\beta}}^{s}$ minimizing the asymptotic risk in (8) is

$$ $\displaystyle{\bm{\beta}}^{s*}=\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1} \bm{\Sigma}_{t}+\frac{\Omega\tau_{t}^{2}}{1-\Omega}\bm{\Sigma}_{t}^{-1}(\bm{ \Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\right)^{-1}{\bm{\beta}_{\star}}.$ $$

Note that, under the setting of Theorem 1, Proposition 1 can be extended to the non-asymptotic risk by applying (10). The corollary below then offers a direct interpretation of the optimal surrogate parameter ${\bm{\beta}}^{s*}$. The proofs of both Corollary 1 and Proposition 1 are in Appendix A.

###### Corollary 1 .

Without loss of generality, suppose that $\bm{\Sigma}_{t}$ is diagonal.(^††\dagger†^††\dagger†$\dagger$If not, there exists an orthogonal matrix ${\bm{U}}\in\mathbb{R}^{p\times p}$ s.t. ${\bm{U}}\bm{\Sigma}_{t}{\bm{U}}^{\top}$ is diagonal. Then, we can consider the covariance matrix as ${\bm{U}}\bm{\Sigma}_{t}{\bm{U}}^{\top}$ and the ground truth parameter as ${\bm{U}}{\bm{\beta}_{\star}}$, which behaves the same as the original parameters, see Observation 2.) Let $(\lambda_{i})_{i=1}^{p}$ be the eigenvalues of $\bm{\Sigma}_{t}$ in non-increasing order and let $\zeta_{i}=\frac{\tau_{t}}{\lambda_{i}+\tau_{t}}$ for $i\in[p]$. Then, the following results hold:

- 1.
$\beta^{s*}_{i}=(\beta_{*})_{i}\left((1-\zeta_{i})+\zeta_{i}\frac{\Omega}{1-
\Omega}\frac{\zeta_{i}}{1-\zeta_{i}}\right)^{-1}$ for every $i\in[p]$.
- 2.
$|\beta^{s*}_{i}|>|(\beta_{*})_{i}|$ if and only if $1-\zeta_{i}>\Omega=\frac{\sum_{j=1}^{p}(1-\zeta_{j})^{2}}{\sum_{j=1}^{p}(1-
\zeta_{j})}$ for every $i\in[p]$.
- 3.
${\bm{\beta}}^{s*}={\bm{\beta}_{\star}}$ if and only if the covariance matrix $\bm{\Sigma}_{t}=c{\bm{I}}$ for some $c\in\mathbb{R}$.

The first part shows that the optimal surrogate parameter is fully characterized by $(\zeta_{i})_{i=1}^{p}$, which only depends on the covariance spectrum (via $\lambda_{i}$) and the sample size $n$ (via $\tau_{t}$). As the eigenvalues $(\lambda_{i})_{i=1}^{p}$ are ordered, the
$\zeta_{i}$’s are ordered as well, and the second part of the corollary identifies a threshold behavior: before the transition point $1-\zeta_{i}=\Omega$, the entries of the surrogate are amplified w.r.t. the ground-truth parameter ${\bm{\beta}_{\star}}$, while they experience shrinkage after the transition. The threshold corresponds to the ratio of the sample second moment to the sample first moment of the random variable whose realization is given by $(1-\zeta_{i})$, and it arises from the optimization of the trade-off between the bias and variance terms in (8). Finally, the third part of the corollary shows that, unless the eigenvalues of the covariance matrix are constant, there is potential for improvement by tuning the surrogate parameter.

###### Observation 1 .

###### Definition 1 .

###### Theorem 1 .

###### Proposition 1 .

###### Corollary 1 .

### 3.1 Weak-to-strong generalization

To connect with knowledge distillation (Hinton et al., ) and weak-to-strong generalization (Burns et al., ), we allow the surrogate model to use fewer features, $p_{s}<p$, by introducing a mask operation $\mathcal{M}(\bm{x})$, where $\mathcal{M}(\bm{x})\in\mathbb{R}^{p_{s}}$ selects $p_{s}$ features from the full set of $p$ features in $\bm{x}\in\mathbb{R}^{p}$. Alongside this mask, we adjust the distributions for both the surrogate and target models as

$$ $\displaystyle(\mathcal{M}(\tilde{\bm{x}}),\tilde{y})\sim{\cal{D}}_{s}^{p_{s}} \text{ follows }\tilde{y}=\mathcal{M}(\tilde{\bm{x}})^{\top}\mathcal{M}({\bm{ \beta}_{\star}})+{\tilde{z}},\text{ where }{\bm{\beta}_{\star}}\in\mathbb{R}^{ p},\tilde{\bm{x}}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{s}),{\tilde{z}}\sim \mathcal{N}(0,{\sigma}_{s}^{2}),$ $\displaystyle(\bm{x},y^{s})\sim{\cal{D}}_{t}^{p_{s}}({\bm{\beta}}^{s})\text{ follows }y^{s}=\mathcal{M}(\bm{x})^{\top}{\bm{\beta}}^{s}+z,\text{ where }{\bm {\beta}}^{s}\in\mathbb{R}^{p_{s}},\bm{x}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{t} ),z\sim\mathcal{N}(0,{\sigma}_{s}^{2}).$ $$

Then, ${\bm{\beta}}^{s2t}$ is estimated based on the samples from ${\cal{D}}_{t}^{p_{s}}({\bm{\beta}}^{s})$, and the risk $\mathcal{R}({\bm{\beta}}^{s2t})$ is still calculated with respect to the standard target model distribution ${\cal{D}}_{t}({\bm{\beta}_{\star}})$ as defined in (5). As we focus on analyzing the model shift case, we assume that the covariance matrices $\bm{\Sigma}_{s}$ and $\bm{\Sigma}_{t}$ are identical.

In this formulation, the surrogate model is considered weak because it has access to fewer features, while the target model is the strong model. We now address the following question: Can the surrogate-to-target model outperform the standard target model, provided in (6), in the absence of model shift ($\mathcal{M}({\bm{\beta}_{\star}})={\bm{\beta}}^{s}$)?
The absence of model shift corresponds to the case where the surrogate model has infinitely many data. The next proposition provides a sufficient condition to answer the question above in the affirmative, and it derives the optimal selection of features.

###### Proposition 2 .

Consider the target model in (6), assume that $\bm{\Sigma}_{t}$ is diagonal, and recall the definitions of $\zeta_{i}$ and $\Omega$. Then, the following results hold:

- 1.
If the mask operation $\mathcal{M}$ selects all the features that satisfy $1-\zeta_{i}^{2}>\Omega$, then the surrogate-to-target model outperforms the standard target model in the asymptotic risk in (8).
- 2.
Let $\bm{M}$ represent the set of all possible $\mathcal{M}$, where $|\bm{M}|=2^{p}$. The optimal $\mathcal{M}^{*}$ for the asymptotic risk in (8) within $\bm{M}$ is the one that selects all features satisfying $1-\zeta_{i}^{2}>\Omega$.

The proof of Proposition 2 is provided in Appendix A, and the result can be extended to the non-asymptotic risk by applying Theorem 1. Similarly to Corollary 1, the result above identifies a threshold behavior: the entries of the surrogate are masked (i.e., set to 0 0) after the transition point $1-\zeta_{i}^{2}=\Omega$, while they coincide with the ground-truth parameter ${\bm{\beta}_{\star}}$ otherwise. The transition point changes with respect to Corollary 1 and, as $1-\zeta_{i}^{2}>1-\zeta_{i}$, it is shifted to the right:
the optimal mask includes not only features whose magnitude increases, but also features whose magnitude decreases while selecting the optimal surrogate ${\bm{\beta}}^{s*}$.

In Figure 1(a), we illustrate the optimal $\mathcal{M}$ and ${\bm{\beta}}^{s}$, showing that the threshold associated to the optimal $\mathcal{M}$ is larger than the threshold associated to the transition from amplification to shrinkage in the optimal ${\bm{\beta}}^{s}$. In addition, we note that the ratio between the green curve and the blue curve in Figure 1(a) is not monotone with respect to $\lambda_{i}$. In Figure 1(b), we also present a comparison of their associated risks.

In Figure 2(a), we examine the surrogate-to-target model in the context of image classification. Specifically, we fine-tune a pretrained ResNet-50 model (He et al., ) using both ground-truth labels and predictions from a surrogate (weak) model on the CIFAR-10 dataset (Krizhevsky & Hinton, ). The surrogate models are shallow, 3-layer convolutional neural networks with varying parameter sizes. In all cases, surrogate-to-target models consistently outperform surrogate models across different model sizes. However, in this setting, surrogate-to-target models do not outperform the standard target (strong) model. This is in agreement with the weak-to-strong results in Burns et al. (), where the GPT-4 model trained with GPT-2 labels performs comparably to GPT-3.5. This suggests that the feature selection mechanism, as characterized in Proposition 2, is crucial for surpassing the performance of the target model.

###### Proposition 2 .

## 4 Fundamental limits and scaling laws

We now study the fundamental limits of the surrogate-to-target model with the optimal surrogate parameter ${\bm{\beta}}^{s*}$ (see Proposition 1) and the optimal mask operator $\mathcal{M}^{*}$ (see Proposition 2). Our analysis shows that, when eigenvalues ($\lambda_{i}$) and signal coefficients ($\lambda_{i}\beta_{i}^{2}$) follow a power law, the risk of the surrogate-to-target model under the optimal selection of the parameters ${\bm{\beta}}^{s*}$ and $\mathcal{M}^{*}$ scales the same as that of the target model (even though there is a strict improvement in the risk, as per Corollary 1). By Observation 1, this also indicates that the gain obtained by the covariance shift model, as outlined in Mallinar et al. (), does not change the scaling law. We start our analysis with the definition of the omniscient test risk estimate.

Figure: (a) Weak-to-strong on CIFAR-10
Refer to caption: x3.png

###### Definition 2 (Omniscient test risk estimate) .

Fix $p>n\geq 1$. Given a covariance $\bm{\Sigma}={\bm{U}}\operatorname{diag}(\bm{\lambda}){\bm{U}}^{\top}$, ${\bm{\beta}_{\star}}$, and the noise term $\sigma$, set $\bar{\bm{\beta}}={\bm{U}}^{\top}{\bm{\beta}_{\star}}$ and define $\tau\in\mathbb{R}$ as the unique non-negative solution of $n=\sum_{i=1}^{p}\frac{\lambda_{i}}{\lambda_{i}+\tau}.$ Then, the excess test risk estimate is the following:

$$ $\displaystyle\mathcal{R}(\hat{\bm{\beta}})$ $\displaystyle\approx\mathbb{E}_{\hat{\bm{\beta}}\sim D({\bm{\beta}_{\star}})} \left[(y-\bm{x}^{\top}\hat{\bm{\beta}})^{2}\right]-\sigma^{2}=\frac{\sigma^{2} \Omega+\mathcal{B}(\bar{\bm{\beta}})}{1-\Omega},$ (11) where $\displaystyle\zeta_{i}=\frac{\tau}{\lambda_{i}+\tau},\quad\Omega=\frac{1}{n} \sum_{i=1}^{p}(1-\zeta_{i})^{2},\quad\mathcal{B}(\bar{\bm{\beta}})=\sum_{i=1}^ {p}\lambda_{i}\zeta_{i}^{2}\bar{\beta_{i}}^{2}.$ $$

The above test risk estimate yields exact results (and, hence, $\approx$ in (11) becomes $=$) in the proportional limit via the analysis of Section 3. Specifically, suppose that the empirical distributions of $\bar{\bm{\beta}}$ and $\bm{\lambda}$ converge as $p\to\infty$ having fixed the ratio $p/n=\kappa$. Then, the risk obtained in Theorem 1 converges to the omniscient risk estimate given in (11), as proved in Appendix B. We will use this omniscient risk estimate in the limit of $p\to\infty$, as considered in several papers (Cui et al., ; Simon et al., ; Wei et al., ). Yet, our empirical validations in Figure 1 demonstrate that this framework yields consistent results even when applied to scenarios with moderately sized $p$ and $n$.

Throughout the section, we analyze the case where the surrogate parameter ${\bm{\beta}}^{s}$ is given, therefore we need to take into account only the target covariance matrix $\bm{\Sigma}_{t}$. Without loss of generality, we assume that the covariance matrix $\bm{\Sigma}_{t}$ is diagonal by Observation 2. From now on, we will consider the particular case of power-law eigenstructure, that is $\bm{\Sigma}_{i,i}=\lambda_{i}=i^{-\alpha}$. The omniscient risk under this structure still depends on the parameters $\tau_{t}$ and $\Omega$, and the next proposition analyzes them asymptotically. Its proof is in Appendix B.

###### Proposition 3 (Asymptotic analysis of $\tau_{t}$ and $\Omega$ ) .

Let $\bm{\Sigma}\in\mathbb{R}^{p\times p}$ be diagonal and $\bm{\Sigma}_{i,i}=\lambda_{i}=i^{-\alpha}$ for $1<\alpha$. Recall that, as $p\rightarrow\infty$, $\tau_{t}$ and $\Omega$ are given by the equations

$$ $\displaystyle\sum_{i=1}^{\infty}\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}}=n, \qquad n\Omega=\sum_{i=1}^{\infty}\left(\frac{i^{-\alpha}}{i^{-\alpha}+\tau_{t }}\right)^{2}.$ $$

Then, the following results hold

$$ $\begin{split}\tau_{t}&=cn^{-\alpha}\left(1+O(n^{-1})\right),\qquad\mbox{for }c =\left(\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\right)^{\alpha},\\ \Omega&=\dfrac{\alpha-1}{\alpha}-O(n^{-1}).\end{split}$ (12) $$

Recall from Corollary 1 and Proposition 2 that the cut-off indices for the optimal surrogate parameter ${\bm{\beta}}^{s*}$ and the optimal mask operation $\mathcal{M}^{*}$ are respectively $\zeta_{i}<1-\Omega$ and $\zeta_{i}^{2}<1-\Omega$, which depend on $\tau_{t},\Omega$. Armed with the asymptotic expressions in (12), we now identify the cut-off indices as a function of the sample size $n$.

###### Proposition 4 .

Set the constants $C_{1}:=\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\alpha-1)^{1/\alpha}}$ and $C_{2}:=\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\sqrt{\alpha}-1)^{1/\alpha}}$ and assume the power-law eigenstructure $\lambda_{i}=i^{-\alpha}$ for $1<\alpha$. Then, the indices $i$ for which $\zeta_{i}<1-\Omega$ are $i<nC_{1}+O(1)$; while the indices $i$ for which is $\zeta_{i}^{2}<1-\Omega$ are $i<nC_{2}+O(1)$.

The result above is proved in Appendix B and it shows that, as the sample size $n$ decreases, the cut-off indices of both the optimal surrogate parameter ${\bm{\beta}}^{s*}$ and the optimal mask $\mathcal{M}^{*}$ shift to the left linearly in $n$. This also implies that, with less data, optimal surrogate models tend to be more sparse.

Next, we address the question of how the excess test risk scales with respect to the sample size $n$, when the surrogate parameter is the optimal ${\bm{\beta}}^{s*}$. Specifically, Proposition 5 below shows that, under a power-law decay of both the eigenvalues ($\lambda_{i}$) and the signal coefficients ($\lambda_{i}\beta_{i}^{2}$), the excess test risk of the optimal surrogate-to-target model scales the same as the standard target model.

Before stating the result, we make a comment on the noise assumption needed
to ensure that the scaling law of the excess test risk remains unaffected by the introduction of noise, which allows us to analyze the model’s inherent error. Specifically, we choose the variance of the noise term ${\sigma}_{t}^{2}$ to be at most of the order of the scaling law of the excess test risks when ${\sigma}_{t}^{2}=0$. This corresponds to ${\sigma}_{t}^{2}=O(n^{-\gamma})$, where $\gamma$ is the exponent of the scaling law in the noiseless setting. Conversely, a fixed noise variance ${\sigma}_{t}^{2}=\Theta(1)$ that does not decay with $n$ would cause the noise to overshadow the uncaptured part of the signal, which scales down with $n$. In this unintended scenario, the noise would dominate our observations.

###### Proposition 5 (Scaling law) .

Assume that both eigenvalues $\lambda_{i}$ and signal coefficients $\lambda_{i}\beta_{i}^{2}$ follow a power-law decay, i.e., $\lambda_{i}\beta_{i}^{2}=i^{-\beta}$ and $\lambda_{i}=i^{-\alpha}$ for $\alpha$, $\beta>1$. Let the optimal surrogate parameter ${\bm{\beta}}^{s*}$ be given by Proposition 1 and define the minimum surrogate-to-target risk attained by ${\bm{\beta}}^{s*}$ as $\mathcal{R}^{*}({\bm{\beta}}^{s2t})=\underset{}{\mathrm{min}}\ \mathcal{R}({
\bm{\beta}}^{s2t}).$ Then, in the limit of $p\to\infty$, the excess test risk of the surrogate-to-target model with an optimal surrogate parameter scales the same as that of the standard target model. Specifically, we have

$$ $\displaystyle\mathcal{R}^{*}({\bm{\beta}}^{s2t})$ $\displaystyle=\Theta(n^{-(\beta-1)})=\mathcal{R}({\bm{\beta}}^{t}),$ $\displaystyle\text{if }\beta<2\alpha+1,$ $\displaystyle\mathcal{R}^{*}({\bm{\beta}}^{s2t})$ $\displaystyle=\Theta(n^{-2\alpha})=\mathcal{R}({\bm{\beta}}^{t}),$ $\displaystyle\text{if }\beta>2\alpha+1.$ $$

Since $\mathcal{R}^{*}({\bm{\beta}}^{s2t})$ is a lower bound on $\mathcal{R}({\bm{\beta}}^{s2t})$, we have that the scaling law of the excess test risk of the surrogate-to-target model cannot be improved beyond that of the standard target model, even with the freedom to choose ${\bm{\beta}}^{s}$. This also indicates that the optimal selection of the mask does not improve the scaling law, see Proposition 6 in Appendix B for details.

The proof of Proposition 5 is deferred to Appendix B. Here, we note that we utilize the expression $\frac{\Omega}{1-\Omega}\left(\sum_{i=1}^{p}\lambda_{i}(\beta_{i}^{s})^{2}\zeta
_{i}^{2}\right)$ as a lower bound for the risk, while proving the scaling law for the optimal surrogate-to-target model. Although this expression alone is insufficient to determine an asymptotic lower bound for an arbitrary ${\bm{\beta}}^{s}$, it becomes particularly useful when considering the optimal surrogate parameter ${\bm{\beta}}^{s*}$ provided in Proposition 1. We then leverage the fact that the optimal surrogate parameter yields the minimum test risk to characterize its asymptotic behavior.

Finally, we provide in Appendix B also a non-asymptotic analysis of $\tau_{t}$ and $\Omega$ (see Propositions 7 and 8, respectively), which complements the asymptotic one in Proposition 3 above. This allows to characterize a region with finite $n$ and $p$ where the surrogate-to-target model strictly outperforms the standard target model, see Proposition 9.

###### Definition 2 (Omniscient test risk estimate) .

###### Proposition 3 (Asymptotic analysis of τ t subscript 𝜏 𝑡 \tau_{t} italic_τ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT and Ω Ω \Omega roman_Ω ) .

###### Proposition 4 .

###### Proposition 5 (Scaling law) .

## 5 Risk characterization for the two-stage model

Until now, we have examined the behavior of the surrogate-to-target model when ${\bm{\beta}}^{s}$ is given. In this section, we characterize the non-asymptotic risk of the surrogate-to-target model when ${\bm{\beta}}^{s}$ is the solution of the surrogate problem (3) where $\kappa_{s}=p/m>1$. Our analysis includes two cases: *(i)* the target model has infinitely many data ($n=\infty$), and *(ii)* the target model is overparametrized, i.e., $\kappa_{t}=p/n>1$.

When the target model has infinitely many data, the estimate of the surrogate-to-target model ${\bm{\beta}}^{s2t}$ is equal to the estimate of the surrogate model ${\bm{\beta}}^{s}$. This means that the correct ground-truth parameter ${\bm{\beta}_{\star}}$ is estimated under a distribution ${\cal{D}}_{s}$ and tested under another distribution ${\cal{D}}_{t}({\bm{\beta}_{\star}})$, which is equivalent to the covariance shift model by definition. By Observation 1, the model shift in the surrogate-to-target model is equivalent to the covariance shift model and, hence, the analysis in Sections 3 and 4 is valid for this case.

Finally, we consider the case where the target model is overparametrized and begin with the following asymptotic risk definition.

###### Definition 3 .

Recall the definition of $\tau_{t}$ and $\gamma_{t}$ in Theorem 1. Let $\kappa_{s}=p/m>1$ and define $\tau_{s}\in\mathbb{R}$ similarly to $\tau_{t}$.
We define the random variable $X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}$ based on $\bm{g}_{s}\sim\mathcal{N}(0,{\bm{I}})$ and the function $\gamma_{s}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ as follows:

$$ $\begin{split}X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}(\bm{\Sigma}_{s},{\bm{\beta}_{ \star}},\bm{g}_{s})&:=(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s} \left[{\bm{\beta}_{\star}}+\frac{\bm{\Sigma}_{s}^{-1/2}\gamma_{s}({\bm{\beta}_ {\star}})\bm{g}_{s}}{\sqrt{p}}\right]\\ \gamma_{s}^{2}({\bm{\beta}_{\star}})&:=\kappa_{s}\left({\sigma}_{s}^{2}+ \operatorname{\mathbb{E}}_{\bm{g}_{s}}[\|{\bm{\Sigma}_{s}^{1/2}(X_{\kappa_{s}, {\sigma}_{s}^{2}}^{s}(\bm{\Sigma}_{s},{\bm{\beta}_{\star}},\bm{g}_{s})-{\bm{ \beta}_{\star}})}\|_{2}^{2}]\right).\end{split}$ $$

Let $\dot{\kappa}=(\kappa_{s},\kappa_{t})$, $\dot{\bm{\Sigma}}=(\bm{\Sigma}_{s},\bm{\Sigma}_{t})$, and $\dot{\sigma}=({\sigma}_{s}^{2},{\sigma}_{t}^{2})$.
Then, we define the asymptotic risk estimate as

$$ $\displaystyle\bar{\mathcal{R}}_{\dot{\kappa},\dot{\sigma}}(\dot{\bm{\Sigma}},{ \bm{\beta}_{\star}})$ $\displaystyle=\|{\bm{\Sigma}_{t}^{1/2}\left({\bm{I}}-(\bm{\Sigma}_{t}+\tau_{t} {\bm{I}})^{-1}\bm{\Sigma}_{t}(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma }_{s}\right){\bm{\beta}_{\star}}}\|_{2}^{2}+\frac{\operatorname{\mathbb{E}}_{{ \bm{\beta}}^{s}\sim X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}}[\gamma_{t}^{2}({\bm{ \beta}}^{s})]}{p}\texttt{tr}\left(\bm{\Sigma}_{t}^{2}(\bm{\Sigma}_{t}+\tau_{t} {\bm{I}})^{-2}\right)$ $\displaystyle+\frac{\gamma_{s}^{2}({\bm{\beta}_{\star}})}{p}\texttt{tr}\left( \bm{\Sigma}_{s}^{1/2}(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{t}( \bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t} {\bm{I}})^{-1}\bm{\Sigma}_{t}(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma }_{s}^{1/2}\right).$ $$

The non-asymptotic characterization of the risk is stated below and proved in Appendix C, which also contains a closed-form expression for $\operatorname{\mathbb{E}}_{{\bm{\beta}}^{s}\sim X_{\kappa_{s},{\sigma}_{s}^{2}
}^{s}}[\gamma_{t}^{2}({\bm{\beta}}^{s})]$ (see Lemma 1).

###### Theorem 2 .

Suppose that, for some constant $M_{t}>1$, we have $1/M_{t}\leq\kappa_{s},{\sigma}_{s}^{2},\kappa_{t},{\sigma}_{t}^{2}\leq M_{t}$ and $\left\|{\bm{\Sigma}_{s}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{s}^{-1}}
\right\|_{\text{op}},$ $\left\|{\bm{\Sigma}_{t}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{s}^{-1}}
\right\|_{\text{op}}\leq M_{t}$. Consider the surrogate-to-target model defined in Section 2, and let $\mathcal{R}({\bm{\beta}}^{s2t})$ represent its risk when ${\bm{\beta}_{\star}}$ is given. Recall the definition of $\dot{\bm{\Sigma}},\dot{\kappa},\dot{\sigma}$ and $\bar{\mathcal{R}}_{\dot{\kappa},\dot{\sigma}}$ in Definition 3. Then, there exists a constant $C=C(M_{t})$ such that for any $\varepsilon\in(0,1/2]$, the following holds when $R+1<M_{t}$:

$$ $\displaystyle\sup_{{\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)}\operatorname{ \mathbb{P}}(\left|\mathcal{R}({\bm{\beta}}^{s2t})-\bar{\mathcal{R}}_{\dot{ \kappa},\dot{\sigma}}(\dot{\bm{\Sigma}},{\bm{\beta}_{\star}})\right|\geq \varepsilon)\leq Cpe^{-p\varepsilon^{4}/C}.$ $$

We implement the surrogate-to-target model with a synthetic dataset and demonstrate that our risk characterization agrees well with the experimental two-stage linear regression in Figure 2(b).

###### Definition 3 .

###### Theorem 2 .

## 6 Concluding remarks

We have provided a sharp characterization of knowledge distillation for high-dimensional linear regression when labels are generated by a surrogate model and, additionally, characterized the risk of the two-stage process where the surrogate model is the outcome of an initial ERM. These results shed light on the form of the optimal surrogate model, reveal an amplify-to-shrink phase transition as a function of the eigenspectrum, and draw connections to weak-to-strong generalization. Specifically, we have shown that the labels coming from the optimal surrogate model strictly allow for improving the performance of the target model, unless the covariance is a multiple of the identity. However, even though there is a strict improvement in the risk, the scaling behavior of the two-stage process with labels coming from the optimal surrogate model remains unchanged compared to the standard target model that utilizes ground-truth labels.

We outline two interesting directions for future research. The first is to extend the two-stage process to multiple stages, establishing whether this further improves the risk. The second is to apply the two-stage learning to data pruning, using the surrogate model to decide whether to keep or discard each $(\bm{x},y)$ pair during the training of the target model.

## Acknowledgements

M.E.I., E.O.T., S.O. are supported by the NSF grants CCF-2046816, CCF-2403075, the Office of Naval Research grant N000142412289, an OpenAI Agentic AI Systems grant, and gifts by Open Philanthropy and Google Research. M. M. is funded by the European Union (ERC, INF^2, project number 101161364). Views and opinions expressed are however
those of the author(s) only and do not necessarily reflect those of the European Union or the
European Research Council Executive Agency. Neither the European Union nor the granting
authority can be held responsible for them.

## Appendix A Proofs for Section 3

See 1

###### Proof.

By Observation 2, we assume that $\bm{\Sigma}_{t}$ and $\bm{\Sigma}_{s}$ are diagonal matrices. As $\bm{\Sigma}_{t}$ and $\bm{\Sigma}_{s}$ are jointly diagonalizable, there exists a unique diagonal matrix ${\bm{A}}\in\mathbb{R}^{p\times p}$ such that

$$ $\displaystyle\bm{\Sigma}_{s}={\bm{A}}^{\top}\bm{\Sigma}_{t}{\bm{A}}.$ $$

Then, consider the model shift discussed in Section 3. Take the case where ${\bm{\beta}}^{s}={\bm{A}}{\bm{\beta}_{\star}}$ and labels are generated as $y=\bm{x}^{\top}{\bm{\beta}}^{s}+z$, where $\bm{x}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{t})$ and $z\sim\mathcal{N}(0,{\sigma}_{t}^{2})$. This is equivalent to the case where $y=(\bm{x}^{\top}{\bm{A}}){\bm{\beta}_{\star}}+z=\bar{\bm{x}}^{\top}{\bm{\beta}
_{\star}}+z$ such that $\bm{x}\sim\mathcal{N}(\bm{0},\bm{\Sigma}_{s})$ and $z\sim\mathcal{N}(0,{\sigma}_{t}^{2})$. Note that *(i)* the transformed inputs
and the labels are identical in both scenarios, and *(ii)* the estimators are computed in the same way. Thus, it follows that the risks $\mathcal{R}({\bm{\beta}}^{s2t})$ and $\mathcal{R}^{\text{cs}}(\hat{\bm{\beta}})$ are equivalent. The other way follows from an almost identical argument.
∎

###### Observation 2 .

For any covariance matrix $\bm{\Sigma}\in\mathbb{R}^{p\times p}$, there exists an orthonormal matrix ${\bm{U}}\in\mathbb{R}^{p\times p}$ such that the transformation of $\bm{x}\to{\bm{U}}^{\top}\bm{x}$ and ${\bm{\beta}}\to{\bm{U}}^{\top}{\bm{\beta}}$ does not affect the labels $\bm{y}$ but ensures that the covariance matrix is diagonal.

###### Proof.

Since the covariance matrix $\bm{\Sigma}$ is PSD, its unit-norm eigenvectors are orthogonal. Consider the matrix ${\bm{U}}$ whose columns are the eigenvectors of $\bm{\Sigma}$. Then, $\bm{\Sigma}$ can be expressed as $\bm{\Sigma}={\bm{U}}{\bm{\Lambda}}{\bm{U}}^{\top}$, where ${\bm{\Lambda}}$ is the diagonal matrix containing the eigenvalues of $\bm{\Sigma}$. Consider now the transformation

$$ $\displaystyle{\bm{z}}={\bm{U}}^{\top}\bm{x}\implies\mathbb{E}\left[{\bm{z}}{ \bm{z}}^{\top}\right]=\mathbb{E}\left[{\bm{U}}^{\top}\bm{x}\bm{x}^{\top}{\bm{U }}\right]={\bm{U}}^{\top}\mathbb{E}\left[\bm{x}\bm{x}^{\top}\right]{\bm{U}}={ \bm{U}}^{\top}{\bm{U}}{\bm{\Lambda}}{\bm{U}}^{\top}{\bm{U}}={\bm{\Lambda}}.$ $$

In this way, the covariance matrix is diagonalized. Thus, the transformation $(\bm{x},{\bm{\beta}_{\star}})\to({\bm{U}}^{\top}\bm{x},{\bm{U}}^{\top}{\bm{
\beta}_{\star}})$ works as intended since the labels are preserved. ∎

See 1

See 1

###### Proof.

Even though the claim readily follows from Theorem 2, we give a proof for the sake of completeness.

Define a function $f_{1}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ as $f_{1}(\bm{x})=\|{\bm{\Sigma}_{t}^{1/2}(\bm{x}-{\bm{\beta}_{\star}})}\|_{2}^{2}$. The gradient of this function is

$$ $\displaystyle\|{\nabla f_{1}(\bm{x})}\|_{2}=\|{2\bm{\Sigma}_{t}(\bm{x}-{\bm{ \beta}_{\star}})}\|_{2}\leq 2\left\|{\bm{\Sigma}_{t}}\right\|_{\text{op}}\|{ \bm{x}-{\bm{\beta}_{\star}}}\|_{2}.$ $$

Using Corollary 2, there exists an event $E$ with $\operatorname{\mathbb{P}}(E^{c})\leq C_{t}e^{-p/C_{t}}$ where $C_{t}=C_{t}(M_{t},\frac{M_{t}-R}{2})$ (with the definition of $M_{t}$ in Corollary 2), such that $f_{1}({\bm{\beta}}^{s2t})$ is $2M_{t}^{2}$-Lipschitz if ${\bm{\beta}_{\star}},{\bm{\beta}}^{s}\in{\bm{B}}_{p}(R)$. Applying Theorem 3 on the target model, there exists a constant $\bar{C}_{s}=\bar{C}_{s}(M_{t})$ such that for any $\varepsilon\in(0,1/2]$, we obtain

$$ $\displaystyle\sup_{{\bm{\beta}}^{s}\in{\bm{B}}(\frac{M_{t}+R}{2})} \operatorname{\mathbb{P}}\left(\left|f({\bm{\beta}}^{s2t})-\operatorname{ \mathbb{E}}_{\bm{g}_{t}}[f(X_{\kappa_{t},{\sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t} ,{\bm{\beta}}^{s},\bm{g}_{t}))]\right|\geq\varepsilon\right)\leq Cp\mathrm{e}^ {-p\varepsilon^{4}/C},$ (13) $$

where $f({\bm{\beta}}^{s2t})=\mathcal{R}({\bm{\beta}}^{s2t})$ and

$$ $\displaystyle X_{\kappa_{t},{\sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t},{\bm{\beta}} ^{s},\bm{g}_{t})=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}\left[{ \bm{\beta}}^{s}+\frac{\bm{\Sigma}_{t}^{-1/2}\gamma_{t}({\bm{\beta}}^{s})\bm{g} _{t}}{\sqrt{p}}\right].$ $$

Furthermore,

$$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[f(X_{\kappa_{t},{ \sigma}_{t}^{2}}^{s}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},\bm{g}_{t}))\right]$ $\displaystyle=\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[\|{\bm{\Sigma}_{t}^{ 1/2}\left({\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})-({\bm{I}}-{ \bm{\theta}}_{1}){\bm{\beta}_{\star}}+{\bm{\theta}}_{2}\gamma_{t}({\bm{\beta}} ^{s})\right)}\|_{2}^{2}\right]$ $\displaystyle=({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta}}_{1}^ {\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})+ \gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{ \theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]$ $\displaystyle+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top} \bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{ \star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_ {1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}}),$ (14) $$

where ${\bm{\theta}}_{1}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}$ and ${\bm{\theta}}_{2}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2
}\frac{\bm{g}_{t}}{\sqrt{p}}$. This completes the proof.
∎

See 1

###### Proof.

We have that

$$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[f(X_{\kappa_{t},{ \sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},\bm{g}_{t}))\right]$ $\displaystyle=\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[\|{\bm{\Sigma}_{t}^{ 1/2}\left({\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})-({\bm{I}}-{ \bm{\theta}}_{1}){\bm{\beta}_{\star}}+{\bm{\theta}}_{2}\gamma_{t}({\bm{\beta}} ^{s})\right)}\|_{2}^{2}\right]$ $\displaystyle=({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta}}_{1}^ {\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})+ \gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{ \theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]$ $\displaystyle+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top} \bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{ \star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_ {1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}}),$ $$

where ${\bm{\theta}}_{1}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}$, ${\bm{\theta}}_{2}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2
}\frac{\bm{g}_{t}}{\sqrt{p}}$. Recall from (9) that

$$ $\displaystyle\gamma_{t}^{2}({\bm{\beta}}^{s})$ $\displaystyle=\kappa_{t}\left({\sigma}_{t}^{2}+\bar{\mathcal{R}}^{s2t}_{\kappa _{t},{\sigma}_{t}}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},{\bm{\beta}}^{s})\right)$ $\displaystyle=\kappa_{t}\left({\sigma}_{t}^{2}+\gamma_{t}^{2}({\bm{\beta}}^{s} )\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{\theta}}_{2}^{\top}\bm{\Sigma}_{t }{\bm{\theta}}_{2}]+({\bm{\beta}}^{s})^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{ \top}\bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}}^{s}\right).$ $$

This implies that

$$ $\displaystyle\gamma_{t}^{2}({\bm{\beta}}^{s})$ $\displaystyle=\kappa_{t}\frac{{\sigma}_{t}^{2}+({\bm{\beta}}^{s})^{\top}({\bm{ I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{ \beta}}^{s}}{1-\kappa_{t}\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{\theta}}_ {2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\frac{{\sigma}_{t}^{2}+\tau_{t}^ {2}\|{\bm{\Sigma}_{t}^{1/2}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}{\bm{\beta}} ^{s}}\|_{2}^{2}}{1-\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I }})^{-2}\bm{\Sigma}_{t}^{2}\right)},$ $$

where $(a)$ follows from the fact that ${\bm{I}}-{\bm{\theta}}_{1}={\bm{I}}-(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm
{\Sigma}_{t}=\tau_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}$ and $\kappa_{t}\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{\theta}}_{2}^{\top}\bm{
\Sigma}_{t}{\bm{\theta}}_{2}]=\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+
\tau_{t}{\bm{I}})^{-2}\bm{\Sigma}_{t}^{2}\right)$.

In order to optimize this with respect to ${\bm{\beta}}^{s}$, let’s take the derivative:

$$ $\displaystyle\frac{\partial}{\partial{\bm{\beta}}^{s}}$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[f(X_{\kappa_{t},{ \sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},\bm{g}_{t}))\right]$ $\displaystyle=2{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{ \beta}}^{s}-{\bm{\beta}_{\star}})-2{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}({ \bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}+2\frac{\kappa_{t}\tau_{t}^{2}}{ 1-\Omega}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}{\bm{\beta}}^{s }\frac{\texttt{tr}\left(\bm{\Sigma}_{t}^{2}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^ {-2}\right)}{p}$ $\displaystyle=2{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{ \beta}}^{s}-2\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{\beta}_{\star}}+2\frac{ \kappa_{t}\tau_{t}^{2}}{1-\Omega}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{ I}})^{-2}{\bm{\beta}}^{s}\frac{\texttt{tr}\left(\bm{\Sigma}_{t}^{2}(\bm{\Sigma }_{t}+\tau_{t}{\bm{I}})^{-2}\right)}{p}$ $\displaystyle=2{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{ \beta}}^{s}-2\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{\beta}_{\star}}+2\frac{ \kappa_{t}\tau_{t}^{2}}{1-\Omega}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{ I}})^{-2}{\bm{\beta}}^{s}\frac{n\Omega}{p}$ $\displaystyle\implies{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}{ \bm{\beta}}^{s*}-\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{\beta}_{\star}}+\frac{ \Omega\tau_{t}^{2}}{1-\Omega}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}}) ^{-2}{\bm{\beta}}^{s*}=0$ $\displaystyle\implies({\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1} +\frac{\Omega\tau_{t}^{2}}{1-\Omega}{\bm{\theta}}_{1}(\bm{\Sigma}_{t}+\tau_{t} {\bm{I}})^{-1}){\bm{\beta}}^{s*}=\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{\beta}_{ \star}}$ $$

Hence, the claimed result follows.
∎

See 1

###### Proof.

When the definition of $\zeta_{i}$ and $\Omega$ is plugged in Proposition 1, the first claim is obtained. Using the diagonalization assumption on $\bm{\Sigma}_{t}$, let’s analyze only the $i$-th component of the optimal surrogate given in Proposition 1:

$$ $\displaystyle\beta^{s*}_{i}=\frac{1}{\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}}+ \frac{\Omega}{1-\Omega}\frac{\tau_{t}^{2}}{\lambda_{i}(\lambda_{i}+\tau_{t})}} (\beta_{*})_{i}$ $\displaystyle\iff\beta^{s*}_{i}=\frac{\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}} }{\left(\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}}\right)^{2}+\frac{\Omega}{1- \Omega}\left(\frac{\tau_{t}}{\lambda_{i}+\tau_{t}}\right)^{2}}(\beta_{*})_{i}$ $\displaystyle\iff\beta^{s*}_{i}=(\beta_{*})_{i}\frac{(1-\zeta_{i})}{\left(1- \zeta_{i}\right)^{2}+\frac{\Omega}{1-\Omega}\zeta_{i}^{2}}$ $\displaystyle\iff\beta^{s*}_{i}=(\beta_{*})_{i}\frac{1}{\left(1-\zeta_{i} \right)+\frac{\Omega}{1-\Omega}\frac{\zeta_{i}}{1-\zeta_{i}}\zeta_{i}}.$ $$

It’s now clear that $\zeta_{i}>1-\Omega$ if and only if $|\beta^{s*}_{i}|<|(\beta_{*})_{i}|$.

Let’s now check when the ratio between them is $1$. Algebraic manipulations give:

$$ $\displaystyle\frac{(1-\zeta_{i})}{\left(1-\zeta_{i}\right)^{2}+\frac{\Omega}{1 -\Omega}\zeta_{i}^{2}}=1$ $\displaystyle\iff(1-\zeta_{i})-\left(1-\zeta_{i}\right)^{2}=\frac{\Omega}{1- \Omega}\zeta_{i}^{2}$ $\displaystyle\iff\zeta_{i}=1-\Omega\iff 1-\zeta_{i}=\Omega\text{ where }\Omega =\frac{\sum_{i=1}^{p}(1-\zeta_{i})^{2}}{\sum_{i=1}^{p}(1-\zeta_{i})}.$ $$

This gives that ${\bm{\beta}}^{s*}={\bm{\beta}_{\star}}$ if all $\zeta_{i}$’s are equal, which implies that all $\lambda_{i}$’s are equal. Concluding, the covariance matrix is a multiple of the identity if and only if ${\bm{\beta}}^{s*}={\bm{\beta}_{\star}}$.
∎

See 2

###### Proof.

For the purposes of analysis, we assume, without loss of generality, that the first $p_{s}$ dimensions are selected from ${\bm{\beta}_{\star}}$ in $\mathcal{M}({\bm{\beta}_{\star}})={\bm{\beta}}^{s}\in\mathbb{R}^{p_{s}}$. Based on this, we no longer need to have the decreasing order for the corresponding $\lambda_{i}$’s. From the excess test risk formula in Definition 2, we have that

$$ $\displaystyle\mathcal{R}({\bm{\beta}}^{t})=\mathbb{E}\left[\left(y-\bm{x}^{ \top}{\bm{\beta}}^{t}\right)^{2}\right]-{\sigma}_{t}^{2}=\frac{\mathcal{B}({ \bm{\beta}_{\star}})+{\sigma}_{t}^{2}\Omega}{1-\Omega}.$ (15) $$

Now, consider the zero-padded vector $\bar{{\bm{\beta}}^{s}}=\begin{bmatrix}{\bm{\beta}}^{s}\\
\mathbf{0}_{p-p_{s}}\end{bmatrix}\in\mathbb{R}^{p}$. This way, we consider the labels in the second training phase as $y^{s}=\bm{x}^{\top}\bar{{\bm{\beta}}^{s}}+z$, where $z\sim\mathcal{N}(0,{\sigma}_{t}^{2})$. Next, using the asymptotic risk estimate in Equation (8), we write the excess test risk formula for the surrogate-to-target model with respect to the original ground truth labels:

$$ $\displaystyle\begin{split}\bar{\mathcal{R}}^{s2t}_{\kappa_{t},{\sigma}_{t}}( \bm{\Sigma}_{t},{\bm{\beta}_{\star}},\bar{{\bm{\beta}}^{s}})&:=(\bar{{\bm{ \beta}}^{s}}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{ t}{\bm{\theta}}_{1}(\bar{{\bm{\beta}}^{s}}-{\bm{\beta}_{\star}})+\gamma_{t}^{2 }(\bar{{\bm{\beta}}^{s}})\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{\theta}}_ {2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]\\ &+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t }({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{\star}}^{\top} ({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}(\bar{{\bm{ \beta}}^{s}}-{\bm{\beta}_{\star}}),\end{split}$ $$

where ${\bm{\theta}}_{1}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}$, ${\bm{\theta}}_{2}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2
}\frac{\bm{g}_{t}}{\sqrt{p}}$, and $\bm{g}_{t}\sim\mathcal{N}(\bm{0},{\bm{I}}_{p})$. Algebraic manipulations give:

$$ $\displaystyle\mathcal{R}({\bm{\beta}}^{s2t})$ $\displaystyle=\mathbb{E}\left[\left(y-\bm{x}^{\top}{\bm{\beta}}^{s2t}\right)^{ 2}\right]-\sigma_{t}^{2}$ $\displaystyle\approx\bar{\mathcal{R}}^{s2t}_{\kappa_{t},{\sigma}_{t}}(\bm{ \Sigma}_{t},{\bm{\beta}_{\star}},\bar{{\bm{\beta}}^{s}})$ $\displaystyle=(\bar{{\bm{\beta}}^{s}}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta} }_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}(\bar{{\bm{\beta}}^{s}}-{\bm{\beta }_{\star}})+\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}(\bar{{\bm{\beta}}^{s }})^{\top}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bar{{\bm{ \beta}}^{s}}}{1-\Omega}\frac{\texttt{tr}\left(\bm{\Sigma}_{t}^{2}(\bm{\Sigma}_ {t}+\tau_{t}{\bm{I}})^{-2}\right)}{p}$ $\displaystyle+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top} \bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{ \star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_ {1}(\bar{{\bm{\beta}}^{s}}-{\bm{\beta}_{\star}})$ $\displaystyle=\sum_{i=p_{s}+1}^{p}\lambda_{i}\beta_{i}^{2}\left(\frac{\lambda_ {i}}{\lambda_{i}+\tau_{t}}\right)^{2}+\Omega\frac{{\sigma}_{t}^{2}+\sum_{i=1}^ {p_{s}}\lambda_{i}\beta_{i}^{2}\left(\frac{\tau_{t}}{\lambda_{i}+\tau_{t}} \right)^{2}}{1-\Omega}$ $\displaystyle~{}+\sum_{i=1}^{p}\lambda_{i}\beta_{i}^{2}\left(\frac{\tau_{t}}{ \lambda_{i}+\tau_{t}}\right)^{2}+\sum_{i=p_{s}+1}^{p}\lambda_{i}\beta_{i}^{2} \frac{2\tau_{t}\lambda_{i}}{(\lambda_{i}+\tau_{t})^{2}}$ $\displaystyle=\frac{{\sigma}_{t}^{2}\Omega+\sum_{i=1}^{p_{s}}\lambda_{i}\beta_ {i}^{2}\zeta_{i}^{2}}{1-\Omega}+\sum_{i=p_{s}+1}^{p}\lambda_{i}\beta_{i}^{2}$ $\displaystyle=\frac{\mathcal{B}(\bar{{\bm{\beta}}^{s}})+{\sigma}_{t}^{2}\Omega }{1-\Omega}+\sum_{i=p_{s}+1}^{p}\lambda_{i}\beta_{i}^{2},$ (16) $$

Thus, the risk difference between the target and surrogate-to-target models is

$$ $\displaystyle\mathcal{R}({\bm{\beta}}^{t})-\mathcal{R}({\bm{\beta}}^{s2t})$ $\displaystyle=\frac{\mathcal{B}({\bm{\beta}_{\star}})-\mathcal{B}(\bar{{\bm{ \beta}}^{s}})}{1-\Omega}-\sum_{i=p_{s}+1}^{p}\lambda_{i}\beta_{i}^{2}$ $\displaystyle=\frac{\sum_{i=p_{s}+1}^{p}\lambda_{i}\zeta_{i}^{2}\beta_{i}^{2}} {1-\Omega}-\sum_{i=p_{s}+1}^{p}\lambda_{i}\beta_{i}^{2}.$ $$

We observe that each dimension’s contribution to the excess test risk can be analyzed individually. Therefore, if

$$ $\displaystyle\zeta_{i}^{2}>1-\Omega,$ (17) $$

excluding feature $i$ in the feature selection reduces the overall risk $\mathcal{R}({\bm{\beta}}^{s2t})$. Along the same lines, the projection $\mathcal{M}$ that selects all the features $i$ that satisfy $\zeta_{i}^{2}<1-\Omega$ minimizes the asymptotic excess test risk.
∎

###### Proof.

###### Observation 2 .

###### Proof.

###### Proof.

###### Proof.

###### Proof.

###### Proof.

## Appendix B Proofs for Section 4

See 2
In the following proof, we suppose that the empirical distributions of $\bar{\bm{\beta}}$ and $\bm{\lambda}$ converge as $p\to\infty$ having fixed the ratio $p/n=\kappa$. Then, we will prove that the omniscient risk converges to the asymptotic risk defined in (8).

###### Proof for the proportional asymptotic case.

Using Theorem 2.3 of Han & Xu (), we can estimate $\hat{\bm{\beta}}$ as follows:

$$ $\displaystyle\hat{\bm{\beta}}=(\bm{\Sigma}+\tau{\bm{I}})^{-1}\bm{\Sigma}\left( {\bm{\beta}_{\star}}+\frac{\bm{\Sigma}^{-1/2}\gamma_{s}({\bm{\beta}_{\star}}){ \bm{g}}}{\sqrt{p}}\right),$ $$

where

$$ ${\bm{g}}\sim\mathcal{N}(0,{\bm{I}}_{p}),\quad\gamma_{s}^{2}({\bm{\beta}_{\star }})=\kappa\frac{\sigma+\tau^{2}\|{(\bm{\Sigma}+\tau{\bm{I}})^{-1}\bm{\Sigma}^{ 1/2}{\bm{\beta}_{\star}}}\|_{2}^{2}}{1-\frac{1}{n}\texttt{tr}\left((\bm{\Sigma }+\tau{\bm{I}})^{-2}\bm{\Sigma}^{2}\right)},\quad\tau\text{ is the solution to }n=\sum_{i=1}^{p}\frac{\lambda_{i}}{\lambda_{i}+\tau}.$ $$

Let

$$ $\displaystyle{\bm{X}}_{1}=(\bm{\Sigma}+\tau{\bm{I}})^{-1}\bm{\Sigma}\quad, \quad{\bm{X}}_{2}=\frac{(\bm{\Sigma}+\tau{\bm{I}})^{-1}\bm{\Sigma}^{1/2}\gamma _{s}({\bm{\beta}_{\star}})}{\sqrt{p}}.$ $$

Using this estimate, we can calculate the excess test risk as

$$ $\displaystyle\mathcal{R}(\hat{\bm{\beta}})=\quad$ $\displaystyle\mathbb{E}\left[\left(\left({\bm{X}}_{1}-{\bm{I}}\right){\bm{ \beta}_{\star}}+{\bm{X}}_{2}{\bm{g}}\right)^{\top}\bm{\Sigma}\left(\left({\bm{ X}}_{1}-{\bm{I}}\right){\bm{\beta}_{\star}}+{\bm{X}}_{2}{\bm{g}}\right)\right]$ $\displaystyle=\quad$ $\displaystyle{\bm{\beta}_{\star}}^{\top}({\bm{X}}_{1}-{\bm{I}})^{\top}\bm{ \Sigma}({\bm{X}}_{1}-{\bm{I}}){\bm{\beta}_{\star}}+\mathbb{E}\left[{\bm{g}}^{ \top}{\bm{X}}_{2}^{\top}\bm{\Sigma}{\bm{X}}_{2}{\bm{g}}\right]$ $\displaystyle=\quad$ $\displaystyle{\bm{\beta}_{\star}}^{\top}({\bm{X}}_{1}-{\bm{I}})^{\top}\bm{ \Sigma}({\bm{X}}_{1}-{\bm{I}}){\bm{\beta}_{\star}}+\texttt{tr}\left({\bm{X}}_{ 2}^{\top}\bm{\Sigma}{\bm{X}}_{2}\right).$ (18) $$

Then by recalling the eigendecomposition for the covariance matrix $\bm{\Sigma}={\bm{U}}{\bm{\Lambda}}{\bm{U}}^{\top}$, we have

$$ $\displaystyle{\bm{X}}_{1}$ $\displaystyle=({\bm{U}}{\bm{\Lambda}}{\bm{U}}^{\top}+\tau{\bm{U}}{\bm{U}}^{ \top})^{-1}{\bm{U}}{\bm{\Lambda}}U^{\top}$ $\displaystyle={\bm{U}}({\bm{\Lambda}}+\tau{\bm{I}})^{-1}{\bm{U}}^{\top}{\bm{U} }{\bm{\Lambda}}{\bm{U}}^{\top}$ $\displaystyle={\bm{U}}\operatorname{diag}\left(\frac{\bm{\lambda}}{\bm{\lambda }+\tau}\right){\bm{U}}^{\top}.$ $$

Using the diagonalization of ${\bm{I}}$, ${\bm{X}}_{1}-{\bm{I}}$ can now be computed as

$$ $\displaystyle{\bm{X}}_{1}-{\bm{I}}$ $\displaystyle={\bm{U}}\operatorname{diag}\left(\frac{-\tau}{\bm{\lambda}+\tau} \right){\bm{U}}^{\top}.$ $$

Let’s now compute

$$ $\displaystyle{\bm{\beta}_{\star}}^{\top}({\bm{X}}_{1}-{\bm{I}})^{\top}\bm{ \Sigma}({\bm{X}}_{1}-{\bm{I}}){\bm{\beta}_{\star}}$ $\displaystyle={\bm{\beta}_{\star}}^{\top}{\bm{U}}\operatorname{diag}\left( \frac{-\tau}{\bm{\lambda}+\tau}\right){\bm{U}}^{\top}{\bm{U}}{\bm{\Lambda}}{ \bm{U}}^{\top}{\bm{U}}\operatorname{diag}\left(\frac{-\tau}{\bm{\lambda}+\tau} \right){\bm{U}}^{\top}{\bm{\beta}_{\star}}$ $\displaystyle={\bm{\beta}_{\star}}^{\top}{\bm{U}}\operatorname{diag}\left( \frac{\bm{\lambda}\tau^{2}}{(\bm{\lambda}+\tau)^{2}}\right){\bm{U}}^{\top}{\bm {\beta}_{\star}}.$ $$

As $\bar{\bm{\beta}}={\bm{U}}^{\top}{\bm{\beta}_{\star}}$, we obtain that the RHS of the previous expression equals

$$ $\displaystyle\sum_{i=1}^{p}\frac{\lambda_{i}\tau^{2}\bar{\beta}_{i}^{2}}{( \lambda_{i}+\tau)^{2}}=\mathcal{B}(\bar{\bm{\beta}}).$ $$

Next, we write more compactly the terms $\texttt{tr}\left({\bm{X}}_{2}^{\top}\bm{\Sigma}{\bm{X}}_{2}\right)$ and $\gamma_{s}^{2}({\bm{\beta}_{\star}})$. By defining the short-hand notation $\Omega=\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}+\tau{\bm{I}})^{-2}\bm{\Sigma}^
{2}\right)=\frac{1}{n}\sum_{i=1}^{p}(1-\zeta_{i})^{2}$, we have

$$ $\displaystyle\texttt{tr}\left({\bm{X}}_{2}^{\top}\bm{\Sigma}{\bm{X}}_{2}\right)$ $\displaystyle=\frac{\gamma_{s}^{2}({\bm{\beta}_{\star}})}{p}\sum_{i=1}^{p} \left(\frac{\lambda_{i}}{\lambda_{i}+\tau}\right)^{2}=\frac{\gamma_{s}^{2}({ \bm{\beta}_{\star}})n\Omega}{p},$ $\displaystyle\gamma_{s}^{2}({\bm{\beta}_{\star}})$ $\displaystyle=\kappa\frac{\sigma^{2}+\tau^{2}\|{(\bm{\Sigma}+\tau{\bm{I}})^{-1 }\bm{\Sigma}^{1/2}{\bm{\beta}_{\star}}}\|_{2}^{2}}{1-\Omega}=\kappa\frac{ \sigma^{2}+\sum_{i=1}^{p}\frac{\lambda_{i}\tau^{2}\bar{\beta}_{i}^{2}}{( \lambda_{i}+\tau)^{2}}}{1-\Omega}=\kappa\frac{\sigma^{2}+\mathcal{B}(\bar{\bm{ \beta}})}{1-\Omega},$ $$

where $\kappa=\frac{p}{n}$. Hence, putting it all together in (18) gives the desired result.
∎

See 3

###### Proof.

We start with the asymptotic analysis of $\tau_{t}$.
Along the same lines as Simon et al. (), since $\frac{i^{-\alpha}}{i^{-\alpha}+\tau_{t}}$ is a monotonically decreasing function, we have:

$$ $\displaystyle n=\sum_{i=1}^{\infty}\frac{i^{-\alpha}}{i^{-\alpha}+\tau_{t}} \leq\int_{0}^{\infty}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx=\dfrac{\pi}{ \alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}.$ $$

Furthermore,

$$ $\displaystyle\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-1=\int_ {0}^{\infty}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx-1\leq\int_{1}^{\infty }\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx\leq\sum_{i=1}^{\infty}\frac{i^{- \alpha}}{i^{-\alpha}+\tau_{t}}=n.$ $$

Hence, combining these two facts gives

$$ $\displaystyle\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-1\leq n \leq\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}$ $\displaystyle\iff$ $\displaystyle\left(\frac{\left(n+1\right)\alpha\sin{(\pi/\alpha)}}{\pi}\right) ^{-\alpha}\leq\tau_{t}\leq\left(\frac{n\alpha\sin{(\pi/\alpha)}}{\pi}\right)^{ -\alpha},$ $$

which leads to the desired result.

Next, we move to the asymptotic analysis of $\Omega$. We have that

$$ $\displaystyle n\Omega=\sum_{i=1}^{\infty}\left(\frac{i^{-\alpha}}{i^{-\alpha}+ \tau_{t}}\right)^{2}$ $\displaystyle\leq\int_{0}^{\infty}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t }}\right)^{2}\,dx=\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^ {-1/\alpha}.$ $$

Besides, since the summand is monotonically decreasing, we also have

$$ $\displaystyle\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{-1/ \alpha}-1\leq\int_{0}^{\infty}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}} \right)^{2}\,dx-1\leq\int_{1}^{\infty}\left(\frac{x^{-\alpha}}{x^{-\alpha}+ \tau_{t}}\right)^{2}\,dx\leq\sum_{i=1}^{\infty}\left(\frac{i^{-\alpha}}{i^{- \alpha}+\tau_{t}}\right)^{2}=n\Omega.$ $$

Hence,

$$ $\displaystyle\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{-1/ \alpha}-1\leq n\Omega\leq\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})} \tau_{t}^{-1/\alpha}.$ (19) $$

By the hypothesis on $\tau_{t}$, we have that

$$ $\displaystyle\tau_{t}^{-1/\alpha}=n\frac{\alpha\sin{(\pi/\alpha})}{\pi}\left(1 -O\left(n^{-1}\right)\right),$ (20) $$

and plugging this in (19) gives the desired result. ∎

See 4

###### Proof.

Recall from Proposition 2 that we should identify indices $i$ which satisfy the condition $\zeta_{i}^{2}>1-\Omega$ to decide if we’re better off not selecting this dimension $i$ in the surrogate model. Furthermore, Proposition 3 gives that $\Omega=\dfrac{\alpha-1}{\alpha}-O(n^{-1})$. Putting these together, we have

$$ $\displaystyle\zeta_{i}^{2}>1-\Omega$ $\displaystyle\iff$ $\displaystyle\zeta_{i}^{2}>c^{\prime}\quad\text{where }c^{\prime}=\frac{1}{ \alpha}+O(n^{-1})$ $\displaystyle\iff$ $\displaystyle\frac{\tau_{t}^{2}}{(\tau_{t}+i^{-\alpha})^{2}}=\frac{\tau_{t}^{2 }i^{2\alpha}}{(\tau_{t}i^{\alpha}+1)^{2}}>c^{\prime}$ $\displaystyle\iff$ $\displaystyle(1-c^{\prime})\tau_{t}^{2}i^{2\alpha}>2c^{\prime}\tau_{t}i^{ \alpha}+c^{\prime}$ $\displaystyle\iff$ $\displaystyle\left(\sqrt{1-c^{\prime}}\tau_{t}i^{\alpha}-\frac{c^{\prime}}{ \sqrt{1-c^{\prime}}}\right)^{2}>\frac{c^{\prime}}{1-c^{\prime}}$ $\displaystyle\iff$ $\displaystyle i^{\alpha}>\frac{\sqrt{c^{\prime}}}{\tau_{t}(1-\sqrt{c^{\prime}})}$ $\displaystyle\iff$ $\displaystyle i>\tau_{t}^{-1/\alpha}\left(\frac{\sqrt{c^{\prime}}}{1-\sqrt{c^{ \prime}}}\right)^{1/\alpha}$ $$

As $c^{\prime}=\frac{1}{\alpha}+O(n^{-1})$, we get $\left(\frac{\sqrt{c^{\prime}}}{1-\sqrt{c^{\prime}}}\right)^{1/\alpha}=\frac{1}
{(\sqrt{\alpha}-1)^{1/\alpha}}(1+O(n^{-1}))$. Incorporating (20), we achieve that

$$ $\displaystyle\tau_{t}^{-1/\alpha}\left(\frac{\sqrt{c^{\prime}}}{1-\sqrt{c^{ \prime}}}\right)^{1/\alpha}=n\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\sqrt{\alpha }-1)^{1/\alpha}}\left(1+O(n^{-1})\right)=nC_{2}+O(1).$ $$

Similarly, by following the same procedure with the initial inequality $\zeta_{i}>1-\Omega$, we get

$$ $\displaystyle\zeta_{i}>1-\Omega\iff i>nC_{1}+O(1),\quad\text{where}\quad C_{1} =\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\alpha-1)^{1/\alpha}}.$ $$

∎

In Figure 3, we compare the empirical results with theoretical predictions for the number of features that meet the selection criteria in the optimal mask $\mathcal{M}^{*}$ ($\zeta_{i}^{2}<1-\Omega$). The theoretical value, calculated as $n\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\sqrt{\alpha}-1)^{1/\alpha}}$ ignoring the $O(1)$ term, aligns well with the experimental data and the accuracy in estimation increases with $\alpha$.

Figure: Figure 3: Comparison of the empirical and theoretical number of features satisfying the feature selection condition in the optimal mask $\mathcal{M}^{*}$ ($\zeta_{i}^{2}<1-\Omega$). The theoretical value is calculated as $n\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\sqrt{\alpha}-1)^{1/\alpha}}$, ignoring the $O(1)$ in Proposition 4. Notably, our theoretical estimate aligns very closely with the empirical results when the sample size $n$ is small relative to the feature size $p$. Setting: The feature size is $p=500$, and the feature covariance follows the power-law structure $\lambda_{i}=i^{-\alpha}$ for $\alpha=1.5,3.0,$ and $4.5$.
Refer to caption: x5.png

###### Proposition 6 (Scaling law for masked surrogate-to-target model) .

Together with the eigenvalues, also assume now power-law form for $\lambda_{i}\beta_{i}^{2}$, that is $\lambda_{i}\beta_{i}^{2}=i^{-\beta}$ for $\beta>1$. Then, in the limit of $p\to\infty$, the excess test risk for the masked surrogate-to-target model with the optimal dimensionality has the same scaling law as the reference (target) model:

$$ $\displaystyle\mathcal{R}({\bm{\beta}}^{s2t})=\Theta(n^{-(\beta-1)})\quad\text{ if }\beta<2\alpha+1,$ $$

and

$$ $\displaystyle\mathcal{R}({\bm{\beta}}^{s2t})=\Theta(n^{-2\alpha})\quad\text{if }\beta>2\alpha+1.$ $$

###### Proof.

As discussed in Section 4, in order to analyze the model’s inherent error, we need to set ${\sigma}_{t}^{2}=O(n^{-\gamma})$ where $\gamma$ is the exponent characterizing the scaling law of the test risk in the noiseless setting. We will work on this proof in two cases depending on $\beta$ and $2\alpha+1$.

Case 1: $\beta<2\alpha+1$. In this case, it is previously stated by Cui et al. (); Simon et al. () that the test risk of ridgeless overparameterized linear regression can be described in the scaling sense as $\mathbf{err}=\Theta(n^{-\beta+1})$ when $\beta<2\alpha+1$. Consider the optimal mask operation $\mathcal{M}$ mentioned in Proposition 2 that selects all features satisfying $1-\zeta_{i}^{2}>\Omega$. Let $p_{s}$ be the number of selected features. We can then decompose the risk estimate in Definition 2 as follows:

$$ $\displaystyle\frac{\mathcal{B}(\bar{{\bm{\beta}_{\star}}})+{\sigma}_{t}^{2} \Omega}{1-\Omega}=\frac{\sum_{i=1}^{p_{s}}\lambda_{i}\zeta_{i}^{2}\beta_{i}^{2 }+\sum_{i=p_{s}+1}^{p}\lambda_{i}\zeta_{i}^{2}\beta_{i}^{2}+{\sigma}_{t}^{2} \Omega}{1-\Omega}=\frac{\mathbf{err1}+\mathbf{err2}+{\sigma}_{t}^{2}\Omega}{1- \Omega},$ $$

where $\mathbf{err1}$ and $\mathbf{err2}$ are the contributions to the total risk of the target model from dimensions selected and omitted in the surrogate model, respectively. Therefore, we express the total error as:

$$ $\displaystyle\frac{\mathbf{err1}+\mathbf{err2}+{\sigma}_{t}^{2}\Omega}{1- \Omega}=\mathbf{err}=\Theta(n^{-\beta+1}).$ $$

Going back to Proposition 4, we know that, as $p\to\infty$, the criterion for selecting a feature $i$ in the optimal masked surrogate model is given by

$$ $\displaystyle i<nC_{2}+O(1),\quad\text{where}\quad C_{2}=\dfrac{\alpha\sin{( \pi/\alpha)}}{\pi(\sqrt{\alpha}-1)^{1/\alpha}}.$ $$

Define now $\omega_{n}=nC_{2}+O(1)$. The equation (16) tells us that after the optimal mask operation $\mathcal{M}$, $\dfrac{\mathbf{err2}}{1-\Omega}$ is replaced by $\mathbf{err2}^{\prime}$, which is calculated as follows

$$ $\displaystyle\mathbf{err2}^{\prime}=\sum_{i=\omega_{n}+1}^{p}\lambda_{i}\beta_ {i}^{2}=\sum_{i=\omega_{n}+1}^{p}i^{-\beta}.$ $$

Since $x^{-\beta}$ is a monotonically decreasing function, we can bound the summation by the following two integrals:

$$ $\displaystyle\int_{\omega_{n}+1}^{p+1}x^{-\beta}\,dx$ $\displaystyle\leq\sum_{\omega_{n}+1}^{p}i^{-\beta}\leq\int_{\omega_{n}}^{p}x^{ -\beta}\,dx$ $\displaystyle\frac{(\omega_{n}+1)^{-\beta+1}-(p+1)^{-\beta+1}}{\beta-1}$ $\displaystyle\leq\sum_{\omega_{n}+1}^{p}i^{-\beta}\leq\frac{(\omega_{n})^{- \beta+1}-p^{-\beta+1}}{\beta-1}.$ $$

In the limit of $p\to\infty$, we obtain

$$ $\displaystyle\mathbf{err2}^{\prime}=\Theta(n^{-\beta+1}).$ $$

Thus, we have tightly estimated $\mathbf{err2}^{\prime}$. Using the fact from Proposition 3 that $\Omega=\Theta(1)$, and our assumption on the noise variance ${\sigma}_{t}^{2}=O(n^{-\beta+1})$, we conclude that the scaling law doesn’t change for the surrogate-to-target model as

$$ $\displaystyle\mathcal{R}({\bm{\beta}}^{s2t})=\frac{\mathbf{err1}+{\sigma}_{t}^ {2}\Omega}{1-\Omega}+\mathbf{err2}^{\prime}=\Theta(n^{-\beta+1}).$ $$

Case 2: $\beta>2\alpha+1$. In this case, we show that the scaling law is determined by $\mathbf{err1}$, hence changing $\mathbf{err2}$ to $\mathbf{err2}^{\prime}$ has no effect in the scaling sense. From Proposition 3, we have the asymptotic expression $\tau_{t}=cn^{-\alpha}\left(1+O(n^{-1})\right)$, for $c=\left(\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\right)^{\alpha}$. We can argue that there exists positive constants $c_{1}<\dfrac{1}{c}<c_{2}$, such that $c_{1}n^{\alpha}\leq\dfrac{1}{\tau_{t}}\leq c_{2}n^{\alpha}$. We have that

$$ $\displaystyle\mathbf{err1}=\sum_{i=1}^{\omega_{n}}\frac{i^{-\beta}}{(1+\frac{1 }{\tau_{t}}i^{-\alpha})^{2}}$ $\displaystyle\leq\sum_{i=1}^{\omega_{n}}\frac{i^{-\beta}}{(1+c_{1}n^{\alpha}i^ {-\alpha})^{2}}$ $\displaystyle=\sum_{i=1}^{\omega_{n}}\frac{i^{2\alpha-\beta}}{(i^{\alpha}+c_{1 }n^{\alpha})^{2}}\leq\sum_{i=1}^{\omega_{n}}\frac{i^{2\alpha-\beta}}{c_{1}^{2} n^{2\alpha}}.$ $$

This implies $\mathbf{err1}=O(n^{-2\alpha})$. At the same time,

$$ $\displaystyle\mathbf{err1}=\sum_{i=1}^{\omega_{n}}\frac{i^{-\beta}}{(1+\frac{1 }{\tau_{t}}i^{-\alpha})^{2}}$ $\displaystyle\geq\sum_{i=1}^{\omega_{n}}\frac{i^{2\alpha-\beta}}{(i^{\alpha}+c _{2}n^{\alpha})^{2}}$ $\displaystyle\geq\sum_{i=1}^{\omega_{n}}\frac{i^{2\alpha-\beta}}{((\omega_{n}) ^{\alpha}+c_{2}n^{\alpha})^{2}}=\sum_{i=1}^{\omega_{n}}\frac{i^{2\alpha-\beta} }{n^{2\alpha}((\omega_{n}/n)^{\alpha}+c_{2})^{2}}.$ $$

Using $\omega_{n}/n=\Theta(1)$ gives $\mathbf{err1}=\Omega(n^{-2\alpha})$ and we can conclude that $\mathbf{err1}=\Theta(n^{-2\alpha})$. From Cui et al. (), we already know that $\mathbf{err}=\Theta(n^{-2\alpha})$ when $\beta>2\alpha+1$. Using $\Omega=\Theta(1)$, and our assumption on the noise variance ${\sigma}_{t}^{2}=O(n^{-2\alpha})$ allows us to conclude that the scaling is dominated by $\mathbf{err1}$, and thus, the scaling law remains unchanged.
∎

See 5

###### Proof.

From asymptotic risk decomposition in (26), we can write

$$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[f(X_{\kappa_{t},{ \sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},\bm{g}_{t}))\right]$ $\displaystyle=({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta}}_{1}^ {\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})+ \gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{ \theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]$ $\displaystyle+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top} \bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{ \star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_ {1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})$ $\displaystyle\geq\gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{\mathbb{E}}_{ \bm{g}_{t}}[{\bm{\theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}],$ $$

since we can put in the form of $(a-b)^{2}+c^{2}\geq c^{2}$. At the same time, we know that

$$ $\displaystyle\gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{\mathbb{E}}_{\bm{g} _{t}}[{\bm{\theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]$ $\displaystyle=\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}\|{(\bm{\Sigma}_{t} +\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}{\bm{\beta}_{\star}}}\|_{2}^{2}}{1 -\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{\Sigma }_{t}^{2}\right)}\frac{\texttt{tr}\left(\bm{\Sigma}_{t}^{2}(\bm{\Sigma}_{t}+ \tau_{t}{\bm{I}})^{-2}\right)}{p}$ $\displaystyle=\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}\|{(\bm{\Sigma}_{t} +\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}{\bm{\beta}_{\star}}}\|_{2}^{2}}{1 -\Omega}\frac{n\Omega}{p}$ $\displaystyle=\frac{\Omega}{1-\Omega}\left({\sigma}_{t}^{2}+\sum_{i=1}^{p} \lambda_{i}\beta_{i}^{2}\zeta_{i}^{2}\right).$ $$

Recall the optimal surrogate vector discussed in Proposition 1 and the corresponding minimal surrogate-to-target risk $\mathcal{R}^{*}({\bm{\beta}}^{s2t})$. In this case, we can write

$$ $\displaystyle\sum_{i=1}^{p}\lambda_{i}\beta_{i}^{s*2}\zeta_{i}^{2}=\sum_{i=1}^ {p}\lambda_{i}\beta_{i}^{2}\frac{(1-\zeta_{i})^{2}\zeta_{i}^{2}}{\left((1- \zeta_{i})^{2}+\frac{\Omega}{1-\Omega}\zeta_{i}^{2}\right)^{2}}.$ $$

Similar to the previous proposition and as discussed in Section 4, to analyze the model’s inherent error, we set ${\sigma}_{t}^{2}=O(n^{-\gamma})$ where $\gamma$ is the exponent characterizing the scaling law of the test risk in the noiseless setting. It is previously stated by Cui et al. (); Simon et al. () that the test risk of ridgeless overparameterized linear regression can be described in the scaling sense as $\mathbf{err}=\Theta(n^{-\beta+1})$ when $\beta<2\alpha+1$. We will proceed by considering two cases based on the relationship between $\beta$ and $2\alpha+1$.

Case 1: $\beta<2\alpha+1$

Consider the interval of $i$’s satisfying $\zeta_{i}>1-\Omega$ and $\zeta_{i}^{2}<1-\Omega$. By Proposition 4, we have

$$ $\displaystyle\zeta_{i}>1-\Omega\iff i>nC_{1}+O(1),\quad\text{where}\quad C_{1} =\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\alpha-1)^{1/\alpha}}.$ $\displaystyle\zeta_{i}^{2}>1-\Omega\iff i>nC_{2}+O(1),\quad\text{where}\quad C _{2}=\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\sqrt{\alpha}-1)^{1/\alpha}}.$ $$

Let $\omega_{n}$ be defined as in the previous proposition and define $\phi_{n}=nC_{1}+O(1)$. Then, the interval of interest corresponds to the set of indices $i$ such that $\phi_{n}<i<\omega_{n}$. Within this interval, we observe

$$ $\displaystyle(1-\zeta_{i})^{2}\zeta_{i}^{2}$ $\displaystyle\geq\left(1-\sqrt{(1-\Omega)}\right)^{2}(1-\Omega)^{2}=k_{1}$ $\displaystyle(1-\zeta_{i})^{2}+\frac{\Omega}{1-\Omega}\zeta_{i}^{2}$ $\displaystyle\leq 1+\frac{\Omega}{1-\Omega}=k_{2}$ $$

Using the fact from Proposition 3 that $\Omega=\dfrac{\alpha-1}{\alpha}-O(n^{-1})$ tells us $k_{1}=\Theta(1)$ and $k_{2}=\Theta(1)$. Utilizing these bounds, we obtain

$$ $\displaystyle\mathcal{R}^{*}({\bm{\beta}}^{s2t})\geq\frac{\Omega}{1-\Omega} \sum_{i=1}^{p}\lambda_{i}\beta_{i}^{2}\frac{(1-\zeta_{i})^{2}\zeta_{i}^{2}}{ \left((1-\zeta_{i})^{2}+\frac{\Omega}{1-\Omega}\zeta_{i}^{2}\right)^{2}}$ $\displaystyle\geq\frac{\Omega}{1-\Omega}\sum_{i=\phi_{n}}^{\omega_{n}}i^{- \beta}\frac{(1-\zeta_{i})^{2}\zeta_{i}^{2}}{\left((1-\zeta_{i})^{2}+\frac{ \Omega}{1-\Omega}\zeta_{i}^{2}\right)^{2}}$ $\displaystyle\geq\frac{\Omega}{1-\Omega}\sum_{i=\phi_{n}}^{\omega_{n}}i^{- \beta}\frac{k_{1}}{k_{2}}$ $\displaystyle\geq n^{-\beta+1}\frac{k_{1}}{k_{2}}\frac{\Omega}{1-\Omega}\left( \frac{(\phi_{n}/n)^{-\beta+1}-(\omega_{n}/n)^{-\beta+1}}{\beta-1}\right)$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\Theta(n^{-\beta+1}),$ $$

where (a) follows from the fact $\omega_{n}/n=\Theta(1)$, $\phi_{n}/n=\Theta(1)$, and $\Omega=\Theta(1)$. This implies that

$$ $\displaystyle\mathcal{R}^{*}({\bm{\beta}}^{s2t})=\Omega(n^{-\beta+1}).$ $$

Recall that the optimal surrogate-to-target improves over the risk of the standard target model, thus $\mathcal{R}^{*}({\bm{\beta}}^{s2t})=O(n^{-\beta+1})$. We therefore conclude $\mathcal{R}^{*}({\bm{\beta}}^{s2t})=\Theta(n^{-\beta+1})$ for this case.

Case 2: $\beta>2\alpha+1$

In this case, we have

$$ $\displaystyle\mathcal{R}^{*}({\bm{\beta}}^{s2t})\geq\frac{\Omega}{1-\Omega} \sum_{i=1}^{p}\lambda_{i}\beta_{i}^{2}\frac{(1-\zeta_{i})^{2}\zeta_{i}^{2}}{ \left((1-\zeta_{i})^{2}+\frac{\Omega}{1-\Omega}\zeta_{i}^{2}\right)^{2}}$ $\displaystyle=\frac{\Omega}{1-\Omega}\sum_{i=1}^{p}\lambda_{i}\beta_{i}^{2} \zeta_{i}^{2}\frac{(1-\zeta_{i})^{2}}{\left((1-\zeta_{i})^{2}+\frac{\Omega}{1- \Omega}\zeta_{i}^{2}\right)^{2}}$ $\displaystyle\geq\sum_{i:\zeta_{i}<1-\Omega}\lambda_{i}\zeta_{i}^{2}\beta_{i}^ {2}\frac{\Omega^{3}}{\left(1+\frac{\Omega}{1-\Omega}\right)^{2}(1-\Omega)}= \sum_{i=1}^{\phi_{n}}\frac{i^{-\beta}}{(1+\frac{1}{\tau_{t}}i^{-\alpha})^{2}}k _{3},$ $$

where $k_{3}=\frac{\Omega^{3}}{\left(1+\frac{\Omega}{1-\Omega}\right)^{2}(1-\Omega)}=
\Theta(1)$. From Case 2 in Proposition 6, we already know that the same summation – with upper bound $\omega_{n}$ rather than $\phi_{n}$ – scales as $\Theta(n^{-2\alpha})$. Yet, since $\phi_{n}$ and $\omega_{n}$ have the same order $\Theta(n)$, the result remains. This gives $\mathcal{R}^{*}({\bm{\beta}}^{s2t})=\Omega(n^{-2\alpha})$, which eventually yields

$$ $\displaystyle\mathcal{R}^{*}({\bm{\beta}}^{s2t})\leq\mathcal{R}({\bm{\beta}}^{ t})=O(n^{-2\alpha})\implies\mathcal{R}^{*}({\bm{\beta}}^{s2t})=\Theta(n^{-2 \alpha}).$ $$

Hence, this allows us to say that the scaling law doesn’t improve even with the freedom to choose any ${\bm{\beta}}^{s}$.
∎

###### Proposition 7 (Non-asymptotic analysis of $\tau$ ) .

Suppose that $\bm{\Sigma}\in\mathbb{R}^{p\times p}$ is diagonal and $\bm{\Sigma}_{i,i}=\lambda_{i}=i^{-\alpha}$ for $1<\alpha$. Assume that $n<pk$ for $k=\dfrac{3+2^{-\alpha}}{4+2^{-(\alpha-2)}}$. If $\tau_{t}$ satisfies

$$ $\displaystyle\sum_{i=1}^{p}\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}}=n,$ $$

then $cn^{\alpha}\leq\tau_{t}^{-1}\leq c\left(n+1+\frac{p+1}{\alpha-1}\right)^{\alpha}$ for $c=\left(\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi}\right)^{\alpha}$.

###### Proof.

In a similar vein to  Simon et al. (), we have:

$$ $\displaystyle n=\sum_{i=1}^{p}\frac{i^{-\alpha}}{i^{-\alpha}+\tau_{t}}$ $\displaystyle\leq\int_{0}^{p}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx\leq \int_{0}^{\infty}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx=\dfrac{\pi}{ \alpha\sin{(\pi/\alpha)}}\tau^{-1/\alpha}.$ $$

Since the summand is decreasing, we can bound the Riemann sum by an integral, thus:

$$ $\displaystyle n=\sum_{i=1}^{p}\frac{i^{-\alpha}}{i^{-\alpha}+\tau_{t}}$ $\displaystyle\geq\int_{1}^{p+1}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx$ $\displaystyle=\int_{0}^{\infty}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx- \int_{0}^{1}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx-\int_{p+1}^{\infty} \frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx$ $\displaystyle=\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-\int_{ 0}^{1}\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx-\int_{p+1}^{\infty}\frac{x^ {-\alpha}}{x^{-\alpha}+\tau_{t}}\,dx$ $\displaystyle\geq\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-1- \int_{p+1}^{\infty}\frac{1}{1+\tau_{t}x^{\alpha}}\,dx$ $\displaystyle\geq\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-1- \int_{p+1}^{\infty}\frac{1}{\tau_{t}x^{\alpha}}\,dx$ $\displaystyle=\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-1- \left[\frac{\tau_{t}^{-1}x^{-\alpha+1}}{-\alpha+1}\right]_{p+1}^{\infty}$ $\displaystyle=\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-\frac{ \tau_{t}^{-1}(p+1)^{-\alpha+1}}{\alpha-1}-1.$ $$

Recalling that $\alpha>1$ and assuming $\tau_{t}^{-1}<p^{\alpha}$, we derive:

$$ $\displaystyle\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}-1-\frac {p+1}{\alpha-1}\leq n\leq\dfrac{\pi}{\alpha\sin{(\pi/\alpha)}}\tau_{t}^{-1/\alpha}$ $\displaystyle\iff$ $\displaystyle\left(\frac{n\alpha\sin{(\pi/\alpha)}}{\pi}\right)^{\alpha}\leq \tau_{t}^{-1}\leq\left(\frac{\left(n+1+\frac{p+1}{\alpha-1}\right)\alpha\sin{( \pi/\alpha)}}{\pi}\right)^{\alpha}.$ $$

We conclude by proving that $\tau_{t}^{-1}<p^{\alpha}$. For the sake of contradiction, assume that $\tau_{t}^{-1}\geq p^{\alpha}$. Then,

$$ $\displaystyle n=\sum_{i=1}^{p}\frac{1}{1+i^{\alpha}\tau_{t}}$ $\displaystyle=\sum_{i=1}^{p/2}\frac{1}{1+i^{\alpha}\tau_{t}}+\sum_{i=p/2+1}^{p }\frac{1}{1+i^{\alpha}\tau_{t}}$ $\displaystyle\geq\sum_{i=1}^{p/2}\frac{1}{1+\dfrac{1}{2^{\alpha}}}+\sum_{i=p/2 +1}^{p}\frac{1}{1+1}$ $\displaystyle=p\left(\frac{3+\dfrac{1}{2^{\alpha}}}{4+\dfrac{1}{2^{\alpha-2}}} \right),$ $$

which contradicts our assumption that $n<pk$.
∎

###### Proposition 8 (Non-asymptotic analysis of $\Omega$ ) .

Suppose that $\bm{\Sigma}\in\mathbb{R}^{p\times p}$ is diagonal and $\bm{\Sigma}_{i,i}=\lambda_{i}=i^{-\alpha}$ for $1<\alpha$. Let $\tau_{t}$ be defined as in Proposition 7 and assume that $pk_{1}+\dfrac{\alpha^{2}}{(\alpha-1)^{2}}<n<pk_{2}$ for $k_{1}=\dfrac{\alpha}{(\alpha-1)^{2}}$ and $k_{2}=\dfrac{3+2^{-\alpha}}{4+2^{-(\alpha-2)}}$. Let $\Omega$ be the solution to

$$ $\displaystyle n\Omega=\sum_{i=1}^{p}\left(\frac{\lambda_{i}}{\lambda_{i}+\tau_ {t}}\right)^{2}.$ $$

Then,

$$ $\displaystyle\Omega>\dfrac{\alpha-1}{\alpha}-\frac{1}{\alpha}\left(\dfrac{n+1+ \frac{p+1}{\alpha-1}}{p+1}\right)^{2\alpha-1}-\frac{1}{n}.$ $$

###### Proof.

Since the summand is monotonically decreasing, we can lower bound the sum by the following integral and manipulate:

$$ $\displaystyle n\Omega=\sum_{i=1}^{p}\left(\frac{i^{-\alpha}}{i^{-\alpha}+\tau_ {t}}\right)^{2}$ $\displaystyle\geq\int_{1}^{p+1}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}} \right)^{2}\,dx$ $\displaystyle=\int_{0}^{\infty}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}} \right)^{2}\,dx-\int_{0}^{1}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}} \right)^{2}\,dx-\int_{p+1}^{\infty}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{ t}}\right)^{2}\,dx$ $\displaystyle=\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{-1/ \alpha}-\int_{0}^{1}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\right)^{2} \,dx-\int_{p+1}^{\infty}\left(\frac{x^{-\alpha}}{x^{-\alpha}+\tau_{t}}\right)^ {2}\,dx$ $\displaystyle\geq\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{ -1/\alpha}-1-\int_{p+1}^{\infty}\left(\frac{1}{1+\tau_{t}x^{\alpha}}\right)^{2 }\,dx$ $\displaystyle\geq\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{ -1/\alpha}-1-\int_{p+1}^{\infty}\frac{1}{(\tau_{t}x^{\alpha})^{2}}\,dx$ $\displaystyle=\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{-1/ \alpha}-1-\left[\frac{\tau_{t}^{-2}x^{-2\alpha+1}}{-2\alpha+1}\right]_{p+1}^{\infty}$ $\displaystyle=\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\tau_{t}^{-1/ \alpha}-\frac{\tau_{t}^{-2}(p+1)^{-2\alpha+1}}{2\alpha-1}-1.$ $$

Let’s now utilize the upper and lower bounds for $\tau_{t}^{-1}$ from Proposition 7. Substituting, we have

$$ $\displaystyle n\Omega$ $\displaystyle\geq\dfrac{\pi(\alpha-1)}{\alpha^{2}\sin{(\pi/\alpha})}\dfrac{n \alpha\sin{(\pi/\alpha)}}{\pi}-\frac{\tau_{t}^{-2}(p+1)^{-2\alpha+1}}{2\alpha- 1}-1$ $\displaystyle=\dfrac{n(\alpha-1)}{\alpha}-\frac{\tau_{t}^{-2}(p+1)^{-2\alpha+1 }}{2\alpha-1}-1$ $\displaystyle\geq\dfrac{n(\alpha-1)}{\alpha}-\left(\dfrac{\left(n+1+\frac{p+1} {\alpha-1}\right)\alpha\sin{(\pi/\alpha})}{\pi(p+1)}\right)^{2\alpha}\frac{p+1 }{2\alpha-1}-1$ $$

Dividing both left and right-hand side by $n$ gives:

$$ $\displaystyle\implies\Omega$ $\displaystyle>\dfrac{\alpha-1}{\alpha}-\left(\dfrac{\left(n+1+\frac{p+1}{ \alpha-1}\right)\alpha\sin{(\pi/\alpha})}{\pi(p+1)}\right)^{2\alpha}\frac{p+1} {n(2\alpha-1)}-\frac{1}{n}$ $\displaystyle>\dfrac{\alpha-1}{\alpha}-\frac{1}{2\alpha-1}\frac{n+1+\frac{p+1} {\alpha-1}}{n}\left(\dfrac{n+1+\frac{p+1}{\alpha-1}}{p+1}\right)^{2\alpha-1}- \frac{1}{n},$ $\displaystyle>\dfrac{\alpha-1}{\alpha}-\frac{1}{\alpha}\left(\dfrac{n+1+\frac{ p+1}{\alpha-1}}{p+1}\right)^{2\alpha-1}-\frac{1}{n},$ $$

since $\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi}<1$ for $\alpha>1$ and $n+1+\frac{p+1}{\alpha-1}<n\frac{2\alpha-1}{\alpha}$ by our assumption on $n$.
∎

###### Proposition 9 .

Under the assumption that

$$ $\displaystyle\max\left(2\alpha,p\frac{\alpha}{(\alpha-1)^{2}}+\dfrac{\alpha^{2 }}{(\alpha-1)^{2}}\right)<n<\min\left((p+1)\dfrac{\alpha-2}{\alpha},p\left( \dfrac{3+2^{-\alpha}}{4+2^{-(\alpha-2)}}\right),p\dfrac{\pi\left(\sqrt{\frac{2 \alpha}{5}}-1\right)^{1/\alpha}}{\alpha\sin{(\pi/\alpha)}}-\dfrac{p+1}{\alpha- 1}\right)-1$ $$

and $4<\alpha$, we can find a masked surrogate-to-target setting that improves over the risk of the standard target model by selecting all features $i$ such that $\zeta_{i}^{2}>1-\Omega$.

###### Proof.

From Proposition 8, we have

$$ $\displaystyle\Omega>\dfrac{\alpha-1}{\alpha}-\frac{1}{\alpha}\left(\dfrac{n+1+ \frac{p+1}{\alpha-1}}{p+1}\right)^{2\alpha-1}-\frac{1}{n}.$ $$

It’s then enough to show that we can find a set of $i$’s such that

$$ $\displaystyle\zeta_{i}^{2}>\dfrac{1}{\alpha}+\dfrac{1}{\alpha}\left(\dfrac{n+1 +\frac{p+1}{\alpha-1}}{p+1}\right)^{2\alpha-1}+\dfrac{1}{n}.$ $$

From the proof of Proposition 4, we know that

$$ $\displaystyle\zeta_{i}^{2}>c^{\prime}\iff$ $\displaystyle i>\tau_{t}^{-1/\alpha}\left(\frac{\sqrt{c^{\prime}}}{1-\sqrt{c^{ \prime}}}\right)^{1/\alpha}.$ $$

Hence, using the bound on $\tau_{t}^{-1}$ from Proposition 7, it’s enough to find indices $i$ such that

$$ $\displaystyle i>\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi}\left(n+1+\frac{p+1}{ \alpha-1}\right)\left(\frac{\sqrt{c^{\prime}}}{1-\sqrt{c^{\prime}}}\right)^{1/ \alpha}\quad\text{where }c^{\prime}=\dfrac{1}{\alpha}+\dfrac{1}{\alpha}\left( \dfrac{n+1+\frac{p+1}{\alpha-1}}{p+1}\right)^{2\alpha-1}+\dfrac{1}{n}.$ (21) $$

By our assumption $p+1>n+1+\frac{p+1}{\alpha-1}$ and $n>2\alpha$, we obtain that $\frac{5}{2\alpha}>c^{\prime}$. Since $\left(\frac{\sqrt{x}}{1-\sqrt{x}}\right)^{1/\alpha}$ is increasing with $x$ when $0\leq x\leq 1$, we have

$$ $\displaystyle\left(\frac{1}{\sqrt{\frac{2\alpha}{5}}-1}\right)^{1/\alpha}\geq \left(\frac{\sqrt{c^{\prime}}}{1-\sqrt{c^{\prime}}}\right)^{1/\alpha}.$ $$

Then, to ensure the existence of an interval of $i$’s satisfying the above inequality, we choose

$$ $\displaystyle p-(p+1)\dfrac{\alpha\sin{(\pi/\alpha)}}{\pi(\alpha-1)}\left( \frac{1}{\sqrt{\frac{2\alpha}{5}}-1}\right)^{1/\alpha}\geq(n+1)\dfrac{\alpha \sin{(\pi/\alpha)}}{\pi}\left(\frac{1}{\sqrt{\frac{2\alpha}{5}}-1}\right)^{1/\alpha}$ $\displaystyle\iff$ $\displaystyle p\frac{\pi\left(\sqrt{\frac{2\alpha}{5}}-1\right)^{1/\alpha}}{ \alpha\sin{(\pi/\alpha)}}-\frac{p+1}{\alpha-1}\geq n+1$ $$

which follows from our assumption on $n$. Thus, discarding the features $i$ provided in the interval (21) will strictly improve the test risk of the masked surrogate-to-target model over the standard target model.
∎

One can verify that our assumptions are coherent because they ensure a non-empty interval for $n$ when $\alpha>4$ as the coefficients of $p$ are positive and $\dfrac{\alpha}{(\alpha-1)^{2}}$ is smaller compared to the other three coefficients of $p$ in the minimum function.

###### Proof for the proportional asymptotic case.

###### Proof.

###### Proof.

###### Proposition 6 (Scaling law for masked surrogate-to-target model) .

###### Proof.

###### Proof.

###### Proposition 7 (Non-asymptotic analysis of τ 𝜏 \tau italic_τ ) .

###### Proof.

###### Proposition 8 (Non-asymptotic analysis of Ω Ω \Omega roman_Ω ) .

###### Proof.

###### Proposition 9 .

###### Proof.

## Appendix C Proofs for Section 5

###### Theorem 3 (Distributional characterization, Han & Xu ( 2023 ) ) .

Let $\kappa_{s}=p/m>1$ and suppose that, for some $M>1$, $1/M\leq\kappa_{s},{\sigma}_{s}^{2}\leq M$ and $\left\|{\bm{\Sigma}_{s}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{s}^{-1}}
\right\|_{\text{op}}\leq M$. Let $\tau_{s}\in\mathbb{R}$ be the unique solution of the following equation:

$$ $\displaystyle\kappa_{s}^{-1}=\frac{1}{p}\texttt{tr}\left((\bm{\Sigma}_{s}+\tau _{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}\right).$ (22) $$

We define the random variable $X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}(\bm{\Sigma}_{s},{\bm{\beta}_{\star}},\bm{g
}_{s})$ based on $\bm{g}_{s}\sim\mathcal{N}(0,{\bm{I}})$ and the function $\gamma_{s}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ as follows:

$$ $\begin{split}X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}(\bm{\Sigma}_{s},{\bm{\beta}_{ \star}},\bm{g}_{s})&:=(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s} \left[{\bm{\beta}_{\star}}+\frac{\bm{\Sigma}_{s}^{-1/2}\gamma_{s}({\bm{\beta}_ {\star}})\bm{g}_{s}}{\sqrt{p}}\right]\\ \gamma_{s}^{2}({\bm{\beta}_{\star}})&:=\kappa_{s}\left({\sigma}_{s}^{2}+ \operatorname{\mathbb{E}}_{\bm{g}_{s}}[\|{\bm{\Sigma}_{s}^{1/2}(X_{\kappa_{s}, {\sigma}_{s}^{2}}^{s}(\bm{\Sigma}_{s},{\bm{\beta}_{\star}},\bm{g}_{s})-{\bm{ \beta}_{\star}})}\|_{2}^{2}]\right).\end{split}$ (23) $$

Then, for any $L$-Lipschitz function $f:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ where $L<L(M)$, there exists a constant $C=C(M)$ such that for any $\varepsilon\in(0,1/2]$, we have the following:

$$ $\displaystyle\operatorname{\mathbb{P}}(\sup_{{\bm{\beta}_{\star}}\in{\bm{B}}(R )}\left|f({\bm{\beta}}^{s})-\operatorname{\mathbb{E}}_{\bm{g}_{s}}[f(X_{\kappa _{s},{\sigma}_{s}^{2}}^{w}(\bm{\Sigma}_{s},{\bm{\beta}_{\star}},\bm{g}_{s}))] \right|\geq\varepsilon)\leq Cp\mathrm{e}^{-p\varepsilon^{4}/C},$ (24) $$

where $R<M$.

See 3

See 2

###### Proof.

Define a function $f_{1}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ as $f_{1}(\bm{x})=\|{\bm{\Sigma}_{t}^{1/2}(\bm{x}-{\bm{\beta}_{\star}})}\|_{2}^{2}$. The gradient of this function is

$$ $\displaystyle\|{\nabla f_{1}(\bm{x})}\|_{2}=\|{2\bm{\Sigma}_{t}(\bm{x}-{\bm{ \beta}_{\star}})}\|_{2}\leq 2\left\|{\bm{\Sigma}_{t}}\right\|_{\text{op}}\|{ \bm{x}-{\bm{\beta}_{\star}}}\|_{2}.$ $$

Using Proposition 11, there exists an event $E$ with $\operatorname{\mathbb{P}}(E^{c})\leq C_{t}e^{-p/C_{t}}$ where $C_{t}=C_{t}(M_{t},\frac{M_{t}-R}{2})$ with the definition of $M_{t}$ in Proposition 11, such that $f_{1}({\bm{\beta}}^{s2t})$ is $2M_{t}^{2}$-Lipschitz if ${\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)$. Applying Theorem 3 on the target model, there exists a constant $\bar{C}_{s}=\bar{C}_{s}(M_{t})$ such that for any $\varepsilon\in(0,1/2]$, we obtain

$$ $\displaystyle\sup_{{\bm{\beta}}^{s}\in{\bm{B}}(\frac{M_{t}+R}{2})} \operatorname{\mathbb{P}}\left(\left|f({\bm{\beta}}^{s2t})-\operatorname{ \mathbb{E}}_{\bm{g}_{t}}[f(X_{\kappa_{t},{\sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t} ,{\bm{\beta}}^{s},\bm{g}_{t}))]\right|\geq\varepsilon\right)\leq Cp\mathrm{e}^ {-p\varepsilon^{4}/C},$ (25) $$

where $f({\bm{\beta}}^{s2t})=\mathcal{R}({\bm{\beta}}^{s2t})$ and

$$ $\displaystyle X_{\kappa_{t},{\sigma}_{t}^{2}}^{t}(\bm{\Sigma}_{t},{\bm{\beta}} ^{s},\bm{g}_{t})=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}\left[{ \bm{\beta}}^{s}+\frac{\bm{\Sigma}_{t}^{-1/2}\gamma_{t}({\bm{\beta}}^{s})\bm{g} _{t}}{\sqrt{p}}\right].$ $$

Furthermore,

$$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[f(X_{\kappa_{t},{ \sigma}_{t}^{2}}^{s}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},\bm{g}_{t}))\right]$ $\displaystyle=\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[\|{\bm{\Sigma}_{t}^{ 1/2}\left({\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})-({\bm{I}}-{ \bm{\theta}}_{1}){\bm{\beta}_{\star}}+{\bm{\theta}}_{2}\gamma_{t}({\bm{\beta}} ^{s})\right)}\|_{2}^{2}\right]$ $\displaystyle=({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta}}_{1}^ {\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})+ \gamma_{t}^{2}({\bm{\beta}}^{s})\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{ \theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}]$ $\displaystyle+{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top} \bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}-2{\bm{\beta}_{ \star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_ {1}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}}),$ (26) $$

where ${\bm{\theta}}_{1}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}$ and ${\bm{\theta}}_{2}:=(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2
}\frac{\bm{g}_{t}}{\sqrt{p}}$. Let $E(M_{t},\frac{M_{t}-R}{2})$ be the event defined in Proposition 10. Let $f_{2}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ be defined as $f_{2}(\bm{x}):=(\bm{x}-{\bm{\beta}_{\star}})^{\top}{\bm{\theta}}_{1}^{\top}\bm
{\Sigma}_{t}{\bm{\theta}}_{1}(\bm{x}-{\bm{\beta}_{\star}})$. By Proposition 12, the function $f_{2}$ is $2M_{t}^{2}-$Lipschitz if ${\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)$ on the event $E(M_{t},\frac{M_{t}-R}{2})$. Applying Theorem 3 on the surrogate model, there exists a constant $\bar{C}_{w,1}=\bar{C}_{w,1}(M_{t})$ such that for any $\varepsilon\in(0,1/2]$, we obtain

$$ $\displaystyle\sup_{{\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)}\operatorname{ \mathbb{P}}\left(\left|f_{2}({\bm{\beta}}^{s})-{\bm{\beta}_{\star}}^{\top}({ \bm{I}}-{\bm{\Phi}}_{1})^{\top}{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{ \theta}}_{1}({\bm{I}}-{\bm{\Phi}}_{1}){\bm{\beta}_{\star}}-\gamma_{s}^{2}({\bm {\beta}_{\star}})\operatorname{\mathbb{E}}_{\bm{g}_{s}}[{\bm{\Phi}}_{2}^{\top} {\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}{\bm{\Phi}}_{2}]\right |>\varepsilon\right)\leq\bar{C}_{w,1}pe^{-p\varepsilon^{4}/\bar{C}_{w,1}},$ (27) $$

where ${\bm{\Phi}}_{1}:=(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}$ and ${\bm{\Phi}}_{2}:=(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}^{1/2}
\frac{\bm{g}_{s}}{\sqrt{p}}$.

Let $f_{3}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ be defined as $f_{3}(\bm{x}):=\gamma_{t}^{2}(\bm{x}){\bm{\theta}}_{2}^{\top}\bm{\Sigma}_{t}{
\bm{\theta}}_{2}$. By Proposition 13 and Proposition 2.1 in Han & Xu (), the function $f_{3}$ is $4M_{t}^{2}$-Lipschitz if ${\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)$ on the event $E(M_{t},\frac{M_{t}-R}{2})$. Applying Theorem 3 on the surrogate model, there exists a constant $\bar{C}_{w,2}=\bar{C}_{w,2}(M_{t})$ such that for any $\varepsilon\in(0,1/2]$, we obtain

$$ $\displaystyle\sup_{{\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)}\operatorname{ \mathbb{P}}\left(\left|f_{3}({\bm{\beta}}^{s})-\operatorname{\mathbb{E}}_{{\bm {\beta}}^{s}\sim X^{s}}[\gamma_{t}^{2}({\bm{\beta}}^{s})]\operatorname{\mathbb {E}}_{\bm{g}_{t}}[{\bm{\theta}}_{2}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{2}] \right|>\varepsilon\right)\leq\bar{C}_{w,2}pe^{-p\varepsilon^{4}/\bar{C}_{w,2}}.$ (28) $$

Let $f_{4}:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ as $f_{4}(\bm{x}):=-2{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top
}\bm{\Sigma}_{t}{\bm{\theta}}_{1}(\bm{x}-{\bm{\beta}_{\star}})$. By Proposition 14 and Proposition 2.1 in Han & Xu (), the function $f_{4}$ is $2M_{t}^{2}-$Lipschitz if ${\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)$ on the event $E(M_{t},\frac{M_{t}-R}{2})$. Applying Theorem 3 on the surrogate model, there exists a constant $\bar{C}_{w,3}=\bar{C}_{w,3}(M_{t})$ such that for any $\varepsilon\in(0,1/2]$, we obtain

$$ $\displaystyle\sup_{{\bm{\beta}_{\star}}\in{\bm{B}}_{p}(R)}\operatorname{ \mathbb{P}}\left(\left|f_{4}({\bm{\beta}}^{s})-2\left[{\bm{\beta}_{\star}}^{ \top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{ \Phi}}_{1}-{\bm{I}}){\bm{\beta}_{\star}}\right]\right|>\varepsilon\right)\leq \bar{C}_{w,3}pe^{-p\varepsilon^{4}/\bar{C}_{w,3}}.$ (29) $$

By the definition of these functions, we have

$$ $\displaystyle\operatorname{\mathbb{E}}_{\bm{g}_{t}}\left[f(X_{\kappa_{t},{ \sigma}_{t}^{2}}^{s}(\bm{\Sigma}_{t},{\bm{\beta}}^{s},\bm{g}_{t}))\right]-{\bm {\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}({\bm {I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}=f_{2}({\bm{\beta}}^{s})+f_{3}({\bm {\beta}}^{s})-f_{4}({\bm{\beta}}^{s})$ (30) $$

By the definition of ${\bm{\theta}}_{1}$, ${\bm{\theta}}_{2}$, ${\bm{\Phi}}_{1}$, and ${\bm{\Phi}}_{2}$, we have

$$ $\begin{split}\bar{\mathcal{R}}_{\dot{\kappa},\dot{\sigma}}(\dot{\bm{\Sigma}},{ \bm{\beta}_{\star}})-{\bm{\beta}_{\star}}^{\top}({\bm{I}}-{\bm{\theta}}_{1})^{ \top}\bm{\Sigma}_{t}({\bm{I}}-{\bm{\theta}}_{1}){\bm{\beta}_{\star}}&={\bm{ \beta}_{\star}}^{\top}({\bm{I}}-{\bm{\Phi}}_{1})^{\top}{\bm{\theta}}_{1}^{\top }\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{I}}-{\bm{\Phi}}_{1}){\bm{\beta}_{\star} }+\gamma_{s}^{2}({\bm{\beta}_{\star}})\operatorname{\mathbb{E}}_{\bm{g}_{s}}[{ \bm{\Phi}}_{2}^{\top}{\bm{\theta}}_{1}^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}{ \bm{\Phi}}_{2}]\\ &+\operatorname{\mathbb{E}}_{{\bm{\beta}}^{s}\sim X^{s}}[\gamma_{t}^{2}({\bm{ \beta}}^{s})]\operatorname{\mathbb{E}}_{\bm{g}_{t}}[{\bm{\theta}}_{2}^{\top} \bm{\Sigma}_{t}{\bm{\theta}}_{2}]-2\left[{\bm{\beta}_{\star}}^{\top}({\bm{I}}- {\bm{\theta}}_{1})^{\top}\bm{\Sigma}_{t}{\bm{\theta}}_{1}({\bm{\Phi}}_{1}-{\bm {I}}){\bm{\beta}_{\star}}\right].\end{split}$ (31) $$

Using (30)-(31) and applying a union bound on (25), (27), (28), and (29), we obtain the advertised claim.
∎

###### Proposition 10 .

Suppose that, for some $M_{t}>1$, $1/M_{t}\leq\kappa_{s},{\sigma}_{s}^{2}\leq M_{t}$ and $\left\|{\bm{\Sigma}_{s}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{s}^{-1}}
\right\|_{\text{op}}\leq M_{t}$. For every $c_{s}>0$, there exists an event $E(M_{t},c_{s})$ with $\operatorname{\mathbb{P}}((E(M_{t},c_{s}))^{c})\leq C_{s}\mathrm{e}^{-p/C_{s}}$ where $C_{s}=C_{s}(M_{t},c_{s})$ such that

$$ $\displaystyle\|{{\bm{\beta}}^{s}}\|_{2}\leq\|{{\bm{\beta}_{\star}}}\|_{2}+c_{s }\quad\text{and}\quad\|{{\bm{\beta}}^{s}-{\bm{\beta}_{\star}}}\|_{2}\leq\|{{ \bm{\beta}_{\star}}}\|_{2}+c_{s}.$ $$

###### Proof.

By the definition of ${\bm{\beta}}^{s}$, we have

$$ $\displaystyle{\bm{\beta}}^{s}$ $\displaystyle=\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1} \tilde{\bm{y}}$ $\displaystyle=\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1} \tilde{\bm{X}}{\bm{\beta}_{\star}}+\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{ \bm{X}}^{\top})^{-1}{\tilde{\bm{z}}},$ (32) $$

where ${\tilde{\bm{z}}}\sim\mathcal{N}(\bm{0},{\sigma}_{s}^{2}{\bm{I}})$. By triangle inequality, we obtain

$$ $\displaystyle\|{{\bm{\beta}}^{s}}\|_{2}$ $\displaystyle\leq\|{\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top}) ^{-1}\tilde{\bm{X}}{\bm{\beta}_{\star}}}\|_{2}+\|{\tilde{\bm{X}}^{\top}(\tilde {\bm{X}}\tilde{\bm{X}}^{\top})^{-1}{\tilde{\bm{z}}}}\|_{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\|{{\bm{\beta}_{\star}}}\|_{2 }+\|{\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}{\tilde{ \bm{z}}}}\|_{2},$ (33) $$

where $(a)$ in above follows from the fact that $\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}\tilde{\bm{X}}$ is a projection matrix, and so all of its eigenvalues are either 0 0 or $1$. Focusing on the second term of the RHS, we derive

$$ $\displaystyle\|{\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1 }{\tilde{\bm{z}}}}\|_{2}^{2}={\tilde{\bm{z}}}^{\top}(\tilde{\bm{X}}\tilde{\bm{ X}}^{\top})^{-1}{\tilde{\bm{z}}}$ $\displaystyle=\frac{{\tilde{\bm{z}}}^{\top}}{\sqrt{p}}\left(\frac{\tilde{\bm{X }}\tilde{\bm{X}}^{\top}}{p}\right)^{-1}\frac{{\tilde{\bm{z}}}}{\sqrt{p}}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{{\tilde{\bm{z}}}^{\top} {\tilde{\bm{z}}}}{p}\left\|{\left(\frac{\tilde{\bm{X}}\tilde{\bm{X}}^{\top}}{p }\right)^{-1}}\right\|_{\text{op}},$ (34) $$

where $(a)$ in the above inequality follows from Cauchy-Schwarz inequality. Using Bernstein’s inequality, there exists an absolute constant $C_{0}>0$ that depends on ${\sigma}_{s}^{2}$ such that

$$ $\displaystyle\operatorname{\mathbb{P}}\left(\frac{{\tilde{\bm{z}}}^{\top}{ \tilde{\bm{z}}}}{p}-{\sigma}_{s}^{2}>t\right)\leq\text{exp}\left\{{-c\min\left \{\frac{pt^{2}}{4C_{0}^{2}},\frac{pt}{2C_{0}}\right\}}\right\}.$ $$

On the other hand, let $\tilde{\bm{Z}}=\tilde{\bm{X}}\bm{\Sigma}_{s}^{-1/2}$, which means that the entries of $\tilde{\bm{Z}}$ are independent and normally distributed with zero mean and unit variance. Then,

$$ $\displaystyle\left\|{\left(\frac{\tilde{\bm{X}}\tilde{\bm{X}}^{\top}}{p}\right )^{-1}}\right\|_{\text{op}}=\left\|{\left(\frac{\tilde{\bm{Z}}\bm{\Sigma}_{s} \tilde{\bm{Z}}^{\top}}{p}\right)^{-1}}\right\|_{\text{op}}\leq\left\|{\bm{ \Sigma}_{s}^{-1}}\right\|_{\text{op}}\left\|{\left(\frac{\tilde{\bm{Z}}\tilde{ \bm{Z}}^{\top}}{p}\right)^{-1}}\right\|_{\text{op}}.$ (35) $$

Using Theorem 1.1 in Rudelson & Vershynin (), there exist absolute constants $C_{1},C_{2}>0$ such that we have the following for every $\varepsilon>0$

$$ $\displaystyle\operatorname{\mathbb{P}}\left(\left\|{\left(\frac{\tilde{\bm{Z}} \tilde{\bm{Z}}^{\top}}{p}\right)^{-1}}\right\|_{\text{op}}\leq\varepsilon^{2} \left(1-\frac{1}{\kappa_{s}}\right)^{2}\right)\leq(C_{1}\varepsilon)^{p-m+1}+e ^{-pC_{2}}.$ (36) $$

By combining (34), (35), and (36), we obtain that

$$ $\displaystyle\operatorname{\mathbb{P}}\Big{(}\|{{\bm{\beta}}^{s}}\|_{2}\leq\|{ {\bm{\beta}_{\star}}}\|_{2}+\varepsilon(1-\frac{1}{\kappa_{s}})$ $\displaystyle\sqrt{(t+{\sigma}_{s}^{2})\left\|{\bm{\Sigma}_{s}^{-1}}\right\|_{ \text{op}}}\Big{)}$ $\displaystyle\leq(C_{1}\varepsilon)^{p-m+1}+e^{-pC_{2}}+e^{-c\min\left\{\frac{ pt^{2}}{4C_{0}^{2}},\frac{pt}{2C_{0}}\right\}}$ $$

The advertised claim for $\|{{\bm{\beta}}^{s}}\|_{2}$ follows when $\varepsilon$ is selected as $\varepsilon<\frac{1}{C_{1}e}$. For $\|{{\bm{\beta}}^{s}-{\bm{\beta}_{\star}}}\|_{2}$, using the definition of ${\bm{\beta}}^{s}$, we write as follows:

$$ $\displaystyle\|{{\bm{\beta}}^{s}-{\bm{\beta}_{\star}}}\|_{2}$ $\displaystyle=\|{\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{- 1}\tilde{\bm{X}}{\bm{\beta}_{\star}}+\tilde{\bm{X}}^{\top}(\tilde{\bm{X}} \tilde{\bm{X}}^{\top})^{-1}{\tilde{\bm{z}}}-{\bm{\beta}_{\star}}}\|_{2}$ $\displaystyle\leq\|{\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top}) ^{-1}\tilde{\bm{X}}-{\bm{I}}}\|_{2}\|{{\bm{\beta}_{\star}}}\|_{2}+\|{\tilde{ \bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}{\tilde{\bm{z}}}}\|_{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\|{{\bm{\beta}_{\star}}}\|_{2 }+\|{\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}{\tilde{ \bm{z}}}}\|_{2},$ (37) $$

where $(a)$ in the above inequalities follows from the fact that the eigenvalues of $\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}\tilde{\bm{X}}-
{\bm{I}}$ are either $1$ or 0 0 as the eigenvalues of $\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}\tilde{\bm{X}}$ are either $1$ or 0 0. The remaining part of this proof is identical to the previous part.
∎

###### Corollary 2 .

Suppose that ${\bm{\beta}}^{s}\in\mathbb{R}^{p}$ is given, and for some $M_{t}>1$, we have $1/M_{t}\leq\kappa_{t},{\sigma}_{t}^{2}\leq M_{t}$ and $\left\|{\bm{\Sigma}_{t}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{t}^{-1}}
\right\|_{\text{op}}\leq M_{t}$. For every $c_{t}>0$, there exists an event $E(M_{t},c_{t})$ with $\operatorname{\mathbb{P}}((E(M_{t},c_{t}))^{c})\leq C_{t}\mathrm{e}^{-p/C_{t}}$ where $C_{t}=C_{t}(M_{t},c_{t})$ such that

$$ $\displaystyle\|{{\bm{\beta}}^{s2t}}\|_{2}\leq\|{{\bm{\beta}}^{s}}\|_{2}+c_{t} \quad\text{and}\quad\|{{\bm{\beta}}^{s2t}-{\bm{\beta}}^{s}}\|_{2}\leq\|{{\bm{ \beta}}^{s}}\|_{2}+c_{t}.$ $$

###### Proof.

The result directly follows from the proof of Proposition 10.
∎

###### Proposition 11 .

Suppose that, for some $M_{t}>1$, $1/M_{t}\leq\kappa_{t},{\sigma}_{t}^{2}\leq M_{t}$ and $\left\|{\bm{\Sigma}_{t}}\right\|_{\text{op}},\left\|{\bm{\Sigma}_{t}^{-1}}
\right\|_{\text{op}}\leq M_{t}$. For every $c_{t}>0$, there exists an event $E(M_{t},c_{t})$ with $\operatorname{\mathbb{P}}((E(M_{t},c_{t}))^{c})\leq C_{t}\mathrm{e}^{-p/C_{t}}$ where $C_{t}=C_{t}(M_{t},c_{t})$ such that we have the following on this event $E(M_{t},c_{t})$:

$$ $\displaystyle\|{{\bm{\beta}}^{s2t}}\|_{2}\leq\|{{\bm{\beta}_{\star}}}\|_{2}+c_ {t}\quad\text{and}\quad\|{{\bm{\beta}}^{s2t}-{\bm{\beta}_{\star}}}\|_{2}\leq\| {{\bm{\beta}_{\star}}}\|_{2}+c_{t}$ $$

###### Proof.

By the definition of ${\bm{\beta}}^{s2t}$, we have the following:

$$ $\displaystyle{\bm{\beta}}^{s2t}={\bm{X}}({\bm{X}}{\bm{X}}^{\top})^{-1}{\bm{X}} {\bm{\beta}}^{s}+{\bm{X}}^{\top}({\bm{X}}{\bm{X}}^{\top})^{-1}{\bm{z}}$ (38) $$

where ${\bm{z}}\sim\mathcal{N}(\mathbf{0},{\sigma}_{t}^{2}{\bm{I}})$. Plugging (32) into (38), we obtain

$$ $\displaystyle{\bm{\beta}}^{s2t}={\bm{X}}({\bm{X}}{\bm{X}}^{\top})^{-1}{\bm{X}} \left(\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}\tilde{ \bm{X}}{\bm{\beta}_{\star}}+\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}} ^{\top})^{-1}{\tilde{\bm{z}}}\right)+{\bm{X}}^{\top}({\bm{X}}{\bm{X}}^{\top})^ {-1}{\bm{z}}$ (39) $$

Note that ${\bm{X}}({\bm{X}}{\bm{X}}^{\top})^{-1}{\bm{X}}$ and $\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}\tilde{\bm{X}}$ are projection matrices. Multiplication of two projection matrices results in a projection matrix. Using the fact that the eigenvalues of a projection matrix are either $1$ or 0 0 in (39), we have

$$ $\displaystyle\|{{\bm{\beta}}^{s2t}}\|_{2}\leq\|{{\bm{\beta}_{\star}}}\|_{2}+\| {\tilde{\bm{X}}^{\top}(\tilde{\bm{X}}\tilde{\bm{X}}^{\top})^{-1}{\tilde{\bm{z} }}}\|_{2}+\|{{\bm{X}}^{\top}({\bm{X}}{\bm{X}}^{\top})^{-1}{\bm{z}}}\|_{2}$ (40) $$

By a similar reasoning used in (34),(35), and (36); there exist absolute constants $C_{0},C_{1},C_{2},c>0$ such that we have the following for every $\varepsilon,t>0$:

$$ $\displaystyle\operatorname{\mathbb{P}}\Big{(}\|{{\bm{X}}^{\top}({\bm{X}}{\bm{X }}^{\top})^{-1}{\bm{z}}}\|_{2}\leq\varepsilon(1-\frac{1}{\kappa_{t}})$ $\displaystyle\sqrt{(t+{\sigma}_{t}^{2})\left\|{\bm{\Sigma}_{t}^{-1}}\right\|_{ \text{op}}}\Big{)}$ $\displaystyle\leq(C_{1}\varepsilon)^{p-n+1}+e^{-pC_{2}}+e^{-c\min\left\{\frac{ pt^{2}}{4C_{0}^{2}},\frac{pt}{2C_{0}}\right\}}$ (41) $$

Similarly, for every $\tilde{\varepsilon}>0$, there exist absolute constants $\tilde{C}_{0},\tilde{C}_{1},\tilde{C}_{2},\tilde{c}>0$ such that we have the following for every $\tilde{\varepsilon},\tilde{t}$:

$$ $\displaystyle\operatorname{\mathbb{P}}\Big{(}\|{\tilde{\bm{X}}^{\top}(\tilde{ \bm{X}}\tilde{\bm{X}}^{\top})^{-1}{\tilde{z}}}\|_{2}\leq\tilde{\varepsilon}(1- \frac{1}{\kappa_{s}})$ $\displaystyle\sqrt{(\tilde{t}+{\sigma}_{s}^{2})\left\|{\bm{\Sigma}_{s}^{-1}} \right\|_{\text{op}}}\Big{)}$ $\displaystyle\leq(\tilde{C}_{1}\tilde{\varepsilon})^{p-m+1}+e^{-p\tilde{C}_{2} }+e^{-\tilde{c}\min\left\{\frac{pt^{2}}{4\tilde{C}_{0}^{2}},\frac{p\tilde{t}}{ 2\tilde{C}_{0}}\right\}}$ (42) $$

Note that ${\bm{X}},{\bm{z}},\tilde{\bm{X}},$ and ${\tilde{z}}$ are independent of each other. Therefore, we can apply union bound on (41) and (42) with selecting $\varepsilon,t,\tilde{\varepsilon}$, and $\tilde{t}$ such that $\varepsilon<\frac{1}{C_{1}e}$, $\frac{c_{t}}{2}<\varepsilon(1-\frac{1}{\kappa_{t}})\sqrt{(t+{\sigma}_{t}^{2})
\left\|{\bm{\Sigma}_{t}^{-1}}\right\|_{\text{op}}}$, $\tilde{\varepsilon}<\frac{1}{\tilde{C}_{1}}$, and $\frac{c_{t}}{2}<\varepsilon(1-\frac{1}{\kappa_{s}})\sqrt{(\tilde{t}+{\sigma}_{
s}^{2})\left\|{\bm{\Sigma}_{s}^{-1}}\right\|_{\text{op}}}$. As a result, there exists an event $E$ with $\operatorname{\mathbb{P}}(E^{c})\leq C_{t}(M_{t},c_{t})$ such that

$$ $\displaystyle\|{{\bm{\beta}}^{s2t}}\|_{2}\leq\|{{\bm{\beta}_{\star}}}\|_{2}+c_ {t}.$ $$

Using a similar argument in (37), we derive the following on the same event $E_{1}$

$$ $\displaystyle\|{{\bm{\beta}}^{s2t}-{\bm{\beta}_{\star}}}\|_{2}\leq\|{{\bm{ \beta}_{\star}}}\|_{2}+c_{t}.$ $$

This completes the proof.

∎

###### Proposition 12 .

Let $g:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ be a function such that

$$ $\displaystyle g({\bm{\beta}}^{s}):=\|{\bm{\Sigma}_{t}^{1/2}(\bm{\Sigma}_{t}+ \tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}({\bm{\beta}}^{s}-{\bm{\beta}_{\star}})} \|_{2}^{2}$ $$

Then, on the same event $E(M_{t},c_{s})$ in Proposition 10, the function $g$ is $(\|{{\bm{\beta}_{\star}}}\|_{2}+c_{s})\frac{2\lambda_{1}^{3}}{(\lambda_{1}+
\tau_{t})^{2}}-$Lipschitz where $\lambda_{1}$ is the largest eigenvalue of $\bm{\Sigma}_{t}$.

###### Proof.

We take the gradient of the function $g$:

$$ $\displaystyle\|{\nabla g({\bm{\beta}}^{s})}\|_{2}$ $\displaystyle=2\|{\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{ \Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}({\bm{\beta}} ^{s}-{\bm{\beta}_{\star}})}\|_{2}$ $\displaystyle\leq\|{{\bm{\beta}}^{s}-{\bm{\beta}_{\star}}}\|_{2}\max_{i}\frac{ 2\lambda_{i}^{3}}{(\lambda_{i}+\tau_{t})^{2}}$ $\displaystyle=\|{{\bm{\beta}}^{s}-{\bm{\beta}_{\star}}}\|_{2}\max_{i}2\lambda_ {i}\left(1-\frac{\tau_{t}}{\lambda_{i}+\tau_{t}}\right)^{2}$ $\displaystyle=\|{{\bm{\beta}}^{s}-{\bm{\beta}_{\star}}}\|_{2}\frac{2\lambda_{1 }^{3}}{(\lambda_{1}+\tau_{t})^{2}}.$ $$

Combining Proposition 10 on the event $E(M_{t},c_{s})$ with the above inequality provides the advertised claim.
∎

###### Proposition 13 .

Let $g:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ be a function such that

$$ $\displaystyle g({\bm{\beta}}^{s}):=\frac{1}{p}{\|\bm{\Sigma}_{t}^{1/2}(\bm{ \Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}\gamma_{t}({\bm{\beta}} ^{s})\|_{F}^{2}}$ $$

Then, on the same event $E(M_{t},c_{s})$ in Proposition 10, the function $g$ is $L-$Lipschitz where $(\lambda_{i})_{i=1}^{p}$ are the eigenvalues of $\bm{\Sigma}_{t}$ with a descending order and

$$ $\displaystyle L=\frac{4\tau_{t}^{2}}{m}\frac{\lambda_{1}^{3}}{(\lambda_{1}+ \tau_{t})^{4}}\frac{\|{{\bm{\beta}_{\star}}}\|_{2}+c_{s}}{1-\frac{1}{m}\sum_{i =1}^{p}\left(\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}}\right)^{2}}.$ $$

###### Proof.

We take the gradient of the function $g$:

$$ $\displaystyle\nabla g({\bm{\beta}}^{s})=\frac{2}{p}\bm{\Sigma}_{t}^{1/2}(\bm{ \Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm {I}})^{-1}\bm{\Sigma}_{t}^{1/2}\nabla\gamma_{t}^{2}({\bm{\beta}}^{s}).$ $$

Note that

$$ $\displaystyle\gamma_{t}^{2}({\bm{\beta}}^{s})$ $\displaystyle=\kappa_{s}\left({\sigma}_{s}^{2}+\operatorname{\mathbb{E}}_{\bm{ g}_{s}}[\|{\bm{\Sigma}_{s}^{1/2}(X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}(\bm{ \Sigma}_{s},{\bm{\beta}_{\star}},\bm{g}_{s})-{\bm{\beta}_{\star}})}\|_{2}^{2}]\right)$ $\displaystyle=\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}\|{(\bm{\Sigma}_{t} +\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}{\bm{\beta}}^{s}}\|_{2}^{2}}{1- \frac{1}{m}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{\Sigma} _{t}^{2}\right)}.$ $$

Then, we have

$$ $\displaystyle\nabla\gamma_{t}^{2}({\bm{\beta}}^{s})=2\kappa_{t}\frac{\tau_{t}^ {2}\bm{\Sigma}_{t}^{1/2}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{\Sigma}_{t} ^{1/2}{\bm{\beta}}^{s}}{1-\frac{1}{m}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t }{\bm{I}})^{-2}\bm{\Sigma}_{t}^{2}\right)}.$ $$

Plugging $\nabla\gamma_{t}^{2}({\bm{\beta}}^{s})$ into $\nabla g({\bm{\beta}}^{s})$, we obtain that

$$ $\displaystyle\|{\nabla g({\bm{\beta}}^{s})}\|_{2}$ $\displaystyle=\frac{4\tau_{t}^{2}}{m}\frac{\bm{\Sigma}_{t}^{1/2}(\bm{\Sigma}_{ t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2 }\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}{ \bm{\beta}}^{s}}{1-\frac{1}{m}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I }})^{-2}\bm{\Sigma}_{t}^{2}\right)}$ $\displaystyle\leq\frac{4\tau_{t}^{2}}{m}\frac{\lambda_{1}^{3}}{(\lambda_{1}+ \tau_{t})^{4}}\frac{\|{{\bm{\beta}}^{s}}\|_{2}}{1-\frac{1}{m}\sum_{i=1}^{p} \left(\frac{\lambda_{i}}{\lambda_{i}+\tau_{t}}\right)^{2}}.$ $$

Combining Proposition 10 on the event $E(M_{t},c_{s})$ with the above inequality provides the advertised claim.
∎

###### Proposition 14 .

Let $g:\mathbb{R}^{p}\xrightarrow[]{}\mathbb{R}$ be a function such that

$$ $\displaystyle g({\bm{\beta}}^{s}):=2{\bm{\beta}_{\star}}^{\top}\left({\bm{I}}- (\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}\right)^{\top}\bm{\Sigma }_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}({\bm{\beta}}^{s}-{ \bm{\beta}_{\star}}).$ $$

Then, the function $g$ is $2\|{{\bm{\beta}_{\star}}}\|_{2}\tau_{t}\left(\frac{\lambda_{1}}{\lambda_{1}+
\tau_{t}}\right)^{2}-$Lipschitz where $\lambda_{1}$ is the largest eigenvalue of $\bm{\Sigma}_{t}$.

###### Proof.

We take the gradient of the function $g$:

$$ $\displaystyle\|{\nabla g({\bm{\beta}}^{s})}\|_{2}$ $\displaystyle=2\|{\bm{\Sigma}_{t}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{ \Sigma}_{t}\left({\bm{I}}-(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{ t}\right){\bm{\beta}_{\star}}}\|_{2}$ $\displaystyle\leq 2\|{{\bm{\beta}_{\star}}}\|_{2}\tau_{t}\max_{i}\left(1-\frac {\tau_{t}}{\lambda_{i}+\tau_{t}}\right)^{2}$ $\displaystyle=2\|{{\bm{\beta}_{\star}}}\|_{2}\tau_{t}\left(\frac{\lambda_{1}}{ \lambda_{1}+\tau_{t}}\right)^{2},$ $$

and the desired result readily follows.
∎

###### Lemma 1 .

We have that

$$ $\displaystyle\operatorname{\mathbb{E}}_{{\bm{\beta}}^{s}\sim X_{\kappa_{s},{ \sigma}_{s}^{2}}^{s}}[\gamma_{t}^{2}({\bm{\beta}}^{s})]$ $\displaystyle=\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}\|{(\bm{\Sigma}_{t} +\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}((\bm{\Sigma}_{s}+\tau_{s}{\bm{I}} )^{-1}\bm{\Sigma}_{s}{\bm{\beta}_{\star}}}\|_{2}^{2}}{1-\frac{1}{n}\texttt{tr} \left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{\Sigma}_{t}^{2}\right)}$ $\displaystyle+\frac{\kappa_{t}\tau_{t}^{2}\gamma_{s}^{2}({\bm{\beta}_{\star}}) }{p}\frac{\texttt{tr}\left(\bm{\Sigma}_{s}^{1/2}(\bm{\Sigma}_{s}+\tau_{s}{\bm{ I}})^{-1}\bm{\Sigma}_{t}^{1/2}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{ \Sigma}_{t}^{1/2}(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}^{1/2} \right)}{1-\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2} \bm{\Sigma}_{t}^{2}\right)}.$ $$

###### Proof.

The desired claim follows from the following manipulations using the definition of $X_{\kappa_{s},{\sigma}_{s}^{2}}^{s}$ in (3):

$$ $\displaystyle\operatorname{\mathbb{E}}_{{\bm{\beta}}^{s}\sim X_{\kappa_{s},{ \sigma}_{s}^{2}}^{s}}[\gamma_{t}^{2}({\bm{\beta}}^{s})]$ $\displaystyle=\operatorname{\mathbb{E}}_{{\bm{\beta}}^{s}\sim X_{\kappa_{s},{ \sigma}_{s}^{2}}^{s}}\left[\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}\|{( \bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}{\bm{\beta}}^{s}}\| _{2}^{2}}{1-\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2 }\bm{\Sigma}_{t}^{2}\right)}\right]$ $\displaystyle=\operatorname{\mathbb{E}}_{\bm{g}_{s}}\left[\kappa_{t}\frac{{ \sigma}_{t}^{2}+\tau_{t}^{2}\|{(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-1}\bm{ \Sigma}_{t}^{1/2}\left((\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}{ \bm{\beta}_{\star}}+(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}^{1/ 2}\gamma_{s}({\bm{\beta}_{\star}})\bm{g}_{s}/\sqrt{p}\right)}\|_{2}^{2}}{1- \frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{\Sigma} _{t}^{2}\right)}\right]$ $\displaystyle=\kappa_{t}\frac{{\sigma}_{t}^{2}+\tau_{t}^{2}\|{(\bm{\Sigma}_{t} +\tau_{t}{\bm{I}})^{-1}\bm{\Sigma}_{t}^{1/2}((\bm{\Sigma}_{s}+\tau_{s}{\bm{I}} )^{-1}\bm{\Sigma}_{s}{\bm{\beta}_{\star}}}\|_{2}^{2}}{1-\frac{1}{n}\texttt{tr} \left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{\Sigma}_{t}^{2}\right)}$ $\displaystyle+\frac{\kappa_{t}\tau_{t}^{2}\gamma_{s}^{2}({\bm{\beta}_{\star}}) }{p}\frac{\texttt{tr}\left(\bm{\Sigma}_{s}^{1/2}(\bm{\Sigma}_{s}+\tau_{s}{\bm{ I}})^{-1}\bm{\Sigma}_{t}^{1/2}(\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2}\bm{ \Sigma}_{t}^{1/2}(\bm{\Sigma}_{s}+\tau_{s}{\bm{I}})^{-1}\bm{\Sigma}_{s}^{1/2} \right)}{1-\frac{1}{n}\texttt{tr}\left((\bm{\Sigma}_{t}+\tau_{t}{\bm{I}})^{-2} \bm{\Sigma}_{t}^{2}\right)}.$ $$

∎

###### Theorem 3 (Distributional characterization, Han & Xu ( 2023 ) ) .

###### Proof.

###### Proposition 10 .

###### Proof.

###### Corollary 2 .

###### Proof.

###### Proposition 11 .

###### Proof.

###### Proposition 12 .

###### Proof.

###### Proposition 13 .

###### Proof.

###### Proposition 14 .

###### Proof.

###### Lemma 1 .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: In addition to the points above that should be clarified, I have one following question: how is the CIFAR-10 experiment performed? Notably, how are the surrogate and target distributions generated? This is worth expanding in the paper.
soundness: 4
strengths: This paper is rather well-written and contains quite a few interesting theoretical insights. The results draw clear delineations on how knowledge distillation through a surrogate model can help. As someone who hasn't thought about overparameterized linear regression in a while, Proposition 1 and Corollary 1 were rather surprising results, demonstrating that for a dataset size proportional (but smaller) than the number of parameters, the optimal surrogate predictor to use for generating pseudolabels is actually not the ground truth predictor, and that there is (in theory) always room for benefit as long as the covariance is non-isotropic, which implies that a learner benefits from using something other than the actual distribution of labels. 

In addition to the theory, the numerical results on CIFAR-10 also counterintuitively support that a learner only trained on surrogate pseudolabels on the target domain actually outperform the surrogate model itself, which has access to true labels (albeit with a different covariate distribution...?).
summary: In this paper, the authors propose a precise characterization of the benefits of knowledge distillation, mostly in the context of gaussian linear regression. In particular, the main set-up considers the excess risk of linear regression on a given distribution, but with the learner only able to access pseudolabels generated by a surrogate model instead of the true labels. Notably, the authors show that under a covariance-shift model (i.e. the distribution of covariates $x$ may change between the surrogate and target stages, but the underlying predictor $\beta_\star$ remains the same in between), then the optimal surrogate predictor minimizing the (asymptotic) excess risk on the target distribution is a weighted version of the ground-truth predictor $\beta_\star$, which which amplifies entries corresponding to large eigenvalues of the (diagonal) covariance matrix above a certain threshold, and shrinks entries below a threshold. Furthermore, in a masked setting, where the surrogate model is restricted to selecting a subset of the full set of features, then similarly the optimal surrogate predictor selects predictor entries above a certain threshold of covariance eigenvalues. Lastly, the authors show that in a certain asymptotic regime, an optimal surrogate-to-target model (i.e. a model trained on target distribution covariates with surrogate model pseudolabels) has the same excess risk as the least-squares target model trained with the true labels, demonstrating that knowledge distillation in a sense cannot beat out ERM with access to true labels.
weaknesses: Though the theoretical results are interesting, there are a few aspects that are worth clarifying. Notably, even though it is demonstrated that there can exist a surrogate model that induces better risk on the target model than using the true labels, in general a surrogate model is typically not trained with foreknowledge of the task distribution. I believe this is what Section 5 is trying to convey, but it is not clear to me after reading that section how to interpret the result therein. In particular, it should be explained how this relates to, for example, a target model that is trained using strong labels to demonstrate the marginal gain (or loss).

In general, the paper accrues a lot of jargon and notation; it would be very helpful to either create a table containing the different set-ups/regimes considered and the summary conclusion of the relative gain/suboptimality of knowledge distillation and/or a notation table that summarizes what the various risks denote. This would help clarify how to place the scaling law (Proposition 5) and Section 5 with respect to the results of the prior sections.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Can you provide intuition why the risk of the surrogate-to-target model under the optimal selection of the parameters scales the same as that of the target model (even though there is a strict improvement in the risk)? I am wondering why improvement is possible. Is it because, for example, if the tail components of the covariance is zero, then features on these components are essentially useless, therefore an surrogate that omits those components will be better? 

2. your equation (8) involves $\beta^{s2t}$. Does it mean that your asymptotic risk estimate (9) also involves $\beta^{s2t}$ and thus can not be directly computed? I think in the final bound $\beta^{s2t}$ should not appear; otherwise I can just claim the definition of the excess risk of $\beta^{s2t}$ is already an exact characterization of itself.

3. In observation1, you assume jointly diagonalizability. Is there fundamental hardness to remove this assumption?
soundness: 2
strengths: The authors provide comprehensive theoretical results for weak-to-strong generalization, giving a exact characterization of the excess risk of the weak-to-strong estimator. This knowledge distillation problem is important in modern machine learning, indicating the significance of this work.
summary: This paper studies the problem of knowledge distillation under linear regression. In the first stage, data are collected from a surrogate model. In the second stage, a target model is trained using the data generated in the first stage. The authors characterize the non-asymptotic excess risk of the target model under "model shift" setting and "distribution shift" setting. Numerical results are provided, justifying their theory on ridgeless regression and on neural network architectures.
weaknesses: The presentation of this paper is in general not very satisfactory, in the sense that this paper is lack of necessary intuition and explanation. For example, how to interpret the non-asymptotic bounds and what does each term stand for? Why it is possible that weak-to-strong estimator is even better than purely using the strong model?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Can you give me some insights to extend the theory to neural networks (even if two-layer neural network)? I think the authors should also discuss this in the refined version.
soundness: 3
strengths: 1. Knowledge distillation and weak-to-strong generalization are significant topics today, and their theory is very poor. Therefore, this is a meaningful paper for me.
2. The theory is complete and well-written.
3. The derived bounds seem tight because they are matched with empirical results.
summary: The paper provides a sharp characterization for knowledge distillation in the high-dimensional regression setting, including both model shift and distribution shift, cases. Concretely, the paper characterizes the precise risk of the target model in both cases through non-asymptotic bounds in terms of sample size and data distribution under mild conditions. As a consequence, the paper identifies the form of the optimal surrogate model, which reveals the benefits and limitations of such processes. Finally, the paper validates the results by numerical experiments.
weaknesses: 1. The theory only focuses the high-dimensional linear regression setting, which is well-studied in the literature. Besides, the results can not be extended to neural networks directly.
2. A typo in line 134.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: * I would love to see Figure 1b and Figure 2b in LogLog Scale. One of the main points of the authors in Section 4 (Proposition 5) concerns the learning rates in the high-dimensional limit. It would be nice to see them in Figure 1b.
* Is it possible to generalise the result of Section 3.1 to the case where the chosen features are chosen with a matrix $A$ which has a non zero kernel? The masking seems a specific case of this.


* There is a broken citation on page 3.
* On line 334 is it "omniscent test risk estimate"?
soundness: 3
strengths: This paper is mathematically sound and considers an interesting setting. The main strengths are
* The derivations of the different transition values for the covariates is sharp and to my knowledge is a novel finding.
* The theory match with simulations also work at finite dimension even if the result is high-dimensional.
* The model introduced and studied is expressive enough to show different behaviour that characterise performances.
summary: This paper considers the problem of knowledge distillation under the setting of linear ridge-less regression in the teacher student scenario.

The setting considered in this paper is the the proportional regime where $p,n\to\infty$ and their ratio $\kappa_t = p/n$ is kept fixed. They consider $\kappa_t > 1$ in the overparametrised regime.

The models considered in the paper are three:
* _The Surrogate-to-Target model_ where the data is generated from a dataset $\mathcal{D}$ with input $x\in\mathbb{R}^d$ and output $y = x^\top \beta_\star + z$ with $ \beta_\star$ a teacher vector. This data is used to estimate a min norm estimator called $\beta^s$ and generate a second data set $y^s = x^\top \beta^s + z$ and the final estimation is done as $\beta^{s2t}$ from $(x,y^s)$.
* _The Standard Target model_ where the model is evaluated on the generated data $(x, y)$
* _The Covariance Shift model_ here the dataset is generated with a certain choice of covariance and then the population risk evaluated on a different covariance model.

The first part of the paper is devoted to finding the performance conditioned on a specific teacher while the second to last section considers the full _Surrogate-to-Target_ setup.
The authors also consider the procedure of Masking for the surrogate model. In this case the surrogate model has been trained on a masked version of the data and the new labels are generated from the original inputs and the labels of the surrogate model.

The main technical results presented in the main are the characterisation of the population risk for the model conditioned on $\beta^s$ and then for the _Surrogate-to-Target_ model.

For the case conditioned on the target the authors are able to precisely derive the effect of the surrogate model on the final student, showing specific conditions (depending on the covariates and $\beta^\star$) under which a $\beta^{s2t}$ performs better than a _The Standard Target model_. The same is true for the masking.
weaknesses: * Right now the mathematical result is introduced in generality without explaining the idea behind the proof. The authors could briefly explain that to derive the results one should apply the theory from [Han&Xu2023] that relies on the use of of the convex gordon min max theorem.
* The authors provide some numerical simulations on ResNet-50 on a CIFAR10 classification showing a result that qualitatively differs from the theory. Either this limitation is explained in detail or I don't think it is necessary to be shown.
* Is there any reason why the authors consider in their technical results the ridgeless estimator instead of the ridge one? A long series of works (e.g. [Hastie2020, Louriero2022]) considers general loss and provides similar bounds.
* _(Minor)_ Section 4 is presented unclearly. The settings for the propositions are not well explained and need to be introduced more clearly.

[Hastie 2020] Surprises in High-Dimensional Ridgeless Least Squares Interpolation. Annals of Stats 2020

[Loureiro2022] Learning curves of generic features maps for realistic datasets with a teacher-student model. Neurips 2022

Meta-review: The work presents new results on knowledge distillation, with focus on the linear Gaussian setting. The work characterizes excess risk of such linear regression based on pseudo labels  from a surrogate model, based on covariate shift. The work presents a clear set of results, demonstrating how using a surrogate model can help with knowledge distillation. 

All reviewers found the results novel and interesting, and the presentation to be mostly clear. The clarity and novelty of the results can form the basis of future work on the topic.

Decision: Accept (Spotlight)