Technische Samenvatting: Knowledge Graphs en Reasoning LLM's voor het vinden van eenvoudige doch effectieve voorspellers voor transcriptomische perturbaties

1. Probleemstelling

Het voorspellen van de transcriptomische effecten van ongeziene gen-knockout-perturbaties blijft een aanzienlijke uitdaging voor virtuele celmodellen. Hoewel recente benaderingen biologische kennisgrafen (knowledge graphs) inzetten om extrapolatie buiten de trainingsperturbaties te verbeteren, worstelen complexe deep learning-modellen vaak met het presteren ten opzichte van eenvoudige baselines (bijv. het gemiddelde van de trainingsmonsters). De kern van de moeilijkheid ligt in het generaliseren naar nieuwe interventiedistributies, $P(x|do(i))$, waarbij elke perturbatie een unieke distributie induceert die wordt beperkt door complexe biologische processen. Bestaande methoden die gebruikmaken van Graph Neural Networks (GNN's) of grootschalige pretraining tonen veelbelovende resultaten, maar de specifieke mechanismen waarmee biologische kennis generalisatie naar nieuwe interventies mogelijk maakt, blijven onduidelijk. Bovendien is het direct optimaliseren van de onderliggende causale structuur computationeel onhaalbaar vanwege de exponentiële zoekruimte van mogelijke gen-grafen.

2. Methodologie

2.1 Theoretische Fundering: Een Spaars Hiërarchisch Model

De auteurs stellen dat cellulaire reacties op gen-knockouts een spaars hiërarchisch causaal grafiek volgen. In dit model vormen genen eiwitcomplexen, die weer paden (pathways) vormen, wat uiteindelijk leidt tot organellen en cellulaire reacties. Een knockout onderbreept de productie van een initieel eiwit, waardoor de uitval omhoog in de hiërarchie voortplant.

• Implicatie: Perturbaties die dezelfde laag-niveau complex beïnvloeden, zouden vergelijkbare expressieprofielen moeten opleveren.
• Voorspelling: De gemiddelde expressie van een doel-perturbatie kan worden benaderd door een eenvoudig nearest-neighbour gemiddelde over de lokale buurt gedefinieerd door deze causale grafiek.

2.2 Baseline: De Graph Neighbour Predictor

De auteurs evalueren eerst een "Graph Neighbour" voorspeller, die bestaande biologische kennisgrafen (bijv. StringDB, GO) als priors gebruikt.

• Mechanisme: Voor een doelgen $g$ identificeert de voorspeller zijn buren $N(g)$ in de kennisgrafiek. De voorspelde perturbatie delta ($\hat{\delta}_g$) wordt berekend als het gemiddelde van de trainings-deltas van deze buren:
  $$\hat{y}_g = \bar{x}_g + \frac{1}{|N_{Tr}(g)|} \sum_{\tilde{g} \in N_{Tr}(g)} \bar{\delta}_{\tilde{g}}$$
  waarbij $\bar{x}_g$ de gematchte controle-expressie is en $N_{Tr}(g)$ de buren vertegenwoordigt die aanwezig zijn in de trainingsset.
• Bevinding: Deze eenvoudige $K$-nearest neighbour aanpak presteert beter dan complexe deep learning-modellen (bijv. GEARS, scLAMBDA) en nadert de prestaties van state-of-the-art (SOTA) modellen zoals TxPert, met name in out-of-distribution (OOD) settings.

2.3 Verfijning: Reinforcement Learning met Reasoning LLM's

In het besef dat kennisgrafen incompleet zijn en randen (edges) bevatten die suboptimaal zijn, stellen de auteurs voor om Large Language Models (LLM's) te gebruiken om deze buurten te verfijnen.

• Reductie van de Zoekruimte: In plaats van te zoeken in de onhandelbare ruimte van alle mogelijke grafen ($O(2^{\#genes^2})$), wordt de zoektocht beperkt tot kandidaat-randen die plausibel zijn onder de prior-overtuigingen van de LLM.
• RL-formulering:
  • Policy: Een LLM ($\pi_\theta$) fungeert als een policy om acties te selecteren (Toevoegen $A$, Verwijderen $R$) op een initiële set van kennisgrafiek-buren ($N_{KG}$).
  • Actieruimte: De verfijnde set is $S_{LLM} = (N_{KG} \setminus R) \cup A$.
  • Beloningsfunctie: Het model wordt getraind via Group Relative Policy Optimization (GRPO) om een beloning $R(s, a)$ te maximaliseren die bestaat uit:
    1. Formateringsbeloning ($R_{Fmt}$): Zorgt dat de output parseerbaar is.
    2. Geldigheidsbeloning ($R_{Valid}$): Moedigt het toevoegen van genen uit de trainingsset aan en het verwijderen van genen die niet in de oorspronkelijke buurt aanwezig waren.
    3. Prestatiebeloning ($R_\Delta$): Stimuleert verbeteringen in de Pearson Delta metriek vergeleken met de baseline kennisgrafiek-voorspeller.
• Training: De LLM wordt getraind op perturbatiegegevens van Replogle et al. (2022) cellijnen, met een strikte train/test split om evaluatie op ongeziene perturbaties te garanderen.

3. Belangrijkste Bijdragen

1. Demonstratie van Eenvoudige Baselines: Het artikel stelt vast dat een eenvoudige K-nearest neighbour voorspeller gebaseerd op biologische kennisgrafen een zeer concurrerende baseline is voor OOD-perturbatievoorspelling, die de meeste complexe deep learning-methoden overtreft.
2. LLM als Verfijningsprior: De auteurs demonstreren dat LLM's, wanneer ze geoptimaliseerd worden via Reinforcement Learning, effectief kennisgrafiek-buurten kunnen bewerken (het toevoegen van ontbrekende relaties en het wegfilteren van irrelevante relaties) om de voorspellende prestaties te verbeteren.
3. State-of-the-Art Prestaties: De via RL getunede LLM bereikt prestaties die gelijkwaardig zijn aan huidige SOTA-methoden (specifiek TxPert) op de Replogle et al. cellijnen, ondanks dat deze begint vanuit een eenvoudige nearest-neighbour architectuur.
4. Downstream Generalisatie: De studie toont aan dat RL-training de prestaties van de LLM verbetert op downstream taken (specifiek differentieel expressie voorspelling op de PerturbQA benchmark), wat suggereert dat het model generaliseerbare redeneerpatronen leert in plaats van te overfitten op de specifieke voorspellingstaak.

4. Resultaten

• Benchmark Prestaties: Op vier cellijnen (K562, RPE1, HEPG2, Jurkat) uit Replogle et al. (2022) behaalde de via RL getrainde LLM ($LLM + KG + RL$) Pearson Delta scores die vergelijkbaar zijn met TxPert (bijv. 0.625 vs 0.626 op K562) en presteerde significant beter dan de ongetrainde LLM en de ruwe kennisgrafiek-baseline.
• OOD Generalisatie: De methode generaliseert naar ongeziene perturbaties en behoudt een hoge prestatie, zelfs wanneer de specifieke gen-knockouts niet in de trainingsset aanwezig waren.
• Transfer Learning: Wanneer de methode getraind werd op slechts drie van de vier cellijnen, verbeterde deze nog steeds de prestaties op de uitgesloten cellijn (K562) vergeleken met het basismodel, wat wijst op cross-cell-line generalisatie van de redeneerpolicy.
• Downstream Taak: Op de PerturbQA dataset (Wu et al.) presteerde het na de training verbeterde model beter dan de basis-LLM in het voorspellen van differentiële expressie en de richting van verandering, zelfs zonder specifieke fine-tuning op deze taak.

5. Betekenis en Claims

Het artikel claimt dat biologische kennisgrafen effectieve modelpriors dienen, en dat het eenvoudigste model dat gebruikmaakt van deze priors (nearest neighbour) verrassend robuust is. De primaire betekenis ligt in het aantonen dat Reinforcement Learning LLM's kan verfijnen tot generaliseerbare instrumenten voor complexe biologische voorspellingen.

De auteurs stellen dat hoewel RL-training computationeel duur is en niet noodzakelijkerwijs de meest praktische implementatiemethode is voor directe voorspelling, het dient als een proof of concept:

• Het valideert dat LLM's beschikken over betekenisvolle voorspellende signalen met betrekking tot gen-relaties die verder gaan dan wat expliciet is gecodeerd in statische kennisgrafen.
• Het demonstreert dat RL een LLM kan tunen om te fungeren als een posterior over veranderingen in gen-buurten, waardoor effectief een policy wordt geleerd voor het selecteren van voorspellende gen-sets.
• Het benadrukt dat verbeteringen in voorspellende prestaties kunnen overgaan naar downstream redeneertaken, wat wijst op een pad naar meer robuuste "virtuele cel"-modellen.

De auteurs blijven bescheiden en merken beperkingen op, zoals de computationele kosten van RL en de afhankelijkheid van een trainingsset voor de specifieke cellijn, wat echte "zero-shot" voorspellingen in nieuwe cellijnen verhindert. Ze beweren niet het generalisatieprobleem volledig te hebben opgelost, maar hebben een levensvatbaar pad aangetoont om eenvoudige, interpreteerbare modellen te verfijnen met behulp van geavanceerde taalmodel-redenering.