Integrating gene regulatory priors into Transformer attention with scTransformer for interpretable scRNA-seq analysis

Sintesi Tecnica: scTransformer – Integrare i Prior Regolatori nella Attention dei Transformer

Problematica

I dati di sequenziamento dell'RNA a singola cella (scRNA-seq) sono caratterizzati da alta dimensionalità, sparsità e rumore tecnico, il che pone sfide significative per l'apprendimento delle rappresentazioni. Sebbene i modelli foundation basati su Transformer (ad esempio, scGPT, Geneformer) abbiano dimostrato scalabilità pre-addestrandosi su milioni di cellule, essi trattano tipicamente i geni come token indipendenti con auto-attenzione non vincolata. Questo design permette a qualsiasi gene di prestare attenzione a qualsiasi altro, affidandosi esclusivamente alla covariazione statistica per scoprire le dipendenze. Di conseguenza, questi modelli soffrono spesso di:

1. Interpretabilità Limitata: I pattern di attenzione sono difficili da mappare su processi biologici meccanicistici.
2. Instabilità: In regimi di dati scarsi o in presenza di effetti batch, i modelli non vincolati possono sovra-adattarsi (overfitting) a correlazioni specifiche del dataset, producendo previsioni che sono robuste in termini di accuratezza ma instabili nella selezione delle feature (ovvero, diverse esecuzioni di addestramento selezionano set di geni disgiunti).
3. Conoscenza Biologica Ignorata: Gli approcci esistenti ignorano ampiamente le reti regolatorie curate tra fattori di trascrizione (TF) e geni target (TG), che codificano strutture causali plausibili nel controllo trascrizionale.

Metodologia: Architettura scTransformer

Gli autori propongono scTransformer, un encoder Transformer biologicamente informato che integra reti regolatorie TF–TG curate direttamente nel meccanismo di auto-attenzione tramite masking rigido (hard masking).

1. Rappresentazione dell'Input e Embedding:

• Gene Token Embedding: Ogni gene in una cella viene mappato in un token a $d$ dimensioni che combina un embedding di identità del gene addestinabile ($e^{id}$) e un encoding dell'espressione continua ($e^{expr}$).
• Encoding Sinusoidale Continuo: A differenza del binning discreto, i valori di espressione sono codificati utilizzando uno schema sinusoidale continuo ($\sin(x \cdot \omega_k), \cos(x \cdot \omega_k)$). Ciò preserva l'informazione quantitativa attraverso l'intero range dinamico delle misurazioni scRNA-seq, garantendo fluidità e limitatezza.
• Pooling: Un meccanismo di pooling basato sulla similarità aggrega sequenze di token a lunghezza variabile in un embedding cellulare a dimensionalità fissa. Il vettore di contesto è computato come la media degli embedding dei token ed è utilizzato come query in una cross-attention sulla sequenza dei token.

2. Meccanismo di Attention con Prior-Gating:
L'innovazione principale è la sostituzione dell'attenzione non vincolata con un meccanismo prior-gated:

• Rete Regolatoria: Una rete TF–TG curata (da CollecTRI) viene pre-elaborata per filtrare i TF con un numero sufficiente di target e i geni presenti nel dataset.
• Hard Masking: Viene costruita una matrice di permessi booleana ($M_{allow}$) dove $M_{allow}(i, j) = 1$ solo se il gene $i$ è un TF e il gene $j$ è il suo target (o $i=j$ per auto-loop). Tutte le altre interazioni sono impostate a 0.
• Vincolo di Attention: Questa matrice viene convertita in una maschera additiva ($\tilde{M}_{prior}$) dove le interazioni proibite ricevono il valore $-\infty$. Prima dell'operazione softmax nell'attenzione dot-product scalata, questa maschera viene aggiunta ai punteggi di attenzione:
  $$A_h = \text{softmax}\left(\frac{Q_h K_h^\top}{\sqrt{d_k}} + \tilde{M}_{prior} + M_{pad}\right)$$
  Ciò assicura che l'informazione fluisca solo lungo i bordi regolatori noti (e gli auto-loop), vincolando efficacemente lo spazio delle ipotesi verso interazioni biologicamente plausibili.

3. Metriche di Interpretabilità:
Per quantificare la rilevanza biologica dell'attenzione appresa, gli autori definiscono metriche a livello di modulo:

• Punteggio di Concentrazione ($\phi$): Misura quanto selettivamente un TF concentra l'attenzione sui suoi target annotati utilizzando l'entropia normalizzata. Un'entropia inferiore (maggiore $\phi$) indica un'attenzione focalizzata e strutturata.
• Importanza del Modulo: Combina il punteggio di concentrazione con la massa totale di attenzione diretta verso i target per quantificare la forza dei segnali regolatori.

4. Dettagli di Implementazione:

• Dati: Valutato sul atlante di RNA a singola nucleo MSSM (coorte PsychAD), limitato a 17 sottotipi e intersezioni TF–TG.
• Addestramento: Utilizza l'ottimizzatore AdamW con architettura Pre-LN. L'addestramento impiega il batching sparso operando direttamente su matrici Compressed Sparse Row (CSR) per gestire i grandi dataset in modo efficiente.
• Valutazione: I modelli sono stati testati attraverso diversi regimi di dati di addestramento (dall'1% al 100%) per valutare l'efficienza campionaria e la robustezza.

Contributi Chiave

1. Prior-Gated Attention: Il primo framework Transformer che incorpora reti TF–TG curate come vincoli rigidi all'interno del meccanismo di attenzione, trasformando la conoscenza regolatoria in un bias induttivo strutturale.
2. Continuous Expression Encoding: Uno schema di encoding sinusoidale che preserva la natura quantitativa dei dati scRNA-seq senza discretizzazione.
3. Pipeline di Addestramento Sparsa: Un'implementazione efficiente che opera direttamente su matrici CSR, consentendo l'addestramento scalabile su grandi dataset di singola cella.
4. Metriche di Interpretabilità Formalizzate: Un set di metriche per quantificare la concentrazione dell'attenzione e l'importanza regolatoria, permettendo la valutazione diretta della coerenza biologica senza analisi post-hoc.
5. Protocollo Riproducibile: Una configurazione di addestramento distribuito con partizionamento controllato dei dati e valutazione multi-seed per valutare rigorosamente la stabilità.

Risultati

Gli autori hanno valutato scTransformer contro baseline non vincolati in compiti di classificazione dei tipi cellulari:

• Performance in Regimi di Bassi Dati: In scenari con dati di addestramento limitati (1% e 25%), il modello prior-gated ha raggiunto un'accuratezza significativamente superiore rispetto al baseline non vincolato (ad esempio, 0.486 vs 0.464 all'1% di dati). Ciò dimostra una migliore efficienza campionaria.
• Performance in Regimi di Dati Completi: Con la disponibilità del 100% dei dati, il modello prior-gated (utilizzando solo i geni TF–TG) ha raggiunto un'accuratezza comparabile al baseline non vincolato che utilizza tutti i geni (0.911 vs 0.917), nonostante operi in uno spazio di interazione molto più piccolo.
• Stabilità e Riproducibilità: I modelli non vincolati hanno mostrato un'elevata variabilità tra le esecuzioni di addestramento, selezionando set di geni disgiunti (indici di Jaccard vicini allo zero, correlazioni di Spearman negative). Al contrario, i modelli prior-gated hanno mostrato un'alta stabilità, selezionando costantemente subset di geni simili e classificandoli in modo coerente.
• Coerenza Biologica: Le mappe di attenzione dei modelli prior-gated hanno mostrato punteggi di concentrazione ($\phi$) più elevati, indicando che i TF concentravano l'attenzione su subset di target specifici e biologicamente rilevanti. Questi pattern si allineavano con noti programmi regolatori specifici per sottotipo (ad esempio, l'arricchimento di GRHL3 nelle popolazioni endoteliali-simili, NANOG negli OPC), mentre i modelli non vincolati producevano distribuzioni di attenzione diffuse e meno interpretabili.

Significato e Rivendicazioni

Il documento sostiene che il valore primario dell'integrazione dei prior regolatori non sia semplicemente il miglioramento dell'accuratezza predittiva, ma il potenziamento dell'affidabilità epistemica e dell'interpretabilità delle rappresentazioni apprese.

• Risoluzione dell'Ambiguità: In dati ad alta dimensionalità e sparsi, molteplici combinazioni di geni possono spiegare lo stesso fenotipo. I modelli non vincolati possono convergere verso soluzioni arbitrarie e instabili. I prior regolatori vincolano lo spazio delle ipotesi, guidando il modello verso soluzioni biologicamente fondate e coerenti.
• Bias Induttivo: Incorporando la struttura biologica direttamente nell'architettura, il modello apprende rappresentazioni che sono robuste al rumore e agli effetti batch senza sacrificare la flessibilità quando sono disponibili dati sufficienti.
• Fondamento per l'Interpretabilità: L'approccio offre un passo procedurale verso "modelli foundation biologicamente fondati", dove i pesi dell'attenzione sono direttamente interpretabili come programmi regolatori, facilitando l'insight biologico senza richiedere analisi post-hoc.

Gli autori rimangono modesti, notando che le attuali limitazioni includono l'uso di maschere binarie rigide (che impediscono l'apprendimento di nuove interazioni), la restrizione ai compiti supervisionati e la valutazione su un singolo dataset. Suggeriscono che il lavoro futuro potrebbe esplorare prior soft/apprendibili, pre-addestramento auto-supervisionato e architetture ibride.