SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {W F R} _ {\lambda} \left(q, q ^ {\prime}\right) ^ {2} = \Delta t ^ {2} \left| \mu_ {t} ^ {\prime} \right| ^ {2} = \Delta t ^ {2} \int d x q (x) \left[ \left\| v ^ {*} (x) \right\| ^ {2} + \lambda g ^ {*} (x) ^ {2} \right]. \tag {54}
$$
Thus, the MMS problem can be written as
$$
\begin{array}{l} \inf  _ {g _ {t}, v _ {t}} F [ q _ {t} ] - \Delta t \int \mathrm {d} x \nabla_ {x} \cdot (q _ {t} (x) v _ {t} (x)) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) + \int \mathrm {d} x q _ {t} (x) g _ {t} (x) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - F [ q _ {t} ] \\ + \frac {\Delta t}{2} \int \mathrm {d} x q _ {t} (x) \left[ \| v _ {t} (x) \| ^ {2} + \lambda g _ {t} (x) ^ {2} \right], \tag {55} \\ \end{array}
$$
$$
\begin{array}{l} \inf  _ {g _ {t}, v _ {t}} \int \mathrm {d} x q _ {t} (x) \left\langle v _ {t} (x), \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\rangle + \int \mathrm {d} x q _ {t} (x) g _ {t} (x) \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right) \\ + \frac {1}{2} \int \mathrm {d} x q _ {t} (x) \left[ \| v _ {t} (x) \| ^ {2} + \lambda g _ {t} (x) ^ {2} \right], \tag {56} \\ \end{array}
$$
$$
\begin{array}{l} \inf  _ {g _ {t}, v _ {t}} \int \mathrm {d} x q _ {t} (x) \| v _ {t} (x) + \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \| ^ {2} \\ + \lambda \int \mathrm {d} x q _ {t} (x) \left(g _ {t} (x) + \frac {1}{\lambda} \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right)\right) ^ {2}. \tag {57} \\ \end{array}
$$
$$
v _ {t} (x) = - \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x), \quad g _ {t} (x) = - \frac {1}{\lambda} \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right). \tag {58}
$$
Note that different values of  $\lambda$  result in different gradient flows. For instance, considering the limit  $\lambda \to \infty$  we have  $g(x) \to 0$  and Eq. (52) just becomes  $W_{2}$  gradient flow, which is natural since we have an infinite penalty for the mass teleportation in our metric. Setting  $\lambda \to 0$  requires some additional consideration, since then the growth term explodes, and all the mass will be teleported without any cost. Indeed, for  $\lambda \to 0$ , our metric does not penalize for the mass teleportation at all, but our change of density (Eq. (52)) is still able to teleport mass, hence, it will be doing so "for free".
A.3 PDEs demonstrating the convergence
Consider the change of the density  $q_{t}$  under the continuity equation with the vector field  $v_{t}(x)$  and the growth term  $g_{t}(x)$
$$
\frac {\partial q _ {t}}{\partial t} (x) = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) + g _ {t} (x) q _ {t} (x). \tag {59}
$$
Thus, the change of the functional  $F[q]$  is
$$
\begin{array}{l} \frac {\mathrm {d}}{\mathrm {d} t} F [ q _ {t} ] = - \int \mathrm {d} x \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) + \int \mathrm {d} x q _ {t} (x) g _ {t} (x) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) (60) \\ = \int \mathrm {d} x q _ {t} (x) \left\langle v _ {t} (x), \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\rangle + \int \mathrm {d} x q _ {t} (x) g _ {t} (x) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x). (61) \\ \end{array}
$$
From this equation, we can clearly see that  $v_{t}(x)$  and  $g_{t}(x)$  derived in the previous section minimize  $F[q]$ . Indeed, taking
$$
v _ {t} (x) = - \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x), \quad g _ {t} (x) = - \frac {1}{\lambda} \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right), \tag {62}
$$
$$
\frac {\mathrm {d}}{\mathrm {d} t} F [ q _ {t} ] = - \int \mathrm {d} x q _ {t} (x) \left\| \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\| ^ {2} - \frac {1}{\lambda} \int \mathrm {d} x q _ {t} (x) \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y)\right) ^ {2}
$$
$$
- \frac {1}{\lambda} \underbrace {\int \mathrm {d} x q _ {t} (x) \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right)} _ {= 0} \mathbb {E} _ {q _ {t} (z)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (z) \right] \leq 0. \tag {64}
$$
Note that the considered growth term preserves the normalization of the density, i.e.,
$$
\int \mathrm {d} x \frac {\partial q _ {t}}{\partial t} (x) = \int \mathrm {d} x g _ {t} (x) q _ {t} (x) = - \int \mathrm {d} x q _ {t} (x) \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right) = 0. \tag {65}
$$
Thus, our functional  $F[q]$  decreases when the density  $q_{t}$  evolves according to the PDE
$$
\frac {\partial q _ {t}}{\partial t} (x) = \underbrace {- \nabla_ {x} \cdot \left(q _ {t} (x) \left(- \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x)\right)\right)} _ {\text {t h e c o n t i n u i t y e q u a t i o n}} - \frac {1}{\lambda} \underbrace {\left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right)} _ {\text {g r o w t h t e r m}} q _ {t} (x), \tag {66}
$$
and reaches its stationary point when  $\left\| \nabla_{x}\frac{\delta F[q_{t}]}{\delta q_{t}} (x)\right\|^{2} = 0$ , i.e.,  $\frac{\delta F[q_t]}{\delta q_t} (x)\equiv$  constant.
Note, that in the same way, we can consider the continuity equation and the growth term separately, which defines the gradient flows under 2-Wasserstein and Fisher-Rao metrics respectively. The corresponding PDEs are
$$
\frac {\partial q _ {t}}{\partial t} (x) = - \nabla_ {x} \cdot \left(q _ {t} (x) \left(- \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x)\right)\right), \tag {67}
$$
$$
\frac {\partial q _ {t}}{\partial t} (x) = - \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right) q _ {t} (x). \tag {68}
$$
B Imaginary-time Schrödinger equation as the gradient flow under Fisher-Rao Metric
Theorem. Eq. (12) defines the gradient flow of the energy functional  $\mathbb{E}[q]$  under the Fisher-Rao metric.
Proof. First, we derive the functional derivative of the energy functional  $E[q]$ . We denote the differential of the functional  $F(q)$  along the direction  $h$  as
$$
\mathrm {d} F (q) [ h ] = \int \mathrm {d} x h \cdot \frac {\delta F [ q ]}{\delta q}, \tag {69}
$$
where  $\delta F[q] / \delta q$  is the functional derivative.
Consider the energy functional
$$
E [ q ] = \int \mathrm {d} x q \left[ V - \frac {1}{4} \nabla_ {x} ^ {2} \log q - \frac {1}{8} \| \nabla_ {x} \log q \| ^ {2} \right]. \tag {70}
$$
The functional derivative of this functional is as follows
$$
\begin{array}{l} \mathrm {d} E (q) [ h ] = \frac {\partial E (q + \varepsilon \cdot h)}{\partial \varepsilon} \Bigg | _ {\varepsilon = 0} \\ = \int \mathrm {d} x h \left[ V - \frac {1}{4} \nabla_ {x} ^ {2} \log q - \frac {1}{8} \| \nabla_ {x} \log q \| ^ {2} \right] - \int \mathrm {d} x q \left[ \frac {1}{4} \nabla_ {x} ^ {2} \frac {h}{q} + \frac {1}{4} \langle \nabla_ {x} \log q, \nabla_ {x} \frac {h}{q} \rangle \right]. \\ \end{array}
$$
For the last term, we do integration by parts and get
$$
\int \mathrm {d} x q \left[ \frac {1}{4} \nabla_ {x} ^ {2} \frac {h}{q} + \frac {1}{4} \langle \nabla_ {x} \log q, \nabla_ {x} \frac {h}{q} \rangle \right] = - \frac {1}{4} \int \mathrm {d} x \left\langle \nabla_ {x} q, \nabla_ {x} \frac {h}{q} \right\rangle + \frac {1}{4} \int \mathrm {d} x \left\langle \nabla_ {x} q, \nabla_ {x} \frac {h}{q} \right\rangle = 0. \tag {71}
$$
$$
\mathrm {d} E (q) [ h ] = \int \mathrm {d} x h \underbrace {\left[ V - \frac {1}{4} \nabla_ {x} ^ {2} \log q - \frac {1}{8} \| \nabla_ {x} \log q \| ^ {2} \right]} _ {\delta E [ q ] / \delta q}, \tag {72}
$$
and we see that the derivative coincides with the local energy, i.e.,
$$
\frac {\delta E [ q ]}{\delta q} (x) = E _ {\mathrm {l o c}} (x) = V (x) - \frac {1}{4} \nabla_ {x} ^ {2} \log q (x) - \frac {1}{8} \| \nabla_ {x} \log q (x) \| ^ {2}. \tag {73}
$$
Using the results from Section 2.2, the energy-minimizing gradient flow under Fisher-Rao metric is
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \left[ E _ {\mathrm {l o c}} (x) - E [ q _ {t} ] \right] q _ {t} (x). \tag {74}
$$
Second, we derive the PDE for the time-evolution of the density  $q_{t}$  under the imaginary-time Schrödinger equation.
$$
\frac {\partial \psi_ {t}}{\partial t} = \frac {1}{2} \nabla_ {x} ^ {2} \psi_ {t} - (V - E [ q _ {t} ]) \psi_ {t} \tag {75}
$$
$$
2 \psi_ {t} \frac {\partial \psi_ {t}}{\partial t} = \psi_ {t} \nabla_ {x} ^ {2} \psi_ {t} - 2 (V - E [ q _ {t} ]) \psi_ {t} ^ {2} \tag {76}
$$
$$
\frac {\partial q _ {t}}{\partial t} = \psi_ {t} \nabla_ {x} ^ {2} \psi_ {t} - 2 (V - E [ q _ {t} ]) q _ {t} \tag {78}
$$
Using the identity
$$
\begin{array}{l} \psi \nabla_ {x} ^ {2} \psi = \psi \nabla_ {x} \cdot \left(\psi \nabla_ {x} \log | \psi |\right) = \langle \psi \nabla_ {x} \psi , \nabla_ {x} \log | \psi | \rangle + \psi^ {2} \nabla_ {x} ^ {2} \log | \psi | (79) \\ = \frac {1}{4} \left\langle \nabla_ {x} q, \nabla_ {x} \log q \right\rangle + \frac {1}{2} q \nabla_ {x} ^ {2} \log q = \frac {1}{4} q \| \nabla_ {x} \log q \| ^ {2} + \frac {1}{2} q \nabla_ {x} ^ {2} \log q, (80) \\ \end{array}
$$
$$
\frac {\partial q _ {t}}{\partial t} = - 2 \left[ V - \frac {1}{4} \nabla_ {x} ^ {2} \log q - \frac {1}{8} \left\| \nabla_ {x} \log q _ {t} \right\| ^ {2} - E [ q _ {t} ] \right] q _ {t} \tag {81}
$$
$$
\frac {\partial q _ {t} (x)}{\partial t} = - 2 \left[ E _ {\mathrm {l o c}} (x) - E [ q _ {t} ] \right] q _ {t} (x), \tag {82}
$$
which is equivalent to Eq. (74).
C Analogies to generative modeling and variational inference literature
To draw a connection to generative models and variational inference literature, we consider the reverse and forward KL-divergences as an objective to minimize (instead of the energy) and derive projected gradient flows for them.
C.1 Generative modeling: reverse KL divergence functional
Consider the reverse KL divergence objective, and its first variation
$$
\mathcal {F} [ q _ {t} ] = D _ {\mathrm {K L}} (p \| q _ {t}) \tag {83}
$$
$$
\Rightarrow \frac {\delta \mathcal {F}}{\delta q _ {t}} = - \frac {p}{q _ {t}}, \tag {84}
$$
where  $p(x)$  is the data distribution given empirically.
Fisher-Rao Gradient Flow Using equations Eq. (10), the PDE that defines the Fisher-Rao gradient flow is
$$
\frac {\partial q _ {t}}{\partial t} (x) = \left(\frac {p (x)}{q _ {t} (x)} - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {p (y)}{q _ {t} (y)} \right]\right) q _ {t} (x) = p (x) - q _ {t} (x). \tag {85}
$$
$$
q _ {t} = \exp (- t) q _ {0} + (1 - \exp (- t)) p, \tag {86}
$$
where  $q_{0}$  is the initial distribution. This results in a distributional path constructed by the linear average (mixture) of the end-point distributions. Note that changing the mixture parameter corresponds to teleporting the particles.
Projected Fisher-Rao Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding to the projected Fisher-Rao gradient flow using Proposition 3.2 and Eq. (17),
$$
\begin{array}{l} \Delta \theta^ {*} \propto - \int q _ {t} \left(\frac {\delta \mathcal {F}}{\delta q _ {t}}\right) \nabla_ {\theta} \log q (x, \theta) d x (87) \\ = - \int q _ {t} \left(- \frac {p}{q _ {t}}\right) \nabla_ {\theta} \log q (x, \theta) d x (88) \\ = \nabla_ {\theta} \mathbb {E} _ {p (x)} \left[ \log q (x, \theta) \right]. (89) \\ \end{array}
$$
This corresponds to the standard maximum likelihood objective that is used for training of energy-based models (Xie et al., 2016; Du & Mordatch, 2019).
Wasserstein Gradient Flow Using the equation Eq. (9), the PDE that defines the Wasserstein gradient flow is
$$
\frac {\partial q _ {t}}{\partial t} (x) = - \nabla_ {x} \cdot \left(q _ {t} (x) \nabla_ {x} \frac {p (x)}{q _ {t} (x)}\right). \tag {90}
$$
Projected Wasserstein Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding to the projected Wasserstein gradient flow using Proposition 3.2 and Eq. (17).
$$
\begin{array}{l} \Delta \theta^ {*} \propto - \nabla_ {\theta} \int q _ {t} \left\langle \nabla_ {x} \frac {\delta \mathcal {F}}{\delta q _ {t}}, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x (91) \\ = \nabla_ {\theta} \int q _ {t} \left\langle \nabla_ {x} \frac {p}{q _ {t}}, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x (92) \\ = \nabla_ {\theta} \int q _ {t} \left\langle \frac {1}{q _ {t}} \nabla_ {x} p, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x + \nabla_ {\theta} \int q _ {t} \left\langle p \frac {- 1}{q _ {t} ^ {2}} \nabla q _ {t}, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x (93) \\ = \nabla_ {\theta} \int p \left\langle \nabla_ {x} \log p, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x - \nabla_ {\theta} \int p \left\langle \nabla_ {x} \log q _ {t}, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x (94) \\ = - \frac {1}{2} \nabla_ {\theta} \mathbb {E} _ {p (x)} \left[ \left\| \nabla_ {x} \log p (x) - \nabla_ {x} \log q (x, \theta) \right\| ^ {2} \right]. (95) \\ \end{array}
$$
Note that similar to Proposition 3.2, we use  $q_{t}(x) = q(x,\theta)$  as the density equal to the model density but detached from the parameters  $\theta$ . This objective corresponds to the score-matching objective (Hyvärinen & Dayan, 2005), which is widely used in the diffusion-based generative models (Song et al., 2020).
C.2 Variational inference: forward KL divergence functional
Consider the forward KL divergence objective, and its first variation
$$
\mathcal {F} [ q _ {t} ] = D _ {\mathrm {K L}} \left(q _ {t} \| p\right) \tag {96}
$$
$$
\Longrightarrow \frac {\delta \mathcal {F}}{\delta q _ {t}} = 1 + \log \frac {q _ {t}}{p}. \tag {97}
$$
where  $p(x)$  is the target distribution for variational inference.
$$
\begin{array}{l} \frac {\partial}{\partial t} q _ {t} = - \operatorname {g r a d} _ {\mathrm {F R}} \mathcal {F} (98) \\ = q _ {t} \left(- \frac {\delta \mathcal {F}}{\delta q _ {t}} - \mathbb {E} _ {q _ {t}} \left[ - \frac {\delta \mathcal {F}}{\delta q _ {t}} \right]\right) (99) \\ = q _ {t} \left(- 1 - \log \frac {q _ {t}}{p} + \mathbb {E} _ {q _ {t}} \left[ 1 + \log \frac {q _ {t}}{p} \right]\right) (100) \\ = q _ {t} \left(- \log \frac {q _ {t}}{p} + D _ {\mathrm {K L}} \left(q _ {t} \| p\right)\right). (101) \\ \end{array}
$$
$$
\begin{array}{l} \frac {\partial}{\partial t} \log q _ {t} = \log p - \log q _ {t} + D _ {\mathrm {K L}} \left(q _ {t} \| p\right) (102) \\ \Rightarrow \log q _ {t} = \exp (- t) \log q _ {0} + (1 - \exp (- t)) \log p - \log \mathcal {Z} _ {t} (103) \\ \Longrightarrow \boxed {q _ {t} = \frac {1}{\mathcal {Z} _ {t}} q _ {0} ^ {\exp (- t)} p ^ {(1 - \exp (- t))}} \quad \mathcal {Z} _ {t} = \int d x q _ {0} ^ {\exp (- t)} p ^ {(1 - \exp (- t))}, (104) \\ \end{array}
$$
where  $q_{0}$  is the initial distribution, and  $\mathcal{Z}_{t}$  is the partition function. This results in a distributional path constructed by the geometric average of the end-point distributions. Note that changing the geometric parameter corresponds to teleporting particles. These distributional paths are commonly used in Annealed Importance Sampling (Neal, 2001).
$$
\begin{array}{l} \Delta \theta^ {*} \propto - \int q _ {t} \left(\frac {\delta \mathcal {F}}{\delta q _ {t}}\right) \nabla_ {\theta} \log q (x, \theta) d x (105) \\ = - \int q _ {t} \left(1 + \log \frac {q _ {t}}{p}\right) \nabla_ {\theta} \log q (x, \theta) d x (106) \\ = - \int \log \frac {q _ {t}}{p} \nabla_ {\theta} q (x, \theta) d x (107) \\ = - \int \log q _ {t} \nabla_ {\theta} q (x, \theta) \mathrm {d} x + \nabla_ {\theta} \int q (x, \theta) \log p \mathrm {d} x (108) \\ = - \nabla_ {\theta} \int q (x, \theta) \log q (x, \theta) d x + \nabla_ {\theta} \int q (x, \theta) \log p d x (109) \\ = - \nabla_ {\theta} \int q (x, \theta) \log \frac {q (x , \theta)}{p} d x (110) \\ = - \nabla_ {\theta} D _ {\mathrm {K L}} \left(q (x, \theta) \| p\right). (111) \\ \end{array}
$$
This corresponds to the standard variational inference objective.
$$
\begin{array}{l} \frac {\partial}{\partial t} q _ {t} = - \operatorname {g r a d} _ {\mathrm {w}} \mathcal {F} (112) \\ = \nabla_ {x} \cdot \left(q _ {t} \nabla_ {x} \left(\frac {\delta \mathcal {F}}{\delta q _ {t}}\right)\right) (113) \\ = \nabla_ {x} \cdot \left(q _ {t} \nabla_ {x} \left(1 + \log \frac {q _ {t}}{p}\right)\right) (114) \\ = \nabla_ {x} \cdot \left(q _ {t} \nabla_ {x} (\log q _ {t} - \log p)\right) (115) \\ = - \nabla_ {x} \cdot \left(q _ {t} \nabla_ {x} \log p\right) + \nabla \cdot \left(q _ {t} \nabla_ {x} \log q _ {t}\right) (116) \\ = - \nabla_ {x} \cdot \left(q _ {t} \nabla_ {x} \log p\right) + \nabla_ {x} ^ {2} q _ {t} (117) \\ \end{array}
$$
This is the Fokker-Planck equation, which characterizes the movement of the particles according to Langevin dynamics:
$$
\mathrm {d} X _ {t} = \nabla_ {x} \log p \mathrm {d} t + \sqrt {2} \mathrm {d} W _ {t}.
$$
This is a common approach for sampling from the unnormalized density  $p(x)$  via Markov Chain Monte Carlo algorithms.
Projected Wasserstein Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding to the projected Wasserstein gradient flow using Proposition 3.2 and Eq. (17),
$$
\begin{array}{l} \Rightarrow \Delta \theta^ {*} \propto - \nabla_ {\theta} \int q _ {t} \left\langle \nabla_ {x} \frac {\delta \mathcal {F}}{\delta q _ {t}}, \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x (118) \\ = - \nabla_ {\theta} \int q _ {t} \left\langle \nabla_ {x} \left(1 + \log \frac {q _ {t}}{p}\right), \nabla_ {x} \log q (x, \theta) \right\rangle d x (119) \\ = - \nabla_ {\theta} \int q _ {t} \left\langle \nabla_ {x} \left(\log \frac {q _ {t}}{p}\right), \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x (120) \\ = - \frac {1}{2} \mathbb {E} _ {q (x, \theta)} \left[ \nabla_ {\theta} \left\| \nabla_ {x} \log q (x, \theta) - \nabla_ {x} \log p \right\| ^ {2} \right], (121) \\ \end{array}
$$
which resembles the score-matching objective (Hyvarinen & Dayan, 2005) for variational inference.
D c-Wasserstein gradient flow
$c$ -Wasserstein distance with the convex cost function  $c: \mathbb{R}^d \to \mathbb{R}$  is defined
$$
W _ {c} \left(p _ {0}, p _ {1}\right) = \inf  _ {\pi \in \Gamma \left(p _ {0}, p _ {1}\right)} \int \pi (x, y) c (x - y) d x d y, \tag {122}
$$
where  $\Gamma(p_0, p_1)$  is the set of all possible couplings of  $p_0$  and  $p_1$ . The dynamic formulation of this distance is the following
$$
W _ {c} \left(p _ {0}, p _ {1}\right) := \inf  _ {v _ {t}, q _ {t}} \int_ {0} ^ {1} \mathbb {E} _ {q _ {t} (x)} [ c (v _ {t} (x)) ] \mathrm {d} t, \quad \text {s u b j .} \tag {123}
$$
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right), \quad \text {a n d} \quad q _ {0} = p _ {0}, \quad q _ {1} = p _ {1}. \tag {124}
$$
Proposition. The energy-minimizing  $c$ -Wasserstein gradient flow is defined by the following PDE
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) \nabla c ^ {*} \left(- \nabla_ {x} E _ {\mathrm {l o c}} (x)\right)\right), \tag {125}
$$
where  $c^*(\cdot)$  is the convex conjugate function of  $c(\cdot)$ .
Proof. The movement minimizing scheme for  $W_{c}(\cdot, \cdot)$  is the following optimization problem
$$
\inf  _ {q ^ {\prime}} F [ q ^ {\prime} ] - F [ q ] + \frac {1}{\Delta t} W _ {c} (q, q ^ {\prime}). \tag {126}
$$
Assuming that the density changes according to the continuity equation  $q'(x) = q_t(x) - \Delta t \nabla_x \cdot (q_t(x) v_t(x)) + o(\Delta t)$ , and  $\Delta t$  is small enough so that  $v_t(x)$  defines the optimal transportation plan, we have
$$
\begin{array}{l} \inf  _ {v _ {t}} F [ q _ {t} ] - \Delta t \int \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) d x - F [ q _ {t} ] + \frac {1}{\Delta t} \Delta t ^ {2} \mathbb {E} _ {q _ {t} (x)} c (v _ {t} (x)) (127) \\ = \Delta t \inf  _ {v _ {t}} \int q _ {t} (x) \left\langle v _ {t} (x), \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\rangle \mathrm {d} x + \mathbb {E} _ {q _ {t} (x)} c (v _ {t} (x)) (128) \\ = \Delta t \inf  _ {v _ {t}} \int q _ {t} (x) \left[ c \left(v _ {t} (x)\right) - \left\langle v _ {t} (x), - \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\rangle \right] d x (129) \\ = - \Delta t \int q _ {t} (x) c ^ {*} \left(- \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x)\right) d x (130) \\ \end{array}
$$
and the infimum is achieved at
$$
v _ {t} (x) = \nabla c ^ {*} \left(- \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x)\right), \tag {131}
$$
which gives the formula for the vector field. Using the energy gradient from Theorem 3.1  $\frac{\delta E[q_t]}{\delta q_t} (x) = E_{\mathrm{loc}}$ , we get the result.
Proposition. Coordinate-wise application of tanh to the vector field, i.e.
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) \tanh  \left(- \nabla_ {x} E _ {\mathrm {l o c}} (x)\right)\right), \tag {132}
$$
corresponds to gradient descent with  $c$ -Wasserstein distance, where  $c: \mathbb{R}^d \to \mathbb{R}$  is the following cost function
$$
c (x) = \sum_ {i} ^ {d} \frac {1}{2} \left(\left(x _ {i} + 1\right) \log \left(x _ {i} + 1\right) + \left(1 - x _ {i}\right) \log \left(1 - x _ {i}\right)\right) - d \log 2. \tag {133}
$$
Proof. Consider  $c^*(x) = \sum_i \log(\exp(x_i) + \exp(-x_i))$ . It corresponds to applying hyperbolic tangent non-linearity coordinate-wise to the vector field field, i.e.,
$$
\partial_ {i} c ^ {*} (x) = \frac {\exp \left(x _ {i}\right) - \exp \left(- x _ {i}\right)}{\exp \left(x _ {i}\right) + \exp \left(- x _ {i}\right)} = \tanh  \left(x _ {i}\right). \tag {134}
$$
The corresponding cost function  $c(x)$  is the following
$$
\begin{array}{l} c (x) = \sup  _ {y} \langle x, y \rangle - c ^ {*} (y) (135) \\ = \sup  _ {y} \sum_ {i} ^ {d} \left(x _ {i} y _ {i} - \log \left(\exp \left(y _ {i}\right) + \exp \left(- y _ {i}\right)\right)\right) (136) \\ = \sup  _ {y} \sum_ {i} ^ {d} \left((x _ {i} + 1) y _ {i} - \log \left(\exp \left(2 y _ {i}\right) + 1\right)\right) / / y _ {i} = \frac {1}{2} \log \frac {1 + x}{1 - x} (137) \\ = \sum_ {i} ^ {d} \left(\left(x _ {i} + 1\right) \frac {1}{2} \log \left(\frac {1 + x _ {i}}{1 - x _ {i}}\right) - \log \left(\frac {1 + x _ {i}}{1 - x _ {i}} + 1\right)\right) (138) \\ = \sum_ {i} ^ {d} \frac {1}{2} \left(\left(x _ {i} + 1\right) \log \left(x _ {i} + 1\right) + \left(1 - x _ {i}\right) \log \left(1 - x _ {i}\right)\right) - d \log 2. (139) \\ \end{array}
$$
Figure 3: Optimization results for different chemical systems (every column corresponds to a given molecule). The number of electrons is given in the brackets next to systems' names. Throughout the optimization, we monitor three values: the mean value of the local energy (lower is better), the variance of the local energy, and the median value of the gradient norm of the local energy. In the first row of plots, we average (removing  $5\%$  of outliers from both sides) the energy over 1000 iterations and report the relative error to the actual ground-state energy:  $(E - E_0) / E_0$ . In the second row, we report standard deviation averaged over 1000 iterations (removing  $5\%$  of outliers from both sides). In the third row, we report the median gradient norm averaged over 1000 iterations (removing  $5\%$  of outliers from both sides). All the metrics are plotted versus the wall time of computations performed on four Nvidia A40 GPUs.