Neural Manifold Ordinary Differential Equations

Introduction

Deep generative models are a powerful class of neural networks which fit a probability distribution to produce new, unique samples. While latent variable models such as Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs) are capable of producing reasonable samples, computing the exact modeled data posterior is fundamentally intractable. By comparison, normalizing flows are capable of learning rich and tractable posteriors by transforming a simple probability distribution through a sequence of invertible mappings. Formally, in a normalizing flow, a complex distribution $p(x)$ is transformed to a simple distribution $\pi(z)$ via a diffeomorphism $f$ (i.e. a differentiable bijective map with a differentiable inverse) with probability values given by the change of variables:

To compute this update efficiently, $f$ must be constrained to allow for fast evaluation of the determinant, which in the absence of additional constraints takes $\mathcal{O}(D^{3})$ time (where $D$ is the dimension of $z$). Furthermore, to efficiently generate samples, $f$ must have a computationally cheap inverse. Existing literature increases the expressiveness of such models under these computational constraints and oftentimes parameterizes $f$ with deep neural networks . An important recent advancement, dubbed the Continuous Normalizing Flow (CNF), constructs $f$ using a Neural Ordinary Differential Equation (ODE) with dynamics $g$ and invokes a continuous change of variables which requires only the trace of the Jacobian of $g$ .

Since $f$ is a diffeomorphism, the topologies of the distributions $p$ and $\pi$ must be equivalent. Furthermore, this topology must conform with the underlying latent space, which previous work mostly assumes to be Euclidean. However, topologically nontrivial data often arise in real world examples such as in quantum field theory , motion estimation , and protein structure prediction .

Preexisting manifold normalizing flow works (which we present a complete history of in Section 2) do not generalize to arbitrary manifolds. Furthermore, many examples present constructions extrinsic to the manifold. In this work, we solve these issues by introducing Manifold Continuous Normalizing Flows (MCNFs), a manifold analogue of Continuous Normalizing Flows. Concretely, we:

introduce Neural Manifold ODEs as a generalization of Neural ODEs (seen in Figure 1). We leverage existing literature to provide methods for forward mode integration, and we derive a manifold analogue of the adjoint method for backward mode gradient computation.

develop a dynamic chart method to realize Neural Manifold ODEs in practice. This approach integrates local dynamics in Euclidean space and passes between domains using smooth chart transitions. Because of this, we perform computations efficiently and can accurately compute gradients. Additionally, this allows us to access advanced ODE solvers (without manifold equivalents) and augment the Neural Manifold ODE with existing Neural ODE improvements .

construct Manifold Continuous Normalizing Flows. These flows are constructed by integrating local dynamics to construct diffeomorphisms, meaning that they are theoretically complete over general manifolds. Empirically, we find that our method outperforms existing manifold normalizing flows on their specific domain.

Related Work

In this section we analyze all major preexisting manifold normalizing flows. Previous methods are, in general, hampered by a lack of generality and are burdensomely constructive.

Normalizing Flows on Riemannian Manifolds . The first manifold normalizing flow work constructs examples on Riemannian manifolds by first projecting onto Euclidean space, applying a predefined Euclidean normalizing flow, and projecting back. Although simple, this construction is theoretically flawed since the initial manifold projection requires the manifold to be diffeomorphic to Euclidean space. This is not always the case, since, for example, the existence of antipodal points on a sphere necessarily implies that the sphere is not diffeomorphic to Euclidean space. As a result, the construction only works on a relatively small and topologically trivial subset of manifolds.

Our work overcomes this problem by integrating local dynamics to construct a global diffeomorphism. By doing so, we do not have to relate our entire manifold with some Euclidean space, but rather only well-behaved local neighborhoods. We test against on hyperbolic space, and our results produce a significant improvement.

Latent Variable Modeling with Hyperbolic Normalizing Flows . In a recent manifold normalizing flow paper, the authors propose two normalizing flows on hyperbolic space—a specific Riemannian manifold. These models, which they name the Tangent Coupling (TC) and Wrapped Hyperboloid Coupling (WHC), are not affected by the aforementioned problem since hyperbolic space is diffeomorphic to Euclidean space. However, various shortcomings exist. First, in our experiments we find that the methods do not seem to conclusively outperform . Second, these methods do not generalize to topologically nontrivial manifolds. This means that these flows produce no additional topological complexity and thus the main benefit of manifold normalizing flows is not realized. Third, the WHC construction is not intrinsic to hyperbolic space since it relies on the hyperboloid equation.

Our method, by contrast, is derived from vector fields, which are natural manifold constructs. This not only allows for generality, but also means that our construction respects the manifold geometry. We compare against on hyperbolic space and find that our results produce a substantial improvement.

In comparison, our work is general, requires only local diffeomorphisms, produces globally defined densities, and is tractable for higher dimensions. We test against on the sphere and attain markedly better results.

Other Related Work. In , the authors define a probability distribution on Lie Groups, a special type of manifold, and as a by-product construct a normalizing flow. The constructed flow is very similar to that found in , but the authors include a $\tanh$ nonlinearity at the end of the Euclidean flow to constrain the input space and guarantee that the map back to the manifold is injective. We do not compare directly against this work since Lie Groups are not general (the $2$-dimensional sphere is not a Lie Group ) while Riemannian manifolds are.

There are some related works such as that are orthogonal to our considerations as they either (i) develop applications as opposed to theory or (ii) utilize normalizing flows as a tool to study Riemannian metrics.

Concurrent work also investigates the extension of neural ODEs to smoooth manifolds.

Background

In this section, we present background knowledge to establish naming conventions and intuitively illustrate the constructions used for our work. For a more detailed overview, we recommend consulting a text such as .

Tangent space. For an $n$-dimensional manifold $\mathcal{M}$, the tangent space $T_{x}\mathcal{M}$ at a point $x\in\mathcal{M}$ is a higher-dimensional analogue of a tangent plane at a point on a surface. It is an $n$-dimensional real vector space for all points $x\in\mathcal{M}$.

As one might expect, the pushforward is central in the definition of manifold ODEs (analogous to the importance of the common derivative in Euclidean ODEs). We also use it in our dynamic chart method to map tangent vectors of the manifold to tangent vectors of Euclidean space.

2 Riemannian Geometry

While the above theory is general, to concretize some computational aspects (e.g. how to pick charts) and give background on related manifold normalizing flow work, we define relevant concepts from Riemannian geometry.

Exponential Map. The exponential map $\exp_{x}\mathrel{\mathop{\mathchar 58\relax}}T_{x}\mathcal{M}\to\mathcal{M}$ can be thought of as taking a vector $v\in T_{x}\mathcal{M}$ and following the general direction (on the manifold) such that the distance traveled is the length of the tangent vector. Specifically, the distance on the manifold matches the induced tangent space norm $d_{\rho}(x,\exp_{x}(v))=\left\|v\right\|_{\rho}\mathrel{\mathop{\mathchar 58\relax}}=\sqrt{\left\langle v,v\right\rangle_{\rho}}$. Note that $\exp_{x}(0)=x$.

The exponential map is crucial in our construction since it acts as a chart. Specifically, if we identify the chart domain with $T_{x}\mathcal{M}$ then $\exp_{x}$ is a diffeomorphism when restricted to some local set around .

3 Manifold Ordinary Differential Equations

Manifold ODE. Finally, we introduce the key objects of study: manifold ordinary differential equations. A manifold ODE is an equation which relates a curve $\mathbf{z}\mathrel{\mathop{\mathchar 58\relax}}[t_{s},t_{e}]\to\mathcal{M}$ to a vector field $f$ and takes the form

$\mathbf{z}$ is a solution to the ODE if it satisfies Equation 1 with initial condition $z_{s}$. Similarly to the case of classical ODEs, local solutions are guaranteed to exist under sufficient conditions on $f$ .

Neural Manifold Ordinary Differential Equations

To leverage manifold ODEs in a deep learning framework similar to Neural ODEs , we parameterize the dynamics $f$ by a neural network with parameters $\theta$. We define both forward pass and backward pass gradient computation strategies in this framework. Unlike Neural ODEs, forward and backward computations do not perfectly mirror each other since forward mode computation requires explicit manifold methods, while the backward can be defined solely through an ODE in Euclidean space.

The first step in defining our Neural Manifold ODE block is the forward mode integration. We review and select appropriate solvers from the literature; for a more thorough introduction we recommend consulting a text such as . Broadly speaking, these solvers can be classified into two groups:

implicit methods that solve the manifold ODE locally using charts. These methods only require manifold-implicit constructions.

Projection methods are conceptually simple, but ultimately suffer from generality issues. In particular, for manifolds such as the open ball or the upper half-space, there is no principled way to project off-manifold points back on. Furthermore, even in nominally well-defined cases such as the sphere, there may still exist points such as the origin for which the projection is not well-defined.

Implicit methods, by contrast, can be applied to any manifold and do not require a level set representation. Thus, this approach is more amenable to our generality concerns (especially since we wish to work with, for example, hyperbolic space). However, difficulty in defining charts restricts applicability. Due to this reason, implicit schemes often employ additional structure to define manifold variations of step-based solvers . For example, on a Riemannian manifold one can define a variant of an Euler Method solver with update step $z_{t+\epsilon}=\exp_{z_{t}}(\epsilon f(z_{t},t))$ using the Riemannian exponential map. On a Lie group, there are more advanced Runge-Kutta solvers that use the Lie Exponential map and coordinate frames .

2 Backward Mode Adjoint Gradient Computation

In order to fully incorporate manifold ODEs into the deep learning framework, we must also efficiently compute gradients. Similar to , we develop an adjoint method to analytically calculate the derivative of a manifold ODE instead of directly differentiating through the solver. Similar to manifold adjoint methods for partial differential equations, we use an ambient space .

This theorem resembles the adjoint method in precisely because of our ambient space condition. In particular, our curve $\mathbf{z}\mathrel{\mathop{\mathchar 58\relax}}[t_{s},t_{e}]\to\mathcal{M}$ can be considered as a curve in the ambient space. Furthermore, we do not lose any generality since such an embedding always exists by the Whitney Embedding Theorem , if we simply set $d=2n$.

The full proof of Theorem 4.1 can be found in Appendix A.3. Through this adjoint state, gradients can be derived for other parameters in the equation such as $t_{s},t_{e}$, initial condition $z_{s}$, and weights $\theta$.

Dynamic Chart Method

While our above theory is fully expressive and general, in this section we address certain computational issues and augment applicability with our dynamic chart method.

The motivation for our dynamic chart method comes from , where the author introduces a dynamic manifold trivialization technique for Riemannian gradient descent. Here, instead of applying the traditional Riemannian gradient update $z_{t+1}=\exp_{z_{t}}(-\eta\nabla_{z_{t}}f)$ for $N$ time steps, the author repeatedly applies $n\ll N$ local updates. Each update consists of a local diffeomorphism to Euclidean space, $n$ equivalent Euclidean gradient updates, and a map back to the manifold. This allows us to lower the number of expensive exponential map calls and invoke existing Euclidean gradient optimizers such as Adam . This is similar in spirit to , but crucially only relies on local diffeomorphisms rather than a global diffeomorphism.

We can adopt this strategy in our Neural Manifold ODE. Specifically, we develop a generalization of the dynamic manifold trivialization for the manifold ODE forward pass. In a similar spirit, we use a local chart to map to Euclidean space, solve an equivalent Euclidean ODE locally, and project back to the manifold using the chart. The full forward pass is given by Algorithm 1 and is visualized Figure 2.

We present two propositions which highlight that this algorithm is guaranteed to converge to the manifold ODE solution. The first shows that $\varphi_{z}(\mathbf{y})$ solves the manifold differential equation locally and the second shows that we can pick a finite collection of charts such that we can integrate to time $t_{e}$.

Proofs of these propositions are given in Appendix A.1. Note that this forward integration is implicit, indicating the connections between and . We realize this construction by finding a principled way to pick charts and incorporate this method into neural networks by developing a backward gradient computation.

This allows for gradient computation through backpropagation. We may differentiate through the Neural ODE blocks by the Euclidean adjoint method .

Our dynamic chart method is a significant advancement over previous implicit methods since we can easily construct charts as we integrate. Furthermore, it provides many practical improvements over vanilla Neural Manifold ODEs. Specifically, we

can perform faster evaluations. The aforementioned single step algorithms rely on repeated use of the Lie and Riemannian exponential maps. These are expensive to compute and our method can sidestep this expensive evaluation. In particular, the cost is shifted to the derivative of the chart, but by defining dynamics $g$ on the tangent space directly, we can avoid this computation. We use this for our hyperbolic space construction, where we simply solve $\frac{d\mathbf{y}(t)}{dt}=g(\mathbf{y}(t),t)$.

avoid catastrophic gradient instability. If the dimension of $\mathcal{M}$ is less than the dimension of the ambient space, then the tangent spaces are of measure . This means that the induced error from the ODE solver will cause our gradients to leave their domain, resulting in a catastrophic failure. However, since the errors in Neural ODE blocks do not cause the gradient to leave their domain and as our charts induce only a precision error, our dynamic chart method avoids this trap.

access a wider array of ODE advancements. While substantial work has been done for manifold ODE solvers, the vast majority of cutting edge ODE solvers are still restricted to Euclidean space. Our dynamic chart method can make full use of these advanced solvers in the Neural ODE blocks. Additionally, Neural ODE improvements such as can be directly integrated without additional manifold constructions.

Manifold Continuous Normalizing Flows

With our dynamic chart Neural Manifold ODE, we can construct a Manifold Continuous Normalizing Flow (MCNF). The value can be integrated directly through the ODE, so all that needs to be given is the change in log probability. Here, we can invoke continuous change of variables on the Neural ODE block and can use the smooth chart transition property (which guarantees that the charts are diffeomorphisms) to calculate the final change in probability as:

Note that we drop derivative and integration arguments for brevity. A full expression is given in Appendix B.

For our purposes, the last required computation is the determinant of the $D_{v}\exp_{x}$. We find that these can be evaluated analytically, as shown in the cases of the hyperboloid and sphere .

Since the MCNF requires only the local dynamics (which are in practice parameterized by the exponential map), this means that the construction generalizes to arbitrary manifolds. Furthermore, we avoid diffeomorphism issues, such as in the case of two antipodal points on the sphere, simply by restricting our chart domains to never include these conjugate points.

Experiments

2 Variational Inference

We train a hyperbolic VAE and Euclidean VAE for variational inference on Binarized Omniglot and Binarized MNIST . Both of these datasets have been found to empirically benefit from hyperbolic embeddings in prior work . We compare different flow layers in the latent space of the VAEs. For the Euclidean VAE, we compare with RealNVP and CNF . Along with the two hyperbolic baselines used for density estimation, we also compare against the Tangent Coupling (TC) model in . As shown in Table 7.2, our continuous flow regime is more expressive and learns better than all hyperbolic baselines. In low dimensions, along with the other hyperbolic models, our approach tends to outperform Euclidean models on Omniglot and MNIST. However, in high dimension the hyperbolic models do not reap as much benefit; even the baseline HVAE does not consistently outperform the Euclidean VAE.

We have presented Neural Manifold ODEs, which allow for the construction of continuous time manifold neural networks. In particular, we introduce the relevant theory for defining “pure” Neural Manifold ODEs and augment it with our dynamic chart method. This resolves numerical and computational cost issues while allowing for better integration with modern Neural ODE theory. With this framework of continuous manifold dynamics, we develop Manifold Continuous Normalizing Flows. Empirical evaluation of our flows shows that they outperform existing manifold normalizing flow baselines on density estimation and variational inference tasks. Most importantly, our method is completely general as it does not require anything beyond local manifold structure.

We hope that our work paves the way for future development of manifold-based deep learning. In particular, we anticipate that our general framework will lead to other continuous generalizations of manifold neural networks. Furthermore, we expect to see application of our Manifold Continuous Normalizing Flows for topologically nontrivial data in lattice quantum field theory, motion estimation, and protein structure prediction.

As mentioned in the introduction, our method has applications to physics, robotics, and biology. While there are ethical and social concerns with parts of these fields, our work is too theoretical for us to say for sure what the final impact will be. For deep generative models, there are overarching concerns with generating fake data for malicious use (e.g. deepfake impersonations). However, our work is more concerned with accurately modelling data topology rather than generating hyper-realistic vision or audio samples, so we do not expect there to be any negative consequence in this area.

We would like to acknowledge Prof. Austin Benson and Junteng Jia for their insightful comments and suggestions. In addition, we would like to thank Facebook AI for funding equipment that made this work possible. We would also like to thank Joey Bose for providing access to his prerelease code.

Appendix A Proofs of Propositions

We see that if $\mathbf{z}(t)=\varphi_{z}(\mathbf{y}(t))$ then for all $t^{\prime}\in[\tau,\tau+\epsilon]$

A.2 Derivation of Gradient for Adjoint State

In this section we prove Theorem 4.1. The proof follows from the analogous one in , though we replace certain operations with their manifold counterparts.

We set $T_{\epsilon}(\mathbf{z}(t),t)\mathrel{\mathop{\mathchar 58\relax}}=\mathbf{z}(t+\epsilon)$. As in the original adjoint method derivation , we utilize the chain rule

Appendix B Computation of Manifold Continuous Normalizing Flows

There exist a variety of functions $f$ which constrain the log Jacobian determinant to be computationally tractable. An important one is the Continuous Normalizing Flow (CNF). CNFs construct $f$ to be the solution of an ODE . Explicitly, let $t_{0},t_{1}$ be starting and ending times with $t_{0}<t_{1}$, and consider the ordinary differential equation $\frac{d\mathbf{z}(t)}{dt}=g(\mathbf{z}(t),t;\theta)$, where $\theta$ parameterizes the dynamics $g$. For a sample from the base distribution $z\sim\pi$, solving this ODE with initial condition $\mathbf{z}(t_{0})=z$ gives the sample from the target distribution $x=\mathbf{z}(t_{1})$. The change in the log density given by this model satisfies an ordinary differential equation called the instantaneous change of variables formula:

We can therefore quantify the final probability as

B.2 Manifold Continuous Normalizing Flows

For our MCNF we split up the original time $[t_{s},t_{e}]$ into $[t_{i},t_{i+1}]$ for $i\in[k]$ where $t_{s}=t_{1}<t_{2}<\dots<t_{k}<t_{k+1}=t_{e}$. From our curve $\mathbf{z}$ we can select $z_{i}=\mathbf{z}(t_{i})$, and we have charts $\varphi_{i}\mathrel{\mathop{\mathchar 58\relax}}U_{i}\to V_{i}$ s.t. $z_{i},z_{i+1}\in V_{i}$. If our dynamics are determined by $\frac{d\mathbf{z}(t)}{dt}=f(\mathbf{z}(t),t;\theta)$ then this locally takes the form $\varphi_{i}(\widehat{f_{i}}(\varphi_{i}^{-1}(\mathbf{z}(t),t;\theta)))$, in which $\widehat{f_{i}}$ is a CNF. The update after passing through a chart $\varphi_{i}$ and integrating is given by

We can consider our manifold ODE as a composition of these updates. Therefore, we have that

For our cases we will be setting $\varphi_{i}=\exp_{z_{i}}$.

B.3 Base Distributions

Probability Density. The probability density can be found through the composition of the parallel transport map and exponential map. Specifically, we have that

Spherical Space. We could possibly perform a relatively similar wrapped normal distribution . However, we see that this is theoretically flawed since parallel transport between two conjugate points is not well defined.

Sampling. The von Mises-Fisher can be sampled from with efficient methods .

Probability Density. From , we know that the density is given by

Note that $||\mu||^{2}=1$, $\mathcal{C}_{n}(\kappa)$ is the normalizing constant, and that $\mathcal{I}_{v}$ denotes the modified Bessel function of the first kind of order $v$.

These baseline probability distributions are visualized in Figure 5.

B.4 Hyperbolic Space

In addition, there are many useful identities which appear in our pipeline.

Log Determinants. The log determinant of the derivative of the exponential map is given by :

The log determinant of the derivative of the log map is the negation of the above by the inverse function theorem, and the log determinant of the derivative of parallel transport is .

B.4.2 Numerical Stability

In order to ensure numerical stability, we examine several operations which are inherently numerically unstable and present solutions

Arccosh. Arccosh has a domain of $(-\infty,-1)\cup(1,+\infty)$. In practice, we are concerned with the positive case, although the negative case can be similarly handled. Due to numerical instability a value $1+\epsilon$ may be realized as $1$ in our floating point system. To compensate, we clamp the minimum value to be $1+\epsilon_{0}$ for a small fixed $\epsilon_{0}$.

Sinh Division. In the exponential and logarithmic maps, there exist terms of the form $\frac{\sinh(x)}{x}$. When $|x|<\epsilon$ for some small $\epsilon$, this is numerically unstable. We special case this (and the derivative) for stability by explicitly deriving the limit value of $x\rightarrow 0$ for these cases.

B.5 Spherical Space

Log Determinants. The log determinant for the exponential map of the Sphere is given by

The log determinant of the derivative of the log map is the negation of the above and the log determinant of the derivative of parallel transport is .

B.5.2 Numerical Stability

In order to ensure numerical stability, we note that several functions are inherently numerically unstable

Sine division. In the exp and log maps, there are values of the form $\frac{\sin x}{x}$. Note that this is numerically instable when $|x|<\epsilon$. We special case this (and its derivative) to allow for better propagation.

Log Derivative. We find that we need an explicit derivation of $D_{x}\log y$ for our Manifold ODE on the Sphere (see B.7). Note that this value can be computed using backpropagation, but we derive it explicitly instead due to numerical instability of the higher order derivatives of some of our functions.

For $x,y\in\mathcal{S}^{n}$, and $r=x^{T}y$, if $|r|\neq 1$, then

The limit of $D_{y}\log_{x}(y)$ as $r=x^{T}y\to 1$ is $I-xx^{T}$.

We differentiate the equation of the logarithmic map as given in Table 3. First, suppose that $\left|r\right|\neq 1$. By the product rule we have,

To compute the left summand, we use the chain rule and differentiate by $r$

To check that the limit as $x^{T}y\to 1$ is $I-xx^{T}$, it is enough to compute three separate limits that are all finite. It is clear that

Since the limit of the other term in the left summand can be shown to be finite, this means that the left summand is zero in the limit. For the right summand, the only term that depends on $y$ has a limit

where the final equality can be seen by L’Hopital’s rule. Thus, the entire limit is $I-xx^{T}$. ∎

B.6 Backpropagation

Differentiating the dynamics is done with standard backpropagation. The adjoint state $\frac{\partial L}{\partial\mathbf{y}(t)}$ is computed by the solution of the adjoint differential equation (2) for Euclidean space, with initial condition $\frac{\partial L}{\partial\mathbf{y}(t_{e})}$. The derivative of the loss $\frac{\partial L}{\partial\mathbf{y}(t_{e})}$ can be computed directly. For MCNF, $L$ is taken to be the negative log likelihood.

For the hyperbolic and spherical cases, the log determinant terms take simple forms (as in equations 25 and 27), and are thus easy to differentiate through. Moreover, for the hyperbolic VAE models, we train by maximizing the evidence lower bound (ELBO) on the log likelihood, so that differentiation is done with a reparameterization as in .

B.7 Designing Neural Networks

For the spherical case, we use the default construction (with projection), as there is no global diffeomorphism. Note that when passing from the manifold to tangent space dynamics, we require $D_{y}\log_{x}$. We also must invoke a radius of injectivity, as opposed to hyperbolic space. This is $\pi$ for all points.

B.7.2 Existence of a Solution

We construct our networks in such a way that the Picard-Lindelöf theorem holds. Our dynamics are given by $D_{z}\varphi\circ f$ where $\varphi$ is a chart and $f$ is a neural network. These are well behaved since 1) the neural network dynamics are well behaved using tanh and other Lipchitz nonlinearities and 2) the chart is well behaved since we can bound the domain to be compact.

Appendix C Experimental Details

In our code release, we will include functions to generate samples from the target densities that are used for density estimation in section 7.1.

Hyperbolic Density Estimation We detail the hyperbolic densities in each row of Figure 4.

The hyperbolic density in the first row of Figure 4 is a wrapped normal distribution with mean $(-1,1)$ and covariance $\frac{3}{4}I$.

The third density is a projection onto the hyperboloid of a uniform checkerboard density in $\mathcal{T}_{0}\mathcal{M}$. The square in the second row and third column of the checkerboard has its lower-left corner at the origin $(0,0)$. Each square has side length $1.5$.

The fourth density is a mixture of four wrapped normal distributions. Letting $s=1.3$, $\sigma_{1}=.3$ and $\sigma_{2}=1.5$, the wrapped normals are given as:

Spherical Density Estimation Details are given about the densities that were learned in each row of Figure 4.

The density in the first row of Figure 4 is a wrapped normal distribution with mean $\frac{1}{\sqrt{3}}(-1,-1,-1)$ and covariance $\frac{3}{10}I$.

The second density is built from a mixture of 4 wrapped normals, with means $\frac{1}{\sqrt{3}}(1,1,1),\frac{1}{\sqrt{3}}(-1,-1,-1),\frac{1}{\sqrt{3}}(-1,-1,1)$, and $\frac{1}{\sqrt{3}}(1,1,-1)$; all components of the mixture have covariance $\frac{3}{10}I$.

The third density is a uniform checkerboard density in spherical coordinates $(\varphi,\theta)\in[0,2\pi]\times[0,\pi]$. The rectangle in the second row and third column of the checkerboard has its lower-left corner at $(\pi,\pi/2)$. The side length of each rectangle in the $\varphi$-axis is $\pi/2-0.2$, and the side length in the $\theta$-axis is $\pi/4-0.1$.

Variational Inference For variational inference, we dynamically binarize the MNIST and Omniglot images with the same procedure as given in . We resize the Omniglot images to $28\times 28$, the same size as the MNIST images.

C.2 Density Estimation

In section 7.1 we train on batches of 200 samples from the target density (or batches of size 100 for the discrete spherical normalizing flows ). Our MCNF models and the hyperbolic baselines use at most 1,000,000 samples, with early stopping as needed—the hyperbolic baselines sometimes diverge when training for too many batches. As suggested in , we find that the spherical baseline does indeed needed more samples than this (at least 5,000,000), so we allow it to train until the density converges. Although we do not investigate sample efficiency in detail, we find that our MCNF is able to achieve better results than the spherical baseline with, frequently, over an order of magnitude fewer samples than the spherical baseline. Note that all methods use the Adam optimizer . For our MCNF, our dynamics are given by a neural network of hidden dimension 32 and 4 linear layers with tanh activation; for each integration we use a Runge-Kutta 4 solver.

Hyperbolic Normalizing Flows For the hyperbolic discrete normalizing flows, we train with 4 hidden blocks, hidden flow dimension of 32, and tanh activations. The prior distributions used are given in section B.3 and target distributions are given in section C.1.

The prior distributions used for all target distributions are given in section B.3. For the first spherical target distribution in section C.1, we use $k=2$ and $n=1$. For the second spherical target distribution we use $k=6$ and $n=2$. For the final spherical target distribution we use $k=32$ and $n=12$.

C.3 Variational Inference

Appendix D Dynamic Chart Method

The Benefit of Dynamic Charts. Note that our dynamic chart trivialization is performed in the main paper simply by splitting the time interval $[t_{s},t_{e}]$ up uniformly into segments of length $\epsilon$ and learning the solution locally via $\exp$-map charts at the anchor points (endpoints of the segments). Such a splitting allows for the ball of injectivity (induced by the radius of injectivity) to “move" throughout the training process, so that it always surrounds the locally relevant region (centered at the anchor point). This approach allows for transportation of mass that evades the conjugate point problem (which does not allow for).

This becomes especially clear if we consider the case of conjugate points on a sphere, i.e. the case of antipodal points. Consider a task in which we have a prior with mass surrounding one antipodal point and the target density is concentrated around the other point. A one chart approach would have trouble transporting the mass due to the fact that a fixed $\exp$-map can never have a ball of injectivity (induced by the radius of injectivity $\pi$, in this case) that includes both points throughout training. However, our dynamic approach allows this ball of injectivity to shift throughout the training process and makes correct transportation of mass for such a scenario possible. An experiment testifying to this is given in Appendix D.2 (and Figure 7.2).

Non-uniform Time-domain Segmenting. While our approach allows for the ball of injectivity to shift, the dynamic chart method is not limited to a uniform $\epsilon$-segment splitting of the $[t_{s},t_{e}]$ interval. Similar benefits may be derived with an alternative splitting. However, it is not clear what additional benefits a non-uniform splitting might bring without prior knowledge of local manifold topology. Our experiments find that a uniform split works well for the densities we tested.

D.2 Expressivity and Generality of MCNF

Appendix E Generated MNIST Samples

Here, we generate samples from MNIST using our trained hyperbolic MCNF. To do this, we use an MCNF with the same setup as in C.3, except with a slightly larger VAE architecture. We add a linear layer to both the mean and variance networks, add an additional shared linear layer for both of them, and add a linear layer to the decoder. MNIST samples generated with our approach are given in Figure 7.2. With a latent space of dimension 2, the MCNF generates examples that resemble real digits. Interpolating between points in the latent space gives hybrid intermediaries that meaningfully represent semantic change (for instance, going between a “4" and a “7" produces instances of “9").