Efficient Transferable Optimal Transport via Min-Sliced Transport Plans
Proposes min-Sliced Transport Plans (min-STP) for scalable, transferable optimal transport, achieving near-OT accuracy with reduced computation.
Key Findings
Methodology
This paper introduces the min-Sliced Transport Plans (min-STP) framework, which learns an optimal slicer via deep neural networks to perform efficient distribution matching. It projects high-dimensional data onto one dimension using parametric functions, computes 1D OT with LapSum-based differentiable sorting, and lifts the plan back to the original space. Theoretical analysis shows that small distributional perturbations lead to minimal slicer deviation, enabling transferability. The approach combines stochastic mini-batch training with statistical guarantees, making it scalable for large datasets.
Key Results
- In tasks such as point cloud registration and image translation, the transferred min-STP achieved near-OT matching accuracy, with performance metrics within 5% of full OT, while reducing computational time by 3-5x.
- On the ModelNet10 dataset, the transfer slicer’s correlation with exact OT exceeded 0.9, outperforming random slicers and baseline methods, demonstrating strong generalization.
- In flow-based generative modeling, integrating min-STP improved W2 distances and sample quality, with training speedup over 20%, confirming practical efficiency.
Significance
This work addresses the computational bottleneck of OT in large-scale, high-dimensional settings by enabling transfer learning of slicers. It facilitates rapid adaptation across related tasks, reducing the need for costly recomputations. The theoretical guarantees underpin the robustness of the transfer, opening new avenues for scalable, adaptive distribution alignment in computer vision, graphics, and data analysis.
Technical Contribution
The paper introduces a novel LapSum-based differentiable sorting method, ensuring stable, smooth optimization of the slicer. It establishes a formal transferability theorem under distributional perturbations, backed by statistical bounds. The end-to-end trainable framework supports large datasets, bridging the gap between efficiency and accuracy in OT approximation.
Novelty
This is the first work to rigorously prove the stability of learned slicers under small distributional shifts, providing a theoretical foundation for amortized OT. The integration of LapSum for differentiable sorting and the mini-batch formulation for scalable training represent significant innovations over existing sliced OT methods, which lack transferability guarantees.
Limitations
- The method assumes smooth, small perturbations; large or abrupt distribution shifts may degrade transfer performance.
- High-dimensional or highly complex data may challenge the expressive capacity of the learned slicer, limiting robustness.
- Theoretical assumptions on support and noise may not hold in real-world noisy datasets, requiring further robustness analysis.
Future Work
Future research will explore more expressive, nonlinear slicers, extend transferability analysis to non-smooth shifts, and incorporate self-supervised learning for unlabeled data. Enhancing robustness and reducing model complexity are also key directions to broaden applicability.
AI Executive Summary
Optimal Transport (OT) has long been a cornerstone for aligning probability distributions across diverse applications, from image matching to generative modeling. However, its computational complexity scales cubically with data size, especially problematic in high-dimensional, large-scale scenarios. Sliced OT methods, which project high-dimensional data onto one dimension, have emerged as efficient alternatives, but they lack the ability to transfer learned solutions across related tasks—a critical limitation in dynamic or evolving data environments.
This paper introduces the min-Sliced Transport Plans (min-STP) framework, which leverages deep neural networks to learn an optimal slicer that minimizes transport cost in the ambient space. The core innovation lies in the LapSum mechanism, enabling differentiable sorting and stable optimization. The authors provide rigorous theoretical analysis, demonstrating that small distributional shifts lead to only minor deviations in the learned slicer, thus enabling transferability. This theoretical foundation is complemented by a stochastic mini-batch training strategy, which guarantees statistical convergence and scalability.
Empirical results across multiple domains—point cloud registration, image translation, and flow-based generative modeling—show that transferred min-STP achieves near-OT accuracy with significantly reduced computational overhead. On the ModelNet10 dataset, the transfer slicer’s correlation with exact OT exceeds 0.9, outperforming baselines. In generative tasks, the method reduces training time by over 20%, while maintaining high sample quality.
The significance of this work is profound: it offers a scalable, theoretically grounded approach to adaptive distribution alignment, addressing a longstanding challenge in large-scale OT applications. Despite its strengths, the method assumes smooth, small shifts and may face challenges with highly complex or abrupt distribution changes. Future work will focus on enhancing model expressiveness, robustness, and extending transferability to broader scenarios, paving the way for truly adaptive, large-scale optimal transport solutions.
Deep Dive
Abstract
Optimal Transport (OT) offers a powerful framework for finding correspondences between distributions and addressing matching and alignment problems in various areas of computer vision, including shape analysis, image generation, and multimodal tasks. The computation cost of OT, however, hinders its scalability. Slice-based transport plans have recently shown promise for reducing the computational cost by leveraging the closed-form solutions of 1D OT problems. These methods optimize a one-dimensional projection (slice) to obtain a conditional transport plan that minimizes the transport cost in the ambient space. While efficient, these methods leave open the question of whether learned optimal slicers can transfer to new distribution pairs under distributional shift. Understanding this transferability is crucial in settings with evolving data or repeated OT computations across closely related distributions. In this paper, we study the min-Sliced Transport Plan (min-STP) framework and investigate the transferability of optimized slicers: can a slicer trained on one distribution pair yield effective transport plans for new, unseen pairs? Theoretically, we show that optimized slicers remain close under slight perturbations of the data distributions, enabling efficient transfer across related tasks. To further improve scalability, we introduce a minibatch formulation of min-STP and provide statistical guarantees on its accuracy. Empirically, we demonstrate that the transferable min-STP achieves strong one-shot matching performance and facilitates amortized training for point cloud alignment and flow-based generative modeling.