[논문 리뷰] Proximal Optimal Transport Modeling of Population Dynamics

Charlotte Bunne, Laetitia Meng-Papaxanthos외 2명|arXiv (Cornell University)|2021. 06. 11.

Gaussian Processes and Bayesian Inference인용 수 15

한 줄 요약

이 논문은 희소하고 시간에 따라 분포된 포인트 클라우드 샘플에서 에너지 함수를 학습함으로써, Jordan-Kinderlehrer-Otto (JKO) 플로우를 통한 근접 최적 운반을 사용하여 인구 동역학을 모델링하는 미분 가능 신경망인 JKOnet을 제안한다. 이는 단일세포 RNA-seq 데이터에서 궤적 복원과 생물학적 이질성 포착에 있어 직접 전진 모델링보다 뛰어나며, 노이즈에 강건하고 선조 분화 예측을 정확하게 수행함을 보여준다.

ABSTRACT

We propose a new approach to model the collective dynamics of a population of particles evolving with time. As is often the case in challenging scientific applications, notably single-cell genomics, measuring features for these particles requires destroying them. As a result, the population can only be monitored with periodic snapshots, obtained by sampling a few particles that are sacrificed in exchange for measurements. Given only access to these snapshots, can we reconstruct likely individual trajectories for all other particles? We propose to model these trajectories as collective realizations of a causal Jordan-Kinderlehrer-Otto (JKO) flow of measures: The JKO scheme posits that the new configuration taken by a population at time $t+1$ is one that trades off an improvement, in the sense that it decreases an energy, while remaining close (in Wasserstein distance) to the previous configuration observed at $t$. In order to learn such an energy using only snapshots, we propose JKOnet, a neural architecture that computes (in end-to-end differentiable fashion) the JKO flow given a parametric energy and initial configuration of points. We demonstrate the good performance and robustness of the JKOnet fitting procedure, compared to a more direct forward method.

연구 동기 및 목표

• 주기적이고 파괴적인 샘플만 제공될 때 인구의 개별 입자 궤적을 복원하는 것.
• 에너지 최소화와 이전 상태에 대한 워샤르슈타인 거리 유사도를 균형 잡는 인과적, 에너지 기반 JKO 플로우로 입자 구성의 진화를 모델링하는 것.
• 종료형 미분 가능 최적화를 사용하여 관측된 샘플에서 파arametric 에너지 함수를 학습하는 것.
• 노이즈가 많고 고차원적인 환경에서 직접 전진 모델링 접근법에 비해 강건성과 예측 정확도를 향상시키는 것.
• 시간에 따라 분포된 단일세포 데이터에서 줄기세포 분화 중 선조 분열과 같은 생물학적 이질성을 포착하는 것.

제안 방법

• JKOnet는 에너지 함수를 매개변수화하고 볼록성을 보장하기 위해 입력 볼록 신경망(ICNNs)을 사용하여 JKO 단계를 계산한다. 이는 미분 가능한 최적 운반 플로우를 가능하게 한다.
• 모델은 JKO 단계를 측도 공간에서의 근접 연산자로 프레임화하여, 에너지 함수를 최소화하면서도 이전 구성과 워샤르슈타인 거리에서 가까이 있도록 한다.
• 에너지 매개변수를 학습하기 위해 시간에 따라 예측된 포인트 클라우드 샘플과 관측된 샘플 간의 워샤르슈타인 거리를 최소화하는 이중 최적화 문제를 사용한다.
• 아키텍처는 완전히 미분 가능하여, JKO 단계를 통해 역전파하여 에너지 함수 매개변수를 업데이트할 수 있다.
• 이 방법은 Brenier 정리를 활용하여 플로우가 볼록 잠재함수의 기울기에서 유래되며, 최적 운반 원칙과 일치함을 보장한다.
• 모델은 합성 데이터와 인간 배아 줄기세포의 실제 단일세포 RNA-seq 데이터를 대상으로 평가되었으며, 직접 전진 모델링 기반선과의 비교가 이루어졌다.

실험 결과

연구 질문

• RQ1 미분 가능하고 근접 최적 운반 기반의 프레임워크는 인구의 주기적이고 파괴적인 샘플만으로도 개별 입자 궤적을 복원할 수 있는가?
• RQ2 노이즈와 제한된 데이터 하에서 JKOnet의 성능은 직접 전진 모델링 접근법에 비해 어떻게 비교되는가?
• RQ3 JKOnet는 시간에 따라 분포된 단일세포 데이터에서 줄기세포 분화 중 선조 분열과 같은 복잡한 생물학적 이질성을 어느 정도 포착할 수 있는가?
• RQ4 장기 예측, 특히 다수 단계를 앞서 예측할 경우에도 JKO 기반 접근법은 안정적인 예측 오차를 유지하는가?
• RQ5 모델은 복잡한 생물학적 시스템의 기저 동역학을 반영하는 의미 있고 해석 가능한 에너지 함수를 학습할 수 있는가?

주요 결과

• JKOnet는 모든 시간 단계에서 1단계 예측 및 전체 단계 예측에 대해 $W_{ε}$ 거리 측정치에서 직접 전진 방법을 능가하며, 일주일 27일까지 안정적인 성능을 유지한다.
• 전체 궤적 예측(모든 단계 예측) 시에도 $W_{ε}$ 이질성이 낮게 유지되어 오차 누적에 대한 강건성을 보여준다.
• JKOnet와 전진 방법 모두 인간 배아 줄기세포 분화에서 주요 선조 분열(중순엽, 내중엽, 신경외배엽, 신경주머니)을 포착하지만, 적은 수의 클래스는 정확도가 낮다.
• 예측된 선조 분포와 진짜 선조 분포 간의 $\ell_1$ 및 헬링거 거리 측정치는, 전체 궤적 예측보다 단일 단계 예측에서 성능이 떨어지는 것을 보여준다.
• PCA 및 PHATE 시각화에서 예측된 선조 분포가 진짜 값과 잘 일치함으로써, JKOnet가 생물학적 이질성을 더 잘 유지함을 입증한다.
• 모델는 강력한 일반화 및 안정성을 보이며, 노이즈가 있거나 희소한 데이터에서도 장기 예측 시 예측 오차가 유의미하게 증가하지 않는다.