Optimal Self-Distillation for Rectified Flow via Linear Probing
Machine Learning
Summary
The authors study how to improve generative models that learn from their own outputs, focusing on a technique called self-distillation for a method called rectified flow. They prove that by mixing teacher model outputs with true signals in the right way, the student model can reliably outperform the teacher, correcting either too much or too little regularization. Their method includes a formula to find the best mixing balance without expensive tuning. Experiments confirm that this approach leads to better model accuracy and generation quality.
Authors
Saptarshi Roy, Debepsita Mukherjee, Pratik Patil
Abstract
Modern generative models are increasingly trained using model-generated signals, creating both opportunities for self-improvement and risks of collapse. We study optimal self-distillation (SD) for rectified flow (RF): given a suboptimal teacher velocity field, can a student trained on a mixture of true RF velocities and teacher velocities provably improve the teacher? For linear RF with ridge regularization on fixed interpolation pairs, we prove an exact affine path identity, derive the optimal mixing coefficient in closed form, and show strict improvement in integrated velocity risk whenever the teacher risk is nonstationary along the regularization path. The optimal coefficient obeys a sign rule: positive mixing corrects under-regularized teachers, while negative mixing corrects over-regularized teachers. We also give one-shot generalized cross-validation (GCV) and validation tuning procedure that avoids grid search over mixing weights and repeated refitting. Combining this theorem with RF Wasserstein convergence bounds, we show that optimal self-distillation improves the velocity estimation terms controlling continuous-time and finite-step generation error. Experiments with Gaussian models, Gaussian mixtures, and image data show that optimal self-distillation improves velocity risk, mode recovery, and finite-step generation relative to both the teacher and pure distillation.