[논문 리뷰] Simulation-free Schrödinger bridges via score and flow matching

Alexander Tong, Nikolay Malkin외 6명|arXiv (Cornell University)|2023. 07. 07.

Model Reduction and Neural Networks인용 수 6

한 줄 요약

이 논문은 [SF]2M을 소개한다. [SF]2M은 엔트로픽 최적 운반을 통해 스코어 매칭과 플로우 매칭을 통합함으로써, 임의의 소스 및 타겟 분포 간의 확률적 역학을 시뮬레이션 없이 학습하는 방법이다. 이는 고차원 세포 역학의 정확하고 확장 가능한 모델링을 가능하게 하며, SDE 시뮬레이션을 필요로 하지 않고도 단일세포 데이터로부터 기존에 알려진 유전자 조절 네트워크를 복원할 수 있다.

ABSTRACT

We present simulation-free score and flow matching ([SF]$^2$M), a simulation-free objective for inferring stochastic dynamics given unpaired samples drawn from arbitrary source and target distributions. Our method generalizes both the score-matching loss used in the training of diffusion models and the recently proposed flow matching loss used in the training of continuous normalizing flows. [SF]$^2$M interprets continuous-time stochastic generative modeling as a Schrödinger bridge problem. It relies on static entropy-regularized optimal transport, or a minibatch approximation, to efficiently learn the SB without simulating the learned stochastic process. We find that [SF]$^2$M is more efficient and gives more accurate solutions to the SB problem than simulation-based methods from prior work. Finally, we apply [SF]$^2$M to the problem of learning cell dynamics from snapshot data. Notably, [SF]$^2$M is the first method to accurately model cell dynamics in high dimensions and can recover known gene regulatory networks from simulated data. Our code is available in the TorchCFM package at https://github.com/atong01/conditional-flow-matching.

연구 동기 및 목표

• 기존 스코어 기반 및 플로우 기반 모델이 가우시안 소스 또는 시뮬레이션 중심 학습에 의존하는 한계를 해결한다.
• 학습 중에 확률적 과정을 시뮬레이션하지 않고도 임의의 소스 및 타겟 분포 간의 확률적 역학을 학습할 수 있도록 한다.
• 반복적인 SDE 시뮬레이션을 필요로 하는 전통적 슈뢰딩거 브리지 방법의 확장성 및 수치적 비효율성을 극복한다.
• 스코어 매칭(확산 모델)과 플로우 매칭(노멀라이징 플로우)을 임의의 소스 분포로 일반화하는 통합 프레임워크를 개발한다.
• 쌍방향 스크래치 데이터에서 고차원 생물학적 시스템, 예를 들어 단일세포 유전자 발현 역학을 정확하게 모델링할 수 있도록 한다.

제안 방법

• 엔트로픽 최적 운반을 사용해 슈뢰딩거 브리지 문제를 엔트로픽 최적 운반에 기반한 브라운 운반의 혼합으로 공식화함으로써 정적이고 시뮬레이션 없는 목표 함수를 도출한다.
• 싱크호른 알고리즘을 사용해 소스 및 타겟 분포 간의 정적 엔트로픽 최적 운반 맵을 효율적으로 계산함으로써 동적 SDE 시뮬레이션을 대체한다.
• 신경망을 학습시켜 정합된 확률적 회귀 목표 함수 내에서 ODE 드리프트(플로우 매칭)와 조건부 스코어(스코어 매칭)를 동시에 예측한다.
• 중간 지점들을 브라운 운반의 주변 분포에서 샘플링하고, 모델을 드리프트 및 스코어 함수에 대해 정합하는 확률적 회귀 목표 함수를 정의한다.
• 신경 그래픽 모델(NGM)을 ODE 드리프트에 통합하여 유전자 간 상호작용 네트워크를 매개변수화하고, 희박성 강화를 위해 L1 정규화를 적용한다.
• 학습 시 ODE 통합이 필요 없는 시뮬레이션 없는 목표 함수를 사용하여 고차원 데이터에서 효율적이고 확장 가능한 학습을 가능하게 한다.

실험 결과

연구 질문

• RQ1 시뮬레이션 없는 목표 함수가 임의의 분포 간의 확률적 역학을 학습하기 위해 스코어 매칭과 플로우 매칭을 통합할 수 있는가?
• RQ2 엔트로픽 최적 운반을 사용해 슈뢰딩거 브리지 문제의 정적이고 시뮬레이션 없는 대체 수단을 구성할 수 있는가?
• RQ3 [SF]2M이 쌍방향 스크래치 데이터에서 고차원 생물학적 역학, 예를 들어 단일세포 유전자 발현을 정확하게 모델링할 수 있는가?
• RQ4 [SF]2M이 시뮬레이션된 단일세포 데이터로부터 기존에 알려진 유전자 조절 네트워크를 복원할 수 있으며, 기존 베이스라인보다 우수한 성능을 보일 수 있는가?
• RQ5 플로우 및 스코어 역학을 함께 모델링하는 것이 별도의 스코어 또는 플로우 모델 대비 일반화 및 강인성을 향상시키는가?

주요 결과

• [SF]2M 방법은 SDE 시뮬레이션 없이도 기존의 시뮬레이션 기반 방법보다 더 정확하고 효율적인 슈뢰딩거 브리지 문제 해결을 달성한다.
• 합성 데이터에서 [SF]2M는 추론 시에 산란 스케줄을 조절하여 ODE와 SDE 역학 사이를 성공적으로 보간함으로써, 확률적 성격을 모델링하는 데 있어 유연성을 입증한다.
• 고차원 단일세포 데이터에서 [SF]2M는 유전자 공간 전체에서 세포 역학을 정확하게 모델링할 수 있는 최초의 방법으로, 시뮬레이션된 데이터로부터 기존에 알려진 유전자 조절 네트워크를 복원한다.
• 분열 및 삼중 분열 유전자 조절 네트워크 시뮬레이션에서 [SF]2M는 각각 AUC-ROC 점수 0.92와 0.89를 기록하여 스피어만, 피어슨, DREMI, 그랑저 베이스라인을 모두 능가한다.
• 단일세포 데이터의 첫 1000개의 고변동 유전자에 대해 예측된 분포와 진짜 분포 사이의 1-Wasserstein 거리는 0.12로 측정되어 강력한 분포 일치를 보여준다.
• NGM의 인접 행렬에 L1 정규화를 적용함으로써 진짜 GRN과 유사한 희박하고 해석 가능한 유전자 간 상호작용 네트워크를 얻을 수 있었다.