유전자 변형에 대한 세포 반응 예측을 위한 검색 증강 생성

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)을 기존 언어 모델 응용 분야를 넘어 세포 생물학으로 확장한 새로운 프레임워크인 PT-RAG(Perturbation-aware Two-stage Retrieval-Augmented Generation)를 소개합니다. 사전 훈련된 대형 언어 모델(LLM)을 위한 텍스트 검색을 위해 설계된 표준 RAG 시스템과 달리, 변형 검색 분야에는 정립된 유사도 메트릭이 부재하며 관련 맥락을 구성하는 요소를 학습해야 하므로 미분 가능 검색이 필수적입니다. PT-RAG는 이를 두 단계 파이프라인으로 해결합니다: 첫째, 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.