[논문 리뷰] Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein

Marco Cuturi, Laetitia Meng-Papaxanthos외 4명|arXiv (Cornell University)|2022. 01. 28.

Asphalt Pavement Performance EvaluationEngineering인용 수 18

한 줄 요약

OTT-JAX는 엔트로피 정규화와 낮은 랭크 근사치를 사용하여 효율적이고 미분 가능한 최적 운반 계산을 가능하게 하는 JAX 기반 파이썬 도구상자입니다. 선형 및 이차 OT 문제, 바리센터, 그로모프-워서스타인, 가우시안 혼합 매칭을 지원하며, 확장 가능한 솔버와 기계학습 응용을 위한 자동 미분을 제공합니다.

ABSTRACT

Optimal transport tools (OTT-JAX) is a Python toolbox that can solve optimal transport problems between point clouds and histograms. The toolbox builds on various JAX features, such as automatic and custom reverse mode differentiation, vectorization, just-in-time compilation and accelerators support. The toolbox covers elementary computations, such as the resolution of the regularized OT problem, and more advanced extensions, such as barycenters, Gromov-Wasserstein, low-rank solvers, estimation of convex maps, differentiable generalizations of quantiles and ranks, and approximate OT between Gaussian mixtures. The toolbox code is available at exttt{https://github.com/ott-jax/ott}

연구 동기 및 목표

• 대규모 및 기계학습 응용을 위한 최적 운반(Optimal Transport, OT)의 계산 및 미분 가능성 도전 과제를 해결하기 위해.
• 점군, 히스토그램, 측도 간의 정규화된 OT 문제를 해결하기 위한 통합적이고 고성능 프레임워크를 제공하기 위해.
• JAX의 자동 미분과 JIT 컴파일을 통해 최적 운반 계산의 미분 가능성을 보장하고 딥러닝 파이프라인 내에서 엔드 투 엔드 학습을 가능하게 하기 위해.
• 표준 워서스타인 거리 이외의 OT 기능을 확장하여 바리센터, 그로모프-워서스타인, 소프트 정렬 연산을 포함하기 위해.
• 명시적 행렬 저장 없이 저랭크 근사치와 기하학적 감지 비용 계산을 통해 효율적인 계산을 지원하기 위해.

제안 방법

• JAX의 자동 미분과 JIT 컴파일을 활용하여 CPU 및 TPU/GPU에서 고성능의 미분 가능한 OT 솔버를 구현합니다.
• Sinkhorn 알고리즘을 통한 엔트로피 정규화를 구현하여 최적 운반 계획을 부드럽게 하고 효율적이고 미분 가능한 최적화를 가능하게 합니다.
• 저랭크 Sinkhorn 솔버를 도입하여 운반 행렬을 랭크-r 요소로 근사화함으로써 메모리와 계산 비용을 절감합니다.
• 기하학 클래스를 사용하여 비용 행렬을 암시적으로 계산하고 명시적 저장을 피합니다. 예를 들어, 커널 기반 연산 또는 격자 기반 구조를 통해 점군에 대해 계산합니다.
• 반복 선형화를 통한 그로모프-워서스타인 지원으로, 이차 OT 문제를 선형 OT 문제의 시퀀스로 줄입니다.
• 입력-볼록 신경망(ICNN)을 통합하여 볼록 매핑과 기계학습 기반 재정의를 통한 소프트 정렬 연산을 학습합니다.

실험 결과

연구 질문

• RQ1 현대 딥러닝 프레임워크를 활용하여 대규모에서 미분 가능한 최적 운반 문제를 어떻게 효율적으로 해결할 수 있는가?
• RQ2 저랭크 근사치와 기하학적 감지 계산은 최적 운반에서 메모리와 시간 복잡도를 얼마나 줄일 수 있는가?
• RQ3 미분 가능한 OT를 사용하여 바리센터나 분포 간 매핑과 같은 구조적 표현을 학습할 수 있는가?
• RQ4 최적 운반 계획을 통한 암시적 미분은 기계학습 파이프라인의 학습 안정성에 어떻게 기여하는가?
• RQ5 그로모프-워서스타인과 소프트 정렬과 같은 고급 OT 변형은 완전한 미분 가능성과 확장성과 함께 효율적으로 구현될 수 있는가?

주요 결과

• OTT-JAX는 JAX의 자동 미분를 완전히 지원하여 운반 계획을 통한 역전파를 가능하게 합니다.
• 저랭크 Sinkhorn 솔버는 정확도를 희생시키지 않으면서도 대규모 문제에서 뚜렷한 메모리 및 계산 절감 효과를 보입니다.
• 도구상자는 바리센터, 그로모프-워서스타인 거리, 소프트 정렬된 배열의 엔드 투 엔드 미분 가능한 계산을 지원합니다.
• 기하학 클래스를 통해 비용 행렬(예: 점군 또는 격자)을 암시적으로 계산할 수 있어 명시적 저장을 피하고 격자에서는 O(dn^{d+1}) 연산을 가능하게 합니다.
• Delon과 Desolneux(2020)의 기여를 기반으로 한 미분 가능한 근사치를 사용하여 가우시안 혼합 분포 간의 워서스타인 유사도 거리 계산을 효율적으로 지원합니다.
• 도구상자는 생산용으로 사용 가능하며, 의료 영상 데이터에서의 등온 바리센터 계산과 복잡한 다양체(예: 나선형, 스위스 롤) 간의 형태 매칭 등 고급 응용에 활용되고 있습니다.