Computing high-dimensional optimal transport by flow neural networks
Flow-based models are widely used in generative tasks, including normalizing
flow, where a neural network transports from a data distribution $P$ to a
normal distribution. This work develops a flow-based model that transports from
$P$ to an arbitrary $Q$ where both distributions are only accessible via finite
samples. We propose to learn the dynamic optimal transport between $P$ and $Q$
by training a flow neural network. The model is trained to optimally find an
invertible transport map between $P$ and $Q$ by minimizing the transport cost.
The trained optimal transport flow subsequently allows for performing many
downstream tasks, including infinitesimal density ratio estimation (DRE) and
distribution interpolation in the latent space for generative models. The
effectiveness of the proposed model on high-dimensional data is demonstrated by
strong empirical performance on high-dimensional DRE, OT baselines, and
image-to-image translation.
Chen Xu, Xiuyuan Cheng, Yao Xie
# Optimal transport flow and infinitesimal density ratio estimation

###### Abstract

Continuous normalizing flows are widely used in generative tasks, where a flow network transports from a data distribution \(P\) to a normal distribution. A flow model that transports from \(P\) to an arbitrary \(Q\), where both \(P\) and \(Q\) are accessible via finite samples, is of various application interests, particularly in the recently developed telescoping density ratio estimation (DRE) which calls for the construction of intermediate densities to bridge between the two densities. In this work, we propose such a flow by a neural-ODE model which is trained from empirical samples to transport invertibly from \(P\) to \(Q\) (and vice versa) and optimally by minimizing the transport cost. The trained flow model allows us to perform infinitesimal DRE along the time-parametrized log-density by training an additional continuous-time network using classification loss, whose time integration provides a telescopic DRE. The effectiveness of the proposed model is empirically demonstrated on high-dimensional mutual information estimation and energy-based generative models of image data.

## 1 Introduction

The problem of finding a transport map between two general distributions \(P\) and \(Q\) in high dimension is essential in statistics, optimization, and machine learning. When both distributions are only accessible via finite samples, the transport map needs to be learned from data. In spite of the modeling and computational challenges, this setting has applications in many fields. For example, transfer learning in domain adaption aims to obtain a model on the target domain at a lower cost by making use of an existing pre-trained model on the source domain (Courty et al., 2014, 2017), and this can be achieved by transporting the source domain samples to the target domain using the transport map. The (optimal) transport has also been applied to achieve model fairness (Jiang et al., 2020; Silvia et al., 2020). By transporting distributions corresponding to different sensitive attributes to a common distribution, an unfair model is calibrated to match certain desired fairness criteria. The transport map can also be used to provide intermediate interpolating distributions between \(P\) and \(Q\). In density ratio estimation (DRE), this bridging facilitates the so-called "telescopic" DRE (Rhodes et al., 2020) which has been shown to be more accurate when \(P\) and \(Q\) significantly differ.

This work focuses on a continuous-time formulation of the problem where we are to find an invertible transport map \(T_{t}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}\) continuously parametrized by time \(t\in[0,1]\) and satisfying
that \(T_{0}=\text{Id}\) and \((T_{1})_{\#}P=Q\). Here we denote by \(T_{\#}P\) the push-forward of distribution \(P\) by a mapping \(T\), such that \((T_{\#}P)(\cdot)=P(T^{-1}(\cdot))\). Suppose \(P\) and \(Q\) have densities \(p\) and \(q\) respectively in \(\mathbb{R}^{d}\) (we also use the push-forward notation \({}_{\#}\) on densities), the transport map \(T_{t}\) defines

\[\rho(x,t):=(T_{t})_{\#}p,\quad\text{ s.t. }\quad\rho(x,0)=p,\quad\rho(x,1)=q.\]

We will adopt the neural Ordinary Differential Equation (ODE) approach Chen et al. (2018) where we represent \(T_{t}\) as the solution map of an ODE, which is further parametrized by a continuous-time residual network. The resulting map \(T_{t}\) is invertible, and the inversion can be computed by integrating the neural ODE reverse in time. Our model learns the flow from two sets of finite samples from \(P\) and \(Q\). The velocity field in the neural ODE will be optimized to minimize the transport cost so as to approximate the optimal velocity in dynamic optimal transport (OT) formulation, i.e. Benamou-Brenier equation.

The neural-ODE model has been intensively developed in Continuous Normalizing Flows (CNF) Kobyzev et al. (2020). In CNF, the continuous-time flow model, usually parametrized by a neural ODE, transports from a data distribution \(P\) (accessible via finite samples) to a terminal analytical distribution which is typically the normal one \(\mathcal{N}(0,I_{d})\), per the name "normalizing". The study of normalizing flow dated back to non-deep models with statistical applications Tabak and Vanden-Eijnden (2010), and deep CNFs have recently developed into a popular tool for generative models and likelihood inference of high dimensional data. CNF models rely on the analytical expression of the terminal distribution in training. Since our model is also a flow model that transports from data distribution \(P\) to a general (unknown) data distribution \(Q\), both accessible via empirical samples, we name our model "Q-malizing" flow which is inspired by the CNF literature.

After developing a general approach of OT Q-malizing flow model, in the second half of the paper, we focus on the application to telescopic DRE. After training a Q-malizing flow model (the "flow net"), we leverage the intermediate densities \(\rho(x,t)\), which is accessed by finite sampled by pushing the \(P\) samples by \(T_{t}\), to train an additional continuous-time classification network (the "ratio net") over time \(t\in[0,1]\). The ratio net estimates the infinitesimal change of the log-density \(\log\rho(x,t)\) over time, and its time-integral from 0 to 1 yields an estimate of the (log) ratio \(q/p\). The efficiency of the proposed OT Q-malizing flow net and the infinitesmal DRE net is experimentally demonstrated on high dimensional simulated and image data.

In summary, the contributions of the work include:

* We develop a flow-based model _Q-flow_ net to learn a continuous invertible transport map between arbitrary pair of distributions \(P\) and \(Q\) in \(\mathbb{R}^{d}\) from two sets of samples of the distributions. We propose to train a neural ODE model to minimize the transport cost such that the flow approximates the optimal transport in dynamic OT. The end-to-end training of the model refines an initial flow that may not attain the optimal transport, e.g., obtained by training two CNFs or other interpolating schemes.
* Leveraging a trained Q-flow net, we propose to train a separate continuous-time network _Q-flow-ratio_ net to perform infinitesimal DRE from \(P\) to \(Q\) given finite samples. The Q-flow-ratio net is trained by minimizing classification loss on neighboring time stamps of a discrete-time grid, which is computationally more efficient than prior models.
* We show the effectiveness of the approach on simulated and real data. The optimal transport flow obtained by Q-flow net leads to improved DRE in high dimension, demonstrated by state-of-the-art performance on high-dimensional mutual information estimation and energy-based generative models of image data.

### Related works

Normalizing flows.When the target distribution \(Q\) is an isotropic Gaussian \(\mathcal{N}(0,I_{d})\), normalizing flow models have demonstrated vast empirical successes in building an invertible transport \(T_{t}\) between \(P\) and \(\mathcal{N}(0,I_{d})\)(Kobyzev et al., 2020). The transport is parametrized by deep neural networks, whose parameters are trained via minimizing the KL-divergence between transported distribution \((T_{1})_{\#}P\) and \(\mathcal{N}(0,I_{d})\). Various continuous (Grathwohl et al., 2019; Finlay et al., 2020) and discrete (Dinh et al., 2016; Behrmann et al., 2019) normalizing flow models have been developed, along with proposed regularization techniques (Onken et al., 2021; Xu et al., 2022a,b) that facilitate the training of such models in practice.

Since our \(Q\)-malizing flow is in essence a transport-regularized flow between \(P\) and \(Q\), we further review related works on building normalizing flow models with transport regularization. (Finlay et al., 2020) trained the flow trajectory with regularization based on \(\ell_{2}\) transport cost and Jacobian norm of the network-parametrized velocity field. (Onken et al., 2021) proposed to regularize the flow trajectory by \(\ell_{2}\) transport cost and the deviation from the HJB equation. These regularization have shown to effectively improve over un-regularized models at a reduced computational cost. Regularized normalizing flow models have also been used to solve high dimensional Fokker-Planck equations (Liu et al., 2022) and mean-field games (Huang et al., 2023).

Transport between general distributions.There has been a rich literature on establishing transport between general distributions, especially in the context of optimal transport (OT). The problem of OT dates back to the work by Gaspard Monge (Monge, 1781), and since then many mathematical theories and computational tools have been developed to tackle the question (Villani et al., 2009; Benamou and Brenier, 2000; Peyre et al., 2019). Several works have attempted to make computational OT scalable to high dimensions, including (Lavenant et al., 2018) which applied convex optimization using Riemannian structure of the space of discrete probability distributions, and (Lee et al., 2021) by \(L^{1}\) and \(L^{2}\) versions of the generalized unnormalized OT solved by Nesterov acceleration. Several deep approaches have also been developed recently. (Coeurdoux et al., 2023) leveraged normalizing flow to learn an approximate transport map between two distributions from finite samples, where the flow model has a restricted architecture and the OT constraint is replaced with sliced-Wasserstein distance which may not computationally scale to high dimensional data. (Albergo and Vanden-Eijnden, 2023) proposed to use a stochastic interpolant map between two arbitrary distributions and train a neural network parametrized velocity field to transport the distribution along the interpolated trajectory. Same as in (Lipman et al., 2023; Albergo and Vanden-Eijnden, 2023), our neural-ODE based approach also computes a deterministic probability transport map, in contrast to SDE-based diffusion models (Song et al., 2021). Notably, the interpolant mapping used in (Albergo and Vanden-Eijnden, 2023) was also adopted earlier in Rhodes et al. (2020); Choi et al. (2022) (see Eq. (16) in Appendix B), which is generally not the optimal transport interpolation. In comparison, our proposed Q-malizing flow
optimizes the interpolant mapping parametrized by a neural ODE and approximates the optimal velocity in dynamic OT (see Section 2). Generally, the flow attaining optimal transport can lead to improved model efficiency and generalization performance Huang et al. (2023). In this work, we experimentally show that the optimal transport flow benefits high-dimensional DRE.

Density ratio estimation.Density ratio estimation between distributions \(P\) and \(Q\) is a fundamental problem in statistics and machine learning (Meng and Wong, 1996; Sugiyama et al., 2012). It has direct applications in important fields such as importance sampling (Neal, 2001), change-point detection (Kawahara and Sugiyama, 2009), outlier detection (Kato and Teshima, 2021), mutual information estimation (Belghazi et al., 2018), etc. Various techniques have been developed, including probabilistic classification (Qin, 1998; Bickel et al., 2009), moment matching (Gretton et al., 2009), density matching (Sugiyama et al., 2008), etc. Deep NN models have been leveraged in classification approach Moustakides and Basioti (2019) due to their expressive power. However, as has been pointed out in (Rhodes et al., 2020), the estimation accuracy by a single classification may degrade when \(P\) and \(Q\) differ significantly.

To overcome this issue, (Rhodes et al., 2020) introduced a telescopic DRE approach by constructing intermediate distributions to bridge between \(P\) and \(Q\). (Choi et al., 2022) further proposed to train an infinitesimal, continuous-time ratio net via the so-called time score matching. Despite their improvement over the prior classification methods, both approaches rely on construction of the intermediate distributions between \(P\) and \(Q\) that is not optimal. In contrast, our proposed Q-malizing flow network leverages the expressiveness of deep networks to construct the intermediate distributions by the continuous-time flow transport, and the flow trajectory is regularized to minimize the transport cost in dynamic OT. The model empirically improves the DRE accuracy (see Section 5). In computation, (Choi et al., 2022) applies score matching to compute the infinitesimal change of log-density. The proposed Q-flow-ratio net is based on classification loss training using a fixed time grid which avoids score matching and is computationally lighter (Section 4.2).

## 2 Preliminaries

Neural ODE and CNF.Neural ODE Chen et al. (2018) parametrized an ODE in \(\mathbb{R}^{d}\) by a residual network. Specifically, let \(x(t)\) be the solution of

\[\dot{x}(t)=f(x(t),t;\theta),\quad x(0)\sim p. \tag{1}\]

where \(f(x,t;\theta)\) is a velocity field parametrized by the neural network. Since we impose a distribution \(P\) on the initial value \(x(0)\), the value of \(x(t)\) at any \(t\) also observes a distribution \(p(x,t)\) (though \(x(t)\) is deterministic given \(x(0)\)). In other words, \(p(\cdot,t)=(T_{t})_{\#}p\), where \(T_{t}\) is the solution map of the ODE, namely \(T_{t}(x)=x+\int_{0}^{t}f(x(s),s;\theta)ds\), \(x(0)=x\). In the context of CNF (Kobyzev et al., 2020), the training of the flow network \(f(x,t;\theta)\) is to minimize the KL divergence between the terminal density \(p(x,T)\) at some \(T\) and a target density \(p_{Z}\) which is the normal distribution. The computation of the objective relies on the expression of normal density and can be computed on finite samples of \(x(0)\) drawn from \(p\).
Dynamic OT (Benamou-Brenier).The Benamou-Brenier equation below provides the dynamic formulation of OT Villani et al. (2009), Benamou and Brenier (2000)

\[\begin{split}&\inf_{\rho,v}\mathcal{T}:=\int_{0}^{1}\mathbb{E}_{x(t) \sim\rho(\cdot,t)}\|v(x(t),t)\|^{2}dt\\ & s.t.\quad\partial_{t}\rho+\nabla\cdot(\rho v)=0,\quad\rho(x,0) =p(x),\quad\rho(x,1)=q(x),\end{split} \tag{2}\]

where \(v(x,t)\) is a velocity field and \(\rho(x,t)\) is the probability mass at time \(t\) satisfying the continuity equation with \(v\). The action \(\mathcal{T}\) is the transport cost. Under regularity conditions of \(p\), \(q\), the minimum \(\mathcal{T}\) in (2) equals the squared Wasserstein-2 distance between \(p\) and \(q\), and the minimizer \(v(x,t)\) can be interpreted as the optimal control of the transport problem.

Telescopic and infinitesmal DRE.To circumvent the problem of DRE distinctly different \(p\) and \(q\), the _telescopic DRE_(Rhodes et al., 2020) proposes to "bridge" the two densities by a sequence of intermediate densities \(p_{k}\), \(k=0,\cdots,L\), where \(p_{0}=p\) and \(p_{L}=q\). The consecutive pairs of \((p_{k},p_{k+1})\) are chosen to be close so that the DRE can be computed more accurately, and then by

\[\log(q(x)/p(x))=\log p_{L}(x)-\log p_{0}(x)=\sum_{k=0}^{L-1}\log p_{k+1}(x)- \log p_{k}(x), \tag{3}\]

the log-density ratio between \(q\) and \(p\) can be computed with improved accuracy than a one-step DRE. The _infinitesmal DRE_(Choi et al., 2022) considers a time continuity version of (3). Specifically, suppose the time-parametrized density \(p(x,t)\) is differentiable on \(t\in[0,1]\) with \(p(x,0)=p\) and \(p(x,1)=q\), then

\[\log(q(x)/p(x))=\log p(x,1)-\log p(x,0)=\int_{0}^{1}\partial_{t}\log p(x,t)dt. \tag{4}\]

The quantity \(\partial_{t}\log p(x,t)\) was called the "time score" and can be parametrized by a neural network.

## 3 OT Q-malizing flow network

We introduce the formulation and training objective of the proposed OT Q-malizing flow net in Section 3.1. The training technique consists of the end-to-end training (Section 3.2) and the construction of the initial flow (Section 3.3).

### Formulation and training objective

Given two sets of samples \(\mathbf{X}=\{X_{i}\}_{i=1}^{N}\) and \(\mathbf{\tilde{X}}=\{\tilde{X}_{j}\}_{j=1}^{M}\), where \(X_{i}\sim P\) and \(\tilde{X}_{j}\sim Q\) i.i.d., we train a neural ODE model \(f(x,t;\theta)\) (1) to represent the transport map \(T_{t}\). The formulation is symmetric from \(P\) to \(Q\) and vice versa, and the loss will also have symmetrically two parts. We call \(P\to Q\) the forward direction and \(Q\to P\) the reverse direction. Our training objective is based on the dynamic OT (2) on time \([0,1]\), where we solve the velocity field \(v(x,t)\) by \(f(x,t;\theta)\)
The terminal condition \(\rho(\cdot,1)=q\) is relaxed by a KL divergence (see, e.g., [Ruthotto et al., 2020]). The training loss in forward direction is written as

\[\mathcal{L}^{P\to Q}=\mathcal{L}^{P\to Q}_{\mathrm{KL}}+\gamma\mathcal{L}^{P \to Q}_{T}, \tag{5}\]

where \(\mathcal{L}_{\mathrm{KL}}\) represents the relaxed terminal condition and \(\mathcal{L}_{T}\) is the Wasserstein-2 transport cost to be specified below; \(\gamma>0\) is a weight parameter, and with small \(\gamma\) the terminal condition is enforced.

KL loss.We define the solution mapping of (1) from \(s\) to \(t\) as

\[T^{t}_{s}(x;\theta)=x(s)+\int_{s}^{t}f(x(t^{\prime}),t^{\prime};\theta)dt^{ \prime}, \tag{6}\]

which is also parametrized by \(\theta\), and we may omit the dependence below. By the continuity equation in (2), \(\rho(\cdot,t)=(T^{t}_{0})_{\#}p\). The terminal condition \(\rho(\cdot,1)=q\) is relaxed by minimizing

\[\mathrm{KL}(p_{1}||q)=\mathbb{E}_{x\sim p_{1}}\log(p_{1}(x)/q(x)),\quad p_{1} :=(T^{1}_{0})_{\#}p.\]

The expectation \(\mathbb{E}_{x\sim p_{1}}\) is computed by the sample average over \((X_{1})_{i}\) which observes density \(p_{1}\) i.i.d., where \((X_{1})_{i}:=T^{1}_{0}(X_{i})\) is computed by integrating the neural ODE from time \(0\) to \(1\). It remains to have an estimator of \(\log(p_{1}/q)\) to compute \(KL(p_{1}||q)\), and we propose to train a logistic classification network \(r_{1}(x;\varphi_{r})\) for this. The inner-loop training of \(r_{1}\) is by

\[\min_{\varphi_{r}}\frac{1}{N}\sum_{i=1}^{N}\log(1+e^{r_{1}(T^{1}_{0}(X_{i}; \theta);\varphi_{r})})+\frac{1}{M}\sum_{j=1}^{M}\log(1+e^{-r_{1}(\widehat{X}_{ j};\varphi_{r})}). \tag{7}\]

The functional optimal \(r_{1}^{*}\) of the population version of loss (7) equals \(\log(q/p_{1})\) by direct computation, and as a result, \(\mathrm{KL}(p_{1}||q)=-\mathbb{E}_{x\sim p_{1}}r_{1}^{*}(x)\). The finite sample KL loss is then written as

\[\mathcal{L}^{P\to Q}_{\mathrm{KL}}(\theta)=-\frac{1}{N}\sum_{i=1}^{N}r_{1}(T^{ 1}_{0}(X_{i};\theta);\hat{\varphi}_{r}), \tag{8}\]

where \(\hat{\varphi}_{r}\) is the computed minimizer of (7) solved by inner loops. In practice, when the density \(p_{1}\) is close to \(q\), the DRE by training classification net \(r_{1}\) can be efficient and accurate. We will apply the minimization (7) after the flow net is properly initialized which guarantees the closeness of \(p_{1}=(T^{1}_{0})_{\#}p\) and \(q\) to begin with.

\(W_{2}\) regularization.To compute the transport cost \(\mathcal{T}\) in (2) with velocity field \(f(x,t;\theta)\), we use a time grid on \([0,1]\) as \(0=t_{0}<t_{1}<\ldots<t_{K}=1\). The choice of the time grid is algorithmic (since the flow model is parametrized by \(\theta\) throughout time) and may vary over experiments, see more details in Section 3.2. Define \(h_{k}=t_{k}-t_{k-1}\), and \(X_{i}(t;\theta):=T^{t}_{0}(X_{i};\theta)\), the \(W_{2}\) regularization is written as

\[\mathcal{L}^{P\to Q}_{T}(\theta)=\sum_{k=1}^{K}\frac{1}{h_{k}}\left(\frac{1}{ N}\sum_{i=1}^{N}\|X_{i}(t_{k};\theta)-X_{i}(t_{k-1};\theta)\|^{2}\right). \tag{9}\]
It can be viewed as a time discretization of \(\mathcal{T}\). Meanwhile, since (omitting dependence on \(\theta\)) \(X_{i}(t_{k})-X_{i}(t_{k-1})=T_{t_{k-1}}^{t_{k}}(X_{i}(t_{k-1}))\), the population form of (9) \(\sum_{k=1}^{K}\mathbb{E}_{x\sim\rho(\cdot,t_{k-1})}\|T_{t_{k-1}}^{t_{k}}(x; \theta)\|^{2}/h_{k}\) in minimization can be interpreted as the discrete-time summed (square) Wasserstein-2 distance (Xu et al., 2022a)

\[\sum_{k=1}^{K}W_{2}(\rho(\cdot,t_{k-1}),\rho(\cdot,t_{k}))^{2}/h_{k}.\]

The \(W_{2}\) regularization encourages a smooth flow from \(P\) to \(Q\) with small transport cost, which also guarantees the invertibility of the model in practice when the trained neural network flow approximates the optimal flow in (2).

Both directions.The formulation in the reverse direction is similar, where we transport \(Q\)-samples \(\tilde{X}_{i}\) from \(1\) to \(0\) using the same neural ODE integrated in reverse time. Specifically, \(\mathcal{L}^{Q\to P}=\mathcal{L}^{Q\to P}_{\mathrm{KL}}+\gamma\mathcal{L}^{Q \to P}_{T}\), and \(\mathcal{L}^{Q\to P}_{\mathrm{KL}}(\theta)=-\frac{1}{M}\sum_{j=1}^{M}\tilde{r }_{0}(T_{1}^{0}(\tilde{X}_{j};\theta);\hat{\varphi}_{\tilde{r}})\), where \(\hat{\varphi}_{\tilde{r}}\) is obtained by inner-loop training of another classification net \(\tilde{r}_{0}(x,\varphi_{\tilde{r}})\) via

\[\min_{\varphi_{\tilde{r}}}\frac{1}{M}\sum_{j=1}^{M}\log(1+e^{\tilde{r}_{0}(T_ {1}^{0}(\tilde{X}_{j};\theta);\varphi_{\tilde{r}})})+\frac{1}{N}\sum_{i=1}^{N} \log(1+e^{-\tilde{r}_{0}(X_{i};\varphi_{\tilde{r}})}); \tag{10}\]

Define \(\tilde{X}_{j}(t;\theta):=T_{1}^{t}(\tilde{X}_{j};\theta)\), the reverse-time \(W_{2}\) regularization is

\[\mathcal{L}^{Q\to P}_{T}(\theta)=\sum_{k=1}^{K}\frac{1}{h_{k}}\left(\frac{1}{ M}\sum_{j=1}^{M}\|\tilde{X}_{j}(t_{k-1};\theta)-\tilde{X}_{j}(t_{k};\theta)\|^{2} \right).\]

### End-to-end training algorithm

In the end-to-end training, we assume that the Q-flow net has already been initiated as an approximate solution of the desired Q-malizing flow, see more in Section 3.3. We then minimize \(\mathcal{L}^{P\to Q}\) and \(\mathcal{L}^{Q\to P}\) in an alternative fashion per "flip", and the procedure is given in Algorithm 1. Hyperparameter choices and network architectures are further detailed in Appendix A.

Time integration of flow.In the losses (8) and (9), one need to compute the transported samples \(X_{i}(t;\theta)\) and \(\tilde{X}_{j}(t;\theta)\) on time grid points \(\{t_{k}\}_{k=0}^{K}\). This calls for integrating the neural ODE on \([0,1]\), which we conduct on a fine time grid \(t_{k,s}\), \(s=0,\dots,S\), that divides each subinterval \([t_{k-1},t_{k}]\) into \(S\) mini-intervals. We compute the time integration of \(f(x,t;\theta)\) using a fixed-grid four-stage Runge-Kutta method on each mini-interval. The fine grid is used to ensure the numerical accuracy of ODE integration and the numerical invertibility of the Q-flow net, i.e., the error of using reverse-time integration as the inverse map (see inversion errors in Table A.1). It is also possible to first train the flow \(f(x,t;\theta)\) on a time grid to warm start the later training on a refined grid, so as to improve convergence. We also find that the \(W_{2}\) regularization can be computed at a coarser grid \(t_{k}\) (\(S\) is usually 3-5 in our experiments) without losing the effectiveness of Wasserstein-2 regularization. Finally, one can adopt an adaptive time grid, e.g., by enforcing equal \(W_{2}\) movement on each subinterval \([t_{k-1},t_{k}]\) Xu et al. (2022b), so that the representative points are more evenly distributed along the flow trajectory and the learning of the flow model can be further improved.
```
0: Pre-trained initial flow network \(f(x(t),t;\theta)\); training data \(\mathbf{X}\sim P\) and \(\widetilde{\mathbf{X}}\sim Q\); hyperparameters: \(\{\gamma,\{t_{k}\}_{k=1}^{K},\texttt{Tot},E,E_{0},E_{\text{in}}\}\).
0: Refined flow network \(f(x(t),t;\theta)\)
1:for flip = \(1,\ldots,\texttt{Tot}\)do
2: (If flip = 1) Train \(r_{1}\) by minimizing (7) for \(E_{0}\) epochs.
3:for epoch = \(,1,\ldots,E\)do {\(\triangleright\)\(P\to Q\) refinement}
4: Update \(\theta\) of \(f(x(t),t;\theta)\) by minimizing \(\mathcal{L}^{P\to Q}\).
5: Update \(r_{1}\) by minimizing (7) for \(E_{\text{in}}\) epochs.
6:endfor
7: (If flip = 1) Train \(\widetilde{r}_{0}\) by minimizing (10) for \(E_{0}\) epochs.
8:for epoch = \(,1,\ldots,E\)do {\(\triangleright\)\(Q\to P\) refinement}
9: Update \(\theta\) of \(f(x(t),t;\theta)\) by minimizing \(\mathcal{L}^{Q\to P}\).
10: Update \(\widetilde{r}_{0}\) by minimizing (10) for \(E_{\text{in}}\) epochs.
11:endfor
12:endfor
```

**Algorithm 1** OT \(Q\)-malizing flow refinement

Inner-loop training of \(r_{1}\) and \(\widetilde{r}_{0}\).Suppose the flow net has been successfully warm-started, the transported distributions \((T_{0}^{1})_{\#}P\approx Q\) and \((T_{1}^{0})_{\#}Q\approx P\). The two classification nets are first trained for \(E_{0}\) epochs before the loops of training the flow model and then updated for \(E_{\text{in}}\) inner-loop epochs in each outer-loop iteration. We empirically find that the diligent updates of \(r_{1}\) and \(\widetilde{r}_{0}\) in lines 5 and 10 of Algorithm 1 are crucial for successful end-to-end training of Q-flow net. As we update the flow model \(f(x,t;\theta)\), the push-forwarded distributions \((T_{0}^{1})_{\#}P\) and \((T_{1}^{0})_{\#}Q\) are consequently changed, and then one will need to retrain \(r_{1}\) and \(\widetilde{r}_{0}\) timely to ensure an accurate estimate of the log-density ratio and consequently the KL loss. Compared with training the flow parameter \(\theta\), the computational cost of the two classification nets is light which allows potentially a large number of inner-loop iterations if needed.

Computational complexity.We measure the computational complexity by the number of function evaluations of \(f(x(t),t;\theta)\) and of the classification nets \(\{r_{1},\widetilde{r}_{0}\}\). Suppose the total number of epochs in outer loop training is \(O(E)\), the dominating computational cost lies in the neural ODE integration, which takes \(O(8KS\cdot E(M+N))\) function evaluations of \(f(x,t;\theta)\). We remark that the Wasserstein-2 regularization (9) incurs no extra computation, since the samples \(X_{i}(t_{k};\theta)\) and \(\tilde{X}_{j}(t_{k};\theta)\) are available when computing the forward and reverse time integration of \(f(x,t;\theta)\). The training of the two classification nets \(r_{1}\) and \(\widetilde{r}_{0}\) takes \(O(4(E_{0}+EE_{\text{in}})(M+N))\) additional evaluations of the two network functions since the samples \(X_{i}(1;\theta)\) and \(\tilde{X}_{j}(0;\theta)\) are already computed.

### Flow initialization via CNF

We propose to initialize the Q-flow net by a good approximation of a Q-malizing flow, which matches the transported distributions with the target distributions in both directions and not
necessarily minimizes the transport cost. Such an initialization will significantly accelerate the convergence of the end-to-end training, which can be viewed as a refinement of the initial flow.

The initial flow \(f(x,t;\theta)\) may be specified using prior knowledge of the problem if available. Generally, when only two data sets \(\mathbf{X},\mathbf{\tilde{X}}\) are given, any reasonable interpolating scheme may provide such an initial flow. For example, one can use the linear interpolant mapping in (Rhodes et al., 2020; Choi et al., 2022; Albergo and Vanden-Eijnden, 2023) (see Appendix B), and train the neural network velocity field \(f(x,t;\theta)\) to match the interpolation (Albergo and Vanden-Eijnden, 2023). Here we propose an alternative approach that constructs the initial flow as a concatenation of two CNF models (Figure 1).

Specifically, as has been explained in see Section 2, a CNF model transports from a data distribution (accessible by finite samples) to the normal distribution \(\mathcal{N}(0,I_{d})\). We first train two separate CNFs represented by flow networks \(f(x(t),t;\theta_{P})\) and \(f(x(t),t;\theta_{Q})\), which transport \(P\) to \(\mathcal{N}(0,I_{d})\) on time \([0,1/2]\), and \(Q\) to \(\mathcal{N}(0,I_{d})\) reversely on time \([1/2,1]\), respectively. We then concatenate the two flows to be a single flow model on time \([0,1]\) that transports from \(P\) to \(Q\) (through an intermediate distribution which is normal). The network parameter \(\theta=\{\theta_{P},\theta_{Q}\}\) contains parameters in both CNFs and in computation cost, the overall computation of flow initialization is the sum of training the two CNF models. Any existing neural-ODE CNF models may be adopted for our purpose. In this work, we adopt the JKO-iFlow approach in (Xu et al., 2022b).

## 4 Infinitesimal density ratio estimation

We use a trained Q-malizing flow network \(f(x,t;\theta)\) for infinitesimal DRE as a focused application.

Figure 1: Illustration of building the OT \(Q\)-malizing flow (blue) from CNF initial flow (black). The flow is initialized by concatenating two CNF models as black arrows. The refined flow as the blue two-sided arrow refines the trajectory between \(P\) and \(Q\) to approximate the dynamic OT.
### Training by logistic classification

Let \(p(x,t)=(T_{t})_{\#}p\) and \(T_{t}\) is the transport induced by the trained Q-flow net in Section 3. Using (4), we propose to parametrize the time score \(\partial_{t}\log p(x,t)\) by a neural network \(r(x,t;\theta_{r})\) with parameter \(\theta_{r}\), called the _Q-flow-ratio_ net. The training is by logistic classification applied to transported data distributions on consecutive time grid points: Given a deterministic time grid \(0=t_{0}<t_{1}<\ldots<t_{L}=1\) (which again is an algorithmic choice, see Section 4.2), we expect that the integral

\[R_{k}(x;\theta_{r}):=\int_{t_{k-1}}^{t_{k}}r(x,t;\theta_{r})dt\approx\int_{t_{ k-1}}^{t_{k}}\partial_{t}\log p(x,t)dt=\log(p(x,t_{k})/p(x,t_{k-1})). \tag{11}\]

By that logistic classification recovers the log density ratio as has been used in Section 3.1, this suggests the loss on interval \([t_{k-1},t_{k}]\) as follows, where \(X_{i}(t):=T_{0}^{t}(X_{i})\) and \(T_{0}^{t}\) is computed by integrating the trained Q-flow net,

\[L_{k}^{P\to Q}(\theta_{r})=\frac{1}{N}\sum_{i=1}^{N}\log(1+e^{R_{k}(X_{i}(t_{ k-1});\theta_{r})})+\frac{1}{N}\sum_{i=1}^{N}\log(1+e^{-R_{k}(X_{i}(t_{k}); \theta_{r})}). \tag{12}\]

When \(k=L\), the distribution of \(X_{i}(t_{L})\) may slightly differ from that of \(Q\) due to the error in matching the terminal densities in Q-flow net. Thus by replacing the 2nd term in (12) with an empirical average over the \(Q\)-samples \(\tilde{X}_{j}\) may be beneficial. In the reverse direction, define \(\tilde{X}_{j}(t):=T_{1}^{t}(\tilde{X}_{j})\), we similarly have \(L_{k}^{Q\to P}(\theta_{r})=\frac{1}{M}\sum_{j=1}^{M}\log(1+e^{R_{k}( \tilde{X}_{j}(t_{k-1});\theta_{r})})+\frac{1}{M}\sum_{j=1}^{M}\log(1+e^{-R_{k} (\tilde{X}_{j}(t_{k});\theta_{r})})\), and when \(k=1\), we replace the 1st term with an empirical average over the \(P\)-samples \(X_{i}\). The training of the Q-flow-ratio net is by

\[\min_{\theta_{r}}\sum_{k=1}^{L}L_{k}^{P\to Q}(\theta_{r})+L_{k}^{Q \to P}(\theta_{r}). \tag{13}\]

When trained successfully, the integral of \(r(x,t;\theta_{r})\) over \(t\in[0,1]\) yields the desired log density ratio \(\log(q/p)\) by (4), and furtherly the integral \(\int_{s}^{t}r(x,t^{\prime};\theta_{r})dt^{\prime}\) provides an estimate of \(\log(p(x,t)/p(x,s))\) for any \(s<t\) on \([0,1]\).

### Algorithm and computational complexity

The details of minimizing (13) is given in Algorithm 2. We use an evenly spaced time grid \(t_{k}=k/L\) in all experiments. In practice, one can also progressively refine the time grid in training, starting from a coarse grid to train a Q-flow-ratio net \(r(x,t;\theta_{r})\) and use it as a warmstart for training the network parameter \(\theta_{r}\) on a refined grid. When the time grid is fixed, it allows us to compute the transported samples \(\{X_{i}(t_{k})\}_{i=1}^{N},\{\tilde{X}_{j}(t_{k})\}_{j=1}^{M}\) on all \(t_{k}\) once before the training loops of Q-flow-ratio net (line 1-3). This part takes \(O(8KS(M+N))\) function evaluations of the pre-trained Q-flow net \(f(x,t;\theta)\). Suppose the training loops of line 4-6 conducts \(E\) epochs in total. Assume each time integral in \(R_{k}\) (11) is computed by a fixed-grid four-stage Runge-Kutta method, then \(O(4LE(M+N))\) function evaluations of \(r(x,t;\theta_{r})\) is needed to compute the overall loss (13).
## 5 Experiments

We examine the proposed Q-flow and Q-flow-ratio nets on several examples. We denote our method as "Ours", and compare against three baselines of DRE in high dimensions. The baseline methods are: 1 ratio (by training a single classification network using samples from \(P\) and \(Q\)), TRE [Rhodes et al., 2020], and DRE-\(\infty\)[Choi et al., 2022]. We denote \(P_{t_{k}}\) with density \(p(\cdot,t_{k})\) as the pushforward distribution of \(P\) by the Q-flow transport over the interval \([0,t_{k}]\). The set of distributions \(\{P_{t_{k}}\}\) for \(k=1,\ldots,L\) builds a bridge between \(P\) and \(Q\). Code can be found at [https://github.com/hamrel-cxu/Q-flow-ratio](https://github.com/hamrel-cxu/Q-flow-ratio).

### 2d Gaussian mixtures

We simulate \(P\) and \(Q\) as two Gaussian mixture models with three and two components, respectively, see additional details in Appendix A.1. We compute ratio estimates \(\hat{r}(x)\) with the true value \(r(x)\), which can be computed using the analytic expressions of the densities. The results are shown in Figure 2. We see from the top panel that the mean absolute error (MAE) of Ours is evidently smaller than those of the baseline methods, and Ours also incurs a smaller maximum error \(|\hat{r}-r|\) on test samples. This is consistent with the closest resemblance of Ours to the ground truth (first column) in the bottom panel. In comparison, DRE-\(\infty\) tends to over-estimate \(r(x)\) on the support of \(Q\), while TRE and 1 ratio can severely under-estimate \(r(x)\) on the support of \(P\). As both the DRE-\(\infty\) and TRE models use the linear interpolant scheme (16), the result suggests the benefit of training an optimal-transport flow for DRE.

Figure 2: Estimated log density ratio between 2D Gaussian mixture distributions \(P\) (three components) and \(Q\) (two components). **Top:** (a) training samples from \(P\) and \(Q\). (b)-(d) histograms of errors \(\log(|r(x)-\hat{r}(x)|)\) computed at 10K test samples shown in log-scale. The MAE (15) are shown in the captions. **Bottom:** true and estimated \(\log(q/p)\) from different models shown under shared colorbars.
### Non-Gaussian two-dimensional flow

We design two densities in \(\mathbb{R}^{2}\) where \(P\) represents the shape of two moons and \(Q\) represents a checkerboard, see additional details in Appendix A.2. For this more challenging case, the linear interpolation scheme (16) creates a bridge between \(P\) and \(Q\) as shown in Figure A.3. The flow visually differs from the one obtained by the trained Q-flow net, as shown in Figure 3(a), and the latter is trained to minimize the transport cost.

The result of Q-flow-ratio net is shown in Figure 3(b). The corresponding density ratio estimates of \(\log p(x,t_{k})-\log p(x,t_{k-1})\) visually reflect the actual differences in the two neighboring densities.

### Mutual Information estimation for high-dimensional data

We evaluate different methods on estimating the mutual information (MI) between two correlated distributions from given samples. In this example, we let \(P\) and \(Q\) be two correlated, high-dimensional Gaussian distributions following the setup in (Rhodes et al., 2020; Choi et al., 2022), where we vary the data dimension \(d\) in the range of \(\{40,80,160,320\}\). Additional details can be found in Appendix A.3. Figure 4 shows the results by different methods, where the baselines are trained under their proposed default settings. We find that the estimated MI by our method almost perfectly aligns with the ground truth MI values, reaching nearly identical performance as DRE-\(\infty\). Meanwhile, Ours outperforms the other two baselines

Figure 4: Estimated MI between two correlated high-dimensional Gaussian distributions.

Figure 3: Flow between arbitrary 2D distributions and corresponding log-ratio estimation. **Top**: the intermediate distributions by Q-flow net. **Bottom**: corresponding log-ratio estimated by Q-flow-ratio net. Bluer color indicates smaller estimates of the difference \(\log(p(x,t_{k})/p(x,t_{k-1}))\) evaluated at the common support of the neighboring densities.
increase as the dimension \(d\) increases.

### Energy-based modeling of MNIST

We apply our approach in evaluating and improving an energy-base model (EBM) on the MNIST dataset [LeCun and Cortes, 2005]. We follow the prior setup in [Rhodes et al., 2020, Choi et al., 2022], where \(P\) is the empirical distribution of MNIST images, and \(Q\) is the generated image distributions by three given pre-trained energy-based generative models: a Gaussian noise model, a Gaussian copula model, and a Rational Quadratic Neural Spline Flow model (RQ-NSF) [Durkan et al., 2019]. Specifically, the images are in dimension \(d=28^{2}=784\), and each of the pre-trained models provides an invertible mapping \(F:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}\), where \(Q=F_{\#}\mathcal{N}(0,I_{d})\). We train a Q-flow net between \((F^{-1})_{\#}P\) and \((F^{-1})_{\#}Q\), the latter by construction equals \(\mathcal{N}(0,I_{d})\). Using the trained Q-flow net, we go back to the input space and train the Q-flow-ratio net using the intermediate distributions between \(P\) and \(Q\). Additional details are in Appendix A.4.

The trained Q-flow-ratio \(r(x,s;\hat{\theta}_{r})\) provides an estimate of the data density \(p(x)\) by \(\hat{p}(x)\) defined as \(\log\hat{p}(x)=\log q(x)-\int_{0}^{1}r(x,s;\hat{\theta}_{r})ds\), where \(\log q(x)\) is given by the change-of-variable formula using the pre-trained model \(F\) and the analytic expression of \(\mathcal{N}(0,I_{d})\). As a by-product, since our Q-flow net provides an invertible mapping \(T_{0}^{1}\), we can use it to obtain an improved generative model on top of \(F\). Specifically, the improved distribution \(\tilde{Q}:=(F\circ T_{1}^{0})_{\#}\mathcal{N}(0,I_{d})\), that is, we first use Q-flow to transport \(\mathcal{N}(0,I_{d})\) and then apply \(F\). The performance of the improved generative model can be measured using the "bits per dimension" (BPD) metric:

\[\text{BPD}=\frac{1}{N^{\prime}}\sum_{i=1}^{N^{\prime}}[-\log\hat{p}(X_{i})/(d \log 2)], \tag{14}\]

where \(X_{i}\) are \(N^{\prime}\)=10K test images drawn from \(P\). BPD has been a widely used metric in evaluating the performance of generative models [Theis et al., 2015, Papamakarios et al., 2017]. In our setting, the BPD can also be used to compare the performance of the DRE.

The results show that Ours reaches the state-of-the-art performance in Table 1: it consistently reaches smaller BPD than the baseline methods across all choices of \(Q\). Meanwhile, we also note computational benefits in training: Ours took approximately 8 hours to converge, when TRE took 24 hours and DRE-\(\infty\) took approximately 2 days. We further show improvement in generated digits using the trained Q-flow as shown in Figure 5 for RQ-NSF and Figure A.1 for the other two specifications of \(Q\). The trajectory of improved samples from \(Q\) to \(\tilde{Q}\) for RQ-NSF is further shown in Figure A.2.

\begin{table}
\begin{tabular}{c|c c c c|c c c c|c c c} \hline \hline
**Choice of \(Q\)** & \multicolumn{4}{c|}{RQ-NSF} & \multicolumn{4}{c|}{Copula} & \multicolumn{4}{c}{Gaussian} \\ \hline
**Method** & Ours & DRE-\(\infty\) & TRE & 1 ratio & Ours & DRE-\(\infty\) & TRE & 1 ratio & Ours & DRE-\(\infty\) & TRE & 1 ratio \\ \hline
**BPD (\(\downarrow\))** & **1.05** & 1.09 & 1.09 & 1.09 & **1.14** & 1.21 & 1.24 & 1.33 & **1.31** & 1.33 & 1.39 & 1.96 \\ \hline \hline \end{tabular}
\end{table}
Table 1: DRE performance on the energy-based modeling task for MNIST, reported in BPD and lower is better. Results for DRE-\(\infty\) are from [Choi et al., 2022] and results for 1 ratio and TRE are from [Rhodes et al., 2020].
## 6 Discussion

In this work, we develop Q-malizing flow neural-ODE model which smoothly and invertibly transports between a pair of arbitrary distributions \(P\) and \(Q\), trained to approximate the optimal transport flow, and is learned from finite samples from both distributions. The Q-flow net can be initialized by any interpolant mapping between \(P\) and \(Q\) or concatenating two CNF models, and is subsequently refined by end-to-end training that minimizes the KL-divergence between transported and target distributions plus the Wasserstein-2 transport cost. The trained flow model facilitates the infinitesimal DRE where we train another Q-flow-ratio net by logistic classification between transported empirical data distributions on consecutive time grid points. The proposed method shows strong empirical performance on simulated and image data.

For future directions, it would be of interest to investigate alternative initialization schemes that may be more efficient than the current strategy of training CNFs. The algorithms of training the Q-flow net and the Q-flow-ratio net can be further enhanced by advanced adaptive schemes of time discretization. Other applications beyond DRE and on real datasets would also be useful extensions.

## Acknowledgement

The work is supported by NSF DMS-2134037. C.X. and Y.X. are partially supported by an NSF CAREER CCF-1650913, NSF DMS-2134037, CMMI-2015787, CMMI-2112533, DMS-1938106, and DMS-1830210. X.C. is partially supported by NSF CAREER DMS-2237842, Simons Foundation 814643, the Alfred P. Sloan Foundation.

Figure 5: Improvement in generated samples from \(Q\), where \(Q\) is given by RQ-NSF, a pre-trained normalizing flow model \(F\) that yields \(Q=F_{\#}\mathcal{N}(0,I_{d})\).