text:
\section{Additional details for information geometry}

In this section we provide additional information regarding information geometry. We note that many of these proposition are already know in the literature, however, we include them for completion and for the paper to be standalone.

The Fisher-Rao metric is positive definite only if it is non-singular, and then, defines a Riemannian metric \citep{nielsen:2020}. In this paper, we assume that the observation $\vec{x}\in\mathcal{X}$ is a random variable following a probability distribution $p(\vec{x})$ such that $\vec{x}\sim p(\vec{x}|\eta)$, and any smooth changes of the parameter $\eta$ would alter the observation $\vec{x}$. This way, the Fisher-Rao metric used in our paper is non-singular and the statistical manifold $\mathcal{H}$ is a Riemannian manifold.

A known result in \emph{information geometry} \citep{nielsen:2020, amari:2016} is that the Fisher-Rao metric is the first order approximation of the KL-divergence, as recall in Proposition \ref{prop:FRapproxKL}. Using this fact, we can define the Fisher-Rao distance and energy in function of the KL-divergence, leading to Proposition \ref{prop:FRenergydistance}.

\begin{proposition}\label{prop:FRapproxKL}
The Fisher-Rao metric is the first order approximation of the KL-divergence between perturbed distributions:
\begin{equation}
\mathrm{KL}(p(\vec{x}|\eta),~p(\vec{x}|\eta+\delta\eta)) = \frac{1}{2} \delta\eta^{\top} \mathbf{I}_\HH(\eta) \delta\eta + o(\delta\eta^2),\nonumber
\end{equation}
with $\mathbf{I}_\HH(\eta) = \int p(\vec{x}|\eta) \left[\nabla_\eta\log p(\vec{x}|\eta) \nabla_\eta\log p(\vec{x}|\eta)\T\right] \dif\vec{x}.$
\end{proposition}

\begin{proof}
Let's decompose $\log p(\vec{x}|\eta+\delta\eta)$ using the Taylor expansion:
\[\log p(\vec{x}|\eta+\delta\eta) = \log p(\vec{x}|\eta) + \nabla_\eta \log p(\vec{x}|\eta)\T \delta\eta + \frac{1}{2} \delta\eta\T \text{Hess}_\eta \left[\log p(\vec{x}|\eta)\right] \delta\eta + o(\delta\eta^2),\]
where the Hessian is $\text{Hess}_\eta \left[\log p(\vec{x}|\eta)\right] = \frac{\text{Hess}_\eta [p(\vec{x}|\eta)]}{p(\vec{x}|\eta)} - \nabla_\eta\log p(\vec{x}|\eta) \nabla_\eta\log p(\vec{x}|\eta)\T$ and the $\nabla_\eta \log p(\vec{x}|\eta) = \frac{\nabla_\eta p(\vec{x}|\eta)}{p(\vec{x}|\eta)}$.

Also $\int \nabla_\eta p(\vec{x}|\eta) \dif \vec{x} = \nabla_\eta\int p(\vec{x}|\eta) \dif\vec{x} = 0$ and $\int \text{Hess}_\eta[p(\vec{x}|\eta)] \dif\vec{x} = \text{Hess}_\eta[\int p(\vec{x}|\eta) \dif \vec{x}] = 0.$

Replacing all those expressions to the first equation finally gives:
\begin{align*}
     \mathrm{KL}(p(\vec{x}|\eta),&~p(\vec{x}|\eta+\delta\eta)) = \int p(\vec{x}|\eta) \log p(\vec{x}|\eta) \dif\vec{x} - \int p(\vec{x}|\eta) \log p(\vec{x}|\eta+\delta\eta) \dif\vec{x} \\
    {} & = - \int p(\vec{x}|\eta) \left( \nabla_\eta \log p(\vec{x}|\eta)\T \delta\eta + \frac{1}{2} \delta\eta\T \text{Hess}_\eta [\log p(\vec{x}|\eta)] \delta\eta + o(\delta\eta^2) \right) \dif\vec{x} \\
    {}& = \frac{1}{2} \delta\eta\T \left[ \int p(\vec{x}|\eta) \left[\nabla_\eta\log p(\vec{x}|\eta) \nabla_\eta\log p(\vec{x}|\eta)\T\right] \dif\vec{x} \right] \delta\eta + o(\delta\eta^2).
\end{align*}

\end{proof}

\begin{definition}
We consider a curve $\gamma(t)$ and its derivative $\dot{\gamma}(t)$ on the statistical manifold such that, $\forall t\in[0,1], \gamma(t)=\eta_t\in\mathcal{H}$. The manifold is equipped with the Fisher-Rao metric. The length and the energy functionals are defined with respect to the metric $\mathbf{I}_\HH(\eta)$:
\[ \mathrm{Length}(\gamma) = \int_{0}^{1} \sqrt{\dot{\gamma}(t)\T \mathbf{I}_\HH(\eta) \dot{\gamma}(t)} \dif t \quad\text{and}\quad \mathrm{Energy}(\gamma) = \int_{0}^{1} \dot{\gamma}(t)\T \mathbf{I}_\HH(\eta) \dot{\gamma}(t) \dif t.\]

Locally length-minimising curves between two connecting points are called geodesics. These can be found by minimizing the energy using the Euler-Lagrange equations which gives the following system of 2\textsuperscript{nd} order nonlinear ordinary differential equations (ODEs) \citep{arvanitidis:iclr:2018}
\begin{equation}
\label{eq:ode_system}
    \ddot{\gamma}(t)
    =-\frac{1}{2}\mathbf{I}_{\HH}^{-1}(\gamma(t))\Big[2 (\dot{\gamma}(t)\T \otimes \mathbb{I}_d) \parder{\vectorize{\mathbf{I}_{\HH}(\gamma(t))}}{\gamma(t)}\dot{\gamma}(t)
    - \parder{\vectorize{\mathbf{I}_{\HH}(\gamma(t))}}{\gamma(t)}\T (\dot{\gamma}(t) \otimes \dot{\gamma}(t))\Big].
\end{equation}
\end{definition}

\begin{proposition} \label{prop:FRenergydistance}
The KL-divergence between two close elements of the curve $\gamma$ is defined as: $\mathrm{KL}(p_{t},~p_{t+\delta t})=\mathrm{KL}(p(\vec{x}|\gamma(t)),~p(\vec{x}|\gamma(t+\delta t)))$.
The length and the energy functionals can be approximated with respect to this KL-divergence:
\[\mathrm{Length}(\gamma) \approx\sqrt{2\textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})}  \quad\text{and}\quad \mathrm{Energy}(\gamma) \approx \frac{\delta t}{2} \textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})\]

\end{proposition}

\begin{proof}
On the statistical manifold, we have $\gamma(t+\delta t)=\gamma(t)+\delta t\dot{\gamma}(t)$. The KL-divergence between perturbed distributions can be defined as: $\mathrm{KL}(p_{t},p_{t+\delta t}) = \mathrm{KL}(p(\vec{x}|\gamma(t)),p(\vec{x}|\gamma(t+\delta t)))= \mathrm{KL}(p(\vec{x}|\eta_t),p(\vec{x}|\eta_t+\delta\eta_t)),$ with $\eta_t=\gamma(t)$ and $\delta\eta_t=\delta t\ \dot{\gamma}(t)$.
Then, we obtain:
\[\mathrm{KL}(p_{t},p_{t+\delta t}) = \frac{1}{2} \delta t^2\ \dot{\gamma}(t)\T\mathbf{I}_\HH(\eta_t)\dot{\gamma}(t) + o(\delta t^2).\]

The length and energy terms appear in the following equations:
\begin{align*}
    \int_0^{1}\mathrm{KL}(p_{t},p_{t+\delta t}) \dif t &= \frac{\delta t^2}{2} \int_0^{1} \dot{\gamma}(t)\T\mathbf{I}_\HH(\eta_t)\dot{\gamma}(t) \dif t + o(\delta t^2) = \frac{\delta t^2}{2}\ \mathrm{Energy}(\gamma)+ o(\delta t^2),\\
    \int_0^{1}\sqrt{\mathrm{KL}(p_{t},p_{t+\delta t})} \dif t &= \frac{\delta t}{\sqrt{2}} \int_0^{1} \sqrt{\dot{\gamma}(t)\T \mathbf{I}_\HH(\eta_t)\dot{\gamma}(t)} \dif t + o(\delta t^2) = \frac{\delta t}{\sqrt{2}}\ \mathrm{Length}(\gamma)+ o(\delta t^2).\\
\end{align*}

If we want approximate any continuous function $f$ with a discrete sequence, by partitioning it in $T$ small segments, such that: $\delta t \approx \frac{1}{T}$, we have: $\int_0^{1} f(t) dt \approx \sum_{t=1}^{T} f(t) \delta t$, which in our case gives:

\[\mathrm{Length}(\gamma) \approx \sqrt{2\textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})} \quad \text{and} \quad \mathrm{Energy}(\gamma) \approx \frac{2}{\delta t} \textstyle \sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t}). \]

\end{proof}

\subsection{The Fisher-Rao metric for several distributions} \label{appendix:FR:distributions}

\begin{table}[h]
\begin{center}
\begin{tabular}{llll}
\hline
Distributions&Probability density functions&Parameters& Fisher-Rao matrix\\[0.2cm] \hline
Normal& $\frac{1}{\sqrt{2\pi\sigma^2}} \exp{-\frac{(\vec{x}-\mu)^2}{2\sigma^2}}$ & $\mu, \sigma^2$& $\mathbf{I}_{\mathcal{N}}(\mu,\sigma^2)$\\[0.3cm]
Bernoulli & $\theta^{\vec{x}}(1-\theta)^{1-\vec{x}}$& $\theta$& $\mathbf{I}_{\mathcal{B}}(\theta)$\\[0.3cm]
Categorical& $\prod_{k=1}^{K}\theta_{k}^{\vec{x}_k}$& $\theta_1, \dots, \theta_K$ & $\mathbf{I}_{\mathcal{C}}(\theta_1, \dots, \theta_K)$ \\[0.3cm]
Gamma& $\frac{\beta^{\alpha} \vec{x}^{\alpha-1}e^{-\beta \vec{x}}}{\Gamma(\alpha)}$    & $\alpha, \beta$& $\mathbf{I}_{\mathcal{G}}(\alpha,\beta)$\\[0.3cm]
Von Mises-Fisher, for $\mathbb{S}^2$ & $\frac{\kappa}{4\pi\sinh{\kappa}}\exp{(\kappa \mu\T \vec{x})}$& $\kappa, \mu$& $\mathbf{I}_{\mathcal{S}}(\kappa,\mu)$\\[0.3cm]
Beta & $\frac{\Gamma(\alpha)\Gamma(\beta)}{\Gamma(\alpha + \beta)} \vec{x}^{\alpha - 1}(1-\vec{x})^{\beta -1}$& $\alpha, \beta$ & $\mathbf{I}_{\mathcal{B}}(\alpha,\beta)$ \\[0.3cm] \hline
\end{tabular}
\end{center}
\caption{List of distributions}
\label{tab:distributions}
\end{table}

With the notations of Table \ref{tab:distributions}, the Fisher-Rao matrices of the the univariate Normal, Bernoulli and Categorical are:
\[
\mathbf{I}_{\mathcal{N}}(\mu,\sigma^2) = \begin{pmatrix} \frac{1}{\sigma^2} & 0 \\ 0 & \frac{1}{2\sigma^4}\\\end{pmatrix}, \quad
\mathbf{I}_{\mathcal{B}}(\theta) = \frac{1}{\theta(1-\theta)}, \quad
\mathbf{I}_{\mathcal{C}}(\theta_1, \dots, \theta_K) = \begin{pmatrix}
  \sfrac{1}{\theta_1} & 0 &\dots & 0 \\
  0 & \sfrac{1}{\theta_2} & \dots & 0 \\
  \vdots & \vdots & \ddots & \vdots \\
  0 & 0 & \dots & \sfrac{1}{\theta_K} \\
\end{pmatrix}
\]
In addition, the Fisher-Rao matrices of the Gamma, Von Mises-Fisher and the Beta distributions are:
\begin{align*}
    \mathbf{I}_{\mathcal{G}}(\alpha,\beta) &= \begin{pmatrix} \frac{\alpha}{\beta^2} & -\frac{1}{\beta} \\ -\frac{1}{\beta} & \Psi_1(\alpha) \\\end{pmatrix}, \\
    \mathbf{I}_{\mathcal{S}}(\kappa,\mu) &= \begin{pmatrix} \kappa K(\kappa)(\mathds{1}-3\mu\mu^{\top})+\kappa^2\mu\mu^{\top} & (\kappa K(\kappa)^2-\frac{2}{k}K(\kappa)+1)\mu \\ (\kappa K(\kappa)^2-\frac{2}{k}K(\kappa)+1)\mu^{\top} & 3K(\kappa)^2-\frac{2}{\kappa}K(\kappa)+1 \\\end{pmatrix}, \\
    \mathbf{I}_{\mathcal{B}}(\alpha,\beta) &= \begin{pmatrix} \Psi_1(\alpha)-\Psi_1(\alpha+\beta) & -\Psi_1(\alpha+\beta) \\ -\Psi_1(\alpha+\beta) & \Psi_1(\beta) - \Psi_1(\alpha+\beta) \\\end{pmatrix},
\end{align*}
with $\Psi_1(\alpha) = \frac{\partial^2 \ln\Gamma(\alpha)}{\partial \alpha}$ the trigamma function, and $K(\kappa) = \coth{\kappa}-\frac{1}{\kappa}$.

\begin{proof}
The univariate Normal, Bernoulli and Categorical have already been studied by \cite{tomczak:2012}, and the Beta distribution by \cite{brigant:2019}. We will then focus our proof on the Gamma and the Von-Mises Fisher distributions.

In order to bypass unnecessary details, we will use the following notations, we redefine the Fisher-Rao as: $\mathbf{I}(\eta)=\mathbb{E}_{x}[g(\eta, x)g(\eta, x)\T]$, with $g(\eta,x) = \nabla_{\eta} \ln p(x|\eta)$ the Fisher score. We call $G = g(\eta, x)g(\eta, x)^{\top}$, and $G_{ij}$ the matrix elements.

\textbf{Gamma distribution:}

We have $p(x|\alpha, \beta) = \Gamma(\alpha)^{-1} \beta^{\alpha} x^{\alpha-1}e^{-\beta x}$, which leads to:
\begin{align*}
    \ln p(x|\alpha, \beta) &= - \ln \Gamma(\alpha) + \alpha \ln \beta + (\alpha - 1) \ln x - \beta x, \\
    \frac{\partial \ln p}{\partial \alpha} & = - \Psi(\alpha) + \ln \beta + \ln x, \\
    \frac{\partial \ln p}{\partial \beta} & = \frac{\alpha}{\beta} - x.
\end{align*}
Then:
\begin{align*}
    G_{11} &= \left(\frac{\partial \ln p}{\partial \alpha}\right)^2 = (\Psi_{0}(\alpha) + \ln \beta)^2 + 2 (\Psi(\alpha) + \ln \beta ) ln x + \ln^2 x, \\
    G_{22} &= \left(\frac{\partial \ln p}{\partial \beta}\right)^2  = \left(\frac{\alpha}{\beta}\right)^2 - 2 \frac{\alpha}{\beta} x + x^2,  \\
    G_{12} = G_{21} &= \frac{\partial \ln p}{\partial \alpha} \cdot \frac{\partial \ln p}{\partial \beta}  = \left(\Psi(\alpha) + \ln \beta\right) \left(\frac{\alpha}{\beta} - x\right) + \frac{\alpha}{\beta} \ln x - x \ln x.
\end{align*}

We know that $\mathbb{E}[x] = \frac{\alpha}{\beta}$. We can compute, using your favorite symbolic computation software, the following moments:
\begin{align*}
    \mathbb{E}[\ln x] &= -\ln \beta + \Psi(\alpha) \\
    \mathbb{E}[x \ln x] &= \frac{\alpha}{\beta} \left(\Psi(\alpha + 1) - \ln \beta\right) \\
    \mathbb{E}[\ln^2 x] &= \left(\ln \beta - \Psi(\alpha)\right)^2 + \Psi_1(\alpha)
\end{align*}
Replacing the moments for the following equations: $\mathbb{E}[G_{11}]$, $\mathbb{E}[G_{22}]$ and $\mathbb{E}[G_{12}]$ will finally give the Fisher-Rao matrix.\\

\textbf{Von Mises Fisher distribution, for $\mathbb{S}^2$:}

We have $p(\vec{x}|\mu, \kappa) = C_{3}(\kappa) \exp (\kappa \mu\T \vec{x})$, with $C_{3}(\kappa) = \kappa (4 \pi \sinh \kappa)^{-1}$. Here, $\mu$ is a 3-dimensional vector with $\lVert \mu \rVert = 1$.
\begin{align*}
    \ln p(\vec{x}|\mu, \kappa) &= \ln \kappa - \ln 4\pi - \ln \sinh(\kappa) + \kappa \mu\T \vec{x} \\
    \nabla_{\mu} \ln p & = \kappa \vec{x}  \\
    \frac{\partial \ln p}{\partial \kappa} & =  \kappa^{-1} - \text{coth}(\kappa) + \mu\T \vec{x}.
\end{align*}
Here, the Fisher-Rao matrix $\mathbf{I}_{\mathcal{S}}$ will be composed of block matrices, such that: $\mathbf{I}_{\mathcal{S}} = \mathbb{E}[G]$, with $G_{11}$ a $3\times 3$-matrix, $G_{22}$ a scalar, and $G_{12}=G_{21}^{\top}$ a 3-dimensional vector.
\begin{align*}
    G_{11} &= \nabla_{\mu}\ln p \nabla_{\mu}\ln p^{\top} = \kappa^2 \vec{x} \vec{x}\T \\
    G_{22} &= \left(\frac{\partial \ln p}{\partial \kappa}\right)^2 = K(\kappa)^2 + 2 K(\kappa) \mu^{\top} \vec{x} + (\mu\T \vec{x})^2  \\
    G_{12} = G_{21}^{\top} &= \frac{\partial \ln p}{\partial \kappa} \cdot \nabla_{\mu}\ln p = \left(K(\kappa) + \mu\T \vec{x} \right) \kappa \vec{x},
\end{align*}
with $K(\kappa) = \text{coth}(\kappa) - \frac{1}{\kappa}$.

We know from \cite{hillen:2016} that the mean and variance of the Von Mises Fisher distribution in the 3-dimensional case is: $\mathbb{E}[\vec{x}] = K(\kappa) \mu$ and $\text{Var}[\vec{x}] = \frac{1}{\kappa} K(\kappa) \mathds{1} + (1- \frac{\text{coth}(\kappa)}{\kappa} + \frac{2}{\kappa^2} - \text{coth}^2(\kappa)) \mu \mu^{\top}$. We can then deduce the following meaningful moments:
\begin{align*}
    \mathbb{E}[\vec{x}\x^{\top}] & = \text{Var}[\vec{x}] + \mathbb{E}[\vec{x}]\mathbb{E}[\vec{x}]^{\top} = \left(1-\frac{3}{\kappa} K(\kappa)\right) \mu\mu^{\top} + \frac{1}{\kappa} K(\kappa) \mathds{1}, \\
    \mathbb{E}[\mu\T \vec{x}] & = \mu^{\top} \mathbb{E}[\vec{x}] = K(\kappa) \mu^{\top} \mu = K(\kappa) \\
    \mathbb{E}[(\mu\T \vec{x})^2] & = \mu^{\top} \text{Var}[\vec{x}] \mu + \mathbb{E}[\mu\T \vec{x}]^2 = 1-\frac{2}{\kappa}K(\kappa), \\
    \mathbb{E}[\mu\T \vec{x} \vec{x}] &= \mathbb{E}[\mu \vec{x} \vec{x}\T] = \mathbb{E}[\vec{x} \vec{x}\T] \mu = \left(1 - \frac{3}{\kappa} K(\kappa)\right) \mu + \frac{1}{\kappa} K(\kappa) \mu.
\end{align*}
Replacing those moments in the following expressions: $\mathbb{E}[G_{11}]$, $\mathbb{E}[G_{22}]$, $\mathbb{E}[G_{12}]$ directly gives the Fisher-Rao metric.
\end{proof}

\section{Curve energy approximation for categorical data}
\label{appendix:categorical_details}

In this section we present the details of the example in Section~\ref{sec:example_categorical}. In particular, we the steps to derive an approximation to the energy of a latent curve in closed form, which is suitable for applying automatic differentiation. This is particularly useful for our setting, since it allows us to consider our framework as a Black Box Random Geometry processing toolbox.

Let a random variable $\vec{x}\in\R^D$ that follows a generalized Bernoulli likelihood $p(\vec{x}|\eta)$, so the vector $\vec{x}\in\R^D$ is of the form $\b{x} = (0, \cdots, 1, \cdots, 0)$ with $\sum_i x_i = 1$. The parameters $\eta\in\R^D$ are given as $\eta = h(\vec{z})$, with $\eta_i \geq 0~\forall i$ and $\sum_i \eta_i = 1$ so we know that the parameters lie on the unit simplex. Actually, they represent the probability the corresponding dimension to be 1 on a random draw. Also, the $p(\vec{x}|\vec{z}) = \eta_1^{[x_1]}\cdots \eta_D^{[x_D]}$, where $[x_i] = 1$ if $x_i=1$ else $[x_i] = 0$ which can be seen as an indicator function. The $\log p(\vec{x}|\eta) = \sum_i [x_i] \log(\eta_i)$ and $\nabla_\eta \log p(\vec{x}|\eta) = \left(\frac{[x_1]}{\eta_1}, \dots, \frac{[x_D]}{\eta_D}\right)$. Due to the outer product we have to compute the following expectations
\begin{align}
    &\mathbb{E}_\b{x}\left[ \frac{[x_i]}{\eta_i} \frac{[x_j]}{\eta_j}\right] = 0, ~\quad \text{if}\quad i\neq j,\\
    &\mathbb{E}_\b{x}\left[ \left(\frac{[x_i]}{\eta_i}\right)^2\right] = \frac{1}{\eta_i}, ~\quad \text{if}\quad i=j,
\end{align}
because the $[x_i]$ and $[x_j]$ cannot be 1 on the same time, while the $\mathbb{E}_\b{x}[[x_i]^2] = \eta_i$ as it shows the number of times $x_i=1$. So the Fisher-Rao metric of $\HH$ is equal to $\b{I}_\HH(\eta) = \mathrm{diag}\left(\sfrac{1}{\eta_1},\dots,\sfrac{1}{\eta_D}\right)$. Note that the shortest paths between two distributions must be on the unit simplex in $\HH$, while on the same time respecting the geometry of the Fisher-Rao metric.

We can easily parametrize the unit simplex by $[\eta_1,\dots,\eta_{D-1}, \widetilde{\eta}_D]$ with
\begin{equation}
    \widetilde{\eta}_D(\eta_1,\dots, \eta_{D-1}) = 1 - \sum_{i=1}^{D-1} \eta_i.
\end{equation}
This allows to pullback the Fisher-Rao metric in the latent space $[\eta_1,\dots,\eta_{D-1}]$ as we have described in this paper. Intuitively, the $\vec{z} = [\eta_1,\dots,\eta_{D-1}]$ and the function $h$ is the parametrization of the simplex. Hence, we are able to compute the shortest path using the induced metric.

However, there is a simpler way to compute this path. We know that the element-wise square root of the parameters $\eta$ gives a point on the positive orthonant of the unit sphere as $y_i = \sqrt{\eta_i}~\Rightarrow~\sum_i y_i^2 = \sum_i \sqrt{\eta_i}^2 = 1$. We also know that the shortest path on a sphere is the great-circle. Therefore, the distance between two distributions parametrized by $\eta$ and $\eta'$ on the unit simplex in $\HH$, can be equivalently measured using the great-circle distance between their square roots as
\begin{equation}
    \mathrm{dist}(\eta, \eta') = \arccos\sqrt{\eta}\T\sqrt{\eta'}.
\end{equation}

In this way, we can approximate the energy of a curve $c(t)$ in the latent space as follows
\begin{align}
\label{eq:approx_energy_categorical}
    \mathrm{Energy}[c]&\approx \sum_{n=1}^{N-1}\mathrm{dist}^2(h(c(\sfrac{n}{N})),h(c(\sfrac{n+1}{N}))) = \sum_{n=1}^{N-1}\arccos^2 \sqrt{h(c(\sfrac{n}{N}))}\T\sqrt{h(c(\sfrac{n+1}{N}))}\nonumber\\
    &=\sum_{n=1}^{N-1}\left(2 - 2\sqrt{h(c(\sfrac{n}{N}))}\T\sqrt{h(c(\sfrac{n+1}{N}))}\right),
\end{align}
where we used at the last step the small angle approximation $\cos\theta \approx 1 - \frac{\theta^2}{2} \Leftrightarrow \theta^2 \approx 2 - 2\cos\theta$. Note that this formulation is suitable for our proposed method to compute shortest paths (see Section~\ref{sec:geodesic_computation}).

The derivation above represents the conceptual strategy, while in general we proposed to use the KL divergence approximation result \eqref{eq:KL_approx_inner_product} in place of the great-circle distance. Intuitively, when the KL divergence has an analytic solution, we can derive an analogous energy approximation. Even if the solution of the KL is intractable, we can still use our approach as long as we can estimate the KL using Monte Carlo and propagate the gradient through the samples using a re-parametrization scheme or a score function estimator.

\section{Information geometry in generative modeling}

In this section we present the additional technical information related to the pullback Fisher-Rao metric in the latent space of a VAE.

\subsection{Details for the pullback metric in the latent space}

We call $f$ the non linear function, typically parametrized as deep neural networks, that maps the variables from the latent space $\mathcal{Z}$ to the parameter space $\mathcal{H}$, such that: $f(\vec{z})=\eta$, with $\vec{z}\in\mathcal{Z}$ and $\eta\in\mathcal{H}$. Furthermore, the data $\vec{x} \in \mathcal{X}$ is reconstructed such that it follows a specific distribution: $\vec{x} \sim p(\vec{x}|\eta)$, with $p(\vec{x}|\eta)$ being for instance a Bernoulli or Gaussian distribution. The parameter space $\mathcal{H}$ is a statistical manifold equipped with Fisher-Rao metric: $\b{I}_\HH(\eta) \overset{\Delta}{=} \int p(\vec{x}|\eta) \left[\nabla_\eta\log p(\vec{x}|\eta) \nabla_\eta\log p(\vec{x}|\eta)\T\right] \dif \vec{x}$. We denote by $\b{J}_f$ the Jacobian of $f$.

\begin{proposition}
The latent space $\mathcal{Z}$ is equipped with the pullback metric:
\[ \b{M}(\vec{z}) \overset{\Delta}{=} \b{J}_{f}(\vec{z})\T \b{I}_\HH(f(\vec{z})) \b{J}_{f}(\vec{z}). \]
\end{proposition}

\begin{proof}
The parameter space is a statistical manifold equipped with the Fisher-Rao metric $\b{I}_\HH(\eta)$, thus the scalar product at $\eta$ between two vectors $d\eta_1, d\eta_2 \in \mathcal{H}$ is: $\langle d\eta_1, d\eta_2 \rangle_{\b{I}_\HH(\eta)} =  d\eta_1\T \b{I}_\HH (\eta) d\eta_2$. Furthermore, we have $f:\mathcal{Z}\to\mathcal{H}$ a smooth immersion. For two vectors $d\vec{z}_1, d\vec{z}_2 \in\mathcal{Z}$, we have at $\eta = f(\vec{z})$ that: $\langle d\eta_1, d\eta_2 \rangle_{\b{I}_\HH(\eta)} =  \langle \b{J}_f(\vec{z}) d\vec{z}_1, \b{J}_f(\vec{z}) d\vec{z}_2 \rangle_{\b{I}_\HH(\eta)} = d\vec{z}_1\T (\b{J}_{f}(\vec{z})\T \b{I}_\HH(f(\vec{z})) \b{J}_{f}(\vec{z})) d\vec{z}_2$.
\end{proof}

\begin{proposition}
Our pullback metric $\b{M}(\vec{z})$ is actually equal to the Fisher-Rao metric obtained over the parameter space $\Z$:
\[\b{M}(\vec{z}) = \b{I}_\Z(\vec{z}) \overset{\Delta}{=} \int p(\vec{x}|\vec{z}) \left[\nabla_\vec{z} \log p(\vec{x}|\vec{z}) \nabla_\vec{z} \log p(\vec{x}|\vec{z})\T\right] \dif \vec{x} \]
\end{proposition}

\begin{proof}
We will show that $\b{I}_\Z(\vec{z}) = \b{J}_{f}(\vec{z})\T \b{I}_\HH(\eta) \b{J}_{f}(\vec{z})$.
Let's consider the definition of the Fisher-Rao metric in $\Z$ :
\begin{align}
    \b{I}_\Z(\b{z})
    &= \int \nabla_{\b{z}} \log p(\b{x}~|~\b{z}) \cdot \nabla_{\b{z}} \log p(\b{x}~|~\b{z})\T p(\b{x}~|~\b{z})\dif \b{x}\\
    &= \int \b{J}_f(\b{z})\T \nabla_{\eta} \log p(\b{x}|\eta)\nabla_{\eta} \log p(\b{x}|\eta)\T \b{J}_f(\b{z}) p(\b{x}|\eta)\dif\b{x} \\
    &= \b{J}_f(\b{z})\T \left[ \int_\X \nabla_{\eta} \log p(\b{x}|\eta)\nabla_{\eta} \log p(\b{x}|\eta)\T  p(\b{x}|\eta)\dif\b{x} \right]\b{J}_f(\b{z})\\
    &= \b{J}_f(\b{z})\T \b{I}_\HH(f(\vec{z})) \b{J}_f(\b{z})=\b{M}(\vec{z}) \nonumber
\end{align}
where we use the fact that $\eta=f(\b{z})$ so the $\nabla_{\b{z}}\log p(\b{x}|f(\b{z})) = \b{J}_f(\b{z})\T \cdot \nabla_{\eta} \log p(\b{x}|\eta)$

The same argument can be proved as follows:
\begin{align}
    &\inner{d\eta}{\b{I}_\HH(\eta) d\eta} = \inner{\b{J}_f(\b{z})d\b{z}}{\b{I}_\HH(f(\b{z}))\b{J}_f(\b{z})d\b{z}} \\
    &= \inner{\b{J}_f(\b{z})d\b{z}}{\int \nabla_\eta \log p(\b{x}|\eta) \nabla_\eta \log p(\b{x}|\eta) \T p(\b{x}|\eta)\dif\b{x} ~\b{J}_f(\b{z})d\b{z}}\\
    &=\inner{d\b{z}}{\int \b{J}_f(\b{z})\T \cdot \nabla_\eta \log p(\b{x}|\eta) ~\nabla_\eta \log p(\b{x}|\eta)\T \cdot \b{J}_f(\b{z}) p(\b{x}|\eta)\dif\b{x} ~d\b{z}}\\
    &=\inner{d\b{z}}{\int \nabla_{\b{z}}\log p(\b{x}|\b{z}) \nabla_{\b{z}}\log p(\b{x}|\b{z})\T p(\b{x}|\b{z}) \dif\b{x}~d\b{z}} = \inner{d\b{z}}{\b{I}_\Z(\b{z}) d\b{z}}
\end{align}
\end{proof}

In section \ref{appendix:FR:distributions}, we have seen how to derive a close-form expression of the Fisher-Rao metric for a one-dimensional observation $x$ that follows a specific distribution. In practice, $\vec{x}\in\mathcal{X}\subset\mathbb{R}^D$ is a multi-dimensional variable where each dimension represents, for instance, a pixel when working with images or a feature when working with tabular data. Each feature, $x_i$ with $i=1\cdots D$, is obtained for a specific set of parameters $\{\eta_i\}$. We assume that the features follow the same distribution $\mathcal{D}$, such that: $x_i \sim p(x_i|\eta_i)$, and $p(\vec{x}|\eta) = \prod_{i=1}^{D}p(x_i|\eta_i)$.

\begin{proposition} \label{prop:blockmatrix}
The Fisher-Rao metric $\mathbf{I}_{\HH}(\eta)$ is a block matrix where the diagonal terms are the Fisher-Rao matrices $\mathbf{I}_{\HH,i}$ obtained for each data feature $x_i$:
\[
\mathbf{I}_{\HH}(\eta) = \begin{pmatrix}
  \b{I}_{\HH,1} & 0 &\dots & 0 \\
  0 & \b{I}_{\HH,2} & \dots & 0 \\
  \vdots & \vdots & \ddots & \vdots \\
  0 & 0 & \dots & \b{I}_{\HH, D} \\
\end{pmatrix}
\]
\end{proposition}
\begin{proof}
We have $x_i\sim p(x_i|\eta_i)$ and  $\mathbf{I}_{\HH,i}=\int p(x_i|\eta_i) \left[\nabla_{\eta_i}\log p(x_i|\eta_i) \nabla_{\eta_i}\log p(x_i|\eta_i)\T\right] dx_i$.
Also, we assumed: $p(\vec{x}|\eta)=\prod_{i=1}^{D}p(x_i|\eta_i)$.
We then have: $\log{p(\vec{x}|\eta)} = \sum_{i=1}^{D} \log{p(x_i|\eta_i)}$, and the Fisher score: $\nabla_{\eta}\log{p(\vec{x}|\eta)} = \nabla_{\eta} \sum_{i=1}^{D} \log{p(x_i|\eta_i)} = \left[\nabla\eta_1 \ln{p(x_1|\eta_1)}, \dots, \nabla\eta_D \ln{p(x_1|\eta_D)}\right]\T$. \\

\noindent
The matrix $\mathbf{I}_{\HH}(\eta)$ is thus a $D \times D$ block matrix, where the $(i,j)$-block element is:
\[I_{ij}=\int p(x_i|\eta_i) \left[\nabla_{\eta_i} \log p(x_i|\eta_i) \nabla_{\eta_i}\log p(x_j|\eta_j)\T\right] dx_i.\]
Let's note that:
\[\int p(x_i|\eta_i) \nabla_{\eta_i}\log p(x_i|\eta_i) dx_i = \int p(x_i|\eta_i) \frac{\nabla_{\eta_i} p(x_i|\eta_i)}{p(x_i|\eta_i)}dx_i = \nabla_{\eta_i} \int p(x_i|\eta_i) dx_i = 0.\]
When $i=j$, we have $\mathbf{I}_{ii}=\mathbf{I}_{\HH,i}$, with $\mathbf{I}_{\HH,i}$ being the Fisher-Rao metric obtained for: $x_i\sim p(x_i|\eta_i)$. \\
\noindent
When $i\neq j$, we have: $\mathbf{I}_{ij}=\nabla\log p(x_j|\eta_j)\T \int p(x_i|\eta_i) \nabla_{\eta_i}\log p(x_i|\eta_i) dx_i = 0$.
\end{proof}

Then, for example, if we are dealing with binary images, and make the assumption that each pixel $x_i$ follows a Bernoulli distribution: $p(x_i|\eta_i) = \eta^{x_i}(1-\eta_i)^{1-x_i}$, then according to Section \ref{appendix:FR:distributions} and Proposition \ref{prop:blockmatrix}, the Fisher-Rao matrix that endows the parameter space $\mathcal{H}$ is:
\[
\b{I}_\HH(\eta) = \begin{pmatrix}
              \frac{1}{\eta_1(1-\eta_1)} & 0 &\dots & 0 \\
              0 & \frac{1}{\eta_2(1-\eta_2)} & \dots & 0 \\
              \vdots & \vdots & \ddots & \vdots \\
              0 & 0 & \dots & \frac{1}{\eta_D(1-\eta_D)} \\
            \end{pmatrix}.
\]

We have seen that in theory, we can obtain a close form expression for the pullback metric, if the probability distribution is known. In practice, we can directly infer the metric using the approximation of the KL-divergence.

\begin{proposition}
We define perturbations vectors as: $\delta e_{i} = \varepsilon \cdot \b{e_i}$, with $\varepsilon\in\mathbb{R}_{+}$ a small infinitesimal quantity, and ($\b{e_i}$) a canonical basis vector in $\mathbb{R}^d$. For better clarity, we rename $\mathrm{KL}(p(\vec{x}|\vec{z}), p(\vec{x}|\vec{z}+\delta \vec{z})) = \mathrm{KL}_{\vec{z}}(\delta \vec{z})$ and we note $\b{M}_{ij}$ the components of $\b{M}(\vec{z})$.
We can then approximate by a system of equations the diagonal and non-diagonal elements of the metric:
    \begin{align*}
        \b{M}_{ii} &\approx 2 \ \mathrm{KL}_{\vec{z}}(\delta \b{e_i})/ \varepsilon^2 \ \\
        \b{M}_{ij} = \b{M}_{ji} &\approx  \left(\mathrm{KL}_{\vec{z}}(\delta \b{e_i} + \delta \b{e_j})- \mathrm{KL}_{\vec{z}}(\delta \b{e_i}) -\mathrm{KL}_{\vec{z}}(\delta \b{e_j}) \right)/ \varepsilon^2.
    \end{align*}
\end{proposition}

\begin{proof}
From Proposition \ref{prop:FRapproxKL}, we know that:
\[\mathrm{KL}_{\vec{z}}(\delta \vec{z}) = \frac{1}{2} \delta\vec{z}^{\top} \b{M}(\vec{z}) \delta\vec{z} + o(\delta\vec{z}^2).\]

Let's take $\delta e_{i} = \varepsilon \cdot \b{e_i}$. On one hand, we have: $\delta e_{i}^{\top} \b{M}(\vec{z}) \delta e_{i} = \varepsilon^2 \b{M}_{ii}$. On the second hand, we also have: $ \delta e_{i}^{\top} \b{M}(\vec{z}) \delta e_{i} \approx 2 \mathrm{KL}_{\vec{z}}(\delta \b{e_i})$, which gives us the equation to infer the diagonal elements of the metric.

Now, let's take $\delta e_{i} + \delta e_j = \varepsilon \cdot (\b{e_i}+\b{e_j})$. Then, we have: $(\delta e_{i} + \delta e_j)^{\top} \b{M}(\vec{z}) (\delta e_{i} + \delta e_j) = \varepsilon^2 (\b{M}_{ii} + \b{M}_{jj} + \b{M}_{ij} + \b{M}_{ji})$. We also know that $\b{M}_{ji} =\b{M}_{ij}$. Again, we also have: $(\delta e_{i} + \delta e_j)^{\top} \b{M}(\vec{z}) (\delta e_{i} + \delta e_j) \approx 2 \mathrm{KL}_{\vec{z}}(\delta \b{e_i} + \delta \b{e_j})$.

We can replace the terms $\b{M}_{ii}$ and $\b{M}_{jj}$ in the equation obtained above with the KL-divergence for the diagonal terms. Which finally gives us: $\b{M}_{ij} = \b{M}_{ji} \approx  \left(\mathrm{KL}_{\vec{z}}(\delta \b{e_i} + \delta \b{e_j})- \mathrm{KL}_{\vec{z}}(\delta \b{e_i}) -\mathrm{KL}_{\vec{z}}(\delta \b{e_j}) \right)/ \varepsilon^2$.

\end{proof}

\subsection{Uncertainty quantification and regularization}
\label{appendix:sec:UQ}

As discussed in the main text, we carefully design our mappings from latent space to parameter space such that they model the training codes according to the learned decoders, and extrapolate to uncertainty outside the support of the data. This, we refer to as \textbf{uncertainty regularization}. In this section we explain it in detail. The core idea of this uncertainty regularization is imposing a ``slider'' that forces the distribution $p(\vec{x}|\vec{z})$ to change when $\vec{z}$ is far from the training latent codes. For this, we use a combination of KMeans and the sigmoid activation function.

We start by encoding our training data, arriving at a set of latent codes $\{\vec{z}_n\}_{n=1}^{N}\subseteq\mathcal{Z}$. We then train KMeans($k$) on these latent codes (where $k$ is a hyperparameter that we tweak manually), arriving at $k$ cluster centers $\{\b{c}_j\}_{j=1}^k$. These cluster centers serve as a proxy for "closeness" to the data: we know that a latent code $\vec{z}\in\mathcal{Z}$ is near the support if $D(\vec{z}) := \min_j\left\{\|\vec{z} - \b{c}_j\|^2\right\}$ is close to 0.

The next step in our regularization process is to reweight our decoded distributions such that we decode to high uncertainty when $D(\vec{z})$ is large, and we decode to our learned distributions when $D(\vec{z}) \approx 0$. This mapping from $[0, \infty)\to (0, 1)$ can be constructed using a modified sigmoid function \cite{detlefsens:proteins:2020,detlefsen:2019:reliable}, consider indeed
\begin{equation}
\label{appendix:eq:translated_sigmoid}
\tilde{\sigma}_{\beta}(d) = \text{Sigmoid}\left(\frac{d - c\cdot\text{Softplus}(\beta)}{\text{Softplus}(\beta)}\right),
\end{equation}
where $\beta\in\mathbb{R}$ is another hyperparameter that we manually tweak, and $c\approx 7$.

With this translated sigmoid, we have that $\tilde{\sigma}_\beta(D(\vec{z}))$ is close to 0 when $\vec{z}$ is close to the support of the data (i.e. close to the cluster centers), and it converges to 1 when $D(\vec{z})\to\infty$.  $\tilde{\sigma}_\beta(D(\vec{z}))$ serves, then, as a slider that indicates closeness to the training codes. This reweighting takes the following form:
\begin{equation}
\label{appendix:eq:reweight}
\text{reweight}(\vec{z}) = (1 - \tilde{\sigma}_\beta(D(\vec{z})))h(\vec{z}) + \tilde{\sigma}_\beta(D(\vec{z}))\,\text{extrapolate}(\vec{z}),
\end{equation}
where  $h(\vec{z}) = \eta\in\mathcal{H}$ represents our learned networks in parameter space,  and $\text{extrapolate}(\vec{z})$ returns the parameters of the distribution that maximize uncertainty (e.g. $\sigma\to\infty$ in the case of an isotropic Gaussian, $p\to 1/2$ in the case of a Bernoulli, and $\kappa\to 0$ in the case of the von Mises-Fisher).

For the particular case of the experiment in which we pull back the Fisher-Rao metric from the parameter space of several distributions (see \ref{sec:toy_latent_space_experiment}), Table \ref{appendix:tab:extrapolation_mechanisms} provides the exact extrapolation mechanisms and implementations of $h(\vec{z})$.

\begin{table}[]
    \centering
    \resizebox{\textwidth}{!}{
    \begin{tabular}{lll}
        Distribution & {$h\colon\mathcal{Z}_{\text{toy}}\to\mathcal{H}$} & Extrapolation mechanism \\
        \toprule
        Normal & $\mu(\vec{z})=10\cdot f_3(\vec{z})$, ~ $\sigma(\vec{z}) = 10\cdot \text{Softplus}(f_3(\vec{z}))$ & $\sigma(\vec{z})\to\infty$ \\[0.1cm]
            Bernoulli & {$p(\vec{z}) = \text{Sigmoid}(f_{15}(\vec{z}))$} & $p(\vec{z}) = 1/2$ \\[0.1cm]
            Beta & $\alpha(\vec{z})=10\cdot \text{Softplus}(f_3(\vec{z}))$, ~  $\beta(\vec{z})=10\cdot \text{Softplus}(f_3(\vec{z}))$ & $(\alpha(\vec{z}), \beta(\vec{z})) = (1, 1)$ \\[0.1cm]
            Dirichlet & $\alpha(\vec{z})=\text{Softplus}(f_3(\vec{z}))$  & $\alpha(\vec{z}) = 1$ \\[0.1cm]
            Exponential & $\lambda(\vec{z})=\text{Softplus}(f_3(\vec{z}))$  & $\lambda(\vec{z}) \to 0$ \\[0.2cm]
    \end{tabular}
    }
    \caption{This table shows the implementations of the decode and extrapolate functions in Eq. (\ref{appendix:eq:reweight}) for all the distributions studied in our second experiment (see Sec. \ref{sec:toy_latent_space_experiment}). Here we represent a randomly initialized neural network with $f_i$, where $i$ represents the size of the co-domain. For example, in the case of the Dirichlet distribution, we use a randomly initialized neural network to compute the parameters $\alpha$ of the distribution and, since these have to be positive, we pass the output of this network through a Softplus activation; moreover, since the Dirichlet distribution is approximately uniform when all its parameters equal 1, our extrapolation mechanism consists of replacing the output of the network with a constant vector of ones.}
    \label{appendix:tab:extrapolation_mechanisms}
\end{table}

\section{Details for our implementation and experiments}

In this section we present the technical details that we used in our implementation and experiments. We will publish our code upon acceptance \href{}{\texttt{here}}.

\subsection{Shortest path approximation with cubic splines}
\label{sec:geodesic_solver_details}

As we have described in the main paper, we propose an approximate solution for the shortest paths based on cubic splines. Let a cubic spline $c(t) = [1,t,t^2,t^3]\T[\psi_0, \psi_1, \psi_2, \psi_3]$ with parameters $\bs{\psi}\in\R^4$, which we optimize using the approximation of the curve energy $
    \psi^* = \argmin_\psi \mathrm{Energy}[c_\psi]
$.
In general, we can use Prop.~\ref{prop:FRenergydistance} as long as we can propagate the gradient through the KL or as in \eqref{eq:approx_energy_categorical} if an explicit closed form solution exists. In this case, we are able to use automatic differentiation for the optimization of the parameters.

\subsection{Models used}

\begin{table}
\centering
\begin{tabular}{@{}rl@{}}
\multicolumn{2}{c}{\textbf{Pulling back the Euclidean vs. Fisher-Rao} (Sec. \ref{sec:compare_metrics_with_gaussians_generators_uq})}\\
\toprule
\textbf{Module}  & \multicolumn{1}{c}{\textbf{MLP}}     \\ \midrule
\multicolumn{2}{c}{Encoder}  \\
$\mu$ & Linear(728, 2) \\
\midrule
\multicolumn{2}{c}{Decoder}  \\
$\mu$ & Linear(2, 728) \\
$\sigma_{\text{UR}}$ & RBF(), PosLinear(500, 1), Reciprocal(), PosLinear(1, 728) \\
$\sigma_{\text{no UR}}$ & Linear(2, 728), Softplus() \\
\midrule
Optimizer  & \multicolumn{1}{@{}l@{}}{Adam ($\alpha = 1 \times 10^{-5}$)}  \\
Batch size             & \multicolumn{1}{@{}l@{}}{32} \\
\bottomrule
\end{tabular}
\vspace{0.2cm}
\caption{ This table shows the Variational Autoencoder used in our first experiment (see Sec. \ref{sec:compare_metrics_with_gaussians_generators_uq}). The network for approximating the standard deviation $\sigma$ leverages ideas from \cite{arvanitidis:iclr:2018}, in which an RBF network is trained on latent codes using centers positioned through KMeans. The operation PosLinear($a,b$) represents the usual Linear transformation with $a$ inputs and $b$ outputs, but considering only positive weights. To compare between having and not having uncertainty regularization, we use two different approximations of the standard deviation in the decoder: $\sigma_{\text{UR}}$ when performing meaningful uncertainty quantification, and $\sigma_{\text{ no UR}}$ otherwise.}
\label{appendix:tab:UQ_experiment_networks}
\end{table}

\begin{table}
    \centering
    \begin{tabular}{rllrr}
        \multicolumn{5}{c}{\textbf{Toy latent spaces} (Sec. \ref{sec:toy_latent_space_experiment})}\\
        \toprule
        \textbf{Distribution} & \textbf{Module} & \textbf{MLP} & \textbf{Seed for randomness} & $\beta$ \textbf{in} $\tilde{\sigma}_\beta$ \\
         \midrule
        \multirow{2}{*}{Normal} & $\mu$ & Linear(2,3) & \multirow{2}{*}{1} & \multirow{2}{*}{-2.5} \\
         & $\sigma$ & Linear(2,3), Softplus() & & \\
         \midrule
         Bernoulli & $p$ & Linear(2,15), Sigmoid() & 1 & -3.5 \\
         \midrule
         \multirow{2}{*}{Beta} & $\alpha$ & Linear(2,3), Softplus() & \multirow{2}{*}{1} & \multirow{2}{*}{-4.0} \\
         & $\beta$ & Linear(2,3), Softplus() & & \\
         \midrule
         Dirichlet & $\alpha$ & Linear(2,3), Softplus() & 17 & -4.0 \\
         \midrule
         Exponential & $\lambda$ & Linear(2,3), Softplus() & 17 & -4.0 \\
         \bottomrule \\[0.1cm]
    \end{tabular}
    \caption{This table describes the neural networks used for the experiment presented in Sec. \ref{sec:toy_latent_space_experiment}. Following the notation of PyTorch, Linear($a,b$) represents an MLP layer with $a$ input nodes and $b$ output nodes. In each of these networks, we implement the reweighting operation described in Sec. \ref{appendix:sec:UQ}, and we describe the $\beta$ hyperparameter present in the modified sigmoid function (Eq. (\ref{appendix:eq:translated_sigmoid})). This networks were not trained in any way, and they were initialized using the provided seed.}
    \label{appendix:tab:toy_experiment_networks}
\end{table}

\begin{table}
\centering
\begin{tabular}{@{}rl@{}}
\multicolumn{2}{c}{\textbf{Decoding to a von Mises-Fisher Distribution} (Sec \ref{sec:motion-capture}, \ref{sec:experiment_numerical_approx_metric}, \ref{sec:LAND_experiment})} \\
\toprule
\textbf{Module}  & \multicolumn{1}{c}{\textbf{MLP}}    \\ \midrule
\multicolumn{2}{c}{Encoder (Normal dist.)}  \\[0.1cm]
$\mu$ & Linear($3\times 26$, 90), Linear(90, 2) \\
$\sigma$ & Linear($3\times 26$, 90), Linear(90, 2), Softplus() \\
\midrule
\multicolumn{2}{c}{Decoder (vMF dist.)}  \\[0.1cm]
$\mu$ & Linear(2, 90), Linear(90, $3\times 26$), Linear($3\times 26$, $3\times 26$)  \\
$\kappa$ & Linear(2, 90), Linear(90, $3\times 26$), Linear($3\times 26$, $26$), Softplus()  \\
\midrule
Optimizer  & \multicolumn{1}{@{}l@{}}{Adam ($\alpha = 1 \times 10^{-3}$)}  \\
Batch size             & \multicolumn{1}{@{}l@{}}{16} \\
$\beta$ in $\tilde{\sigma}_\beta$             & \multicolumn{1}{@{}l@{}}{-5.5} \\
\bottomrule
\end{tabular}
\vspace{0.2cm}
\caption{This table shows the Variational Autoencoder used in our last two experiments (see Sec. \ref{sec:motion-capture}, \ref{sec:experiment_numerical_approx_metric}). Our motion capture data tracked 26 different bones, and thus we decode to a product of 26 different von Mises-Fisher distributions. }
\label{appendix:tab:vMF_experiments}
\end{table}

In this section we describe, in detail, the models that we used for our experiments (see Sec.\ref{sec:experiments}). All the networks that we used are Multi-Layer Perceptrons implemented in PyTorch.

First, Table \ref{appendix:tab:UQ_experiment_networks} shows the Variational Autoencoder implemented for the experiment described in Sec. \ref{sec:compare_metrics_with_gaussians_generators_uq}. In the computations \textit{without} uncertainty regularization, we used a simpler model for the uncertainty quantification (namely, a single Linear layer, followed by a Softplus activation). For our second experiment involving a toy latent space, we also provide the implementation of the respective MLPs in Table \ref{appendix:tab:toy_experiment_networks}. Finally, Table \ref{appendix:tab:vMF_experiments} presents the VAE trained for the experiments related to motion capture data: pulling back the Fisher-Rao metric from a product of von Mises-Fisher distributions (Sec. \ref{sec:motion-capture}), approximating the metric locally by solving a system of linear equations (Sec. \ref{sec:experiment_numerical_approx_metric}), and fitting a LAND density in this latent space (Sec. \ref{sec:LAND_experiment}).

\subsection{Metric approximation and KL by sampling}

When visualising our latent space as a statistical manifold, we can obtain a direct approximation of the metric using the KL-divergence between two close distributions (Proposition \ref{prop:approximating_metric}). We will show here, in simple cases, how our metric approximation compares to close-form expressions.

In the following experiment, our statistical manifold is the parameter space of known distributions (Beta and Normal). Their Fisher-Rao matrices are well-known (Sec. \ref{appendix:FR:distributions}), and we approximate them by computing the KL-divergence of sampled distributions. We call $\b{M}_{t}$ the theoretical metric and $\b{M}_{a}$ the approximated metric, and we note $\varepsilon_{r} = \frac{\lVert \b{M}_t - \b{M}_a \rVert}{\lVert \b{M}_t\rVert}$ the relative error between the theoretical and approximated matrices. Here, $\lVert \cdot \rVert$ denotes the Frobenius norm. For the Normal distribution, we empirically obtain: $\varepsilon_{r} =  5.32 \cdot 10^{-4} \pm 9.63 \cdot 10^{-4}$, and for the Beta distribution, we have: $\varepsilon_{r} =  1.73 \cdot 10^{-5} \pm 1.17 \cdot 10^{-5}$.

\subsection{Information for fitting the LAND model}

The locally adaptive normal distribution (LAND) \citep{arvanitidis:nips:2016} is the extension of the normal distribution on Riemannian manifolds learned from data. \citet{pennec:jmiv:2006} first derived this distribution on predefined manifolds as the sphere and also showed that it is the maximum entropy distribution given a mean and a precision matrix. The flexibility of this probability density relies on the shortest paths. However, the computational demand to fit this model is relatively high, especially in our case, since we need to use an approximation scheme to find the shortest paths.

In particular, we compute the \emph{logarithmic map} $\b{v}=\mathrm{Log}_\vec{x}(\b{y})$ by first finding the shortest path between $\vec{x}$ and $\b{y}$, and then, rescaling the initial velocity as $\b{v} = \frac{\dot{c}(0)}{||\dot{c}(0)||} \mathrm{Length}(c)$, which ensures that $||\b{v}|| = \mathrm{Length}(c)$. In addition, for the estimation of the normalization constant we use the \emph{exponential map} $\mathrm{Exp}_\vec{x}(\b{v})=c_\b{v}(t)$, which is the inverse operator that generates the shortest path with $c(1)=\b{y}$ taking the rescaled initial velocity $\b{v}$ as input. Also, we should be able to evaluate the metric. While the logarithmic map can be approximated using our approach (Section~\ref{sec:geodesic_solver_details}), for the exponential map we need to solve the ODEs system \eqref{eq:ode_system} as an initial value problem (IVP). Note that we fit the LAND using gradient descent, which implies that the computation of these operators is the main computational bottleneck.

We provided a method in Proposition~\ref{prop:approximating_metric}, which enables us to approximate the pullback metric in the latent space of a generative model using the corresponding KL divergence. Even if this is a sensible approach, in practice, the computational cost is relatively high as we might need to estimate the KL using Monte Carlo. For example, this is the case when the likelihood is the von Mises-Fisher. This further implies that fitting the LAND using this approach is prohibited due to the computational cost. Especially, since we need to evaluate many times the metric and its derivative for the computation of each exponential map. Hence, in order to fit the LAND efficiently, we used the following approximation based on \citet{hauberg:nips:2012}.

First we construct a uniformly spaced grid in the latent space. Then, we evaluate the metric using Proposition~\ref{prop:approximating_metric} for each point on the grid getting a set $\{\vec{z}_s, \b{M}_s\}_{s=1}^S$ of metric tensors. Thus, we can estimate the metric at any point $\vec{z}$ as
\begin{equation}
    \b{M}(\vec{z}) = \sum_{s=1}^S \widetilde{w}_s(\vec{z}) \b{M}_s, ~~ \text{with} ~~ \widetilde{w}_s(\vec{z}) = \frac{w_s(\vec{z})}{\sum_{j=1}^S w_s(\vec{z})} ~~ \text{and} ~~ w_s(\vec{z}) = \exp\left(-\frac{ ||\vec{z}_s - \vec{z}||^2}{2\sigma^2}\right)
\end{equation}
where $\sigma >0$ the bandwidth parameter. This is by definition a Riemannian metric as a weighted sum of Riemannian metrics with a smooth weighting function. In this way, we can approximate the pullback of the Fisher-Rao metric in the latent space $\Z$ in order to perform the necessary computations more efficiently.

\clearpage

\section{Introduction}
\label{sec:intro}

Generative models such as \emph{variational autoencoders (VAEs)} \citep{kingma:iclr:2014, rezende:icml:2014} and \emph{generative adversarial networks (GANs)} \citep{goodfellow:neurips:2014} provide state-of-the-art density estimators for high dimensional data. The underlying assumption is that data $\vec{x} \in \X$ lie near a low-dimensional manifold $\M \subset \X$, which is parametrized through a low-dimensional \emph{latent representation} $\vec{z} \in \Z$. As data is finite and noisy, we only recover a probabilistic estimate of the true manifold, which, in VAEs, is represented through a decoder distribution $p(\vec{x} | \vec{z})$. Our target is the geometry of this \emph{random manifold}.

The geometry of the manifold has been shown to carry great value when systematically interacting with the latent representations, as it provides a stringent solution to the \emph{identifiability problem} that plagues latent variable models \citep{tosi:uai:2014, arvanitidis:iclr:2018, hauberg:only:2018}. For example, this geometry has allowed VAEs to discover latent evolutionary signals in proteins \citep{detlefsens:proteins:2020}, provide efficient robot controls \citep{scannellTrajectory2021, chen2018active}, improve latent clustering abilities \citep{yang:arxiv:2018, arvanitidis:iclr:2018} and more. The fundamental issue with these geometric approaches is that the studied manifold is inherently a stochastic object, but classic differential geometry only supports the study of \emph{deterministic} manifolds. To bridge the gap, \citet{eklund:arxiv:2019} have shown how VAEs with a Gaussian decoder family can be viewed as a random projection of a deterministic manifold, thereby making the classic theories applicable to the random manifold.

A key strength of VAEs is that they can model data from diverse modalities through the choice of decoder distribution $p(\vec{x} | \vec{z})$. For discrete data, we use categorical decoders, while for continuous data we may opt for a Gaussian, a Gamma or whichever distribution best suits the data. But for non-Gaussian decoders, there exists no useful approach for treating the associated random manifold as deterministic, which prevents us from systematically interacting with the latent representations without being subject to identifiability issues.

\textbf{In this paper}, we provide a general framework that allows us to interact with the geometry of almost any random manifold. The key, and simple idea is to reinterpret the decoder as spanning a deterministic manifold in the space of probability distributions $\HH$, rather than a random manifold in the observation space (see Fig.~\ref{fig:teaser_img}). Calling on classical \emph{information geometry} \citep{amari:2016, nielsen:2020}, we show that the learned manifold is a Riemannian submanifold of $\HH$, and provide corresponding computational tools. The approach is applicable to any family of decoders for which the KL-divergence can be differentiated, allowing us to work with a wide array of models from a single codebase.\looseness=-1

\begin{SCfigure}
    \includegraphics[width=4cm]{imgs/concept_v1_mnist.pdf}

    \includegraphics[width=4cm]{imgs/concept_v1_probabilities2.pdf}
    \caption{Traditionally (left), we view the learned manifold as a stochastic manifold in the observation space. We propose (right) to view the learned manifold as a deterministic manifold embedded in the space of decoder distributions, which has an information geometry.}
    \label{fig:teaser_img}
\end{SCfigure}

\section{The geometry of generative models}
\label{sec:riemannian_geometry}

As a starting point, consider the deterministic generative model given by a prior $p(\vec{z})$ and a decoder $f: \Z=\R^d \rightarrow \X=\R^D$, which is assumed to be a smooth immersion. The latent representation $\vec{z}$ of an observation $\vec{x}$ is generally not \emph{identifiable}, meaning that one can recover different latent representations that give rise to equally good density estimates. For example, let $g: \Z \rightarrow \Z$ be a smooth invertible function such that $\vec{z} \sim p(\vec{z}) \Leftrightarrow g(\vec{z}) \sim p(\vec{z})$, then the latent representation $g(\vec{z})$ coupled with the decoder $f \circ g^{-1}$ gives the same density estimate as $\vec{z}$ coupled with $f$ \citep{hauberg:only:2018}. Practically speaking, the identifiability issue implies that it is improper to view the latent space $\Z$ as being Euclidean, as any reasonable view of $\Z$ should be invariant to reparametrizations $g$.

  The classic geometric solution to the identifiability problem is to define any quantity of interest in the observation space $\X$ rather than the latent space $\Z$. For example, the length of a curve $\gamma: [0, 1] \rightarrow \Z$ in the latent space can be defined as its length measured in $\X$ on the manifold $\M=f(\Z)$ as:
  \begin{align}
    \mathrm{Length}(\gamma)
      &= \lim_{N \rightarrow \infty} \sum_{n=1}^{N-1} \| f(\sfrac{n}{N}) - f(\sfrac{n+1}{N}) \|
       = \int_0^1 \| \partial_t f(\gamma(t)) \| \dif t \\

      &= \int_0^1 \sqrt{ \dot{\gamma}(t)\T \mat{J}_f({\gamma(t)})\T \mat{J}_f({\gamma(t)}) \dot{\gamma}(t) } \dif t,
     \label{eq:curve_length}
  \end{align}
  where we used the chain rule $\partial_t f(\gamma(t)) = \b{J}_f(\gamma(t))\dot{\gamma}(t)$ with $\dot{\gamma}(t) = \partial_t \gamma(t)$ being the curve derivative, and $\mat{J}_f({\gamma(t)})\in\R^{D\times d}$ the Jacobian of $f$ at $\gamma(t)$. This construction shows how we may calculate lengths in the latent space with respect to the metric of the observation space, which is typically assumed to be the Euclidean, but other options exist \citep{arvanitidis:aistats:2021}. In this way, the symmetric positive definite matrix $\mat{J}_f({\gamma(t)})\T \mat{J}_f({\gamma(t)})$ is denoted by $\b{M}({\gamma(t)})\in\R_{\succ 0}^{d\times d}$ and captures the geometry of $\M$ in $\Z$. This is known as the \emph{pullback metric} as it pulls the Euclidean metric from $\X$ into $\Z$. As the Jacobian spans the $d$-dimensional tangent space at the point $\vec{x} = f(\vec{z})$, we may interpret $\b{M}(\vec{z})$ as an inner product $\inner{\vec{u}}{\vec{v}}_{\b{M}} = \vec{u}\T\b{M}(\vec{z})\vec{v}$ over this tangent space, given us all the ingredients to define \emph{Riemannian manifolds}:
  \begin{definition}
    A Riemannian manifold is a smooth manifold $\M$ together with a Riemannian metric $\b{M}({\vec{z}})$, which is a positive definite matrix that changes smoothly throughout space and defines an inner product on the tangent space $\mathcal{T}_{\vec{z}}\M$.

  \end{definition}

We see that the decoder naturally spans a Riemannian manifold and the latent space $\Z$ can be considered as the \emph{intrinsic coordinates}. Technically, we can consider any Euclidean space as the intrinsic coordinates of an abstract $\M$ using a suitable metric $\b{M}(\vec{z})$, which is implicitly induced by an abstract $f$. Since the Riemannian length of a latent curve \eqref{eq:curve_length}, by construction, is invariant to reparametrizations, it is natural to extend this view with a notion of \emph{distance}. We say that the distance between two points $\vec{z}_0, \vec{z}_1 \in \Z$ is simply the length of the shortest connecting path, $\mathrm{dist}(\vec{z}_0, \vec{z}_1) = \min_{\gamma} \mathrm{Length}(\gamma)$. Calculating distances implies finding the shortest path. One can show \citep{gallot:2004} that length minimizing curves also have minimal \emph{energy}:
  \begin{align}
    \mathrm{Energy}(\gamma)
      &= \int_0^1 \| \partial_t f(\gamma(t)) \|^2 \dif t
       = \int_0^1 \dot{\gamma}(t)\T \b{M}({\gamma(t)}) \dot{\gamma}(t) \dif t,
      \label{eq:curve_energy}
  \end{align}
  which is a locally convex functional. Shortest paths can then be found by direct energy minimization \citep{yang:arxiv:2018} or by solving the associated system of ordinary differential equations (ODEs) \citep{hennig:aistats:2014, arvanitidis:aistats:2019} (see supplementary materials).

\begin{SCfigure}
    \begin{overpic}[width=4cm]{imgs/obs_pullback.pdf}
    \put(0, 80){$\M\subset\X$}
    \put(22,22){\color{black}\vector(1,1){20}}
    \put(10, 10){$\mu(\gamma(t))$}
    \end{overpic}

    \qquad

    \begin{overpic}[width=3cm]{imgs/latent_pullback.pdf}
    \put(0, 80){$\Z$}
    \put(45,50){\color{black}\vector(1,-1){20}}
    \put(25, 50){$\gamma(t)$}
    \end{overpic}

    \caption{A conceptual example of a Riemannian manifold $\M=\mu(\Z)$ lying in $\X$ and the corresponding latent space $\Z$, together with an associated shortest paths.}
    \label{fig:riem_geom_example}
\end{SCfigure}

 \subsection{Stochastic decoders}

As previously discussed, deterministic decoders directly induce a Riemannian geometry in the latent space. However, most models of interest are stochastic and there is significant evidence that this stochasticity is important to faithfully capture the intrinsic structure of data \citep{hauberg:only:2018}. When the decoder is a smooth stochastic process, e.g.\ as in the Gaussian Process Latent Variable Model (GP-LVM) \citep{lawrence:2005:jmlr}, \citet{tosi:uai:2014} laid the ground works for modeling a stochastic geometry. Most contemporary models, such as VAEs, are however not stochastically smooth as the noise is assumed independent, and this theory is inapplicable. \citet{arvanitidis:iclr:2018} proposed an extension to VAEs with Gaussian decoders which takes the form
\begin{align}
    f(\vec{z})
      &= \mu(\vec{z}) + \sigma(\vec{z}) \odot \vec{\epsilon}
       = \begin{bmatrix}
           \I_{D} & \mathrm{diag}(\vec{\epsilon})
         \end{bmatrix}
         \begin{bmatrix}
           \mu(\vec{z}) \\ \sigma(\vec{z})
         \end{bmatrix}
      = \mat{P}_\varepsilon~ h(\vec{z}),
    \quad \text{where } \vec{\epsilon} \sim \mathcal{N}(\vec{0}, \I_{D}).
\end{align}
Here we have written the Gaussian decoder in its reparametrized form. This can be viewed as a random projection of a deterministic manifold spanned by $h$ with projection matrix $\b{P}_\epsilon$ \citep{eklund:arxiv:2019}, which can easily be given a geometry. The associated Riemannian metric,
\begin{align}
\label{eq:vae_metric}
    \b{M}({\b{z}}) = \b{J}_\mu(\vec{z})\T\b{J}_\mu(\vec{z}) + \b{J}_\sigma(\vec{z})\T\b{J}_\sigma(\vec{z}),
\end{align}
gives shortest paths that follow the data as distances grow with the model uncertainty \citep{arvanitidis:iclr:2018, hauberg:only:2018}. An example of a shortest path $\gamma(t)\in\Z$ computed under this metric is shown in Fig.~\ref{fig:riem_geom_example} and the respective curve on the corresponding expected manifold $\mu(\gamma(t))\in\M\subset\X$.

Previous work has, thus, focused on \emph{pulling back} the Euclidean metric from the observation space to the latent space using the reparametrization of the Gaussian decoder. This is, however, intrinsically linked with the simple reparametrization of the Gaussian, and this strategy can only extend to location-scale distributions. We propose an alternative, principled way of dealing with stochasticity by changing the focus from observations to distributions, leveraging the metrics defined in classical information geometry.\looseness=-1

\section{An information geometric latent metric}
\label{sec:information_geometry}

So far we have seen how we can endow the latent space $\Z$ with meaningful distances only when our stochastic decoders are reparameterizable and their codomain is the observation space $\X$. Ideally, we would like a more general framework of computing shortest path distances for a more general class of distributions.

We first note that the codomain of a VAE decoder is the parameter space $\HH$ of a probability density function. In particular, depending on the type of data a likelihood $p(\vec{x}|\eta)$ is specified with parameters $\eta\in\HH$ and the mapping $h:\Z\rightarrow\HH$ allows us to write the likelihood as $p(\vec{x}|\vec{z})$.

With this in mind, we can ask what is a natural distance in the latent space $\Z$ between two infinitesimally near points $\vec{z}_1$ and $\vec{z}_2 = \vec{z}_1\!+\!\epsilon$ when measured in $\HH$. Since our latent codes map to distributions we can define the (infinitesimal) distance through the KL-divergence:
\begin{equation}
    \mathrm{dist}^2(\vec{z}_1, \vec{z}_2) = \mathrm{KL}(p(\vec{x}~|~\vec{z}_1),~p(\vec{x}~|~\vec{z}_2)).
    \label{eq:kl_dist}
\end{equation}

We can then define the length of a latent curve $\gamma:[0, 1]\rightarrow \Z$ as
\begin{align}
  \mathrm{Length}(\gamma)
    &= \lim_{N \rightarrow \infty} \sum_{n=1}^{N-1} \sqrt{\mathrm{KL}(p(\vec{x}~|~\gamma(\sfrac{n}{N})),~ p(\vec{x}~|~\gamma(\sfrac{n\!+\!1}{N})))},
\end{align}
and distances could be defined as before.
This would satisfy our desiderata of a deterministic notion of similarity in the latent space that is applicable to wide range of decoder distributions.

This construction may seem arbitrary, but in reality it carries deeper geometric meaning. \emph{Information geometry} \citep{nielsen:2020} considers families of probabilistic densities $p(\vec{x}|\eta)$ as represented by their parameters $\eta \in \HH$, such that $\HH$ is constructed as a statistical manifold equipped with the Fisher-Rao metric, which infinitesimally coincides with the KL divergence in Eq.~\ref{eq:kl_dist}. This is known to be a Riemannian metric over $\HH$ that takes the following form:

\begin{equation}
\label{eq:fisher_rao_metric}
    \b{I}_{\HH}(\eta) = \int_\X \left[\nabla_{\eta} \log p(\b{x}~|~\eta) \cdot \nabla_{\eta} \log p(\b{x}~|~\eta)\T\right] p(\b{x}~|~\eta)\dif \b{x}.
\end{equation}
When the parameter space $\HH$ is equipped with this metric, we call it a \emph{statistical manifold}.
\begin{definition}
  A statistical manifold consists of the parameter space $\HH$ of a probability density function $p(\b{x}|\eta)$ equipped with the Fisher-Rao information matrix $\b{I}_{\HH}(\eta)$ as a metric.
\end{definition}
Note that the geometry induced by the Fisher-Rao metric is predefined and can be seen as a modeling decision, since it does not change with data.
As mentioned above, the Fisher-Rao metric coincides with the KL-divergence locally:
\begin{proposition}
  The Fisher-Rao metric is the first order approximation of the KL-divergence between perturbed distributions:
  \begin{equation}
  \label{eq:KL_approx_inner_product}
    \mathrm{KL}(p(\b{x}~|~\eta),~p(\b{x}~|~\eta+\delta\eta)) = \frac{1}{2} \delta\eta\T \b{I}_{\HH}(\eta) \delta\eta + o(\delta\eta^2).
  \end{equation}
\end{proposition}
The central idea put forward in this paper is to consider the decoder as a map $h: \Z \rightarrow \HH$ instead of $f:\Z\rightarrow\X$, and let $\HH$ be equipped with the appropriate Fisher-Rao metric. The VAE can then be interpreted as spanning a manifold $h(\Z)$ in $\HH$ and the latent space $\Z$ can be endowed with the corresponding metric. We detail this approach in the sequel.

\subsection{The Riemannian pull-back metric}
Our construction implies that the length of a latent curve $\gamma:[0,1]\rightarrow \Z$ when mapped through $h$ can be measured in the parameter space $\HH$ using the Fisher-Rao metric therein as
\begin{align}
    \mathrm{Length}(\gamma)
      &= \int_0^1 \sqrt{\partial_t h(\gamma(t))\T \b{I}_{\HH}(h(\gamma(t))) \partial_t h(\gamma(t))} \dif t,
\end{align}
  %
which is obtained from the pullback metric:
  %
\begin{proposition}
    \label{prop:fisher_pullback}
    Let $h:\Z\rightarrow\HH$ be an immersion that parametrizes the likelihood. Then, the latent space $\mathcal{Z}$ is equipped with the pull-back metric $\b{M}(\b{z}) = \b{J}_{h}\T(\b{z}) \b{I}_{\HH}(h(\b{z})) \b{J}_{h}(\b{z})$.
\end{proposition}

Instead of considering the parameters $\eta\in\HH$ of the probabilistic density function that models the data, we can consider the latent variable $\b{z}$ as the parameters of the model. This view is equivalent to the one explained above, and the corresponding pull-back metric is directly the Fisher-Rao metric endowed in the latent space $\Z$:
  %
\begin{proposition}
    The pullback metric $\b{M}(\b{z})$ is identical to the Fisher-Rao metric obtained over the parameter space $\Z$ as $\b{M}(\b{z}) = \int_{\X} \left[\nabla_\b{z}\log p(\b{x}|\b{z}) \cdot \nabla_\b{z} \log p(\b{x}|\b{z})\T\right] p(\b{x}|\b{z}) \dif\b{x}. $
\end{proposition}

Therefore, pulling back the Fisher-Rao metric from $\HH$ into $\Z$ enables us to compute length minimizing curves which are indentifiable. The advantange of this approach is that it applies to any type of decoders and data, as the actual distance is measured over the manifold spanned by $h$ in the parameter space $\HH$. So shortest paths between probability distributions move optimally on this manifold while taking the geometry of $\HH$ into account through the Fisher-Rao metric.

Computing shortest paths directly in $\HH$ need not result in a sensible sequence of probability density functions $p(\b{x}|\eta)$. To ensure that the shortest paths computed under our metric stay within the support of the data, we carefully design our decoder $h$ to extrapolate to uncertain distributions outside the support of the data (see supplements).

\begin{wrapfigure}[14]{r}{0.3\textwidth}
    \vspace{-5mm}
    \begin{center}
        \includegraphics[width=0.3\textwidth]{imgs/new_comp.pdf}
    \end{center}

    \caption{The distributions $p(\b{x}|\gamma(t))$ under $\b{M}(\vec{z})$ respect the structure of data.

    The curve with minimal length in $\HH$ does not as it leaves $\M$.}
    \label{fig:inf_geom_compare_ours_vs_fisher}
\end{wrapfigure}
In Fig.~\ref{fig:inf_geom_compare_ours_vs_fisher} we compare a shortest path  $\gamma:[0,1]\rightarrow\Z$ under the proposed metric $\b{M}(\vec{z})$ against a curve $c:[0,1]\rightarrow\HH$ with minimal length. We consider a Gaussian likelihood with isotropic covariance. We show the resulting sequence of means for both interpolants color-coded by the corresponding variances. As expected $c(t)$ does not take into account the given data, but only respects the geometry of $\HH$ implied by the likelihood.

\subsection{Efficient geodesic computation}
\label{sec:geodesic_computation}

  An essential task in computational geometry is to compute shortest paths.
  This can be achieved by minimizing curve energy \eqref{eq:curve_energy} or solving the corresponding system of ODEs (see Supplementary Material). The latter, however, requires inordinate computational resources, since the evaluation of the system relies on the Jacobian of the decoder and its derivatives.

  Bearing in mind that the metric is an approximation of the KL divergence between perturbations \eqref{eq:KL_approx_inner_product}, the energy is directly expressed as a sum of KL divergence terms along a discretized curve $\gamma$:
  \begin{equation}
    \label{eq:our_energy_with_kl}
        \mathrm{Energy}(\gamma) = \lim_{N \rightarrow \infty} \frac{2}{N} \ \sum_{n=1}^{N-1} \ \mathrm{KL}(p(\vec{x}~|~\gamma(\sfrac{n}{N})),~ p(\vec{x}~|~\gamma(\sfrac{n\!+\!1}{N}))).
  \end{equation}
  A simple algorithm for computing shortest paths is to minimize Eq.~\ref{eq:our_energy_with_kl} with respect to the parameters of a curve $\gamma$. Here we represent $\gamma$ as a cubic spline with fixed end-points. Then standard free-form optimization can be applied to minimize energy.

\subsection{Example: categorical decoders}
\label{sec:example_categorical}

The motivation for our approach is that while several options for decoders exist in VAEs depending on the type of the given data, we could only capture and use the learned geometry in a principled way with Gaussian decoders. Our proposed methodology is significantly more general.

For a constructive example, assume that $\vec{x}$ is a categorical variable. We can select a generalized Bernoulli likelihood $p(\vec{x}|\vec{z})$, such that $h(\vec{z}) = (\eta_1, \cdots, \eta_D)$ where each $\eta_i$ represents the probability of $x_i$ being 1. Therefore, we know that the parameters $\eta$ lie on the unit simplex $\HH$. The Fisher-Rao distance between such categorical variables are known to coincide with spherical distance between points $\sqrt{\eta}$ on the unit sphere (here the square root is element-wise),
\begin{equation}
    \mathrm{dist}(\eta,\eta') = \mathrm{arccos}\left(\sqrt{\eta}\T\sqrt{\eta'}\right).
\end{equation}
We derive this previously known result in the supplementary materials.

Given a curve $\gamma:[0,1]\rightarrow\Z$ we can approximate the energy by using the small angle approximation $\cos\theta \approx 1 - \sfrac{\theta^2}{2} \Leftrightarrow \theta^2 \approx 2 - 2\cos\theta$ to give
\begin{equation}
    \mathrm{Energy}(\gamma) = \sum_{n=1}^{N-1}\left(2 - 2\sqrt{h(\gamma(\sfrac{n}{N}))}\T\sqrt{h(\gamma(\sfrac{n+1}{N}))}\right),
\end{equation}
for sufficiently fine discretization. This gives a particular simple expression for the energy, which can be minimized to compute geodesics.

\subsection{Black-box random geometry}\label{sec:blackbox}
  While it is generally possible to derive suitable expressions for computing metrics and energies for a given decoder family, doing so is tedious, error-prone and time-consuming. This limits the practical use of the developed theory. Drawing inspiration from \emph{black-box variational inference} \citep{ranganath:aistats:2014}, we propose a notion of \emph{black-box random geometry}. Assuming we have access to a differentiable KL divergence for our choice of decoder distribution, we can apply the algorithm from Sec.~\ref{sec:geodesic_computation} to compute geodesics. In practice, modern libraries such as PyTorch \citep{paszke:neurips:2019} have this functionality implemented for several distributions, so the route is easily followed. For some distributions, we do not have closed-form expression for the KL divergence, in which case we can resort to Monte Carlo estimates thereof. More specifically, we can estimate the KL divergence by generating samples from the likelihood based on the re-parametrization trick, which allows us to get derivatives with automatic differentiation. Such an approach is e.g.\ also used in the probabilistic programming language Pyro \citep{bingham:jmlr:2019}.

  Interestingly, apart from finding the shortest path through the KL formulation, we can also approximate the actual metric tensor $\b{M}(\vec{z})$. As we have discussed above, evaluating explicitly this metric is not a trivial task in many cases. One problem is that we need access to the Jacobian of the parametrization $h$, which is typically a deep neural network, so the computation is not always straightforward. Alternatively, one could use that the Fisher-Rao metric is the Hessian of the KL-divergence \eqref{eq:KL_approx_inner_product}, but such approaches fare poorly with current tools for automatic differentiation, where higher-order derivatives are often incompatible with batching. Furthermore, the Fisher-Rao metric itself may be intractable depending on the chosen likelihood $p(\vec{x}|\eta)$. Nevertheless, we show that the KL formulation \eqref{eq:KL_approx_inner_product} allows us to approximate the latent metric as follows:
    \begin{proposition}
    \label{prop:approximating_metric}
    We define perturbations vectors as: $\delta e_{i} = \varepsilon \cdot \b{e_i}$, with $\varepsilon\in\mathbb{R}_{+}$ a small infinitesimal quantity, and ($\b{e_i}$) a canonical basis vector in $\mathbb{R}^d$. For better clarity, we rename $\mathrm{KL}(p(\vec{x}|\vec{z}), p(\vec{x}|\vec{z}+\delta \vec{z})) = \mathrm{KL}_{\vec{z}}(\delta \vec{z})$ and we note $\b{M}_{ij}$ the components of $\b{M}(\vec{z})$.
    We can then approximate by a system of equations the diagonal and non-diagonal elements of the metric:
        \begin{align*}
            \b{M}_{ii} &\approx 2 \ \mathrm{KL}_{\vec{z}}(\delta \b{e_i})/ \varepsilon^2 \ \\
            \b{M}_{ij} = \b{M}_{ji} &\approx  \left(\mathrm{KL}_{\vec{z}}(\delta \b{e_i} + \delta \b{e_j})- \mathrm{KL}_{\vec{z}}(\delta \b{e_i}) -\mathrm{KL}_{\vec{z}}(\delta \b{e_j}) \right)/ \varepsilon^2.
        \end{align*}
    \end{proposition}
  Note that this formulation only requires $h$ to be a smooth immersion. This is particularly useful, as the metric is used for other purposes on a Riemannian manifold and not exclusively for computing shortest paths. For example, we can compute the exponential map by solving the corresponding ODE system as an initial value problem relying on $\b{M}(\vec{z})$. Assuming a fully differentiable KL-divergence, then the approximated metric is also differentiable. This is all that is required for practical usage of differential geometry, and we have a reasonable notion of \emph{black-box random geometry}.

\section{Experiments}
\label{sec:experiments}

 \begin{figure}[t!]
      \resizebox{\textwidth}{!}{
      \begin{tabular}{cccc}
        \includegraphics[height=5cm]{imgs/Gaussian_euc.pdf} &
        \includegraphics[height=5cm]{imgs/Gaussian_euc_badUQ.pdf} &
        \includegraphics[height=5cm]{imgs/Gaussian_FI.pdf} &
        \includegraphics[height=5cm]{imgs/Gaussian_FI_badUQ.pdf} \\
      \end{tabular} }

      \caption{Pulling back the Euclidean and Fisher-Rao metrics with Gaussian decoders. Left to right: Euclidean pull-back with regularized uncertainty, Euclidean pull-back with a NN to model uncertainty, Fisher-Rao pull-back with regularized uncertainty, Fisher-Rao pull-back with a NN to model uncertainty.}
      \label{fig:gaussian_decoders_UQ}
  \end{figure}

\subsection{Pulling back Euclidean or Fisher-Rao metric with Gaussian decoders}
\label{sec:compare_metrics_with_gaussians_generators_uq}

We start our experiments by comparing our proposed way of inducing geometry in latent spaces with the existing theory: pulling back the Euclidean metric using a stochastic Gaussian decoder (see Eq.~\ref{eq:vae_metric}). We also include in this comparison the effect of regularizing the uncertainty quantification in the learned geometries. In this regularization, we use transition networks \citep{detlefsen:2019:reliable} to ensure high uncertainty outside the support of the data (see supplements).

In this experiment, we train four VAEs on a subset of the MNIST dataset composed of only the digits with label 1. Two of these VAEs implement a standard Gaussian decoder, and we induce a metric in the latent space by pulling the Euclidean metric back using the Jacobian of the decoder. In the other two, we consider the co-domain as a stochastic manifold and approximate the pullback of the Fisher-Rao metric by using the KL divergence locally. In these two sets, one of the decoder implements the uncertainty regularization described above.

\begin{figure}[t!]
    \centering
    \begin{subfigure}{0.19\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/latent_space_normal.png}
    \end{subfigure}

    \begin{subfigure}{0.19\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/latent_space_bernoulli.png}
    \end{subfigure}

    \begin{subfigure}{0.19\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/latent_space_beta.png}
    \end{subfigure}

    \begin{subfigure}{0.19\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/latent_space_dirichlet.png}
    \end{subfigure}

    \begin{subfigure}{0.19\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/latent_space_exp.png}
    \end{subfigure}
    \caption{Pulling back the metric from different parameter spaces. From left to right: Normal, Bernoulli, Beta, Dirichlet and Exponential. White areas represent low entropy of the decoded distribution, while blue areas represent higher entropy. Notice that the Bernoulli latent space is darker blue (i.e.\@ more entropic) because distributions with parameters around $\sfrac{1}{2}$ are near uniform.}
    \label{fig:latent_spaces_toy_experiment}
\end{figure}

Figure \ref{fig:gaussian_decoders_UQ} shows the latent spaces of these four decoders, illuminated by the volume measure. In each of this latent spaces, we analyze the geometry induced by the respective pullbacks by computing and plotting several geodesics. This figure illustrates two key findings: (1) Our approach is on par with the existing literature in learning geometric structure, which can be seen by comparing the first and third latent spaces (Euclidean vs. Fisher Rao, respectively), and (2) Performing uncertainty regularization plays an instrumental role on learning a sensible geometric structure, which can be seen when comparing the first and second latent spaces (both coming from the Euclidean pullback, with and without regularization respectively), and similarly for the third and fourth.

\subsection{Pulling back geometry for a toy latent space}
\label{sec:toy_latent_space_experiment}

For our second experiment, we induced a geometry on a known latent space (given by noisy circular data in $\mathcal{Z}_{\text{toy}} = \mathbb{R}^2$) by \textit{pulling back} the Fisher-Rao metric from the parameter space of different distributions, showcasing the potential for computing geodesics efficiently, even in non-Gaussian settings. These latent codes were constructed by sampling 200 angles $\{\theta_i\}_{i=1}^{200}$, constructing vectors $[\cos(\theta), \sin(\theta)] + 0.1\epsilon$, where $\epsilon$ is sampled from a unit Gaussian. The stochastic manifolds from which we pull the metric are associated with multivariate versions of the Normal, Bernoulli, Beta, Dirichlet and Exponential distributions. For this approximation to follow the support of the data we need to ensure that our mapping $\mathcal{Z}_{\text{toy}}\to\mathcal{H}$ extrapolates to high uncertainty outside our training codes (see Fig. \ref{fig:gaussian_decoders_UQ}). To do so, we perform uncertainty regularization for each one of the decoded distributions (see supplementary materials for details on uncertainty regularization).

Figure~\ref{fig:latent_spaces_toy_experiment} shows the toy latent space alongside several geodesics computed using the pulled-back the Fisher-Rao metric from the stochastic manifolds associated with the Gaussian, Bernoulli, Beta, Dirichlet and Exponential distributions. These geodesics were embedded as splines by minimizing their energy (as presented in Eq. (\ref{eq:our_energy_with_kl})), using automatic differentiation. These results show that the approximated pulled-back metric induces a meaningful geometry in this latent space, which recovers the initial circular structure of the data. In the case of the Bernoulli manifold, we notice that some of the geodesics fail to converge. We hypothesize that our uncertainty regularization (which decodes to the uniform outside the support) is not strong enough since Bernoulli distributions with parameters close to $1/2$ are already highly entropic.

It is worth stating that these computations were performed using the analytical expression for the KL divergence when minimizing the geodesics' energy (Eq. (\ref{eq:our_energy_with_kl})). We also experimented with computing these divergences by sampling, arriving at similar results when the number of samples was above the threshold of $5\times10^3$. Notice how this opens the doors for approximating the geometry of the statistical manifold of any distribution that allows sampling with derivatives.

\subsection{Motion capture data with products of von Mises-Fisher distributions}
\label{sec:motion-capture}

As a further demonstration of our black-box random geometry, we consider a model of human motion capture data. Here we observe a time series, where each time point represent a `skeleton' corresponding to a human pose. As only pose, and not shape, changes over time, individual limbs on the body only change position and orientation, but not length. Each limb is then a point on a sphere in $\mathbb{R}^3$ with radius given by the limb length. Following \citet{tournier2009motion} we view the skeleton representation space as a product of spheres. From this, we build a VAE where the decoder distribution is a product of von Mises-Fisher distributions. To ensure a sensible uncertainty estimates in the decoder, we enforce that the concentration parameter extrapolate to a small constant.

In this, we do not have easily accessible Fisher-Rao metrics, and instead lean on the KL-formulation from Sec.~\ref{sec:blackbox}. However, the KL-divergence does not have a closed-form expression for the von Mises-Fisher distribution, so we resort to a Monte Carlo estimate thereof. This is realisable with off-the-shelf tools \citep{s-vae18}.

\begin{wrapfigure}[12]{r}{0.3\textwidth}
    \vspace{-5mm}
    \includegraphics[width=\linewidth]{imgs/vmf_geodesics.pdf}
    \caption{Shortest paths in the latent space of a motion capture model with a product of von Mises-Fisher decoder.}
    \label{fig:vmf_geodesics}
\end{wrapfigure}
Figure~\ref{fig:vmf_geodesics} show the latent representation of a motion capture sequence of a person walking\footnote{Seq.\ \texttt{69\_06} from \url{http://mocap.cs.cmu.edu/}.} with geodesics superimposed. As before, we see that geodesics follow the trend of the data, and reflect the underlying periodic nature of the observed walking motion. We pick two random points in the latent space, and traverse both the geodesic and the straight line implied by a Euclidean interpretation of the latent space. As we traverse, we sample from the decoder distribution, thereby producing two new motion sequences, which appear in Fig.~\ref{fig:mocap_interp}. As can be seen, the straight line traverses uncharted territory of the latent space and end up creating an implausible motion. This is in contrast to the shortest path, that consistently generate meaningful poses.

\begin{SCfigure}
    \centering
    \includegraphics[width=0.75\textwidth]{imgs/mocap_pieces.pdf}
    \caption{Geodesic (top, green) versus linear (bottom, red) interpolation in a motion capture model. The geodesic follows the trend of the data manifold, while the linear path traverses regions with no data support.}
    \label{fig:mocap_interp}
\end{SCfigure}

\subsection{Numerical approximation of the pullback metric}
\label{sec:experiment_numerical_approx_metric}

Proposition~\ref{prop:approximating_metric} provide an approximation to the metric and we test its accuracy as per Eq.~\ref{eq:KL_approx_inner_product}. We discretize the latent space for the just-described von Mises-Fisher decoder and, for each $z$ in this grid, we both approximate $\b{M}(\vec{z})$ and compute the expected value of $\lVert\mathrm{KL}(p(\vec{x}|\vec{z}), p(\vec{x}|\vec{z}+\delta \vec{z})) - \frac{1}{2}\delta\vec{z}^\top \b{M}(\vec{z})\delta\vec{z}\rVert$ for several samples of $\delta\vec{z}$, uniformly distributed around the circle of radius $\varepsilon = 0.1$. Notice that we do not have a ground truth to compare against, and that this error will always be off by $o(\delta \vec{z}^2)$.
Figure~\ref{fig:approx_error_and_land} shows the latent space plotted against the average error. We can see that the approximate metric is well-estimated both within and outside the support of the data. The error, however, grows at the boundaries of the support, where the distribution is changing from a concentrated von Mises-Fisher to a uniform distribution.
It is worth mentioning that we observe some approximated metrics have negative determinant, showing that our numerical approximations are imprecise at the boundary. These results warrant further research on more stable ways of approximating pulled back metrics under our proposed approach.

\begin{SCfigure}
    \includegraphics[width=4cm]{imgs/approx_error.pdf}

    \includegraphics[width=4cm]{imgs/land_density_2_image.pdf}
    \caption{\emph{Left}: Average error of the approximated metric in the von Mises-Fisher latent space. Darker colors indicate lower error (less than $\varepsilon^2$), while higher values are clear. \emph{Right}: The LAND density well-adapts to the nonlinear structure of the latent representations due to the shortest paths behavior.}
    \label{fig:approx_error_and_land}
\end{SCfigure}

\subsection{Statistical models on Riemannian manifolds}
\label{sec:LAND_experiment}
To demonstrate the usefulness of the approximated metrics, we fit a distribution to data in the latent space, which require normalization according to the measure induced by the metric. In particular, we fit a locally adaptive normal distribution (LAND) \citep{arvanitidis:nips:2016}, which extends the Gaussian distribution to learned manifolds. The probability density function is $\rho(\b{z}) = C(\bs{\mu},\bs{\Sigma}) \cdot \exp\left(-0.5\cdot \mathrm{Log}_{\bs{\mu}}(\vec{z})\T \bs{\Gamma}\mathrm{Log}_{\bs{\mu}}(\vec{z}) \right)$, where $\bs{\mu}\in\R^d$ is the mean, $\bs{\Gamma}\in\R_{\succ 0}^{d\times d}$ is the precision matrix and $C(\bs{\mu}, \bs{\Gamma})$ the normalization constant. The operator $\mathrm{Log}_{\bs{\mu}}(\vec{z})$ returns the scaled initial velocity $\mathbf{v}=\dot{\gamma}(0)\in\R^d$ of the shortest connecting path with $\gamma(1)=\vec{z}$ and $||\mathbf{v}||=\mathrm{Length}(\gamma)$. In Figure~\ref{fig:approx_error_and_land} we show the LAND density on the learned latent representations under the approximated Riemannian metric from Sec.~\ref{sec:experiment_numerical_approx_metric}. Since shortest paths follow the data, so does the density $\rho$.

\section{Related work}
    Besides stochastic decoder models there is rich literature on models with deterministic decoders. Autoencoders \citep{rumelhart:nature:1986} and generative adversarial networks \citep{goodfellow:neurips:2014} are two prominent examples. Our work is not applicable to this setting, and a series of papers have investigated such deterministic decoders \citep{shao2018riemannian, chen2018metrics, laine2018feature}. As demonstrated in Sec.~\ref{sec:compare_metrics_with_gaussians_generators_uq} stochasticity is essential to shape the latent space according to the data manifold. \citet{hauberg:only:2018} argues model uncertainty plays a role much akin to topology in classic geometry, in that it, practically, allow us to deviate from the Euclidean topology of the latent space.

    Our constructions rely on information geometry, and Fisher-Rao metrics in particular \citep{nielsen:2020}. Our work is well within the spirit of information geometry, but does not represent typical usage of this theory. The most prominent examples of information geometry is within optimization with \emph{natural gradients} \citep{martens2014new, martens2015optimizing} Markov Chain Monte Carlo methods \citep{girolami:2011} and hypothesis testing \citep{nielsen:2020}. Information geometry provides a rich family of alternative divergences over the here-applied KL-divergence. We did not investigate their usage in our context.

    To make use of the here-developed tools, we may lean on techniques for statistics on manifolds. These provide generalizations of a long list of classic statistical algorithms \citep{zhang2013probabilistic, hauberg:tpami:princurve, fletcher2011geodesic}. We refer the reader to \citet{pennec:jmiv:2006} for a gentle introduction to this line of research.

\section{Conclusion and discussion}
\label{sec:conclusion}
We have proposed a new approach for getting a well-defined and useful geometry in the latent space of generative models with stochastic decoders. The theory is easy to apply and readily generalize to a large family of decoder distributions. The latent geometry gives access to a series of operations on latent variables that are invariant to reparametrizations of the latent space, and therefore are not subject to a large class of identifiability issues. Such operational representations have already shown great value in applications ranging from biology \citep{detlefsens:proteins:2020} to robotics \citep{scannellTrajectory2021}.
We have here focused on the Fisher-Rao metric, but other geometries over distributions may apply equally well, e.g.\@ the Wasserstein geometry may be interesting to explore.

\paragraph{Limitations.} The largest practical hurdle with the proposed methodology, is that it only works well for decoders with well-calibrated uncertainties. That is, the decoder should yield high entropy in regions of little training data to ensure that geodesics follow the trend of the data. This constraint is shared with existing approaches \citep{arvanitidis:iclr:2018}. Some heuristics exists \citep{detlefsen:2019:reliable}, but principled approaches are currently lacking.

\paragraph{Negative societal impact.} The root purpose of this work is to build a rich framework for interacting with latent variables. As latent variable models are sadly being used to generate misleading content, e.g.\@ deepfakes, then our tools can regretfully also be used for such purposes. We are, however, optimistic that the positive uses outweigh the negative.

\begin{ack}

This work was supported in part by a research grant (15334) from VILLUM FONDEN. This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement n\textsuperscript{o} 757360).

\end{ack}

\section{Fisher-Rao metric: definition, proof and discussion}

\subsection{Definitions} \label{appendix:FR:definitions}
\begin{definition}
A statistical manifold is a smooth manifold equipped with a Riemannian metric tensor, whose points are probability distributions.
\end{definition}
\begin{definition}
The Fisher-Rao metric, also called Fisher information matrix, is a symmetric and semi-positive definite matrix that is defined on a statistical manifold as:
\[I:\mathcal{H}\to\mathcal{R}^{|\mathcal{H}|\times|\mathcal{H}|} :\eta\to\int p(x|\eta) \left[\nabla\ln p(x|\eta) \nabla\ln p(x|\eta)^{\top}\right] dx
\]
The inner product between two distributions on the statistical manifold $\mathcal{H}$ is then defined as: \[g_I:\mathcal{H}\times\mathcal{H}\to\mathcal{R}:(\eta_1, \eta_2)\to \eta_1^{\top} I(\eta) \eta_{2} \overset{\Delta}{=} \langle\eta_1,\eta_2\rangle_{I} \]
\end{definition}

The Fisher-Rao metric is positive definite only if it is non-singular, and then, defines a Riemannian metric. In this paper, we assume that the observation $x\in\mathcal{X}$ is a random variable following a probability distribution $\mathcal{D}$ such that $x\sim\mathcal{D}(\eta)$, and any smooth changes of the parameter $\eta$ would alter the observation $x$. This way, the Fisher-Rao metric used in our paper is non-singular and the statistical manifold $\mathcal{H}$ is a Riemannian manifold.

\subsection{Geometrical properties}
A known result in Information Geometry \cite{amari:2016} is that the Fisher-Rao metric is the first order approximation of the KL-divergence, as recall in Proposition \ref{prop:FRapproxKL}. Using this fact, we can define the Fisher-Rao distance and energy in function of the KL-divergence, leading to Proposition \ref{prop:FRenergydistance}.

\begin{proposition}\label{prop:FRapproxKL}
The Fisher-Rao metric is the first order approximation of the KL-divergence between perturbed distributions:
\[\mathrm{KL}(p(x|\eta),p(x|\eta+\delta\eta)) = \frac{1}{2} \delta\eta^{\top} I(\eta) \delta\eta + o(\delta\eta^2).\]
with: $I(\eta) = p(x|\eta) \left[\nabla\ln p(x|\eta) \nabla\ln p(x|\eta)^{\top}\right] dx.$
\end{proposition}

\begin{proof}
Let's decompose $\ln p(x|\eta+\delta\eta)$ using the Taylor expansion:
\[\ln p(x|\eta+\delta\eta) = \ln p(x|\eta) + \nabla \ln p(x|\eta)^{\top} \delta\eta + \frac{1}{2} \delta\eta^{\top} \text{Hess} \left[\ln p(x|\eta)\right] \delta\eta + o(\delta\eta^2),\]
with: $\text{Hess} \left[\ln p(x|\eta)\right] = \frac{\text{Hess} \{p(x|\eta)\}}{p(x|\eta)} - \nabla\ln p(x|\eta) \nabla\ln p(x|\eta)^{\top}$ and $\nabla \ln p(x|\eta) = \frac{\nabla p(x|\eta)}{p(x|\eta)}$.

Also: $\int \nabla p(x|\eta) dx = \nabla\int p(x|\eta) dx = 0$ and $\int \text{Hess}[p(x|\eta)] dx = \text{Hess}[\int p(x|\eta) dx] = 0.$

Replacing all those expressions to the first equation finally gives:
\begin{align*}
     \mathrm{KL}(p(x|\eta),p(x|\eta+\delta\eta)) &= \int p(x|\eta) \ln p(x|\eta) dx - \int p(x|\eta) \ln p(x|\eta+\delta\eta) dx \\
    {} & = - \int p(x|\eta) \left( \nabla \ln p(x|\eta)^{\top} \delta\eta + \frac{1}{2} \delta\eta^{\top} \text{Hess} [\ln p(x|\eta)] \delta\eta + o(\delta\eta^2) \right) dx \\
    {}& = \frac{1}{2} \delta\eta^{\top} \left[ \int p(x|\eta) \left[\nabla\ln p(x|\eta) \nabla\ln p(x|\eta)^{\top}\right] dx \right] \delta\eta + o(\delta\eta^2).
\end{align*}

\end{proof}

\begin{definition}
We consider a curve $\gamma(t)$ and its derivative $\dot{\gamma}(t)$ on the statistical manifold such that, $\forall t\in[0,1], \gamma(t)=\eta_t\in\mathcal{H}$. The manifold is equipped with the Fisher-Rao metric. The length $\mathcal{L}$ and the energy $\mathcal{E}$ functionals are defined with respect to the metric $I(\eta)$:
\[ \mathcal{L}(\gamma) = \int_{0}^{1} \sqrt{\dot{\gamma}(t)^{\top} I(\eta) \dot{\gamma}(t)} dt \quad\text{and}\quad \mathcal{E}(\gamma) = \int_{0}^{1} \dot{\gamma}(t)^{\top} I(\eta) \dot{\gamma}(t) dt.\]

Locally length-minimising curves between two connecting points are called Geodesics.
\end{definition}

\begin{proposition} \label{prop:FRenergydistance}
The KL-divergence between two close elements of the curve $\gamma$ is defined as: $\mathrm{KL}(p_{t},p_{t+\delta t})=\mathrm{KL}(p(x|\gamma(t)),p(x|\gamma(t+\delta t)))$.
The length $\mathcal{L}$ and the energy $\mathcal{E}$ functionals can be approximated with respect to this KL-divergence:
\[\mathcal{L}(\gamma) \approx\sqrt{2\textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})}  \quad\text{and}\quad \mathcal{E}(\gamma) \approx \frac{\delta t}{2} \textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})\]

\end{proposition}

\begin{proof}
On the statistical manifold, we have $\gamma(t+\delta t)=\gamma(t)+\delta t\dot{\gamma}(t)$. The KL-divergence between perturbed distributions can be defined as: $\mathrm{KL}(p_{t},p_{t+\delta t}) = \mathrm{KL}(p(x|\gamma(t)),p(x|\gamma(t+\delta t)))= \mathrm{KL}(p(x|\eta_t),p(x|\eta_t+\delta\eta_t)),$ with $\eta_t=\gamma(t)$ and $\delta\eta_t=\delta t\ \dot{\gamma}(t)$.
Then, we obtain:
\[\mathrm{KL}(p_{t},p_{t+\delta t}) = \frac{1}{2} \delta t^2\ \dot{\gamma}(t)^{\top}I(\eta_t)\dot{\gamma}(t) + o(\delta t^2).\]

The length and energy terms appear in the following equations:
\begin{align*}
    \int_0^{1}\mathrm{KL}(p_{t},p_{t+\delta t}) dt &= \frac{\delta t^2}{2} \int_0^{1} \dot{\gamma}(t)^{\top}I(\eta_t)\dot{\gamma}(t) dt + o(\delta t^2) = \frac{\delta t^2}{2}\ \mathcal{E}(\gamma)+ o(\delta t^2),\\
    \int_0^{1}\sqrt{\mathrm{KL}(p_{t},p_{t+\delta t})} dt &= \frac{\delta t}{\sqrt{2}} \int_0^{1} \sqrt{\dot{\gamma}(t)^{\top}I(\eta_t)\dot{\gamma}(t)} dt + o(\delta t^2) = \frac{\delta t}{\sqrt{2}}\ \mathcal{L}(\gamma)+ o(\delta t^2).\\
\end{align*}

If we want approximate any continuous function $f$ with a discrete sequence, by partitioning it in $T$ small segments, such that: $\delta t \approx \frac{1}{T}$, we have: $\int_0^{1} f(t) dt \approx \sum_{t=1}^{T} f(t) \delta t$, which in our case gives:

\[\mathcal{L}(\gamma) \approx \sqrt{2\textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})} \quad \text{and} \quad \mathcal{E}(\gamma) \approx \frac{2}{\delta t} \textstyle \sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t}). \]

\end{proof}

\subsection{Fisher-Rao metric for different distributions} \label{appendix:FR:distributions}

\begin{table}[h]
\begin{center}
\begin{tabular}{llll}
\hline
Distributions&Probability density functions&Parameters& Fisher-Rao matrix\\[0.2cm] \hline
Normal& $\frac{1}{\sqrt{2\pi\sigma^2}} \exp{-\frac{(x-\mu)^2}{2\sigma^2}}$ & $\mu, \sigma^2$& $I_{\mathcal{N}}(\mu,\sigma^2)$\\[0.3cm]
Binomial& $\theta^{x}(1-\theta)^{1-x}$& $\theta$& $I_{\mathcal{B}}(\theta)$\\[0.3cm]
Categorical& $\prod_{k=1}^{K}\theta_{k}^{x_k}$& $\theta_1, \dots, \theta_K$ & $I_{\mathcal{C}}(\theta_1, \dots, \theta_K)$ \\[0.3cm]
Gamma& $\frac{\beta^{\alpha} x^{\alpha-1}e^{-\beta x}}{\Gamma(\alpha)}$    & $\alpha, \beta$& $I_{\mathcal{G}}(\alpha,\beta)$\\[0.3cm]
Von Mises-Fisher, for $\mathbb{S}^2$ & $\frac{\kappa}{4\pi\sinh{\kappa}}\exp{(\kappa \mu^{\top}x)}$& $\kappa, \mu$& $I_{\mathcal{S}}(\kappa,\mu)$\\[0.3cm]
Beta & $\frac{\Gamma(\alpha)\Gamma(\beta)}{\Gamma(\alpha + \beta)} x^{\alpha - 1}(1-x)^{\beta -1}$& $\alpha, \beta$ & $I_{\mathcal{B}}(\alpha,\beta)$ \\[0.3cm] \hline
\end{tabular}
\end{center}
\caption{List of distributions}
\label{tab:distributions}
\end{table}

With the notations of Table \ref{tab:distributions}, the Fisher-Rao matrices of the the univariate Normal, Binomial and Categorical are:
\[
I_{\mathcal{N}}(\mu,\sigma^2) = \begin{pmatrix} \frac{1}{\sigma^2} & 0 \\ 0 & \frac{1}{2\sigma^4}\\\end{pmatrix}, \quad
I_{\mathcal{B}}(\theta) = \frac{1}{\eta(1-\eta)}, \quad
I_{\mathcal{C}}(\theta_1, \dots, \theta_K) = \begin{pmatrix}
  \theta_1 & 0 &\dots & 0 \\
  0 & \theta_2 & \dots & 0 \\
  \vdots & \vdots & \ddots & \vdots \\
  0 & 0 & \dots & \theta_K \\
\end{pmatrix}
\]
In addition, the Fisher-Rao matrices of the Gamma, Von Mises-Fisher and the Beta distributions are:
\[
I_{\mathcal{G}}(\alpha,\beta) = \begin{pmatrix} \frac{\alpha}{\beta^2} & -\frac{1}{\beta} \\ -\frac{1}{\beta} & \Psi_1(\alpha) \\\end{pmatrix}, \quad
I_{\mathcal{S}}(\kappa,\mu) = \begin{pmatrix} \kappa K(\kappa)(\mathds{1}-3\mu\mu^{\top})+\kappa^2\mu\mu^{\top} & (\kappa K(\kappa)^2-\frac{2}{k}K(\kappa)+1)\mu \\ (\kappa K(\kappa)^2-\frac{2}{k}K(\kappa)+1)\mu^{\top} & 3K(\kappa)^2-\frac{2}{\kappa}K(\kappa)+1 \\\end{pmatrix} ,
 \]
 \[I_{\mathcal{B}}(\alpha,\beta) = \begin{pmatrix} \Psi_1(\alpha)-\Psi_1(\alpha+\beta) & -\Psi_1(\alpha+\beta) \\ -\Psi_1(\alpha+\beta) & \Psi_1(\beta) - \Psi_1(\alpha+\beta) \\\end{pmatrix} \]
 \\
with $\Psi_1(\alpha) = \frac{\partial^2 \ln\Gamma(\alpha)}{\partial \alpha}$ the trigamma function, and $K(\kappa) = \coth{\kappa}-\frac{1}{\kappa}$.

\begin{proof}
The univariate Normal, Binomial and Categorical have already been studied by \cite{tomczak:2012}, and the Beta distribution by \cite{brigant:2019}. We will then focus our proof on the Gamma and the Von-Mises Fisher distributions.

In order to bypass unnecessary details, we will use the following notations, we redefine the Fisher-Rao as: $I(\eta)=\mathbb{E}_{x}[g(\eta, x)g(\eta, x)^{\top}]$, with $g(\eta,x) = \nabla_{\eta} \ln p(x|\eta)$ the Fisher score. We call $G = g(\eta, x)g(\eta, x)^{\top}$, and $G_{ij}$ the matrix elements.

\textbf{Gamma distribution:}

We have $p(x|\alpha, \beta) = \Gamma(\alpha)^{-1} \beta^{\alpha} x^{\alpha-1}e^{-\beta x}$, which leads to:
\begin{align*}
    \ln p(x|\alpha, \beta) &= - \ln \Gamma(\alpha) + \alpha \ln \beta + (\alpha - 1) \ln x - \beta x \\
    \frac{\partial \ln p}{\partial \alpha} & = - \Psi(\alpha) + \ln \beta + \ln x \\
    \frac{\partial \ln p}{\partial \beta} & = \frac{\alpha}{\beta} - x
\end{align*}
Then:
\begin{align*}
    G_{11} &= \left(\frac{\partial \ln p}{\partial \alpha}\right)^2 = (\Psi_{0}(\alpha) + \ln \beta)^2 + 2 (\Psi(\alpha) + \ln \beta ) ln x + \ln^2 x \\
    G_{22} &= \left(\frac{\partial \ln p}{\partial \beta}\right)^2  = \left(\frac{\alpha}{\beta}\right)^2 - 2 \frac{\alpha}{\beta} x + x^2  \\
    G_{12} = G_{21} &= \frac{\partial \ln p}{\partial \alpha} \cdot \frac{\partial \ln p}{\partial \beta}  = \left(\Psi(\alpha) + \ln \beta\right) \left(\frac{\alpha}{\beta} - x\right) + \frac{\alpha}{\beta} \ln x - x \ln x
\end{align*}

We know that $\mathbb{E}[x] = \frac{\alpha}{\beta}$. We can compute, using your favorite symbolic computation software, the following moments:
\begin{align*}
    \mathbb{E}[\ln x] &= -\ln \beta + \Psi(\alpha) \\
    \mathbb{E}[x \ln x] &= \frac{\alpha}{\beta} \left(\Psi(\alpha + 1) - \ln \beta\right) \\
    \mathbb{E}[\ln^2 x] &= \left(\ln \beta - \Psi(\alpha)\right)^2 + \Psi_1(\alpha)
\end{align*}
Replacing the moments for the following equations: $\mathbb{E}[G_{11}]$, $\mathbb{E}[G_{22}]$ and $\mathbb{E}[G_{12}]$ will finally give the Fisher-Rao matrix.\\

\textbf{Von Mises Fisher distribution, for $\mathrm{S}^2$:}

We have $p(x|\mu, \kappa) = C_{3}(\kappa) \exp (\kappa \mu^{\top} x)$, with $C_{3}(\kappa) = \kappa (4 \pi \sinh \kappa)^{-1}$. Here, $\mu$ is a 3-dimensional vector with $\lVert \mu \rVert = 1$.
\begin{align*}
    \ln p(x|\mu, \kappa) &= \ln \kappa - \ln 4\pi - \ln \sinh(\kappa) + \kappa \mu^{\top}x \\
    \nabla_{\mu} \ln p & = \kappa x  \\
    \frac{\partial \ln p}{\partial \kappa} & =  \kappa^{-1} - \text{coth}(\kappa) + \mu^{\top} x.
\end{align*}
Here, the Fisher-Rao matrix $I_{\mathcal{S}}$ will be composed of block matrices, such that: $I_{\mathcal{S}} = \mathbb{E}[G]$, with $G_{11}$ a $3\times 3$-matrix, $G_{22}$ a scalar, and $G_{12}=G_{21}^{\top}$ a 3-dimensional vector.
\begin{align*}
    G_{11} &= \nabla_{\mu}\ln p \nabla_{\mu}\ln p^{\top} = \kappa^2 x x^{\top} \\
    G_{22} &= \left(\frac{\partial \ln p}{\partial \kappa}\right)^2 = K(\kappa)^2 + 2 K(\kappa) \mu^{\top} x + (\mu^{\top}x)^2  \\
    G_{12} = G_{21}^{\top} &= \frac{\partial \ln p}{\partial \kappa} \cdot \nabla_{\mu}\ln p = \left(K(\kappa) + \mu^{\top} x \right) \kappa x,
\end{align*}
with $K(\kappa) = \text{coth}(\kappa) - \frac{1}{\kappa}$.

We know from \cite{hillen:2016} that the mean and variance of the Von Mises Fisher distribution in the 3-dimensional case is: $\mathbb{E}[x] = K(\kappa) \mu$ and $\text{Var}[x] = \frac{1}{\kappa} K(\kappa) \mathds{1} + (1- \frac{\text{coth}(\kappa)}{\kappa} + \frac{2}{\kappa^2} - \text{coth}^2(\kappa)) \mu \mu^{\top}$. We can then deduce the following meaningful moments:
\begin{align*}
    \mathbb{E}[xx^{\top}] & = \text{Var}[x] + \mathbb{E}[x]\mathbb{E}[x]^{\top} = \left(1-\frac{3}{\kappa} K(\kappa)\right) \mu\mu^{\top} + \frac{1}{\kappa} K(\kappa) \mathds{1}, \\
    \mathbb{E}[\mu^{\top}x] & = \mu^{\top} \mathbb{E}[x] = K(\kappa) \mu^{\top} \mu = K(\kappa) \\
    \mathbb{E}[(\mu^{\top}x)^2] & = \mu^{\top} \text{Var}[x] \mu + \mathbb{E}[\mu^{\top}x]^2 = 1-\frac{2}{\kappa}K(\kappa), \\
    \mathbb{E}[\mu^{\top}x x] &= \mathbb{E}[\mu x x^{\top}] = \mathbb{E}[xx^{\top}] \mu = \left(1 - \frac{3}{\kappa} K(\kappa)\right) \mu + \frac{1}{\kappa} K(\kappa) \mu.
\end{align*}
Replacing those moments in the following expressions: $\mathbb{E}[G_{11}]$, $\mathbb{E}[G_{22}]$, $\mathbb{E}[G_{12}]$ directly gives the Fisher-Rao metric.
\end{proof}

\subsection{Pullback metric for the VAE case}

We call $f$ the non linear function (neural networks) that maps the variables from the latent space $\mathcal{Z}$ to the parameter space $\mathcal{H}$, such that: $f(z)=\eta$, with $z\in\mathcal{Z}$ and $\eta\in\mathcal{H}$. Furthermore, the data $x \in \mathcal{X}$ is reconstructed such that it follows a specific distribution: $x \sim \mathcal{D}(\eta)$, with $\mathcal{D}$ being for instance a Bernoulli or Gaussian distribution. The parameter space $\mathcal{H}$ is a statistical manifold equipped with Fisher-Rao metric: $I(\eta) \overset{\Delta}{=} \int p(x|\eta) \left[\nabla\ln p(x|\eta) \nabla\ln p(x|\eta)^{\top}\right] dx$. We define $J_f$ is the Jacobian of $f$.

\begin{proposition}
The latent space $\mathcal{Z}$ is equipped with the pullback metric:
\[ M(z) \overset{\Delta}{=} J_{f}^{\top} I(\eta) J_{f}. \]
\end{proposition}

\begin{proof}
The parameter space is a statistical manifold equipped with the Fisher-Rao metric $I(\eta)$, thus the scalar product between two parameters $\eta_1, \eta_2 \in \mathcal{H}$ is: $\langle \eta_1, \eta_2 \rangle_{I(\eta)} =  \eta_1^{\top}I(\eta) \eta_2$. Furthermore, we have $f:\mathcal{Z}\to\mathcal{H}$ a differentiable map. With $z_1, z_2 \in\mathcal{Z}$, we have: $\langle \eta_1, \eta_2 \rangle_{I(\eta)} =  \langle J_f z_1, J_f z_2 \rangle_{I(\eta)} = z_1^{\top} (J_{f}^{\top} I(\eta) J_{f}) z_2$.
\end{proof}

\begin{proposition}
Our pullback metric $M(z)$ is actually equal to the Fisher-Rao metric obtained over the parameter space $\Z$:
\[M(z) = I(z) \overset{\Delta}{=} \int p(x|z) \left[\nabla\ln p(x|z) \nabla\ln p(x|z)^{\top}\right] dx \]
\end{proposition}

\begin{proof}
We will show that $I(z) = J_{f}^{\top} I(\eta) J_{f}$.
Let's consider the definition of the Fisher-Rao metric in $\Z$ :
\begin{align}
    I(\b{z})
    &= \int_\X \nabla_{\b{z}} \log p(\b{x}~|~\b{z}) \cdot \nabla_{\b{z}} \log p(\b{x}~|~\b{z})\T p(\b{x}~|~\b{z})\dif \b{x}\\
    &= \int_\X \b{J}_f(\b{z})\T \nabla_{\eta} \log p(\b{x}|\eta)\nabla_{\eta} \log p(\b{x}|\eta)\T \b{J}_f(\b{z}) p(\b{x}|\eta)\dif\b{x} \\
    &= \b{J}_f(\b{z})\T \left[ \int_\X \nabla_{\eta} \log p(\b{x}|\eta)\nabla_{\eta} \log p(\b{x}|\eta)\T  p(\b{x}|\eta)\dif\b{x} \right]\b{J}_f(\b{z})\\
    &= \b{J}_f(\b{z})\T I(\eta) \b{J}_f(\b{z}) \nonumber
\end{align}
where we use the fact that $\eta=f(\b{z})$ so the $\nabla_{\b{z}}\log p(\b{x}|\b{f}(\b{z})) = \b{J}_f(\b{z})\T \cdot \nabla_{\eta} \log p(\b{x}|\eta)$

We could probably prove the same as:
\begin{align}
    &\inner{d\eta}{I(\eta) d\eta} = \inner{\b{J}_f(\b{z})d\b{z}}{I(f(\b{z}))\b{J}_f(\b{z})d\b{z}} \\
    &= \inner{\b{J}_f(\b{z})d\b{z}}{\int \nabla_\eta \log p(\b{x}|\eta) \nabla_\eta \log p(\b{x}|\eta) \T p(\b{x}|\eta)\dif\b{x} ~\b{J}_f(\b{z})d\b{z}}\\
    &=\inner{d\b{z}}{\int \b{J}_f(\b{z})\T \cdot \nabla_\eta \log p(\b{x}|\eta) ~\nabla_\eta \log p(\b{x}|\eta)\T \cdot \b{J}_f(\b{z}) p(\b{x}|\eta)\dif\b{x} ~d\b{z}}\\
    &=\inner{d\b{z}}{\int \nabla_{\b{z}}\log p(\b{x}|\b{z}) \nabla_{\b{z}}\log p(\b{x}|\b{z})\T p(\b{x}|\b{z}) \dif\b{x}~d\b{z}} = \inner{d\b{z}}{I(\b{z}) d\b{z}}
\end{align}
\end{proof}

In sections \ref{appendix:FR:definitions} and \ref{appendix:FR:distributions}, we have seen how to derive a close-form expression of the Fisher-Rao metric for a one-dimensional observation $x$ that follows a specific distribution. In practice, $x\in\mathcal{X}\subset\mathbb{R}^D$ is a multi-dimensional variable where each dimension represents, for instance, a pixel when working with images or a feature when working with tabular data. Each feature, $x_i$ with $i=1\cdots D$, is obtained for a specific set of parameters $\{\eta_i\}$. We assume that the features follow the same distribution $\mathcal{D}$, such that: $x_i \sim \mathcal{D}(\eta_i)$, and $p(x|\eta) = \prod_{i=1}^{D}p(x_i|\eta_i)$.

\begin{proposition} \label{prop:blockmatrix}
The Fisher-Rao metric $I(\eta)$ is a block matrix where the diagonal terms are the Fisher-Rao matrices $I_i$ obtained for each data feature $x_i$:
\[
I(\eta) = \begin{pmatrix}
  I_1 & 0 &\dots & 0 \\
  0 & I_2 & \dots & 0 \\
  \vdots & \vdots & \ddots & \vdots \\
  0 & 0 & \dots & I_D \\
\end{pmatrix}
\]
\end{proposition}
\begin{proof}
We have $x_i\sim\mathcal{D}(\eta_i)$ and we define $I_i=\int p(x_i|\eta_i) \left[\nabla\ln p(x_i|\eta_i) \nabla\ln p(x_i|\eta_i)^{\top}\right] dx_i$.
Furthermore, we assumed: $p(x|\eta)=\prod_{i=1}^{D}p(x_i|\eta_i)$.
We then have: $\ln{p(x|\eta)} = \sum_{i=1}^{D} \ln{p(x_i|\eta_i)}$, and the Fisher score: $\nabla_{\eta}\ln{p(x|\eta)} = \nabla_{\eta} \sum_{i=1}^{D} \ln{p(x_i|\eta_i)} = \left[\nabla\eta_1 \ln{p(x_1|\eta_1)}, \dots, \nabla\eta_D \ln{p(x_1|\eta_D)}\right]^{\top}$. \\

\noindent
The matrix $I(\eta)$ is thus a $D \times D$ block matrix, where the $(i,j)$-block element is:
\[I_{ij}=\int p(x_i|\eta_i) \left[\nabla\ln p(x_i|\eta_i) \nabla\ln p(x_j|\eta_j)^{\top}\right] dx_i.\]
Let's note that:
\[\int p(x_i|\eta_i) \nabla\ln p(x_i|\eta_i) dx_i = \int p(x_i|\eta_i) \frac{\nabla p(x_i|\eta_i)}{p(x_i|\eta_i)}dx_i = \nabla \int p(x_i|\eta_i) dx_i = 0.\]
When $i=j$, we have $I_{ii}=I_{i}$, with $I_i$ being the Fisher-Rao metric obtained for: $x_i\sim\mathcal{D}(\eta_i)$. \\
\noindent
When $i\neq j$, we have: $I_{ij}=\nabla\ln p(x_j|\eta_j)^{\top} \int p(x_i|\eta_i) \nabla\ln p(x_i|\eta_i) dx_i = 0$.
\end{proof}

Then, for example, if we are dealing with binary images, and make the assumption that each pixel $x_i$ follows a Bernoulli distribution: $p(x_i|\eta_i) = \eta^{x_i}(1-\eta_i)^{1-x_i}$, then according to Section \ref{appendix:FR:distributions} and Proposition \ref{prop:blockmatrix}, the Fisher-Rao matrix that endows the parameter space $\mathcal{H}$ is:
\[
I(\eta) = \begin{pmatrix}
              \frac{1}{\eta_1(1-\eta_1)} & 0 &\dots & 0 \\
              0 & \frac{1}{\eta_2(1-\eta_2)} & \dots & 0 \\
              \vdots & \vdots & \ddots & \vdots \\
              0 & 0 & \dots & \frac{1}{\eta_D(1-\eta_D)} \\
            \end{pmatrix}.
\]

\section{Uncertainty Quantification and Regularization}
\label{appendix:sec:UQ}

As discussed in the main text, we carefully design our mappings from latent space to parameter space such that they model the training codes according to the learned decoders, and extrapolate to uncertainty outside the support of the data. This, we refer to as \textbf{uncertainty regularization}. In this section we explain it in detail. The core idea of this uncertainty regularization is imposing a ``slider'' that forces the distribution $p(x|z)$ to change when $z$ is far from the training latent codes. For this, we use a combination of KMeans and the sigmoid activation function.

We start by encoding our training data, arriving at a set of latent codes $\{z_i\}_{i=1}^{N}\subseteq\mathcal{Z}$. We then train KMeans($k$) on these latent codes (where $k$ is a hyperparameter that we tweak manually), arriving at $k$ cluster centers $\{c_j\}_{j=1}^k$. These cluster centers serve as a proxy for "closeness" to the data: we know that a latent code $z\in\mathcal{Z}$ is near the support if $D(z) := \min_j\left\{\|z - c_j\|^2\right\}$ is close to 0.

The next step in our regularization process is to reweight our decoded distributions such that we decode to high uncertainty when $D(z)$ is large, and we decode to our learned distributions when $D(z) \approx 0$. This mapping from $[0, \infty)\to (0, 1)$ can be constructed using a modified sigmoid function \cite{detlefsens:proteins:2020,detlefsen:2019:reliable}, consider indeed
\begin{equation}
\label{appendix:eq:translated_sigmoid}
\tilde{\sigma}_{\beta}(d) = \text{Sigmoid}\left(\frac{d - c\cdot\text{Softplus}(\beta)}{\text{Softplus}(\beta)}\right),
\end{equation}
where $\beta\in\mathbb{R}$ is another hyperparameter that we manually tweak, and $c\approx 7$.

With this translated sigmoid, we have that $\tilde{\sigma}_\beta(D(z))$ is close to 0 when $z$ is close to the support of the data (i.e. close to the cluster centers), and it converges to 1 when $D(z)\to\infty$.  $\tilde{\sigma}_\beta(D(z))$ serves, then, as a slider that indicates closeness to the training codes. This reweighting takes the following form:
\begin{equation}
\label{appendix:eq:reweight}
\text{reweight}(z) = (1 - \tilde{\sigma}_\beta(D(z)))h(z) + \tilde{\sigma}_\beta(D(z))\,\text{extrapolate}(z),
\end{equation}
where  $h(z) = \eta(z)\in\mathcal{H}$ represents our learned networks in parameter space,  and $\text{extrapolate}(z)$ returns the parameters of the distribution that maximize uncertainty (e.g. $\sigma\to\infty$ in the case of an isotropic Gaussian, $p\to 1/2$ in the case of a Bernoulli, and $\kappa\to 0$ in the case of the von Mises-Fisher).

For the particular case of the experiment in which we pull back the Fisher-Rao metric from the parameter space of several distributions (see \ref{sec:toy_latent_space_experiment}), Table \ref{appendix:tab:extrapolation_mechanisms} provides the exact extrapolation mechanisms and implementations of $h(z)$.

\begin{table}[]
    \centering
    \resizebox{\textwidth}{!}{
    \begin{tabular}{lll}
        Distribution & {$h\colon\mathcal{Z}_{\text{toy}}\to\mathcal{H}$} & Extrapolation mechanism \\
        \toprule
        Normal & $\mu(z)=10\cdot f_3(z)$, ~ $\sigma(z) = 10\cdot \text{Softplus}(f_3(z))$ & $\sigma(z)\to\infty$ \\[0.1cm]
            Bernoulli & {$p(z) = \text{Sigmoid}(f_{15}(z))$} & $p(z) = 1/2$ \\[0.1cm]
            Beta & $\alpha(z)=10\cdot \text{Softplus}(f_3(z))$, ~  $\beta(z)=10\cdot \text{Softplus}(f_3(z))$ & $(\alpha(z), \beta(z)) = (1, 1)$ \\[0.1cm]
            Dirichlet & $\alpha(z)=\text{Softplus}(f_3(z))$  & $\alpha(z) = 1$ \\[0.1cm]
            Exponential & $\lambda(z)=\text{Softplus}(f_3(z))$  & $\lambda(z) \to 0$ \\[0.2cm]
    \end{tabular}
    }
    \caption{This table shows the implementations of the decode and extrapolate functions in Eq. (\ref{appendix:eq:reweight}) for all the distributions studied in our second experiment (see Sec. \ref{sec:toy_latent_space_experiment}). Here we represent a randomly initialized neural network with $f_i$, where $i$ represents the size of the co-domain. For example, in the case of the Dirichlet distribution, we use a randomly initialized neural network to compute the parameters $\alpha$ of the distribution and, since these have to be positive, we pass the output of this network through a Softplus activation; moreover, since the Dirichlet distribution is approximately uniform when all its parameters equal 1, our extrapolation mechanism consists of replacing the output of the network with a constant vector of ones.}
    \label{appendix:tab:extrapolation_mechanisms}
\end{table}

\section{Categorical}
\label{appendix:categorical_details}

Let a random variable $\vec{x}\in\R^D$ that follows a generalized Bernoulli likelihood $p(\vec{x}|\eta)$, so the vector $\vec{x}\in\R^D$ is of the form $\b{x} = (0, 0, \cdots, 1, \cdots)$ with $\sum_i x_i = 1$. The parameters $\eta\in\R^D$ are given as $\eta = h(\vec{z})$, with $\eta_i \geq 0~\forall i$ and $\sum_i \eta_i = 1$ so we know that the parameters lie on the unit simplex. Actually, they represent the probability the corresponding dimension to be 1 on a random draw. Also, the $p(\vec{x}|\vec{z}) = \eta_1^{[x_1]}\cdots \eta_D^{[x_D]}$, where $[x_i] = 1$ if $x_i=1$ else $[x_i] = 0$. The $\log p(\vec{x}|\eta) = \sum_i [x_i] \log(\eta_i)$ and $\nabla_\eta \log p(\vec{x}|\eta) = \left(\frac{[x_1]}{\eta_1}, \dots, \frac{[x_D]}{\eta_D}\right)$. Due to the outer product we have the following expectations
\begin{align}
    &\mathbb{E}_\b{x}\left[ \frac{[x_i]}{\eta_i} \frac{[x_j]}{\eta_j}\right] = 0, ~\quad \text{if}\quad i\neq j,\\
    &\mathbb{E}_\b{x}\left[ \left(\frac{[x_i]}{\eta_i}\right)^2\right] = \frac{1}{\eta_i}, ~\quad \text{if}\quad i=j,
\end{align}
because the $[x_i]$ and $[x_j]$ cannot be 1 on the same time, while the $\mathbb{E}_\b{x}[[x_i]^2] = \eta_i$. So the Fisher-Rao metric of $\HH$ is equal to $\b{I}_\HH(\eta) = \mathrm{diag}\left(\frac{1}{\eta_1},\dots,\frac{1}{\eta_2}\right)$. Note that the shortest paths between two distributions must be on the unit simplex, while respecting the geometry of the Fisher-Rao metric.

We can easily parametrize the unit simplex by $[\eta_1,\dots,\eta_{D-1}, \widetilde{\eta}_D]$ with
\begin{equation}
    \widetilde{\eta}_D(\eta_1,\dots, \eta_{D-1}) = 1 - \sum_{i=1}^{D-1} \eta_i.
\end{equation}
This allows to pullback the Fisher-Rao metric in the latent space $[\eta_1,\dots,\eta_{D-1}]$ as we have described in this paper. Intuitively, the $\vec{z} = [\eta_1,\dots,\eta_{D-1}]$ and the function $h$ is the parametrization of the simplex. Hence, we are able to compute the shortest path using the induced metric.

However, there is a simpler way to compute this path. We know that the element-wise square root of the parameters $\eta$ gives a point on the positive orthonant of the unit sphere as $y_i = \sqrt{\eta_i}~\Rightarrow~\sum_i y_i^2 = \sum_i \sqrt{\eta_i}^2 = 1$. We also know that the shortest path on a sphere is the great-circle. Therefore, the distance between two distributions parametrized by $\eta$ and $\eta'$ can be measured using the great-circle distance between their square roots as
\begin{equation}
    \mathrm{dist}(\eta, \eta') = \arccos\sqrt{\eta}\T\sqrt{\eta'}.
\end{equation}

In this way, we can approximate the energy of a curve $c(t)$ in the latent space as follows
\begin{align}
    \mathrm{Energy}[c]&\approx \sum_{n=1}^{N-1}\mathrm{dist}^2(h(c(\sfrac{n}{N})),h(c(\sfrac{n+1}{N}))) = \sum_{n=1}^{N-1}\arccos^2 \sqrt{h(c(\sfrac{n}{N}))}\T\sqrt{h(c(\sfrac{n+1}{N}))}\nonumber\\
    &=\sum_{n=1}^{N-1}\left(2 - 2\sqrt{h(c(\sfrac{n}{N}))}\T\sqrt{h(c(\sfrac{n+1}{N}))}\right),
\end{align}
where we used at the last step the small angle approximation $\cos\theta \approx 1 - \frac{\theta^2}{2} \Leftrightarrow \theta^2 \approx 2 - 2\cos\theta$.

Note that this formulation is suitable for our proposed method to compute shortest paths (see Section~\ref{sec:geodesic_computation}). This represents the conceptual strategy, while we use the KL divergence approximation result \eqref{eq:KL_approx_inner_product} in place of the great-circle distance. Intuitively, when the KL divergence has an analytic solution we end up to an analogous energy approximation. Even if the solution of the KL is intractable, we can still use our approach as long as we can estimate the KL using Monte Carlo and propagate the gradient through the samples using a re-parametrization scheme or score function estimator.
\section{Introduction}
\label{sec:intro}

Generative models such as \emph{variational autoencoders (VAEs)} \citep{kingma:iclr:2014, rezende:icml:2014} and \emph{generative adversarial networks (GANs)} \citep{goodfellow:neurips:2014} provide state-of-the-art density estimators for high dimensional data. The underlying assumption is that data $\vec{x} \in \X$ lie near a low-dimensional manifold $\M \subset \X$, which is parametrized through a low-dimensional \emph{latent representation} $\vec{z} \in \Z$. As data is finite and noisy, we only recover a probabilistic estimate of the true manifold, which, in VAEs, is represented through a decoder distribution $p(\vec{x} | \vec{z})$. Our target is the geometry of this \emph{random manifold}.

The geometry of the manifold has been shown to carry great value when systematically interacting with the latent representations, as it provides a stringent solution to the \emph{identifiability problem} that plagues latent variable models \cite{tosi:uai:2014, arvanitidis:iclr:2018, hauberg:only:2018}. For example, this geometry has allowed VAEs to discover latent evolutionary signals in proteins \cite{xxx}, provide efficient robot controls \cite{xxx, yyy, zzz}, improve latent clustering abilities \cite{yang:arxiv:2018, arvanitidis:iclr:2018} and more. The fundamental issue with these geometric approaches is that the studied manifold is inherently a stochastic object, but classic differential geometry only supports the study of \emph{deterministic} manifolds. To bridge the gap, \citet{eklund:arxiv:2019} have shown how VAEs with a Gaussian decoder family can be viewed as a random projection of a deterministic manifold, thereby making the classic theories applicable to the random manifold.

A key strength of VAEs is that they can model data from diverse modalities through the choice of decoder distribution $p(\vec{x} | \vec{z})$. For discrete data, we use categorical decoders, while for continuous data we may opt for a Gaussian, a Gamma or whichever distribution best suits the data. But for non-Gaussian decoders, there exists no useful approach for treating the associated random manifold as deterministic, which prevents us from systematically interacting with the latent representations without being subject to identifiability issues.

\textbf{In this paper}, we provide a general framework that allows us to interact with the geometry of almost any random manifold. The key, and simple idea is to reinterpret the decoder as spanning a deterministic manifold in the space of probability distributions $\mathcal{P}$, rather than a random manifold in the observation space (see Fig.~\ref{fig:teaser_img}). Calling on classical \emph{information geometry} \cite{amari:2016}, we show that the learned random manifold is a Riemannian submanifold of $\mathcal{P}$, and provide corresponding computational tools. The approach is applicable to any family of decoders for which the KL-divergence can be differentiated, allowing us to work with a wide array of models from a single codebase.\looseness=-1

\begin{SCfigure}
    \includegraphics[width=4cm]{imgs/concept_v1_mnist.pdf}
    \includegraphics[width=4cm]{imgs/concept_v1_probabilities.pdf}
    \caption{Traditionally (left), we view the learned manifold as a stochastic manifold in the observation space. We propose (right) to view the learned manifold as a deterministic manifold embedded in the space of decoder distributions, which has an information geometry.}
    \label{fig:teaser_img}
\end{SCfigure}

\section{The geometry of generative models}
\label{sec:riemannian_geometry}

\begin{wrapfigure}{r}{0.35\textwidth}
    \vspace{-25pt}
    \begin{center}
    \includegraphics[width=0.35\textwidth]{imgs/teaser_2.pdf}
  \end{center}
  \caption{An example of $\M$ with $\Z$ and shortest path.}
  \vspace{-25pt}
  \label{fig:riem_geom_example}
\end{wrapfigure}

  As a starting point, consider the deterministic generative model given by a prior $p(\vec{z})$ and a decoder $f: \Z \rightarrow \X$. The latent representation $\vec{z}$ of an observation $\vec{x}$ is generally not \emph{identifiable}, meaning that one can recover different latent representations that give rise to equally good density estimates. For example, let $g: \Z \rightarrow \Z$ be a smooth invertible function such that $\vec{z} \sim p(\vec{z}) \Leftrightarrow g(z) \sim p(\vec{z})$, then the latent representation $g(\vec{z})$ coupled with the decoder $f \circ g^{-1}$ gives the same density estimate as $\vec{z}$ coupled with $f$ \cite{hauberg:only:2018}. Practically speaking, the identifiability issue imply that it is improper to view the latent space $\Z$ as being Euclidean, as any reasonable view of $\Z$ should be invariant to reparametrizations $g$.

  The classic geometric solution to the identifiability problem is to define any quantity of interest in the observation space $\X$ rather than the latent space $\Z$. For example, the length of a curve $\gamma: [0, 1] \rightarrow \Z$ in the latent space can be defined as its length measured in $\X$:
  \begin{align}
    \mathrm{Length}[\gamma]
      &= \int_0^1 \| \partial_t f(\gamma(t)) \| \dif t
       = \int_0^1 \| \mat{J}_{\gamma(t)} \dot{\gamma(t)} \| \dif t
       = \int_0^1 \sqrt{ \dot{\gamma(t)}\T \mat{J}_{\gamma(t)}\T \mat{J}_{\gamma(t)} \dot{\gamma(t)} } \dif t,
     \label{eq:curve_length}
  \end{align}
  where $\dot{\gamma(t)} = \partial_t \gamma(t)$ is the curve derivative, and $\mat{J}_{\gamma(t)}$ is the Jacobian of $f$ at $\gamma(t)$. This construction shows how we may calculate lengths in the latent space with respect to the metric of the observation space. The symmetric positive definite matrix $\mat{J}_{\gamma(t)}\T \mat{J}_{\gamma(t)}$ often appears and we denote it $\b{M}_{\gamma(t)}$. As the Jacobian spans the tangent space of $f$, we may interpret $\b{M}$ as an inner product $\inner{\vec{u}}{\vec{v}}_{\b{M}} = \vec{u}\T\b{M}\vec{v}$ over this tangent space, given us all the ingredients to define \emph{Riemannian manifolds}:
  \begin{definition}
    A Riemannian manifold is a smooth manifold $\M$ together with a Riemannian metric $\b{M}_{\vec{z}}$, which is a positive definite matrix that changes smoothly throughout the space and defines an inner product on the tangent space $\mathcal{T}_{\vec{z}}\M$.
  \end{definition}

  We see that the decoder naturally spans a Riemannian manifold. As the Riemannian length of a latent curve \eqref{eq:curve_length}, by construction, is invariant to reparametrizations, it is natural to extend this view with a notion of \emph{distance}. We say that the distance between two points $\vec{z}_0, \vec{z}_1 \in \Z$ is simply the length of the shortest connecting path, $\mathrm{dist}(\vec{z}_0, \vec{z}_1) = \min_{\gamma} \mathrm{Length}(\gamma)$. Calculating distances, then imply first finding the shortest path. One can show \cite{xxx} that length minimizing curves also have minimal \emph{energy}
  \begin{align}
    \mathrm{Energy}(\gamma)
      &= \int_0^1 \| \partial_t f(\gamma(t)) \|^2 \dif t
       = \int_0^1 \dot{\gamma(t)}\T \b{M}_{\gamma(t)} \dot{\gamma(t)} \dif t,
  \end{align}
  which is a locally convex functional. Shortest paths can then be found by direct energy minimization or by solving a corresponding system of ordinary differential equations (see supplements).

 \subsection{Stochastic decoders}

As previously discussed, deterministic decoders directly define a Riemannian geometry in the latent space. However, most models of interest are stochastic and there is significant evidence that this stochasticity is important to faithfully capture the intrinsic structure of data \cite{hauberg:only:2018}. When the decoder is a smooth stochastic process, e.g.\ as in the Gaussian process latent variable model \cite{lawrence:2005:jmlr}, \citet{tosi:uai:2014} laid the ground works for modeling a stochastic geometry. Most contemporary models, such as VAEs, are however not stochastically smooth as the noise is assumed independent, and this theory is inapplicable. \citet{arvanitidis:iclr:2018} proposed an extension to VAEs with Gaussian decoders
\begin{align}
    f(\vec{z})
      &= \mu(\vec{z}) + \sigma(\vec{z}) \odot \vec{\epsilon}
       = \begin{bmatrix}
           \mat{I} & \mathrm{diag}(\vec{\epsilon})
         \end{bmatrix}
         \begin{bmatrix}
           \mu(\vec{z}) \\ \sigma(\vec{z})
         \end{bmatrix}
      = \mat{P} h(\vec{z}),
    \quad \text{where } \vec{\epsilon} \sim \mathcal{N}(\vec{0}, \mat{I}).
\end{align}
Here we have written the Gaussian decoder in its reparametrized form. This can be viewed as a random projection of a deterministic manifold \cite{eklund:arxiv:2019} spanned by $h$, which can easily be given a geometry. The associated metric,\vspace{-3mm}
\begin{align}
\label{eq:vae_metric}
    \b{M}_{\b{z}} = \b{J}_{\mu(\b{z})}\T \b{J}_{\mu(\b{z})} + \b{J}_{\sigma(\b{z})}\T \b{J}_{\sigma(\b{z})},
\end{align}
gives shortest paths that follow the data as distances grow with the model uncertainty \cite{arvanitidis:iclr:2018, hauberg:only:2018}.

Previous work has, thus, focused on \emph{pulling back} the Euclidean distance from the observation space to the latent space using the reparametrization of the Gaussian decoder. This is, however, intrinsically linked with the simple reparametrization of the Gaussian, and the strategy does not extend to non-Gaussian decoder distributions. We propose an alternative, principled way of dealing with uncertainty by changing the focus from observations to distributions, leveraging the metrics defined in classical information geometry.

\section{An information geometric latent metric}
So far we have seen how we can endow the latent space $\Z$ with meaningful distances when our stochastic decoder mappings are reparameterizable and their codomain lies in the observation space (as in the case of a Gaussian). Ideally, we would like a more general framework of computing meaningful shortest path distances even when our decoders parameterize discrete and/or non-reparameterizable distributions (such as a Bernoulli distribution).

We first note that the codomain of a VAE decoder is the parameter space $\HH$ of a distribution. With this in mind, we can ask what is a natural distance in the latent space $\Z$ between two points $\vec{z}_1$ and $\vec{z}_2 = \vec{z}_1\!+\!\epsilon$ when measured in $\HH$. Since our latent codes correspond to distributions we can define the (infinitesimal) distance through the KL-divergence:
\begin{equation}
    \mathrm{dist}^2(\vec{z}_1, \vec{z}_2) = \mathrm{KL}(p(\vec{x}|\vec{z}_1), p(\vec{x}|\vec{z}_2)).
    \label{eq:kl_dist}
\end{equation}
We can then define the length of a latent curve $\gamma$ as
\begin{align}
  \mathrm{Length}(\gamma)
    &= \lim_{N \rightarrow \infty} \sum_{n=1}^{N-1} \sqrt{\mathrm{KL}(p(\vec{x}|\gamma(\sfrac{n}{N}), p(\vec{x}|\gamma(\sfrac{n\!+\!1}{N}))},
\end{align}
and distances could be defined as before.
This would satisfy our desiderata of a deterministic notion of similarity in the latent space that is applicable to wide array of decoder distributions.

This construction may seem arbitrary, but in reality it carries deeper geometric meaning. \emph{Information geometry} \citep{nielsen:2020} considers families of probabilistic densities $p(\vec{x}|\eta)$ as represented by their parameters $\eta \in \HH$, such that $\HH$ is constructed as a statistical manifold equipped with the Fisher-Rao metric, which infinitesimally coincides with the KL divergence as in Eq.~\ref{eq:kl_dist}. This is known to be a Riemannian metric over $\HH$ which takes the following form:

\begin{equation}
\label{eq:fisher_rao_metric}
    \b{I}_{\HH}(\eta) = \int_\X \left[\nabla_{\eta} \log p(\b{x}~|~\eta) \cdot \nabla_{\eta} \log p(\b{x}~|~\eta)\T\right] p(\b{x}~|~\eta)\dif \b{x}.
\end{equation}
When the parameter space $\HH$ is equipped with this metric, we call it a \emph{statistical manifold}.
\begin{definition}
  A \emph{statistical manifold} whose points are probability distributions is the parameter space $\HH$ of a probability density function $p(\b{x}|\eta)$ that is a smooth manifold, equipped with a Riemannian metric tensor which is the Fisher-Rao information matrix \eqref{eq:fisher_rao_metric}.
\end{definition}
This definition implies that the geometry that is induced by the Fisher-Rao metric is predefined and is a modeling decision, since it does not depend on any learnable parameters given the data.

As mentioned, the Fisher-Rao metric coincides with the KL-divergence locally:
\begin{proposition}
  The Fisher-Rao metric is the first order approximation of the KL-divergence between perturbed distributions:
  \begin{equation}
  \label{eq:KL_approx_inner_product}
    \mathrm{KL}(p(\b{x}|\eta),~p(\b{x}|\eta+\delta\eta)) = \frac{1}{2} \delta\eta\T \b{I}_{\HH}(\eta) \delta\eta + o(\delta\eta^2).
  \end{equation}
\end{proposition}
The central idea put forward in this paper is to consider the decoder as a map $f: \Z \rightarrow \HH$, and let $\HH$ be equipped with the appropriate Fisher-Rao metric. The VAE can then be interpreted as spanning a submanifold of $\HH$ and the latent space can be endowed with the corresponding metric. We detail this approach in the sequel.

\subsection{The Riemannian pull-back metric}
  Our construction implies that the length of a latent curve $\gamma$ is measured as
  \begin{align}
    \mathrm{Length}(\gamma)
      &= \int_0^t \sqrt{(\partial_t \gamma(t))\T \b{I}_{\HH}(\gamma(t) (\partial_t \gamma(t))} \dif t.
  \end{align}
  %
  Applying the chain-rule then trivially gives:
  %
  \begin{proposition}
    \label{prop:fisher_pullback}
    Let $f:\Z\rightarrow\HH$ be an immersion that parametrizes the likelihood. Then, the latent space $\mathcal{Z}$ is equipped with the pull-back metric $\b{M}(\b{z}) = \b{J}_{f}\T(\b{z}) \b{I}_{\HH}(f(\b{z})) \b{J}_{f}(\b{z})$.
  \end{proposition}

  Instead of considering the parameters $\eta\in\HH$ of the probabilistic density function that models the data, we can consider the latent variable $\b{z}$ as the parameters of the model. This view is equivalent to the one explained above, and the corresponding pull-back metric is directly the Fisher-Rao metric endowed in the latent space $\Z$:
  %
  \begin{proposition}
    The pullback metric $\b{M}(\b{z})$ is identical to the Fisher-Rao metric obtained over the parameter space $\Z$: $\b{M}(\b{z}) = \b{I}_{\HH}(\b{z}) = \int_{\X} \left[\nabla_\b{z}\log p(\b{x}|\b{z}) \nabla_\b{z} \log p(\b{x}|\b{z})\T\right] p(\b{x}|\b{z}) \dif\b{x}. $
  \end{proposition}

\begin{figure}[t]
    \centering

    \begin{subfigure}{0.3\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/dist_X.pdf}
    \end{subfigure}
    ~
    \begin{subfigure}{0.3\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/dist_H.pdf}
    \end{subfigure}
    ~
    \begin{subfigure}{0.3\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/dist_Z.pdf}
    \end{subfigure}

    \caption{Here we compare in $\Z$ the shortest path on the approximated data manifold $g(\Z)\subset\X$ with the proposed shortest path on $f(\Z)\subset\HH$. \gear{Maybe is interesting to show the sequence of $p(\b{x}|\b{z})$ if the shortest path in $\HH$ is computed using only the Fisher-Rao and not constrained to be on our learned surface. Maybe in the experiments?}}
    \label{fig:proposed_method_example}
\end{figure}

In practice, our mapping $f(\cdot)$ immerses a deterministic surface into the parameter space $\HH$ of probability distributions, where the metric of the space is the Fisher-Rao. The resulting $\b{M}(\cdot)$ is the pull-back of the $\b{I}_{\HH}(\cdot)$. In other words, this allows us to compute the shortest path between two distributions that lie on the surface $f(\Z)$, while on the same time respecting the geometry of the parameter space $\HH$ that is defined by $\b{I}_{\HH}(\cdot)$. While a common practice is to use the Fisher-Rao metric that is induced by the modeling choice, in our method the metric $\b{M}(\cdot)$ is additionally based on the data through the training of the generative model. This is in contrast to other approaches that might simply use the Fisher-Rao metric of likelihood, which is agnostic of the data and depends only on modeling assumptions.

Also, note that the distances in $\Z$ between two points computed under the pull-back metric $\b{M}(\cdot)$ is identifiable \citep{hauberg:only:2018}. This means that even if we re-parametrize the latent space $\Z$ with a bijective function $h:\Z\rightarrow\Z$, as long as the function $f(\cdot)$ is the same the distances are preserved. In particular, the $\mathrm{length}[c] = \mathrm{length}[c']$ for $c(t)\in\Z$ computed under the $\b{M}(\cdot)$ and $c(t) = h(c'(t))$ computed under the metric $\b{J}_h\T(\cdot)\b{M}(h(\cdot))\b{J}_h\T(\cdot)$.

The Fisher-Rao metric has closed-form solution for a one-dimensional observation $\b{x}$ that follows a specific distribution. In practice, $\b{x}\in\mathcal{X}$ is a multi-dimensional variable where each dimension represents, for instance, a pixel when working with images or a feature when working with tabular data. Each feature, $x_i$ with $i=1\cdots D$, is obtained for a specific set of parameters $\{\eta_i\}$. We assume that the features follow the same distribution $\mathcal{D}$, such that: $x_i \sim \mathcal{D}(\eta_i)$, and $p(\b{x}|\eta) = \prod_{i=1}^{D}p(x_i|\eta_i)$. This assumption enables us to define a general Fisher-Rao metric that represents the multi-dimensional variable $\b{x}$.

\begin{proposition} \label{prop:blockmatrix}
The Fisher-Rao metric $\b{I}_{\HH}(\eta)$ is a block diagonal matrix where the diagonal terms are the Fisher-Rao matrices $\b{I}_{{\HH}_i}(\eta_i)$ obtained for each data feature $x_i$.
\end{proposition}

Hence, we are able to compute shortest paths in $\Z$ that respect the metric $\b{M}(\cdot)$ from Prop.~\ref{prop:fisher_pullback}, even if the likelihood is not re-parametrizable as the Gaussian. Of course, the resulting curve is different from the curve in the latent space that we find if we use the pull-back of the generator. In Fig.~\ref{fig:proposed_method_example} we compare the resulting curves in $\Z$ when the generator is the Gaussian, between the pull-back of the Euclidean and the Fisher-Rao. \gear{What is the difference in behavior? Can be theoretically analyzed to some extend?} \alison{good question!}  Of course, the two paths are not exactly the same, however, they tend to have the same behavior. This can be achieved by using the same idea of stochasticity in the Gaussian decoder. In particular, when we move away from the latent codes we want the $f(\b{z})$ to increase similar to the Gaussian's uncertainty. \gear{we do this right?}

\iffalse
\begin{wrapfigure}{r}{0.2\textwidth}
    \vspace{-15pt}
    \begin{center}
    \includegraphics[width=0.2\textwidth]{imgs/teaser_3.pdf}
  \end{center}
  \caption{Example of shortest path between Gaussians.}
  \vspace{-15pt}
  \label{fig:inf_geom_example}
\end{wrapfigure}

The Fisher-Rao metric has been shown to satisfy the properties of a Riemannian metric \citep{rao:1945}. This defines the geometry of the parameter space $\HH$ and enables us to define distances between two distributions.

Let's take for example the univariate Gaussian distribution: $\N(\mu,\sigma^2 )$. This distribution can be represented by its parameters $\eta = [\mu, \sigma]$, which forms a vector in the parameter space $\HH\subset\R^{2}$. This parameter space is equipped with the Fisher-Rao metric: $I([\mu, \sigma])=\mathrm{diag}(\frac{1}{\sigma^2}, \frac{2}{\sigma^2})$ \cite{tomczak:2012}. It has been shown that this metric is similar to the Poincare Half-plane metric up to a change of variables \cite{costa:2014}, and so the distance between two univariate gaussians follows an arc length, as represented in Figure (?). \alison{Is it the right place for this paragraph? Is it a bit too overkill to talk about the poincare half plane? But it nicely shows that the distance between two gaussians follows an arc lenght :-)}
\dimitris{I think it nicely highlights the wisdom of thinking of statistical manifolds as proper deterministic manifolds and fits nicely with smoothly links the section on Riemm. geom. with the current one.}

\begin{definition}
A statistical manifold whose points are probability distributions is the parameter space $\HH$ of a probability density function $p(\b{x}|\eta)$ that is a smooth manifold, equipped with a Riemannian metric tensor that is the Fisher-Rao information matrix.
\end{definition}
This definition implies that the geometry that is induced by the Fisher-Rao metric is predefined and is a modeling decision, since it does not depend on any learnable parameters given the data.

In most cases, the Fisher-Rao metric has a closed form expression (see Appendix, Section.~\ref{appendix:FR:distributions}). When the computation of the expectation Eq.~\ref{eq:fisher_rao_metric} is intractable, the integral can be estimated empirically using the training data, but this can lead to a degenerate metric and may not represent the same geometrical structure as the true metric \cite[\S 4.1.4]{nielsen:2020}.

As we mentioned above this is a Riemannian metric, so it defines a smoothly changing inner product $\inner{\delta\eta}{\b{I}_{\HH}(\eta)\delta\eta}$ between tangent vectors $\delta\eta \in \mathcal{T}_{\eta}\HH$. Note that we consider the whole $\HH$ as a smooth manifold, so the tangent space $\mathcal{T}_{\eta}\HH$ is the same space centered at $\eta$ and the metric simply changes the way we measure distances locally. Therefore, the length of a curve $\gamma:[0,1]\rightarrow\HH$ can be computed again with the integral $\mathrm{length}[\gamma] = \int_0^1 \sqrt{\inner{\dot{\gamma}(t)}{\b{I}_{\HH}(\gamma(t))\dot{\gamma}(t)}}\dif t$. Therefore, we are able to compute the length minimizing curve between two distributions by minimizing the energy functional as we described in Sec.~\ref{sec:riemannian_geometry}, which results in a system of ODEs. For example, in Fig.~\ref{fig:inf_geom_example} we show such a path with minimal length between two multivariate Gaussians with diagonal covariance.

This means that when the close form solution of the Fisher-Rao metric is intractable, we are not able to find the shortest path using the system of ODEs. However, the following proposition shows that we can approximate instead the local inner product using the Kullback–Leibler (KL) divergence.

\begin{proposition}
The Fisher-Rao metric is the first order approximation of the KL-divergence between perturbed distributions:
\begin{equation}
\label{eq:KL_approx_inner_product}
    \mathrm{KL}(p(\b{x}|\eta),~p(\b{x}|\eta+\delta\eta)) = \frac{1}{2} \delta\eta\T \b{I}_{\HH}(\eta) \delta\eta + o(\delta\eta^2).
\end{equation}
\end{proposition}
This allows us to approximate the length of the curve on the statistical manifold as $\mathrm{length}[\gamma]=\int_0^1\sqrt{\mathrm{KL}(p(\b{x}|\gamma(t)),~p(\b{x}|\gamma(t+\delta t)))} dt$, where $\delta t \rightarrow 0$.
This is particularly useful because in some cases the KL divergence between two distributions can be analytically derived. In this way, we can approximate the length of the curve by discretizing it and computing the piecewise quantities.

\begin{proposition}
The KL-divergence between two close elements of the curve $\gamma$ is defined as: $\mathrm{KL}(p_{t},p_{t+\delta t})=\mathrm{KL}(p(\b{x}|\gamma(t)),~p(\b{x}|\gamma(t+\delta t)))$.
The length $\mathcal{L}$ and the energy $\mathcal{E}$ functionals can be approximated with respect to this KL-divergence:
\begin{equation}
\label{eq:eq:discretized_length_energy}
    \mathcal{L}(\gamma) \approx\sqrt{2\textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})}  \quad\text{and}\quad \mathcal{E}(\gamma) \approx \frac{\delta t}{2}\textstyle\sum_{t=1}^{T} \mathrm{KL}(p_{t},p_{t+\delta t})
\end{equation}
\end{proposition}

So far we have seen how we can transform the parameter space $\HH$ of a probability distribution $p(\b{x}|\eta)$ into a Riemannian manifold, which allows us to compute shortest paths between points. In the next section, we show how we can exploit the geometric structure of $\HH$ in order to compute shortest paths in the latent space of generative models, when the conditional density is not necessarily reparameterizable like e.g. a Gaussian density.
\fi

\section{Proposed methodology}
\label{sec:proposed_method}

As discussed in Sec.~\ref{sec:generative_models}, we consider that the generator is learnt to stochastically embed the data from lower dimensional latent space $\Z=\R^d$ to an observational space $\X=\R^{D}$, such that $\b{x} = g(\b{z})$, with $g(\cdot)$ a smooth stochastic function. The latent space being endowed with the stochastic pull-back metric $\b{J}_{g}(\b{z})\T \b{J}_{g}(\b{z})$, we can then redefine the length and the energy of curve used to compute geodesics. For practical purposes, we are usually seeking a deterministic approximation of the metric, leading us to take its expectation $\mathbb{E}[\b{J}_{g}(\b{z})\T\b{J}_{g}(\b{z})]$.

\paragraph{Shortest path computation.} In practice, we are able to compute the shortest paths in $\Z$ with respect to $\b{M}(\b{z})$ using the same system of ODEs (see App.~\ref{}). However, evaluating the metric is usually computationally expensive and even if a close form solution is available, the Jacobian of $f(\cdot)$ and its derivatives are needed. Instead, we can use Prop.~\ref{prop:discretized_length_energy} and optimize the length similar to \cite{yang:arxiv:2018}. In particular, we can parametrize a curve $c_\psi(t)$ in $\Z$ using a cubic spline with parameters $\psi$. Then, this curve can be mapped into the statistical manifold $\HH$ using the trained function $\gamma(t)=f_\theta(c(t))$ of the VAE. We discretize the interval $[0,1]$, and hence, we approximate the length using the finite sum of KL divergences. The KL divergence in many cases can be computed analytically. This allows to optimize the parameters $\psi$ using automatic differentiation, simply by considering the energy as our cost function.

\section{Experiments}
\label{sec:experiments}

\begin{figure}[t]
    \centering

    \begin{subfigure}{0.3\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/land_1.pdf}
    \end{subfigure}
    ~
    \begin{subfigure}{0.3\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/land_1.pdf}
    \end{subfigure}
    ~
    \begin{subfigure}{0.3\textwidth}
	    \includegraphics[width=1\linewidth]{imgs/compare_paths.pdf}
    \end{subfigure}

    \caption{\emph{Left}: We compare the LAND fits for Gaussian and Bernouli decoder using MNIST data (should be similar). \emph{Right}: We compare the shortest paths for different decoders when the latent codes are given.}
    \label{fig:ex:1}
\end{figure}

\subsection{Compare generators pull-back with Fisher-Rao pull-back using Gaussians}

We consider a synthetic example where the data lie in a continuous space. We train a VAE and we compare in the latent space the result of the generator's pull-back and the result of the Fisher-Rao pull-back. In addition, we show the result of the LAND. \gear{We expect to be approximately the same.}

\subsection{Some toy examples with Gaussian, Bernoulli, Gamma, etc.}

We consider multiple synthetic datasets. We use the same latent codes (e.g. a donut in $\Z$) and we use different generators (Gaussian, Bernoulli, Gamma, von Mises-Fisher, etc). We see how the shortest paths behave for each case. \gear{Probably in all cases the shortest paths approximately should be the same.}

\subsection{MNIST with Bernoulli.}

We use the MNIST data (0,1,2) and we train a VAE with Bernoulli decoder. In the latent space we fit a LAND. Should we do the same with Gaussian decoders and compare the fits? We can also show the decoded sequence only using the Fisher-Rao and not the learned surface.

\subsection{MoCap data with von Mises-Fisher}

We use a more funky dataset and decoder. The resulting latent space is ``complicated'' where we are able to compute shortest paths.

\section{Conclusion and Discussion}
\label{sec:conclusion}

We have provided a way to compute shortest paths in the latent space of a generative model, which allows any kind of generator to be used. In practice, we compute the shortest path on the manifold that is immersed in the space of probability distributions, and we pull-back the Fisher-Rao metric of this particular probability distribution. This is different from the baseline where the generator has to be Gaussian in order to pull-back the metric of the data space. Interestingly, we do not have to compute the Fisher-Rao metric, but we can use the $\mathrm{KL}$ between infitesimally close probability distributions together with automatic differentiation in order to compute the shortest path. In future work, it would be interesting to see what happens if we replace the metric of the probability space i.e. replace the Fisher-Rao with Wasserstein?

\clearpage
\section*{Notations}
For us, it can be modified and will not appear in the final draft. :)

\begin{tabular}{cp{0.8\textwidth}}
  $z\in\mathcal{Z}\subset\mathbb{R}^{q}$ & Latent variable ($z$) in the latent space ($\mathcal{Z}$) \\
  $x\in\mathcal{X}\subset\mathbb{R}^{D}$ & Data variable ($x$) in the observational space ($\mathcal{X}$) \\
  $\eta\in\mathcal{H}$ & Parameter ($\eta$) of a known distribution, in the statistical manifold ($\mathcal{H}$)\\
  $I(\eta)$ & Fisher-Rao metric obtained in $\mathcal{H}$. (Notation borrowed from \cite{nielsen:2020})  \\
  $J_f$& Jacobian of the decoder $f$ \\
  $\mathcal{L}, \mathcal{E}, \mathcal{V}$ & Length, energy and volume of a metric. \\
\end{tabular}

redpajama_set_name: RedPajamaArXiv