PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling

Motivation

• Known: 기존 단일 세포 perturbation 예측 모델들은 관측된 cellular context와 perturbation type이 주어지면 control/perturbed cell distribution이 하나로 고정된다고 가정하고, unpaired된 control과 perturbed 세포 집단 간의 mapping을 학습해왔다(예: linear models, scGPT, CellOT, STATE, CellFlow, Squidiff).
• Gap: 실제로는 microenvironmental fluctuation, batch effect 등 관측되지 않는 latent factor로 인해 동일한 관측 조건에서도 반응 분포가 체계적으로 달라지는 manifold를 형성하는데, 기존 방법들은 이를 무시하고 하나의 정적인 반응 분포로 marginalize하여 unseen perturbation이나 cellular context에 대한 일반화가 제한된다.
• Why: Virtual Cell 구축은 시스템 생물학의 오랜 목표이며, single-cell sequencing이 destructive하여 동일 세포의 전후 관측이 불가능한 근본적 제약 속에서, 관측되지 않는 latent factor에 의한 분포 변동성까지 포착하는 것은 unseen perturbation 및 cellular context에 대한 일반화 성능을 크게 향상시킬 수 있어 중요하다.
• Approach: population distribution 자체를 하나의 random variable로 취급하여, cell distribution을 RKHS 상의 kernel mean embedding으로 임베딩하고 이 함수 값(function-valued) 표현 위에서 직접 작동하는 diffusion 생성 과정을 정의함으로써 개별 세포가 아닌 분포들의 분포(distribution over distributions)를 모델링한다.

Achievement

Figure 3. Perturbation prediction results across methods, metrics, and datasets. Each axis

1. PerturbDiff 프레임워크: 세포 분포를 개별 데이터 포인트가 아닌 RKHS 상의 함수 값 표현으로 임베딩하고, 그 위에서 diffusion 과정을 정의하여 관측되지 않은 latent factor에 의한 population-level 반응 분포의 변동성을 포착하는 새로운 생성 프레임워크를 제안했다.
2. MMD 기반 학습 목적함수 도출: 함수 공간에서의 diffusion 구성이 자연스럽게 Maximum Mean Discrepancy(MMD) 목적함수를 유도하도록 하여, 표준 diffusion의 pointwise MSE 손실보다 sparse하고 고차원적인 single-cell 특징에 더 적합한 학습 신호를 제공한다.
3. 대규모 벤치마크에서 SOTA 성능 달성: PBMC, Tahoe100M, Replogle 등 signaling, 약물, 유전적 perturbation을 포괄하는 대규모 벤치마크의 14개 다양한 지표에서 state-of-the-art 성능과 균형 잡힌 결과를 달성했으며, 특히 unseen perturbation에 대한 일반화 성능이 크게 향상되었다.
4. 새로운 사전학습 전략: perturbed 데이터의 희소성(제한된 세포 유형 등) 문제를 해결하기 위해, perturbation 데이터와 대규모 unperturbed RNA-seq 데이터를 결합하여 marginal single-cell distribution을 학습하는 사전학습 전략을 도입, non-trivial한 zero-shot 예측 능력과 데이터가 부족할 때 scratch 학습 대비 일관된 성능 향상을 보였다.

How

Figure 2. Overview of the PerturbDiff framework. (a) Distribution-valued random variables Dc,τ and Dc,τ in cell space ar

• 유전자 집합 G에 대해 세포를 gene expression vector로 표현하고, control 세포군 Xctrl과 perturbed 세포군 Xpert는 unpaired 상태로 주어지며, 각각 cellular context label c와 (c, τ)로 조건화된 분포 Pc, Pc,τ를 다룬다.
• 세포 분포를 kernel mean embedding을 통해 RKHS 상의 단일 함수 값 포인트로 표현하여, 분포 자체를 diffusion의 대상 데이터로 취급한다.
• 표준 diffusion의 forward Markov chain q(xt|xt-1)과 reverse process를 함수 값 표현 공간으로 확장하고, x0-prediction 기반의 denoising objective를 사용한다.
• classifier-free guidance(CFG)를 도입하여 조건부/비조건부 입력을 혼합 학습하고, 샘플링 시 조건 방향으로 예측을 조정(linear extrapolation)하여 조건 충실도를 강화한다.
• RKHS 임베딩 구성으로부터 자연스럽게 유도되는 MMD 목적함수를 diffusion 학습에 사용하여 source와 target 세포군 분포를 정렬한다.
• 대규모 unperturbed RNA-seq 데이터를 활용한 사전학습 후 perturbation 데이터로 fine-tuning하는 전략을 적용한다.

Originality

• 기존 연구들이 개별 세포 또는 고정된 하나의 반응 분포를 학습 대상으로 삼은 것과 달리, population distribution 자체를 random variable로 보고 "분포들의 분포(distribution over distributions)"를 모델링하는 관점 전환을 제시했다.
• 세포 분포를 RKHS의 kernel mean embedding으로 임베딩하여 함수 값 표현 공간 위에서 diffusion 과정을 직접 정의한 것은 기존 CellOT(optimal transport), CellFlow(flow matching), Squidiff/Unlasting(개별 세포 또는 bridge 기반 diffusion) 등과 차별화되는 새로운 생성 모델링 패러다임이다.
• diffusion 구성으로부터 MMD 손실을 이론적으로 자연스럽게 유도해 pointwise MSE 대신 분포 정렬에 적합한 학습 목적을 제공한 점이 독창적이다.
• perturbation 데이터와 대규모 unperturbed RNA-seq 데이터를 결합한 사전학습 전략으로 데이터 희소성 문제를 완화하는 접근도 새롭다.

Limitation & Further Study

• 발췌된 내용만으로는 latent factor의 구체적 정의나 그 수를 어떻게 결정하는지, RKHS 커널 선택이 성능에 미치는 민감도에 대한 상세한 분석이 부족해 보인다.
• 함수 값(function-valued) diffusion의 계산 복잡도나 대규모 유전자 수(수만 차원)에서의 확장성에 대한 논의가 본문 발췌에서 명확히 드러나지 않는다.
• 사전학습에 사용한 unperturbed RNA-seq 데이터의 배치 효과나 도메인 차이가 fine-tuning 성능에 미치는 영향에 대한 심층 분석이 필요해 보인다.
• 후속 연구로는 다양한 organism이나 tissue로의 일반화, 그리고 latent factor에 대한 명시적 해석 가능성을 높이는 방향이 고려될 수 있다.

Evaluation

총평: 단일 세포가 아닌 분포 수준에서 diffusion을 정의하는 참신한 관점 전환과 이를 뒷받침하는 MMD 기반 이론적 근거를 갖추었으며, 대규모 벤치마크에서 SOTA 성능과 unseen perturbation에 대한 강한 일반화를 보여준 견실한 연구이다.