기술 요약: 단순하면서도 효과적인 전사체 섭동 예측기를 찾기 위한 지식 그래프 및 추론 LLM

1. 문제 정의

보지 못한 유전자 녹아웃(knockout) 섭동의 전사체 효과를 예측하는 것은 가상 세포 모델(virtual cell models)에 있어 여전히 중요한 과제입니다. 최근의 접근 방식들은 생물학적 지식 그래프를 활용하여 훈련된 섭동 범위를 넘어선 외삽(extrapolation) 능력을 개선하려 노력하고 있지만, 복잡한 딥러닝 모델들은 종종 단순한 베이스라인(예: 훈련 샘플의 평균값)보다 우수한 성능을 내는 데 어려움을 겪습니다. 핵심적인 어려움은 각 섭동이 복잡한 생물학적 과정에 의해 제약되는 고유한 분포 $P(x|do(i))$를 유도하기 때문에, 새로운 개입 분포(interventional distributions)로 일반화하는 데 있습니다. 그래프 신경망(GNN)이나 대규모 사전 학습을 활용하는 기존 방법들은 유망한 결과를 보여주었으나, 생물학적 지식이 새로운 개입에 대한 일반화를 가능하게 하는 구체적인 메커니즘은 여전히 불분명합니다. 또한, 기저의 인과 구조(causal structure)를 직접 최적화하는 것은 가능한 유전자 그래프의 지수적인 탐색 공간 때문에 계산적으로 불가능합니다.

2. 방법론

2.1 이론적 토대: 희소 계층 모델 (A Sparse Hierarchical Model)

저자들은 유전자 녹아웃에 대한 세포 반응이 희소 계층적 인과 그래프(sparse hierarchical causal graph)를 따른다고 가정합니다. 이 모델에서 유전자는 단백질 복합체를 형성하고, 이는 다시 경로(pathway)를 형성하며, 궁극적으로 세포 소기관 및 세포 반응으로 이어집니다.

• 시사점: 동일한 저수준 복합체에 영향을 미치는 섭동은 유사한 발현 프로파일을 나타낼 것입니다.
• 예측: 타겟 섭동의 평균 발현은 이 인과 그래프에 의해 정의된 국소 이웃(local neighbourhood)에 대한 단순한 최근접 이웃 평균으로 근사될 수 있습니다.

2.2 베이스라인: 그래프 이웃 예측기 (The Graph Neighbour Predictor)

저자들은 먼저 기존의 생물학적 지식 그래프(예: StringDB, GO)를 사전 지식(prior)으로 활용하는 "그래프 이웃" 예측기를 평가합니다.

• 메커니즘: 타겟 유전자 $g$에 대해, 예측기는 지식 그래프에서의 이웃 $N(g)$를 식별합니다. 예측된 섭동 델타($\hat{\delta}_g$)는 다음과 같이 이웃들의 훈련 델타 평균으로 계산됩니다:
  $$\hat{y}_g = \bar{x}_g + \frac{1}{|N_{Tr}(g)|} \sum_{\tilde{g} \in N_{Tr}(g)} \bar{\delta}_{\tilde{g}}$$
  여기서 $\bar{x}_g$는 매칭된 대조군 발현이며, $N_{Tr}(g)$는 훈련 세트에 존재하는 이웃들을 나타냅니다.
• 결과: 이 단순한 $K$-최근접 이웃 방식은 복잡한 딥러닝 모델(예: GEARS, scLAMBDA)보다 우수한 성능을 보였으며, SOTA 모델인 TxPert의 성능에 근접했습니다. 특히 분포 외(OOD) 설정에서 그러했습니다.

2.3 정교화: 추론 LLM을 이용한 강화 학습 (Reinforcement Learning with Reasoning LLMs)

지식 그래프는 불완전하고 최적이 아닌 엣지(edge)를 포함할 수 있다는 점을 인식하여, 저자들은 LLM을 사용하여 이러한 이웃 관계를 정교화하는 방법을 제안합니다.

• 탐색 공간 축소: 모든 가능한 그래프의 다루기 힘든 공간($O(2^{\#genes^2})$)을 탐색하는 대신, LLM의 사전 신념(prior beliefs) 하에서 타당한 후보 엣지로 탐색을 제한합니다.
• 강화 학습 공식화:
  • 정책(Policy): LLM($\pi_\theta$)은 초기 지식 그래프 이웃($N_{KG}$)에 대해 행동(추가 $A$, 제거 $R$)을 선택하는 정책 역할을 합니다.
  • 행동 공간(Action Space): 정교화된 집합은 $S_{LLM} = (N_{KG} \setminus R) \cup A$입니다.
  • 보상 함수(Reward Function): 모델은 다음으로 구성된 보상 $R(s, a)$를 최대화하도록 그룹 상대 정책 최적화(GRPO)를 통해 훈련됩니다:
    1. 포맷 보상 ($R_{Fmt}$): 출력이 파싱 가능한지 확인합니다.
    2. 유효성 보상 ($R_{Valid}$): 훈련 세트에 있는 유전자를 추가하고 원래 이웃에 없는 유전자를 제거하도록 장려합니다.
    3. 성능 보상 ($R_\Delta$): 베이스라인 지식 그래프 예측기에 비해 Pearson Delta 지표의 개선을 독려합니다.
• 훈련: 모델은 Replogle et al. (2022)의 세포주 데이터를 사용하여 훈련되며, 미보지 섭동에 대한 평가를 보장하기 위해 엄격한 훈련/테스트 분할을 적용합니다.

3. 주요 기여

1. 단순 베이스라인의 입증: 본 논문은 생물학적 지식 그래프에 기반한 단순한 K-최근접 이웃 예측기가 OOD 섭동 예측에 있어 매우 경쟁력 있는 베이스라인이며, 대부분의 복잡한 딥러닝 방법보다 우수함을 입증했습니다.
2. 정교화 도구로서의 LLM: 저자들은 LLM이 강화 학습을 통해 최적화될 때, 지식 그래프의 이웃(누락된 관계를 추가하고 불필요한 것을 제거)을 효과적으로 편집할 수 있음을 보여주었습니다.
3. 최첨단 성능(SOTA): RL로 튜닝된 LLM은 단순한 최근접 이웃 아키텍처에서 시작했음에도 불구하고, Replogle et al. 세포주에 대해 현재의 SOTA 방법(특히 TxPert)과 대등한 성능을 달고 있습니다.
4. 다운스트림 일반화: 연구는 RL 훈련이 LLM의 다운스트림 작업 성능(구체적으로 PerturbQA 벤치마크의 차등 발현 예측)을 향상시킨다는 것을 보여주며, 이는 모델이 특정 예측 작업에 과적합되는 것이 아니라 일반화 가능한 추론 패턴을 학습하고 있음을 시사합니다.

4. 결과

• 벤치마크 성능: 네 가지 세포주(K562, RPE1, HEPG2, Jurkat)에 대해, RL로 훈련된 LLM ($LLM + KG + RL$)은 TxPert와 대등한 Pearson Delta 점수(예: K562에서 0.625 vs 0.626)를 기록했으며, 훈련되지 않은 LLM 및 원본 지식 그래프 베이스라인보다 유의미하게 높은 성능을 보였습니다.
• OOD 일반화: 이 방법은 미보지 섭동에 대해서도 일반화되며, 특정 유전자 녹아웃이 훈련 세트에 없더라도 높은 성능을 유지합니다.
• 전이 학습: 네 가지 세포주 중 세 가지만으로 훈련했을 때도, 모델은 홀드아웃(held-out) 세포주(K562)에서 베이스 모델보다 개선된 성능을 보였으며, 이는 추론 정책의 교차 세포주 일반화를 나타냅니다.
• 다운스트림 작업: PerturbQA 데이터셋(Wu et al.)에서, 사후 훈련된 모델은 특정 미세 조정 없이도 차등 발현 및 변화 방향을 예측하는 데 있어 기본 LLM보다 우수한 성능을 보였습니다.

5. 의의 및 주장

본 논문은 생물학적 지식 그래프가 효과적인 모델 사전 지식(prior) 역할을 한다는 점과, 이러한 사전 지식을 활용하는 가장 단순한 모델(최근접 이웃)이 놀라울 정도로 견고하다는 점을 주장합니다. 주요 의의는 강화 학습이 LLM을 일반화 가능한 도구로 정교화할 수 있음을 보여준 데 있습니다.

저자들은 RL 훈련이 계산적으로 비용이 많이 들며 즉각적인 예측을 위한 실용적인 배포 방법은 아닐 수 있다고 언급하며, 이를 하나의 **개념 증명(proof of concept)**으로 제시합니다:

• 이는 LLM이 명시적으로 기록된 정적 지식 그래프 이상의 유의미한 예측 신호(유전자 관계에 관한)를 보유하고 있음을 검증합니다.
• 이는 RL이 LLM을 유전자 이웃 변화에 대한 사후 확률(posterior)로 튜닝하여, 예측적인 유전자 집합을 선택하는 정책을 학습할 수 있음을 보여줍니다.
• 이는 예측 성능의 향상이 다운스트림 추론 작업으로 전이될 수 있음을 강조하며, 더 강력한 "가상 세포" 모델을 향한 경로를 제시합니다.

저자들은 RL의 계산 비용과 특정 세포주에 대한 훈련 세트 의존성(이로 인해 새로운 세포주에 대한 진정한 '제로샷' 예측이 불가능함)과 같은 한계를 언급하며 겸허한 태도를 유지합니다. 그들은 이 문제가 완전히 해결되었다고 주장하는 것이 아니라, 고급 언어 모델의 추론을 사용하여 단순하고 해석 가능한 모델을 정교화할 수 있는 실행 가능한 경로를 보여주었다고 주장합니다.