Riepilogo Tecnico: FLUX – Adattamento del Flusso Longitudinale Consapevole della Geometria con Mixture of Experts

1. Enunciazione del Problema

Molti sistemi biologici, come i circuiti neurali e le popolazioni cellulari in differenziazione, evolvono attraverso dinamiche locali continue pur passando tra regimi latenti definiti dall'apprendimento, dal contesto degli stimoli o dalle fasi dello sviluppo. Tuttavia, questi sistemi sono frequentemente osservati solo come istantanee longitudinali non accoppiate. In scenari come l'imaging del calcio a campo largo o il sequenziamento dell'RNA a cellula singola, le stesse cellule o neuroni individuali non possono essere tracciati attraverso i punti temporali; invece, i ricercatori osservano una sequenza di distribuzioni marginali ($\mu_0, \mu_1, \dots, \mu_{T-1}$).

Questa impostazione presenta due sfide accoppiate:

1. Vincoli Geometrici: Le misurazioni biologiche ad alta dimensionalità spesso si concentrano vicino a varietà non lineari curve a bassa dimensionalità. Gli interpolanti euclidei standard tra gli estremi non accoppiati possono passare attraverso regioni a bassa densità, violando i vincoli biofisici e non rispettando la geometria intrinseca dello spazio degli stati.
2. Scoperta dei Regimi: Il meccanismo di trasporto sottostante (il "flusso") cambia spesso nel tempo. Ad esempio, l'attività neurale può spostarsi da un'aspettativa generale di stimolo-ricompensa a una discriminazione specifica, oppure le cellule possono transitare da programmi pluripotenti a programmi impegnati in un lignaggio. Questi sono cambiamenti nelle dinamiche locali (regimi), non semplici cluster spaziali. I metodi esistenti spesso si concentrano sull'accuratezza del trasporto ma falliscono nel decomporre il campo di velocità per identificare questi regimi dinamici latenti.

2. Metodologia: FLUX

Gli autori introducono FLUX (Flow matching for Unpaired longitudinal data with miXture-of-experts), un framework che modella congiuntamente il trasporto continuo e la scoperta di regimi non supervisionata. Il metodo opera in tre fasi sequenziali:

A. Percorsi Condizionali Consapevoli della Geometria

FLUX estende il Metric Flow Matching a contesti longitudinali. Invece di utilizzare percorsi euclidei rettilinei tra le marginali adiacenti, FLUX apprende una metrica dipendente dai dati dalle osservazioni raggruppate (inclusi campioni non etichettati).

• Apprendimento della Metrica: Un backend geometrico (utilizzando metriche di kernel profondo RBF o RBF-MLP) stima il punteggio di distanza dalla varietà, definendo una metrica $G(x)$ che penalizza i percorsi che lasciano la varietà dei dati.
• Rete di Piega (Bend Network): Una rete $B$ parametrizza un percorso approssimato consapevole della geometria $\gamma_k(\alpha)$ tra gli estremi $x_k$ e $x_{k+1}$. Questo percorso si piega per rimanere sulla varietà, fornendo punti di training $z_{k,\alpha}$ e tangenti target $\dot{z}_{k,\alpha}$ che rispettano la geometria locale.

B. Campi di Velocità con Interruzione dei Regimi

FLUX scompone il campo di velocità globale $v_\theta$ in una Mixture of Experts (MoE).

• Campi Vettoriali degli Esperti: Il modello apprende $M$ campi vettoriali di esperti sparsi $f_m$.
• Router: Una rete router $g(t, x)$ mappa tempo e stato verso i logit degli esperti.
• Meccanismo di Instradamento: Durante l'addestramento, un meccanismo Straight-Through Gumbel-Softmax seleziona un esperto attivo per ciascun campione, consentendo un'ottimizzazione differenziabile delle decisioni di instradamento discrete. All'inferenza, il regime è determinato dall'assegnazione deterministica rigida (argmax).
• Obiettivo: Il campo di velocità è addestrato per corrispondere alle tangenti dei percorsi consapevoli della geometria. Il router è regolarizzato per incoraggiare assegnazioni sparse e coerenti nel tempo e prevenire il collasso degli esperti.

C. Pipeline di Addestramento

L'addestramento procede in tre fasi per garantire stabilità:

1. Apprendimento della Metrica: Addestramento del backend geometrico su campioni raggruppati.
2. Addestramento della Rete di Piega: Addestramento della rete di piega con la metrica congelata per generare interpolanti consapevoli della geometria.
3. Addestramento del Campo di Velocità: Congelamento delle reti geometriche e di piega; addestramento degli esperti di velocità e del router sui percorsi e sulle tangenti generati.

3. Contributi Chiave

1. Estensione alle Marginali Longitudinali: Gli autori estendono l'adattamento del flusso consapevole della geometria a catene ordinate di distribuzioni marginali, enfatizzando il ruolo dell'apprendimento di varietà non supervisionato quando le etichette di traiettoria sono scarse ma sono disponibili osservazioni non etichettate.
2. Scoperta dei Regimi tramite Decomposizione della Velocità: Introducono un campo di velocità con interruzione dei regimi in cui l'instradamento degli esperti sparsi scopre regimi dinamici latenti direttamente dalla struttura di trasporto locale, piuttosto che clusterizzare osservazioni statiche a posteriori.
3. Interazione Geometria-Regimi: Il lavoro dimostra che l'apprendimento geometrico è critico per la scoperta dei regimi, in particolare quando i regimi sono definiti da cambiamenti nelle dinamiche locali piuttosto che da una separazione spaziale statica.
4. Validazione Biologica: Il framework è validato su due contesti biologici distinti:
   • Dati Neurali: Separazione della fase iniziale/intermedia di addestramento dalla fase tardiva nell'apprendimento associativo, coincidente con la divergenza comportamentale.
   • Biologia dello Sviluppo: Recupero di regimi di evoluzione dell'espressione associati a popolazioni cellulari pluripotenti rispetto a differenziate in corpi embrionali.

4. Risultati

Il documento valuta FLUX su quattro dataset: Stanford Bunny (varietà controllata), Attrattore di Lorenz (dinamiche con interruzione dei regimi), Imaging del Calcio a Campo Largo (apprendimento neurale) e Differenziazione del Corpo Embrionale.

• Stanford Bunny: FLUX preserva il trasporto superficiale in embedding ad alta dimensionalità dove i metodi euclidei falliscono nel rimanere sulla varietà.
• Attrattore di Lorenz: FLUX ottiene un recupero perfetto dei regimi (ARI=1.0, NMI=1.0) identificando l'interruzione tra regimi caotici e subcritici. Crucialmente, FLUX senza Apprendimento della Varietà (utilizzando percorsi euclidei) collassa in un singolo esperto, fallendo nel scoprire il regime nonostante la stessa architettura MoE. Questo dimostra che i percorsi consapevoli della geometria espongono le curvature e le firme di velocità necessarie per la specializzazione degli esperti.
• Dati Neurali: FLUX identifica una transizione nelle dinamiche corticali che si allinea con l'emergere della discriminazione comportamentale (CS+ vs CS−), superando l'instradamento euclideo e le linee di base dinamiche con interruzione.
• Corpo Embrionale: Mentre l'instradamento euclideo performa bene qui (poiché le fasi sono separabili spazialmente), FLUX con apprendimento geometrico rimane competitivo, rafforzando il concetto che la geometria è più critica quando i regimi sono codificati nelle dinamiche locali.

Risultati delle Ablazioni: Il documento afferma esplicitamente che l'instradamento da solo della mixture-of-experts è insufficiente. Senza apprendimento geometrico, il modello può adattarsi al trasporto locale ma fallisce nel recuperare i regimi quando tali regimi sono codificati nelle dinamiche locali piuttosto che in cluster statici.

5. Significato e Affermazioni

Gli autori affermano che la decomposizione della velocità consapevole della geometria fornisce una strategia generale per scoprire transizioni di stati biologici latenti da istantanee longitudinali non accoppiate.

• Impatto Scientifico: FLUX sposta il focus dall'ottimizzazione puramente dell'accuratezza del trasporto (minimizzando la distanza di Wasserstein) al recupero di strutture dinamiche interpretabili. Suggerisce che i sistemi biologici si riorganizzano spesso in regimi dinamici distinti durante processi come l'apprendimento e la differenziazione.
• Limitazioni: Gli autori riconoscono che FLUX è computazionalmente più costoso dell'adattamento del flusso euclideo a causa dell'addestramento di tre reti distinte. Notano anche che il numero di esperti è fissato in anticipo e che i regimi inferiti dovrebbero essere interpretati con cautela quando le dinamiche locali sono rumorose o scarsamente campionate.
• Conclusione: Il documento conclude che FLUX offre un framework per la modellazione dinamica interpretabile di istantanee biologiche non accoppiate, colmando il divario tra l'apprendimento di varietà geometriche e i sistemi dinamici con interruzione dei regimi.