PRiMeFlow: capturing complex expression heterogeneity in perturbation response modelling

Motivation

• Known: 기존 perturbation response 모델들(CPA, SAMS-VAE 등)은 PCA나 single-cell foundation model에서 유래한 고정된 pretrained latent embedding을 활용해 저차원 공간에서 예측을 수행해왔으며, PerturBench와 같은 벤치마크를 통해 이러한 모델들이 체계적으로 평가되어 왔다.
• Gap: PerturBench의 실증 결과에 따르면 CPA의 adversarial loss나 SAMS-VAE의 sparsity constraint 같은 강한 inductive bias를 제거한 단순한 모델이 오히려 더 나은 성능을 보이는데, 이는 latent embedding 기반 다단계 파이프라인이 생물학적 신호 손실과 오류 누적을 야기할 수 있음을 시사하며 gene expression space에서 직접 작동하는 end-to-end 접근법의 필요성을 제기한다.
• Why: in-silico perturbation 효과 예측은 세포 행동의 driver를 대규모로 식별하고 신약 개발을 가속화할 수 있는 핵심 기술이며, 단일세포 발현 분포를 정확히 근사하는 능력은 이질적인 세포 반응을 반영한 신뢰성 있는 counterfactual screen을 가능케 한다.
• Approach: PRiMeFlow는 latent embedding 없이 gene expression space에서 직접 velocity field를 U-Net으로 parameterize하는 conditional flow matching 기반 생성모델을 제안하고, 다중 데이터셋 atlas에 대한 pretraining-finetuning 전략으로 확장성을 검증한다.

Achievement

Figure 2

1. 분포 적합 성능 입증: PerturBench 내 다양한 baseline 대비 LogFC cosine similarity, RMSE, MMD, DEG recall 등 분포 기반 지표에서 유의미한 개선을 달성하여 실제 단일세포 이질성을 재현함을 보였다.
2. 핵심 설계 선택 검증: ablation study를 통해 (i) latent space가 아닌 gene expression space에서의 velocity field parameterization, (ii) MLP 대신 U-Net 아키텍처, (iii) optimal transport coupling 대신 independent coupling 사용이 성능에 크게 기여함을 확인했다.
3. 대규모 확장 및 실전 성능: 여러 perturbation 데이터셋을 아우르는 atlas로 모델을 스케일링하고 pretraining-finetuning 전략을 적용하여, ARC Virtual Cell Challenge의 H1 human embryonic stem cell 벤치마크에서 뛰어난 성능을 입증했다.

How

Figure 1

• Flow matching 프레임워크를 기반으로, source distribution p0(예: Gaussian)에서 empirical single-cell gene expression 분포 pdata로의 매핑을 conditional flow matching loss(LCFM)로 학습
• velocity field vθt(x)를 U-Net 아키텍처로 parameterize하여 gene expression space에서 직접 학습 (latent embedding 우회)
• PerturBench 플랫폼의 covariate transfer(Srivatsan20, Jiang24) 및 combo prediction(Norman19) task에서 벤치마킹
• pseudobulk 지표(LogFC cosine similarity, RMSE)와 distribution 기반 지표(MMD in GEX/PCA space, DEG recall) 및 rank 기반 변형 지표로 평가
• 대규모 perturbation data atlas에 대한 pretraining 후 ARC Virtual Cell Challenge의 H1 hESC 데이터셋에 finetuning하는 전략 적용
• ablation study로 latent space vs. gene expression space, MLP vs. U-Net, optimal transport coupling vs. independent coupling 비교

Originality

• perturbation response modeling에 flow matching을 end-to-end로 최초 적용한 사례로, 기존의 PCA/foundation model 기반 고정 latent embedding 의존성을 배제
• gene expression space에서 직접 velocity field를 학습하는 접근은 다단계 파이프라인의 오류 누적 문제를 회피하는 참신한 설계
• U-Net을 velocity field parameterization에 도입하고 independent coupling의 우수성을 실증적으로 검증한 점이 기존 CNF/flow matching 응용과 차별화됨
• 여러 perturbation 데이터셋을 통합한 atlas 기반 pretraining-finetuning 전략을 대규모 실전 챌린지(ARC VCC)에 적용한 확장성 있는 시도

Limitation & Further Study

• 본문 발췌만으로는 U-Net 및 independent coupling 선택의 이론적 근거와 계산 비용(추론 시 ODE 적분에 따른 효율성)에 대한 상세 분석이 부족해 보임
• ARC VCC의 H1 hESC 단일 세포 유형에 대한 결과만 강조되어, 다양한 세포 유형·조직에 대한 일반화 가능성 검증이 제한적일 수 있음
• 고차원 gene expression space에서 직접 작동함에 따른 scalability와 메모리/연산 비용에 대한 정량적 논의가 abstract 및 발췌본에서는 확인되지 않음
• 후속 연구로는 다양한 세포 타입 및 perturbation 조합(combo)에 대한 일반화 성능, 그리고 다른 foundation model embedding과의 hybrid 접근 비교가 필요해 보임

Evaluation

총평: flow matching을 gene expression space에 직접 적용해 latent embedding의 한계를 극복하고 PerturBench와 ARC VCC에서 실증적으로 우수한 성능을 보인 견실한 연구이나, 확장성과 일반화에 대한 추가 검증이 필요하다.