One-Sided Quantile Coupling for Flow Matching

Technische Samenvatting: One-Sided Quantile Coupling voor Flow Matching

Probleemstelling
Flow Matching (FM) traint continue-tijd generatieve modellen door de snelheid van een waarschijnlijkheidspad te regresseren dat een eenvoudige bronverdeling (meestal een standaard Gaussiaanse verdeling) verbindt met een doelverdeling van data. De kwaliteit van de koppeling — de pairing van bron- en doelstalen — beïnvloedt de geometrie van de trainingspaden cruciaal. Standaard FM gebruikt onafhankelijke willekeurige pairing, wat vaak resulteert in lange, conflicterende paden die de variantie van het doelveld van de snelheid verhogen en de efficiëntie van de training hinderen.

Hoewel gestructureerde koppelingen geïnspireerd door Optimal Transport (OT) deze paden kunnen verstevigen en de generatiekwaliteit kunnen verbeteren, vertrouwen bestaande methoden (bijv. OT-CFM) op het oplossen van een mini-batch toewijzingsprobleem. Deze benaderingen vereisen het construeren van paarwijze kostenmatrices en het oplossen van toewijzingsproblemen (bijv. via het Hongaarse algoritme of Sinkhorn-iteraties), wat leidt tot computationele kosten die ten minste kwadratisch (of kubisch voor exacte oplossingen) schalen met de batchgrootte. Dit beperkt de schaalbaarheid en kan instabiliteit introduceren. Het artikel adresseert de vraag of de geometrische voordelen van gestructureerde koppeling bereikt kunnen worden zonder twee vooraf gesamplede batches te matchen, waardoor de hoge kosten van batchwise transport worden vermeden.

Methodologie: Quantile Coupling Flow Matching (QC-FM)
De auteurs stellen Quantile Coupling Flow Matching (QC-FM) voor, een lichtgewicht, eenzijdige koppelingmethode. In tegenstelling tot traditionele methoden die zowel ruis- als databatches samplen en deze vervolgens opnieuw paren, samplen QC-FM alleen de databatch en construeert de gepaarde bronstalen direct.

1. Eenzijdige Constructie:

   • Gegeven een databatch $\mathcal{B}_x$, projecteert de methode datapunten op $k$ willekeurige orthogonale richtingen (slices) gedefinieerd door een matrix $U$.
   • Voor elke richting worden de datapunten binnen de batch gerangschikt. Deze empirische rangordes worden gemapt naar Gaussische kwantielen met behulp van de inverse standaard normale CDF.
   • De bron latente code $e^{(i)}$ wordt geconstrueerd om aan deze geprojecteerde kwantielbeperkingen te voldoen. Specifiek voldoet de projectie van de bron op de slice-richtingen aan de Gaussische kwantielen afgeleid van de data-rangordes.
   • De resterende dimensies (het orthogonale complement van de slice-subruimte) worden ingevuld door te samplen uit een conditionele Gaussische verdeling. Dit zorgt ervoor dat de bron stokastisch blijft in onbeperkte richtingen.

2. Gaussian-Remainder Hybride Extensies:
   Omdat de volledige QC-koppeling (waarbij alle bronpunten via rangordes worden geconstrueerd) afwijkt van de standaard i.i.d. Gaussische prior voor $k \ge 2$ vanwege afhankelijkheden tussen de slice-codes, introduceren de auteurs twee hybride strategieën om de voordelen van QC te behouden terwijl de prior behouden blijft:

   • QC-FM-Mixture: Past de QC-constructie toe op een "anker"-subset van de batch (bijv. 20%) en koppelt de resterende "remainder" bronstalen met de resterende datastalen met behulp van onafhankelijke willekeurige pairing (de Baseline-benadering). Dit creëert een mengeling van gestructureerde en ongestructureerde koppelingen.
   • QC-FM-Adjacency: Past QC toe op de anker-subset, maar koppelt de remainder bronstalen op basis van de lokale geometrie geïnduceerd door de ankers. Resterende datapunten worden toegewezen aan hun dichtstbijzijnde anker (in de gestructureerde subruimte), en de resterende bronstalen worden binnen deze groepen gekoppeld. Dit propageert de gestructureerde bias zonder doelpunten te dupliceren.

Belangrijkste Bijdragen

• QC-FM Algoritme: Een nieuwe, batch-lokale gestructureerde koppeling die monotone eendimensionale transport langs willekeurige slices afdwingt. Het elimineert de noodzaak voor paarwijze kostenmatrices en toewijzingsoplossers, waardoor de dominante computationele kosten worden teruggebracht naar $O(kB \log B)$ (sorteren) in plaats van $O(B^2)$ of $O(B^3)$.
• Theoretische Analyse:
  • Variantereductie: De auteurs bewijzen dat voor een vaste frame, de QC-koppeling de onherleidbare regressievariantie langs elke geselecteerde slice elimineert, waardoor de ideale flow exact recht is in die richtingen.
  • Prior Preservatie: Ze kwantificeren dat voor $k=1$, de bronmarginaal exact Gaussisch is. Voor $k \ge 2$ is de afwijking van de Gaussische prior beperkt tot de copula van de slice-codes, en de transportkosten van deze afwijking zijn begrensd.
• Hybride Trainingsschema's: De introductie van Mixture en Adjacency varianten die een balans vinden tussen het injecteren van geometrische bias en het behoud van het expliciete Gaussische signaal van de Baseline.

Experimentele Resultaten
De methode werd geëvalueerd op CIFAR-10, CelebA-64, FFHQ-64 en ImageNet-64 onder gelijke trainingsbudgetten en architecturen.

• Prestaties: QC-FM-Mixture presteerde consistent beter dan de onafhankelijke koppelings-Baseline en de exacte Hongaarse OT-CFM over alle vier de datasets. De verbeteringen ten opzichte van de Baseline varieerden van 6,3% (CIFAR-10) tot 12,9% (FFHQ) in termen van FID-50K.
• Efficiëntie: In termen van computationele kosten is QC-FM aanzienlijk sneller dan OT-gebaseerde methoden. Bij een batchgrootte van 2048 was QC-FM-Mixture meer dan 800 keer sneller dan de Hongaarse OT-solver en 35 keer sneller dan de Sinkhorn-solver.
• Toy Diagnostiek: Op synthetische datasets (Checkerboard en 8-Gaussians) produceerden QC-FM varianten bijna rechte geleerde trajecten (Path-Length Ratio $\approx 1$) en een significant lagere doel-snelheidsvariantie vergeleken met de Baseline, wat de geometrische voordelen van de koppeling bevestigt.

Betekenis en Claims
Het artikel claimt dat het behouden van de geprojecteerde rangorde een eenvoudige, schaalbare en effectieve manier is om nuttige geometrische bias in Flow Matching koppelingen te injecteren zonder een mini-batch transportprobleem op te lossen.

• Schaalbaarheid: QC-FM biedt een aantrekkelijk middenpad tussen de ongestructureerde Baseline en dure batchwise transportmethoden. Het is computationeel lichtgewicht en controleerbaar via het aantal projecties ($k$).
• Praktijkgerichtheid: De methode vereist geen geleerde componenten voor inferentie; generatie begint nog steeds vanaf de standaard Gaussische verdeling. De koppeling wordt on-the-fly geconstrueerd tijdens de training.
• Beperkingen: De auteurs merken bescheiden op dat QC-FM een batch-lokale surrogaat is voor gestructureerde koppeling, en geen benadering van globale Optimal Transport. De kwantielen worden geschat binnen mini-batches, en de theorie controleert de per-slice variantie in plaats van een globale traject-rechtheid of specifieke FID-verbeteringen te garanderen. De selectie van het aantal slices ($k$) en de ankerratio ($p$) berust momenteel op heuristieken en dataset-specifieke afstemming.

Concluderend toont dit werk aan dat eenzijdig quantile coupling de steekproefkwaliteit en trainingsefficiëntie in Flow Matching kan verbeteren door de ruis-data pairing te structureren via rangbehoud, wat een praktische alternatief biedt voor computationeel prohibitieve optimal transport methoden.