Understanding the Gains from Repeated Self-Distillation

Thank you for your review. We address your concerns below.

Concern 1: Relaxing the strong assumption of $\theta^\star$ being perfectly parallel to $\mathbf{u}_1$

Regarding the condition of $\theta^\star$ being perfectly parallel to $\mathbf{u}_1$, we address this in the shared response, since this was a common concern among reviewers. We agree that it is a strong assumption -- and we presented it as such in the interest of a clean exposition of Theorem 1. At the time of writing, the focus was on showing that there exists a regime where $r$-step SD can order-wise outperform $1$-step SD (i.e., not just by a $O(1)$ factor).

But the result does hold true more generally. This is clarified in detail in the shared response via the below two arguments, and a new version of Theorem 1 that incorporates both of these points.

• First, we note that alignment with any one of $\mathbf{u}_j, j \in [r]$ is sufficient for the separation (as noted in lines 203-205 in the manuscript also).
• Second, based on your valuable recommendation, we have indeed derived a version of Theorem 1 with a relaxed Assumption 2.2, where the closeness of $\theta^\star$ and the singular vector $\mathbf{u}_j$ is controlled with a parameter $\beta$, and setting $\beta = 0$ recovers the special case that we presented.

Through the revised Theorem 1, we argue that the regime of separation is more than just an atypical corner case. Overall, we show that the separation holds whenever the ground-truth $\theta^\star$ is "$\beta$-close" to any of the $\mathbf{u}_j, j \in [r]$; and for small $\beta$. We believe this is a fairly general regime and not just a special case. We hope this provides resolution for your concern.

Concern 2: Insights for self-distillation applied to classification

We understand that linear regression is a relatively simple task, but we believe that the $\Omega(r)$ order-wise separation between $r$-step SD and $1$-step SD, although in the simple framework of linear regression, is a somewhat surprising result worth sharing with the community.

The de-facto loss for classification problems is the cross-entropy loss, which introduces non-linearity and additional technical challenges for theoretical analysis. But practically, there's a large body of work on empirically analyzing multi-step SD in classification problems. Many researchers have observed the phenomenon of multi-step SD providing performance gains in various settings [1,2,3]. Our work is an attempt to theoretically characterize the gain, and it's surprising that linear regression itself exhibits a non-trivial gain from self-distillation. A detailed analysis for classification with non-linearity falls outside the scope of this work.

[1] T. Furlanello, Z. Lipton, M. Tschannen, L. Itti, and A. Anandkumar. Born again neural networks. In International Conference on Machine Learning, pages 1607–1616. PMLR, 2018.
[2] Y. Li, J. Yang, Y. Song, L. Cao, J. Luo, and L.-J. Li. Learning from noisy labels with distillation. In Proceedings of the IEEE international conference on computer vision, pages 1910–1918, 2017.
[3] Z. Zhang and M. Sabuncu. Self-distillation as instance-specific label smoothing. Advances in Neural Information Processing Systems, 33:2184–2195, 2020.

Thank you for your review! We address your concerns and questions below.

Concern #1.1: Comparison to another relevant paper [1]

Self-distillation has been explored in [1,2,3], and we do provide a comparison of our results with [2,3]. Thank you for pointing out the highly relevant reference [1] also! There are definitely similarities, but the main conclusions do differ. Let us elaborate more below.

Indeed [1] also studies multi-step SD similar to us (Fig 1 in both papers are analogous), and they also obtain an analytical form for the k-step SD (Theorem 4.1 in theirs is analogous to Eq. (8) in ours).

The crucial difference is the freedom of the $\xi$ parameters being different at each step of self-distillation. In Lemma 4.2 and Theorem 4.3, [1] assumes the $\xi$ values at each step are equal (denoted $\alpha^{(2)} = \alpha^{(3)} = \cdots = \alpha^{(\tau)}$ in their paper). Consequently, they conclude that subsequent steps of SD progressively sparsify the basis set of obtainable solutions. This means that after a point, running more steps of SD will result in a poorer performing model (similar to [3]), as you pointed out in your review.

Our main result (Theorem 1) is different. It says that subsequent steps of SD strictly provide more freedom, and that the best multi-step SD can outperform the best 1-step SD by an $\Omega(r)$ factor. This separation relies on the freedom of $\xi$s being different at each step, which required careful analysis.

We believe this is a significant novel contribution since an order-wise separation of $\Omega(r)$ (and not just an $O(1)$ difference) is somewhat surprising. But we will surely cite the contributions of [1] as they are highly relevant.

[1] Kenneth Borup and Lars N.Andersen. Even your Teacher Needs Guidance: Ground-Truth Targets Dampen Regularization Imposed by Self-Distillation. In Advances in neural information processing systems, 2020.

[2] Rudrajit Das and Sujay Sanghavi. Understanding self-distillation in the presence of label noise. In Proceedings of the 40th International Conference on Machine Learning, 2023.

[3] Hossein Mobahi, Mehrdad Farajtabar, and Peter Bartlett. Self-distillation amplifies regularization in hilbert space. In Advances in Neural Information Processing Systems, 2020.

Concern #1.2: $k$-step SD for $k>2$

In the regression tasks of Section 5.3, what was challenging for higher $k$-step SD ($k>2$) was numerical instability. In particular, from Theorem 5 we know that the optimal choice of $\xi$ involves computing $M^{-1}$ for the matrix $M \in \mathbb{R}^{k \times k}$, which was unstable to invert for $k > 2$ in the real-world datasets of Section 5.3. There might be more stable ways to approximate the optimal solutions, which will make our approach more practical for higher order SD. We leave this as a future research direction.

As a proof of concept though, we refer you to Figure 2 of the rebuttal pdf, where we run $k$-step SD upto $k=5$ on a synthetic problem. We show that for that specific example, self-distillation beyond $2$-step SD is indeed helpful. In the real-world experiments, we ran only till $2$-step SD because $3$-step SD started becoming numerically unstable for all three of the datasets.

Concern #2: Assumption 2.2 being too strong

We agree that Assumption 2.2 as presented is indeed a strong condition, but we do not really need that strong condition to see the gains of multi-step SD. Let us elaborate why:

1. Theoretical argument: In general, we require significantly weaker conditions on $\theta^\star$ for the separation result (Theorem 1). This is clarified in great detail in the shared response (since this was a common concern among reviewers), where we (i) note that alignment with any one of $\mathbf{u}_j, j \in [r]$ is sufficient for the separation, and (ii) also present a relaxed version of Theorem 1 with a weaker Assumption 2.2 that does not require exact alignment. Please refer to the shared response window for a full discussion.
2. Empirical argument: The relaxed condition of Assumption 2.2 is fairly reasonable, as shown empirically on the Air Quality and Airfoil regression tasks, where $2$-step SD beats $\{1,0\}$-step SD (especially Air Quality where the gap in performance is significant). We included AEP dataset as a negative example where multi-step SD does not provide gains. Further, Figure 2 of the rebuttal pdf provides insight into why AEP saw no gains. It shows that the $\theta^\star$ is indeed strongly aligned with one of the $\mathbf{u}_j, j \in [d]$ for Air Quality and Airfoil datasets, but NOT for the AEP dataset.

We presented Assumption 2.2 as such in the interest of a clean exposition of Theorem 1. At the time of writing, the focus was on showing that there exists a regime where $r$-step SD can order-wise outperform $1$-step SD (i.e., not just by a $O(1)$ factor). But the above shows, both theoretically and empirically, that multi-step SD can provide gains in much more general settings.

Question #2: Models and datasets used

Since our setting is regression, CIFAR-10 is a less relevant dataset for us. We demonstrate the empirical results on regression datasets taken from the commonly used UCI repository. The experiments presented are for linear estimators of the type $\hat{\theta} \in \mathbb{R}^d$. Non-linear networks were not used to keep the experiments in accordance with the theory.

Thank you for your review! We address your questions and concerns below.

Question 1: Analysis of a more complex model family

We understand that linear regression is a relatively simple task, but we believe that the $\Omega(r)$ order-wise separation between $r$-step SD and $1$-step SD, although in the simple framework of linear regression, is an interesting and perhaps surprising result.

On extensions, kernel regression is a fairly direct extension in terms of technical tools. The challenge for more expressive models (eg. MLPs) is the introduction of $(i)$ non-linearity and $(ii$) non-convexity, which pose significant additional technical challenges for theoretical analysis. The solution of the ERM problem becomes hard to characterize. Further, there is interplay with optimization since the objective can become non-convex.

Because of these difficulties, researchers generally resort to empirical studies for more complex models. Many past works have observed the phenomenon of multi-step SD providing performance gains in various settings [1,2,3]. Our work is an attempt to theoretically characterize the gain, and it's somewhat surprising that linear regression itself exhibits a non-trivial order-wise gain from self-distillation.

[1] T. Furlanello, Z. Lipton, M. Tschannen, L. Itti, and A. Anandkumar. Born again neural networks. In International Conference on Machine Learning, pages 1607–1616. PMLR, 2018.
[2] Y. Li, J. Yang, Y. Song, L. Cao, J. Luo, and L.-J. Li. Learning from noisy labels with distillation. In Proceedings of the IEEE international conference on computer vision, pages 1910–1918, 2017.
[3] Z. Zhang and M. Sabuncu. Self-distillation as instance-specific label smoothing. Advances in Neural Information Processing Systems, 33:2184–2195, 2020.

Question 2: Interpretation of Assumption 2.2 in more complex models

At a high-level, $\theta^\star$ is akin to the parameterization of the ground-truth generative process (which could be non-linear in general). And $\mathbf{u}_j$ is akin to a basis direction of the observed data (which is computed simply via SVD for the linear case, but the notion of basis could involve non-linearity for more general models). These notions are far from precise though, and it is not clear what the interpretation of their "closeness" (as in Assumption 2.2) would mean in a more complex regime. Characterizing this more explicitly would be a very interesting direction of future work, but falls outside the scope of our current work.

Question 3: Insight into performance on the AEP dataset

Note that $\{1, 2\}$-step SD achieve roughly the same test MSE as ridge on the AEP dataset (Table 1 of manuscript), so there is no drop in performance, it's just flat. In the attached pdf (Figure 1), we provide an explanation for why this is the case. Unlike the other two datasets, AEP happens to have a $\theta^\star$ which is not strongly aligned with any particular basis direction $\mathbf{u}_j, j \in [d]$. This violates Assumption 2.2, where we require $\theta^\star$ to be well-aligned with one of the $\mathbf{u}_j$ directions for observing large gains from multi-step SD.

Please also refer to the shared response for a detailed discussion on Assumption 2.2. In particular, we present a relaxed version of Theorem 1 where $\theta^\star$ needs to be closely (and not just exactly) aligned with any of the $\mathbf{u}_j$, and not necessarily $\mathbf{u}_1$.

To be more explicit, let us break it into two cases:

• Case 1: From the new version of Theorem 1, we know that closeness of $\theta^\star$ to some $\mathbf{u}_j, j \in [d]$ (which translates to the existence of one high peak in the Figure 1 bar chart) implies large gains from self-distillation.
• Case 2: On the other hand, from Theorem 3 of the manuscript, we know that if $\theta^\star$ is equally aligned to all of $\mathbf{u}_j$s (which translates to a flat distribution in the Figure 1 bar chart), it provably implies no gain from self-distillation.

Since the AEP dataset happened to be closer to the second case, we see no gains from self-distillation. We thank you for this suggestion, and will aim to include this explanation in the revised manuscript.

Other Concerns: Hyperparameter selection

The comments also raised a concern about principled hyperparameter selection methods. We would like to highlight that there is a principled one-shot tuning method we explain in Section 5.2 (and Appendix K) for the $\xi$ parameters, leveraging the theoretical insights from section 4.4. In particular, from Theorem 5 in section 4.4, we know that the excess risk is quadratic in the $\bar{\xi}$ parameters (one-to-one reparameterization of the $\xi$ parameters). Using the quadratic nature, one can evaluate optimal $\xi$s for any regularization strength $\lambda$, leaving only one parameter ($\lambda$) for the grid search.

Smaller points to improve

Thank you for these comments. We address them below.

1. Assumption 2.2 as stated actually implies that $\theta^\star$ is orthogonal to every eigenvector except the first one. We state Assumption 2.2 as $\measuredangle (\theta^\star, \mathbf{u}_1) = 0$, which means that $\theta^\star$ is perfectly parallel to $\mathbf{u}_1$. Since $\mathbf{u}_j, j \in [d]$ are all eigenvectors of $\mathbf{X} \mathbf{X}^T$, they form an orthonormal basis of $\mathbb{R}^d$ (also stated in Appendix A), i.e. $\langle \mathbf{u}_i, \mathbf{u}_j \rangle = 0$ for $i, j \in [d], i \neq j$. So indeed, the assumption as stated implies that $\langle \theta^\star, \mathbf{u}_j \rangle = 0$ for $j \in \{2, 3, \cdots, d\}$. There is a separate discussion on whether this condition is too strong, which we address in the shared (global) response window.
2. Thank you for catching this! We agree that a clearer way of stating this would explicitly call out the sub-assumptions 2.1 and 2.2. We will work this out in the final draft.
3. Totally agreed, we will modify the appendix to have relevant equations handy.
4. Also a valid point. We will rework this to a friendlier plotting style (for example, plotting with different marker styles as opposed to different colors).

Thank you for your encouraging review! Let us try to resolve some of your questions and concerns.

Question 1: Gap between $2$-step SD and $r$-step SD

The power of two choices in load balancing is related to two choices being the first non-trivial departure from the standard one choice algorithm. Qualitatively (and intuitively), that departure happens perhaps at $1$-step SD in our setting, where the self-distillation is first used. This is why we focused on the separation between $1$-step SD and $r$-step SD in our manuscript, along with the $0$-step (i.e. ridge) versus $r$-step SD. Quantitatively, it is challenging to quantify the gap between $k$-step SD and $r$-step SD for a general $k > 1$. Our educated guess is that the gap between each additional step of SD can be upto $O(1)$ (and accumulating the gains over $r$ steps gives an $O(r)$ total gain), but this is not currently proved.

As a proof of concept though, we refer you to Figure 2 of the rebuttal pdf, where we run $k$-step SD upto $k=5$ on a synthetic problem. We show that for that specific example, self-distillation beyond $2$-step SD is indeed helpful. In the real-world experiments, we ran only till $2$-step SD because $3$-step SD started becoming numerically unstable for all three of the datasets.

Question 2: Justification of using excess risk as the metric

In linear regression literature, there are two reasons why excess risk in the $\hat{\Sigma}_n$-norm is the preferred metric. The first is that it corresponds to the test error on a fresh sample, which is typically the evaluation metric of choice in many settings. The second is that it is invariant to the basis of the input $x$. We explain each of these reasons in detail below.

Test error calculation

In many cases, we care about the average test error, i.e. the expected error incurred on a fresh sample drawn from the underlying joint distribution $P$. Let's write that mathematically,

$

{\rm TotalRisk} (\hat{\theta}) &= \mathbb{E}_{(x,y) \sim P} \left[ \left( \langle \hat{\theta}, x\rangle - y\right)^2 \right] .

$

We first break the joint expectation into a conditional expectation on $y \sim P_{Y|X=x}$ followed by $x \sim P_X$.

From Assumption 1 of the paper, we know that $P_{Y|X=x}$ is a distribution with mean $\langle \theta^\star, x \rangle$ and variance $\gamma^2$. So we can write that that $y = \langle \theta^\star, x \rangle + \eta$, where $\eta$ denotes the noise and satisfies $\mathbb{E}[\eta]=0, {\rm Var}[\eta] = \gamma^2$. Using this in the total risk expression, we get

$

{\rm TotalRisk} (\hat{\theta}) &= \mathbb{E}_{x \sim P_X} \left[ \langle \hat{\theta} - \theta^\star, x \rangle^2 + \gamma^2 \right] .

$

This simplifies to

$

{\rm TotalRisk} (\hat{\theta}) &= \left( \hat{\theta} - \theta^\star \right)^\top \mathbb{E}_{x \sim P_X} \left[ xx^\top \right] \left( \hat{\theta} - \theta^\star \right) + \gamma^2 \text{ }.

$

Since $P_X$ is the density with point masses at the $n$ training points in the fixed design setting, we get $\mathbb{E}_{x \sim P_X} \left[ xx^\top \right] = \frac{1}{n} \cdot \mathbf{X} \mathbf{X}^\top = \hat{\Sigma}_n$. Using this, we obtain

$

{\rm TotalRisk} (\hat{\theta}) &= || \hat{\theta} - \theta^\star ||^2_{\hat{\Sigma}_n} + \gamma^2 \text{ }.

$

The above calculation shows that excess risk in the $\hat{\Sigma}_n$-norm (i.e. the first term) is the right metric that determines how good an estimator is, above the "noise floor" of $\gamma^2$. We remark in Appendix B that some previous works have used the $\ell_2$-norm instead of the more natural $\hat{\Sigma}_n$-norm.

Invariance

Another reason is that we prefer a measure of performance that is invariant to what basis we choose for the input $x$. Precisely, we want the performance of an algorithm to be the same whether we apply it to data coming from $(\mathbf{X}, \theta^\star)$ or $(A \mathbf{X}, A^{-\top}\theta^\star)$ for any $d\times d$ invertible matrix $A$. This follows from the fact that one can always apply any $A$ to the input before feeding it to the algorithm, and if the performance changes based on $A$, then it makes it hard to draw a fair comparison. We are assuming there is no numerical instability to focus on the invariance property. It happens that the excess risk $|| \hat{\theta} - \theta^\star ||^2_{\hat{\Sigma}_n}$ is invariant to such matrix multiplication (which is the change and/or scaling of the basis). We believe this is also one of the reasons excess risk in the $\hat{\Sigma}_n$-norm is preferred over the $\ell_2$ distance. This is what we meant when we said "ensures that the signal-to-noise ratio is uniform in all directions."

Larger weaknesses

1. Please see the above justification for using excess risk as the metric.
2. There is a one-shot tuning method we explain in Section 5.2 (and Appendix K), where we delineate a principled hyperparameter selection method for the $\xi$ parameters, leveraging the theoretical insights from section 4.4. In particular, from Theorem 5 in section 4.4, we know that the excess risk is quadratic in the $\bar{\xi}$ parameters (one-to-one reparameterization of the $\xi$ parameters). Using the quadratic nature, one can evaluate optimal $\xi$s for any regularization strength $\lambda$, leaving only one parameter ($\lambda$) for the grid search. In practice, what was challenging for higher $k$-step SD ($k>2$) was numerical instability. In particular, from Theorem 5 we know that the optimal choice of $\xi$ would involve computing $M^{-1}$ for the matrix $M \in \mathbb{R}^{k \times k}$, which was unstable to invert for $k > 2$ in the real-world datasets of Section 5.3. There might be more stable ways to approximate the optimal solutions, which will make our approach more practical for higher order SD. We leave this as a future research direction.