ICML 2024
Conference paper

Quantum Theory and Application of Contextual Optimal Transport

View code


Optimal Transport (OT) has fueled machine learning (ML) applications across various domains. In cases where paired data measurements (µ, ν) are coupled to a context variable p_i , one may aspire to learn a global transportation map, parameterized through the context to facilitate prediction of transport plans even from unseen context. Existing approaches for this task leverage Brenier’s theorem and utilize Neural OT. Here, we develop a quantum computing formulation for learning transportation plans parameterized by a context variable. This is achieved through exploiting a natural link between doubly stochastic matrices and unitary operators which relates to recent results in quantum learning theory suggesting intrinsic advantages in modelling constrained problems with quantum. We verify our methodology on synthetic and real data, by predicting variations in cell type distributions parameterized through drug dosage as context. Our comparisons to several baselines reveal that our method can capture dose-induced variations in cell distributions, even to some extent when dosages are extrapolated and sometimes with performance similar to the best classical models. In summary, this is a first step toward learning to predict contextualized transportation plans through quantum.