Source-linked AI summary
Normalizing Flows on Tori and Spheres
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racanière, Michael S. Albergo, Gurtej Kanwar, Phiala E. Shanahan, Kyle Cranmer
TL;DR
Existing normalizing flows primarily target Euclidean data, leaving applications with circular, toroidal, or spherical geometry underserved. This paper constructs and compares recursively built flows for compact connected manifolds, showing expressive and numerically stable models with exact density evaluation and sampling. The study also identifies practical trade-offs among circle and sphere flow constructions and boundaries for extending the approach to other manifolds.
Problem
Existing normalizing flows mainly use Euclidean spaces, whereas applications involving angles, directions, and periodic locations require non-Euclidean geometries such as circles, tori, and spheres.
Method
The paper constructs flows recursively on circles, tori, spheres, and products of these compact connected manifolds, including non-compact projection, recursive, and exponential-map approaches.
Results
The proposed flows provide arbitrarily flexible models with tractable and exact density evaluation and sampling algorithms, while avoiding numerical instabilities of previous approaches.
Takeaways & Limitations
Circle-flow choice depends on use: CS is preferable when both density evaluation and sampling are required, whereas Möbius or NCP may be preferable otherwise.
Takeaways & Limitations
Extending the recursive construction to SU(D) and U(D) requires unresolved details such as appropriate boundary conditions for explicit coordinate systems.
Abstract
from arXiv · showhide
Normalizing flows are a powerful tool for building expressive distributions in high dimensions. So far, most of the literature has concentrated on learning flows on Euclidean spaces. Some problems however, such as those involving angles, are defined on spaces with more complex geometries, such as tori or spheres. In this paper, we propose and compare expressive and numerically stable flows on such spaces. Our flows are built recursively on the dimension of the space, starting from flows on circles, closed intervals or spheres.
1. Introduction
Normalizing flows are well suited to flexible density modelling, but standard Euclidean formulations do not match data living on circles, tori, spheres, or other non-Euclidean spaces. The paper proposes expressive flows for compact connected manifolds that avoid numerical instabilities affecting earlier projection-based approaches.
- Background: Normalizing flows define complex distributions by transforming samples from a simple base density through an invertible, differentiable map.The diffeomorphism constraint enables target-density evaluation through the change-of-variables formula.
- Motivation: Euclidean flows are unsuitable when variables have intrinsic geometries such as circular angles or periodic particle locations on a torus.The paper identifies circles and three-dimensional tori as natural domains for these examples.
- Motivation: Non-Euclidean probabilistic modelling arises in applications involving angles, axes, and directions, including protein structure, geology, and robotics.Such spaces are also used as latent spaces in variational autoencoders.
- Prior limitations: Manifold-to-manifold flows require diffeomorphic source and target spaces, while naive projection methods can introduce singularities and numerical instability.Exponential-map approaches on Lie groups are also difficult to compose because their maps are not from the manifold to itself.
- Contribution: The paper constructs arbitrarily flexible flows on circles, tori, spheres, and their products for compact connected differentiable manifolds.The flows are intended for settings where the manifold structure is known and are tested on synthetic sharp and multimodal densities.
2. Methods
The paper builds flows on compact manifolds recursively from circle, interval, and sphere transformations, emphasizing expressivity, invertibility, and numerical stability. It presents recursive and exponential-map sphere constructions alongside three circle diffeomorphisms, with distinct flexibility and computational trade-offs.
- Recursive construction: Flows on tori and spheres are constructed recursively from transformations on circles, closed intervals, and lower-dimensional spheres.Tori use Cartesian products of circles, while spheres use cylindrical coordinates based on S^(D−1) × [−1,1].
- Circle flows: Circle flows include Möbius transformations, circular splines, and non-compact projections, each designed to satisfy circle-specific diffeomorphism conditions.These conditions enforce consistent endpoint mapping, strict monotonicity, and matching Jacobians at the identified boundary.
- Circle flows: Circular splines are arbitrarily flexible as the segment count K increases and can be inverted exactly using binary search and an analytic quadratic solution.This avoids the convex combinations required to increase Möbius expressivity.
- Recursive construction: The recursive sphere flow composes cylindrical-to-spherical, autoregressive cylindrical, and spherical-to-cylindrical transformations, with finite and efficiently computed density updates.The construction is invertible almost everywhere because the cylinder and sphere have different topologies, while spline boundary conditions guarantee finite density.
- Exponential-map flows: Exponential-map flows use the exponential map of a sphere’s gradient field, which is guaranteed to be a diffeomorphism and can incorporate a radial scalar field.The radial field addresses difficulty learning concentrated multimodal spherical densities with polynomial and high-frequency fields.
- Exponential-map flows: For S2, the recursive flow outperforms exponential-map alternatives: the latter improve substantially with the proposed radial field but remain noncompetitive.The exponential-map density update is practical only for small D because its computational cost scales poorly.
3. Related Work
Existing approaches define distributions on tori and spheres through wrapping, projecting, conditioning, Lie-group exponential maps, or embedded coordinate charts. These methods offer tractability or composability in restricted settings but face flexibility, computational, topological, or numerical limitations on compact manifolds.
- Traditional approaches: Wrapping, projecting, and conditioning are the traditional strategies for defining distributions on tori and spheres.Their density evaluation and sampling are tractable only in special cases, which typically produce simple distributions with limited flexibility.
- Traditional approaches: Wrapped distributions can require an analytically intractable infinite sum, while projected distributions may require an intractable integral over the radial coordinate.These issues arise respectively when the Euclidean base has unbounded support and when probability mass is projected from an embedding space onto the manifold.
- Traditional approaches: Conditioning embedded distributions on a manifold generally requires integrating over the manifold to compute a normalizing constant.Examples include multivariate von Mises distributions obtained by conditioning Gaussian densities on tori.
- Flexible alternatives: Mixtures can increase the flexibility of simple base distributions, but may need many components for sharp complex distributions and can be harder to fit than compositional flows.The paper notes that mixtures could serve as base distributions for flows on tori and spheres.
- Lie-group methods: Lie-group exponential-map methods generalize wrapping but are difficult to compose and can be expensive in high dimensions.For tori, the number of terms in the wrapped sum scales exponentially with dimension even when the sum is finitely truncated.
- Embedded-coordinate methods: Embedded-coordinate methods produce composable transformations but are better suited to manifolds homeomorphic to R^D; on compact tori and spheres, undefined global charts can cause numerical instabilities.The coordinate chart cannot be defined everywhere on compact manifolds that are not homeomorphic to R^D.
4. Experiments
The experiments evaluate flows on tori and spheres against sharp, correlated, and multi-modal targets using ESS, KL, and learned-density visualizations. Recursive flows learn the tested torus and sphere targets, while circular-flow performance and inversion trade-offs differ by transformation.
- Evaluation setup: The experiments vary inverse temperature β to test increasingly concentrated targets, using KL minimization and ESS as a target-matching evaluation.Higher ESS indicates a better match when reliably estimated; ESS is reported as a percentage of the actual sample size.
- Experiments on T2: T2 experiments cover unimodal, three-component multi-modal, and correlated von Mises targets across different concentration levels.The reported results show that multi-modal and correlated densities can be learned reliably.
- Experiments on S2: On S2, the recursive Möbius-spline flow is compared with exponential-map polynomial and radial flows on a four-mode target using KL and ESS.The experiment also reports a substantial improvement over the polynomial scalar field, while exponential-map flows remain less competitive with the recursive flow.
- Experiments on T2: NCP and Möbius performed best among circle flows, whereas CS performed less well for highly concentrated targets.The comparison is based on the T2 experiments across varying β values.
- Experiments on SU(2) ≅ S3: A recursive flow learns a density on SU(2) ≅ S3, with the reported example achieving ESS = 84% and KL = 0.14.The model uses a Möbius transform for the circle and spline transforms for the two intervals.
5. Discussion
The proposed flows are flexible and numerically stable on tori and spheres of any dimension, with exact density evaluation and sampling. The discussion emphasizes practical trade-offs among circle transformations and the recursive flow’s computational advantage on spheres, while identifying extensions to higher-dimensional unitary groups as future work.
- Scope and contribution: The proposed flows on tori and spheres can be made arbitrarily flexible while retaining tractable and exact density evaluation and sampling.The construction applies to compact connected manifolds including circles, tori, spheres, and products of these spaces.
- Comparison, scope and limitations: For circle flows, Möbius and NCP perform best, while CS is preferable when both density evaluation and sampling are required because it is analytically invertible.Möbius and NCP with multiple components require root-finding for inversion, whereas CS gains expressivity by adding spline segments.
- Comparison, scope and limitations: On S^D, the recursive flow performs better than the exponential-map flow and scales better because its density can be computed efficiently.The exponential-map flow’s intrinsic formulation did not provide a practical benefit in the reported experiments.
- Towards higher-dimensional unitary groups: The construction demonstrates expressive flows on U(1) ≅ S1 and SU(2) ≅ S3, motivated by the relevance of these unitary Lie groups to fundamental physics.The authors identify higher-dimensional SU(D) and U(D) flows as a direction for future work because their coordinate systems and boundary conditions are more involved.
A. Density Transformations on Manifolds
The section derives density transformations between embedded manifolds using tangent-space volume changes, then specializes them to sphere–cylinder maps used in recursive flows. Although the coordinate maps are only almost-everywhere diffeomorphic, paired transformations cancel boundary divergences and preserve finite densities.
- General manifold transformations: Density transformation on an embedded manifold is computed from the source density and the tangent-map volume factor.The factor uses an orthonormal tangent basis E and the ambient Jacobian J through det((JE)ᵀJE).
- General manifold transformations: When M=N=R^D, the manifold formula reduces to the familiar Euclidean density update.
- Sphere specialization: For the sphere, the map T_c→s transforms S^{D−1} × (−1,1) into S^D, with its Jacobian evaluated using cylindrical coordinates and rotational symmetry.The proof simplifies the determinant by rotating points to a canonical form and choosing a convenient tangent basis.
- Numerical stability: Boundary terms that individually diverge cancel when T_c→s and T_s→c are combined with the interval spline, yielding a well-behaved finite-density flow.Implementation achieves stability by omitting the terms that cancel; continuity of the resulting density is not guaranteed.
- Recursive construction: The recursive sphere construction wires conditional maps according to the recursion, while other variable orders or compositions of autoregressive flows are also permitted.
D. Fourier Transformations on S1
This section introduces Fourier transformations for circle flows and describes their frequency parameterization. Empirically, the family is not competitive with the other transformations, particularly for highly concentrated densities.
- Fourier transformations: Fourier transformations provide an alternative family of circle transformations parameterized by fixed frequencies in the Fourier basis.
- Recursive sphere flow: Figure 6 illustrates the recursive sphere flow by transforming the sphere to S^1 × [−1,1]^{D−1}, applying an autoregressive cylinder flow, and mapping back.
- Empirical behavior: The Fourier transformation family is not competitive with the other considered transformations, especially for highly concentrated densities.
- Comparison model: The polynomial exponential-map flow uses a scalar field with parameters constrained by ∥µ∥_1 + ∥A∥_1 ≤ 1.
F. Target Densities Used in Experiments
The experiments use structured target densities on T^2, S^2, and SU(2) ≅ S^3, including von Mises targets, multimodal sphere mixtures, and illustrative circle densities formed from Möbius transformations.
- Torus targets: Experiments on the torus T^2 use target densities from the von Mises family.
- Sphere targets: The S^2 target is a mixture with four specified component locations in the spherical experiment.
- Circle targets: Figure 7 shows densities formed by convex combinations of 15 Möbius transformations applied to a uniform base on S^1.Each distribution uses 30 = 15 × 2 parameters.
- Model comparisons: Figure 8 compares the corresponding models while adding KL and Fourier-transform values; bracketed numbers identify used frequencies and repetitions.
- Sphere targets: The SU(2) ≅ S^3 target is a four-component mixture with specified locations in the embedded R^4 coordinates.
G. Misaligned Density on S2
The paper tests whether axis choices in the recursive sphere coordinates create artifacts by comparing an unrotated flow with the same flow composed with a learned rotation. On S^2, the two model sizes show no significant difference, while higher-dimensional effects remain open.
- Experimental design: The experiment probes possible axis-related artifacts using a nonaxis-aligned ring target and a learned rotation comparison.
- Results: No significant differences were observed between large (K_s = 32, K_m = 12) and small (K_s = 3, K_m = 3) autoregressive Möbius-Spline flows on S^2.
- Scope: More experiments are needed to investigate the potential axis-choice effect in higher dimensions.
H. NCP as a complex M¨obius transformation
The paper shows that a constrained Möbius transformation on the unit circle can reproduce an NCP transformation, linking the two parameterizations through trigonometric algebra.
- A Möbius transformation defines a diffeomorphism on S1 only under a constrained parameter form.The constraint preserves the unit circle and leaves two real-valued free parameters, ℜ(a) and ℑ(a).
- Choosing ℑ(a) = 0 and ℜ(a) = −1−α yields an NCP transform with scale parameter α and offset β = 0.The equivalence is stated for w, θ ∈ (−π, π), with z = e^iθ.
- The derivation expands the Möbius expression into trigonometric quantities and isolates w through the numerator v.The final algebra uses tan(2x) = 2 tan(x)/(1−tan(x)^2).
- The construction relies on the fact that only the ratio of the imaginary and real parts of the relevant expression determines tan(w).This ratio-based step connects the complex transformation to the angular NCP rule.
I. Application: Multi-Link Robot Arm
The paper applies torus flows to infer the posterior over six joint angles of a planar robot arm subject to a multimodal tip-position constraint.
- The robot arm’s six joint angles form points in T6, providing a concrete torus-valued inference problem.The application approximates the posterior density over θ1,...,6.
- The experiment uses an autoregressive Möbius flow on the torus to learn the posterior density over the robot’s joint angles.Figure 10 illustrates the learned posterior for the six-link arm.
- The tip-position constraint is modeled by a two-component Gaussian-mixture likelihood, while the prior over joint angles is uniform on T6.The link lengths are lk = 0.2 and the arm is affixed at r0 = (0, 0).
J. Application: Learning from samples
The paper demonstrates flow training from direct samples on a sphere and examines model behavior on spherical and toroidal density-learning tasks.
- Unlike most experiments, this setting trains the flow directly from data samples rather than a target density known only up to normalization.The sample-based objective is maximum likelihood and requires an explicit inverse map.
- The spherical dataset follows a density shaped like Earth’s continental map, creating a nontrivial target on S2.The construction is intended to test whether inverse-map computation remains feasible for a complex spherical density.
- The sample-trained model stacks two autoregressive flows using circular splines and interval splines, with Ks = 80 segments and 100,000 training steps.Samples and a heat map of the learned density are shown in Figure 11.
- On S2, composing a Möbius-spline flow with a learnable rotation produced no substantial performance differences from the unrotated alternatives.The comparison includes full and small Möbius-spline flows, with Ks = 32, Km = 12 for MS and Ks = 3, Km = 3 for SMS.
- The torus and sphere applications illustrate flows learning complex densities in angular spaces, including robot-arm posteriors and sample-based spherical targets.The robot-arm experiment models a six-angle posterior on T6, while the sample-based figure compares target, model samples, and learned density.