[논문 리뷰] Learning High Dimensional Wasserstein Geodesics

Liu Shu, Shaojun Ma외 3명|arXiv (Cornell University)|2021. 02. 05.

Seismic Imaging and Inversion Techniques참고 문헌 42인용 수 5

한 줄 요약

이 논문은 라그랑주 승수와 카우시-쿠른-터커(KKT) 조건을 통해 최적 운반 문제를 미니맥심 경사 문제로 재구성함으로써, 딥 뉴럴 네트워크를 사용해 고차원 워셔스타인 기하선을 계산하는 딥 러닝 기반 방법을 제안한다. 이 방법은 운반 지도와 속도장을 딥 뉴럴 네트워크로 매개변수화하여, 명시적 정규화나 차원 축소 없이도 샘플 기반, 이중 방향 학습을 가능하게 하며, 워셔스타인 거리, 최적 지도, 기하선 경로를 동시에 추정한다.

ABSTRACT

We propose a new formulation and learning strategy for computing the Wasserstein geodesic between two probability distributions in high dimensions. By applying the method of Lagrange multipliers to the dynamic formulation of the optimal transport (OT) problem, we derive a minimax problem whose saddle point is the Wasserstein geodesic. We then parametrize the functions by deep neural networks and design a sample based bidirectional learning algorithm for training. The trained networks enable sampling from the Wasserstein geodesic. As by-products, the algorithm also computes the Wasserstein distance and OT map between the marginal distributions. We demonstrate the performance of our algorithms through a series of experiments with both synthetic and realistic data.

연구 동기 및 목표

• 기존의 PDE 기반 방법이 차원의 저주에 시달리는 고차원 공간에서 워셔스타인 기하선을 계산하는 데 도전하는 문제를 해결한다.
• 공간 이산화와 정규화를 피하는 스케일러블하고 샘플 기반의 알고리즘을 개발하여 고차원 환경에서의 적용을 가능하게 한다.
• 통합된 프레임워크를 통해 분포 간 워셔스타인 거리, 최적 운반 지도, 기하선 경로를 동시에 계산한다.
• 기하선을 따라 자연스러운 샘플링 메커니즘을 제공함으로써 머신러닝, 로봇공학, 제어 분야의 실용적 응용을 가능하게 한다.
• 거리나 지도 추정에만 집중하는 기존의 운반 이론(OT) 방법의 한계를 극복하고, 전체 기하선 경로를 함께 고려한다.

제안 방법

• 라그랑주 승수를 사용해 동적 최적 운반 문제를 미니맥심 문제로 재구성함으로써, 사다리꼴 문제로 변환한다.
• KKT 조건을 적용해 검색 공간을 축소하고 최적화를 단순화함으로써, 더 다루기 쉬운 형태로 전환한다.
• 속도장과 이중 잠재함수를 딥 뉴럴 네트워크로 매개변수화하여 고차원에서의 민첩한 함수 근사가 가능하도록 한다.
• 앞서와 뒤로의 운반 경로를 번갈아 최적화하는 이중 방향, 샘플 기반 학습 알고리즘을 설계한다.
• 예비 조정(preconditioning)과 이중 방향 역학을 통합해 고차원 공간에서의 학습 안정성과 수렴성을 향상시킨다.
• 미니배치 샘플링을 사용해 네트워크를 종단 간(end-to-end)으로 학습시켜 스케일러블한 추론과 기하선 샘플링을 가능하게 하되, 명시적 정규화 제약 조건 없이 수행한다.

실험 결과

연구 질문

• RQ1 공간 이산화나 명시적 정규화에 의존하지 않고 딥 러닝 기반 접근이 고차원 워셔스타인 기하선을 계산할 수 있는가?
• RQ2 기하선 계산 과정에서 최적 운반 지도와 워셔스타인 거리를 동시에 추정할 수 있는가?
• RQ3 제안된 이중 방향 학습 전략이 고차원 OT 문제에서 안정성과 수렴성에 얼마나 기여하는가?
• RQ4 이 방법은 기하선 경로를 따라 실제적인 중간 샘플을 생성할 수 있는가? 이는 이미지 간 변환 및 데이터 생성 응용에 기여하는가?
• RQ5 이 방법은 복잡한 고차원 분포(예: 이미지, 손글씨 숫자)를 포함한 실제 데이터에서 어떻게 성능을 발휘하는가?

주요 결과

• 제안된 방법은 차원 축소 없이도 고차원 공간(예: 28×28 MNIST 숫자, 256×256 이미지)에서 워셔스타인 기하선을 성공적으로 계산한다.
• 이미지 색상 이동 및 숫자 변환 작업에서 확인된 바와 같이, 원천 및 목표 분포 사이의 시각적으로 현실적인 중간 샘플을 생성한다.
• 이미지 번역 작업(예: 여름에서 가을 숲 풍경)에서 생성된 이미지는 지표와 유사하며, 색상 분포에서 높은 유사성을 보였다.
• 합성 및 실제 데이터 세트에서 POT 라이브러리와의 비교를 통해, 워셔스타인 거리와 최적 운반 지도를 높은 정확도로 계산한다.
• 이중 방향 학습 전략은 안정성과 수렴성을 향상시켜 복잡한 고차원 환경에서도 신뢰할 수 있는 기하선 샘플링을 가능하게 한다.
• 일반적인 볼록 비용 함수(예: Lp 노름 포함)에 대해 샘플 기반 기하선 계산에서 최신 기술 성능을 달성하며, 정방향 및 역방향 운반을 모두 지원한다.