Technische Samenvatting: FlowConsist

Probleemstelling

Snelle flow-modellen, waaronder Consistency Models en MeanFlow, beogen het generatieve samplingsproces te versnellen door direct de integraal van de Ordinary Differential Equation (ODE) paden te leren voorspellen, wat één-stap of enkele-stap generatie mogelijk maakt. De auteurs stellen echter dat huidige trainingsparadigma's lijden onder twee fundamentele problemen die de prestaties verslechteren:

1. Systematische Trajectdrift: Huidige methoden construeren conditionele snelheden door willekeurig ruis ($\epsilon$) en data ($x$) samples aan elkaar te koppelen. Hoewel de verwachting van deze conditionele snelheden gelijk is aan de marginale snelheid, snijden individuele conditionele paden meerdere verschillende marginale trajecten. Training op deze conditionele snelheden zorgt ervoor dat het optimalisatiedoel afdrijft van het ware, unieke marginale ODE-pad, wat leidt tot inconsistente trajecten.
2. Foutaccumulatie: Benaderingsfouten die worden geïntroduceerd door het model accumuleren over grote tijdsintervallen. In snelle flow-modellen die vertrouwen op zelf-voorspelde snelheden voor training, stapelen deze fouten zich op, wat leidt tot significante afwijkingen tussen de gegenereerde samples en de grondwaarheid van de datadistributie, met name bij single-step (1-NFE) generatie.

Methodologie

De auteurs stellen FlowConsist voor, een trainingsframework ontworpen om trajectconsistentie af te dwingen en foutaccumulatie te corrigeren zonder een vooraf getraind docentmodel (teacher model) te vereisen.

1. Afdwingen van Trajectconsistentie (Vervanging van Conditionele Snelheden)

De kern van het theoretische inzicht is dat het vertrouwen op conditionele snelheden ($v_t = \epsilon - x$) een variantieterm ($L_{var}$) in de verliesfunctie introduceert, die het model dwingt om gradiënten in hoog-variante richtingen te onderdrukken en het getrouw fitten van gekromde trajecten verhindert.

• Theoretische Verschuiving: In plaats van de willekeurig gekoppelde conditionele snelheid te gebruiken als surrogaat voor de marginale snelheid, vervangt FlowConsist de conditionele snelheid in de consistentie-objectief door de marginale snelheid die door het model zelf wordt voorspeld ($u_\theta$).
• Implementatie: De auteurs leiden een nieuwe trainingsdoelstelling af (Eq. 7) die een standaard Flow Matching-term combineert met een trajectconsistentie-regularisatieterm.
  • De doelstelling minimaliseert het verschil tussen de voorspelde gemiddelde snelheid en een doel dat afgeleid is van de eigen marginale snelheidspredictie van het model.
  • Cruciaal is dat de conditionele snelheid alleen wordt gebruikt in een term die convergeert naar de marginale snelheid in verwachting zodra het model getraind is, waardoor de systematische drift tijdens het optimalisatieproces effectief wordt geëlimineerd.
• Classifier-Free Guidance (CFG): De methode incorporeert CFG door de guidance scale $\omega$ te behandelen als een secundaire input, waardoor het model een verenigde representatie kan leren die tijdens de inferentie aangepast kan worden.

2. Trajectrectificatie (Aanpakken van Foutaccumulatie)

Om de monotone divergentie aan te pakken die wordt veroorzaakt door cumulatieve benaderingsfouten over tijdstappen heen, introduceren de auteurs een trajectrectificatiestrategie.

• Marginale Snelheid Alignment: In tegen tegenstelling tot eerdere methoden die alleen de initiële ruis corrigeren, brengt FlowConsist de marginale distributie van door het model gegenereerde samples in lijn met de echte datadistributie op elk tijdstip langs het traject.
• Teacher-Free Mechanisme: In plaats van te vertrouwen op een vooraf getraind docentmodel (zoals in DMD) om echte marginale snelheden te schatten, gebruikt FlowConsist een hulpmodel ($G_\psi$) dat getraind is om de marginale snelheid van door het hoofdmodel ($F_\theta$) gegenereerde samples te schatten.
• Correctieproces: Het hoofdmodel wordt getraind om zijn voorspelde marginale snelheid uit te lijnen met de snelheid die door het hulpmodel wordt geschat. Dit "trekt" de voorspellingen van het model effectief terug naar de echte datadistributie, waardoor de cumulatieve fouten die anders zouden accumuleren in een single-step mapping, worden gecorrigeerd.

Belangrijkste Bijdragen

1. Identificatie van Trajectdrift: Het artikel levert een theoretisch bewijs (Theorem 1 & 2) dat aantoont dat conditionele snelheden in Flow Matching inherent een niet-nul covariantie bezitten, wat leidt tot onvermijdelijke trajectdrift wanneer ze worden gebruikt als instantane snelheidsschattingen in snelle flow-modellen.
2. Nieuwe Trainingsdoelstelling: De auteurs stellen een principieel alternatief voor (Eq. 7) dat conditionele snelheden vervangt door zelf-voorspelde marginale snelheden, waardoor de optimalisatie verankerd blijft aan een consistent marginaal traject.
3. Trajectrectificatiestrategie: Een teacher-free methode (Eq. 11) wordt geïntroduceerd om marginale distributies op elk tijdstip uit te lijnen, wat de accumulatie van benaderingsfouten over lange intervallen vermindert.
4. State-of-the-Art Prestaties: De methode bereikt een nieuwe state-of-the-art op ImageNet 256×256 met een Fréchet Inception Distance (FID) van 1.52 met slechts 1 sampling stap (1-NFE), zonder distillatie van een docentmodel.

Experimentele Resultaten

• Dataset & Setup: Experimenten werden uitgevoerd op ImageNet 256×256 met een DiT-architectuur en een vooraf getrainde VAE latent space.
• Ablatieonderzoeken:
  • Het vervangen van de standaard MeanFlow loss door de voorgestelde marginale trajectconsistentie loss (Eq. 7) verbeterde de FID van 5.86 naar 4.99.
  • Het toevoegen van de marginale snelheid alignment loss (Eq. 11) verlaagde de FID verder naar 4.31.
  • De methode toonde robuustheid over verschillende CFG-schalen en integratieratio's van de twee verliescomponenten.
• Vergelijking: FlowConsist-XL/2 bereikte een FID van 1.52 (1-NFE), waarmee het de vorige beste non-distilled fast flow model (iMF-XL/2 met 1.72) overtreft en de prestaties van gedistilleerde methoden benadert. Het was ook superieur aan multi-step diffusiemodellen in termen van efficiëntie (1 stap vs. 250+ stappen) terwijl het concurrerende kwaliteit behield.

Betekenis

Het artikel beweert dat FlowConsist een nieuw perspectief biedt voor snelle flow-modellering door de grondoorzaken van prestatiebeperkingen in bestaande frameworks aan te pakken: de theoretische mismatch tussen conditionele en marginale snelheden, en de praktische accumulatie van benaderingsfouten. Door deze problemen te corrigeren, tonen de auteurs aan dat high-fidelity, single-step generatie haalbaar is door enkel verfijnde trainingsdoelstellingen, waarmee de kloof tussen teacher-free en teacher-distilled paradigma's wordt overbrugd. Het werk benadrukt dat trajectconsistentie een kerneis is voor snelle flow-modellen om succesvol lineaire mappings tussen tijdstappen uit te voeren.