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

TL;DR

This paper provides a detailed theoretical analysis of knowledge distillation in high-dimensional regression, revealing how surrogate models influence target model risk, and demonstrating conditions under which weak-to-strong generalization can outperform traditional training.

Contribution

It offers the first sharp, non-asymptotic bounds for high-dimensional knowledge distillation, characterizes the optimal surrogate model, and clarifies the limits of weak-to-strong generalization in data scaling.

Findings

W2S training can outperform strong label training under the same data budget.

Optimal surrogate models depend on data distribution and feature relevance.

W2S does not improve data scaling laws despite performance benefits.

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.