Optimal Self-Distillation for Rectified Flow via Linear Probing

2026-07-16

AI Summary

This work addresses self-distillation under suboptimal teacher models, aiming to prevent performance degradation and achieve strict improvement over the teacher. The authors propose training the student model by blending the ground-truth and teacher velocity fields within a linear rectified flow augmented with ridge regularization, and derive a closed-form expression for the optimal mixing coefficient. The sign of this coefficient serves as a diagnostic indicator distinguishing under-regularization from over-regularization, enabling a one-shot, grid-search-free hyperparameter tuning strategy. Combining generalized cross-validation with Wasserstein convergence analysis, the theoretical framework significantly reduces velocity field risk and enhances mode recovery and finite-step generation quality across Gaussian, Gaussian mixture, and image datasets.

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.

Problem

• self-distillation
• rectified flow
• velocity field
• model collapse
• generative models

Innovation

• optimal self-distillation
• rectified flow
• velocity field
• generalized cross-validation
• ridge regularization