[논문 리뷰] Scalable Inference in SDEs by Direct Matching of the Fokker-Planck-Kolmogorov Equation

Arno Solin, Ella Tamir외 1명|arXiv (Cornell University)|2021. 10. 29.

Gaussian Processes and Bayesian Inference인용 수 4

한 줄 요약

이 논문은 스토케스틱 미분 방정식(SDEs)에 대한 확장 가능하고 샘플링이 불필요한 추론 방법을 제안한다. 이는 Fokker-Planck-Kolmogorov(FPK) 방정식의 첫 번째 두 모멘트(평균과 공분산)를 직접 일치시키는 방식으로, 가우시안 가정 밀도 근사법을 사용한다. SDE 해를 모멘트 상의 결정론적 상미분방정식(ODE)으로 공식화함으로써, Euler–Maruyama와 같은 샘플링 기반 방법에 비해 수십만 배 빠른 속도 향상을 이룩하면서도, 시계열 및 생성 모델링 과제에서 경쟁 가능한 성능을 유지를 한다.

ABSTRACT

Simulation-based techniques such as variants of stochastic Runge-Kutta are the de facto approach for inference with stochastic differential equations (SDEs) in machine learning. These methods are general-purpose and used with parametric and non-parametric models, and neural SDEs. Stochastic Runge-Kutta relies on the use of sampling schemes that can be inefficient in high dimensions. We address this issue by revisiting the classical SDE literature and derive direct approximations to the (typically intractable) Fokker-Planck-Kolmogorov equation by matching moments. We show how this workflow is fast, scales to high-dimensional latent spaces, and is applicable to scarce-data applications, where a non-parametric SDE with a driving Gaussian process velocity field specifies the model.

연구 동기 및 목표

• 샘플링 기반 추론의 계산 비효율성, 특히 고차원 및 데이터가 부족한 환경에서의 문제를 해결하기 위해.
• 해결이 불가능한 Fokker-Planck-Kolmogorov(FPK) 방정식의 해를 직접적으로, 샘플링 없이 근사화하는 방법을 개발하기 위해.
• 잠재 공간 SDE 모델에서 빠르고 확장 가능한 추론을 가능하게 하기 위해, 시간에 따른 주변 분포에 대해 가우시안 가정을 활용하기 위해.
• 드라이프에 GP 사전분포를 결합하고 모멘트 일치 추론을 사용함으로써, SDE의 머신러닝 응용 가능성을 확장하기 위해.
• 경로 기반 샘플링을 피하면서도 분포 정확도를 유지하는 결정론적 도구로서, 확률적 룬게-쿠타 방법의 대안을 제공하기 위해.

제안 방법

• 잠재 역학을 민감하게 모델링할 수 있도록, 드라이프 벡터 필드에 가우시안 프로세스 사전분포를 도입하여 SDE를 공식화한다.
• 해의 분포의 첫 번째 두 모멘트(평균과 공분산)를 일치시킴으로써, Fokker-Planck-Kolmogorov(FPK) 방정식에 대한 닫힌 형태 근사식을 유도한다.
• 유도된 모멘트 동역학을 결정론적 상미분방정식(ODE)으로 풀어, 궤적 샘플링이 필요 없도록 한다.
• 시간에 따른 주변 분포를 표현하기 위해 가우시안 가정 밀도 근사를 사용하여, 고차원 잠재 공간에서의 효율적 계산을 가능하게 한다.
• 변분 추론 프레임워크를 사용하여 모델을 구현하며, 인코더는 초기 모멘트를 출력하고 디코더는 모멘트 일치 SDE 동역학을 사용한다.
• 자동 미분를 통해 모멘트 ODE를 통해 기울기를 계산함으로써, 가능도 최대화 또는 변분 목표를 최적화한다.

실험 결과

연구 질문

• RQ1 Fokker-Planck-Kolmogorov 방정식의 직접적 모멘트 일치가, SDE에서 샘플링 기반 추론의 확장 가능한 대안이 될 수 있는가?
• RQ2 고차원 및 저데이터 환경에서, 모멘트 일치 추론은 Euler–Maruyama 샘플링에 비해 정확도와 속도 면에서 어떻게 비교되는가?
• RQ3 비선형 드라이프와 분산을 가진 SDE의 진정한 시간 주변 분포를 가우시안 가정 밀도 근사가 어느 정도 잘 포괄할 수 있는가?
• RQ4 GP 사전분포를 가진 SDE와 모멘트 일치를 조합하면, 데이터가 부족한 응용 분야에서 효과적인 모델링이 가능한가?
• RQ5 제안된 방법은 계산 비용을 줄이면서도, 시계열 및 생성 모델링 벤치마크에서 경쟁 가능한 성능을 유지하는가?

주요 결과

• 동일한 아키텍처를 사용할 때, 모멘트 일치 방법은 MOCAP 데이터셋에서 테스트 MSE 7.55±0.05를 기록하여 기준선 잠재 SDE(4.03±0.20)를 능가했으며, 이는 가정 밀도 제약 조건 하에서도 강건함을 시사한다.
• 벽시계 시간 측정 결과, 단일 평가에 대해 GPU에서는 6.0±0.1초, CPU에서는 1.8±0.1초가 소요되었고, 200개의 경로를 사용한 Euler–Maruyama는 37.1±0.1초가 소요되었다.
• Rotating MNIST 작업에서, 모멘트 일치 방법은 MSE 0.051±0.007과 NLPD 52.7±9.5를 기록하여, 궤적 샘플링을 피하면서도 경쟁 가능한 성능을 입증했다.
• 이 방법은 고차원 잠재 공간으로도 효율적으로 확장되었으며, 단계당 차원 수에 비례해 계산 비용이 선형적으로 증가(O(d))하나, 샘플링 방법은 차원 증가에 따라 성능이 급격히 떨어지는 경향을 보였다.
• 샘플링 기반 SDE 솔버에 비해 추론 시간을 수십만 배 단위로 감소시켜, 실시간 및 대규모 응용 분야에서 실용 가능성을 입증했다.
• 이 방법은 데이터가 부족한 환경에서도 효과적임을 입증했으며, 드라이프 필드에 대한 GP 사전분포가 강력한 인덕티브 바이어스를 제공해 일반화 성능을 향상시켰다.