Optimal Self-Distillation for Rectified Flow via Linear Probing

2026-07-16 · Source: Machine Learning · Field: Technology & Digital — Artificial Intelligence & Machine Learning · Depth: Expert, quick

Summary

Optimal Self-Distillation for Rectified Flow via Linear Probing introduces a method to provably enhance modern generative models, specifically rectified flow (RF), by training a student model on a mixture of true RF velocities and suboptimal teacher velocities. The research establishes an exact affine path identity for linear RF with ridge regularization, deriving a closed-form optimal mixing coefficient. This coefficient strictly improves integrated velocity risk when the teacher risk is nonstationary. A key finding is a sign rule: positive mixing corrects under-regularized teachers, while negative mixing addresses over-regularized ones. The authors also present a one-shot generalized cross-validation (GCV) and validation tuning procedure to avoid extensive grid search. This optimal self-distillation improves velocity estimation terms, controlling continuous-time and finite-step generation error, as demonstrated through experiments with Gaussian models, Gaussian mixtures, and image data, showing improvements over both the teacher and pure distillation.

Key takeaway

For machine learning engineers developing generative models based on rectified flow, this research offers a direct path to improving model performance. You should investigate implementing optimal self-distillation to enhance velocity estimation, mode recovery, and finite-step generation quality, especially when starting with suboptimal teacher models. Utilize the proposed one-shot generalized cross-validation (GCV) procedure to efficiently determine the optimal mixing coefficient, streamlining your development workflow and avoiding costly hyperparameter grid searches.

Key insights

Optimal self-distillation provably improves rectified flow models by deriving and applying an optimal mixing coefficient for teacher and true velocities.

Principles

• Optimal mixing coefficient corrects regularization state.
• Positive mixing for under-regularized teachers.
• Negative mixing for over-regularized teachers.

Method

Derive an optimal mixing coefficient using an affine path identity for linear rectified flow, then apply a one-shot generalized cross-validation (GCV) or validation tuning procedure.

In practice

• Improve rectified flow velocity risk.
• Enhance mode recovery and generation quality.
• Avoid grid search for mixing weights.

Topics

• Rectified Flow
• Self-Distillation
• Generative Models
• Linear Probing
• Velocity Estimation
• Model Improvement

Best for: Research Scientist, Computer Vision Engineer, AI Scientist, Machine Learning Engineer

Related on AIssential

• Prediction-Only Distillation in Linear and Logistic Regression — Prediction-only distillation improves model risk by combining teacher and student predictions, even with out-of-distribution data.
• Velocity Scheduled Flow Matching — Velocity Scheduled Flow Matching improves sample quality and efficiency by dynamically adjusting the interpolation speed between noise and data.
• Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation — Wasserstein Gradient Flows offer a scalable, regularized approach to computing Wasserstein barycenters.
• Hugging Face Journal Club: Embarrassingly Simple Self-Distillation Improves Code Generation — Self-distillation via temperature-tuned SFT improves model performance by refining token distribution.
• SGMD: Score Gradient Matching Distillation for Few-Step Video Diffusion Distillation — SGMD accelerates video diffusion distillation by directly optimizing fake scores with stable Fisher objectives and dual potentials.
• Self-Distillation is Optimal Among Spectral Shrinkage Estimators in Spiked Covariance Models — Self-distillation is statistically optimal among spectral shrinkage estimators in spiked covariance models.

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Machine Learning.