Sintesi Tecnica: Knowledge Graph e LLM di Ragionamento per la Ricerca di Predittori di Perturbazione Trascrizionale Semplici ma Efficaci

1. Definizione del Problema

Prevedere gli effetti trascrizionali di perturbazioni di knockout genico non viste rimane una sfida significativa per i modelli di cellula virtuale. Sebbene approcci recenti sfruttino i grafi della conoscenza biologica per migliorare l'estrapolazione oltre le perturbazioni di addestramento, i complessi modelli di deep learning spesso faticano a superare i semplici baseline (ad esempio, la media dei campioni di addestramento). La difficoltà principale risiede nel generalizzare a nuove distribuzioni interventive, $P(x|do(i))$, dove ogni perturbazione induce una distribuzione unica vincolata da processi biologici complessi. I metodi esistenti che utilizzano Reti Neurali su Grafi (GNN) o pre-addestramento su larga scala hanno mostrato potenziale, ma i meccanismi specifici attraverso i quali la conoscenza biologica abilita la generalizzazione a nuove intervezioni rimangono poco chiari. Inoltre, ottimizzare direttamente la struttura causale sottostante è computazionalmente intrattabile a causa dello spazio di ricerca esponenziale dei possibili grafi genici.

2. Metodologia

2.1 Fondamento Teorico: Un Modello Gerarchico Sparso

Gli autori postulano che le risposte cellulari ai knockout genici seguano un grafo causale gerarchico sparso. In questo modello, i geni formano complessi proteici, che formano pathway, portando infine a organelli e risposte cellulari.

• Implicazione: Le perturbazioni che colpiscono lo stesso complesso di basso livello dovrebbero produrre profili di espressione simili.
• Predizione: L'espressione media di una perturbazione target può essere approssimata da una semplice media dei vicini più prossimi (nearest-neighbour) definita dal vicinato locale di questo grafo causale.

2.2 Baseline: Il Predittore del Vicino di Grafo (Graph Neighbour Predictor)

Gli autori valutano innanzitutto un predittore "Graph Neighbour", che sfrutta l'esistente conoscenza biologica (es. StringDB, GO) come prior.

• Meccanismo: Per un gene target $g$, il predittore identifica i suoi vicini $N(g)$ nel grafo della conoscenza. La delta della perturbazione predetta ($\hat{\delta}_g$) è calcolata come la media delle delta di addestramento di questi vicini:
  $$\hat{y}_g = \bar{x}_g + \frac{1}{|N_{Tr}(g)|} \sum_{\tilde{g} \in N_{Tr}(g)} \bar{\delta}_{\tilde{g}}$$
  dove $\bar{x}_g$ è l'espressione di controllo corrispondente e $N_{Tr}(g)$ rappresenta i vicini presenti nel set di addestramento.
• Risultato: Questo semplice approccio K-nearest neighbour supera i complessi modelli di deep learning (es. GEARS, scLAMBDA) e si avvicina alle prestazioni dei modelli allo stato dell'arte (SOTA) come TxPert, particolarmente in contesti out-of-distribution (OOD).

2.3 Raffinamento: Reinforcement Learning con LLM di Ragionamento

Riconoscendo che i grafi della conoscenza sono incompleti e contengono archi subottimali, gli autori propongono di utilizzare i Large Language Models (LLM) per raffinare questi vicinati.

• Riduzione dello Spazio di Ricerca: Invece di cercare nello spazio intrattabile di tutti i possibili grafi ($O(2^{\#genes^2})$), la ricerca è limitata a candidati archi plausibili sotto le credenze a priori dell'LLM.
• Formulazione RL:
  • Policy: Un LLM ($\pi_\theta$) agisce come una policy per selezionare azioni (Aggiungi $A$, Rimuovi $R$) su un set iniziale di vicini del grafo della conoscenza ($N_{KG}$).
  • Spazio delle Azioni: Il set raffinato è $S_{LLM} = (N_{KG} \setminus R) \cup A$.
  • Funzione di Reward: Il modello viene addestrato tramite Group Relative Policy Optimization (GRPO) per massimizzare un reward $R(s, a)$ composto da:
    1. Reward di Formattazione ($R_{Fmt}$): Assicura che l'output sia parsabile.
    2. Reward di Validità ($R_{Valid}$): Incoraggia l'aggiunta di geni dal set di addestramento e la rimozione di geni non presenti nel vicinato originale.
    3. Reward di Performance ($R_\Delta$): Incentiva il miglioramento del metrico Pearson Delta rispetto al baseline del predittore del grafo della conoscenza.
• Addestramento: L'LLM è addestrato su dati di perturbazione dalle linee cellulari di Replogle et al. (2022), con una rigorosa divisione train/test per garantire la valutazione su perturbazioni non viste.

3. Contributi Chiave

1. Dimostrazione di Baseline Semplici: Il documento stabilisce che un semplice predittore K-nearest neighbour basato su grafi della conoscenza biologica è un baseline altamente competitivo per la predizione di perturbazioni OOD, superando la maggior parte dei complessi metodi di deep learning.
2. L'LLM come Prior di Raffinamento: Gli autori dimostrano che gli LLM, quando ottimizzati tramite Reinforcement Learning, possono modificare efficacemente i vicinati dei grafi della conoscenza (aggiungendo relazioni mancanti e rimuovendo quelle irrilevanti) per migliorare la performance predittiva.
3. Prestazioni allo Stato dell'Arte: L'LLM raffinato tramite RL raggiunge prestazioni equivalenti ai metodi SOTA attuali (specificamente TxPert) sulle linee cellulari di Replogle et al. (2022), partendo da una semplice architettura di nearest-neighbour.
4. Generalizzazione Downstream: Lo studio mostra che l'addestramento RL migliora le prestazioni dell'LLM su task downstream (specificamente la predizione dell'espressione differenziale sul benchmark PerturbQA), suggerendo che il modello apprenda pattern di ragionamento generalizzabili piuttosto che fare overfitting sul compito di predizione specifico.

l. Risultati

• Performance nei Benchmark: Su quattro linee cellulari (K562, RPE1, HEPG2, Jurkat) di Replogle et al. (2022), l'LLM addestrato con RL ($LLM + KG + RL$) ha ottenuto punteggi Pearson Delta comparabili a TxPert (es. 0.625 vs 0.626 su K562) e ha superato significativamente l'LLM non addestrato e il baseline del grafo della conoscenza grezzo.
• Generalizzazione OOD: Il metodo generalizza a perturbazioni non viste, mantenendo alte prestazioni anche quando i knockout genici specifici non erano presenti nel set di addestramento.
• Transfer Learning: Quando addestrato su sole tre delle quattro linee cellulari, il modello ha comunque migliorato le prestazioni sulla linea cellulare esclusa (K562) rispetto al modello base, indicando una generalizzazione della policy di ragionamento tra diverse linee cellulari.
• Task Downstream: Sul dataset PerturbQA (Wu et al.), il modello post-addestrato ha superato l'LLM base nella predizione dell'espressione differenziale e della direzione del cambiamento, anche senza un fine-tuning specifico per questo compito.

5. Significato e Rivendicazioni

Il documento afferma che i grafi della conoscenza biologica fungono da efficaci prior per il modello, e che il modello più semplice che sfrutta tali prior (nearest neighbour) è sorprendentemente robusto. Il significato primario risiede nel dimostrare che il Reinforcement Learning può raffinare gli LLM in strumenti generalizzabili per la complessa predizione biologica.

Gli autori sostengono che, sebbene l'addestramento RL sia computazionalmente costoso e non necessariamente il metodo di deployment più pratico per la predizione immediata, esso funge da prova di concetto (proof of concept):

• Valida che gli LLM possiedano segnali predittivi significativi riguardo alle relazioni geniche oltre a quanto esplicitamente codificato nei grafi della conoscenza statici.
• Dimostra che il RL può sintonizzare un LLM per agire come una distribuzione posteriore sui cambiamenti del vicinato genico, imparando efficacementmente una policy per la selezione di set genici predittivi.
• Evidenzia come i miglioramenti nelle prestazioni predittive possano trasferirsi a task di ragionamento downstream, suggerendo una via verso modelli di "cellula virtuale" più robusti.

Gli autori mantengono un tono modesto, citando limitazioni come il costo computazionale del RL e la dipendenza dall'avere un set di addestramento per la specifica linea cellulare, il che impedisce predizioni "zero-shot" reali in nuove linee cellulari. Non pretendono di aver risolto interamente il problema della generalizzazione, ma hanno mostrato una via percorribile per raffinare modelli semplici e interpretabili utilizzando il ragionamento avanzato dei modelli linguistici.