Scaling an Autoregressive Transformer for Single-Cell Generation

Sintesi Tecnica: Scalabilità di un Transformer Autoregressivo per la Generazione di Singole Cellule

Definizione del Problema
Il paper affronta la sfida di generare vettori di espressione genica realistici per specifici tipi cellulari. Mentre i metodi esistenti rappresentano le linee cellulari tramite nomi categorici, singoli vettori di esempio o sequenze di vettori, è necessario un framework scalabile e auto-supervisionato in grado di generare profili di espressione genica sintetici per aumentare le misurazioni scarse e consentire la sperimentazione in-silico. Inoltre, il campo manca di una chiara comprensione di come la perdita di pre-addestramento (pretraining loss) scala rispetto alla dimensione del modello ($N$) e alla dimensione dei dati ($D$) per i modelli di fondazione single-cell, con segnalazioni contrastanti sul fatto che questi modelli seguano una legge di potenza prevedibile o si livellino (plateau) a causa dei vincoli di diversità dei dati.

Metodologia
Gli autori propongono un'architettura transformer autoregressiva accoppiata a un tokenizer VAE (Variational Autoencoder) quantizzato e appreso. L'approccio prevede due componenti principali addestrate sequenzialmente:

1. Tokenizer (RQ-VAE): Un VAE a quantizzazione residua (Residual Quantized VAE) mappa i vettori continui di espressione genica single-cell (18.080 geni) in una sequenza di 8 token discreti. L'encoder comprime l'input in un vettore latente a 256 dimensioni, che viene poi quantizzato da grossolano a fine attraverso otto code di codici (codebooks) residue. Questo tokenizer viene addestrato una sola volta sul dataset scBaseCount e poi congelato per servire come strato di rappresentazione discreta per il transformer.
2. Modello Transformer: Un transformer causale (basato sull'architettura LLaMA) modella le sequenze di questi token discreti. Una "frase" viene costruita concatenando blocchi di token di 32 cellule dello stesso tipo cellulare, separati da token di separazione appresi. Il modello viene addestrato tramite la predizione del next-token utilizzando la perdita di cross-entropy.
3. Griglia di Scalabilità (Scaling Grid): Per investigare le leggi di scalabilità, gli autori hanno costruito una griglia controllata variando indipendentemente la dimensione del modello ($N$) e la dimensione dei dati ($D$).
   • Dimensioni del Modello: 5 configurazioni da 1,31M a 83,9M di parametri, mantenendo un rapporto fisso tra larghezza e profondità.
   • Dimensioni dei Dati: Budget di addestramento da 38,1M a 362M di cellule, campionati da un sottoinsieme fisso di circa 200M di cellule dal dataset scBaseCount con campionamento proporzionale per preservare la composizione naturale dei tipi cellulari.
   • Totale: Sono stati valutati 22 punti distinti $(N, D)$.

Contributi Chiave

1. Formulazione del Task: Definizione di un compito di apprendimento auto-supervisionato per generare vettori di espressione genica sintetici dato un contesto di linea cellulare, utilizzando la tokenizzazione discreta.
2. Metriche di Valutazione: Stabilimento di metriche downstream per la fedeltà biologica, inclusa la correlazione di Pearson dell'espressione media, l'Errore Assoluto Medio (MAE) e un punteggio di discriminabilità del tipo cellulare.
3. Correlazione della Perdita: Dimostrazione che la perdita di cross-entropy di pre-addestramento correla fortemente con le metriche di qualità della generazione downstream.
4. Legge di Scalabilità Congiunta: Applicazione di una legge di scalabilità con due esponenti $L(N, D) = E + A N^{-\alpha} + B D^{-\beta}$ alla perdita di pre-addestramento. Per quanto ne sappiano gli autori, questo è il primo tale fit per un transformer single-cell su una griglia $N \times D$ variata indipendentemente.
5. Frontiera Computazionalmente Ottimale: Derivazione di una strategia di allocazione computazionalmente ottimale basata sugli esponenti calcolati, suggerendo che, alle scale attuali, le risorse computazionali sono meglio spese nell'aumentare il volume dei dati piuttosto che i parametri del modello.

Risultati

• Qualità della Generazione: Il modello genera vettori specifici per tipo cellulare che corrispondono strettamente ai replicati empirici. Al checkpoint di perdita di validazione più bassa, la popolazione generata ha raggiunto una correlazione di Pearson di 0,960 e un MAE di 0,041 rispetto ai dati empirici tenuti in riserva (held-out), avvicinandosi quasi al baseline del confronto tra due sottoinsiemi empirici disgiunti (0,964 e 0,038, rispettivamente). La discriminabilità del tipo cellulare era di 0,990, indicando che il generatore non collassa verso un singolo modo.
• Fit della Legge di Scalabilità: Il fit della legge di scalabilità ha prodotto esponenti $\alpha \approx 0,81$ (dimensione del modello) e $\beta \approx 0,57$ (dimensione dei dati), con una perdita minima irreversibile $E \approx 2,90$ nats. Il fit ha raggiunto un $R^2$ di 0,977.
• Potere Predittivo: La legge calcolata ha predetto le esecuzioni tenute in riserva (variando sia $N$ che $D$) con errori entro lo 0,5% dai valori misurati, validando la capacità di estrapolazione del modello.
• Frontiera Computazionalmente Ottimale: Poiché $\alpha > \beta$, l'allocazione ottimale pende verso i dati ($D^* \propto C^{0,59}$ rispetto a $N^* \propto C^{0,41}$).

Significatività e Rivendicazioni
Il paper sostiene di risolvere un dibattito riguardante le leggi di scalabilità single-cell. Mentre studi precedenti (il cluster "diversità sopra la scala") riportavano che i modelli single-cell si livellano a grandi scale, attribuendo ciò a una mancanza di diversità del contesto, gli autori sostengono che si tratti di un'interpretazione errata di colli di bottiglia della capacità. I loro risultati mostrano che, quando $N$ e $D$ sono variati congiuntamente, emerge una chiara legge di potenza. I plateau osservati nei lavori precedenti sono attribuiti all'aumento dei dati a una dimensione del modello fissa o non controllata, dove il modello manca della capacità di sfruttare ulteriori dati una volta che il termine dei dati scende al di sotto del termine della capacità.

Gli autori pongono questo lavoro come un passo fondamentale per i modelli di fondazione single-cell, fornendo una frontiera computazionalmente ottimale per l'addestramento. Discutono anche un potenziale' applicazione di "Fase 2": il fine-tuning del modello pre-addestrato per la predizione della risposta alle perturbazioni. Tuttavia, notano un'osservazione durante il pre-addestramento in cui una perdita su un benchmark di perturbazione tenuto in riserva traccia una traiettoria a "forma di segno di spunta" (diminuisce poi aumenta) distinta dalla perdita di scalabilità principale, suggerendo un mismatch di composizione tra il corpus generale di pre-addestramento e i contesti specifici di perturbazione. Questo rimane una questione aperta per il lavoro futuro piuttosto che un risultato riportato.

Il paper conclude che l'architettura proposta sintetizza con successo vettori di espressione genica con alta fedeltà biologica e che il comportamento di scalabilità dei transformer single-cell segue leggi di potenza congiunte prevedibili, simili a quelle del linguaggio naturale, a condizione che la griglia di addestramento sia sufficientemente controllata.