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

M. Emrullah Ildiz^∗ ¹ Halil Alperen Gozeten^∗ ¹ Ege Onur Taga¹
Marco Mondelli^† ² Samet Oymak^† ¹ ¹ University of Michigan, Ann Arbor
{eildiz,alperen,egetaga,oymak}@umich.edu
² Institute of Science and Technology Austria
marco.mondelli@ist.ac.at

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.

^*^*footnotetext: Equal Contribution^$\dagger$^$\dagger$footnotetext: Equal Advising

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., 2015). Variations of this approach include generating synthetic data from powerful language models (Wang et al., 2023; Gunasekar et al., 2023; Abdin et al., 2024), weak-to-strong generalization to obtain stronger models under weak supervision (Burns et al., 2023), and filtering/curating ML datasets via a smaller model to train a larger model (Fang et al., 2023; Lin et al., 2024b). 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 denote the target distribution, be sampled i.i.d. from this distribution, and denote the dimension of the features . Given a surrogate model , we create the synthetic labels and obtain the target model by minimizing the empirical risk, i.e.,

(1)

where denotes the hypothesis class and the loss function. This procedure motivates a few fundamental questions regarding (1): (i) What is the excess risk of compared to that of the minimizer of the population risk ? (ii) What is the optimal surrogate model that minimizes such excess risk? (iii) Can the optimal strictly outperform using true labels or setting in (1)? Furthermore, in practice, itself is the outcome of an empirical risk minimization (ERM) procedure. Specifically, let be the surrogate distribution, be sampled i.i.d. from (having feature dimension ), and assume

(2)

where denotes the surrogate hypothesis class. Thus, as a final (and more challenging question), one may ask: (iv) How do the sample sizes and the data distributions affect the performance of ?

(a) Ground-truth and surrogate model weights

(b) Test risks as a function of sample size

Figure 1: Structure and performance of optimal surrogate models. (a): We compare the weights of the optimal surrogate model (green) with the ground-truth (blue). This reveals a transition from amplification to shrinkage as we move from principal to tail eigenvalues. The yellow curve displays the optimal 0-1 masking of the ground-truth where we either keep or discard a feature. (b): Associated test risks as a function of sample size. The theoretical bounds (full lines) match the experiments (markers). Setting: The feature size is ; the sample size is in (a) and variable in (b); the feature covariance follows the power-law structure , ; is the covariance statistics (see Corollary 1) governing the optimal surrogate’s structure.

Main contributions. We address these questions in the context of high-dimensional ridgeless regression, where are the class of linear models and is the quadratic loss. We provide a sharp non-asymptotic characterization of the test risk of in two core settings:

1.
The surrogate is provided and we solve via (1), which addresses the first three questions above;
2.
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 and the feature dimension are all proportional and also allow for a distribution shift between and , which correspond to the two ERM stages. Our theoretical guarantees in Section 3 precisely characterize the regression coefficients of the optimal surrogate model in terms of the corresponding coefficients of the population risk minimizer and the feature covariance. This is depicted in Figure 1(a) where and 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 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 , that summarize the role of feature covariance and finite sample size in the test risk. There is a well-defined transition point, in Fig. 1(a), where the gain passes from strict amplification () to strict shrinkage () 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 is restricted to be a 0-1 mask on the entries of so that the surrogate has the direct interpretation of feature pruning. We show that beyond the transition point , 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 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 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., 2020) or double descent (Belkin et al., 2019). Specifically, asymptotic and non-asymptotic risk characterizations for the minimum -norm interpolator (i.e., ridgeless regression with ) have been provided in a recent line of work (Hastie et al., 2020; Cheng & Montanari, 2024; Han & Xu, 2023; Wu & Xu, 2020; Richards et al., 2021; Loureiro et al., 2022).

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., 2021; Montanari et al., 2023; Han & Xu, 2023) utilize Gaussian process theory (Thrampoulidis et al., 2015) 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., 2024; Song et al., 2024; Yang et al., 2023; Mallinar et al., 2024). 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. (2020) provide a theoretical analysis of self-distillation, whereas Menon et al. (2020); Nagarajan et al. (2023); Harutyunyan et al. (2023) 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., 2022; Oymak & Gulcu, 2021; Wei et al., 2022b). 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. (2024) investigates 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. (2023) consider the problem of surrogate-based data selection. Finally, Charikar et al. (2024); Lang et al. (2024) aims 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., 2017; Kaplan et al., 2020). 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., 2020; Simon et al., 2023; Bahri et al., 2024; Paquette et al., 2024; Bordelon et al., 2024b; Lin et al., 2024a; Maloney et al., 2022). Bordelon et al. (2024a) analyze a random feature model trained with gradient descent via dynamical mean field theory. Jain et al. (2024) consider scaling laws with surrogate data, whereas Sorscher et al. (2022) study the benefits of data pruning.

2 Problem setup

Notation. Let denote the set for an integer . We use lower-case and upper-case bold letters (e.g., ) to represent vectors and matrices, respectively; denotes the -th entry of the vector , the pseudo-inverse of the matrix , and the trace of . For further notations, please see Table 1.

We consider a two-stage linear learning problem. In the first stage, pairs of labels and input features from the distribution 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 . 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 following the linear model , where , and is independent of . Let be the dataset for the surrogate model drawn i.i.d. from . We analyze both under- and over-parametrized settings: in the former, we estimate by minimizing the quadratic loss; in the latter, we estimate as the minimum norm interpolator. As a result, the estimator of the surrogate model can be written as follows:

(3)

where and .

Stage 2: Target model. Given , we consider another data distribution following the linear model , where and . Let be the dataset for the target model drawn i.i.d. from . As for the surrogate model, the estimator for the target model is defined as

(4)

where and . Our analysis will generally apply to an arbitrary choice and will not require it to be the outcome of (3). Finally, we define the excess (population) risk for a given estimator as

(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 with respect to the standard target model, which has access to the ground-truth parameter through labeling. Specifically, consider the dataset drawn i.i.d. from ; then, the estimation is

(6)

where and . We compare the excess risks of the surrogate-to-target model with that of the standard target model .

Reference 2: Covariance shift model (Mallinar et al., 2024; Patil et al., 2024). Given , let be a dataset drawn i.i.d. from , where , , and ; then, using the same notation and , the estimation is . The test risk of is calculated under covariance shift. Let be a distribution such that , , and . Then, the excess transfer risk is

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 . First, we provide a non-asymptotic bound on the risk conditioned on , and then we optimize this quantity with respect to (finding also a closed-form expression of the corresponding optimal value of ). In addition, we build a connection between our surrogate-to-target model and knowledge distillation (Hinton et al., 2015), as well as weak-to-strong generalization (Burns et al., 2023). Finally, we extend our analysis to ridge regression in Appendix A.1.

We note that when , the calculations in this section simplify to the non-asymptotic risk characterization of the minimum -norm interpolator, recently studied by Hastie et al. (2020); Cheng & Montanari (2024); Han & Xu (2023). When , 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., 2024). Formally, given any and any jointly diagonalizable covariance matrices , there exists a unique such that the risk of the surrogate-to-target problem with is equivalent to the risk of the covariance shift model with .

The joint diagonalizability of and is also required by Mallinar et al. (2024) (see their Assumption 2.1), where upper and lower bounds for the bias and variance are provided by adapting results from (Bartlett et al., 2020; Tsigler & Bartlett, 2022). In contrast, our approach utilizes the non-asymptotic characterization of the -norm interpolator by Han & Xu (2023). This allows us to directly characterize the non-asymptotic risk (instead of giving upper and lower bounds, as in (Mallinar et al., 2024)) and, thus, to obtain the optimal surrogate parameter . We begin by defining the asymptotic risk of the surrogate-to-target model, given the surrogate parameter , in the proportional regime where and the ratio is kept fixed.

Definition 1.

Let and be the unique solution of the following equation

(7)

Let and where . Now, define the asymptotic characterization of the minimum -norm interpolator and the function as the following fixed point equation based on the asymptotic risk

				(8)
				(9)

where the asymptotic risk is defined as

(10)

In the asymptotic risk, the term corresponds to a part of the bias risk caused by the model shift (), and the implicit regularization term where the eigenvalues of are less than 1. The term corresponds to the remaining part of the bias and variance risks. We now state our non-asymptotic characterization of the risk.

Theorem 1.

Suppose that, for some constant , we have and . Recall from (5) that represents the risk of the surrogate-to-target model given . Then, there exists a constant such that, for any , the following holds with :

(11)

The proof of Theorem 1 utilizes the non-asymptotic characterization of the minimum norm interpolator in Han & Xu (2023), which is based on the convex Gaussian min-max theorem (Gordon, 1988; Thrampoulidis et al., 2015). The proof is deferred to Appendix A.

The surrogate parameter that minimizes the individual term in (10) is . On the other hand, the surrogate parameter that minimizes the individual term in (10) is the zero vector, which follows from (16) in Appendix A. Now, we are going to jointly minimize the asymptotic risk in the next proposition. The optimal surrogate parameter is visualized as the green curve in Figure 1.

Proposition 1.

Let . The optimal surrogate minimizing the asymptotic risk in (10) is

Note that, under the setting of Theorem 1, Proposition 1 can be extended to the non-asymptotic risk by applying (11). The corollary below then offers a direct interpretation of the optimal surrogate parameter . The proofs of both Corollary 1 and Proposition 1 are in Appendix A.

Corollary 1.

Without loss of generality, suppose that is diagonal.^{††\dagger†}^{††\dagger†}If not, there exists an orthogonal matrix s.t. is diagonal. Then, we can consider the covariance matrix as and the ground truth parameter as , which behaves the same as the original parameters, see Observation 2. Let be the eigenvalues of in non-increasing order and let for . Then, the following results hold:

1.
for every .
2.
if and only if for every .
3.
if and only if the covariance matrix for some .

The first part shows that the optimal surrogate parameter is fully characterized by , which only depends on the covariance spectrum (via ) and the sample size (via ). Note that the spectrum characterizes the risk in both linear regression and random features regression, shown in Ildiz et al. (2024) and Simon et al. (2023), respectively. As the eigenvalues are ordered, the ’s are ordered as well, and the second part of the corollary identifies a threshold behavior: before the transition point , the entries of the surrogate are amplified w.r.t. the ground-truth parameter , 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 , and it arises from the optimization of the trade-off between the bias and variance terms in (10). 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.

The intuition behind improving the performance of the standard target model by utilizing a surrogate parameter different from is associated with the implicit regularization of the minimum norm interpolator in the over-parametrized region . As long as the covariance matrix eigenvalues are not constant, there is a way to mitigate the bias risk caused by the implicit regularization. This implicit regularization term is specific to the over-parametrized region. Indeed, in the next proposition (proved in Appendix A), we will show that the optimal surrogate parameter is when the target model is under-parametrized:

Proposition 2.

The optimal surrogate parameter that minimizes the asymptotic risk in the under-parametrized region () is equivalent to the ground truth parameter . In other words, for any , the surrogate-to-target model cannot outperform the standard target model in the asymptotic risk.

In other words, the result above shows that the improvement in the surrogate-to-target model compared to the standard target model is special to the over-parameterized region.

3.1 Weak-to-strong generalization

To connect with knowledge distillation (Hinton et al., 2015) and weak-to-strong generalization (Burns et al., 2023), we allow the surrogate model to use fewer features, , by introducing a mask operation , where selects features from the full set of features in . Alongside this mask, we adjust the distributions for both the surrogate and target models as

Then, is estimated based on the samples from , and the risk is still calculated with respect to the standard target model distribution as defined in (5). As we focus on analyzing the model shift case, we assume that the covariance matrices and 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 ()? 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 3.

Consider the target model in (6), assume that is diagonal, and recall the definitions of and . Then, the following results hold:

1.
If the mask operation selects all the features that satisfy , then the surrogate-to-target model outperforms the standard target model in the asymptotic risk in (10).
2.
Let represent the set of all possible , where . The optimal for the asymptotic risk in (10) within is the one that selects all features satisfying .

The proof of Proposition 3 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 ) after the transition point , while they coincide with the ground-truth parameter otherwise. The transition point changes with respect to Corollary 1 and, as , 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 .

In Figure 1(a), we illustrate the optimal and , showing that the threshold associated with the optimal is larger than the threshold associated with the transition from amplification to shrinkage in the optimal . 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 . 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., 2015) using both ground-truth labels and predictions from a surrogate (weak) model on the CIFAR-10 dataset (Krizhevsky & Hinton, 2009). 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. (2023), where the GPT-4 model trained with GPT-2 labels performs comparably to GPT-3.5. The reason why the surrogate-to-target model underperforms the standard target model is the surrogate model is not able to follow the feature selection mechanism characterized in Proposition 3. This suggests that the feature selection mechanism is crucial for surpassing the performance of the standard target model. We provide further experimental details in Appendix A.2.

4 Fundamental limits and scaling laws

We now study the fundamental limits of the surrogate-to-target model with the optimal surrogate parameter (see Proposition 1) and the optimal mask operator (see Proposition 3). Our analysis shows that, when eigenvalues () and signal coefficients () follow a power law, the risk of the surrogate-to-target model under the optimal selection of the parameters and 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. (2024), does not change the scaling law. We start our analysis with the definition of the omniscient test risk estimate.

(a) Weak-to-strong on CIFAR-10

(b) Comparison of theoretical and experimental risks

Figure 2: (a): On the CIFAR-10 dataset, we fine-tune a ResNet50 model using the ground-truth labels (target) and the predictions of three weak convolutional models (surrogate) with different capacities: big (b), medium (m), and small (s). We observe that surrogate-to-target models consistently outperform surrogate models’ accuracies, even though they are trained on the surrogate models’ predictions. (b): We compare the experimental two-stage risk with our estimated theoretical risk. In the experimental setup, , and we vary from 1 to 100. Both feature covariances follow the power-law structure for and ; the ground truth parameter is specified as .

Definition 2 (Omniscient test risk estimate).

Fix . Given a covariance , , and the noise term , set and define as the unique non-negative solution of Then, the omniscient excess test risk estimate is the following:

				(12)
	where

The above test risk estimate yields exact results (and, hence, in (12) becomes ) in the proportional limit via the analysis of Section 3. In other words, the omniscient test risk estimate is identical to , which can be derived by substituting into Equation 10. Specifically, suppose that the empirical distributions of and converge as having fixed the ratio . Then, the risk obtained in Theorem 1 converges to the omniscient risk estimate given in (12), as proved in Appendix B. We will use this omniscient risk estimate in the limit of , as considered in several papers (Cui et al., 2022; Simon et al., 2024; Wei et al., 2022a). Yet, our empirical validations in Figure 1 demonstrate that this framework yields consistent results even when applied to scenarios with moderately sized and .

Throughout the section, we analyze the case where the surrogate parameter is given, therefore we need to take into account only the target covariance matrix . Without loss of generality, we assume that the covariance matrix is diagonal by Observation 2. From now on, we will consider the particular case of power-law eigenstructure, that is . The omniscient risk under this structure still depends on the parameters and , and the next proposition analyzes them asymptotically. Its proof is in Appendix B.

Proposition 4 (Asymptotic analysis of and ).

Let the covariance matrix be diagonal and for . Recall from Definition 2 that, as , and are given by the equations

Then, the following results hold

(13)

Recall from Corollary 1 and Proposition 3 that the cut-off indices for the optimal surrogate parameter and the optimal mask operation are respectively and , which depend on . Armed with the asymptotic expressions in (13), we now identify the cut-off indices as a function of the sample size .

Proposition 5.

Set the constants and and assume the power-law eigenstructure for . Let and be the solutions given by Proposition 4 and define . Then, the indices for which are ; while the indices for which is are .

The result above is proved in Appendix B, and it shows that, as the sample size decreases, the cut-off indices of both the optimal surrogate parameter and the optimal mask shift to the left linearly in . 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 , when the surrogate parameter is the optimal . Specifically, Proposition 6 below shows that, under a power-law decay of both the eigenvalues () and the signal coefficients (), 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 to be at most of the order of the scaling law of the excess test risks when . This corresponds to , where is the exponent of the scaling law in the noiseless setting. Conversely, a fixed noise variance that does not decay with would cause the noise to overshadow the uncaptured part of the signal, which scales down with . In this unintended scenario, the noise would dominate our observations.

Proposition 6 (Scaling law).

Let the covariance matrix be diagonal with eigenvalues , and let the ground-truth parameter have components corresponding to each feature. Assume that both eigenvalues and signal coefficients follow a power-law decay, i.e., and for , . Let the optimal surrogate parameter be given by Proposition 1 and define the minimum surrogate-to-target risk attained by as where is described in Definition 2. Then, in the limit of , 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

Since is a lower bound on , 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 . This also indicates that the optimal selection of the mask does not improve the scaling law, see Proposition 7 in Appendix B for details.

The proof of Proposition 6 is deferred to Appendix B. Here, we note that we utilize the expression 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 , it becomes particularly useful when considering the optimal surrogate parameter 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 and (see Propositions 8 and 9, respectively), which complements the asymptotic one in Proposition 4 above. This allows us to characterize a region with finite and where the surrogate-to-target model strictly outperforms the standard target model, see Proposition 10.

5 Risk characterization for the two-stage model

Until now, we have examined the behavior of the surrogate-to-target model when is given. In this section, we characterize the non-asymptotic risk of the surrogate-to-target model when is the solution of the surrogate problem (3) where . Our analysis includes two cases: (i) the target model has infinitely many data (), and (ii) the target model is overparametrized, i.e., .

When the target model has infinitely many data, the estimate of the surrogate-to-target model is equal to the estimate of the surrogate model . This means that the correct ground-truth parameter is estimated under a distribution and tested under another distribution , 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 and in Theorem 1. Let and define similarly to . We define the random variable based on and the function as follows:

Let , , and . Then, we define the asymptotic risk estimate as

The non-asymptotic characterization of the risk is stated below and proved in Appendix C, which also contains a closed-form expression for (see Lemma 1).

Theorem 2.

Suppose that, for some constant , we have and . Consider the surrogate-to-target model defined in Section 2, and let represent its risk when is given. Recall the definition of and in Definition 3. Then, there exists a constant such that for any , the following holds when :

In the proof of Theorem 2, we apply the distributional characterization of the minimum norm interpolator (Han & Xu, 2023) twice. The main technical difficulty is to satisfy the Lipschitz condition of the distributional characterization in the second step, which is handled by bounding the intermediate step’s interpolator in a ball.

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). Furthermore, we extend our analysis of the two-stage model to the under-parametrized region in Appendix C.1. This extension enables us to analyze Section 3.1 thoroughly because this analysis allows the surrogate model to have finite data.

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 three 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 pair during the training of the target model. The third is to go beyond linear regression towards neural network models. In this regard, the precise asymptotics of the test error of the ERM solution were provided by (Mei & Montanari, 2022) for the random features model and by (Montanari & Zhong, 2022) for two-layer neural networks in the NTK regime. However, a non-asymptotic characterization (similar to that given by (Han & Xu, 2023) for linear regression) remains an open problem. The resolution of this open problem, as well as the analysis of the phenomena of knowledge distillation and weak-to-strong generalization, represent exciting future directions.

Acknowledgements

M.E.I., H.A.G., 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², 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.

References

Abdin et al. (2024) Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, et al. Phi-3 technical report: A highly capable language model locally on your phone, 2024. URL https://arxiv.org/abs/2404.14219.
Bahri et al. (2024) Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27):e2311878121, 2024.
Bartlett et al. (2020) Peter L. Bartlett, Philip M. Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117(48):30063–30070, 2020.
Belkin et al. (2019) Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical bias–variance trade-off. Proceedings of the National Academy of Sciences, 116(32):15849–15854, 2019.
Bordelon et al. (2024a) Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws, 2024a. URL https://arxiv.org/abs/2402.01092.
Bordelon et al. (2024b) Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How feature learning can improve neural scaling laws, 2024b. URL https://arxiv.org/abs/2409.17858.
Burns et al. (2023) Collin Burns, Pavel Izmailov, Jan Hendrik Kirchner, Bowen Baker, Leo Gao, Leopold Aschenbrenner, Yining Chen, Adrien Ecoffet, Manas Joglekar, Jan Leike, Ilya Sutskever, and Jeff Wu. Weak-to-strong generalization: Eliciting strong capabilities with weak supervision, 2023. URL https://arxiv.org/abs/2312.09390.
Chang et al. (2021) Xiangyu Chang, Yingcong Li, Samet Oymak, and Christos Thrampoulidis. Provable benefits of overparameterization in model compression: From double descent to pruning neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 6974–6983, 2021.
Charikar et al. (2024) Moses Charikar, Chirag Pabbaraju, and Kirankumar Shiragur. Quantifying the gain in weak-to-strong generalization, 2024. URL https://arxiv.org/abs/2405.15116.
Cheng & Montanari (2024) Chen Cheng and Andrea Montanari. Dimension free ridge regression, 2024. URL https://arxiv.org/abs/2210.08571.
Cui et al. (2022) Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Generalization error rates in kernel regression: the crossover from the noiseless to noisy regime. Journal of Statistical Mechanics: Theory and Experiment, 2022(11):114004, 2022.
Fang et al. (2023) Alex Fang, Albin Madappally Jose, Amit Jain, Ludwig Schmidt, Alexander Toshev, and Vaishaal Shankar. Data filtering networks, 2023. URL https://arxiv.org/abs/2309.17425.
Frei et al. (2022) Spencer Frei, Difan Zou, Zixiang Chen, and Quanquan Gu. Self-training converts weak learners to strong learners in mixture models. In International Conference on Artificial Intelligence and Statistics, pp. 8003–8021. PMLR, 2022.
Gordon (1988) Y. Gordon. On milman’s inequality and random subspaces which escape through a mesh in rn. In Geometric Aspects of Functional Analysis, pp. 84–106, Berlin, Heidelberg, 1988. Springer Berlin Heidelberg.
Gunasekar et al. (2023) Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harkirat Singh Behl, Xin Wang, Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need, 2023. URL https://arxiv.org/abs/2306.11644.
Han & Xu (2023) Qiyang Han and Xiaocong Xu. The distribution of ridgeless least squares interpolators, 2023. URL https://arxiv.org/abs/2307.02044.
Harutyunyan et al. (2023) Hrayr Harutyunyan, Ankit Singh Rawat, Aditya Krishna Menon, Seungyeon Kim, and Sanjiv Kumar. Supervision complexity and its role in knowledge distillation, 2023. URL https://arxiv.org/abs/2301.12245.
Hastie et al. (2020) Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J. Tibshirani. Surprises in high-dimensional ridgeless least squares interpolation, 2020. URL https://arxiv.org/abs/1903.08560.
He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL https://arxiv.org/abs/1512.03385.
Hestness et al. (2017) Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md. Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically, 2017. URL https://arxiv.org/abs/1712.00409.
Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015. URL https://arxiv.org/abs/1503.02531.
Ildiz et al. (2024) Muhammed E. Ildiz, Zhe Zhao, and Samet Oymak. Understanding inverse scaling and emergence in multitask representation learning. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 4726–4734. PMLR, 02–04 May 2024. URL https://proceedings.mlr.press/v238/ildiz24a.html.
Jain et al. (2024) Ayush Jain, Andrea Montanari, and Eren Sasoglu. Scaling laws for learning with real and surrogate data, 2024. URL https://arxiv.org/abs/2402.04376.
Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URL https://arxiv.org/abs/2001.08361.
Kolossov et al. (2023) Germain Kolossov, Andrea Montanari, and Pulkit Tandon. Towards a statistical theory of data selection under weak supervision, 2023. URL https://arxiv.org/abs/2309.14563.
Krizhevsky & Hinton (2009) Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical Report 0, University of Toronto, Toronto, Ontario, 2009. URL https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf.
Lang et al. (2024) Hunter Lang, David Sontag, and Aravindan Vijayaraghavan. Theoretical analysis of weak-to-strong generalization, 2024. URL https://arxiv.org/abs/2405.16043.
Lin et al. (2024a) Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter L. Bartlett, and Jason D. Lee. Scaling laws in linear regression: Compute, parameters, and data, 2024a. URL https://arxiv.org/abs/2406.08466.
Lin et al. (2024b) Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, and Weizhu Chen. Rho-1: Not all tokens are what you need, 2024b. URL https://arxiv.org/abs/2404.07965.
Loureiro et al. (2022) Bruno Loureiro, Cédric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mézard, and Lenka Zdeborová. Learning curves of generic features maps for realistic datasets with a teacher-student model*. Journal of Statistical Mechanics: Theory and Experiment, 2022(11):114001, November 2022. ISSN 1742-5468. doi: 10.1088/1742-5468/ac9825. URL http://dx.doi.org/10.1088/1742-5468/ac9825.
Mallinar et al. (2024) Neil Mallinar, Austin Zane, Spencer Frei, and Bin Yu. Minimum-norm interpolation under covariate shift, 2024. URL https://arxiv.org/abs/2404.00522.
Maloney et al. (2022) Alexander Maloney, Daniel A. Roberts, and James Sully. A solvable model of neural scaling laws, 2022. URL https://arxiv.org/abs/2210.16859.
Mei & Montanari (2022) Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4):667–766, 2022.
Menon et al. (2020) Aditya Krishna Menon, Ankit Singh Rawat, Sashank J. Reddi, Seungyeon Kim, and Sanjiv Kumar. Why distillation helps: a statistical perspective, 2020. URL https://arxiv.org/abs/2005.10419.
Mobahi et al. (2020) Hossein Mobahi, Mehrdad Farajtabar, and Peter Bartlett. Self-distillation amplifies regularization in hilbert space. Advances in Neural Information Processing Systems, 33:3351–3361, 2020.
Montanari & Zhong (2022) Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. The Annals of Statistics, 50(5):2816–2847, 2022.
Montanari et al. (2023) Andrea Montanari, Feng Ruan, Youngtak Sohn, and Jun Yan. The generalization error of max-margin linear classifiers: Benign overfitting and high dimensional asymptotics in the overparametrized regime, 2023. URL https://arxiv.org/abs/1911.01544.
Nagarajan et al. (2023) Vaishnavh Nagarajan, Aditya K Menon, Srinadh Bhojanapalli, Hossein Mobahi, and Sanjiv Kumar. On student-teacher deviations in distillation: does it pay to disobey? Advances in Neural Information Processing Systems, 36:5961–6000, 2023.
Oymak & Gulcu (2021) Samet Oymak and Talha Cihad Gulcu. A theoretical characterization of semi-supervised learning with self-training for gaussian mixture models. In International Conference on Artificial Intelligence and Statistics, pp. 3601–3609. PMLR, 2021.
Paquette et al. (2024) Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 phases of compute-optimal neural scaling laws, 2024. URL https://arxiv.org/abs/2405.15074.
Patil et al. (2024) Pratik Patil, Jin-Hong Du, and Ryan J. Tibshirani. Optimal ridge regularization for out-of-distribution prediction, 2024. URL https://arxiv.org/abs/2404.01233.
Richards et al. (2021) Dominic Richards, Jaouad Mourtada, and Lorenzo Rosasco. Asymptotics of ridge (less) regression under general source condition. In International Conference on Artificial Intelligence and Statistics, pp. 3889–3897. PMLR, 2021.
Rudelson & Vershynin (2009) Mark Rudelson and Roman Vershynin. The smallest singular value of a random rectangular matrix, 2009. URL https://arxiv.org/abs/0802.3956.
Simon et al. (2023) James B Simon, Madeline Dickens, Dhruva Karkada, and Michael Deweese. The eigenlearning framework: A conservation law perspective on kernel ridge regression and wide neural networks. Transactions on Machine Learning Research, 2023.
Simon et al. (2024) James B. Simon, Dhruva Karkada, Nikhil Ghosh, and Mikhail Belkin. More is better in modern machine learning: when infinite overparameterization is optimal and overfitting is obligatory, 2024. URL https://arxiv.org/abs/2311.14646.
Song et al. (2024) Yanke Song, Sohom Bhattacharya, and Pragya Sur. Generalization error of min-norm interpolators in transfer learning, 2024. URL https://arxiv.org/abs/2406.13944.
Sorscher et al. (2022) Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 35:19523–19536, 2022.
Spigler et al. (2020) Stefano Spigler, Mario Geiger, and Matthieu Wyart. Asymptotic learning curves of kernel methods: empirical data versus teacher–student paradigm. Journal of Statistical Mechanics: Theory and Experiment, 2020(12):124001, 2020.
Thrampoulidis et al. (2015) Christos Thrampoulidis, Samet Oymak, and Babak Hassibi. Regularized linear regression: A precise analysis of the estimation error. In Proceedings of The 28th Conference on Learning Theory, volume 40 of Proceedings of Machine Learning Research, pp. 1683–1709, Paris, France, 03–06 Jul 2015. PMLR.
Tsigler & Bartlett (2022) A. Tsigler and P. L. Bartlett. Benign overfitting in ridge regression, 2022. URL https://arxiv.org/abs/2009.14286.
Wang et al. (2023) Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi. Self-instruct: Aligning language models with self-generated instructions. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 13484–13508, 2023.
Wei et al. (2022a) Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize, 2022a. URL https://arxiv.org/abs/2203.06176.
Wei et al. (2022b) Colin Wei, Kendrick Shen, Yining Chen, and Tengyu Ma. Theoretical analysis of self-training with deep networks on unlabeled data, 2022b. URL https://arxiv.org/abs/2010.03622.
Wu & Xu (2020) Denny Wu and Ji Xu. On the optimal weighted regularization in overparameterized linear regression. Advances in Neural Information Processing Systems, 33:10112–10123, 2020.
Yang et al. (2023) Fan Yang, Hongyang R. Zhang, Sen Wu, Christopher Ré, and Weijie J. Su. Precise high-dimensional asymptotics for quantifying heterogeneous transfers, 2023. URL https://arxiv.org/abs/2010.11750.

Appendix A Proofs for Section 3

Below, we summarize the notations used throughout this paper:

Category	Symbol	Description
Risk		Excess population test risk.
		Asymptotic excess test risk.
		Asymptotic omniscient excess test risk, which is equivalent to the two-stage risk when .
Parameters		Estimator of the surrogate model after the stage 1.
		Estimator of the target model after the stage 2.
		Estimator of the standard target model.
		Ground truth parameter.
Covariances		Covariance matrix of the data distribution used in surrogate model.
		Covariance matrix of the data distribution used in target and standard target models.
Variables		’th eigenvalue of the covariance matrix in decreasing order.
		The unique solution to the fixed point equation:
		The unique solution to the fixed point equation:
		Covariance statistics given by .
		.
		Asymptotic characterization of the minimum interpolator obtained after training surrogate model.
		Asymptotic characterization of the minimum interpolator obtained after training target model.
		.
		.

Table 1: Summary of notations used in the paper.

See 1

Proof.

By Observation 2, we assume that and are diagonal matrices. As and are jointly diagonalizable, there exists a unique diagonal matrix such that

Then, consider the model shift discussed in Section 3. Take the case where and labels are generated as , where and . This is equivalent to the case where such that and . 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 and are equivalent. The other way follows from an almost identical argument. ∎

Observation 2.

For any covariance matrix , there exists an orthonormal matrix such that the transformation of and does not affect the labels but ensures that the covariance matrix is diagonal.

Proof.

Since the covariance matrix is PSD, its unit-norm eigenvectors are orthogonal. Consider the matrix whose columns are the eigenvectors of . Then, can be expressed as , where is the diagonal matrix containing the eigenvalues of . Consider now the transformation

In this way, the covariance matrix is diagonalized. Thus, the transformation 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 as . The gradient of this function is

Using Corollary 2, there exists an event with where (with the definition of in Corollary 2), such that is -Lipschitz if . Applying Theorem 4 on the target model, there exists a constant such that for any , we obtain

(14)

where and

Furthermore,



				(15)

where and . This completes the proof. ∎

See 1

Proof.

We have that

where , . Recall from (9) that

This implies that

				(16)

where follows from the fact that and .

In order to optimize this with respect to , let’s take the derivative:

Hence, the claimed result follows. ∎

See 1

Proof.

When the definition of and is plugged in Proposition 1, the first claim is obtained. Using the diagonalization assumption on , let’s analyze only the -th component of the optimal surrogate given in Proposition 1:

It’s now clear that if and only if .

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

This gives that if all ’s are equal, which implies that all ’s are equal. Concluding, the covariance matrix is a multiple of the identity if and only if . ∎

See 3

Proof.

The proof for Proposition 3 was correct, yet we made changes to it to improve clarity by explicitly stating all intermediate steps. For the purposes of analysis, we assume, without loss of generality, that the first dimensions are selected from in . Based on this, we no longer need to have the decreasing order for the corresponding ’s. Let and represent the asymptotic risk of the target and the surrogate-to-target models, respectively. From the excess test risk formula in Definition 1, we have that

(17)

where . Now, consider the zero-padded vector . This way, we consider the labels in the second training phase as , where . Next, using the asymptotic risk estimate in Equation (10), we write the excess test risk formula for the surrogate-to-target model with respect to the original ground truth labels:

where , , and . Algebraic manipulations give:








				(18)

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

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

(19)

excluding feature in the feature selection reduces the overall risk . Along the same lines, the projection that selects all the features that satisfy minimizes the asymptotic excess test risk. ∎

See 2

Proof.

According to Theorem 3 of Chang et al. (2021), when the second stage is in the under-parameterized regime, the estimator can be expressed asymptotically as:

where . Then, the excess test risk estimate of this estimate is:

As is positive semi-definite, the expression is non-negative and takes its minimum value of at . Note that at the expression corresponds to the risk of the standard target model. Hence, we conclude that it is not possible to surpass the performance of the standard target model with the surrogate-to-target model when the target model is under-parameterized (). ∎

A.1 Analysis for Model Shift under Ridge Regression

In this subsection, we analyze the behavior of the surrogate-to-target model under ridge regression. First, we redefine the surrogate and target models below.

Stage 1: Surrogate model. We consider a data distribution following the linear model , where , and is independent of . Let be the dataset for the surrogate model drawn i.i.d. from . The estimator of the surrogate model can be written as follows:

(20)

where and .

Stage 2: Target model. Given , we consider another data distribution following the linear model , where and . Let be the dataset for the target model drawn i.i.d. from . As for the surrogate model, the estimator for the target model is defined as

(21)

where and . Our analysis will generally apply to an arbitrary choice and will not require it to be the outcome of (20). Finally, we define the excess (population) risk for a given estimator as defined in (5). Throughout the section, we compare the surrogate-to-target model with the following reference model.

Reference Model: Standard target model. We study the generalization performance of with respect to the standard target model, which has access to the ground-truth parameter through labeling. Specifically, consider the dataset drawn i.i.d. from ; then, the estimation is

(22)

where and . We compare the excess risks of the surrogate-to-target model with that of the standard target model .

Now, we will obtain a result similar to Theorem 1 under ridge regression. For this purpose, we have the following definition:

Definition 4.

Let and be the unique solution of the following equation

(23)

Then, define , the function , and the asymptotic risk as in Definition 1 based on .

Theorem 3.

Suppose that, for some constant , we have and . Then, there exists a constant such that, for any , the following holds with :

(24)

Proof.

The proof directly follows from the proof of Theorem 1 with only one modification in (14):

(25)

∎

In Propositions 1, 3 and Corollary 1, we treat as a variable, which is a fixed point equation based on the covariance matrix and ratio between numbers of features and samples under ridgeless regression. However, under ridge regression, we modify the parameter and define the asymptotic risk the same as we define under ridgeless regression based on this parameter. This means that we can directly apply the results of Propositions 1, 3 and Corollary 1 to ridge regression by only modifying the parameter . Similarly, Theorem 2 can also be extended to ridge regression.

A.2 Experimental Details

In the CIFAR-10 experiment, we initially trained the surrogate models on the training portion of the CIFAR-10 dataset. While training the surrogate-to-target models, we employed the predictions from the surrogate models. Conversely, when training the standard target model, we utilized the ground truth labels. During testing, all models were evaluated using the test portion of the CIFAR-10 dataset. It is important to note that the distributions for both surrogate and target are identical since we trained and tested the models on the same training and testing sets.

We employ three distinct surrogate model sizes: big, medium, and small. The big model contains 127,094 parameters, the medium model 58,342 parameters, and the small model 28,286 parameters. All three models are shallow, three-layer convolutional networks that follow the same architectural specifications. For the small model, ; for the medium model, ; and for the big model, .

We initialize the optimizer for our model using stochastic gradient descent (SGD) provided by the optim module of PyTorch. The optimizer is configured with the following parameters: learning rate set to , momentum to , and weight decay to . Additionally, we define a learning rate scheduler, specifically a cosine annealing scheduler, which adjusts the learning rate using a cosine function over 200 iterations, denoted T_max. We use a batch size of 32 and trained all models over 60 epochs.

Appendix B Proofs for Section 4

See 2 In the following proof, we suppose that the empirical distributions of and converge as having fixed the ratio . Then, we will prove that the omniscient risk converges to the asymptotic risk defined in (10).

Proof for the proportional asymptotic case.

Using Theorem 2.3 of Han & Xu (2023), we can estimate as follows:

where

Let

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



				(26)

Then by recalling the eigendecomposition for the covariance matrix , we have

Using the diagonalization of , can now be computed as

Let’s now compute

As , we obtain that the RHS of the previous expression equals

Next, we write more compactly the terms and . By defining the short-hand notation , we have

where . Hence, putting it all together in (26) gives the desired result. ∎

See 4

Proof.

We start with the asymptotic analysis of . Along the same lines as Simon et al. (2024), since is a monotonically decreasing function, we have:

Furthermore,

Hence, combining these two facts gives

which leads to the desired result.

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

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

Hence,

(27)

By the hypothesis on , we have that

(28)

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

See 5

Proof.

Recall from Proposition 3 that we should identify indices which satisfy the condition to decide if we’re better off not selecting this dimension in the surrogate model. Furthermore, Proposition 4 gives that . Putting these together, we have

As , we get . Incorporating (28), we achieve that

Similarly, by following the same procedure with the initial inequality , we get

∎

Figure 3: Comparison of the empirical and theoretical number of features satisfying the feature selection condition in the optimal mask (). The theoretical value is calculated as , ignoring the in Proposition 5. Setting: The feature size is , and the feature covariance follows the power-law structure for and .

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 (). The theoretical value, calculated as ignoring the term, aligns well with the experimental data and the accuracy in estimation increases with . Notably, our theoretical estimate also closely matches the empirical results when the sample size is small relative to the feature size .

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

Together with the eigenvalues, also assume now power-law form for , that is for . Then, in the limit of , 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:

and

Proof.

As discussed in Section 4, in order to analyze the model’s inherent error, we need to set where 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 and .

Case 1: . In this case, it is previously stated by Cui et al. (2022); Simon et al. (2024) that the test risk of ridgeless overparameterized linear regression can be described in the scaling sense as when . Consider the optimal mask operation mentioned in Proposition 3 that selects all features satisfying . Let be the number of selected features. We can then decompose the risk estimate in Definition 2 as follows:

where and 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:

Going back to Proposition 5, we know that, as , the criterion for selecting a feature in the optimal masked surrogate model is given by

Define now . The equation (18) tells us that after the optimal mask operation , is replaced by , which is calculated as follows

Since is a monotonically decreasing function, we can bound the summation by the following two integrals:

In the limit of , we obtain

Thus, we have tightly estimated . Using the fact from Proposition 4 that , and our assumption on the noise variance , we conclude that the scaling law doesn’t change for the surrogate-to-target model as

Case 2: . In this case, we show that the scaling law is determined by , hence changing to has no effect in the scaling sense. From Proposition 4, we have the asymptotic expression , for . We can argue that there exists positive constants , such that . We have that

This implies . At the same time,

Using gives and we can conclude that . From Cui et al. (2022), we already know that when . Using , and our assumption on the noise variance allows us to conclude that the scaling is dominated by , and thus, the scaling law remains unchanged. ∎

See 6

Proof.

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

since we can put in the form of . At the same time, we know that

Recall the optimal surrogate vector discussed in Proposition 1 and the corresponding minimal surrogate-to-target risk . In this case, we can write

(a) Test risks in log-log scale when

(b) Test risks in log-log scale when

Figure 4: Scaling law behavior of the test risks of optimal surrogate models. (a): Associated test risks as a function of sample size in log-log scale. Setting: The feature size is ; the sample size changes from to with increments of ; the feature covariance follows the power-law structure , (b): Associated test risks as a function of sample size in log-log scale when . Setting: Same as in (a) except that .

Similar to the previous proposition and as discussed in Section 4, to analyze the model’s inherent error, we set where is the exponent characterizing the scaling law of the test risk in the noiseless setting. It is previously stated by Cui et al. (2022); Simon et al. (2024) that the test risk of ridgeless overparameterized linear regression can be described in the scaling sense as when . We will proceed by considering two cases based on the relationship between and .

Case 1:

Consider the interval of ’s satisfying and . By Proposition 5, we have

Let be defined as in the previous proposition and define . Then, the interval of interest corresponds to the set of indices such that . Within this interval, we observe

Using the fact from Proposition 4 that tells us and . Utilizing these bounds, we obtain

where (a) follows from the fact , , and . This implies that

Recall that the optimal surrogate-to-target improves over the risk of the standard target model, thus . We therefore conclude for this case.

Case 2:

In this case, we have

where . From Case 2 in Proposition 7, we already know that the same summation – with upper bound rather than – scales as . Yet, since and have the same order , the result remains. This gives , which eventually yields

Hence, this allows us to say that the scaling law doesn’t improve even with the freedom to choose any . ∎

Figure 5: We compare the experimental two-stage risk with our estimated theoretical risk in log-log scale to demonstrate the scaling law. Setting: The feature size is ; the sample sizes change from to with increments of ; both feature covariances follow the power-law structure for and ; the ground truth parameter is specified as .

In Figures 4 and 5, we illustrate the scaling-law behavior of the test risk. In Figure 4(a), the sample size varies from small values up to values close to . Even though a linear trend is observable for smaller values of when in the log-log scale, as approaches , it is less apparent due to finite-sample effects of . We note that the asymptotic approximations in Propositions 4 and 5 rely on being significantly large compared to . Thus, Figure 4(b) considers the same experiment with a larger dimension () and the same range of values to satisfy . In this scenario, we observe a clear linear behavior in the log-log plot, which is consistent with the scaling-law results presented in Proposition 6.

In Figure 5, we analyze the scaling-law behavior of the test risk under a two-stage framework where the sample sizes vary relative to a fixed dimension . In this scenario, we observe a similar linear behavior in the log-log scale when is sufficiently small compared to . Whereas, as the gap between and diminishes, the linearity in the figure disappears as the asymptotic approximations in Propositions 4 and 5 require to be significantly larger compared to . These observations align well with our expectations and results in Section 4.

Proposition 8 (Non-asymptotic analysis of ).

Suppose that is diagonal and for . Assume that for . If satisfies

then for .

Proof.

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

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

Recalling that and assuming , we derive:

We conclude by proving that . For the sake of contradiction, assume that . Then,

which contradicts our assumption that . ∎

Proposition 9 (Non-asymptotic analysis of ).

Suppose that is diagonal and for . Let be defined as in Proposition 8 and assume that for and . Let be the solution to

Then,

Proof.

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

Let’s now utilize the upper and lower bounds for from Proposition 8. Substituting, we have

Dividing both left and right-hand side by gives:

since for and by our assumption on . ∎

Proposition 10.

Under the assumption that

and , we can find a masked surrogate-to-target setting that improves over the risk of the standard target model by selecting all features such that .

Proof.

From Proposition 9, we have

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

From the proof of Proposition 5, we know that

Hence, using the bound on from Proposition 8, it’s enough to find indices such that

(29)

By our assumption and , we obtain that . Since is increasing with when , we have

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

which follows from our assumption on . Thus, discarding the features provided in the interval (29) 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 when as the coefficients of are positive and is smaller compared to the other three coefficients of in the minimum function.

Appendix C Proofs for Section 5

Theorem 4 (Distributional characterization, Han & Xu (2023)).

Let and suppose that, for some , and . Let be the unique solution of the following equation:

(30)

We define the random variable based on and the function as follows:

(31)

Then, for any -Lipschitz function where , there exists a constant such that for any , we have the following:

(32)

where .

See 3

See 2

Proof.

Define a function as . The gradient of this function is

Using Proposition 12, there exists an event with where with the definition of in Proposition 12, such that is -Lipschitz if . Applying Theorem 4 on the target model, there exists a constant such that for any , we obtain

(33)

where and

Furthermore,



				(34)

where and . Let be the event defined in Proposition 11. Let be defined as . By Proposition 13, the function is Lipschitz if on the event . Applying Theorem 4 on the surrogate model, there exists a constant such that for any , we obtain

(35)

where and .

Let be defined as . By Proposition 14 and Proposition 2.1 in Han & Xu (2023), the function is -Lipschitz if on the event . Applying Theorem 4 on the surrogate model, there exists a constant such that for any , we obtain

(36)

Let as . By Proposition 15 and Proposition 2.1 in Han & Xu (2023), the function is Lipschitz if on the event . Applying Theorem 4 on the surrogate model, there exists a constant such that for any , we obtain

(37)

By the definition of these functions, we have

(38)

By the definition of , , , and , we have

(39)

Using (38)-(39) and applying a union bound on (33), (35), (36), and (37), we obtain the advertised claim. ∎

Proposition 11.

Suppose that, for some , and . For every , there exists an event with where such that

Proof.

By the definition of , we have


				(40)

where . By triangle inequality, we obtain


				(41)

where in above follows from the fact that is a projection matrix, and so all of its eigenvalues are either or . Focusing on the second term of the RHS, we derive


				(42)

where in the above inequality follows from Cauchy-Schwarz inequality. Using Bernstein’s inequality, there exists an absolute constant that depends on such that

On the other hand, let , which means that the entries of are independent and normally distributed with zero mean and unit variance. Then,

(43)

Using Theorem 1.1 in Rudelson & Vershynin (2009), there exist absolute constants such that we have the following for every

(44)

By combining (42), (43), and (44), we obtain that

The advertised claim for follows when is selected as . For , using the definition of , we write as follows:



				(45)

where in the above inequalities follows from the fact that the eigenvalues of are either or as the eigenvalues of are either or . The remaining part of this proof is identical to the previous part. ∎

Corollary 2.

Suppose that is given, and for some , we have and . For every , there exists an event with where such that

Proof.

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

Proposition 12.

Suppose that, for some , and . For every , there exists an event with where such that we have the following on this event :

Proof.

By the definition of , we have the following:

(46)

where . Plugging (40) into (46), we obtain

(47)

Note that and 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 or in (47), we have

(48)

By a similar reasoning used in (42),(43), and (44); there exist absolute constants such that we have the following for every :


				(49)

Similarly, for every , there exist absolute constants such that we have the following for every :


				(50)

Note that and are independent of each other. Therefore, we can apply union bound on (49) and (50) with selecting , and such that , , , and . As a result, there exists an event with such that

Using a similar argument in (45), we derive the following on the same event

This completes the proof.

∎

Proposition 13.

Let be a function such that

Then, on the same event in Proposition 11, the function is Lipschitz where is the largest eigenvalue of .

Proof.

We take the gradient of the function :

Combining Proposition 11 on the event with the above inequality provides the advertised claim. ∎

Proposition 14.

Let be a function such that

Then, on the same event in Proposition 11, the function is Lipschitz where are the eigenvalues of with a descending order and

Proof.

We take the gradient of the function :

Note that

Then, we have

Plugging into , we obtain that

Combining Proposition 11 on the event with the above inequality provides the advertised claim. ∎

Proposition 15.

Let be a function such that

Then, the function is Lipschitz where is the largest eigenvalue of .

Proof.

We take the gradient of the function :

and the desired result readily follows. ∎

Lemma 1.

We have that

Proof.

The desired claim follows from the following manipulations using the definition of in (3):

∎

C.1 Analysis of Two-stage Model in Under-parametrized Region

Proposition 16.

Consider the two-stage model where the surrogate model is under-parameterized with features () and the target model is over-parameterized () with features. Then, the difference in the asymptotic risks of surrogate-to-target and standard target models is:

where , , and

Proof.

Let where . Then, by Equation (10), the asymptotic risk estimate for the single-stage linear regression is given as:

(51)

For the ease of the analysis, define

Let be obtained by setting the last entries of to . Then, by Theorem 3 of Chang et al. (2021), since the surrogate model is under-parameterized, the following asymptotic estimate for the surrogate model obtained after this stage holds:

where . Let , , and . Thus, plugging this surrogate parameter in the asymptotic risk estimate for the second stage gives:

where

Further modifications give:

Then, using the fact that since is orthogonal to (thus, their supports do not overlap), it follows that the difference between this surrogate-to-target risk and the standard target model risk is:

(52)

This completes the proof. ∎

Proposition 17.

Consider the two-stage model where the surrogate model is under-parameterized with features () and the second stage is over-parameterized () with features. Under this setting, the surrogate-to-target model cannot outperform the standard target model in terms of asymptotic risk.

Proof.

Let and where . Then, by Equation (10), the asymptotic risk estimate for the single-stage linear regression is given as:

(53)

When the surrogate model is under-parameterized, we can use the following asymptotic estimate for the surrogate model obtained after this stage by Theorem 3 of Chang et al. (2021):

where . Let , , and . Then, by plugging this in Equation (53), the asymptotic risk estimate is the following:

where

Further simplifications give:

which is the asymptotic risk estimate of the standard target model (follows from Definition 2 or by simply plugging in Equation (53)). Hence, we conclude that improvement over the excess test risk is not possible when the first stage is under-parameterized with -dimensions. ∎

Corollary 3.

Consider the setting of Proposition 3. Let be the covariance matrix of the target model and let be the mask such that the surrogate-to-target model with and outperforms the standard target model in the asymptotic risk. Let where . Then, for any , there exists such that the surrogate-to-target model with two stages given , , , and outperforms the standard target model.

Proof.

We make the following three observations. First, the expression in (52) is continuous with respect to when . Note that is a natural number in our setting, but we consider the expression in (52) as a function of where . Second, the expression in (52) is monotone with respect to since and . Third, by definition, we know that the standard target model with and outperforms the standard target model in the asymptotic setting. This means that the standard target model with two stages given , and while approaches outperforms the standard target model. By combining these three observations, the desired result follows. ∎

Search across your past Websets

Research Papers: knowledge distillation, transfer learning, distribution shift

Search Criteria

Enrichments

Find more results