Displacement-Sparse Neural Optimal Transport

2025-02-03

AI Summary

Optimal transport (OT) methods often suffer from poor interpretability of learned transport maps and ambiguous displacement structures. Method: This paper proposes a displacement-sparse neural optimal transport framework. Its core innovation is the first formulation of the displacement vector Δ(x) = T(x) − x as an interpretable regularization target, enforced via ℓ₁/ℓ₀ sparsity constraints. We design a dual-path sparsity control strategy—dynamic parameter adaptation in low-dimensional subspaces and explicit dimension-wise sparsity constraints in high dimensions—and integrate input-convex neural networks (ICNNs) to enable minimax optimization of the Wasserstein distance. Results: Evaluated on synthetic single-cell RNA-seq data and real-world 4i cellular perturbation datasets, our method significantly improves transport feasibility and biological interpretability, outperforming state-of-the-art neural OT approaches.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Optimization
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods

Abstract

Optimal Transport (OT) theory seeks to determine the map $T:X o Y$ that transports a source measure $P$ to a target measure $Q$, minimizing the cost $c(mathbf{x}, T(mathbf{x}))$ between $mathbf{x}$ and its image $T(mathbf{x})$. Building upon the Input Convex Neural Network OT solver and incorporating the concept of displacement-sparse maps, we introduce a sparsity penalty into the minimax Wasserstein formulation, promote sparsity in displacement vectors $Delta(mathbf{x}) := T(mathbf{x}) - mathbf{x}$, and enhance the interpretability of the resulting map. However, increasing sparsity often reduces feasibility, causing $T_{#}(P)$ to deviate more significantly from the target measure. In low-dimensional settings, we propose a heuristic framework to balance the trade-off between sparsity and feasibility by dynamically adjusting the sparsity intensity parameter during training. For high-dimensional settings, we directly constrain the dimensionality of displacement vectors by enforcing $dim(Delta(mathbf{x})) leq l$, where $l<d$ for $X subseteq mathbb{R}^d$. Among maps satisfying this constraint, we aim to identify the most feasible one. This goal can be effectively achieved by adapting our low-dimensional heuristic framework without resorting to dimensionality reduction. We validate our method on both synthesized sc-RNA and real 4i cell perturbation datasets, demonstrating improvements over existing methods.

Problem

• Minimize cost in Optimal Transport maps
• Balance sparsity and feasibility trade-off
• Enhance interpretability of transport maps

Innovation

• Sparsity penalty in OT
• Dynamic sparsity adjustment
• Dimensionality constraint enforcement