A Machine Learning Framework for Geodesics Under Spherical Wasserstein–Fisher–Rao Metric and Its Application for Weighted Sample Generation

Abstract

Wasserstein–Fisher–Rao (WFR) distance is a family of metrics to gauge the discrepancy of two Radon measures, which takes into account both transportation and weight change. Spherical WFR distance is a projected version of WFR distance for probability measures so that the space of Radon measures equipped with WFR can be viewed as metric cone over the space of probability measures with spherical WFR. Compared to the case for Wasserstein distance, the understanding of geodesics under the spherical WFR is less clear and still an ongoing research focus. In this paper, we develop a deep learning framework to compute the geodesics under the spherical WFR metric, and the learned geodesics can be adopted to generate weighted samples. Our approach is based on a Benamou–Brenier type dynamic formulation for spherical WFR. To overcome the difficulty in enforcing the boundary constraint brought by the weight change, a Kullback–Leibler divergence term based on the inverse map is introduced into the cost function. Moreover, a new regularization term using the particle velocity is introduced as a substitute for the Hamilton–Jacobi equation for the potential in dynamic formula. When used for sample generation, our framework can be beneficial for applications with given weighted samples, especially in the Bayesian inference, compared to sample generation with previous flow models.

Appendix A: Implementation Details for Solving SWFR Distance with Optimization Method

We can apply a primal-dual hybrid algorithm to compute the SWFR distance. The algorithm is also used to compute the earth mover’s distance in [31]. To begin with, we reformulate the optimization problem (3.4) as a convex form:

$$\begin{aligned} \min _{\rho , m, \xi }\left\{ \frac{1}{2} \int _0^1 \int _{\mathbb {R}^d} \left( \frac{|m|^2}{\rho }+\alpha \frac{\xi ^{2}}{\rho } \right) d x d t, \partial _t \rho +\nabla \cdot m =\xi , \int _{\mathbb {R}^d} \xi d x=0\right\} . \end{aligned}$$

(A.1)

Then we solve the corresponding min-max problem:

$$\begin{aligned}{} & {} \max _{\phi ,\lambda }\min _{\rho ,m,\xi } \int _0^1 \int _{\Omega } \left( \frac{|m|^2}{\rho }+\alpha \frac{\xi ^{2}}{\rho } \right) d x d t\nonumber \\{} & {} \quad + \int _0^1\int _{\Omega } \phi (x,t)\left( \partial _t \rho +\nabla \cdot m-\xi \right) dx dt+\int _0^1\lambda (t) \left( \int _{\Omega } \xi dx \right) dt. \end{aligned}$$

(A.2)

We consider the space domain $\Omega =[0,1]^{d}$ and the time domain [0, 1] for simplicity. Let $\Omega _{h}$ be the discrete space mesh-grid of $\Omega$ with step size h, i.e. $\Omega _{h}=\{0,h,2h,\cdots ,1\}^d$. The time domain [0, 1] is discreted with step size $\Delta t$ . Let $N_x=1/h$ be the space grid size and $N_t=1/\Delta t$ be the time grid size. All optimization variables ($\rho$, $\xi$, m, $\phi$ and $\lambda$) are defined on the grid.

We employ the same discrete scheme for divergence operator and boundary conditions as [58]. We give some definitions on the discrete space $\Omega _h$:

$$\begin{aligned} \begin{aligned} \int _{\Omega _h}f(x)dx&:=\sum _{x \in \Omega _{h},x_i \ne 0} f(x) h^d, \quad \left\langle f, g \right\rangle _h:= \int _{\Omega _h}f(x) g(x) dx,\\ \nabla _{h} \cdot m(x)&:=\sum _{i=1}^{d} D_{h,i} m(x), x\in \Omega _h, \end{aligned} \end{aligned}$$

where $D_{h,i}$ denotes discrete differential operator for i-th component in i-th dimension with step size h:

$$\begin{aligned} D_{h, i} m(x)=\left( m_i \left( x_1,\cdots , x_i, \cdots , x_d\right) -m_i \left( x_1,\cdots , x_i-h, \cdots , x_d\right) \right) / h, \quad h \le x_i \le 1. \\ \end{aligned}$$

Then the discretization of (A.2) can be written as:

$$\begin{aligned}{} & {} \max _{\phi ,\lambda }\min _{\rho ,m,\xi }L\left( m, \xi , \rho ,\phi ,\lambda \right) = \sum _{t=1}^{N_t} \left\{ \int _{\Omega _h} \left( \frac{|m_{t}|^2}{2\rho _{t}}+\alpha \frac{\xi _{t}^{2}}{2\rho _{t}}\right) d x + \left\langle \phi _{t},\frac{\rho _{t+1}-\rho _{t}}{\Delta t}+\nabla _h \cdot m_{t}-\xi _{t} \right\rangle _h \right. \nonumber \\{} & {} \quad \left. +\lambda _t \left( \int _{\Omega _h} \xi _{t} dx \right) \right\} . \end{aligned}$$

(A.3)

From the discretized form, we can describe the sizes of the optimization variables individually. The sizes of discretized variables are : $(N_t+1) \times (N_x+1)^d$ for $\rho$, $N_t\times (N_x+1)^{d} \times d$ for m, $N_t\times (N_x)^{d}$ for $\xi$ and $\phi$, and $N_t$ for $\lambda$. The boundary conditions are given as $\rho _1=\mu$, $\rho _{N_t+1}=\nu$ and $m(x)=0$ for all $x \in \partial \Omega _h$. Then the primal-dual hybrid algorithm gives as follows for variables with superscripts k:

$$\begin{aligned} \begin{aligned}&\left( m_{t}^{k+1},\xi _t^{k+1},\rho _t^{k+1}\right) =\underset{m^{\star },\xi ^{\star },\rho ^{\star }}{\arg \min } L\left( m^{\star },\xi ^{\star },\rho ^{\star },\phi ^{k},\lambda ^{k} \right) \\&\qquad \quad +\frac{1}{2 \mu }\left( \left\| m^{\star }-m^{k}_{t}\right\| _{2}^2+\alpha \left\| \xi ^{\star }-\xi ^k_t \right\| _{ 2}^2 +\left\| \rho ^{\star }-\rho ^k_t\right\| _{ 2}^2 \right) , \\&\tilde{m}_t^{k+1}=2 m_t^{k+1}-m_t^k, \quad \tilde{\xi }_t^{k+1}=2 \xi _t^{k+1}-\xi _t^k, \quad \tilde{\rho }_t^{k+1}=2 \rho _t^{k+1}-\rho _t^k ,\\&\qquad \left( \phi _t^{k+1},\lambda _t^{k+1}\right) =\underset{\phi ^{\star },\lambda ^{\star }}{\arg \max } L\left( \tilde{m}_t^{k+1}, \tilde{\xi }_h^{k+1}, \tilde{\rho }_h^{k+1},\phi ^{\star }, \lambda ^{\star }\right) -\frac{1}{2 \tau } \left( \left\| \phi ^{\star }-\phi ^{k}_t\right\| _{2}^2 + \left\| \lambda ^{\star }-\lambda ^{k}_t\right\| _{2}^2 \right) . \\&\end{aligned} \end{aligned}$$

(A.4)

The first step of the algorithm is equivalent to solving the following system:

$$\begin{aligned} \left\{ \begin{array}{l} m^{\star }=\frac{\rho ^{\star }(m_t^{k}-\mu \text {div}_{h}^{*} \phi _t^{k})}{\rho ^{\star }+\mu }, \\ \xi ^{\star }=\frac{\rho ^{\star }(\xi _t^{k}+\frac{1}{\alpha }\mu (\phi _t^{k}-\lambda _t^k))}{\rho ^{\star }+\mu }, \\ -\frac{{m^{\star }}^{2}}{2{\rho ^{\star }}^2}-\alpha \frac{{\xi ^{\star }}^{2}}{2{\rho ^{\star }}^2}+\frac{\phi _{t-1}^{k}-\phi ^{k}_{t}}{\Delta t}+\frac{1}{\mu }(\rho ^{\star }-\rho _t^k)=0, \end{array}\right. \end{aligned}$$

(A.5)

where $\text {div}_{h}^{*}$ represents the conjugate operator of divergence operator. By the definition of conjugate operator, we have

$$\begin{aligned} \left\langle \nabla _h \cdot f,g \right\rangle _{h} =\left\langle f,\text {div}_{h}^{*} g \right\rangle _h= \left\langle f, -\nabla _h g \right\rangle _h. \end{aligned}$$

Thus $\text {div}_{h}^{*}=-\nabla _{h}$ and

$$\begin{aligned} \nabla _h u =\left( \partial _{h, 1} u, \partial _{h, 2} u, \cdots , \partial _{h, d} u\right) , \end{aligned}$$

where each $\partial _{h,i}$ denotes discrete differential operator in i-th dimension with step size h:

$$\begin{aligned} \partial _{h, i} u(x)=\left\{ \begin{array}{l} \left( u\left( x_1,\cdots , x_i+h, \cdots , x_d\right) -u\left( x_1,\cdots , x_i, \cdots , x_d\right) \right) / h, \quad 0\le x_i < 1, \\ 0,\quad x_i=1. \end{array}\right. \end{aligned}$$

We also introduce a natural boundary condition $\phi _0=0$, which comes from deriving optimal value condition for updating $\rho$. Solving above system (A.5) requires us to solve the roots for a third order polynomial, where $\rho ^{\star }$ should be the largest real root. The third step of the algorithm is to update dual variables:

$$\begin{aligned} \left\{ \begin{array}{l} \phi _t^{k+1}=\phi _{t}^{k}+ \tau \left( \frac{\tilde{\rho }_{t+1}^{k+1}-\tilde{\rho }_{t}^{k+1}}{\Delta t}-\nabla \cdot \tilde{m}_{t}^{k+1}-\tilde{\xi }_t^{k+1}\right) , \\ \lambda _t^{k+1}=\lambda _{t}^{k}+\tau \left( \int _{\Omega }\tilde{\xi }_t^{k+1}dx\right) . \end{array}\right. \end{aligned}$$

(A.6)