Riepilogo Tecnico: Flow Matching per Dati di Conteggio (count-FM)

Enunciato del Problema

I dati di conteggio ad alta dimensionalità, prevalenti in campi come il sequenziamento dell'RNA a cellula singola (scRNA-seq) e l'analisi delle sequenze di spike neuronali, presentano sfide uniche per la modellazione generativa. Questi tipi di dati sono caratterizzati da discretezza, sparsità e strutture di correlazione complesse. Gli approcci generativi profondi esistenti spesso faticano in questo dominio:

• Approcci a Stato Categorico: I metodi che trattano i conteggi come variabili categoriche (ad esempio, token) richiedono dimensioni di output che scalano con il massimo intervallo di conteggio. Ciò porta a una scarsa efficienza dei parametri e tratta i conteggi adiacenti come stati non correlati, ignorando la geometria intrinseca dello spazio dei conteggi.
• Trasformazioni Continue: Gli approcci che mappano i conteggi in spazi continui tramite dequantizzazione o rappresentazioni latenti sostituiscono la massa di probabilità discreta con densità continue, potenzialmente offuscando la struttura discreta sottostante.
• Limitazioni dei Modelli di Conteggio Esistenti: Sebbene esistano alcuni modelli di salto specifici per i conteggi (ad esempio, Poisson-JUMP), spesso non sono progettati per il trasporto tra distribuzioni di conteggio arbitrarie o mancano delle proprietà di addestramento e generazione efficienti dei moderni metodi basati su flow.

Il documento affronta la necessità di un framework generativo che operi direttamente nello spazio dei conteggi, rispetti la geometria discreta dei dati e consenta un trasporto efficiente tra distribuzioni sorgente e target arbitrarie senza il sovraccarico parametrico dei modelli a stato categorico.

Metodologia: count-FM

Gli autori propongono count-FM, un framework di flow matching adattato per i dati di conteggio basato su un processo di nascita-morte a tempo continuo con salti unitari locali.

1. Processo di Nascita-Morte a Tempo Continuo

Il modello definisce la transizione tra una distribuzione sorgente $p_0$ e una distribuzione target $p_1$ come un processo di salto di Markov a tempo continuo (CTMC) sullo spazio dei vettori di conteggio $\mathbb{N}_0^d$. Il processo evolve tramite nascite unitarie locali ($+1$) e morti ($-1$) per ciascuna coordinata $i$:

• Tasso di Nascita: $\lambda_{t,i}(x)$
• Tasso di Morte: $\mu_{t,i}(x)$
• Dinamica: In un piccolo intervallo $h$, la probabilità di transizione dallo stato $x$ a $x + e_i$ è $h \lambda_{t,i}(x)$, e a $x - e_i$ è $h \mu_{t,i}(x)$. La morte non è consentita quando un conteggio è zero.

Questa parametrizzazione è altamente efficiente: invece di prevedere una distribuzione categorica completa su tutti i possibili livelli di conteggio (che scala con il massimo conteggio), il modello prevede solo due tassi per dimensione ($2d$ output), indipendentemente dall'intervallo di conteggio.

2. Ponte Binomiale Condizionale e Addestramento

Per addestrare il modello, gli autori definiscono un ponte condizionale tra estremi fissi $(x_0, x_1)$. Per ciascuna coordinata $i$, il ponte è costruito utilizzando una distribuzione Binomiale:
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
Questo ponte garantisce che la media condizionale si sposti linearmente da $x_0$ a $x_1$ mantenendo stati intermedi a valori interi.

Imponendo la conservazione della massa (l'equazione di Kolmogorov in avanti) su questo ponte condizionale, gli autori derivano espressioni in forma chiusa per i tassi di nascita e morte condizionali:
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
dove $(\cdot)^+$ denota la rettificazione positiva.

3. Obiettivo di Addestramento

Il modello viene addestrato minimizzando una divergenza Kullback-Leibler (KL) nello spazio dei percorsi tra il processo del ponte condizionale e il processo appreso dal modello. Ciò si decompone in un obiettivo locale di matching dei tassi:
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
dove $\ell(u, v) = v - u \log v$ è un termine KL generalizzato. Ciò consente un addestramento senza simulazione dei tassi di transizione marginali apprendendo i tassi condizionali identificabili.

4. Campionamento e Accoppiamento

• Campionamento: I campioni sono generati simulando il processo di nascita-morte appreso in avanti da $t=\epsilon_t$ a $t=1-\epsilon_t$ utilizzando una discretizzazione di salto locale del primo ordine.
• Accoppiamento degli Estremi: Il documento esplora due strategie di accoppiamento per l'abbinamento di campioni sorgente e target:
  • Accoppiamento Indipendente: Accoppiamento casuale ($p_0(x_0)p_1(x_1)$).
  • Accoppiamento Trasporto Ottimale (OT): Accoppiamento OT su mini-batch utilizzando un costo di Poisson simmetrico. Ciò tende a indurre transizioni con minore curvatura, migliorando l'interpretabilità e l'efficienza del campionamento.

5. Generazione Condizionale

Per compiti condizionali, il modello incorpora covariate $y$ e utilizza la Guidance senza Classificatore (CFG). Durante l'addestramento, le variabili di condizionamento vengono casualmente omesse. Durante il campionamento, i tassi condizionali e incondizionati sono combinati tramite una scala di guida $w$ per controllare la forza della condizione.

Contributi Chiave

1. Framework count-FM: Un framework di flow matching per dati di conteggio basato su dinamiche di nascita-morte locali, che consente il trasporto direttamente nello spazio dei conteggi con stati intermedi che rimangono a valore di conteggio.
2. Schema di Addestramento Tracciabile: Un metodo per apprendere i tassi di transizione marginali tramite un ponte binomiale condizionale, mantenendo l'efficienza del flow matching evitando la complessità della modellazione a stato categorico.
3. Efficienza Parametrica: Una formulazione che prevede solo tassi locali di nascita e morte ($2d$ output) anziché distribuzioni categoriche complete, rendendolo adatto a dati ad alta dimensionalità con ampi intervalli di conteggio.

Risultati Sperimentali

1. Simulazione (Mixture Gamma-Poisson 2D)

• Prestazioni: count-FM e count-FM-OT hanno raggiunto la migliore qualità del campione (minima distanza 2-Wasserstein e MMD$^2$) rispetto alle baseline, tra cui D3PM, discrete-FM, tauLDR, SEDD e Poisson-JUMP.
• Efficienza: count-FM ha raggiunto questi risultati con un numero sostanzialmente inferiore di parametri addestrabili (2.372) rispetto alle baseline a stato categorico (14.000–67.000+).
• Geometria: Gli stati intermedi si sono evoluti in modo fluido nello spazio dei conteggi. La versione accoppiata OT ha seguito un percorso di transizione visibilmente più rettilineo e ha raggiunto la qualità target con un minor numero di valutazioni della funzione (NFE).

2. RNA-seq a Cellula Singola (Gyrus Dentato)

• Generazione Incondizionata: Su un dataset di 2.930 cellule e 13.913 geni, count-FM ha superato le baseline a variabili latenti (scVI) e i modelli di diffusione latente (scDiffusion, scLDM) nella qualità del campione (W2 e MMD$^2$) utilizzando meno parametri (circa 9,8M contro 29M–39M per i concorrenti).
• Trasporto (P12 a P35): Il modello ha modellato con successo il trasporto dello sviluppo dal giorno postnatale 12 al giorno 35. Utilizzando un accoppiamento OT a linea di discendenza ristretta, i percorsi di trasporto hanno rispettato le strutture di sviluppo note (ad esempio, la progressione da neuroblasto a cellula granulare), producendo distribuzioni intermedie interpretabili che sono rimaste nello spazio dei conteggi.

3. Sequenze di Spike Neuronali (Ippocampo/Entorinale)

• Generazione Condizionale: Il modello è stato applicato per generare conteggi di spike condizionati alla posizione linearizzata e alla direzione di corsa.
• Confronto: count-FM con scala di guida $w=1$ ha superato un regressore MLP deterministico e una baseline Poisson MLP.
• Struttura di Correlazione: A differenza della Poisson MLP, che sottostimava le correlazioni di popolazione, count-FM ha preservato meglio la struttura di dipendenza tra neuroni.
• Compromesso della Guida: $w=1$ ha fornito il miglior equilibrio tra sintonizzazione spaziale e calibrazione distribuzionale. $w=2$ ha affinato i campi di luogo ma ha distorto la calibrazione dell'ampiezza.

Significato e Affermazioni

Il documento afferma che count-FM fornisce un framework efficace per la modellazione di dati di conteggio ad alta dimensionalità consentendo la generazione e il trasporto direttamente nello spazio dei conteggi. Il suo significato principale risiede in:

• Fedeltà Geometrica: Modellando il trasporto tramite salti unitari locali, il metodo rispetta la geometria discreta dei dati di conteggio, evitando gli effetti di offuscamento delle approssimazioni continue.
• Scalabilità: L'efficienza parametrica lo rende praticabile per applicazioni ad alta dimensionalità (ad esempio, migliaia di geni) dove i metodi a stato categorico diventano computazionalmente proibitivi.
• Interpretabilità: Gli stati intermedi nel percorso di trasporto sono campioni di conteggio validi, consentendo l'ispezione di distribuzioni transizionali significative in applicazioni come la biologia dello sviluppo.
• Versatilità: Il framework supporta la generazione incondizionata, il trasporto tra distribuzioni arbitrarie e la generazione condizionale con guida controllabile.

Gli autori riconoscono le limitazioni, notando che il ponte binomiale condizionale è relativamente semplice e potrebbe non esplorare tutti i percorsi intermedi plausibili, e che l'attuale discretizzazione del primo ordine potrebbe essere migliorata con schemi di simulazione più accurati. Tuttavia, i risultati empirici su dataset di simulazione e biomedici dimostrano che count-FM raggiunge un'alta qualità del campione con una maggiore efficienza di modellazione rispetto agli approcci concorrenti.