Sintesi Tecnica: Residui di Wasserstein: Apprendimento di Flussi di Gradiente da Dinamiche di Popolazione

Definizione del Problema

Il paper affronta il problema inverso della ricostruzione delle dinamiche di popolazione da un insieme sparso di snapshot osservativi. In molti campi scientifici, come la biologia computazionale e la dinamica degli aggregati, l'evoluzione di una popolazione è modellata come un Flusso di Gradiente di Wasserstein (WGF): una curva di distribuzioni di probabilità $\rho_t$ guidata dalla discesa più rapida di un funzionale di energia $F$. L'obiettivo è recuperare il funzionale di energia sottostante $F$ tale che il suo flusso di gradiente si adatti alle distribuzioni marginali osservate $\{q_t\}$ ai tempi di osservazione discreti $t \in \mathcal{T}_{obs}$.

Gli approcci dominanti esistenti si affidano allo schema Jordan–Kinderlehrer–Otto (JKO), che definisce i WGF come una sequenza di passi prossimali che minimizzano un'energia più una penalità di trasporto ottimale (OT). Tuttavia, i metodi basati su JKO soffrono di due limitazioni primarie:

1. Inflessibilità rispetto alla discretizzazione temporale: Richiedono la risoluzione di costosi problemi di OT tra snapshot consecutivi.
2. Sensibilità alle lacune: Faticano a gestire grandi intervalli temporali tra le osservazioni, spesso fallendo nel catturare traiettorie curve (ad esempio, particelle che si muovono lungo una valle sinusoidale) poiché si affidano all'interpolazione lineare tra gli snapshot.

Metodologia: Il Framework dei Residui di Wasserstein

Gli autori propongono un approccio basato sui residui che evita lo schema JKO e i accoppiamenti di trasporto ottimale. Invece di imporre la condizione JKO, essi impongono direttamente le equazioni differenziali fondamentali che caratterizzano i flussi di gradiente tramite una funzione di perdita non negativa.

1. La Formulazione del Residuo

Il paper utilizza le formulazioni Tangente ed Equazione di Dissipazione dell'Energia (EDE) dei flussi di gradiente. Nello specifico, si concentra sul vincolo di velocità:
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
dove $v_t$ è il campo di velocità e $\frac{\delta F}{\delta \rho_t}$ è la variazione prima del funzionale di energia.

Gli autori definiscono un residuo di velocità $R_{vel}$ che è zero se e solo se $\rho$ è un flusso di gradiente di $F$:
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

L'obiettivo globale combina questo residuo con una divergenza di adattamento ai dati $D$ ai tempi osservati:
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
Questo framework unifica metodi esistenti come Path-Finding e Action Matching come istanze di un singolo principio di residuo.

2. L'Algoritmo "Stitching"

Per minimizzare questo obiettivo, gli autori introducono Stitching, un metodo basato su particelle e privo di simulazione (simulation-free), con le seguenti caratteristiche:

• Traiettorie apprendibili: A differenza dei metodi JKO dove la curva $\rho$ è implicita, Stitching parametrizza esplicitamente la curva $\rho_\theta$ come una Stima della Densità Kernel (KDE) su un insieme di $N$ traiettorie di particelle apprendibili $x_{t,k}^\theta$.
• Privo di simulazione: Il metodo non richiede la risoluzione di Neural ODE o l'integrazione di SDE. Le posizioni e le velocità delle particelle sono parametri apprendibili diretti.
• Approssimazione empirica: Per rendere l'aspetto computazionale efficiente, la KDE è approssimata da una misura empirica nei centri delle particelle. Ciò semplifica il residuo di velocità in una somma su particelle:
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• Parametrizzazione funzionale: Il funzionale di energia $F_\theta$ è parametrizzato come una somma di un termine di potenziale $V_\theta$, un termine entropico (log-densità) e un termine di interazione $W_\theta$, tutti rappresentati da reti neurali.

Contributi Chiave

1. Framework dei Residui: Gli autori riconsiderano l'apprendimento di WGF come la minimizzazione di un termine residuo non negativo, fornendo una prospettiva unificata che assorbe Path-Finding e Action Matching.
2. Algoritmo Stitching: Introducono un nuovo metodo particellare privo di simulazione che tratta la curva di traiettoria $\rho$ come una variabile apprendibile di primo ordine. Ciò consente al metodo di tollerare ampi gap tra gli snapshot e recuperare dinamiche curve che i predittori a corda basati su JKO perdono.
3. Prestazioni allo Stato dell'Arte: Il metodo raggiunge risultati superiori nei benchmark di inferenza delle traiettorie, particolarmente in scenari con dati sparsi o non accoppiati (unpaired).

Risultati Sperimentali

Il paper valuta Stitching su tre benchmark principali:

1. Esempio Illustrativo (Flusso Continuo vs JKO)

In un potenziale sintetico a "valle ondulata", le particelle seguono una traiettoria sinusoidale.

• Risultato: I metodi basati su JKO (specificamente JKOnet*) falliscono nel catturare la curvatura tra gli snapshot sparsi ($t=10$ a $t=20$), produvendo interpolazioni lineari. Stitching recupera correttamente la traiettoria curva, portando a un recupero del potenziale sottostante significativamente migliore ($R^2 = 0.62$ vs $0.51$).

2. Recupero di Potenziale Sintetico

Il metodo è testato su 15 potenziali bidimensionali in regimi accoppiati (traiettorie correlate) e non accoppiati (snapshot decorrelati).

• Risultato: Stitching mantiene prestazioni coerenti in entrambi i regimi. Al contrario, i metodi basati su JKO degradano significativamente o collassano nel regime non accoppiato perché si affidano a accoppiamenti OT tra snapshot consecutivi. Gli errori di Stitching rimangono entro un fattore di $\sim 2\times$ tra i regimi, mentre i metodi JKO spesso non riescono a recuperare campi di gradiente informativi nel caso non accoppiato.

3. Inferenza di Traiettoria Single-Cell (Dataset Embryoid Body)

Applicato a dati di differenziazione di cellule staminali embrionali umane (5 snapshot su 27 giorni).

• Risultato: Stitching ottiene le migliori prestazioni (distanze $W_1$ e $W_2$ più basse) rispetto ai baseline inclusi Neural SDE, TrajectoryNet e vari metodi basati su JKO. Supera i competitor sia nell'addestramento con dati completi che in un protocollo di interpolazione temporale "leave-two-out".

4. Recupero della Dinamica di Interazione

Il metodo viene utilizzato per disaggregare i potenziali ambientali ($V$) e i kernel di interazione ($W$) in un sistema di McKean–Vlasov.

• Risultato: Stitching recupera con successo sia il potenziale di confinamento che il kernel di attrazione, ottenendo alti punteggi di pattern $R^2$ ($0.84$ per $V$, $0.81$ per $W$). È inoltre significativamente più veloce per epoca rispetto ai metodi basati su JKO, che richiedono la risoluzione di problemi OT ad ogni passo.

5. Flussi Non-Gradiente

Gli autori dimostrano che il framework dei residui può essere esteso oltre i flussi di gradiente sostituendo il termine di gradiente con un campo vettoriale generale, tracciando con successo dinamiche chirali (non conservative).

Significato e Rivendicazioni

Il paper sostiene che la prospettiva del Residuo di Wasserstein offre un'alternativa più flessibile e robusta allo schema JKO per l'apprendimento delle dinamiche di popolazione. Decoupling (disaccoppiando) l'apprendimento del funzionale di energia dai vincoli degli accoppiamenti di trasporto ottimale, Stitching abilita:

• Robustezza alla Sparsità: La capacità di apprendere dinamiche da osservazioni ampiamente distanziate senza la "perdita di curvatura" inerente all'interpolazione lineare.
• Compatibilità con Dati Non Accoppiati: La capacità di apprendere da dataset dove l'identità delle particelle non è tracciata nel tempo (snapshot non accoppiati), uno scenario comune nell'imaging biologico.
• Efficienza Priva di Simulazione: Evitare l'overhead computazionale della risoluzione di ODE o problemi OT durante l'addestramento.

Gli autori posizionano questo lavoro come un'unificazione di paradigmi esistenti e un avanzamento pratico per la scoperta scientifica in campi che vanno dalla biologia delle singole cellule al comportamento collettivo, pur riconoscendo i limiti come la complessità computazionale $O(N^2)$ per i termini di interazione e la mancanza di garanzie teoriche di convergenza per $N \to \infty$.