Relative Entropy Estimation in Function Space: Theory and Applications to Trajectory Inference

1. 문제 정의 (Problem Definition)

배경:

• 궤적 추론 (Trajectory Inference, TI): 단일 세포 유전체학 (single-cell genomics) 등 파괴적인 측정 방식으로 인해 연속적인 궤적이 아닌, 이산적인 시간점에서의 상태 스냅샷 (snapshot) 만 관측되는 상황에서, 잠재적인 동적 과정을 복원하는 작업입니다.
• 현재의 한계: 기존 TI 방법론들은 주로 보유된 (held-out) 스냅샷 데이터의 한계 (marginals) 를 재구성하는 능력으로 평가됩니다. (예: Earth Mover's Distance, Wasserstein Distance 등).
• 핵심 문제:
  • 유한한 수의 한계 분포 (marginals) 만으로는 서로 다른 시간대의 상태들이 어떻게 연결되는지 (궤적의 동역학) 를 결정할 수 없습니다.
  • 따라서, 서로 다른 동역학을 가진 모델들이 동일한 한계 분포 점수를 얻을 수 있어, 실제 추론된 궤적의 품질을 평가하기 어렵습니다.
  • 기존 평가 지표는 함수 공간 (function space) 상의 전체 확률 측도 (probability measure) 간의 차이를 포착하지 못합니다.

2. 방법론 (Methodology)

이 논문은 함수 공간 (Function Space) 상의 두 확률 측도 간의 Kullback-Leibler (KL) 발산을 추정할 수 있는 일반적인 프레임워크를 제안합니다.

핵심 기술 요소:

1. 함수 흐름 매칭 (Functional Flow Matching, FFM) 기반 접근:

   • 참조 가우시안 법칙과 데이터 유도 법칙 사이의 조건부 가우시안 측도 경로를 구성합니다.
   • 이를 통해 데이터 분포에 대응하는 무조건부 측도의 경로를 유도합니다.

2. 속도장 (Velocity Field) 만을 이용한 KL 발산 유도:

   • 이론적 기반: Cameron-Martin 공간과 Radon-Nikodym 도함수의 존재성을 가정합니다.
   • 주요 결과 (Theorem 3.6): KL 발산을 로그 Radon-Nikodym 도함수의 기울기 (gradient) 와 속도장 ($v_t$) 의 불일치로 표현할 수 있음을 증명합니다.
   • 최종 추정식: 두 확률 분포 $\nu_A$ 와 $\nu_B$ 간의 KL 발산은 두 분포에 대응하는 속도장 ($v^A_t, v^B_t$) 의 차이 제곱을 시간과 공간에 대해 적분한 형태로 표현됩니다.
     $$KL(\nu_A \| \nu_B) = \int_0^1 \int_H \frac{t}{1-t} \| v^A_t(x) - v^B_t(x) \|^2_{H_{\mu_0}} d\mu^A_t(x) dt$$
   • 이 식은 ODE 적분을 통한 전체 생성 역학 시뮬레이션 없이도, 학습된 신경망 속도장만으로 KL 을 추정할 수 있게 합니다.

3. 구현 (Implementation):

   • MINO (Mesh-Informed Neural Operator): 두 개의 속도장을 단일 네트워크로 동시에 근사하기 위해 사용 (이진 조건 변수 $c$로 구분).
   • 안정화 기법: $t \to 1$일 때 발생하는 $\frac{t}{1-t}$ 특이점을 제거하기 위해 경계 조건 $v(x, 1) = x$를 만족하도록 파라미터화합니다.
   • 모의 실험 (Monte Carlo): 학습된 속도장과 샘플을 사용하여 식 (7) 을 수치적으로 적분하여 FKL (Functional KL) 값을 계산합니다.

3. 주요 기여 (Key Contributions)

1. 함수 공간 KL 추정기의 일반적 구성: 절대 연속성 (absolute continuity) 과 trace-class 노이즈 공분산에 대한 엄밀한 처리를 바탕으로, 함수 분포 간의 KL 발산을 추정하는 단순하고 확장 가능한 추정기를 제시했습니다.
2. 이론적 검증: 해석적 KL 발산을 알 수 있는 특수한 경우 (가우시안 측도, 선형 SDE) 에서 추정된 FKL 이 이론적 값과 매우 잘 일치함을 검증했습니다.
3. TI 방법론에 대한 포괄적 벤치마크: 합성 데이터 (Lotka-Volterra, Repressilator, Petal) 와 실제 scRNA-seq 데이터 (EB, hESC, ME, HF) 를 사용하여 기존 5 가지 주요 스냅샷 기반 평가 지표와 새로운 FKL 지표를 비교했습니다.
4. 새로운 평가 기준의 제안: 기존 한계 기반 지표들이 일관되지 않은 순위와 잘못된 평가를 초래할 수 있음을 보여주었으며, FKL 이 동역학의 불일치를 명확히 포착하여 TI 방법론 평가에 더 원리적이고 일관된 기준을 제공함을 입증했습니다.

4. 실험 결과 (Results)

A. 추정 정확도:

• 가우시안 및 선형 SDE 시나리오에서 추정된 FKL 값이 해석적 정답 (Ground Truth) 과 매우 근사하게 일치하여 이론적 타당성을 입증했습니다.

B. 기존 평가 지표의 한계 (Finding 1 & 2):

• 불일치성: 스냅샷 시간점, 메트릭 선택 (EMD, W2, MMD 등), 하이퍼파라미터에 따라 방법론의 순위가 극단적으로 달라집니다.
• 동역학 무시: 한계 분포 (marginals) 를 잘 맞추더라도 실제 궤적의 동역학 (예: 분기, 순환, 분기점) 을 잘못 추론할 수 있습니다. 예를 들어, Repressilator 데이터에서 TIGON 은 한계 분포 점수는 좋았으나 실제 3 유전자 진동 패턴을 재현하지 못했습니다.

C. FKL 의 효과성 (Finding 3 & 4):

• 일관된 순위: FKL 은 시각적 검증과 일치하는 일관된 방법론 순위를 제공합니다. (예: Petal 데이터에서 MSBM 이 FKL 에서는 최상위이지만, 기존 지표에서는 하위 순위였습니다.)
• 동역학 불일치 포착: FKL 은 스냅샷 거리에서는 유사해 보이지만 실제 동역학이 다른 모델들 (예: 결정론적 흐름 vs 확률적 흐름) 을 명확히 구분합니다.
• 방향성 KL 의 의미:
  • Forward KL ($KL(\nu_A \| \nu_B)$): 모드 커버리지 (mode coverage) 를 선호합니다. (참조 분포의 영역을 잘 덮는 모델에 유리, 예: MFL).
  • Reverse KL ($KL(\nu_B \| \nu_A)$): 모드 시킹 (mode seeking) 을 선호합니다. (참조 분포의 영역에 집중하는 모델에 유리, 예: MSBM).
  • 이 비대칭성은 기존 지표가 놓치는 모델의 성향을 드러냅니다.

5. 의의 및 결론 (Significance)

• 평가 패러다임의 전환: 단일 세포 유전체학 및 동적 시스템 모델링 분야에서, "스냅샷의 정확성"에서 **"전체 궤적 분포의 정확성"**으로 평가 기준을 전환할 수 있는 이론적, 실용적 도구를 제공합니다.
• 신뢰성 있는 비교: 부분 관측 (partial observability) 하에서 서로 다른 TI 방법론들을 비교할 때, FKL 은 동역학적 불일치를 드러내어 연구자들이 모델의 실제 성능을 더 정확하게 판단할 수 있게 합니다.
• 확장성: 이 방법은 실제 데이터의 복잡성과 높은 차원성을 처리할 수 있도록 설계되어, 함수 공간에서의 생성 모델 평가에 대한 새로운 표준을 제시합니다.

요약하자면, 이 논문은 **함수 공간 KL 발산 (FKL)**을 도입하여 기존 TI 평가 방법의 근본적인 한계를 해결하고, 동역학적 정확성을 기반으로 한 더 엄밀한 모델 평가 체계를 확립했습니다.