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

We thank the reviewer for the positive evaluation of our work and for the detailed comments. We address all of them below, and we have uploaded a revised version of the manuscript highlighting in red the main changes.

W1: 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 the convex gordon min max theorem.

Response: Thanks for the feedback. We added a short proof sketch for Theorems 1 and 2 in the main text. In addition, we want to note that we give credit to the theory developed by Han & Xu (2023) in lines 237-239 and 477-479.

W2: 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.

Response: Thank you for your insightful feedback. We improved our explanation of not being able to outperform the standard target model in practical settings, unlike our theoretical results. Here is our revised explanation:

“The reason why the surrogate-to-target model underperforms the standard target model is that the surrogate model is not able to follow the feature selection mechanism characterized in Proposition 2 as there’s no notion of feature selection in neural networks, unlike linear models. This suggests that the feature selection mechanism is crucial for surpassing the performance of the target model.”

We believe that this figure is valuable because the phenomenon of weak-to-strong generalization extends beyond large language models and can also be observed in neural network architectures.

W3: 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.

Response: We appreciate the insightful question. Our result can be extended to ridge regression by defining $\tau$ based on the ridge parameter. This extension is achievable due to the characterization of the parameter $\hat{\beta}$ provided by Han & Xu (2023). Indeed, we have made initial efforts in this direction, as discussed in Appendix A.1, where we apply the results of Section 3 to ridge regression. Our choice to perform the analysis for ridgeless regression was motivated by the desire to keep the model as simple as possible while conveying the important insights of the surrogate-to-target model.

Furthermore, we want to highlight that the mentioned previous works (e.g. Hastie et al. (2020), Louriero et al. (2022)) that study the ridge regression provide only the asymptotic risk characterization and do not provide the characterization of the parameter $\hat{\beta}$. Therefore, our work is based on a relatively advanced technique developed by Han & Xu (2023).

W4: (Minor) Section 4 is presented unclearly. The settings for the propositions are not well explained and need to be introduced more clearly.

Response: Thank you for your helpful feedback, we have revised the propositions in Section 4 to improve clarity.

Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J. Tibshirani. Surprises in high-dimensional ridgeless least squares interpolation, 2020.

Qiyang Han and Xiaocong Xu. The distribution of ridgeless least squares interpolators, 2023.

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

I want to thank the authors for the time taken to reply. I acknowledge the changes in the text to make the result more transparent in the text and the figures. I thank the authors for explaining the relationship of their work with the previous ones: it clarified my questions. I also don't think it is necessary to provide a full analysis of ridge regression as already interesting behaviour is found with this simple model.

It seems to us that generalising the results of Section 3.1 to the case where the features are chosen via a matrix A actually corresponds to the optimal surrogate given by Proposition 1. In fact, such an optimal surrogate is obtained by multiplying a certain matrix with the ground-truth parameters. In case we have misunderstood the point of the reviewer, we remain at disposal for additional clarifications.

At the cost of being repetitive, I would like to explain myself better. Considering a projection $A$, one can define in the same spirit the masked s2t model as $\tilde{y} = (A x)^T (A \beta_\star) + \tilde{z}$, and then one can choose to include or not the projection matrix in the $\mathcal{D}_t$, for example, one could consider $y = (A x)^T \beta_s+ z$. I don't see how these cases can be framed (with a change of variables) to the results in Proposition 1, and I am genuinely curious, also because $\beta_s^\star$ doesn't seem to depend on $\Sigma_s$.

I have decided to increase my score.

P.S. A minor typo in line 290 is the variance of the noise. I believe it should be $\sigma_t$ instead of $\sigma_s$.

We thank the reviewer for their prompt response and positive feedback. We also appreciate the clarification for the projection matrix $A$.

For the surrogate model stage, the labels $\tilde{y}$ are generated by the following linear equation: $\tilde{y} = (\mathbf{A \tilde{x}})^\top (\mathbf{A \beta_*})$. Then, the surrogate parameter $\mathbf{\beta^s}$ is estimated using the features $\mathbf{A \tilde{x}}$, as described in line 291. When the surrogate model has infinitely many data, the covariance matrix $\Sigma_s$ does not affect the surrogate parameter $\mathbf{\beta^s}$, so the surrogate parameter $\mathbf{\beta^s}$ is equal to $\mathbf{A \beta_*}$. For the analysis of this setting under finite data, please refer to Appendix C.1.

For the target model stage, the labels $y$ are generated by the following linear equation: $y = (\mathbf{A x})^\top \mathbf{\beta^s}$. Then, the surrogate-to-target parameter $\mathbf{\beta^{s2t}}$ is estimated using the features $\mathbf{x}$, as described in line 292. The label-generation process reduces to the following:

$\qquad y = (\mathbf{A x})^\top \mathbf{\beta^s} = \mathbf{x}^\top (\mathbf{A^\top \beta^s}) = \mathbf{x}^\top (\mathbf{A^\top A \beta_*}).$

Let’s define the matrix $\mathbf{K}$ as follows:

$\qquad \mathbf{K} = ((\Sigma_t + \tau_t \mathbf{I})^{-1} \Sigma_t + \frac{\Omega \tau_t^2}{1- \Omega} \Sigma_t^{-1} (\Sigma_t + \tau_t \mathbf{I})^{-1})^{-1}.$

Now, if the projection matrix $\mathbf{A}$ is rank$-p$, then the optimal projection matrix $\mathbf{A}$ can be found by solving $\mathbf{ A^\top A} = \mathbf{K}$ following Proposition 1.

If the projection matrix $\mathbf{A}$ is rank$-p_s$ where $p_s \leq p$, then we need to select the largest $p_s$ signal coefficients of the labels $y$. This corresponds to the the largest $p_s-$eigen-directions based on $\lambda_i^2 \beta_{*, i}$ where $\lambda_i$ are the eigenvalues of $\Sigma_t$. (Here, WLOG, we assume that $\Sigma_t$ is diagonal by Observation 2). Let $\mathbf{U} \in \mathbb{R}^{p \times p_s}$ be the matrix whose columns are the largest $p_s-$eigen-directions with unit norm. Since the optimal surrogate parameter in Proposition 1 can be decoupled for each dimension, then an optimal projection matrix can be $\mathbf{U^\top K^{1/2}}$ following Proposition 1.

Kindly let us know if this post addresses your question. If the reviewer thinks this generalization is worth writing in the paper, we can incorporate it in the Appendix.

P.S. A minor typo in line 290 is the variance of the noise. I believe it should be $\sigma_t$ instead of $\sigma_s$.

Thanks for the typo notice. We fixed the typo and uploaded the new version.

We thank the reviewer for appreciating the strengths of our work and for the detailed comments. We reply to each of the points raised in the review below. We have also uploaded a revised version of the manuscript highlighting in red the main changes.

W1: 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.

Response: Even if high-dimensional linear regression has been the subject of intense study in the literature, a precise characterization of knowledge distillation and weak-to-strong generalization was still lacking prior to our work. We see the extension to neural networks as an exciting future direction, and, following the suggestion of the reviewer, we now discuss such extension in Section 6 of the revision, see our response to Q1 below.

W2: A typo in line 134.

Response: Thanks for noticing this, we have corrected the broken reference.

Q1: 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.

Response: Going beyond linear regression, 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 directions for future research.

Thank you for this suggestion, we now discuss this point in the final section of the revision.

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, 2022.

Andrea Montanari and Yiqiao Zhong, “The interpolation phase transition in neural networks: Memorization and generalization under lazy training”, Annals of Statistics, 2022.

We thank the reviewer for recognizing the significance of our work, and the comprehensiveness of our results, as well as for the detailed comments. We address all of them below, and we have uploaded a revised version of the manuscript highlighting in red the main changes.

W: 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?

Response: Thank you for raising this issue, which has allowed us to improve the presentation of our results in the revision. We now discuss the interpretation of the terms appearing in the expression for the risk (see l. 227-230 and 240-244 of the revision), and we give the intuition behind improving the strong model with the weak-to-strong estimator (see l. 270-281 of the revision). These changes are discussed in detail in our response to Q1 below.

Q1: 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?

Response: The improvement in the surrogate-to-target model's risk arises because we effectively cut/shrink (depending on the optimal masked surrogate or optimal surrogate model) the tail components of the covariance matrix that contribute less useful information. This is illustrated in Figure 1-a.

However, the scaling law of the risk remains the same because the significant/uncut components (with large corresponding eigenvalues and signal coefficients) of the covariance matrix dominate the behavior of the risk and determine the scaling law exponent. This essentially follows from using the expression discussed in line 430 as the lower bound.

To improve the presentation, we split the risk in (10) into two parts, $a$ and $b$, providing intuition for each of them. Specifically, the terms $a$ and $b$ corresponds to the following:

$\qquad a = || \Sigma_t^{1/2} (\beta_* - \mathbf{\theta}_1 \beta^s ) ||_2^2$

$\qquad b = \gamma_t^2(\beta^s) \mathbb{E}_{\mathbf{g}_t} [ \mathbf{\theta}_2^\top \Sigma_t \mathbf{\theta}_2 ]$

We have added the following explanation in the main text: “In the asymptotic risk, the term $(a)$ corresponds to a part of the bias risk caused by the model shift ($\beta^s$), and the implicit regularization term where the eigenvalues of $\mathbf{\theta}_1$ are less than 1. The term $(b)$ corresponds to the remaining part of the bias and variance risks.”

To answer the question of why there is room for improvement by introducing the surrogate parameter, we have added the following paragraph: “The surrogate parameter $\beta^s$ that minimizes the individual term (a) in 10 is $\mathbf{\theta_1^{-1}} \beta_*$. On the other hand, the surrogate parameter $\beta^s$ that minimizes the individual term $(b)$ 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.”

We have provided the intuition behind a better performance of the surrogate parameter at the end of Section 3 as follows: “The intuition behind improving the performance of the standard target model by utilizing a surrogate parameter $\beta^s$ different from $\beta_*$ is associated with the implicit regularization of the minimum norm interpolator in the over-parametrized region $(p > n)$. 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, we will show that the optimal surrogate parameter $\beta^s$ is $\beta_*$ when the target model is under-parametrized:

Proposition The optimal surrogate parameter $\beta^s$ that minimizes the asymptotic risk in the under-parametrized region ($n > p$) is equivalent to the ground truth parameter $\beta_*$. In other words, for any $\beta^s$, 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.”

We thank the reviewer for their positive evaluation of our work and for the detailed comments. We reply to each of the points raised in the review below. We have also uploaded a revised version of the manuscript highlighting the main changes in red color.

W1: 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).

Response: The reviewer is correct that the optimal surrogate parameter $\beta^s$ does depend on the ground-truth parameter $\beta_*$. By characterizing the optimal surrogate, we derived its form in terms of $\beta_*$ and provided intuitions for the surrogate-to-target model, including the transition points. Crucially, this also revealed that the mapping from the ground truth parameter $\beta_*$ to the optimal surrogate parameter $\beta^{s}$ is purely in terms of the feature covariance and does not depend on the foreknowledge of the task distribution, as stated in Proposition 1. This means that, rather than pruning $\beta_*$, we can alternatively modify the feature covariance (e.g. by truncating tail eigendirections) during the two stage regression in Section 5, as the reviewer has also suggested. If the first stage has access to infinite data (population covariance), two stage regression with covariance truncation is identical to Section 3.1. If we have finite data, then we can use empirical covariance as an approximation.

Indeed, the aim of Section 5 is to fill the gap between theoretical findings and practical applications because this section characterizes the risk of the surrogate-to-target model when the surrogate model has finite data. We have extended the analysis in Section 5 to provide a more comprehensive treatment of the findings from Section 3.1 in Appendix C.1., where we allow the surrogate model to be inside the under-parametrized region.

W2: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.

Response: We thank the reviewer for the question. We have created a notation table at the beginning of Appendix A to bring further clarification. In Section 3, we analyze the case where the surrogate parameter $\beta^s$ is given and optimize the surrogate parameter based on the asymptotic risk. In Section 4, we analyze the scaling performance of the surrogate-to-target model as $p \rightarrow \infty$ and find that the scaling law of the surrogate-to-target model and target model is the same. Finally, in Section 5, we analyze the case where the surrogate parameter $\beta^s$ is a solution to an empirical minimization problem. To further clarify our results in the main text, we have added an explanation of the settings in the statement of Proposition 6 (Proposition 5 in the submitted version). In Theorems 1 and 2, we had already explained the setting and its notation in the statement of the submitted version. We are open to any other suggestions to clarify our statements.

We would like to thank all the reviewers for their constructive comments, which significantly helped improve both the clarity and content of the paper.