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

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.

• The paper demonstrates that weak surrogate models can outperform strong labels by achieving robust generalization under limited data.
• It derives rigorous non-asymptotic risk bounds for both arbitrary surrogate and ERM-based models in model and distribution shift scenarios.
• The findings indicate that optimal surrogate selection preserves scaling law exponents while reducing overall risk, enhancing distillation pipelines.

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

The paper presents a detailed examination of knowledge distillation by analyzing the statistical properties of high-dimensional ridgeless regression under model and distribution shifts. The research focuses on two settings: model shift, involving arbitrary surrogate models, and distribution shift, where surrogate models result from empirical risk minimization (ERM) with out-of-distribution data. This work achieves a precise risk characterization of the target model using non-asymptotic bounds linked to sample size and data distribution, highlighting the optimal surrogate model form and the implications for knowledge distillation.

Technical Summary

The paper investigates scenarios where knowledge distillation is applied, providing a rigorous statistical understanding of its effectiveness. The settings analyzed include:

1. Model Shift: Here, the surrogate model is arbitrary. The researchers derive non-asymptotic risk bounds to characterize this scenario.
2. Distribution Shift: The surrogate model arises from ERM with out-of-distribution data. The paper explores how the two stages of distillation influence the generalization performance of the target model.

Key to this analysis is the high-dimensional regime where sample sizes and feature dimensions are proportional. The authors offer thorough theoretical guarantees, centered on the effective performance of the surrogate-to-target learning process. Specifically, the study highlights the concept of weak-to-strong (W2S) generalization, which suggests that using a surrogate model as a weak form can outperform strong labels under the same data budget. This underscores the potential of a carefully curated surrogate model to enhance learning efficiency.

Strong Numerical Results and Claims

The study identifies that in certain conditions, especially under a power-law decay covariance structure, the surrogate model's risk does not alter the scaling law's exponent, though it can achieve a lower overall risk compared to using true labels. This finding implies the surrogate can enhance performance without modifying fundamental learning curve behaviors.

Implications and Future Directions

The research implications of this study are far-reaching, both practically and theoretically:

• Practical Implications: The insights into optimal surrogate model selection can guide practitioners in designing better knowledge distillation pipelines, especially in settings with large, high-dimensional datasets.
• Theoretical Insights: The work extends the understanding of benign overfitting and double descent phenomena, reflecting on foundational statistical learning principles.
• Future Directions: Potential research extensions include exploring the cascading effects of multi-stage distillation processes and refining data selection criteria within distillation frameworks.

In the field of AI developments, the implications suggest an evolution in handling large-scale data efficiently, paving the way for more robust deep learning models trained via distillation techniques. The detailed exploration of covariance and model behaviors enriches the discourse on scaling laws and presents a refined theoretical model that others can build upon. Such advancements are promising for future AI applications demanding efficient model training with limited or imperfect data labels.