단일세포 교란 예측을 위한 분포 흐름 매칭 모델

초록

scDFM은 조건부 흐름 매칭과 MMD 정규화를 결합한 생성 모델로, 제어 상태와 교란 정보를 입력받아 교란 후 세포 집단의 전체 분포를 직접 학습한다. 유전자 상호작용 그래프 기반 마스크 어텐션과 차별적 트랜스포머(PAD‑Transformer)를 백본으로 사용해 잡음과 희소성을 완화하고, 복합 교란 상황에서도 높은 예측 정확도를 달성한다.

상세 분석

본 논문은 단일세포 전사체 데이터의 특성—고차원, 희소, 잡음—을 고려하여 기존의 셀‑레벨 대응을 전제로 한 모델들의 한계를 지적한다. scDFM은 두 가지 핵심 아이디어로 이를 극복한다. 첫째, 조건부 흐름 매칭(Conditional Flow Matching, CFM)을 이용해 노이즈 분포에서 목표 교란 분포로의 연속적인 변환 경로를 학습한다. 여기서 시간‑조건부 속도장 vθ(x_t|t,c_x,c_p)를 ODE 형태로 정의하고, 선형 보간 경로 π_t를 기준으로 실제 속도와의 L2 차이를 최소화한다. 둘째, CFM만으로는 최종 분포 정합성을 보장하기 어렵다는 점을 인식하고, 다중 커널 MMD 정규화를 도입한다. MMD는 샘플 기반 거리 측정으로 고차원에서 계산 효율성이 높으며, 다양한 밴드폭을 혼합해 서로 다른 스케일의 분포 차이를 포착한다. λ 파라미터로 CFM 손실과 MMD 손실을 균형 있게 결합함으로써, 개별 셀의 미세한 궤적 학습과 전체 집단의 통계적 일치를 동시에 달성한다.

백본인 PAD‑Transformer는 세 가지 혁신을 포함한다. (1) 유전자‑유전자 상관 그래프를 마스크로 활용한 어텐션은 생물학적 연관성을 보존하면서 불필요한 독립 처리로 인한 과적합을 방지한다. (2) 차별적 어텐션 메커니즘은 제어와 교란 조건을 별도 토큰으로 인코딩하고, 두 토큰 간 상호작용을 강조해 교란에 특이적인 변화를 효과적으로 학습한다. (3) 시간 임베딩을 삽입해 연속적인 흐름 매칭 과정에서 시점 정보를 전달함으로써, 속도장 예측의 시간 의존성을 모델링한다.

실험에서는 Norman 유전자 교란 데이터와 Combosiplex 약물 교란 데이터를 사용해 두 가지 평가 설정(추가 학습, 보류된 조합)에서 기존 최첨단 모델(CPA, GEARS, CellFlow 등)을 능가한다. 특히 복합 교란 상황에서 평균 제곱 오차(MSE)를 19.6% 감소시키며, MMD 기반 정합이 전체 분포 재현에 크게 기여함을 확인한다. 전반적으로 scDFM은 셀‑레벨 대응이 불가능한 단일세포 교란 예측 문제에 대해 분포‑레벨 생성 모델링과 그래프‑기반 트랜스포머 설계가 어떻게 시너지를 내는지를 보여준다.