Technical Summary: Knowledge Graphs and Reasoning LLMs for Finding Simple Yet Effective Transcriptomic Perturbation Predictors

1. Problem Statement

Predicting the transcriptomic effects of unseen gene knockout perturbations remains a significant challenge for virtual cell models. While recent approaches leverage biological knowledge graphs to improve extrapolation beyond training perturbations, complex deep learning models often struggle to outperform simple baselines (e.g., the mean of training samples). The core difficulty lies in generalizing to new interventional distributions, $P(x|do(i))$, where each perturbation induces a unique distribution constrained by complex biological processes. Existing methods utilizing Graph Neural Networks (GNNs) or large-scale pretraining have shown promise, but the specific mechanisms by which biological knowledge enables generalization to novel interventions remain unclear. Furthermore, optimizing the underlying causal structure directly is computationally intractable due to the exponential search space of possible gene graphs.

2. Methodology

2.1 Theoretical Foundation: A Sparse Hierarchical Model

The authors posit that cellular responses to gene knockouts follow a sparse hierarchical causal graph. In this model, genes form protein complexes, which form pathways, ultimately leading to organelles and cell responses. A knockout interrupts the production of an initial protein, propagating failure up the hierarchy.

• Implication: Perturbations affecting the same low-level complex should yield similar expression profiles.
• Prediction: The average expression of a target perturbation can be approximated by a simple nearest-neighbour average over the local neighbourhood defined by this causal graph.

2.2 Baseline: The Graph Neighbour Predictor

The authors first evaluate a "Graph Neighbour" predictor, which leverages existing biological knowledge graphs (e.g., StringDB, GO) as priors.

• Mechanism: For a target gene $g$, the predictor identifies its neighbours $N(g)$ in the knowledge graph. The predicted perturbation delta ($\hat{\delta}_g$) is calculated as the mean of the training deltas of these neighbours:
  $$\hat{y}_g = \bar{x}_g + \frac{1}{|N_{Tr}(g)|} \sum_{\tilde{g} \in N_{Tr}(g)} \bar{\delta}_{\tilde{g}}$$
  where $\bar{x}_g$ is the matched control expression and $N_{Tr}(g)$ represents neighbours present in the training set.
• Finding: This simple $K$-nearest neighbour approach outperforms complex deep learning models (e.g., GEARS, scLAMBDA) and approaches the performance of state-of-the-art (SOTA) models like TxPert, particularly in out-of-distribution (OOD) settings.

2.3 Refinement: Reinforcement Learning with Reasoning LLMs

Recognizing that knowledge graphs are incomplete and contain suboptimal edges, the authors propose using Large Language Models (LLMs) to refine these neighbourhoods.

• Search Space Reduction: Instead of searching the intractable space of all possible graphs ($O(2^{\#genes^2})$), the search is constrained to candidate edges plausible under the LLM's prior beliefs.
• RL Formulation:
  • Policy: An LLM ($\pi_\theta$) acts as a policy to select actions (Add $A$, Remove $R$) to an initial set of knowledge graph neighbours ($N_{KG}$).
  • Action Space: The refined set is $S_{LLM} = (N_{KG} \setminus R) \cup A$.
  • Reward Function: The model is trained via Group Relative Policy Optimization (GRPO) to maximize a reward $R(s, a)$ composed of:
    1. Formatting Reward ($R_{Fmt}$): Ensures output is parsable.
    2. Validity Reward ($R_{Valid}$): Encourages adding genes from the training set and removing genes not in the original neighbourhood.
    3. Performance Reward ($R_\Delta$): Incentivizes improvements in the Pearson Delta metric compared to the baseline knowledge graph predictor.
• Training: The LLM is trained on perturbation data from Replogle et al. (2022) cell lines, with a strict train/test split to ensure evaluation on unseen perturbations.

3. Key Contributions

1. Demonstration of Simple Baselines: The paper establishes that a simple K-nearest neighbour predictor based on biological knowledge graphs is a highly competitive baseline for OOD perturbation prediction, outperforming most complex deep learning methods.
2. LLM as a Refinement Prior: The authors demonstrate that LLMs, when optimized via Reinforcement Learning, can effectively edit knowledge graph neighbourhoods (adding missing relationships and pruning irrelevant ones) to improve predictive performance.
3. State-of-the-Art Performance: The RL-tuned LLM achieves performance equivalent to current SOTA methods (specifically TxPert) on Replogle et al. cell lines, despite starting from a simple nearest-neighbour architecture.
4. Downstream Generalization: The study shows that RL training improves the LLM's performance on downstream tasks (specifically differential expression prediction on the PerturbQA benchmark), suggesting the model learns generalizable reasoning patterns rather than overfitting to the specific prediction task.

4. Results

• Benchmark Performance: On four cell lines (K562, RPE1, HEPG2, Jurkat) from Replogle et al. (2022), the RL-trained LLM ($LLM + KG + RL$) achieved Pearson Delta scores comparable to TxPert (e.g., 0.625 vs 0.626 on K562) and significantly outperformed the untrained LLM and the raw knowledge graph baseline.
• OOD Generalization: The method generalizes to unseen perturbations, maintaining high performance even when the specific gene knockouts were not present in the training set.
• Transfer Learning: When trained on only three of the four cell lines, the model still improved performance on the held-out cell line (K562) compared to the base model, indicating cross-cell-line generalization of the reasoning policy.
• Downstream Task: On the PerturbQA dataset (Wu et al.), the post-trained model outperformed the base LLM in predicting differential expression and direction of change, even without specific fine-tuning on this task.

5. Significance and Claims

The paper claims that biological knowledge graphs serve as effective model priors, and that the simplest model leveraging these priors (nearest neighbour) is surprisingly robust. The primary significance lies in showing that Reinforcement Learning can refine LLMs into generalizable tools for complex biological prediction.

The authors argue that while RL training is computationally expensive and not necessarily the most practical deployment method for immediate prediction, it serves as a proof of concept:

• It validates that LLMs possess meaningful predictive signals regarding gene relationships beyond what is explicitly encoded in static knowledge graphs.
• It demonstrates that RL can tune an LLM to act as a posterior over gene neighbourhood changes, effectively learning a policy for selecting predictive gene sets.
• It highlights that improvements in predictive performance can transfer to downstream reasoning tasks, suggesting a path toward more robust "virtual cell" models.

The authors remain modest, noting limitations such as the computational cost of RL and the dependency on having a training set for the specific cell line, which prevents true "zero-shot" predictions in novel cell lines. They do not claim to have solved the generalization problem entirely but have shown a viable path to refining simple, interpretable models using advanced language model reasoning.