Normalizing Flows on Tori and Spheres

Introduction

Normalizing flows are a flexible way of defining complex distributions on high-dimensional data. A normalizing flow maps samples from a base distribution $\pi(u)$ to samples from a target distribution $p(x)$ via a transformation $f$ as follows:

The transformation $f$ is restricted to be a diffeomorphism: it must be invertible and both $f$ and its inverse $f^{-1}$ must be differentiable. This restriction allows us to calculate the target density $p(x)$ via a change of variables:

In practice, $\pi(u)$ is often taken to be a simple density that can be easily evaluated and sampled from, and either $f$ or its inverse $f^{-1}$ are implemented via neural networks such that the Jacobian determinant is efficient to compute.

A normalizing flow implements two operations: sampling via Equation 1, and evaluating the density via Equation 2. These operations have distinct computational requirements: generating samples and evaluating their density requires only $f$ and its Jacobian determinant, whereas evaluating the density of arbitrary datapoints requires only $f^{-1}$ and its Jacobian determinant. Thus, the intended usage of the flow dictates whether $f$, $f^{-1}$ or both must have efficient implementations. For an overview of various implementations and associated trade-offs, see (Papamakarios et al., 2019).

The need for probabilistic modelling of non-Euclidean data often arises in applications where the data is a set of angles, axes or directions (Mardia & Jupp, 2009). Such applications include protein-structure prediction in molecular biology (Hamelryck et al., 2006; Mardia et al., 2007; Boomsma et al., 2008; Shapovalov & Dunbrack Jr, 2011), rock-formation analysis in geology (Peel et al., 2001), and path navigation and motion estimation in robotics (Feiten et al., 2013; Senanayake & Ramos, 2018). Non-Euclidean spaces have also been explored in machine learning, and specifically generative modelling, as latent spaces of variational autoencoders (Davidson et al., 2018; Falorsi et al., 2018; Wang & Wang, 2019; Mathieu et al., 2019; Wang et al., 2019).

Methods

Since and $2\pi$ are identified as the same point, the transformation $f:[0,2\pi]\rightarrow[0,2\pi]$ must satisfy appropriate boundary conditions to be a valid diffeomorphism on the circle. The following conditions are sufficient:

The first two conditions ensure that and $2\pi$ are mapped to the same point on the circle. The third condition ensures that the transformation is strictly monotonic, and thus invertible. Finally, the fourth condition ensures that the Jacobians agree at and $2\pi$, thus the probability density is continuous.

A restriction in the above conditions is that and $2\pi$ are fixed points. Nonetheless, this restriction can be easily overcome by composing a transformation $f$ satisfying these conditions with a phase translation $\theta\mapsto(\theta+\phi)\!\!\mod 2\pi$, where $\phi$ can be a learnable parameter. Such a phase translation is volume-preserving, so it does not incur a volume correction in the calculation of the probability density.

Given a collection $\left\{f_{i}\right\}_{i=1,\ldots,K}$ of transformations satisfying the above conditions, we can combine them into a more complex transformation $f$ that also meets these conditions. One such mechanism for combining transformations is function composition $f=f_{K}\circ\cdots\circ f_{1}$, which can easily be seen to satisfy Equations 3, 5, 4 and 6. Alternatively, we can combine transformations using convex combinations, as any convex combination defined by

Next, we describe three circle diffeomorphisms that by construction satisfy the above conditions: Möbius transformations, circular splines, and non-compact projections.

We define the Möbius transformation $h_{\omega}(z)$ of $z$ with centre $\omega$ to be $-g_{\omega}(z)$.

An explicit formula for $h_{\omega}$ is given by

When $\omega=0$, the transformation $h_{\omega}$ is just the identity. In general, the variables $z$ and $\omega$ in Equation 8 are meant to be $(D+1)$-dimensional real vectors. When $D=1$, the equation also reads correctly if $z$ and $\omega$ are taken to be complex numbers. In this case, Equation 8 and the Möbius transformations of the complex plane are related as explained in formula (6) of Kato & McCullagh (2015).

The transformation $h_{\omega}$ expands the part of the sphere that is close to $\omega$, and contracts the rest. Hence, it can transform a uniform base distribution into a unimodal smooth distribution parameterized by $\omega$. One property of the Möbius transformation is that it does not become more expressive by composing various $h_{\omega}$. This is because the set of transformations $Rh_{\omega}$, where $R$ can be any matrix in $\text{SO}(D+1)$, forms a group under function composition (see Theorem 2 of Kato & McCullagh, 2015). Since composition of two such transformations remains a member of the group, their expressivity is not increased.

1.2 Circular Splines (CS)

where the coefficients $\left\{\alpha_{ki},\beta_{ki}\right\}_{i=0,1,2}$ are chosen so that $f$ is strictly monotonically increasing and $f(x_{k-1})=y_{k-1}$, $f(x_{k})=y_{k}$, and $\nabla f(\theta)|_{\theta=x_{k}}=s_{k}$ for all $k=1\ldots,K$ (see Durkan et al., 2019b, for more details).

We can easily restrict $f$ to be a diffeomorphism from $[0,2\pi]$ to itself, by setting $x_{0}=y_{0}=0$ and $x_{K}=y_{K}=2\pi$. This construction satisfies the first three sufficient conditions in Equations 3, 4 and 5, and can be used to define probability densities on the closed interval $[0,2\pi]$. In addition, by setting $s_{1}=s_{K}$, we satisfy the fourth condition in Equation 6, and hence we obtain a valid circle diffeomorphism which we refer to as a circular spline (CS).

Circular splines can be made arbitrarily flexible by increasing the number of segments $K$. Therefore, unlike Möbius transformations, it is not necessary to combine them via convex combinations to increase their expressivity. An advantage of circular splines is that they can be inverted exactly, by first locating the corresponding segment (which can be done in $\mathcal{O}\mathopen{}\left(\log K\right)\mathclose{}$ iterations using binary search since the segments are sorted), and then inverting the corresponding rational quadratic (which can be done analytically by solving a quadratic equation).

1.3 Non-Compact Projection (NCP)

Even though the expression for $f$ is not defined at the endpoints and $2\pi$, the expression for the gradient is. The transformation satisfies the appropriate boundary conditions,

Therefore, we can extend $f$ to $[0,2\pi]$ by continuity such that $f(0)=0$ and $f(2\pi)=2\pi$, which yields a valid circle diffeomorphism. Although not immediately obvious, NCP flows and Möbius transformations are intimately related, as explained in Appendix H (see also Downs & Mardia, 2002, Section 2.1).

The above boundary conditions are satisfied when the transformation $g$ is affine, but they are not generally satisfied when $g$ is an arbitrary diffeomorphism. This limits the type of flow we can put on the non-compact space. Therefore, instead of making $g$ more expressive, we choose to increase the expressivity of the NCP flow by combining multiple transformations $f$ via convex combinations.

A potential issue with NCP is that, near the endpoints and $2\pi$, evaluating $f$ using Equation 10 directly is numerically unstable. To circumvent this numerical difficulty, we can use equivalent linearized expressions when $\theta$ is near the endpoints. For example, for $\theta$ close to we can approximate $f(\theta)\approx\theta/\alpha$, whereas for $\theta$ close to $2\pi$ we have $f(\theta)\approx 2\pi+(\theta-2\pi)/\alpha$.

The parameters $\psi_{i}$ of the $i$-th transformer are a function of $(\theta_{1},\ldots,\theta_{i-1})$ known as the $i$-th conditioner. In order to guarantee that the conditioners are periodic functions of each $\theta_{i}$, we can make $\psi_{i}$ be a function of $(\cos{\theta_{1}},\sin{\theta_{1}},\ldots,\cos{\theta_{i-1}},\sin{\theta_{i-1}})$ instead. In our experiments, we implemented the conditioners using coupling layers (Dinh et al., 2017). Implementations based on masking (Kingma et al., 2016; Papamakarios et al., 2017) are also possible.

More generally, autoregressive flows can be applied in the same way on any manifold that can be written as a Cartesian product of circles and intervals, such as the $2$-dimensional cylinder. Flows on intervals can be constructed e.g. using regular (non-circular) splines as described in Section 2.1.2. Thus, by taking $f_{\psi_{i}}$ to be either a circle diffeomorphism or an interval diffeomorphism as required, we can handle arbitrary products of circles and intervals.

Finally, the cylinder is transformed back to the sphere by

The update due to $T_{s\rightarrow c}$ is simply the inverse of the above. Additionally, in Section A.1 we prove that the combined density update due to $h=T_{c\rightarrow s}\circ T_{c\rightarrow c}\circ T_{s\rightarrow c}$ does not have any singularity whenever $g$ is such that $g^{\prime}(-1)>0$ and $g^{\prime}(1)>0$. Since we implement $g$ with spline flows (Section 2.1.2), this condition is satisfied by construction, so the flow density is guaranteed to be finite.

3.2 Exponential-Map Flows

Ideally, we would like to specify flows directly on the manifold and avoid mapping between non-diffeomorphic sets. This motivates us to explore exponential-map flows, a mechanism for building flows on spheres proposed by Sei (2013).

satisfy the required conditions, therefore the map

As shown in Appendix A, the density update due to $f$ is

Related Work

The study of distributions on objects such as angles, axes and directions has a long history, and is known as directional statistics (Mardia & Jupp, 2009). According to the taxonomy of Navarro et al. (2017), directional statistics traditionally uses three approaches for defining distributions on tori and spheres: wrapping, projecting, and conditioning.

In general, the above three strategies lead to distributions with tractable density evaluation and sampling algorithms only in special cases, which typically yield simple distributions with limited flexibility. One approach for increasing the flexibility of such simple distributions is via combining them into mixtures (see e.g. Peel et al., 2001; Mardia et al., 2007). Such mixtures could be used as base distributions for the flows on tori and spheres that we present in this work. However, unlike flows whose expressivity increases via composition, mixtures generally require a large number of components to represent sharp and complex distributions, and can be harder to fit in practice.

Experiments

We evaluate and compare our proposed flows based on their ability to learn sharp, correlated and multi-modal target densities. We used targets $p(x)=e^{-\beta u(x)}/Z$ with inverse temperature $\beta$ and normalizer $Z$. We varied $\beta$ to create targets with different degrees of concentration/sharpness. The models $q_{\eta}$ were trained by minimizing the KL divergence

For an additional evaluation of how well the flows match the target, we used $S=20{,}000$ samples $x_{s}$ from the trained flows to estimate the normalizer via importance sampling:

where $w_{s}=e^{-\beta u(x_{s})}/q_{\eta}(x_{s})$. The effective sample size, (Arnaud et al., 2001; Liu, 2008), can be estimated by

Higher ESS indicates that the flow matches the target better (when reliably estimated). We report ESS as a percentage of the actual sample size.

All models were optimized using Adam (Kingma & Ba, 2015) with learning rate $2\times 10^{-4}$, $20{,}000$ iterations, and batch size $256$. The reported error bars are the standard deviation of the average, computed from independent replicas of each experiment with identical hyper-parameters, but with different initialization of the neural network weights. All shown model densities are kernel density estimates using $20{,}000$ samples and Gaussian kernel with bandwidth $0.2$.

Discussion

This work shows how to construct flexible normalizing flows on tori and spheres of any dimension in a numerically stable manner. Unlike many of the distributions traditionally used in directional statistics, the proposed flows can be made arbitrarily flexible, but have tractable and exact density evaluation and sampling algorithms. We conclude with a comparison of the proposed models, a discussion of their limitations, and some preliminary thoughts on how to extend flows to other manifolds of interest to fundamental physics.

Among the flows on the circle, Möbius and NCP performed the best, with CS performing less well for highly concentrated target densities. However, increasing the expressivity of Möbius and NCP required convex combinations, whereas CS can be made more expressive by adding more spline segments. As a result, CS is the cheapest to invert (it can be done analytically), whereas Möbius and NCP (with more than one component) require a root-finding algorithm such as bisection search. Therefore, in practice it may be preferable to use CS if both density evaluation and sampling are required, and use Möbius or NCP otherwise.

2 Towards Normalizing Flows on SU​(D)SU𝐷\text{SU}(D) and U​(D)U𝐷\text{U}(D)

Acknowledgements

We thank Heiko Strathmann and Alex Botev for discussions and feedback. PES is partially supported by the National Science Foundation under CAREER Award 1841699 and PES and GK are partially supported by the U.S. Department of Energy, Office of Science, Office of Nuclear Physics under grant Contract Number DE-SC0011090. KC is supported by the National Science Foundation under the awards ACI-1450310, OAC-1836650, and OAC-1841471 and by the Moore-Sloan data science environment at NYU.

References

Appendix A Density Transformations on Manifolds

In this section, we explain how to update the density of a distribution transformed from one Riemannian manifold to another by a smooth map. We only consider the case where both manifolds are sub-manifolds of Euclidean spaces.

Since for any $x\in\mathcal{M}$, we can always choose $G$ such that $Gx$ is of the form $(\sqrt{1-r^{2}},0,\ldots,0,r)$, we can restrict ourselves to this case. For such a choice, the Jacobian simplifies to

For $E$, we can simply choose the $D+1$ by $D$ matrix made by removing the first column from the identity matrix. Then $JE$ is equal to $J$ with the first column removed:

The product $(JE)^{\top}JE$ is simply a diagonal matrix of size $D$ with diagonal $(1-r^{2},\ldots,1-r^{2},\frac{1}{1-r^{2}})$. Taking the determinant and applying Equation 24 concludes the proof. ∎

At first, 1 might seem worrying since the density ratio in that proposition vanishes when $r$ is $-1$ or $1$. So, as $r$ approaches the boundary of the interval $$, it seems that the correction term to the density will tend to infinity and lead to numerical instability.

As $\epsilon$ goes to , the density corrections coming from $T_{c\rightarrow s}$ and $T_{s\rightarrow c}$ combine to

as $\epsilon$ goes to . In particular, the terms that tend to infinity cancel each other, and the flow is well-behaved. When implementing the flow, numerical stability is achieved by not adding the terms that cancel each other. Finally, we note a subtle point about what we proved: the sequence of transformations $T_{c\rightarrow s}\circ T_{c\rightarrow c}\circ T_{s\rightarrow c}$ will transform a distribution with finite density into another distribution with finite density, but we do not guarantee that the resulting density will be continuous.

In Figure 6, we provide an illustration of the recursive construction in Equations 12, 13 and 14, showing the specific wiring order of the conditional maps inside the flow. This order is the one implied by the recursion. In general, any other order can be used, or a composition of autoregressive flows with multiple orders.

Appendix C Examples of Möbius Transformations

Another family of circle transformations that we considered are Fourier transformations, defined by

We found empirically that this family of transformations is not competitive with the other transformations considered in this paper, especially for highly concentrated densities as shown in Figure 8.

Appendix E Polynomial Exponential Map

The polynomial exponential map of Sei (2013) is the exponential-map flow built using the scalar field

Appendix F Target Densities Used in Experiments

The recursive formulas shown in Equations 12, 13 and 14 require choosing a sequence of axes in order to construct the cylindrical coordinate system. This may introduce artifacts to the density related to this choice of axes. To test if this results in numerical problems, we compare the flow from Equations 12, 13 and 14 on a target density that forms a non-axis-aligned ring against a composition of the same flow with a learned rotation.

More experiments would be necessary to investigate this potential effect in higher dimensions.

Appendix H NCP as a complex Möbius transformation

This form ensures that $h(z)\bar{h}(z)=1$ if $z\bar{z}=1$ and has two real-valued free parameters $\Re(a)$ and $\Im(a)$.

In what follows we show that for the choice $\Im(a)=0$ and $\Re(a)=-\frac{1-\alpha}{1+\alpha}$, the transformation $h(z)$ is equivalent to an NCP transform $w=2\arctan\left(\alpha\tan\left(\frac{\theta}{2}\right)+\beta\right)$ with scale parameter $\alpha$ and offset parameter $\beta=0$ (assuming $w,\theta\in(-\pi,\pi)$). If we define $\theta$ via $z=e^{i\theta}$, the goal is to show that $w$ defined via

We begin by expanding Equation 31 in terms of more basic trigonometric quantities,

In order to isolate $w$, only the numerator $v=2\left(\cos\left(\frac{\theta}{2}\right)+i\alpha\sin\left(\frac{\theta}{2}\right)\right)^{2}$ of the expression above matters as we are only interested in ratios of the imaginary and real parts of this expression, $\tan(w)=\frac{\Im(v)}{\Re(v)}$. The numerator can be expanded as

Using the trigonometric formula $\tan(2x)=\frac{2\tan(x)}{1-\tan(x)^{2}}$, we arrive at the final result

Appendix I Application: Multi-Link Robot Arm

Appendix J Application: Learning from samples

In most of the experiments shown on this paper, we trained the models to fit a target density known up to a normalization constant (i.e. an inference problem). In this experiment we train our flow directly on data samples instead.

We trained a flow built from stacking two autoregressive flows. Each flow in the stack used circular splines and standard splines on the interval. The model was trained to maximize the likelihood of the dataset for $100{,}000$ training steps. Both splines used $K_{s}=80$ segments. The neural networks producing the spline parameters are the same as for the other experiments. In Figure 11 (middle) we show samples from the learned model overlaid on Earth’s map and in Figure 11 (right) we show a heat map of the learned density.