Displacement-Sparse Neural Optimal Transport

TL;DR

This paper introduces a new method for learning displacement-sparse maps in neural optimal transport, improving interpretability and efficiency in high-dimensional biological applications.

Contribution

It proposes a theoretically grounded approach with a novel smoothed $ ext{l}_0$ regularizer and an adaptive sparsity control framework for neural OT.

Findings

Outperforms $ ext{l}_1$ regularization in experiments.

Enhances interpretability of neural OT maps.

Improves accuracy and usability in large-scale applications.

Abstract

Optimal transport (OT) aims to find a map $T$ that transports mass from one probability measure to another while minimizing a cost function. Recently, neural OT solvers have gained popularity in high dimensional biological applications such as drug perturbation, due to their superior computational and memory efficiency compared to traditional exact Sinkhorn solvers. However, the overly complex high dimensional maps learned by neural OT solvers often suffer from poor interpretability. Prior work addressed this issue in the context of exact OT solvers by introducing \emph{displacement-sparse maps} via designed elastic cost, but such method failed to be applied to neural OT settings. In this work, we propose an intuitive and theoretically grounded approach to learning \emph{displacement-sparse maps} within neural OT solvers. Building on our new formulation, we introduce a novel smoothed $\ell_0$ regularizer that outperforms the $\ell_1$ based alternative from prior work. Leveraging Input Convex Neural Network's flexibility, we further develop a heuristic framework for adaptively controlling sparsity intensity, an approach uniquely enabled by the neural OT paradigm. We demonstrate the necessity of this adaptive framework in large-scale, high-dimensional training, showing not only improved accuracy but also practical ease of use for downstream applications.