동적 비균형 최적수송을 위한 원스텝 흐름 매칭 WFR‑MFM

초록

본 논문은 단일 세포 RNA‑seq 데이터에서 관측된 스냅샷을 이용해, 질량 변화를 포함한 동적 시스템을 빠르게 복원하는 방법을 제안한다. 평균 흐름(mean‑flow) 개념을 도입해 시간 구간 전체를 하나의 평균 속도와 평균 성장률 필드로 요약하고, 이를 Wasserstein‑Fisher‑Rao (WFR) 기하학에 맞춰 학습한다. 결과적으로 시뮬레이션 없이 한 번의 연산으로 최종 상태를 예측할 수 있어 기존 방법 대비 수십 배에서 수백 배 빠른 추론 속도를 보이며, 정확도도 유지한다.

상세 분석

WFR‑MFM은 기존의 흐름 매칭(Flow Matching, FM) 프레임워크를 비균형(unbalanced) 상황에 확장한 점이 핵심이다. 전통적인 FM은 질량 보존을 전제로 순간 속도장 uₜ(x)를 직접 회귀하지만, 세포 증식·사멸과 같이 질량이 변하는 경우에는 충분하지 않다. 이를 해결하기 위해 저자들은 평균 흐름 변수 v(x,t,T)와 h(x,t,T)를 정의한다. v는 구간 (t,T) 동안의 평균 이동 속도, h는 평균 질량 성장률을 의미한다. 수식 (2)‑(3)에서 보듯, 이 두 변수는 순간 속도 uₜ와 성장률 gₜ를 적분해 얻으며, 구간을 나누면 가법성(additivity)과 다중 단계 일관성(multi‑step consistency)을 만족한다. 이러한 성질을 이용해 “one‑step update” x₁ = x₀ + v(x₀,0,1), m₁ = m₀·exp(h(x₀,0,1)) 로 직접 최종 상태를 계산한다.

학습 단계에서는 평균 흐름을 신경망으로 파라미터화하고, 순간 속도·성장률과의 관계를 미분 항등식으로 강제한다. 구체적으로 손실 L(θ,ϕ)은 vθ와 hϕ가 순간 uₜ, gₜ와의 관계식 v = uₜ + (T−t)·Dv, h = gₜ + (T−t)·Dh 를 만족하도록 회귀하도록 설계된다. 여기서 D는 전방향 미분 연산이며, stop‑gradient 연산을 통해 목표값이 고정된다. 이 방식은 기존의 CUFM(Conditional Unbalanced Flow Matching)처럼 전체 궤적을 시뮬레이션하지 않아도 되므로, 학습과 추론 모두에서 ODE 통합 비용을 제거한다.

WFR 기하학을 적용함으로써, 평균 흐름이 WFR 거리의 지오데식(geodesic)과 일치하도록 최적화된다. WFR 거리의 정의식 (1)에서 보듯, 이동 비용 ‖u‖²와 질량 변동 비용 δ²‖g‖²가 동시에 최소화된다. 따라서 학습된 v와 h는 실제 세포 집단이 겪는 이동·증식·사멸 과정을 정확히 재현한다.

실험에서는 합성 데이터와 실제 scRNA‑seq 데이터(예: 발달 과정, 약물 처리 후 변이)에서 WFR‑MFM을 평가하였다. 추론 시간은 기존 ODE‑기반 방법(Neural ODE, WFR‑FM 등)에 비해 10⁰‑10² 배 빠르며, KL‑divergence, FID, 및 세포 상태 매핑 정확도 측면에서도 경쟁력 있는 성능을 보였다. 특히 수천 개의 교란 조건을 포함한 대규모 합성 데이터셋에서, WFR‑MFM은 실시간 수준의 응답 예측을 가능하게 하여, 대규모 단일 세포 교란 실험에 실용적인 솔루션을 제공한다.

이 논문의 주요 기여는 (1) 비균형 동적 시스템을 평균 흐름으로 요약하는 이론적 프레임워크, (2) WFR 기하학에 맞춘 평균 흐름 매칭 알고리즘, (3) 시뮬레이션‑프리 학습·추론 파이프라인을 통한 실질적인 속도 향상이다. 향후 연구는 평균 흐름을 다중 모달 데이터에 확장하거나, 확률적 성장률 모델을 도입해 더 복잡한 생물학적 현상을 포착하는 방향으로 진행될 수 있다.