기술 요약: Wasserstein Residuals: 인구 역학으로부터의 경사 흐름 학습 (Wasserstein Residuals: Learning Gradient Flows from Population Dynamics)

문제 정의

본 논문은 희소한 관측 스냅샷 세트로부터 인구 역학(population dynamics)을 재구성하는 역문제(inverse problem)를 다룹다. 계산 생물학 및 군집 역학과 같은 많은 과학 분야에서 인구의 진화는 에너지 범함수 $F$의 최급 하강(steepest descent)에 의해 구동되는 **Wasserstein 경사 흐름(Wasserstein Gradient Flow, WGF)**인 확률 분포 $\rho_t$의 곡선으로 모델링된다. 목표는 관측된 주변 분포 $\{q_t\}$가 이산적 관측 시간 $t \in \mathcal{T}_{obs}$에서 에너지 범함수 $F$의 경사 흐름과 일치하도록 하는 기저 에너지 범함수 $F$를 복구하는 것이다.

기존의 지배적인 접근 방식들은 Jordan–Kinderlehrer–Otto (JKO) 스키마에 의존한다. JKO는 에너지와 최적 운송(optimal transport, OT) 패널티를 최소화하는 근접 단계(proximal steps)의 시퀀스로 WGF를 정의한다. 그러나 JKO 기반 방법론은 다음과 같은 두 가지 주요 한계점을 갖는다:

1. 시간 이산화에 대한 유연성 부족: 연속적인 스냅샷 사이의 비용이 많이 드는 OT 문제를 해결해야 한다.
2. 간극에 대한 민감도: 관측치 사이의 시간 간격이 클 경우 어려움을 겪으며, 스냅샷 사이의 선형 현(chord) 보간에 의존하기 때문에 곡선 궤적(예: 사인파 형태의 골짜기를 따라 움직이는 입자)을 포착하는 데 실패하는 경우가 많다.

방법론: Wasserstein Residual 프레임워크

저자들은 JKO 스키마와 최적 운송 결합(coupling)을 우회하는 **잔차 기반 접근 방식(residual-based approach)**을 제안한다. JKO 조건을 강제하는 대신, 이들은 비음수 손실 함수를 통해 경사 흐름을 특징짓는 근본적인 미분 방정식을 직접적으로 강제한다.

1. 잔차 정식화 (The Residual Formulation)

본 논문은 경사 흐름의 접선(Tangent) 및 에너지 소산 등식(Energy Dissipation Equality, EDE) 정식화를 활용한다. 구체적으로는 다음의 **속도 제약(velocity constraint)**에 집중한다:
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
여기서 $v_t$는 속도장이며, $\frac{\delta F}{\delta \rho_t}$는 에너지 범함수의 1차 변분이다.

저자들은 $\rho$가 $F$의 경사 흐름일 때와 동치인 값이 0이 되는 속도 잔차(velocity residual) $R_{vel}$을 정의한다:
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

전역 목적 함수는 이 잔차와 관측 시간에서의 데이터 적합 발산(divergence) $D$를 결합한다:
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
이 프레임워크는 Path-Finding 및 Action Matching과 같은 기존 방법들을 단일 잔차 원리의 구체적인 사례로 통합한다.

2. "Stitching" 알고리즘

이 목적 함수를 최소화하기 위해, 저자들은 다음과 같은 특성을 가진 시뮬레이션 프리(simulation-free), 입자 기반 방법인 Stitching을 도입한다:

• 학습 가능한 궤적: $\rho$가 암시적인 JKO 방식과 달리, Stitching은 곡선 $\rho_\theta$를 $N$개의 학습 가능한 입자 궤적 $x_{t,k}^\theta$에 대한 **커널 밀도 추정(Kernel Density Estimate, KDE)**으로 명시적으로 매개변수화한다.
• 시뮬레이션 프리 (Simulation-Free): 이 방법은 Neural ODE를 풀거나 SDE를 적분할 필요가 없다. 입자의 위치와 속도는 직접적인 학습 가능 파라미터이다.
• 경험적 근사 (Empirical Approximation): 계산 효율성을 위해 KDE를 입자 중심에서의 경험적 측도로 근사한다. 이를 통해 속도 잔차는 입자들에 대한 합으로 단순화된다:
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• 함수 매개변수화: 에너지 범함수 $F_\theta$는 신경망으로 표현되는 포텐셜 항 $V_\theta$, 엔트로피 항(로그 밀도), 그리고 상호작용 항 $W_\theta$의 합으로 매개변수화된다.

주요 기여

1. 잔차 프레임워크: 저자들은 WGF 학습을 비음수 잔차 항을 최소화하는 것으로 재구성하여, Path-Finding과 Action Matching을 포괄하는 통합된 관점을 제공한다.
2. Stitching 알고리즘: 궤적 곡선 $\rho$를 일급 학습 가능 변수로 취급하는 새로운 시뮬레이션 프리 입자 방법을 도입한다. 이를 통해 JKO 기반의 현 예측(chord predictor)이 놓치는 곡선 역학을 포착하고, 넓은 간격의 관측치로부터 역학을 학습할 수 있는 능력을 부여한다.
3. 최첨단 성능 (State-of-the-Art Performance): Stitching은 희소하거나 짝이 맞지 않는(unpaired) 데이터 시나리오에서 궤적 추론 벤치마크에 대해 우수한 결과를 달성한다.

실험 결과

본 논문은 세 가지 주요 벤치마크에서 Stitching의 성능을 평가한다:

1. 예시적 사례 (연속 흐름 vs. JKO)

합성 "물결치는 골짜기(wavy valley)" 포텐셜에서 입자들은 사인파 궤적을 따른다.

• 결과: JKO 기반 방법(특히 JKOnet*)은 희소한 스냅샷 사이($t=10$에서 $t=20$)의 곡률을 포착하지 못하고 직선 보간을 생성한다. 반면 Stitching은 곡선 궤적을 정확히 복구하며, 이는 기저 포텐셜을 훨씬 더 잘 복구함을 보여준다 ($R^2 = 0.62$ 대 $0.51$).

2. 합성 포텐셜 복구

이 방법은 쌍이 맞는(paired) (상관관계가 있는 궤적) 및 쌍이 맞지 않는(unpaired) (상관관계가 없는 스냅샷) 환경에서의 15가지 2차원 포텐셜에 대해 테스트되었다.

• 결과: Stitching은 두 환경 모두에서 일관된 성능을 유지한다. 반면 JKO 기반 방법은 연속적인 스냅샷 간의 OT 결합에 의존하기 때문에 unpaired 설정에서 성능이 크게 저하되거나 붕괴된다. Stitching의 오차는 두 환경 사이에서 약 $2\times$ 이내를 유지하지만, JKO 방법은 unpaired 케이스에서 유의미한 경사장을 복구하는 데 종종 실패한다.

3. 단일 세포 궤적 추론 (Embryoid Body 데이터셋)

인간 배아 줄기세포 분화 데이터(27일 동안 5개의 스냅샷)에 적용되었다.

• 결-과: Stitching은 Neural SDEs, TrajectoryNet 및 다양한 JKO 기반 방법들을 포함한 베이스라인들과 비교하여 가장 우수한 성능($W_1$ 및 $W_2$ 거리 최저)을 달-성했다. 또한 전체 데이터 학습 및 "leave-two-out" 시간적 보간 프로토콜 모두에서 경쟁 모델들을 압도했다.

4. 상호작용 역학 복구

McKean–Vlasivov 시스템에서 환경 포텐셜($V$)과 상호작용 커널($W$)을 분리하는 데 이 방법이 사용되었다.

• 결과: Stitching은 구속 포텐셜(confining potential)과 인력 상호작용 커널(attractive interaction kernel)을 모두 성공적으로 복구하며, 높은 패턴 $R^2$ 점수($V$에 대해 $0.84$, $W$에 대해 $0.81$)를 달성했다. 또한 매 단계마다 OT 문제를 풀어야 하는 JKO 기반 방법보다 에포크당 계산 속도가 현저히 빠르다.

5. 비경사 흐름 (Non-Gradient Flows)

저자들은 경사 항을 일반적인 벡터장으로 교체함으로써 잔차 프레임워크가 경사 흐름을 넘어 카이랄(chiral, 비보존적) 역학을 추적할 수 있음을 입증한다.

의의 및 주장

본 논문은 Wasserstein Residual 관점이 인구 역학 학습을 위한 JKO 스키마보다 더 유연하고 강력한 대안을 제공한다고 주장한다. 최적 운송 결합의 제약으로부터 에너지 범함수의 학습을 분리함으로써, Stitching은 다음과 같은 기능을 가능하게 한다:

• 희소성에 대한 강건성: 선형 보간에 내재된 "곡률 손실" 없이 넓게 떨어진 관측치로부터 역학을 학습할 수 있다.
• Unpaired 데이터 호환성: 입자의 정체가 시간에 따라 추적되지 않는 데이터셋(생물학적 이미징의 흔한 시나리오)에서도 학습이 가능하다.
• 시뮬레이션 프리 효율성: 훈련 중 ODE나 OT 문제를 푸는 데 드는 계산 오버헤드를 피한다.

저자들은 이 연구를 기존 패러다임의 통합이자, $N \to \infty$일 때의 이론적 수렴 보장 부재 및 상호작용 항에 대한 $O(N^2)$ 계산 복잡도와 같은 한계를 인정하면서도, 단일 세포 생물학에서 집단 행동에 이르는 다양한 분야의 과학적 발견을 위한 실질적인 진보로 자리매김한다.