[논문 리뷰] Trajectory Inference via Mean-field Langevin in Path Space

Lénaïc Chizat, Stephen X. Zhang외 2명|arXiv (Cornell University)|2022. 05. 14.

Single-cell and spatial transcriptomics인용 수 4

한 줄 요약

이 논문은 경로 공간에서 위너 측도에 대한 최소 엔트로피 추정량을 계산하기 위해 슈뢰딩거 브리지와 노이즈가 있는 기울기 하강법을 활용하는 격자 없는 평균장 라ング주인(MFL) 동역학 알고리즘을 제안한다. 이 방법은 종단 간 이론적 보장을 통해 희소한 시간 해상도 데이터로부터 일관된 추론을 가능하게 하며, 단일세포 RNA-시퀀싱 및 질량 변화 모델링에 적용 가능하다.

ABSTRACT

Trajectory inference aims at recovering the dynamics of a population from snapshots of its temporal marginals. To solve this task, a min-entropy estimator relative to the Wiener measure in path space was introduced by Lavenant et al. arXiv:2102.09204, and shown to consistently recover the dynamics of a large class of drift-diffusion processes from the solution of an infinite dimensional convex optimization problem. In this paper, we introduce a grid-free algorithm to compute this estimator. Our method consists in a family of point clouds (one per snapshot) coupled via Schrödinger bridges which evolve with noisy gradient descent. We study the mean-field limit of the dynamics and prove its global convergence to the desired estimator. Overall, this leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We also present how to adapt the method to deal with mass variations, a useful extension when dealing with single cell RNA-sequencing data where cells can branch and die.

연구 동기 및 목표

• 무한차원 볼록 최적화로 인해 계산적으로 비가역적인 경로 공간에서의 최소 엔트로피 추정량을 계산하기 위한 격자 없는 확장 가능한 알고리즘을 개발하는 것.
• 제안된 MFL 동역학이 진정한 최소 엔트로피 추정량으로의 전역 수렴을 증명함으로써 궤적 추론에 종단 간 이론적 보장을 제공하는 것.
• 세포 분열과 사멸과 같은 질량 변화를 다룰 수 있도록 방법을 확장하여, 단일세포 RNA-시퀀싱 응용에서 흔히 발생하는 문제를 해결하는 것.
• 엄격한 수렴 보장을 갖지 못하는 비볼록 딥러닝 기반 궤적 추론 방법에 대한 теор적으로 탄탄하고 해석 가능한 대안을 제공하는 것.

제안 방법

• 시간 스탬프별로 하나의 포인트 클라우드를 구성하고, 이를 엔트로피 정규화된 최적 운반 또는 슈뢰딩거 브리지로 연결하여 시간에 따른 동역학을 모델링한다.
• 관측된 데이터의 적합도와 엔트로피 정규화를 동시에 고려하는 목적 함수에 대해 노이즈가 있는 기울기 하강법을 사용하는 평균장 라ング주인 동역학으로 추론 문제를 설정한다.
• 관측된 주변 분포에 대한 충실도와 노이즈에 의한 부드러움을 상대 엔트로피를 통해 위너 측도에 대해 균형을 이루는 확률적 미분 방정식으로 동역학을 제어한다.
• 안정적인 수렴을 확보하기 위해 단계 크기 $\eta$, 노이즈 수준 $\sigma$, 온도 $\tau$에 대한 안내 스케줄을 사용한다.
• 기본 성장률 $g(t,x)$를 통합하여 질량 변화를 다룰 수 있도록 동역학의 수정된 버전을 도입함으로써 세포 증식과 사멸을 모델링할 수 있도록 한다.
• 명시적인 경계 조건을 부여하지 않고도 관측된 데이터의 지지체 내에서 입자의 안정성을 기반으로 하여 CPU 기반 프레임워크에서 구현된다.

실험 결과

연구 질문

• RQ1 무한차원 볼록 최적화로 인해 계산적으로 비가역적인 경로 공간에서 최소 엔트로피 추정량을 이론적 수렴 보장과 함께 계산할 수 있는 격자 없는 확률적 최적화 방법을 설계할 수 있는가?
• RQ2 평균장 라ング주인 동역학은 어떻게 전역적으로 진정한 궤적 추론 해법으로 수렴하며, 수렴 속도는 어떠한가?
• RQ3 표준 운반 모델이 포착하지 못하는 세포 분열과 사망과 같은 질량 변화를 다룰 수 있도록 방법을 확장할 수 있는가?
• RQ4 희소 샘플링에 대한 복원 정확도와 내성 면에서 기존 방법인 Global Waddington-OT에 비해 제안된 MFL 방법의 성능은 어떠한가?

주요 결과

• MFL 동역학은 지수적 속도로 최소 엔트로피 추정량으로 전역 수렴하며, 궤적 추론에 대해 강력한 이론적 보장을 확립한다.
• 합성 및 실제 단일세포 데이터에서 최신 기술 수준의 성능을 달성하였으며, 진짜 값과의 에너지 거리가 MFL의 경우 $\lambda = 0.025$, gWOT의 경우 $\lambda = 0.01$에서 최소화된다.
• 입자 수가 50개뿐인 조건에서도 복잡한 동역학, 예를 들어 분기 과정과 시간에 따라 변화하는 잠재력 등을 성공적으로 재구성한다.
• 알려진 성장률 $g(t,x)$를 포함함으로써 질량 변화, 예를 들어 세포 증식을 정확하게 모델링할 수 있으며, 생물학적으로 현실적인 상황에서 재구성 정밀도가 향상된다.
• 알려진 경계 조건이 없이도 안정적으로 유지되며, 최적화 과정 전반에 걸쳐 입자들이 관측된 데이터의 지지체 내에 머물러 있다.
• 초기 조건 설정에 대해 강건하며, 에너지 거리와 시각적 재구성 품질을 통해 다양한 데이터셋과 파rameter 설정에서 일관된 성능을 달성한다.