Optimal Self-Distillation for Rectified Flow via Linear Probing

Modern generative models often rely on synthetic data, which can lead to model collapse or performance degradation. The authors investigate whether a student model, trained on a mixture of ground-truth velocity targets and predictions from a suboptimal rectified flow (RF) teacher, can provably improve upon the teacher's performance without the risks associated with recursive training.

The researchers focus on rectified flow using linear probing, where a feature map is fixed and only a linear head is updated. By keeping the interpolation covariates constant, they derive an exact affine path identity that links the teacher, a pure-distilled refit, and the self-distilled student. This allows the integrated RF risk to be expressed as a quadratic function of the mixing coefficient, enabling a closed-form solution for the optimal mixing weight. They also introduce a one-shot tuning procedure using generalized cross-validation (GCV) to avoid expensive grid searches.

The study establishes that optimal self-distillation strictly improves the integrated velocity risk whenever the teacher is nonstationary along the regularization path. A key discovery is the sign rule: positive mixing coefficients correct under-regularized teachers, while negative mixing coefficients are necessary to correct over-regularized teachers. The authors further prove that this reduction in velocity estimation error tightens the upper bounds for both continuous-time and finite-step Euler generation errors. Experiments across Gaussian models, Gaussian mixtures, and image datasets confirm that this one-shot approach effectively recovers mode information and improves sample quality.