遺伝子撹乱に対する細胞応答予測のための検索拡張生成

Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

March 7, 2026

著者: Andrea Giuseppe Di Francesco, Andrea Rubbi, Pietro Liò

cs.AI

要旨

細胞が遺伝的擾乱にどのように応答するかを予測することは、遺伝子機能、疾患メカニズム、治療法開発を理解する上で重要である。近年の深層学習アプローチは単一細胞の擾乱応答のモデル化において有望な成果を示しているが、生成時の文脈情報が限られているため、細胞タイプや擾乱コンテキストを超えた汎化が困難である。本研究では、Retrieval-Augmented Generation（RAG）を従来の言語モデル応用から細胞生物学に拡張する新規フレームワーク、PT-RAG（Perturbation-aware Two-stage Retrieval-Augmented Generation）を提案する。事前学習済み大規模言語モデルを用いたテキスト検索向けに設計された標準的なRAGシステムとは異なり、擾乱検索には確立された類似度指標がなく、関連する文脈の構成要素を学習する必要があるため、微分可能な検索が不可欠である。PT-RAGは、2段階のパイプラインによってこの課題に対処する。まず、GenePT埋め込みを用いて候補擾乱Kを検索し、次に、細胞状態と入力擾乱の両方を条件としたGumbel-Softmax離散サンプリングを通じて適応的に選択を絞り込む。この細胞タイプを考慮した微分可能検索により、検索目的関数と生成を共同でエンドツーエンド最適化することが可能となる。Replogle-Nadig単一遺伝子擾乱データセットにおいて、PT-RAGは同一実験条件下でSTATEおよびバニラRAGを上回り、特に分布類似性指標（W_1, W_2）で顕著な改善を示した。特筆すべきは、バニラRAGの顕著な失敗自体が重要な知見である点である。これは、この領域では微分可能かつ細胞タイプを考慮した検索が本質的に重要であり、単純な検索が積極的に性能を損なうことを実証している。我々の結果は、遺伝子擾乱に対する細胞応答のモデル化において、検索拡張生成が有望なパラダイムであることを立証する。実験を再現するコードはhttps://github.com/difra100/PT-RAG_ICLRで公開されている。

English

Predicting how cells respond to genetic perturbations is fundamental to understanding gene function, disease mechanisms, and therapeutic development. While recent deep learning approaches have shown promise in modeling single-cell perturbation responses, they struggle to generalize across cell types and perturbation contexts due to limited contextual information during generation. We introduce PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), a novel framework that extends Retrieval-Augmented Generation beyond traditional language-model applications to cellular biology. Unlike standard RAG systems designed for text retrieval with pre-trained LLMs, perturbation retrieval lacks established similarity metrics and requires learning what constitutes relevant context, making differentiable retrieval essential. PT-RAG addresses this through a two-stage pipeline: first, retrieving candidate perturbations K using GenePT embeddings, then adaptively refining the selection through Gumbel-Softmax discrete sampling conditioned on both the cell state and the input perturbation. This cell-type-aware differentiable retrieval enables end-to-end optimization of the retrieval objective jointly with generation. On the Replogle-Nadig single-gene perturbation dataset, we demonstrate that PT-RAG outperforms both STATE and vanilla RAG under identical experimental conditions, with the strongest gains in distributional similarity metrics (W_1, W_2). Notably, vanilla RAG's dramatic failure is itself a key finding: it demonstrates that differentiable, cell-type-aware retrieval is essential in this domain, and that naive retrieval can actively harm performance. Our results establish retrieval-augmented generation as a promising paradigm for modelling cellular responses to gene perturbation. The code to reproduce our experiments is available at https://github.com/difra100/PT-RAG_ICLR.