Prediction-Only Distillation in Linear and Logistic Regression

Technical Summary: Prediction-Only Distillation in Linear and Logistic Regression

Problem Statement
Standard self-distillation (SD) typically involves retraining a student model on the teacher's original labeled training data. However, in many practical deployment scenarios, the original labeled training data are unavailable due to privacy, storage, or safety constraints. In these "prediction-only" regimes, practitioners have access only to a trained predictor (the teacher) and a stream of fresh, potentially out-of-distribution (OOD) unlabeled covariates. The central question addressed is whether such fresh, unlabeled data can be leveraged to improve a fixed predictor when the original training data and the teacher's specific regularization parameters are unknown.

Methodology
The authors propose a Prediction-Mixed Self-Distillation (PMSD) framework. The procedure involves three stages:

1. Pure Distillation (PD): A student model is trained on fresh unlabeled covariates using the teacher's predictions as pseudo-labels. This yields a "pure-distilled" student.
2. Affine Mixing: Instead of deploying the PD student alone, the final predictor is constructed as an affine combination of the teacher's predictions and the PD student's predictions: $f_{pmsd}(x) = (1-\xi)f_{teacher}(x) + \xi f_{pd}(x)$. The mixing weight $\xi$ is a real-valued parameter that need not be restricted to the interval $[0, 1]$.
3. Calibration: Since the optimal mixing weight depends on unknown population quantities, the authors propose estimating it using a small, independent labeled calibration set. This estimation is performed in a single post-training step without retraining the models.

Theoretical analysis is conducted under proportional asymptotics ($n, p \to \infty$ with $p/n$ constant) for two settings:

• Ridge Regression: Analyzed under general anisotropic covariance structures and deterministic signals. The authors derive deterministic equivalents for the optimal mixing weight and the resulting prediction risk.
• Logistic Regression: Analyzed in a binary classification setting with hard pseudo-labels, extending the analysis to scenarios where the teacher's error rate is high (even exceeding 50%).

Key Contributions

1. Theoretical Characterization of PMSD Risk: The paper derives exact oracle identities and deterministic equivalents for the optimal PMSD mixing weight and risk in ridge regression. These results hold for general anisotropic covariances where the fresh covariates may have a different distribution (but commuting covariance matrices) than the teacher's training data.
2. Strict Improvement Guarantees: The authors prove that, for almost every pair of teacher and student regularization levels, the optimally mixed PMSD student strictly outperforms the teacher. This holds even when:
   • The fresh covariates are out-of-distribution (e.g., sampled from an isotropic Gaussian distribution).
   • The teacher's regularization level is unknown or suboptimal.
   • The student's regularization is not optimally tuned.
     The only exceptions occur at specific, finite "degeneracy" points where the risks coincide.
3. Non-Monotonicity of Unlabeled Data: Contrary to intuition, the paper demonstrates that increasing the amount of fresh unlabeled data does not monotonically improve performance. In the same-$\lambda$ isotropic setting, the optimal PMSD risk is unimodal with respect to the unlabeled sample size; adding more data beyond a certain point can degrade performance.
4. Consistent Calibration without Retraining: The authors show that the optimal mixing weight cannot be identified from unlabeled data alone (an information-theoretic limitation). However, they prove that a small independent labeled calibration set allows for consistent estimation of the optimal weight in a single post-hoc step, requiring no additional model fitting.
5. Logistic Regression Gains: Extending beyond squared loss, the paper shows that prediction mixing can strictly outperform both the teacher and the pure-distilled student in logistic regression. Notably, even when the teacher and PD student both fail to recover the true labels (achieving 0% accuracy in high-noise regimes), the PMSD student can achieve 100% population accuracy by appropriately extrapolating via the mixing weight.

Results

• Empirical Validation: Experiments on real-world datasets (UCI Blog Feedback, Communities and Crime, Airfoil) and synthetic data confirm the theoretical predictions. The PMSD student consistently achieves lower squared prediction risk than both the teacher and the PD student across various regularization levels and covariance structures (including isotropic and AR1).
• Robustness to OOD: The method remains effective even when fresh covariates are sampled from an isotropic distribution, distinct from the teacher's training distribution.
• Classification Performance: In linear probing experiments on Caltech-101, Caltech-256, and CIFAR-100 with corrupted labels, PMSD consistently improves test accuracy over the teacher and PD student. The optimal mixing weights often lie outside $[0, 1]$, validating the use of affine rather than convex combinations.

Significance and Claims
The paper claims that prediction-only distillation offers a theoretically grounded and practically viable method for improving fixed predictors in data-constrained environments. Its significance lies in:

• Bridging the Data Gap: Providing a mechanism to adapt models when original training data is inaccessible, a common constraint in real-world deployment.
• Generic Improvement: Establishing that strict improvement over the teacher is a generic property of the PMSD scheme, not dependent on fine-tuned hyperparameters or specific distributional alignments.
• Operational Efficiency: Demonstrating that the benefits of distillation can be realized with minimal computational overhead (one-shot calibration) and without access to ground-truth labels for the fresh data.
• Theoretical Novelty: Challenging the assumption that more unlabeled data always helps, revealing a non-monotonic relationship between sample size and risk in the fresh-X setting.

The authors position this work as a complement to existing "same-X" self-distillation literature, highlighting that while same-X methods can recover optimal performance given the original data, PMSD provides the best possible improvement given only the teacher's predictions and fresh covariates.