Manifold-Optimal Guidance: A Unified Riemannian Control View of Diffusion Guidance

1. Het Probleem: Geometrische Mismatch bij Diffusiemodellen

Diffusieprobabilistische modellen zijn de huidige standaard voor generatieve taken (zoals het maken van afbeeldingen en video's). Om deze modellen gestuurd te laten werken (bijvoorbeeld op basis van een tekstprompt), wordt Classifier-Free Guidance (CFG) gebruikt. CFG werkt door de onvoorwaardelijke en voorwaardelijke scores (gradiënten) lineair te extrapoleren met een schaalparameter $w$.

Hoewel CFG effectief is, leidt het gebruik van hoge waarden voor $w$ vaak tot ernstige kwaliteitsproblemen:

• Oversaturatie en textuurartefacten: Beelden worden onnatuurlijk helder of "gebrand".
• Structurele ineenstorting: De structuur van het gegenereerde beeld vervormt.
• De oorzaak: De auteurs identificeren dit als een geometrische mismatch. CFG behandelt de latente ruimte als een isotrope Euclidische ruimte. Door de scores lineair te extrapoleren, duwt CFG de sampling-trajecten echter weg van de hoog-dichtheid data-manifold (de verzameling van realistische data). Dit fenomeen noemen ze "off-manifold drift". Hoe sterker de guidance, hoe sneller het traject de realistische data-verdeling verlaat, wat leidt tot de genoemde artefacten.

Bestaande oplossingen (zoals dynamische drempelwaarden of projectie-methoden) behandelen vaak slechts de symptomen en zijn niet gebaseerd op een uniek, geometrisch onderbouwd doel.

2. Methodologie: Manifold-Optimal Guidance (MOG)

De auteurs introduceren Manifold-Optimal Guidance (MOG), een raamwerk dat guidance herformuleert als een lokaal optimaal control-probleem op de data-manifold, gebruikmakend van Riemanniaanse meetkunde.

Kernconcepten:

• Riemanniaanse Optimalisatie: In plaats van de Euclidische gradiënt te volgen, minimaliseert MOG de conditionele energie terwijl het de sampling-trajecten dicht bij de data-manifold houdt. Dit wordt gedaan door een Riemanniaanse metriek ($M_t$) in te voeren die bewegingen die van de manifold afwijzen, straft.
• De Update-regel: De optimale update $u_t$ wordt afgeleid als de Riemanniaanse natuurlijke gradiënt:
  $$s_{MOG} = s_0 + \beta(t) M_t^{-1} \Delta s$$
  Waarbij:
  • $s_0$ de onvoorwaardelijke score is.
  • $\Delta s$ het verschil is tussen voorwaardelijke en onvoorwaardelijke scores.
  • $M_t^{-1}$ de inverse metriek is die de update "voorbereidt" (preconditions) zodat deze de manifold-geometrie respecteert.
  • $\beta(t)$ de sterkte van de guidance regelt.

Praktische Implementaties:

De auteurs presenteren twee varianten die geen hertraining vereisen:

1. MOG-Score: Gebruikt een anisotrope metriek gebaseerd op de onvoorwaardelijke score $s_0$. De score wordt gezien als de normaalvector naar de manifold. De metriek straft bewegingen langs deze normaal (off-manifold) zwaar, terwijl bewegingen langs de manifold (tangentiële richting) vrij worden gelaten. Dit wordt efficiënt berekend met de Sherman-Morrison-formule ($O(d)$ complexiteit).
2. Auto-MOG: Een dynamische schedulingsmethode die de guidance-strength $\beta(t)$ automatisch aanpast. Het baseert dit op een energie-balans: de guidance-energie moet evenredig zijn met de prior-energie. Dit elimineert de noodzaak voor handmatige hyperparameter-tuning en past zich automatisch aan aan de moeilijkheidsgraad van de prompt en het ruisniveau.

3. Belangrijkste Bijdragen

• Unificatie van Guidance: Het formuleren van guidance als een variatie-optimalisatieprobleem onder een Riemanniaanse metriek, wat leidt tot een gesloten-formule oplossing die optimaal is ten opzichte van een geodesische kost.
• Geometrisch Bewuste Updates: Het introduceren van een methode die off-manifold drift onderdrukt zonder de efficiëntie van de afdaalrichting te verliezen.
• Training-vrije en Efficiënte Implementatie: Zowel MOG-Score als Auto-MOG zijn plug-and-play modules die werken met bestaande diffusiemodellen en geen extra rekentijd vereisen (behalve een verwaarloosbare overhead voor vector-bewerkingen).
• Auto-MOG: Een adaptieve regeling die de manuele zoektocht naar de juiste guidance-schaal ($w$) overbodig maakt.

4. Resultaten

De auteurs hebben hun methode uitgebreid getest op zes verschillende diffusiemodellen (waaronder DiT, EDM2, SD-2.1, SD-XL, SD-3.5 en FLUX.1) en diverse datasets (ImageNet, MS-COCO).

• Kwaliteit en Fidelity: MOG (en specifiek Auto-MOG) bereikt state-of-the-art resultaten op FID (Fréchet Inception Distance), wat aangeeft dat de gegenereerde beelden dichter bij de echte data-verdeling liggen dan bij baselines.
• Artefactenreductie: Er is een aanzienlijke reductie in oversaturatie en textuurartefacten, zelfs bij hoge guidance-waarden waar standaard CFG faalt.
• Alignement: De methode behoudt of verbetert de tekst-beeld alignement (gemeten via CLIP-score en Human Preference Score - HPSv2).
• Gebruikersstudie: In een blind vergelijkend onderzoek (N=37) gaf de meerderheid van de deelnemers (52,5% - 70,6%) de voorkeur aan Auto-MOG boven standaard CFG en andere geavanceerde methoden (zoals CFG++ en APG), vooral op het gebied van kleurtrouw en realisme.
• Efficiëntie: De methode voegt vrijwel geen extra rekentijd toe, omdat de berekeningen lineair zijn met de dimensie van de latent space en de inferentie van het netwerk de bottleneck blijft.

5. Betekenis en Impact

Dit werk biedt een fundamenteel nieuw perspectief op het sturen van diffusiemodellen. In plaats van te vertrouwen op heuristieken of post-hoc correcties, biedt MOG een wiskundig onderbouwd, geometrisch consistent raamwerk.

De implicaties zijn groot:

• Het lost het fundamentele probleem op van "off-manifold drift" bij hoge guidance-waarden.
• Het maakt het mogelijk om zeer sterke guidance te gebruiken voor betere tekst-alignement zonder in te leveren op beeldkwaliteit.
• Het elimineert de noodzaak voor handmatige tuning van guidance-schalen, wat de bruikbaarheid van diffusiemodellen voor eindgebruikers en ontwikkelaars vergroot.
• Het raamwerk is breed toepasbaar op verschillende architecturen (UNet, Transformers, Rectified Flow), wat het een universele oplossing maakt voor het huidige ecosysteem van generatieve AI.

Kortom, MOG transformeert guidance van een ruwe "knop" naar een geoptimaliseerd, geometrisch bewust besturingsmechanisme dat de intrinsieke structuur van de data respecteert.