Prediction-Only Distillation in Linear and Logistic Regression

Machine Learning

Summary

The authors examine a way to improve machine learning predictions when you only have access to a trained model and new unlabeled data, but not the original training labels. They propose training a new model (student) using labels generated by the original model (teacher) on fresh data, then combining both models' predictions to get better results. Their analysis in ridge regression shows this combined approach usually beats just using the teacher alone, even when new data differs from the original. They also demonstrate how to find the best combination using a small amount of labeled data and show similar benefits in logistic regression tasks.

Authors

Hien Dang, Pratik Patil, Alessandro Rinaldo

Abstract

Self-distillation (SD) is typically studied when the student is retrained on the teacher's original training inputs. In many practical deployments, however, the labeled training data are no longer available, and one has access only to the trained predictor and fresh unlabeled covariates. We study SD in this prediction-only regime through a fresh-X prediction-mixed scheme: a pure-distilled student is trained on fresh covariates pseudo-labeled by the teacher, and the final predictor is an affine combination of the teacher and student predictions. For ridge regression under proportional asymptotics, we derive deterministic equivalents for the optimally mixed prediction risk under general anisotropic covariance and deterministic signal. We show that this risk is strictly smaller than the teacher risk for almost every pair of teacher and student regularization levels, including when the fresh covariates are out-of-distribution and even when their covariance is isotropic. We further show that the optimal mixing weight cannot be identified from unlabeled data alone, but can be consistently estimated in a single post-training step using a small independent labeled calibration set, without additional model fitting. Finally, for binary logistic regression, we show that prediction mixing can outperform both the teacher and the pure-distilled classifier.