Technische Samenvatting: Leren op de Manifold met Riemannian Flow Matching

1. Probleemstelling

Recente vooruitgang in generatieve modellering heeft de verschuiving naar het benutten van bevroren, hoogdimensionale semantische representatie-encoders (bijv. DiNOv2, SigLIP) om de synthesegetrouwheid en convergentiesnelheid te verbeteren. Echter, standaard Diffusion Transformers (DiT) en Flow Matching (FM) frameworks slagen er niet in effectief te convergeren wanneer ze direct op deze representaties worden getraind.

Eerder werk, specifiek de Representation Autoencoders (RAE), schreef dit falen toe aan een capaciteitsbottleneck, waarbij werd voorgesteld dat modellen "breedte-schaling" vereisen (het vergroten van de transformer-breedte om overeen te komen met de latente dimensionaliteit) om dit probleem op te lossen. Dit artikel daagt die hypothese uit door te stellen dat de kernoorzaak niet een tekort aan modelcapaciteit is, maar een fundamentele geometrische mismatch, getermd Geometrische Interferentie.

Het kernprobleem ontstaat omdat representatie-encoders kenmerken projecteren op een hypersferische manifold ($S^{d-1}$) met een bijna nul radiale variantie (een "hard shell" geometrie). In tegen tegenstelling hiermee gaat standaard Euclidische Flow Matching uit van een lineair probabilistisch pad (een koord) tussen de data en Gaussische ruis. Deze lineaire interpolatie snijdt door het gebied met lage dichtheid in het binnenste van de hypersfeer, waardoor het model gedwongen wordt om snelheidvelden te leren in ongedefinieerde regio's van de feature-ruimte. Bijgevolg verspilt het model zijn capaciteit aan het minimaliseren van radiale fouten (die structureel nul zijn op de manifold) in plaats van het leren van de semantische angulaire transities.

2. Methodologie: Riemannian Flow Matching met Jacobi Regularisatie (RJF)

Om Geometrische Interferentie op te lossen zonder architecturale wijzigingen zoals breedte-schaling, stellen de auteurs RJF voor, een framework dat het generatieve proces afstemt op de intrinsieke geometrie van de latente ruimte.

2.1 Riemannian Flow Matching (Trajectcorrectie)

In plaats van lineaire interpolatie in de Euclidische ruimte, definieert RJF het probabilistische pad direct op de hypersferische manifold met behulp van Spherical Linear Interpolation (SLERP).

• Geodetische paden: Het traject tussen een datapunt $x$ en ruis $\epsilon$ (beiden geprojecteerd op de eenheidsbol) volgt de geodeet (grote cirkelboog) in plaats van een koord.
• Tangenstiale snelheid: Het doel-snelheidsveld wordt berekend door het SLERP-pad te differentiëren. Dit zorgt ervoor dat de voorspelde snelheidsvector strikt in de tangentiële ruimte van de manifold ligt ($v_t \cdot x_t = 0$), wat het nutteloos leren van radiale componenten elimineert.
• Sampling: Tijdens de inferentie gebruikt het model een Geodetische (Exponential Map) integrator om de latente staat bij te werken, waardoor het traject gedurende het hele generatieproces op het oppervlak van de manifold blijft.

2.2 Jacobi Regularisatie (Curvatuurcorrectie)

De auteurs observeren dat het enkel corrigeren van het traject onvoldoende is, omdat de standaard Mean Squared Error (MSE) loss snelheidfouten uniform behandelt. Op een positief gekromde hypersfeer convergeren geodetische paden (focusseren ze), wat betekent dat snelheidfouten niet-lineair voortplanten.

• Jacobi-velden: Om dit te corrigeren, introduceren de auteurs een weegfactor afgeleid van Jacobi-velden, die de scheiding tussen geodetische paden kwantificeren als gevolg van snelheidsperturbaties.
• Loss Reweighting: De verliesfunctie wordt vermenigvuldigd met een geometrische weegfactor $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$, waarbij $\Omega$ de geodetische afstand is. Dit verlaagt de weging van fouten nabij de data ($t=0$), waar focus perturbaties mitigeert, en geeft prioriteit aan precisie nabij de ruis ($t=1$), waardoor het leersignaal effectief wordt geanneald om aan te sluiten bij de kromming van de manifold.

3. Belangrijkste Bijdragen

1. Geometrische Analyse van Convergentiefalen: Het artikel identificeert Geometrische Interferentie als de primaire bottleneck voor standaard diffusie op representatie-encoders. Het demonstreert dat het falen voortkomt uit het feit dat Euclidische doelstellingen modellen dwingen om door het off-manifold binnenste van de hypersfeer te leren, in plaats van een gebrek aan modelcapaciteit.
2. RJF Framework: Het voorstel van een geometrisch consistent trainingsframework dat Riemannian Flow Matching combineert (voor traject-afstemming) met Jacobi Regularisatie (voor kromming-bewuste foutvoortplanting).
3. Efficiëntie zonder Breedte-schaling: De demonstratie dat standaard DiT-architecturen (bijv. DiT-B met 131M parameters) effectief kunnen convergeren op hoogdimensionale semantische latents zonder de computationeel dure breedte-schaling die door eerdere werken werd voorgesteld.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op ImageNet-1K met een resolutie van 256×256 met gebruik van verschillende modelmaten (DiT-B, DiT-L, DiT-XL) en representatie-encoders (DiNOv2-B, SigLIP, MAE).

• Convergentie en FID-prestaties:
  • DiT-B (131M): Standaard Euclidische Flow Matching op DiNOv2-B features faalde om te convergeren (FID ~24.21). RJF bereikte een FID van 6.77 in 80 epochs en 3.37 (met guidance) in 200 epochs.
  • DiT-XL (677M): RJF bereikte een FID van 3.62 in 80 epochs zonder guidance, waarmee het de Euclidische baseline (FID 4.28) en de VAE-gebaseerde LightingDiT (FID 4.29) overtrof.
  • Schaling: De methode toonde consistente verbeteringen over verschillende modelgroottes, waarbij een state-of-the-art Inception Score (IS 186.2) en Precisie (0.82) werd behaald.
• Ablatie studies:
  • Radiaal vs. Angulair: Het maskeren van de radiale loss leverde slechts marginale winst op (FID 21.99), wat bevestigt dat het lineaire pad zelf (het koord) het primaire probleem is.
  • Traject: Het overschakelen naar enkel Riemannian Flow Matching (SLERP) verbeterde de prestaties drastisch (FID 7.06), wat de hypothese van de traject-mismatch valideert.
  • Regularisatie: Het toevoegen van Jacobi Regularisatie verfijnde de resultaten verder (FID 6.77).
• Generalisatie: De methode convergeerde succesvol op andere encoders met hypersferische beperkingen, waaronder SigLIP (FID 10.39) en MAE (FID 19.82), waar standaard methoden volledig faalden.

5. Betekenis en Claims

Het artikel claimt dat de moeilijkheid bij het trainen van diffusiemodellen op representatie-encoders fundamenteel geometrisch is, niet architecturaal. Door de hypersferische topologie van de latente ruimte te respecteren, kunnen standaard Diffusion Transformers hoogwaardige semantiek efficiënt leren.

De auteurs stellen dat RJF het potentieel van standaard DiT-architecturen ontsluit, waardoor ze grotere, breedte-geschaalde modellen of complexe VAE-pipelines kunnen evenaren of overtreffen. Het werk stelt vast dat efficiënte generatieve modellering op representatie-manifolds geometrische afstemming vereist (correcte trajecten en kromming-bewuste losses) in plaats van brute-force capaciteit-schaling. Deze aanpak maakt synthese met hoge getrouwheid mogelijk met minder parameters en snellere convergentie, zoals bewezen door het DiT-B model dat een FID van 3.37 bereikte waar eerdere methoden niet konden convergeren.