[논문 리뷰] Riemannian Neural Optimal Transport

Alessio Micheli, Yueqi Cao외 2명|arXiv (Cornell University)|2026. 02. 03.

3D Shape Modeling and Analysis인용 수 0

한 줄 요약

리만 기하 다양체 위에서 OT 맵을 직접 학습하기 위한 연속 신경망 프레임워크인 Riemannian Neural OT(RNOT)를 소개합니다. 이는 이산화를 피하고 차원의 저주를 차원에 따른 다항 복잡도로 깨뜨립니다.

ABSTRACT

Computational optimal transport (OT) offers a principled framework for generative modeling. Neural OT methods, which use neural networks to learn an OT map (or potential) from data in an amortized way, can be evaluated out of sample after training, but existing approaches are tailored to Euclidean geometry. Extending neural OT to high-dimensional Riemannian manifolds remains an open challenge. In this paper, we prove that any method for OT on manifolds that produces discrete approximations of transport maps necessarily suffers from the curse of dimensionality: achieving a fixed accuracy requires a number of parameters that grows exponentially with the manifold dimension. Motivated by this limitation, we introduce Riemannian Neural OT (RNOT) maps, which are continuous neural-network parameterizations of OT maps on manifolds that avoid discretization and incorporate geometric structure by construction. Under mild regularity assumptions, we prove that RNOT maps approximate Riemannian OT maps with sub-exponential complexity in the dimension. Experiments on synthetic and real datasets demonstrate improved scalability and competitive performance relative to discretization-based baselines.

연구 동기 및 목표

• 다양체에서 OT의 필요성을 동기부여하고 이산적 접근의 한계(curse of dimensionality)를 확인한다.
• 리만 기하 다양체에서 직접 작동하는 연속적이고 내재된 신경 OT 프레임워크를 제안한다.
• RNOT의 근사 품질과 복잡도에 대한 이론적 보장을 제시한다.
• 합성 및 실제 다양체-valued 데이터에서 확장성 및 경쟁력 있는 성능을 입증한다.

제안 방법

• OT 포텐셜을 암시적 c-concave 클래스에서 c-transform을 통해 구성하여 OT 구조를 구성상으로 강제한다.
• 다양체 특징 맵에서 작용하는 신경망으로 포텐셜을 매개변수화하고 리만지수 지수함수: T(x)=exp_x(-∇φ(x))를 통해 운송 맵을 복구한다.
• Assumption 2.2를 만족하는 injective 특징 맵 φ를 만들기 위해 Gromov 거리-랜드마크 임베딩을 사용한다.
• universality를 증명: 암시적 c-concave 클래스 내의 포텐셜 근사를 통해 유도 맵의 수렴으로 OT 맵이 실제 맵으로 수렴한다.
• 완만한 규칙성(ε-정확도) 하에서 OT 포텐셜 및 맵을 근사하기 위한 뉴럴 네트워크 폭과 깊이에 대해 ε^{-1} 다항식 경계 제공.
• Jacobian determinants 없이 Kantorovich 부분 이중 목표로 엔드투엔드 학습한다.

(a) Source $\mu$ (150 Ma)

실험 결과

연구 질문

• RQ1 연속적이고 고유한 신경 OT 맵이 이산화 없이도 정확한 운송을 달성할 수 있는가?
• RQ2 신경망으로 학습할 때 리만 OT 맵의 복잡도 및 근사 보장은 무엇인가?
• RQ3 RNOT 포텐셜은 샘플 외 일반화 및 차원에 따라 다양체에서 다항적으로 확장 가능한가?
• RQ4 합성 및 실제 다양체 값 데이터에서 RNOT의 성능은 이산화 기반 비교 기반선과 비교해 어떤가?

주요 결과

모델	KL	ESS	시간 (초)
Ours (FPS)	0.03 a 0.00	0.97 a 0.00	986 a 2
Ours (RND)	0.04 a 0.00	0.95 a 0.00	1100 a 1
RCPM	0.0037 a 0.0008	0.996 a 0.000	37.3 a 0.6
RCNF	2.38 a 0.10	0.48 a 0.02	352 a 6
Moser Flow	1.16 a 0.03	0.82 a 0.00	758 a 16

• 다양체에서의 이산 출력 OT 방법은 차원의 저주 때문에: 고정된 정확도는 차원에 따라 지수적 매개변수 증가를 필요로 한다.
• RNOT은 암시적 c-concavity가 내재된 연속적 접근을 제공하여 야코비를 통해 샘플 외 생성 가능.
• 이론적 결과는 완만한 규칙성 하에서 OT 포텐셜과 맵 근사에 대해 ε^{-1}의 다항적 복잡도를 보인다(지수적이 아닌).
• S^2와 T^2에서의 실험은 경쟁력 있는 KL 발산 및 ESS를 보이고, 차원에서 이산화 기반 baselines보다 더 나은 확장성을 보인다.
• S^2에서의 대륙 이동 OT는 지구물리학 직관과 맞물리는 해석 가능한 기하적 운송을 보여준다.

(b) Target $\nu$ (Present)