Quantum Theory and Application of Contextual Optimal Transport

TL;DR

This work proposes a first-of-its-kind quantum computing formulation for amortized optimization of contextualized transportation plans through quantum computing and exploits a direct link between doubly stochastic matrices and unitary operators thus unravelling a natural connection between OT and quantum computation.

Abstract

Optimal Transport (OT) has fueled machine learning (ML) across many domains. When paired data measurements $(\boldsymbolμ, \boldsymbolν)$ are coupled to covariates, a challenging conditional distribution learning setting arises. Existing approaches for learning a $\textit{global}$ transport map parameterized through a potentially unseen context utilize Neural OT and largely rely on Brenier's theorem. Here, we propose a first-of-its-kind quantum computing formulation for amortized optimization of contextualized transportation plans. We exploit a direct link between doubly stochastic matrices and unitary operators thus unravelling a natural connection between OT and quantum computation. We verify our method (QontOT) on synthetic and real data by predicting variations in cell type distributions conditioned on drug dosage. Importantly we conduct a 24-qubit hardware experiment on a task challenging for classical computers and report a performance that cannot be matched with our classical neural OT approach. In sum, this is a first step toward learning to predict contextualized transportation plans through quantum computing.

Figures (12)

• Figure 1: A) Contextual OT is a conditional distribution learning problem. B) Our proposed ansatz receives the context ($\boldsymbol{p}_k$) and the initial distribution $\boldsymbol{\mu}_i$ and produces a DSM that can be rescaled to a transport plan $T$ with marginal distributions $\boldsymbol{\mu}_i$ and $\boldsymbol{\hat{\nu}}_{i,k}$.
• Figure 2: Circuit structures for the transportation map prediction. The registers $\{i_k\} \cup \{j_k\}$ represent the bits for the index $(i, j)$ related to the entry $Q_{i, j}$ of the resulting DSM. The registry $\{a_k\}$ refers to the $2m$ auxiliary qubits as per \ref{['subsection-qc-birkhoff']}. Regarding \ref{['fig-qc-dsm']}, we remark that the registry $i$ has been added for construction reasons, however in practice it can be removed and substituted with a classical uniform sampling (using the computational basis states on $n$ qubits) over the registry $j$. Consequently, the number of required qubits for DSM-encoding can be reduced to $2m +n$. In \ref{['fig-qc-rsm']}, we have applied that trick to embed transportation maps, with $\ket{i_1} \in \{\ket{0}, \ket{1}\}$, as per \ref{['subsection-tmap-embedding']}.
• Figure 3: Application overview. A population of cells treated with varying drug dosages, resulting in ($X_i, Y_i, \bm{p}_i$) where $X_i$ ($Y_i$) represent scRNA-seq measurements before (after) a drug administered with dosage $\bm{p}_i \in [0, 1]$. We cluster the measurements to identify cell types and compute for each batch the distribution of cell types before and after perturbation, i.e., $\boldsymbol{\mu}$ and $\boldsymbol{\nu}$. A classical OT solver computes the ground truth OT plan $T_i$ based on $\boldsymbol{\mu}_i$, $\boldsymbol{\nu}_i$ (not shown). Given our initial cluster distribution before perturbation $\boldsymbol{\mu}_i$ and the dosage $\bm{p}_i$ our ansatz predicts a transport plan $\bar{T}_i$.
• Figure 4: In a $24$-qubit hardware experiment, the performance of QontOT surpasses a $18$-qubit simulation and various classical neural OT models trained with different hyperparameter settings (cf. \ref{['section-classical-embedding']}). The NeuCOT models of size XS, S, M and L optimize respectively 3k, 8k, 81k and 5M parameters with ADAM compared to $124$ gradient-free optimized parameters in our ansatz.
• Figure A1: Transportation maps. The left and top sequences of blobs represent the initial ($\boldsymbol{\mu}$) and final ($\boldsymbol{\nu}$) distributions. The grid blobs denote the mass displaced from row $i$ to column $j$. The principle of mass preservation manifests as maintaining the total area of initial and final distribution blobs. The left quadrant shows a diagonal transportation (without displacement) so $\boldsymbol{\mu}=\boldsymbol{\nu}$.

Theorems & Definitions (8)

• Lemma 2.1
• Lemma 3.1
• Lemma 3.1
• proof
• Lemma 3.2
• proof
• proof : Proof of Lemma \ref{['lemma-u-to-vec-u']}
• proof : Proof of Lemma \ref{['lemma-transp-plan-recovery']}