Technische Samenvatting: Riemannian MeanFlow

Probleemstelling

Diffusie- en stromingsmodellen zijn de dominante paradigma geworden voor generatieve modellering op Riemanniaanse variëteiten, met succesvolle toepassing op taken zoals generatie van eiwitruggraden en ontwerp van DNA-sequenties. Deze methoden vereisen echter doorgaans tientallen tot honderden evaluaties van het neurale netwerk (functie-evaluaties, of NFE's) tijdens inferentie om gewone differentiaalvergelijkingen (ODE's) of stochastische differentiaalvergelijkingen (SDE's) numeriek te integreren. Deze rekenkundige last creëert een knelpunt, met name in wetenschappelijke ontwerppijplijnen waar generatieve modellen dienen als voorstelmechanismen binnen iteratieve lussen voor eigenschapsgeleide optimalisatie. Hoewel er in Euclidische settings is geëxperimenteerd met generatietechnieken in weinig stappen, zoals consistentiemodellen en het leren van stromingskaarten, blijven deze onderbelicht en vaak instabiel in settings met hoogdimensionale Riemanniaanse variëteiten.

Methodologie: Riemannian MeanFlow (RMF)

De auteurs introduceren Riemannian MeanFlow (RMF), een raamwerk ontworpen om stromingskaarten direct op variëteiten te leren, waardoor hoogwaardige generatie mogelijk is met slechts één forward pass. De kern van de methodologie bestaat uit het afleiden en gebruiken van drie equivalente karakterisering van de gemiddelde snelheid op de variëteit, die de stromingskaart tussen willekeurige tijdstippen $s$ en $t$ relateren.

1. Theoretische Fundamenten

Het raamwerk definieert de gemiddelde snelheid $u_{s,t}$ als de constante snelheid die, indien gehandhaafd langs een geodeet van $x_s$ naar $x_t$ over de tijd $t-s$, het punt correct zou vervoeren. De auteurs leiden drie identiteiten af die dienen als basis voor trainingsdoelstellingen:

• Euleriaanse Identiteit: Afgeleid door differentiatie van de definiërende relatie met betrekking tot de bron-tijd $s$. Deze relateert de gemiddelde snelheid aan de covariante afgeleide van het snelheidsveld langs de integraalkromme en de afgeleide van de logaritmische kaart.
• Lagrangeaanse Identiteit: Afgeleid door differentiatie met betrekking tot de doeltijd $t$. Deze relateert de gemiddelde snelheid aan de differentiaal van de logaritmische kaart toegepast op de momentane snelheid en de partiële tijdsafgeleide van de gemiddelde snelheid.
• Semigroep-Identiteit: Een algebraïsche voorwaarde afgeleid uit de semigroep-eigenschap van stromingskaarten ($\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$). Deze dwingt consistentie af tussen de directe stroming van $s$ naar $t$ en de samenstelling van stromingen via een tussenliggend tijdstip $r$.

2. Trainingsdoelstellingen en Stabilisatie

De auteurs converteren deze identiteiten naar regressiedoelstellingen. Een kritieke ontwerpkeuze is de plaatsing van de stop-gradient-operator. In tegenstelling tot eerdere werken (bijv. Generalized Flow Maps) die stop-gradients gedeeltelijk toepassen, past RMF de stop-gradient toe op de volledige regressiedoel. Deze formulering:

• Elimineert de noodzaak om hogere-orde afgeleiden te evalueren (bijv. gradiënten via Jacobiaan-vectorproducten) tijdens backpropagatie.
• Vermindert het geheugengebruik met bijna 2$\times$.
• Verbeterd de trainingsstabiliteit in hoogdimensionale settings.

Om optimalisatieproblemen in hoge dimensies aan te pakken, stelt het artikel diverse stabilisatietechnieken voor:

• Tijdsampling: Geordende tijdparen ($s \le t$) voor Euleriaanse en Semigroep-doelstellingen, en ongeordende paren voor Lagrangeaanse doelstellingen.
• Adaptieve Loss-Weigting: Herwegen van de loss op basis van de grootte van de regressiedoel om hoge variantie in afgeleide-afhankelijke termen te hanteren.
• Tijdsafgeleide-Controle: Het gebruik van laagfrequente tijdembeddings om de grootte van tijdsafgeleiden in de netwerkoutput te begrenzen.

3. Parametrisatie

Het artikel evalueert drie parametrisatieschema's voor de gemiddelde snelheid:

• $v$-predictie: Het voorspellen van de raakvector (snelheid).
• $x_t$-predictie: Het direct voorspellen van het eindpunt.
• $x_1$-predictie: Het voorspellen van het trajecteindpunt bij $t=1$ gegeven de huidige staat en het tijdsverschil.
  De auteurs vinden dat $x_1$-predictie het meest robuust is voor hoogdimensionale variëteiten. Het is compatibel met bestaande wetenschappelijke architecturen (bijv. die SE(3)-frames voor eiwitten uitvoeren) en schaalt goed tot dimensies tot $D=2048$, terwijl $x_t$-predictie faalt en $v$-predictie degradeert naarmate de dimensionaliteit toeneemt.

4. Beloningsgeleide Inferentie

Het raamwerk ondersteunt beloningsgeleide generatie zonder modelhertraining. Door gebruik te maken van het vermogen van de stromingskaart om eindtoestanden te voorspellen vanuit tussenliggende stappen, stellen de auteurs een belonings-look-ahead strategie voor. In plaats van de gradiënt van de beloning bij de huidige ruisige staat $x_t$ te gebruiken, gebruiken ze de gradiënt bij het voorspelde eindpunt $\Phi_{t,1}(x_t)$. Deze aanpak blijkt effectiever te zijn voor het sturen van generaties naar specifieke doelen.

Belangrijkste Resultaten

De auteurs valideren RMF op twee biologische taken: Promoter DNA-ontwerp (sequenties op een simplex in $\mathbb{R}^{1024 \times 4}$) en Eiwitruggraadgeneratie (op $SE(3)^N$ met $N$ tot 128).

• Efficiëntie: RMF bereikt steekproefkwaliteit die vergelijkbaar is met state-of-the-art meerstaps-baselines (bijv. Fisher FM, FrameFlow, FrameDiff), terwijl het tot 10$\times$ minder functie-evaluaties vereist.
  • Bij eiwitruggraadgeneratie bereikt RMF 82% ontwerpbaarheid bij 5 stappen en 35% bij 1 stap, terwijl baselines in deze weinig-stapsregimes sterk dalen of volledig instorten.
  • Bij DNA-ontwerp komen RMF-varianten de prestaties van Fisher FM met 100 stappen overeen met slechts 1 stap.
• Stabiliteit: De voorgestelde stabilisatietechnieken (specifiek de combinatie van Semigroep-doelstellingen en $x_1$-predictie) maken stabiele training mogelijk in hoogdimensionale ruimten waar andere stromingskaart-methoden (zoals GFM) falen om te convergeren of foutrijke steekproeven produceren.
• Beloningsgeleiding: Beloningsgeleide inferentie met de $x_1$ look-ahead strategie verbetert significant de uitlijning met doelreguleringsprofielen (DNA) en secundaire structuursamenstellingen (eiwitten) in vergelijking met naïeve geleiding of ongeleide generatie.

Betekenis en Aanspraken

Het artikel claimt een principiële theoretische onderbouwing en praktische richtlijnen te vestigen voor snelle sampling op Riemanniaanse variëteiten. De betekenis hiervan ligt in:

1. Het Overbruggen van de Kloof: Het breidt generatiecapaciteiten in weinig stappen uit van Euclidische ruimten naar complexe, hoogdimensionale Riemanniaanse variëteiten, een cruciale stap voor wetenschappelijke toepassingen.
2. Schaalbaarheid: Door identiteiten af te leiden die hogere-orde afgeleiden vermijden en $x_1$-predictie te gebruiken, schaalt de methode tot dimensies ($D \approx 2000$) en sequentielengtes die voorheen ontoegankelijk waren voor generatieve modellen voor variëteiten in weinig stappen.
3. Efficiëntie in Wetenschappelijke Werkstromen: De vermindering van inferentiekosten (tot 10$\times$) maakt generatieve modellen meer levensvatbaar als voorstelmechanismen in iteratieve wetenschappelijke ontwerplussen, zoals eigenschapsgeleide optimalisatie voor geneesmiddelenontwikkeling of synthetische biologie.
4. Belonings Look-Ahead: Het toont aan dat stromingskaarten efficiënt beloningsgeleid ontwerp mogelijk maken door eindtoestanden te voorspellen vanuit tussenliggende stappen, een capaciteit die voorheen niet op variëteiten was aangetoond.

De auteurs positioneren RMF niet als vervanging voor alle generatieve modellering, maar als een gespecialiseerd raamwerk dat efficiënte, hoogwaardige generatie mogelijk maakt in geometrisch beperkte wetenschappelijke domeinen waar rekenkosten een beperkende factor zijn.