SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) \left(- \nabla_ {x} E _ {\mathrm {l o c}} (x)\right)\right) \tag {20}
$$
Proof. In Theorem 3.1, we show that  $\delta E[q] / \delta q = E_{\mathrm{loc}}(x)$ . Plugging this into the 2-Wasserstein gradient flow defined in Eq. (9), yields the result in Eq. (20).
$c$ -Wasserstein Metric This result can be further generalized to the  $c$ -Wasserstein metric with any convex cost function  $c: \mathbb{R}^d \to \mathbb{R}$  on the tangent space. The  $c$ -Wasserstein distance between  $p_0$  and  $p_1$  is defined as follows
$$
W _ {c} \left(p _ {0}, p _ {1}\right) := \inf  _ {v _ {t}, q _ {t}} \int_ {0} ^ {1} \mathbb {E} _ {q _ {t} (x)} \left[ c \left(v _ {t} (x)\right) \right] \mathrm {d} t, \quad \text {s u b j .} \tag {21}
$$
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right), \quad \text {a n d} q _ {0} = p _ {0}, \quad q _ {1} = p _ {1}. \tag {22}
$$
Proposition 3.5. The energy-minimizing  $c$ -Wasserstein gradient flow is defined by the following equation
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) \nabla c ^ {*} \left(- \nabla_ {x} E _ {\text {l o c}} (x)\right)\right), \tag {23}
$$
where  $c^*(\cdot)$  is the convex conjugate function of  $c(\cdot)$ , and  $\nabla c^*(y)$  is its gradient at  $y$ .
Theorem 3.4 can be viewed as a special case of Proposition 3.5 where  $c(\cdot) = \frac{1}{2} \| \cdot \| ^2$ . Introducing a different  $c$  than  $L^2$  norm translates to a non-linear transformation of the gradient  $-\nabla_{x}E_{\mathrm{loc}}(x)$ . In Appendix D, we demonstrate how to choose  $c$  such that it corresponds to the coordinate-wise application of tanh to the gradient, which we use in practice.
Finally, using Proposition 3.5 in Eq. (17), we get the expression for the parameter update, i.e.
$$
\Delta \theta^ {*} \propto \int q _ {t} (x) \nabla_ {\theta} \left\langle \nabla c ^ {*} (- \nabla_ {x} E _ {\mathrm {l o c}} (x)), \nabla_ {x} \log q (x, \theta) \right\rangle \mathrm {d} x. \tag {24}
$$
Similar to the discussion of the previous section for QVMC, we can precondition the gradient with the Fisher Information Matrix, exploiting the geometry of the parametric manifold.
In Algorithm 1, we provide a pseudocode for the proposed algorithms. The procedure follows closely QVMC but introduces a different objective. When using gradients both from Eqs. (18) and (24), we follow the gradient flow under the Wasserstein Fisher-Rao metric with the coefficient  $\lambda$ . For  $\lambda \to \infty$ , the cost of mass teleportation becomes infinite and we use only the gradient from Eq. (24), which corresponds to the gradient flow under the  $c$ -Wasserstein metric (we refer to this algorithm as WQMC). For  $\lambda \to 0$ , the cost of mass teleportation becomes negligible compared to the transportation cost and the resulting algorithm becomes QVMC, which uses the gradient from Eq. (18). In practice, we consider the extreme cases ( $\lambda \to 0$ ,  $\infty$ ) and the mixed case  $\lambda = 1$ .
Algorithm 1 W(FR)QMC
Require: samples  $\{x^{(i)}\}_{i = 1}^{N}\sim q_{t = 0}(x)$
Require: potential function  $V(x)$
$$
E _ {\text {l o c}} (x ^ {(i)}) = V (x ^ {(i)}) - \frac {1}{4} \nabla_ {x} ^ {2} \log q (x ^ {(i)}, \theta) - \frac {1}{8} \| \nabla_ {x} \log q (x ^ {(i)}, \theta) \| ^ {2} \tag {seeEq.5}
$$
$$
\nabla_ {x} E _ {\mathrm {l o c}} \left(x ^ {(i)}\right) = \operatorname {s t o p} _ {-} \operatorname {g r a d i e n t} \left(\nabla_ {x} E _ {\mathrm {l o c}} \left(x ^ {(i)}\right)\right)
$$
$$
\Delta \theta^ {*} = \frac {1}{N} \sum_ {i} ^ {N} \nabla_ {\theta} \left\langle \nabla c ^ {*} \left(- \nabla_ {x} E _ {\mathrm {l o c}} \left(x ^ {(i)}\right)\right), \nabla_ {x} \log q \left(x ^ {(i)}, \theta\right) \right\rangle \tag {seeEq.24}
$$
$$
E _ {\text {l o c}} \left(x ^ {(i)}\right) = \operatorname {s t o p} \text {g r a d i e n t} \left(E _ {\text {l o c}} \left(x ^ {(i)}\right)\right)
$$
$$
\Delta \theta^ {*} + = - \frac {1}{\lambda} \frac {1}{N} \sum_ {i} ^ {N} \left[ E _ {\mathrm {l o c}} \left(x ^ {(i)}\right) - \frac {1}{N} \sum_ {j} ^ {N} E _ {\mathrm {l o c}} \left(x ^ {(j)}\right) \right] \nabla_ {\theta} \log q \left(x ^ {(i)}, \theta\right) \tag {seeEq.18}
$$
$$
\theta^ {\prime} = \operatorname {o p t i m i z e r} \left(\theta , \mathcal {F} _ {\theta} ^ {- 1} \Delta \theta^ {*}\right) \tag {seeEq.19}
$$
update  $x^{(i)}$  by sampling from  $q(x,\theta^{\prime})$  via MCMC
return model  $q(x,\theta^{*})$  samples  $\{x^{(i)}\}_{i = 1}^{N}\sim q(x,\theta^{*})$
Figure 2: 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). See the descriptions of methods in the text.
For the empirical study of the proposed method, we consider Born-Oppenheimer approximation of chemical systems. Within this approximation, the wave function of the electrons in a molecule can be studied separately from the wave function of the atomic nuclei. Namely, we consider the following Hamiltonian
$$
H = - \frac {1}{2} \nabla_ {x} ^ {2} + \sum_ {i <   j} \frac {1}{\| x _ {i} - x _ {j} \|} - \sum_ {i, I} \frac {Z _ {I}}{\| x _ {i} - X _ {I} \|} + \sum_ {I <   J} \frac {Z _ {I} Z _ {J}}{\| X _ {I} - X _ {J} \|}, \tag {25}
$$
where  $x_{i}$  are the coordinates of electrons,  $X_{I}, Z_{I}$  are the coordinates and charges of nuclei. The first kinetic term contains derivatives with respect to the electron positions  $x$ . Indeed, the positions of the nuclei are given and fixed, and we target the ground state of the electronic wave function  $\psi(x)$ , which is an explicit function of the electron positions only. Solving the electronic Schrödinger equation is a notoriously difficult task, and is a topic of intense research in quantum chemistry and material sciences.
$^{1}$ Code reproducing experiments is available at github.com/necludov/wqmc. The method is also implemented in the FermiNet library: github.com/google-deepmind/ferminet.
Method name	QVMC	WQMC	W(FR)QMC	QVMC	WQMC	W(FR)QMC
Molecule	Be (4)	B (5)
Relative energy error	1.50e-5	3.79e-6	1.04e-6	9.01e-6	9.58e-6	2.69e-6
Energy variance (Ha2)	5.53e-3	1.08e-3	1.07e-3	1.96e-2	1.84e-2	1.19e-2
Molecule	Li2(6)	H10 (10)
Relative energy error	4.43e-5	3.71e-5	3.66e-5	4.24e-4	3.90e-4	3.88e-4
Energy variance (Ha2)	7.21e-3	5.76e-4	5.59e-4	1.89e-3	2.30e-4	2.45e-4
Table 1: Energy and variances estimates for all systems after 10k iterations (20k for  ${\mathrm{H}}_{10}$  ).
Since electrons are indistinguishable fermions, we restrict the Hilbert space to states  $\psi$  that are antisymmetric under electron permutations (see Section 2.1). This can be achieved by incorporating Slater determinants into the deep neural network, which parametrizes the wave function  $\psi(x,\theta)$ , as proposed in various recent works (Luo & Clark, 2019; Hermann et al., 2020; Pfau et al., 2020; Hermann et al., 2022). The density is then given by the Born rule  $q(x,\theta) = |\psi(x,\theta)|^2$ . For all our experiments, we follow (von Glehn et al., 2022) and use the "psiformer" architecture together with preconditioning the gradients via K-FAC (Martens & Grosse, 2015).
In our method, we apply several tricks which stabilize the optimization and improve convergence speed. Firstly, we have observed that applying a tanh non-linearity coordinate-wise to the gradient  $\nabla_{x}E_{\mathrm{loc}}(x)$  significantly improves convergence speed. This corresponds to a different cost function in the Wasserstein metric, as we discuss in Proposition 3.5 and Appendix D. Also, we remove samples from the batch whose norm  $\| \nabla_x\log q(x,\theta)\|$  significantly exceeds the median value. Namely, we estimate the deviation from the norm as  $\mathbb{E}_{q(x,\theta)}[||\nabla_x\log q(x,\theta)|| - \mathrm{median}(\| \nabla_x\log q(x,\theta)\|)]$  and remove samples whose norm exceeds five deviations from the median. When including the gradient from Eq. (18), we clip the local energy values as proposed in (von Glehn et al., 2022), i.e. by estimating the median value and clipping to five deviations from the median, where the deviation is estimated in the same way as for the norm of the gradient.
We consider different chemical systems and compare against QVMC as a baseline. We run our novel method with the same architecture and hyperparameters as the baseline QVMC-based approach in (von Glehn et al., 2022). For the chemical systems, we consider Be, and B atoms, the  $\mathrm{Li_2}$  molecule and the hydrogen chain  $\mathrm{H}_{10}$  from (Hermann et al., 2020). The exact values of energies for Be, B,  $\mathrm{Li_2}$  are taken from (Pfau et al., 2020), the exact value of the energy for  $\mathrm{H}_{10}$  is from (Hermann et al., 2020). All the hyperparameters and architectural details are provided in the supplementary material.
In Figure 2, we demonstrate the convergence plots for the baseline (QVMC) and the proposed methods (WQMC and W(FR)QM, see Algorithm 1). For all the considered systems, both WQMC and W(FR)QMC yield more precise estimations of the ground state energy (the first row of Figure 2). To assess convergence, we also monitor the variance of the local energy and the gradient norm of the local energy. As we discuss in Eq. (11), both metrics must vanish at the ground state. More fundamentally, the variance of the local energy can be shown to vanish for eigenstates of the Hamiltonian (Wu et al., 2023), referred to as the zero-variance property. First, we point out that obtaining the ground states of the considered molecules with QVMC is challenging, and even with powerful deep-learning architecture, discrepancies remain with the ground state. Since we use existing state-of-the-art architectures as a backbone, our results are also limited by the limitations of the latter (Gao & Gunnemann, 2023). Developing novel architectures is out of the scope of this work.
However, in Figure 2, we clearly observe that both WQMC and W(FR)QMC yield significantly faster convergence of the aforementioned metrics compared to QVMC. In particular, for the larger molecules  $\mathrm{Li}_2$  and  $\mathrm{H}_{10}$ , we observe that we consistently obtain lower energies within  $10\mathrm{k}$  steps (20k for  $\mathrm{H}_{10}$ ) with a more stable convergence. For the smaller molecules we observe that QVMC obtains lower energies in the first few iterations, but its convergence slows down significantly, after which our approach steadily yields improved energies below QVMC. Overall, our experiments demonstrate that taking into account the Wasserstein metric allows for faster convergence to accurate approximations of the ground state. See the final metrics in Table 1.
Since our method requires an extra gradient evaluation compared to QVMC, we include the runtime of all the algorithms in Figure 3, Appendix E. Namely, for a proper comparison, instead of reporting the metrics per iteration, we report them per wall time in seconds. Note that all the claims of the experimental section still hold in terms of wall time. All the models were benchmarked on four A40 GPUs. Third-order derivatives can be efficiently computed using modern deep learning frameworks such as JAX (Bradbury et al., 2018), which we used to implement our method.
Potentially, one can alleviate the extra cost of the iteration by coming up with more efficient Monte Carlo schemes or other updates of the samples. Indeed, in our experiments, we observed that the proposed method requires fewer MCMC steps (not included in the paper). Moreover, one can use the evaluated gradient of the local energy as the proposal vector field for updating the samples. This would allow to decrease the number of MCMC steps at the low cost of additional hyperparameter tuning.
Conclusion In the current paper, we propose a novel approach to solving the quantum many-body Schrödinger equation, by incorporating the Wasserstein metric on the space of Born distributions. Compared to the Fisher-Rao metric, which allows for probability mass "teleportation", the Wasserstein metric constrains the evolution of the density to local changes under the locally-optimal transportation plan, i.e., following fluid dynamics. This property is favorable when the evolution of the parametric model is accompanied by the evolution of samples (performed by an MCMC algorithm). Indeed, by forbidding or regularizing non-local mass "teleportation" in the density change, one prevents the appearance of distant modes in the density model, which would require longer MCMC runs for proper mixing of the samples.
In practice, we demonstrate that following the gradient flow under the Wasserstein (or Wasserstein Fisher-Rao) metric results in better convergence to the ground state wave function. This is expected to be due to our proposed loss, which takes into account the gradient of the local energy and achieves its minimum when the norm of the gradient vanishes, therefore explicitly minimizing the norm of the local energy gradient.
We believe that our new theoretical framework for solving the time-independent Schrödinger equation for time-reversal symmetric Hamiltonians based on optimal transport will open new avenues to develop improved numerical methods for quantum chemistry and physics.
Connection to Energy-Based and Score-Based Generative Models The developed ideas of this paper, i.e., projecting gradient flows under different metrics onto a parametric family, can be extended to generative modeling by swapping the energy functional with the KL-divergence. More precisely, as we show in Appendix C, using the KL-divergence as our objective functional, the Fisher-Rao gradient flow yields energy-based training scheme, while the 2-Wasserstein gradient flow corresponds to the score-matching, which is used for training diffusion generative models.
The authors thank Rob Brekelmans and anonymous reviewers for helpful discussions and feedback. J.N. was supported by Microsoft Research. J.C. acknowledges support from the Natural Sciences and Engineering Research Council (NSERC), the Shared Hierarchical Academic Research Computing Network (SHARCNET), Compute Canada, and the Canadian Institute for Advanced Research (CIFAR) AI Chairs program. A.M. acknowledges support from the Canada CIFAR AI Chairs program. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute www.vectorinstitute.ai/#partners.
McMillan, W. L. Ground state of liquid he 4. Physical Review, 138(2A):A442, 1965.
Neal, R. M. Annealed importance sampling. Statistics and computing, 11:125-139, 2001.
Pescia, G., Han, J., Lovato, A., Lu, J., and Carleo, G. Neural-network quantum states for periodic systems in continuous space. Physical Review Research, 4(2):023138, 2022.
Pescia, G., Nys, J., Kim, J., Lovato, A., and Carleo, G. Message-passing neural quantum states for the homogeneous electron gas, 2023.
Pfau, D., Spencer, J. S., Matthews, A. G., and Foulkes, W. M. C. Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical Review Research, 2(3):033429, 2020.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
Sorella, S. Green function Monte Carlo with stochastic reconfiguration. Phys. Rev. Lett., 80:4558-4561, May 1998. doi: 10.1103/PhysRevLett.80.4558. URL https://link.aps.org/doi/10.1103/PhysRevLett.80.4558.
Stokes, J., Izaac, J., Killoran, N., and Carleo, G. Quantum natural gradient. Quantum, 4:269, 2020.
Vicentini, F., Hofmann, D., Szabo, A., Wu, D., Roth, C., Giuliani, C., Pescia, G., Nys, J., Vargas-Calderon, V., Astrakhantsev, N., et al. Netket 3: Machine learning toolbox for many-body quantum systems. SciPost Physics Codebases, pp. 007, 2022.
von Glehn, I., Spencer, J. S., and Pfau, D. A self-attention ansatz for ab-initio quantum chemistry. arXiv preprint arXiv:2211.13672, 2022.
White, S. R. Density matrix formulation for quantum renormalization groups. Physical review letters, 69(19):2863, 1992.
Wu, D., Rossi, R., Vicentini, F., Astrakhantsev, N., Becca, F., Cao, X., Carrasquilla, J., Ferrari, F., Georges, A., Hibat-Allah, M., et al. Variational benchmarks for quantum many-body problems. arXiv preprint arXiv:2302.04919, 2023.
Xie, J., Lu, Y., Zhu, S.-C., and Wu, Y. A theory of generative convnet. In International Conference on Machine Learning, pp. 2635-2644. PMLR, 2016.
Yuan, X., Endo, S., Zhao, Q., Li, Y., and Benjamin, S. C. Theory of variational quantum simulation. Quantum, 3:191, 2019.
A Gradient flows under Wasserstein Fisher-Rao metric
We, first, remind the concept of the functional derivative. The change of the functional  $F[q]: \mathcal{P}_2 \to \mathbb{R}$  along the direction  $h$  can be expressed as
$$
F [ q + h ] = F [ q ] + \mathrm {d} F [ h ] + o (\| h \|), \quad \mathrm {d} F [ h ] = \int \mathrm {d} x h (x) \underbrace {\frac {\delta F [ q ]}{\delta q} (x)} _ {\text {d e r i v a t i v e}}. \tag {26}
$$
Consider a change of the density in time, the change of the functional can be defined through the differential as
$$
F \left[ q _ {t} + \Delta t \frac {\partial q _ {t}}{\partial t} \right] = F [ q _ {t} ] + \Delta t \cdot \mathrm {d} F \left[ \frac {\partial q _ {t}}{\partial t} \right] + o (\| h \|), \quad \mathrm {d} F \left[ \frac {\partial q _ {t}}{\partial t} \right] = \int \mathrm {d} x \frac {\partial q _ {t} (x)}{\partial t} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x). \tag {27}
$$
$$
\frac {\mathrm {d}}{\mathrm {d} t} F [ q _ {t} ] = \mathrm {d} F \left[ \frac {\partial q _ {t}}{\partial t} \right] = \int \mathrm {d} x \frac {\partial q _ {t} (x)}{\partial t} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x). \tag {28}
$$
A.1 Minimizing movement scheme for 2-Wasserstein distance and Kullback-Leibler divergence
Gradient flow under  $\mathbf{W}_2$  Consider the following minimizing movement scheme (MMS)
$$
\inf  _ {q ^ {\prime}} F [ q ^ {\prime} ] - F [ q ] + \frac {1}{\Delta t} \frac {1}{2} W _ {2} ^ {2} (q, q ^ {\prime}), \tag {29}
$$
where the change of the density is restricted to the continuity equation, i.e.,
$$
\frac {\partial q _ {t}}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right), \quad \text {a n d} \quad q ^ {\prime} (x) = q _ {t} (x) - \Delta t \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) + o (\Delta t). \tag {30}
$$
Using the static formulation of  $W_{2}$  distance, we have
$$
W _ {2} ^ {2} (q, q ^ {\prime}) = \int \mathrm {d} x q (x) \| x - T ^ {*} (x) \| ^ {2} = \Delta t ^ {2} \int \mathrm {d} x q (x) \| v ^ {*} (x) \| ^ {2}, \tag {31}
$$
where  $T^{*}(x)$  is the optimal transportation plan, and  $v^{*}(x)$  is the corresponding optimal gradient field. Thus, we can rewrite the MMS problem as
$$
\inf  _ {v} F [ q ] - \Delta t \int \mathrm {d} x \nabla_ {x} \cdot (q (x) v (x)) \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - F [ q ] + \frac {\Delta t}{2} \int \mathrm {d} x q (x) \| v (x) \| ^ {2}, \tag {32}
$$
$$
\inf  _ {v} \int \mathrm {d} x q (x) \left\langle v (x), \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\rangle + \frac {1}{2} \int \mathrm {d} x q (x) \| v (x) \| ^ {2}, \tag {33}
$$
$$
\inf  _ {v} \int \mathrm {d} x q (x) \left\| v (x) + \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\| ^ {2}. \tag {34}
$$
From the last optimization problem, we have
$$
v (x) = - \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x). \tag {35}
$$
Gradient flow under KL Consider the following minimizing movement scheme (MMS)
$$
\inf  _ {q ^ {\prime}} F [ q ^ {\prime} ] - F [ q ] + \frac {1}{\Delta t} D _ {\mathrm {K L}} \left(q ^ {\prime} \| q\right), \tag {36}
$$
where the change of the density is restricted to the following weighting scheme
$$
\frac {\partial q _ {t}}{\partial t} = g _ {t} (x) q _ {t} (x), \quad \text {h e n c e}, \tag {37}
$$
$$
q ^ {\prime} (x) = q _ {t} (x) + \Delta t q _ {t} (x) g _ {t} (x) + o (\Delta t), \quad \text {a n d} \tag {38}
$$
$$
\log q ^ {\prime} (x) = \log q _ {t} (x) + \Delta t g _ {t} (x) - \frac {\Delta t ^ {2}}{2} g _ {t} ^ {2} (x) + o (\Delta t ^ {2}). \tag {39}
$$
The KL-divergence is then
$$
\begin{array}{l} D _ {\mathrm {K L}} \left(q ^ {\prime} \| q _ {t}\right) = \int \mathrm {d} x q ^ {\prime} (x) \left(\Delta t g _ {t} (x) - \frac {\Delta t ^ {2}}{2} g _ {t} ^ {2} (x)\right) + o \left(\Delta t ^ {2}\right) (40) \\ = \Delta t \int \mathrm {d} x q _ {t} (x) g _ {t} (x) - \frac {\Delta t ^ {2}}{2} \int \mathrm {d} x q _ {t} (x) g _ {t} ^ {2} (x) (41) \\ + \Delta t ^ {2} \int \mathrm {d} x q _ {t} (x) g _ {t} ^ {2} (x) + o (\Delta t ^ {2}) (42) \\ = \frac {\Delta t ^ {2}}{2} \int \mathrm {d} x q _ {t} (x) g _ {t} ^ {2} (x) + o (\Delta t ^ {2}). (43) \\ \end{array}
$$
In the last equation we are using
$$
\int q _ {t} g _ {t} \mathrm {d} x = 0. \tag {44}
$$
Thus, we can rewrite the MMS problem as
$$
\inf  _ {g} F [ q _ {t} ] + \Delta t \int \mathrm {d} x g _ {t} (x) q _ {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) g _ {t} (x) ^ {2}, \tag {45}
$$
$$
\inf  _ {g} \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) g _ {t} (x) ^ {2}, \tag {46}
$$
$$
\left. \inf  _ {g} \int \mathrm {d} x q _ {t} (x) \left[ 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)\right) \right] ^ {2}. \tag {47}
$$
$$
g _ {t} (x) = - \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right). \tag {48}
$$
Note, however, that  $D_{\mathrm{KL}}(\cdot \| \cdot)$  is not the same as the Fisher-Rao metric. The derivations here demonstrate that the Fisher-Rao gradient flow can be derived as the MMS scheme with the KL-divergence.
A.2 Minimizing movement scheme for the Wasserstein Fisher-Rao metric
Consider the Wasserstein Fisher-Rao distance
$$
\operatorname {W F R} _ {\lambda} \left(p _ {0}, p _ {1}\right) ^ {2} := \inf  _ {v _ {t}, g _ {t}, q _ {t}} \int_ {0} ^ {1} \mathbb {E} _ {q _ {t} (x)} \left[ \| v _ {t} (x) \| ^ {2} + \lambda g _ {t} (x) ^ {2} \right] \mathrm {d} t, \quad \text {s u b j . t o} \tag {49}
$$
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) + g _ {t} (x) q _ {t} (x), \quad q _ {0} (x) = p _ {0} (x), \quad q _ {1} (x) = p _ {1} (x). \tag {50}
$$
The minimizing movement scheme (MMS) for this distance is
$$
\inf  _ {q ^ {\prime}} F \left[ q ^ {\prime} \right] - F [ q ] + \frac {1}{\Delta t} \frac {1}{2} \mathrm {W F R} _ {\lambda} \left(q, q ^ {\prime}\right) ^ {2}, \tag {51}
$$
where the change of the density is given by the continuity equation with the growth term
$$
\frac {\partial q _ {t}}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) + g _ {t} (x) q _ {t} (x). \tag {52}
$$
For close enough  $q$  and  $q'$ ,  $\mathrm{WFR}_{\lambda}(q, q')^2$  can be estimated via the metric derivative  $|\mu_t'|^2$  that is defined as
$$
\left| \mu_ {t} ^ {\prime} \right| ^ {2} = \left(\lim  _ {\Delta t \rightarrow 0} \frac {\operatorname {W F R} _ {\lambda} \left(q _ {t} , q _ {t + \Delta t}\right)}{\Delta t}\right) ^ {2}, \tag {53}
$$