FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

Sintesi Tecnica: FreeBridge – Variational Schrödinger Bridges per la Dinamica di Transizione Cellulare

Problema
I saggi di imaging ad alto contenuto sono essenziali per quantificare le risposte cellulari a perturbazioni chimiche e genetiche. Tuttavia, esiste un limite fondamentale: le cellule sono chimicamente fissate al momento dell'acquisizione delle immagini, rendendo le singole traiettorie temporali non osservabili. Di conseguenza, la modellazione della perturbazione si riduce all'inferenza del trasporto stocastico tra le popolazioni di controllo e trattate, dove sono note solo le distribuzioni marginali finali.

Sebbene i recenti modelli generativi (es. modelli di diffusione, flow matching) raggiungano un forte allineamento tra gli endpoint di controllo e quelli perturbati, la sola coerenza degli endpoint è insufficiente per determinare l'evoluzione intermedia. Infiniti processi stocastici possono connettere marginali identiche; senza ulteriori vincoli, questi modelli possono attraversare regioni dello spazio latente non supportate dalle morfologie cellulari osservate, portando a stati intermedi biologicamente implausibili. I metodi esistenti spesso ottimizzano obiettivi a livello di distribuzione che lasciano sottodeterminate le strutture cellulari intermedie, riuscendo potenzialmente a soddisfare metriche come FID o KID ma generando artefatti.

Metodologia: Il Framework FreeBridge
Gli autori propongono FreeBridge, un framework di trasporto stocastico vincolato formulato come un "Motore Cellulare" che separa esplicitamente la definizione di stati cellulari validi dal meccanismo di transizione. Il metodo opera sotto supervisione di soli endpoint e consiste in tre componenti principali:

1. Unità di Stato (Geometria Single-Cell Esplicita):

   • Invece di trattare le immagini come miscele a livello di popolazione, FreeBridge definisce gli stati atomici come rappresentazioni single-cell segmentate per istanza (utilizzando Cellpose).
   • Un encoder pre-addestrato e fisso $\phi$ imbede questi ritagli in uno spazio latente $z = \phi(x)$. Questo encoder è congelato durante l'apprendimento del trasporto, stabilendo una geometria latente fissa definita prima dell'apprendimento della regola di transizione.
   • Una banca di supporto empirico $Z_{emp}$ viene costruita a partire dagli embedding latenti delle cellule osservate. Questa banca definisce il "manifold" delle morfologie biologicamente valide.

2. Regola di Transizione (Trasporto Schrödinger Bridge):

   • La transizione tra la distribuzione di controllo $\mu_z$ e la distribuzione perturbata $\nu_{c,z}$ è modellata come un'equazione differenziale stocastica (SDE) controllata: $dZ_t = u_t(Z_t) dt + \sigma dW_t$.
   • L'obiettivo è trovare un campo di drift $u_t$ che trasporti la distribuzione iniziale alla distribuzione target al tempo $t=1$.
   • A differenza del trasporto ottimo standard, questa formulazione incorpora un costo di stato specifico per il task per regolarizzare il percorso.

3. Regolarizzazione Geometrica (Vincolo del Supporto Empirico):

   • Per evitare che il processo stocastico attraversi i "vuoti" nello spazio latente, FreeBridge introduce un costo di supporto $V(z) = \lambda \cdot d_{emp}(z)^2$, dove $d_{emp}(z)$ è la distanza dal vicino più prossimo nella banca di supporto empirico $Z_{emp}$.
   • L'obiettivo totale minimizza l'energia cinetica (magnitudo del drift) e il costo di supporto:
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • Questa regolarizzazione ancora gli stati intermedi al manifold morfologico empiricamente osservato senza richiedere traiettorie accoppiate a livello di cellula.

Contributi Chiave

• Formulazione: Introduce una formulazione di Schrödinger Bridge per la modellazione delle transizioni single-cell che separa esplicitamente la specifica dello stato dal trasporto stocastico.
• Fondamento Geometrico: Propone un termine di regolarizzazione innovativo basato sulla stima non parametrica della densità (distanza del vicino più prossimo rispetto a una banca di supporto empirico) per vincolare le traiettorie intermedie all'interno di regioni morfologicamente plausibili.
• Protocollo Unificato: Stabilisce un protocollo di valutazione unificato attraverso molteplici dataset (BBBC021, RxRx1, JUMP) utilizzando ritagli single-cell segmentati per istanza, garantendo un confronto equo tra i metodi.
• Performance: Dimostra che il vincolo della geometria migliora non solo la fedeltà dell'endpoint, ma anche la ritenzione semantica dei meccanismi di perturbazione (Mechanism-of-Action).

Risultati
Valutato sui dataset BBBC021, RxRx1 e JUMP, FreeBridge dimostra quanto segue:

• Fedeltà dell'Endpoint: Raggiunge punteggi Fréchet Inception Distance (FID) e Kernel Inception Distance (KID) competitivi o migliorati rispetto ai baseline allo stato dell'arte (PhenDiff, IMPA, CellFlux). Su BBBC021, raggiunge il FID più basso (15.35) e il KID più basso (1.10).
• Consistenza Intermedia: Riduce significativamente il Tasso di Violazione del Supporto ($R_{viol}$), una metrica che misura la proporzione di stati intermedi che cadono al di fuori del supporto morfologico empirico. Su BBBC021, FreeBridge riduce $R_{viol}$ a 0.11, rispetto allo 0.38 di PhenDiff e allo 0.31 di CellFlux.
• Ritenzione Semantica: Gli studi di ablazione mostrano che la rimozione della regolarizzazione geometrica diminuisce l'accuratezza della classificazione del Mechanism-of-Action (da 72.08% a 65.08% su BBBC021), indicando che il vincolo dell'evoluzione intermedia aiuta a preservare i segnali meccanicistici.
• Analisi Qualitativa: Le visualizzazioni confermano che FreeBridge preserva meglio le caratteristiche morfologiche specifiche dei composti (es. compattazione nucleare, restringimento citoplasmatico) rispetto ai baseline, che spesso mostrano smoothing della texture o pattern incompleti.

Significato e Rivendicazioni
L'articolo sostiene che modellare le transizioni biologiche richiede più del semplice allineamento dei confini; richiede l'esplicito fondamento geometrico dell'evoluzione intermedia. Ancorando il trasporto stocastico a un manifold cellulare empiricamente supportato, FreeBridge fornisce un approccio rigoroso per generare dinamiche di perturbazione biologicamente interpretabili. Gli autori affermano che questa formulazione offre una base robusta per la simulazione di "cellule virtuali", affrontando il gap critico in cui i modelli allineati agli endpoint falliscono nel catturare stati cellulari intermedi realistici. Il lavoro evidenzia come, senza vincoli geometrici, i modelli generativi possano produrre endpoint statisticamente validi che attraversano regioni dello spazio latente biologicamente non supportate.

Limitazioni
Gli autori riconoscono che FreeBridge dipende da un encoder congelato; i fattori morfologici assenti dal pre-addestramento dell'encoder non possono essere recuperati tramite il solo apprendimento del trasporto. Inoltre, il calcolo della distanza di supporto scala quadraticamente ($O(B^2)$) con la dimensione del minibatch, il che potrebbe richiedere una ricerca del vicino più prossimo approssimata per banche di supporto più grandi.