Technisch overzicht: Optimale zelf-distillatie voor Rectified Flow via Lineaire Probing

Probleemstelling

Moderne generatieve modellen vertrouwen steeds vaker op door het model gegenereerde signalen voor training, wat een tweesnijdend zwaard creëert: hoewel zelfverbetering mogelijk is, brengt recursieve training op synthetische data het risico op "model collapse" met zich mee, gekenmerkt door verlies van diversiteit en distributionele drift. Daartegenover staat zelf-distillatie (SD) in supervised learning—het hertrainen van een model op een mengeling van grondwaarheid-labels en voorspellingen van de leraar (teacher)—wat vaak prestatiewinst oplevert.

Dit onderzoek onderzoekt of er een bewijsbaar zelfverbeteringsmechanisme bestaat voor Rectified Flow (RF), een generatief modelleringsframework dat een tijd-geconditioneerd snelheidsveld leert om bronruis naar data te transporteren via een Ordinary Differential Equation (ODE). Specifiek vragen de auteurs: gegeven een suboptimale RF-teacher velocity field, kan een student die getraind is op een mengeling van ware RF-velocities en teacher-velocities de teacher aantoonbaar verbeteren?

Methodologie

De studie richt zich op lineaire RF-modellen getraind op vaste interpolatieparen (vaste bron-ruis en data-eindpunten). De methodologie maakt gebruik van lineaire probing, waarbij een vaste feature map $\phi(t, x)$ wordt gebruikt en alleen een lichtgewicht lineaire kop $W$ wordt bijgewerkt.

1. Opzet:

   • Teacher: Een ridge-geregulariseerd RF-model $\hat{W}_\lambda$ dat getraind is op ware velocity-targets $u = X_1 - X_0$.
   • Pure Distillation (PD): Een student die enkel wordt hertraind op de voorspellingen $\hat{u}$ van de teacher.
   • Self-Distillation (SD): Een student die wordt getraind op een gemengde target $u(\xi) = (1-\xi)u + \xi\hat{u}$, waarbij $\xi \in \mathbb{R}$ een mengcoëfficiënt is.

2. Theoretisch Kader:

   • De auteurs bewijzen dat voor lineaire RF met ridge-regularisatie de zelf-gedistilleerde student op een exact affien pad ligt tussen de teacher en de pure-distilled refit.
   • Deze geometrische eigenschap impliceert dat de geïntegreerde RF-risico (velocity-schattingsfout) een exacte kwadratische functie is van de mengcoëfficiënt $\xi$.
   • Consequentelijk kan de optimale mengcoëfficiënt $\xi^*$ in gesloten vorm worden afgeleid.

3. Tuning Procedure:

   • Om grid search en herhaaldelijk refitten te vermijden, stelt het paper een one-shot tuning procedure voor met behulp van Generalized Cross-Validation (GCV) of een enkele validatiestap. Dit schat de optimale $\xi$ in zonder de student voor elke kandidaat-gewicht opnieuw te trainen.

Belangrijkste Bijdragen

1. Theoretische Garanties voor Optimale SD

• Affine Path Identity: Het paper stelt vast dat $\hat{W}_{sd, \lambda, \xi} = (1-\xi)\hat{W}_\lambda + \xi\hat{W}_{pd, \lambda}$. Dit koppelt de teacher, de pure-distilled refit en de zelf-gedistilleerde student lineair aan elkaar.
• Strict Improvement Theorem: Het geïntegreerde RF-risico wordt geminimaliseerd bij een $\xi^*$ in gesloten vorm. De auteurs bewijzen dat strikte verbetering ten opzichte van de teacher gegarandeerd is wanneer de RF-risico niet-stationair is langs het ridge-regularisatiepad (d.w.z. $R'(\lambda) \neq 0$).
• De Sign Rule: De optimale coëfficiënt volgt een specifieke tekenregel:
  • Positieve $\xi^*$: Corrigeert onder-geregulariseerde teachers (waar het risico afneemt naarmate de regularisatie toeneemt).
  • Negatieve $\xi^*$: Corrigeert over-geregulariseerde teachers (waar het risico toeneemt naarmate de regularisatie toeneemt). Cruciaal is dat dit de student in staat stelt om weg van het pure-distilled veld te extrapoleren om herhaalde krimp (shrinkage) tegen te gaan, een mechanisme dat verschilt van standaard convexe distillatie.

2. Implicaties voor Generatiefouten

Het paper verbindt de reductie in velocity-risico aan generatieve prestatiebounds:

• Continu-tijd: Met behulp van Wasserstein-convergentiebounds laten de auteurs zien dat optimale SD de bovenste grens op de continue-tijd generatiefout strikt verbetert ten opzichte van zowel de teacher als pure distillatie.
• Finite-Step Sampling: Voor Euler-discretisatie vermindert de optimale SD de velocity-schattingsfoutterm in de generatiebound, wat theoretisch de foutenbound voor finite-step sampling aanscherpt.

3. Empirische Validatie

Experimenten werden uitgevoerd op Gaussische modellen, Gaussische mengsels en echte beelddatasets (Handwritten digits, Fashion-MNIST, CIFAR-10):

• Gaussische Modellen: In over-geregulariseerde regimes verslechterde pure distillatie de prestaties, terwijl optimale SD (met negatieve $\xi$) het velocity-risico aanzienlijk verminderde en de mode-recovery verbeterde.
• Real Images: Gecontroleerde stress-tests omvatten het omlaag schalen van de output van de teacher om miskalibratie te simuleren. Pure distillatie versterkte de degradatie, terwijl optimale SD negatieve mengcoëfficiënten selecteerde om herkenbare samples te herstellen.
• Neurale Extensies: Op CIFAR-10 met een time-conditioned U-Net verbeterde de methode de FID- en KID-scores zelfs voorbij de lineaire probing setting, hoewel de strikte theoretische garanties gelden voor het lineaire geval.

Resultaten

• Risicoreductie: In Gaussische settings verminderde optimale SD het velocity-risico aanzienlijk (bijv. van 1.67 naar 1.14 in specifieke over-geregulariseerde gevallen) vergeleken met de teacher.
• Mode Recovery: In niet-lineaire Gaussische mengsels produceerde de teacher vervaagde distributies, terwijl optimale SD met negatieve mixing de afzonderlijke modi herstelde (sliced $W_2^2$ verbeterde van ~1.02 naar ~0.09).
• Beeldgeneratie: Op Fashion-MNIST en CIFAR-10 herstelde optimale SD hoogwaardige samples van gedegradeerde teachers, terwijl pure distillatie en geconstrangeerde SD (geclipt naar $[0,1]$) er niet in slaagden de miskalibratie te corrigeren.
• One-Shot Tuning: De GCV-gebaseerde one-shot tuning volgde de oracle optimale mixing nauwgezet, wat de efficiëntie van de voorgestelde procedure valideert.

Betekenis en Claims

Het paper claimt een one-shot mechanisme te hebben geïdentificeerd waarmee RF-modellen kunnen verbeteren van hun eigen voorspellingen zonder recursief de eindpunten-koppelingen te veranderen (in tegenstelling tot Reflow).

• Theoretische Vernieuwing: Het draagt de structurele geometrie van optimale onbeperkte zelf-distillatie uit supervised ridge regressie over naar RF velocity matching, waarbij exacte affine pad-identiteiten en closed-form optimaliteit worden geboden.
• Praktisch Inzicht: Het werk benadrukt dat negatieve mixing essentieel is voor het corrigeren van over-geregulariseerde teachers, een bevinding die de conventionele convexe distillatie-paradigma uitdaagt waar menggewichten doorgaans beperkt zijn tot $[0,1]$.
• Beperkingen en Reikwijdte: De auteurs zijn bescheiden over hun claims. De strikte verbeteringsstelling is exact alleen voor lineaire RF met ridge-regularisatie. Voor neurale netwerken wordt de affine familie beschouwd als een empirische extensie. Het paper claimt niet dat elke downstream metriek zal verbeteren onder willekeurige misspecificatie, en merkt op dat finite-step generatie ook afhangt van Lipschitz-constanten, flow-rechtlijnigheid en ODE-stabiliteit. De primaire garantie is de reductie van het geïntegreerde RF-velocity risico, wat theoretisch de generatiefoutbounds aanscherpt.

Samenvattend biedt het paper een theoretisch onderbouwde, computationeel efficiënte strategie om suboptimale RF-teachers te "repareren" door ware en voorspelde velocities te mengen, waarbij wordt aangetoond dat optimale zelf-distillatie de velocity-schatting en daarmee de generatieve prestaties strikt kan verbeteren.