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

2026-06-07

AI Summary

Predicting the transcriptional effects of out-of-distribution gene knockout perturbations remains a significant challenge. This work proposes a novel approach that integrates biological knowledge graphs with reasoning-capable large language models. The method first constructs a k-nearest neighbors (k-NN) base predictor grounded in a knowledge graph, which substantially outperforms existing methods in out-of-distribution settings. Subsequently, reinforcement learning is employed to optimize the language model’s strategy for selecting relevant neighbors, thereby indirectly enhancing downstream differential expression prediction performance. This framework provides the first empirical validation of knowledge graph–guided k-NN for perturbation response prediction and achieves state-of-the-art results through a reinforcement learning–driven inference mechanism, markedly improving generalization to tasks not directly seen during training.

Abstract

Predicting the effect of an unseen gene knockout perturbation on transcriptomic gene expression remains a highly challenging problem for virtual cell models. Recent progress has been made by leveraging biological knowledge graphs to provide a notion of similar perturbation, allowing for improved extrapolation beyond the set of training perturbations. In this work, we demonstrate that the simplest model to leverage these assumptions - a K-nearest neighbour from the knowledge graph - achieves highly competitive performance on this task, and that this can be improved further using LLMs optimised via reinforcement learning (RL) for predictive performance. Specifically, we find that the K-nearest neighbour approach beats almost all methods on out-of-distribution perturbation prediction, and when a reasoning LLM is trained via RL to make changes to the neighbourhood, it obtains equivalent performance to current state of the art methods on the cell lines from Replogle et al. (2022). We also demonstrate that the RL training improves the LLM's performance on the downstream task of differential expression prediction, despite not being trained on this directly. Overall, these findings demonstrate the efficacy of knowledge graphs as model priors, and show early signs that RL can refine LLMs into generalizable tools for predicting complex biological responses.

Problem

• transcriptomic perturbation prediction
• gene knockout
• out-of-distribution prediction
• virtual cell models
• biological knowledge graphs

Innovation

• Knowledge Graph
• Reinforcement Learning
• Large Language Models
• Transcriptomic Perturbation Prediction
• K-nearest Neighbour