[논문 리뷰] Large-Scale Wasserstein Gradient Flows

Petr Mokrov, Alexander Korotin외 4명|arXiv (Cornell University)|2021. 06. 01.

Advanced Neuroimaging Techniques and Applications참고 문헌 59인용 수 10

한 줄 요약

이 논문은 입력-볼록 신경망(ICNN)을 사용하여 JKO 스킴 내 최적 운반 지도를 매개변수화함으로써 대규모 워셔스타인 기울기 흐름을 스케일러블하게 다루는 방법을 제안한다. 이는 도메인 이산화나 입자 시뮬레이션 없이도 효율적인 샘플링과 밀도 추정을 가능하게 하며, 확률적 기울기 하강법을 통해 정확한 확산 모델링과 비정규화된 밀도 샘플링을 달성한다. 이 방법은 포커-플랭크 역학, 베이지안 추론 및 비선형 필터링에서 성공적으로 적용되었다.

ABSTRACT

Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated to each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate our algorithm's performance by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.

연구 동기 및 목표

• 고차원 확산 과정에서 워셔스타인 기울기 흐름의 JKO 단계를 풀이하는 데 있어 계산적으로 비가능한 문제를 해결하기 위해.
• 샘플과 밀도에 모두 접근 가능한 조건에서 시간에 따라 변화하는 확률 측도의 연속적이고 미분 가능한 근사화를 가능하게 하기 위해.
• 도메인 이산화나 입자 기반 방법의 한계를 극복하기 위해 ICNN를 통해 매개변수화된 운반 지도를 사용하기 위해.
• 비정규화된 사후 샘플링 및 비선형 필터링과 같은 머신러닝 응용 분야에서 스케일러블하고 미분 가능한 동역학을 지원하기 위해.
• 브레니에의 정리와 ICNN를 활용하여 고비용 최적 운반 계산을 피하는 일반적인 워셔스타인 기울기 흐름 해법 프레임워크를 제공하기 위해.

제안 방법

• 최적 운반 지도를 볼록 잠재함수의 기울기로 표현하기 위해 브레니에의 정리를 활용하며, 이를 입력-볼록 신경망(ICNN)으로 매개변수화한다.
• 워셔스타인 공간 내 정규화된 자유 에너지 기능을 최소화함으로써 각 JKO 단계를 최적화하기 위해 확률적 기울기 하강법(SGD)을 사용한다.
• 시간 이산화된 포커-플랭크 동역학을 연속적인 ICNN 기반 운반 지도로 매개변수화함으로써 측도의 진화를 순차적으로 학습할 수 있도록 한다.
• ICNN 기반 지도의 역행성 구조를 활용하여 진화하는 측도 $\rho_t$ 에서의 샘플링과 밀도 추정을 모두 가능하게 한다.
• 워셔스타인 거리의 명시적 계산을 피하기 위해 ICNN의 기울기 구조와 볼록성을 활용하여 잘 정의된 최적화를 보장한다.
• 이중 표현을 통한 효율적 최적화를 가능하게 하기 위해 JKO 단계에 엔트로피 정규화를 적용한다.

실험 결과

연구 질문

• RQ1 입력-볼록 신경망(ICNN)을 사용하여 고차원 공간에서 워셔스타인 기울기 흐름의 JKO 단계를 효율적이고 스케일러블하게 근사화할 수 있는가?
• RQ2 도메인 이산화나 입자 시뮬레이션 없이도 진화하는 측도 $\rho_t$ 에서 동시에 샘플링과 밀도 추정을 수행할 수 있는가?
• RQ3 ICNN 기반 매개변수화 방식이 워셔스타인 기울기 흐름의 기하학적 구조를 얼마나 잘 유지하면서도 확률적 최적화에 적합한가?
• RQ4 이 방법은 머신러닝 분야의 비정규화된 사후 샘플링 및 비선형 필터링 작업에 효과적으로 적용될 수 있는가?
• RQ5 전통적인 최적 운반 해법기나 입자 기반 방법과 비교할 때, JKO 단계에 ICNN를 사용하는 데서 발생하는 계산적 및 통계적 트레이드오프는 무엇인가?

주요 결과

• 제안된 ICNN 기반 JKO 방법은 도메인 이산화나 입자 시뮬레이션 없이도 진화하는 측도 $\rho_t$ 에서 효율적인 샘플링과 밀도 추정을 가능하게 한다.
• SGD를 통한 순차적 운반 지도 학습을 통해 포커-플랭크 역학을 정확하게 근사화하며, 명시적 워셔스타인 거리 계산이 필요로 하지 않는다.
• 이 방법은 비정규화된 사후 샘플링 및 비선형 필터링 작업을 지원하여 베이지안 머신러닝 분야에서 실용적인 유용성을 입증한다.
• ICNN의 사용을 통해 측도 진화의 연속적이고 미분 가능한 모델링이 가능해져 매 시간 단계에서 생성과 밀도 평가가 모두 가능해진다.
• 이 방법은 고차원으로 확장 가능하며 격자 기반 이산화의 지수적 비용을 피하기 때문에 복잡한 실제 응용의 확산 과정에 적합하다.
• 분자 발견 및 인구 역학에 대한 초보적 결과는 더 넓은 적용 가능성을 시사하며, MOSES 분자 데이터셋에서 유망한 성능을 보였다.