[논문 리뷰] A Computational Framework for Solving Wasserstein Lagrangian Flows

Kirill Neklyudov, Rob Brekelmans외 4명|arXiv (Cornell University)|2023. 10. 16.

Model Reduction and Neural Networks인용 수 4

한 줄 요약

이 논문은 라그랑주 작용 함수의 이중 표현을 활용하여 워샤르스타인 라그랑주 흐름을 해결하기 위한 통합적인 딥러닝 프레임워크를 제안한다. 이는 궤적 역전파 또는 최적의 커플링에 대한 접근 없이도 효율적인 최적화를 가능하게 하며, 단일 세포 역행 추론에서 이전의 접근 방식을 능가한다. 이 방법은 사용자 정의 가능한 운동 에너지 및 위치 에너지 항을 통해 다양한 사전 지식을 민첩하게 통합할 수 있다.

ABSTRACT

The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schrödinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.

연구 동기 및 목표

• 단일 세포 RNA-seq 데이터에서 희소하고 국소적인 관측치로부터 유도되는 궤적 추론 문제의 잘 정의되지 않은 성격을 해결하기 위해.
• 스체로딩어 브리지, 비균형 OT, 물리적으로 제약된 OT와 같은 다양한 최적 운반 변종들을 단일 계산 프레임워크로 통합하기 위해.
• 라그랑주 표현에서 사용자 정의 가능한 운동 에너지 및 위치 에너지 항을 통해 역학에 사전 지식을 민감하게 통합하기 위해.
• 개별 입자 궤적을 시뮬레이션하거나 역전파하지 않는, 미분 가능하고 확장 가능한 방법을 개발하기 위해.
• 이중 최적화를 활용하여 단일 세포 역행 추론 및 이미지 생성에서 뛰어난 성능을 입증하기 위해.

제안 방법

• 궤적 추론을 운동 에너지(워샤르스타인-2 또는 워샤르스타인-피셔-레이어 기하학)와 위치 에너지(예: 엔트로피 또는 공간 잠재력)로 구성된 라그랑주 작용 함수의 최소화 문제로 공식화한다.
• 라그랑주의 이중 표현을 사용하여 실현 가능한 최적화 목표를 도출함으로써, 개별 입자 궤적을 시뮬레이션하거나 미분하는 데 필요한 과정을 피한다.
• 딥 뉴럴 네트워크를 사용하여 벡터 필드 $ s_t(x, \theta) $ 와 밀도 경로 $ \rho_t(x, \eta) $ 를 매개변수화하고, 국소 제약 조건을 충족시키기 위해 이중 목표로 훈련한다.
• 직접 궤적 전개 또는 커플링 추정 없이도 엔드 투 엔드 훈련이 가능한 이중 목표를 통해 최적화를 수행한다.
• 이미지 생성을 위해, 영속 가속도가 0임을 의미하는 하밀토니안-자코비 조건을 통해 단일 스텝 샘플링이 가능하며, 이는 $ X_1 = X_0 + \nabla s_0(X_0) $ 를 의미한다.
• 에너지 항을 조합하여 다양한 구성 요소를 지원함으로써, 특정 생물학적 또는 물리적 제약 조건에 대한 인덕티브 바이어스를 제공할 수 있다.

실험 결과

연구 질문

• RQ1 통합적인 딥러닝 프레임워크가 비균형 OT, 스체로딩어 브리지, 물리적으로 제약된 OT와 같은 다양한 최적 운반 문제를 단일 공식으로 해결할 수 있는가?
• RQ2 세포 역학에 대한 사전 지식를 변분 프레임워크 내에서 어떻게 민감하게 통합할 수 있는가?
• RQ3 라그랑주 작용의 이중 최적화가 궤적 역전파의 계산 부담을 피하면서도 정확도를 유지할 수 있는가?
• RQ4 기존 방법에 비해 이 프레임워크가 단일 세포 RNA-seq 데이터에서 궤적 추론 성능을 얼마나 향상시키는가?
• RQ5 이중 최적화를 통해 반복적인 ODE 해법 없이도 최적 운반 맵을 이용한 단일 스텝 이미지 생성이 가능한가?

주요 결과

• 제안된 프레임워크는 특히 맞춤형 라그랑주를 통해 생물학적 사전 지식를 통합할 때 기존 방법보다 뛰어난 성능을 보였다.
• Embryoid Body, CITE-seq, Multiome 데이터를 포함한 저차원 및 고차원 단일 세포 데이터셋에서 최신 기술 성능을 달성했다.
• 이미지 생성 분야에서, 하밀토니안-자코비 조건을 통해 단일 스텝 샘플링이 가능했으며, 단 한 번의 함수 평가로 반복적 ODE 해법의 품질에 도달했다.
• 이중 최적화 접근법은 궤적 역전파 및 커플링 추정을 피함으로써 계산 비용을 크게 감소시켰다.
• 라그랑주 구성 요소를 조정함으로써, 비균형 OT 및 스체로딩어 브리지와 같은 다양한 최적 운반 변종에 대해 프레임워크가 성공적으로 일반화되었다.
• MNIST 및 CIFAR-10에서의 실험 결과, 학습된 벡터 필드를 통한 단일 스텝 생성이 반복적 확산 모델과 비교해도 높은 품질의 샘플을 생성하는 데 성공했다.