Technische Samenvatting: Verbeteren van Rectified Flow met Randvoorwaarden

Probleemstelling

Rectified Flow (RF), ook bekend als Flow Matching, is een krachtig generatief modelleringsframework dat een snelheidsveld leert om een ruisverdeling te transformeren naar een doelverdeling van data via een gewone differentiaalvergelijking (ODE). Ondanks het succes in taken zoals tekst-naar-beeld en video-generatie, identificeren de auteurs een kritieke beperking in standaard RF-implementaties: het geleerde snelheidsveld voldoet vaak niet aan de theoretisch afgeleide randvoorwaarden.

Specifiek, onder lineaire interpolatie tussen ruis ($t=0$) en data ($t=1$), moet het optimale snelheidsveld voldoen aan:

1. Linker randvoorwaarde ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Rechter randvoorwaarde ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

Standaard neurale netwerkparameterisaties van het snelheidsveld dwingen deze beperkingen niet inherent af. Zoals in het artikel gevisualiseerd, wijkt de voorspelde snelheid bij $t=1$ vaak af van de verwachte identiteitsafbeelding. Deze schending is bijzonder schadelijk tijdens stochastische sampling (bijv. met Langevin-dynamica). Aangezien de scorefunctie-schatting (via Tweedie's formule) een deling door $(1-t)$ bevat, veroorzaakt een onnauwkeurige snelheid nabij $t=1$ dat de score-schatting divergeert. Dit leidt tot instabiele sampling, wat resulteert in overgesmoorde beelden of "cartoonachtige" artefacten, vooral bij het gebruik van grotere stapgroottes of ruisschalen.

Methodologie

Om dit aan te pakken, stellen de auteurs het Boundary-enforced Rectified Flow Model (Boundary RF Model) voor. Deze aanpak incorporeert expliciet randvoorwaarden in de parameterisatie van het snelheidsveld via twee afzonderlijke varianten:

1. Masker-gebaseerd Boundary RF Model

Deze variant dwingt beide randvoorwaarden ($t=0$ en $t=1$) af door het snelheidsveld te structureren als een lineaire combinatie van randtermen en een leerbaar neuraal netwerkoutput:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Waarbij:

• $m_\theta(\mathbf{x}, t)$ een standaard neuraal netwerk is (bijv. U-Net of DiT).
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ scalaire functies zijn die ontworpen zijn om te verdwijnen of te activeren bij specifieke grenzen (bijv. $g(1)=f(0)=h(0)=h(1)=0$ en $g(0)=f(1)=1$).
• De auteurs suggereren een "Standaard Cosinus" keuze: $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, en $h(t) = \sin(\pi t)$.

Deze constructie garandeert dat de randvoorwaarden door ontwerp worden voldaan, ongeacht de output van het neurale netwerk.

2. Subtractie-gebaseerd Boundary RF Model

Gesterkt door de observatie dat het afdwingen van de $t=1$ conditie alleen al significante empirische voordelen oplevert, stellen de auteurs een eenvoudigere, enkele-randvariant voor:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
Deze formulering voldoet inherent aan $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ omdat de term $m_\theta(\mathbf{x}, 1)$ wegvalt bij $t=1$. Hoewel deze aanpak ontwerpsimpliciteit en flexibiliteit biedt, brengt het een computationele kostenpost met zich mee, aangezien er een extra forward pass nodig is om $m_\theta(\mathbf{x}, 1)$ voor elke evaluatie te berekenen.

Impact op Stochastische Sampling

Door $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ af te dwingen, zorgt het Boundary RF Model ervoor dat de teller in Tweedie's formule ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) naar nul nadert wanneer $t \to 1$. Dit voorkomt dat de scorefunctie $\nabla \log \rho_t$ divergeert, waardoor de stochastische samplers (SDE samplers) worden gestabiliseerd en robuuste generatie mogelijk wordt, zelfs met grotere ruisschalen of stapgroottes.

Belangrijkste Bijdragen

1. Identificatie van een Beperking: Het artikel benadrukt dat onbeperkte snelheidsvelden in Rectified Flow de theoretische randvoorwaarden schenden, wat leidt tot onnauwkeurige ODE-trajecten en instabiele stochastische sampling.
2. Nieuwe Parameterisatie: De introductie van het Boundary RF Model, dat twee praktische varianten (Masker-gebaseerd en Subtractie-gebaseerd) biedt om deze condities af te dwingen met minimale wijzigingen in de code.
3. Stabilisatie van Scorefuncties: Het aantonen dat het afdwingen van randvoorwaarden de scorefunctie nabij $t=1$ direct stabiliseert, wat de divergentieproblemen die de vanilla RF teisteren, vermindert.
4. Empirische Validatie: Uitgebreide experimenten die consistente verbeteringen laten zien in de kwaliteit van generatie en de stabiliteit van sampling over verschillende datasets en modelformaten.

Experimentele Resultaten

De auteurs hebben hun modellen geëvalueerd op CIFAR-10 en ImageNet (256$\times$256 en 512$\times$512) met behulp van Diffusion Transformer (DiT) architecturen.

• Kwantitatieve Verbeteringen:

  • Op ImageNet 256$\times$256 bereikte het Subtraction-gebaseerde Boundary RF Model een FID van 6.32 (ODE sampling) en 6.38 (SDE sampling), vergeleken met de vanilla RF baseline van respectievelijk 6.87 en 6.79.
  • Het Masker-gebaseerde model bereikte een FID van 6.63 (ODE) en 6.18 (SDE).
  • Verbeteringen werden ook waargenomen in Inception Score (IS), Precision en Recall over beide datasets.
  • De Subtraction-gebaseerde variant behaalde consequent de hoogste IS-scores.

• Sampling Robuustheid:

  • In experimenten met de Overshooting Sampler met variërende ruisschalen ($c$), vertoonden de Boundary RF-modellen een significant stabielere prestatie en minder degradatie in de kwaliteit van de samples vergeleken met de vanilla RF naarmate de stochastiek toenam.
  • Kwalitatieve resultaten toonden aan dat Boundary RF-modellen scherpere beelden produceren met minder artefacten en een betere objectgetrouwheid, vooral in stochastische sampling-omgevingen waar vanilla RF de neiging heeft tot overmatige smoothing.

• Schaalbaarheid:

  • De voordelen van het Boundary RF Model waren consistent bij het opschalen naar grotere modellen (DiT-L, DiT-XL) en hogere resoluties (512$\times$512).

Betekenis en Claims

Het artikel beweert dat het Boundary RF Model een eenvoudige, gemakkelijk te implementeren oplossing biedt voor een fundamentele beperking in Rectified Flow. Door de randvoorwaarden expliciet af te dwingen, biedt de methode:

1. Verbeterde Generatiekwaliteit: Presteert consequent beter dan vanilla RF op basis van FID en andere metrieken.
2. Maakt Stabiele Stochastische Sampling Mogelijk: Lost de divergentie van de scorefunctie nabij de eindtijd op, wat robuustere samplingstrategieën mogelijk maakt.
3. Is Eenvoudig Aan te passen: Kan met minimale architecturale wijzigingen in bestaande RF-modellen (zoals die met U-Net of DiT) worden geïntegreerd.

De auteurs positioneren dit werk als een noodzakelijke verfijning voor de praktische inzet van Rectified Flow, met name in scenario's die hoge kwaliteit stochastische sampling vereisen. Ze suggereren dat hoewel hun focus op RF ligt, het kernconcept van het afdwingen van randvoorwaarden potentieel met kleine aanpassingen kan worden uitgebreid naar diffusiemodellen. Toekomstig werk wordt voorgesteld om adaptieve randafdwinging en toepassing op tekst-naar-video generatie te verkennen.