Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

Probleemstelling

Consistency Models (CMs) zijn een nieuwe generatie generatieve modellen die beelden kunnen genereren in slechts één of enkele stappen, wat ze aanzienlijk sneller maakt dan traditionele Diffusion Models (DMs). Een cruciale techniek voor het beheersen van de generatie in DMs is Classifier-Free Guidance (CFG). CFG stelt gebruikers in staat om een afweging te maken tussen de trouw (fidelity) aan de prompt en de diversiteit van de gegenereerde samples door een "guidance scale" ($\omega$) te gebruiken.

Echter, er bestaat een significant probleem bij het toepassen van CFG op CMs:

1. Afhankelijkheid van een leraar: Bestaande methoden voor CFG in CMs vereisen kennisdistillatie van een vooraf getrainde DM als "leraar". Dit beperkt CFG tot CMs die via Consistency Distillation (CD) zijn getraind.
2. Gebrek aan post-hoc oplossing: CMs die via Consistency Training (CT) zijn getraind (een volledig datagedreven methode zonder DM-leraar), kunnen momenteel geen effectieve guidance ontvangen zonder opnieuw te worden getraind met een DM-leraar.
3. Complexiteit: Het trainen van een nieuwe CM met guidance is complexer dan het simpele aanpassen van een ODE-traject, omdat het een aparte, ongeleide versie vereist voor een eerlijke vergelijking.

Het doel van dit paper is een methode te ontwikkelen die post-hoc guidance mogelijk maakt voor reeds getrainde CMs (zowel CD als CT), zonder afhankelijk te zijn van een externe DM-leraar.

Methodologie: Joint Flow Distribution Learning (JFDL)

De auteurs introduceren Joint Flow Distribution Learning (JFDL), een lichtgewicht alignatiemethode die guidance mogelijk maakt door de richting van geconditioneerde en onvoorwaardelijke verdelingen te interpoleren.

Kernconcepten:

• CM als ODE-oplosser: De auteurs behandelen een vooraf getrainde CM als een solver voor de Probability Flow ODE (PF-ODE). Ze verifiëren dat de snelheidsvelden (velocity fields) die voortkomen uit onvoorwaardelijke en conditionele verdelingen, een Gaussische ruisimplicatie hebben.
• Hybride ODE-oplossing: In plaats van een aparte onvoorwaardelijke DM te gebruiken, construeert JFDL een "hybride" onvoorwaardelijke ankerpunt ($y^{\emptyset, t}_0$). Dit wordt gedaan door een sample $x^c_0$ eerst naar een ruisniveau $t$ te brengen via de conditionele stroom, en dit vervolgens terug te projecteren naar $t=0$ via de onvoorwaardelijke stroom (of een geschatte variant daarvan).
• Pseudo-Ruis Analyse: De auteurs tonen theoretisch en empirisch aan dat de "gemengde signalen" die ontstaan bij deze hybride projectie, gedomineerd worden door de Gaussische ruis ($\sigma_{max}z$). Hierdoor gedraagt de onvoorwaardelijke richting zich als een zuivere Gaussische verdeling, wat essentieel is voor de geldigheid van CFG.
• De Loss Functie: JFDL introduceert een extra loss-term ($L_{JFDL}$) naast de standaard Consistency Training loss ($L_{ECT}$). Deze loss forceert het model om te leren in de richting van een geëxtrapoleerde vector:
  $$u_\omega = \omega \cdot u_{cls} + (1 - \omega) \cdot u_{\emptyset}$$
  Waarbij $u_{cls}$ de conditionele denoising-richting is en $u_{\emptyset}$ de onvoorwaardelijke (of geschatte) richting.

Varianten:

1. Naive JFDL: Vereist dat het CM al getraind is om de onvoorwaardelijke ($\emptyset$) ODE op te lossen.
2. Random JFDL: Een innovatieve variant die de noodzaak van een expliciet $\emptyset$-solver elimineert. In plaats van een $\emptyset$-label te gebruiken, wordt een willekeurig label ($c'$) uit de dataset getrokken om de "onvoorwaardelijke" anker te construeren. Dit werkt verrassend goed en is toepasbaar op modellen die puur conditioneel zijn getraind.

Belangrijkste Bijdragen

1. Onafhankelijke Post-Hoc Guidance: JFDL is de eerste methode die effectieve guidance mogelijk maakt voor CMs zonder een DM-leraar. Dit sluit de kloof voor CT-modellen die eerder geen guidance konden gebruiken.
2. Theoretische Onderbouwing: De auteurs bieden inzicht in waarom JFDL werkt door de connectie met Flow-based generative models (FMs) te leggen en te bewijzen dat de hybride stroom de marginale data-verdeling behoudt en dat de resulterende "pseudo-ruis" Gaussisch is.
3. Flexibiliteit: De methode werkt zelfs op CMs die niet expliciet zijn ontworpen voor onvoorwaardelijke sampling (via Random JFDL).
4. Efficiëntie: Het fine-tunen met JFDL vereist slechts een klein deel van de trainingsdata en GPU-tijd vergeleken met het volledig opnieuw trainen van een model met guidance.

Resultaten

De methode is getest op CIFAR-10 en ImageNet 64x64 met als baseline het ECT (Consistency Training) model.

• FID Verbetering:
  • Op CIFAR-10 (1-stap sampling) verbeterde JFDL de FID van 3.40 (baseline) naar 3.29 (Naive) en 3.24 (Random) bij een optimale guidance scale ($\omega \approx 1.15$).
  • Op ImageNet 64x64 (1-stap) verbeterde de FID van 5.84 naar 4.38 (Naive) en 4.68 (Random).
• Gedrag van Guidance: De FID-curves tonen het klassieke gedrag van CFG: verbetering van de beeldkwaliteit bij lage tot middelhoge guidance waarden ($\omega < 2$), gevolgd door een daling bij te hoge waarden door verlies aan diversiteit.
• Kwalitatieve Resultaten: Visualisaties tonen dat JFDL leidt tot scherper contrast en betere perceptuele kwaliteit (bijv. duidelijker details in "hotdogs" op ImageNet), vergelijkbaar met CFG in Diffusion Models.
• Efficiëntie: Het fine-tuning proces verbruikte slechts ongeveer 7,5% van de trainingsdata en 30% van de GPU-uren van de oorspronkelijke ECT-training.

Significantie

Dit paper is een doorbraak voor de efficiëntie en bruikbaarheid van Consistency Models:

• Democratisering van Guidance: Het maakt geavanceerde controle over generatieve processen toegankelijk voor modellen die via snelle, datagedreven methoden (CT) zijn getraind, zonder de noodzaak van zware DM-distillatie.
• Post-Hoc Aanpasbaarheid: Het biedt een praktische route om bestaande, getrainde CMs "op te fleuren" met guidance-knoppen, wat essentieel is voor praktische toepassingen waar hertraining onhaalbaar is.
• Theoretische Unificatie: Het paper versterkt het theoretische fundament door Diffusion Models, Flow-based Models en Consistency Models te verenigen in één raamwerk voor guidance, wat nieuwe inzichten biedt voor toekomstig onderzoek naar snelle generatieve modellen.

Kortom, JFDL lost een fundamentele beperking op in de Consistency Model-ecosysteem en maakt ze even controleerbaar als Diffusion Models, maar met de snelheid van één-stap generatie.