LoCoFlow: Local-Context Flow Matching for Single-Cell Dynamics and Fate Inference

Sintesi Tecnica: LoCoFlow

Problematica

Il sequenziamento a singola cellula fornisce istantanee ad alta risoluzione degli stati cellulari in momenti discreti, ma la natura distruttiva della misurazione impedisce l'osservazione delle vere corrispondenze longitudinali tra le singole cellule. La ricostruzione di traiettorie di sviluppo continue da queste istantanee è un problema mal posto con molteplici soluzioni possibili. Inoltre, i metodi esistenti spesso trattano le cellule come vettori isolati, fallendo nell'هrappresentare adeguatamente la geometria del manifold locale dello spazio degli stati cellulari o l'incertezza intrinseca nella diversificazione del lignaggio. Sebbene cellule con stati trascrizionali simili possano esibire tendenze di destino divergenti, gli attuali approcci di optimal transport e flow-matching tipicamente inferiscono accoppiamenti temporali a livello di singola cellula, mancando di un framework robusto per modellare l'evoluzione dei vicinati cellulari locali come unità coese.

Metodologia: LoCoFlow

Gli autori introducono LoCoFlow (Local-Context Flow Matching), un framework che integra la geometria dei point-cloud locali con l'incertezza del destino condizionale per modellare la dinamica a singola cellula. Il metodo opera in due fasi primarie:

1. Costruzione di Patch Locali e Corrispondenza Gerarchica

Invece di accoppiare cellule isolate, LoCoFlow costruisce "patch" sovrapposte (vicinati locali) attorno a cellule ancora (anchor cells).

• Definizione di Patch: Per i dati di espressione, le patch sono vicinati geodesici su un grafo di k-vicini più prossimi; per i dati spaziali, sono vicinati di dimensione fissa in coordinate misurate.
• Trasporto Gerarchico: Il metodo impiega uno schema di optimal transport (OT) a due livelli. In primo luogo, risolve un problema di trasporto a livello di patch utilizzando la distanza di Bures per allineare i centri e le covarianze delle patch (catturando sia la posizione che la dispersione). In secondo luogo, all'interno delle coppie di patch mantenute, calcola un piano di trasporto a livello di cellula regolarizzato dall'entropia.
• Assegnazione Soft: L'aggregazione di questi piani produce una distribuzione condizionale delle cellule target per ogni sorgente, piuttosto che un'assegnazione hard, preservando molteplici potenziali corrispondenze di successione.

2. Pipeline di Addestramento in Due Fasi

• Fase 1: Trasporto di Bridge Regolarizzato dalla Geometria:
  • Utilizzando i set di endpoint accoppiati derivati dal trasporto gerarchico, il modello apprende i "bridge" (percorsi intermedi) tra le patch sorgente e target.
  • Una rete neurale predice un campo di velocità residuo $f_\theta$ che, combinato con l'interpolazione lineare, assicura che il percorso passi esattamente attraverso gli endpoint.
  • Prior Geometrici: La funzione di perdita include termini di energia cinetica e una penalità radiale che vincola la deformazione relativa (espansione/contrazione) della patch senza limitare la traslazione comune. Penalizza inoltre le velocità ortogonali alla geometria del manifold locale stimata.
• Fase 2: Distillazione della Velocità Condizionale:
  • Poiché lo stato futuro è sconosciuto durante l'inferenza, il bridge (che richiede entrambi gli endpoint) non può essere utilizzato direttamente. LoCoFlow destilla le velocità del bridge in un campo di velocità condizionale $V_\omega$.
  • Latenti Gerarchici: Il modello apprende a generare codici target $c_i$ da un contesto sorgente $Z_0$ utilizzando una struttura latente gerarchica: un latente condiviso dalla patch $\eta_P$ (che cattura gli esiti a livello di popolazione) e latenti specifici per la cellula $\eta_i$ (che catturano la variazione individuale).
  • Addestramento: Un encoder target apprende prima di mappare gli endpoint osservati in codici $c^\star$, quindi un generatore apprende a produrre questi codici dal contesto sorgente e dai latenti. La rete di velocità viene addestrata per corrispondere alle velocità del bridge condizionate su questi codici generati.
  • Inferenza: In fase di inferenza, i codici latenti vengono estratti da un prior (senza osservazioni future) e mantenuti fissi durante l'integrazione. Ciò consente il campionamento di traiettorie alternative che rappresentano l'incertezza del destino.

Contributi Chiave

1. Rappresentazione del Contesto Locale: Sposta l'unità fondamentale della modellazione dinamica dai vettori di espressione isolati agli stati della popolazione locale (point cloud), modellando esplicitamente l'evoluzione della geometria del vicinato.
2. Trasporto Gerarchico: Introduce un meccanismo di optimal transport accoppiato a livello di patch e di cellula che rispetta sia il movimento globale delle popolazioni cellulari sia i loro cambiamenti strutturali interni (es. eterogeneità, espansione).
3. Incertezza del Destino Condizionale: Sviluppa un processo di distillazione in due fasi che separa i vincoli geometrici del percorso dall'incertezza dell'esito, consentendo il campionamento di molteplici traiettorie di destino plausibili senza osservazioni future.
4. Framework di Analisi delle Perturbazioni: Fornisce un meccanismo per distinguere tra effetti focalizzati sulla cellula e sugli effetti del vicinato (background) nelle simulazioni di perturbazione genica, rivelando come il contesto locale modelli le risposte predette.

Risultati Sperimentali

Il paper valuta LoCoFlow su tre dataset: differenziazione di corpi embrioidi (EB), riprogrammazione di iPSC, sviluppo di C. elegans e rigenerazione cerebrale di Axolotl.

• Ricostruzione Held-Out: Sui dataset EB e iPSC, LoCoFlow raggiunge distanze di Wasserstein ($W_1, W_2$) inferiori rispetto ai baseline come OT-CFM, MFM e PACE nella maggior parte delle configurazioni ad alta dimensionalità (es. PCA50). Tuttavia, il paper nota che LoCoFlow a volte presenta una Maximum Mean Discrepancy (MMD) più elevata, suggerendo che le predizioni possano essere più concentrate in regioni specifiche dello spazio degli stati cellulari.
• Predizione del Destino (Dati LARRY): Utilizzando i dati di tracciamento del lignaggio LARRY, LoCoFlow riduce l'errore medio Brier da 0.7990 (scDiffEq) a 0.6620, indicando una migliore accuratezza nella predizione dei vettori di probabilità del destino completo.
• Risposte alle Perturbazioni:
  • Patch Geometriche: Perturbando i geni in patch definite geometricamente, le risposte dell'intero membro e del solo background producono risultati allineati (es. spostamento del destino da monociti a neutrofili).
  • Minibatches Casuali: Al contrario, i modelli che utilizzano minibatch assemblati casualmente mostrano risposte opposte tra cellula focale e background, evidenziando l'importanza del contesto geometrico.
• Interpolazione Spaziale (Axolotl): LoCoFlow migliora la predizione spaziale ($W_2$) per la rigenerazione cerebrale di Axolotl rispetto a OT-CFM, sebbene gli autori notino che le distribuzioni spaziali predette rimangono visivamente simili al time point sorgente, suggerendo la necessità di un ulteriore allineamento spaziale.

Significato e Rivendicazioni

Il paper sostiene che LoCoFlow fornisca un framework robusto per studiare come la struttura della popolazione cellulare locale modelli le traiettorie inferite, rappresentando esplicitamente l'incertezza del destino cellulare. Gli autori sottolineano che:

• Il Contesto Locale è Critico: I distinti pattern di risposta osservati tra patch geometriche e minibatch casuali supportano l'ipotesi che il contesto locale sia una rappresentazione utile per l'inferenza probabilistica del destino.
• Assunzioni di Modellazione: Gli autori mantengono una posizione modesta, riconoscendo che i dati snapshot non identificano univocamente le storie cellulari. I componenti geometrici e latenti di LoCoFlow impongono assunzioni di modellazione piuttosto che rivelare vere relazioni di lignaggio.
• Ipotesi Testabili: Il framework genera ipotesi testabili su come il contesto locale modelli le risposte predette alle perturbazioni, distinguendo tra cambiamenti cellulari intrinseci e influenze ambientali.

Il lavoro non pretende di risolvere l'ambiguità fondamentale dell'inferenza di traiettoria, ma offre un metodo che si allinea meglio alla geometria del manifold degli stati cellulari e fornisce un modo strutturato per quantificare e campionare l'incertezza del destino.