MapPFN: 맥락 기반 인과 교란 지도 학습

초록

MapPFN은 합성 인과 모델로 사전 학습된 Prior‑Data Fitted Network(PFN)를 이용해, 관측된 단일세포 데이터와 소수의 교란 실험을 컨텍스트로 삼아 새로운 생물학적 환경에서의 유전자 교란 효과를 한 번의 전방패스만으로 예측한다. 실제 실험 데이터에 대한 파인튜닝 없이도 차등 발현 유전자를 정확히 찾아내며, 기존 방법과 동등한 성능을 보인다.

상세 분석

본 논문은 단일세포 교란 예측을 “분포‑대‑분포” 매핑 문제로 정의하고, 이를 메타‑러닝 프레임워크인 Prior‑Data Fitted Network(PFN)와 결합한다. 저자들은 인공적인 구조적 인과 모델(SCM) 혹은 합성 유전자조절망(GRN)을 사전 분포로 설정하고, 이들로부터 관측 분포 Y_obs와 여러 교란 분포 Y_int_k를 생성한다. 핵심 아이디어는 이러한 합성 데이터셋을 대규모로 사전 학습시켜, 트랜스포머 기반의 Multimodal Diffusion Transformer(MMDiT)가 “컨텍스트‑조건부”로 후교란 분포 Y_int_q를 직접 출력하도록 하는 것이다.

학습 과정에서 모델은 (Y_obs, C)라는 입력 시퀀스와 질의 교란 t_q를 동시에 받으며, 이를 통해 인코더‑디코더 구조 없이도 한 번의 전방패스로 후처리 분포를 생성한다. 이때 사용된 손실은 Conditional Flow Matching(CFM) 기반의 확률 경로 손실로, 연속적인 시간 스케일 τ에 따라 샘플을 보간하고, Gaussian 확률 밀도 차이를 최소화한다.

특히 저자들은 “paired vs. unpaired” 사전 학습을 비교했는데, 동일한 노이즈 N_k를 공유하는 쌍(pair) 방식이 실제 데이터 전이 성능을 유의미하게 향상시킨다는 점을 발견했다. 이는 교란 전후 셀 간의 잠재적 대응 관계를 모델이 암묵적으로 학습하도록 돕는 효과로 해석된다.

실험에서는 두 가지 평가 환경을 제시한다. 첫 번째는 선형 SCM을 이용한 완전 통제된 합성 벤치마크로, few‑shot(몇 개의 교란 관측) 및 zero‑shot(교란 관측 없음) 상황에서 MapPFN이 기존 최첨단 방법들을 능가하거나 동등한 AUROC를 기록한다. 두 번째는 실제 단일세포 Perturb‑Seq 데이터(Frangieh et al., 2021)로, 사전 학습만으로도 차등 발현 유전자를 정확히 탐지했으며, 실제 데이터에 대해 파인튜닝된 모델들과 비교해도 성능 격차가 거의 없었다.

이러한 결과는 “synthetic‑to‑real” 전이가 가능함을 보여준다. 저자들은 사전 분포가 충분히 다양하고, 교란 종류(유전자 knockout)와 구조적 특성이 실제 생물학적 시스템을 포괄한다면, 복잡한 인과 관계를 명시적으로 모델링하지 않아도 충분히 강건한 교란 예측기가 구축될 수 있음을 입증한다. 또한, 기존 PFN이 개별 샘플의 스칼라 레이블만 예측하던 것과 달리, 본 연구는 전체 셀 집단의 다변량 분포를 직접 출력함으로써 단일세포 데이터의 비정형 특성을 효과적으로 다룬다.

전반적으로 MapPFN은 (1) 합성 인과 사전을 활용한 대규모 메타‑학습, (2) 멀티모달 디퓨전 트랜스포머를 통한 분포 예측, (3) 인-컨텍스트 학습을 통한 테스트‑타임 적응이라는 세 축을 결합해, 기존 교란 예측 모델이 갖던 “훈련‑테스트 도메인 격차”와 “그라디언트 기반 파인튜닝 필요성”을 극복한다는 점에서 의미가 크다.