[논문 리뷰] Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Guillaume Huguet, D S Magruder외 5명|arXiv (Cornell University)|2022. 06. 29.

Traffic Prediction and Management Techniques인용 수 19

한 줄 요약

MIOFlow는 정적이고 시간 간격이 넓은 샘플에서 연속적이고 확률적인 인구 동역학을 학습하기 위해 다양체 학습과 최적 운반을 결합한 신경 ODE 기반 프레임워크를 제안한다. 지오데식 오토에인코더를 사용하여 거리가 새로운 다중 척도 확산 지오데식 거리에 반영되는 잠재 공간에 데이터를 임bedding함으로써, 복잡한 다양체에서 정확한 보간과 궤적 추론이 가능해지며, 정규화된 플로우와 슈뢰딩거 다리를 능가한다.

ABSTRACT

We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) that learns stochastic, continuous population dynamics from static snapshot samples taken at sporadic timepoints. MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differential equations (Neural ODE) to interpolate between static population snapshots as penalized by optimal transport with manifold ground distance. Further, we ensure that the flow follows the geometry by operating in the latent space of an autoencoder that we call a geodesic autoencoder (GAE). In GAE the latent space distance between points is regularized to match a novel multiscale geodesic distance on the data manifold that we define. We show that this method is superior to normalizing flows, Schrödinger bridges and other generative models that are designed to flow from noise to data in terms of interpolating between populations. Theoretically, we link these trajectories with dynamic optimal transport. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.

연구 동기 및 목표

• 고차원 데이터의 정적이고 시간 순서가 지정된 샘플에서 연속적이고 확률적인 인구 동역학을 추론하는 데 도전하는 것.
• 단일세포 RNA 시퀀싱과 같이 종단적 추적을 수행할 수 없는 시스템에서 궤적 추론을 모델링하는 것.
• 추론된 궤적이 잠재된 데이터 다양체의 내재 기하학을 존중하도록 보장하는 것, 주로 고차원 공간이 아닌 내재 기하학을 기반으로 하는 것.
• 기존 생성 모델에서 흔한 제약인 가우시안 사전분포와 결정론적 플로우를 피하는 계산적으로 효율적인 방법을 개발하는 것.
• 복잡한 생물학적 시스템에서 관측되지 않은 시간대와 개별 세포 궤적의 정확한 보간 및 재구성 가능하게 하는 것.

제안 방법

• MIOFlow는 서로 다른 시간대의 정적 인구 샘플 간의 연속적 플로우를 모델링하기 위해 신경 미분방정식(Neural ODEs)을 사용한다.
• 지오데식 오토에인코더(GAE)는 데이터 다양체 상에서 새로운 다중 척도 확산 지오데식 거리를 유지하도록 잠재 공간을 정규화한다.
• 확산 지오데식 거리는 커널 기반 그래프 라플라시안을 통해 정의되며, 무한한 샘플 수에 도달할 경우 진짜 다양체 상의 지오데식 거리로 수렴한다.
• 플로우는 다양체 기준 거리로 최적 운반을 통해 페널티를 부여하여, 시간대 간 운반 비용이 데이터의 내재 기하학과 일치하도록 보장한다.
• 방법은 GAE의 잠재 공간에서 작동하여 계산 비용을 줄이고, 저차원 다양체에서의 동역학 학습을 통해 일반화 성능을 향상시킨다.
• 잠재 공간에서의 초기 샘플링을 통해 확률성이 자연스럽게 도입되며, 결정론적 플로우에서처럼 명시적 노이즈 주입이 필요로 하지 않는다.

실험 결과

연구 질문

• RQ1 가우시안 사전분포를 가정하지 않고도 정적이고 시간 간격이 넓은 샘플에서 연속적이고 확률적인 인구 동역학을 학습할 수 있는가?
• RQ2 추론된 궤적이 잠재된 데이터 다양체의 내재 기하학적 구조를 존중하도록 보장할 수 있는가?
• RQ3 학습된 잠재 공간에서 진짜 다양체 상의 지오데식 거리로 수렴하는 데이터 기반의 다중 척도 거리 측정법을 구성할 수 있는가?
• RQ4 MIOFlow는 정규화된 플로우, 슈뢰딩거 다리, TrajectoryNet과 같은 기존 방법과 비교해 복잡한 동역학(예: 분열과 융합)을 모델링하는 데 얼마나 우수한가?
• RQ5 MIOFlow는 개별 세포 궤적을 정확히 재구성하고, 분화나 치료 반응과 같은 생물학적 과정 중 유전자 발현의 동적 변화를 포착할 수 있는가?

주요 결과

• MIOFlow는 분열과 융합이 있는 시뮬레이션 데이터에서 정규화된 플로우와 슈뢰딩거 다리보다 보간 성능이 뛰어나며, 더 매끄럽고 지오데식으로 정확한 궤적을 생성한다.
• 시뮬레이션 데이터에서 MIOFlow는 분리된 궤적을 성공적으로 모델링하며, 플로우가 잠재 다양체의 가지치기 형태를 정확히 포착한다.
• 배아 유사체 분화의 단일세포 RNA-seq 데이터에서 MIOFlow는 알려진 발달 마커와 일치하는 생물학적으로 타당한 선조 궤적을 재구성한다.
• 급성 백혈병(AML) 치료 데이터에서 MIOFlow는 생존 세포가 후기 시간대에 백혈병 줄기세포(LSC) 마커로 전이됨을 드러내며, SLPI, MYB, B2M의 발현이 증가함을 보였다.
• 확산 지오데식 거리가 적용된 지오데식 오토에인코더는 정확한 다양체 재구성을 가능하게 하며, 잠재 공간의 플로우는 생물학적으로 해석 가능한 유전자 발현 패tern으로 복원되는 의미 있는 생물학적 거리를 유지한다.
• MIOFlow는 자코비안 추적 계산을 피하기 때문에 정규화된 플로우보다 기능 평가당 더 빠른 추론을 달성하며, 고차원 차원 k에 대해 O(k)로 스케일링되며, O(k^2)가 아닌 방식이다.