Riepilogo Tecnico: Riemannian MeanFlow

Enunciato del Problema

I modelli di diffusione e di flusso sono diventati il paradigma dominante per la modellazione generativa su varietà Riemanniane, applicati con successo a compiti come la generazione di scheletri proteici e la progettazione di sequenze di DNA. Tuttavia, questi metodi richiedono tipicamente da decine a centinaia di valutazioni della rete neurale (valutazioni della funzione, o NFE) durante l'inferenza per integrare numericamente equazioni differenziali ordinarie (ODE) o equazioni differenziali stocastiche (SDE). Questo onere computazionale crea un collo di bottiglia, in particolare nelle pipeline di progettazione scientifica dove i modelli generativi fungono da meccanismi di proposta all'interno di loop iterativi per l'ottimizzazione guidata dalle proprietà. Sebbene tecniche di generazione in pochi passi, come i modelli di consistenza e l'apprendimento delle mappe di flusso, siano state esplorate in contesti euclidei, rimangono poco esplorate e spesso instabili in contesti di varietà Riemanniane ad alta dimensionalità.

Metodologia: Riemannian MeanFlow (RMF)

Gli autori introducono Riemannian MeanFlow (RMF), un framework progettato per apprendere mappe di flusso direttamente sulle varietà, abilitando una generazione di alta qualità con al più un singolo passaggio in avanti. Il nucleo della metodologia consiste nel derivare e utilizzare tre caratterizzazioni equivalenti della velocità media sulla varietà, che relazionano la mappa di flusso tra punti temporali arbitrari $s$ e $t$.

1. Fondamenti Teorici

Il framework definisce la velocità media $u_{s,t}$ come la velocità costante che, se mantenuta lungo una geodetica da $x_s$ a $x_t$ nell'intervallo di tempo $t-s$, trasporterebbe correttamente il punto. Gli autori derivano tre identità che servono come base per gli obiettivi di addestramento:

• Identità Euleriana: Derivata differenziando la relazione definitoria rispetto al tempo sorgente $s$. Relaziona la velocità media alla derivata covariante del campo di velocità lungo la curva integrale e alla derivata della mappa logaritmica.
• Identità Lagrangiana: Derivata differenziando rispetto al tempo target $t$. Relaziona la velocità media al differenziale della mappa logaritmica applicato alla velocità istantanea e alla derivata parziale temporale della velocità media.
• Identità Semigruppo: Una condizione algebrica derivata dalla proprietà di semigruppo delle mappe di flusso ($\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$). Impone coerenza tra il flusso diretto da $s$ a $t$ e la composizione dei flussi attraverso un tempo intermedio $r$.

2. Obiettivi di Addestramento e Stabilizzazione

Gli autori convertono queste identità in obiettivi di regressione. Una scelta progettuale critica è il posizionamento dell'operatore stop-gradient. A differenza di lavori precedenti (ad esempio, Generalized Flow Maps) che applicano lo stop-gradient parzialmente, RMF applica lo stop-gradient all'intero target di regressione. Questa formulazione:

• Elimina la necessità di valutare derivate di ordine superiore (ad esempio, gradienti attraverso prodotti vettore-Jacobiano) durante la retropropagazione.
• Riduce l'utilizzo di memoria di quasi 2$\times$.
• Migliora la stabilità dell'addestramento in contesti ad alta dimensionalità.

Per affrontare le difficoltà di ottimizzazione in alta dimensionalità, il documento propone diverse tecniche di stabilizzazione:

• Campionamento Temporale: Coppie di tempi ordinate ($s \le t$) per gli obiettivi Euleriani e Semigruppo, e coppie non ordinate per gli obiettivi Lagrangiani.
• Pesatura Adattiva della Loss: Ripesatura della loss basata sulla grandezza del target di regressione per gestire l'alta varianza nei termini dipendenti dalle derivate.
• Controllo della Derivata Temporale: Utilizzo di embedding temporali a bassa frequenza per limitare la grandezza delle derivate temporali nell'output della rete.

3. Parametrizzazione

Il documento valuta tre schemi di parametrizzazione per la velocità media:

• Predizione $v$: Predizione del vettore tangente (velocità).
• Predizione $x_t$: Predizione diretta del punto finale.
• Predizione $x_1$: Predizione del punto finale della traiettoria a $t=1$ dato lo stato corrente e l'intervallo temporale.
  Gli autori riscontrano che la predizione $x_1$ è la più robusta per varietà ad alta dimensionalità. È compatibile con architetture scientifiche esistenti (ad esempio, quelle che producono frame SE(3) per le proteine) e scala bene fino a dimensioni $D=2048$, mentre la predizione $x_t$ fallisce e la predizione $v$ si degrada all'aumentare della dimensionalità.

4. Inferenza Guidata dalla Ricompensa

Il framework supporta la generazione guidata dalla ricompensa senza riaddestramento del modello. Sfruttando la capacità della mappa di flusso di prevedere stati terminali da passaggi intermedi, gli autori propongono una strategia di look-ahead della ricompensa. Invece di utilizzare il gradiente della ricompensa allo stato rumoroso corrente $x_t$, utilizzano il gradiente al punto finale previsto $\Phi_{t,1}(x_t)$. È stato dimostrato che questo approccio è più efficace per indirizzare le generazioni verso obiettivi specifici.

Risultati Chiave

Gli autori validano RMF su due compiti biologici: Progettazione di DNA Promotore (sequenze su un simpletto in $\mathbb{R}^{1024 \times 4}$) e Generazione di Scheletri Proteici (su $SE(3)^N$ con $N$ fino a 128).

• Efficienza: RMF raggiunge una qualità del campione comparabile a quella delle basi multi-step all'avanguardia (ad esempio, Fisher FM, FrameFlow, FrameDiff) richiedendo fino a 10$\times$ meno valutazioni della funzione.
  • Nella generazione di scheletri proteici, RMF raggiunge l'82% di progettabilità in 5 passi e il 35% in 1 passo, mentre le basi crollano bruscamente o collassano completamente in questi regimi a pochi passi.
  • Nella progettazione di DNA, le varianti di RMF eguagliano le prestazioni a 100 passi di Fisher FM con appena 1 passo.
• Stabilità: Le tecniche di stabilizzazione proposte (in particolare la combinazione di obiettivi Semigruppo e predizione $x_1$) abilitano un addestramento stabile in spazi ad alta dimensionalità dove altri metodi basati su mappe di flusso (come GFM) non convergono o producono campioni ad alto errore.
• Guida tramite Ricompensa: L'inferenza guidata dalla ricompensa utilizzando la strategia di look-ahead $x_1$ migliora significativamente l'allineamento con i profili regolatori target (DNA) e le composizioni di struttura secondaria (proteine) rispetto a una guida ingenua o a una generazione non guidata.

Significato e Affermazioni

Il documento afferma di stabilire una fondazione teorica fondata e linee guida pratiche per il campionamento rapido su varietà Riemanniane. Il suo significato risiede in:

1. Colmare il Divario: Estende le capacità di generazione in pochi passi dagli spazi euclidei a varietà Riemanniane complesse e ad alta dimensionalità, un passo critico per le applicazioni scientifiche.
2. Scalabilità: Derivando identità che evitano derivate di ordine superiore e utilizzando la predizione $x_1$, il metodo scala a dimensioni ($D \approx 2000$) e lunghezze di sequenza precedentemente inaccessibili ai modelli generativi su varietà in pochi passi.
3. Efficienza nei Flussi di Lavoro Scientifici: La riduzione del costo di inferenza (fino a 10$\times$) rende i modelli generativi più fattibili come meccanismi di proposta in loop iterativi di progettazione scientifica, come l'ottimizzazione guidata dalle proprietà per la scoperta di farmaci o la biologia sintetica.
4. Look-Ahead della Ricompensa: Dimostra che le mappe di flusso abilitano una progettazione efficiente guidata dalla ricompensa prevedendo stati terminali da passaggi intermedi, una capacità non precedentemente mostrata sulle varietà.

Gli autori posizionano RMF non come un sostituto per tutta la modellazione generativa, ma come un framework specializzato che abilita una generazione efficiente e di alta qualità in domini scientifici vincolati geometricamente, dove il costo computazionale è un fattore limitante.