Equivariant Manifold Flows

1 Introduction

• Manifold density learning addresses non-Euclidean data relevant to quantum field theory [44], robotics, and protein-structure prediction.Many natural-science applications also require symmetries that manifold CNFs cannot properly model.Source: §1 · p. 1
• Examples of required invariance include axial rotations on a sphere and conjugation by SU(n) for lattice-QFT flow-based sampling [3].Figure 1 depicts an SU(3) density invariant to conjugation, with eigenvalue angles θ1 and θ2 on its axes.Source: Figure 1 · p. 1, §1 · p. 2
• A quotient approach can be difficult to construct, may have a boundary that prevents manifold CNFs, or may introduce density abnormalities, as in Boyda et al. [3].For the sphere modulo axial rotations, the quotient is a closed interval and a CNF would “flow out” at its boundary.Source: §1 · p. 2
• The proposed framework learns flows on arbitrary Riemannian manifolds with explicit symmetries and proves universal approximation for distributions on closed manifolds.The approach constructs equivariant vector fields and uses them to obtain symmetry-invariant flows.Source: §1 · p. 2
• Symmetry-aware flows produce higher-quality densities that respect required symmetries, and on SU(n) avoid the unnatural artifacts of quotient construction.The authors compare against general manifold density-learning approaches that are unaware of the problem's symmetries.Source: §1 · p. 2

Normalizing Flows on Manifolds

• Manifold CNFs extend continuous normalizing flows to Riemannian manifolds; this work further generalizes them to account for target-density symmetries.Earlier constructions included methods specific to tori, spheres, and hyperbolic space, as well as general Riemannian-manifold methods.Source: §2 · p. 3

Equivariant Machine Learning

• Equivariant normalizing flows had modeled symmetric densities in Euclidean space, including physical-chemistry settings [26].The cited work notes its approach is better suited to some such settings than general-purpose normalizing flows.Source: §2 · p. 3

Symmetries on Manifolds

• Before this work, manifold symmetry-invariant flow learning beyond conjugation-invariant SU(n) densities [3] had received little attention.The paper presents its approach as the first general equivariant manifold normalizing-flow model for arbitrary manifolds and symmetries.Source: §2 · p. 3

3 Background

• The background introduces manifold geometry and manifold normalizing flows as the setting for learning probability densities beyond Euclidean space.Source: §3 · p. 3

3.1 Riemannian Geometry

• Riemannian geometry equips each tangent space with a smoothly varying inner product, which induces a distance on the manifold.Source: §3.1 · p. 3

3.2 Normalizing Flows on Manifolds

• A manifold normalizing flow is a diffeomorphism that transforms a prior density into a model density, while a continuous flow follows a manifold ODE from a base point.The flow map sends each base point to the ODE solution at a chosen time.Source: §3.2 · pp. 3–4

3.3 Equivariance and Invariance

• An isometry subgroup acts on the manifold; equivariant functions, vector fields, and flows commute appropriately with that action, while invariant densities retain the same value under it.Source: §3.3 · p. 4

4 Invariant Densities from Equivariant Flows

• The section organizes the construction as a chain from invariant potentials to equivariant vector fields and flows, then to invariant densities.With an invariant prior, constructing an invariant potential is presented as simpler than directly parameterizing an invariant density or equivariant flow.Source: §4 · p. 4

4.1 Equivariant Gradient of Potential Function

• A smooth G-invariant potential has a G-equivariant gradient, which provides the vector field used to construct the flow.This construction is for isometry groups or their subgroups.Source: §4.1 · p. 5

4.2 Constructing Equivariant Manifold Flows from Equivariant Vector Fields

• A time-dependent vector field is G-equivariant if and only if its induced flow is G-equivariant.Source: §4.2 · p. 5

4.3 Invariant Manifold Densities from Equivariant Flows

• A G-invariant prior passed through a G-equivariant flow yields a G-invariant density, provided G is an isometry subgroup.The result does not hold for general diffeomorphism subgroups.Source: §4.3 · p. 5

4.4 Sufficiency of Flows Generated via Invariant Potentials

• On a closed Riemannian manifold, flows generated by invariant potentials can learn any smooth invariant target distribution in KL divergence.The theorem's target is smooth and non-vanishing.Source: §4.4 · p. 6

5.1 Equivariant Manifold Flow Model

• The model differentiates a G-invariant potential to obtain its vector field, then integrates on the manifold and computes density changes with the Riemannian CNF framework.Source: §5.1 · p. 6

5.2 Constructing G-invariant Potential Functions

• The potential constructions parameterize a neural network on the part of the manifold left free by the symmetry action.The two examples illustrate the method's versatility but are not exhaustive.Source: §5.2 · p. 6

5.2.1 Isotropy Invariance on S2

• For isotropy invariance on S2, the network takes the z-coordinate and integration time as inputs, so learned density variation is along the z-axis.The experiments use the uniform sphere density as their invariant prior.Source: §5.2.1 · p. 7

5.2.2 Conjugation Invariance on SU(n)

• For conjugation invariance on SU(n), a permutation-invariant DeepSet acts on the matrix eigenvalues, assigning equal density to similar matrices.The target distributions in lattice QFT must be invariant under SU(n) conjugation [3].Source: §5.2.2 · p. 7

5.3 Training Paradigms for Equivariant Manifold Flows

• With an exact target density, training minimizes KL divergence on samples from the Haar prior; with target samples, training uses a negative log-likelihood loss against a tractable prior.Source: §5.3 · pp. 7–8
• Figure 2 compares learned densities on SU(2) and SU(3), with each density normalized to a maximum value of 1.Source: Figure 2 · p. 8

6.1 SU(n) Gauge Equivariant Neural Network Flows

• The SU(n) experiments learn conjugation-invariant densities for lattice quantum field theory, and the model outperforms the construction of Boyda et al. [3].The setup follows Boyda et al. [3] and learns to sample from target distributions; potentials are parameterized as DeepSet networks on eigenvalues with a Haar-uniform prior.Source: §6 · p. 8, §6.1 · p. 9

6.1.1 SU(2)

• On SU(2), both models match high-density regions, but the equivariant flow better captures lower-density tails and produces smoother densities than Boyda et al. [3].The Boyda et al. [3] model retains mass near ±π, possibly due to its construction.Source: §6.1.1 · p. 9

6.1.2 SU(3)

• On SU(3), the learned densities fit targets more accurately and follow their geometry more naturally than Boyda et al. [3], with smoother curves instead of sharp corners.On SU(2), the model also improves in lower-density tails and shows fewer unnecessary bumps in high-density regions.Source: §6.1.2 · p. 9, §6.1.1 · p. 9

6.2 Asteroid Impact Dataset Bias Correction

• For asteroid impacts, an isotropy-invariant S2 flow models locations from a land-only dataset while accounting for the expected rotational invariance of impacts.Figure 3 contrasts recorded land impacts with the modelled density; a regular manifold flow would model impacts only on land.Source: §6.2 · pp. 9–10, Figure 3 · p. 9

6.3 Modelling Invariance Matters

• On the isotropy-invariant sphere dataset, the equivariant flow captures the distribution better than NMODE [29] despite having fewer parameters.The comparison uses flows of the same width; the equivariant flow has input dimension 1, and the authors attribute its performance to its symmetry-aware inductive bias.Source: §6.3 · p. 10, Figure 4 · p. 10

7 Conclusion

• The work introduces a general theory for equivariant manifold flows and demonstrates their use for conjugation-invariant SU(2) and SU(3) densities and asteroid-impact modelling.The SU(n) models adhere more naturally to target geometry than prior work, and the paper compares equivariant and regular manifold flows for asteroid impacts.Source: §7 · p. 10
• The experiments demonstrate the utility of equivariant flows for SU(n) lattice gauge theory sampling and compare their treatment of asteroid impacts with regular manifold flows.The conclusion also presents the asteroid application as evidence for modelling existing invariances.Source: §7 · p. 10
• The theory and implementations require substantial tuning, and future work could explore constructions beyond invariant potentials and whether they can learn over open manifolds.The authors also state that they cannot confidently assess future societal impacts because the theory and learning algorithm are too abstract.Source: §7 · p. 10

A Proof of Theorems

• The appendix frames these proofs as foundations for the models and notes that prior work addresses some of the formalized results.Source: A Proof of Theorems · p. 14

A.1 Proof of Theorem 1

• For a Riemannian manifold and an isometry group or subgroup, a smooth invariant potential gives the commuting diagram stated in Theorem 1.The proof identifies the corresponding gradient relation as tight: it holds only when the group action is by isometries.Source: §A.1 · p. 14
• The proof uses the Riemannian gradient chain rule and the isometry property to relate gradients under the group action.The resulting differential satisfies (DuLg)⊤◦(DuLg) = I.Source: §A.1 · p. 14

A.2 Proof of Theorem 2

• A time-dependent vector field is G-equivariant if and only if its flow is G-equivariant for every stated time T.The forward direction uses the flow of the pushed-forward vector field; the reverse direction follows from the chain rule.Source: §A.2 · p. 15
• For a diffeomorphism F:M→N and smooth vector field X on M with flow θ, the flow of F_*X is η_t=F∘θ_t∘F^{-1}, with domain N_t=F(M_t) for each t∈R.Source: §A.2 · p. 15

A.3 Proof of Theorem 3

• If ρ is a G-invariant density and f is a G-equivariant diffeomorphism, then the transformed density ρf remains G-invariant.The proof uses invariance of ρ, equivariance of f, and the unit absolute Jacobian determinant of an isometry.Source: §A.3 · pp. 15–16

A.4 Proof of Theorem 4

• The flow converges to π in KL, and its limiting map from ρ0 to π is well-defined under the stated differential-equation construction.The proof first treats uniform π via the heat equation, then extends the convergence argument to general π.Source: §A.4 · pp. 16–18
• For uniform π, the heat-equation solution converges in L2 to the constant potential, while G-invariant initialization preserves invariance and yields a G-equivariant flow.The harmonic expansion has exponentially decaying terms and a constant term.Source: §A.4 · p. 18

A.5 Conjugation by SU(n) is an Isometry

• Conjugation by SU(n) is an isometry subgroup because its action preserves orthonormality of tangent frames.The proof establishes this using the unitary matrix-conjugation action.Source: §A.5 · p. 18

B Manifold Details for the Special Unitary Group SU(n)

• SU(n) is a smooth Lie manifold whose identity tangent space consists of traceless skew-Hermitian matrices, with metric tr(A⊤B).Its elements are unitary matrices with determinant 1.Source: B Manifold Details for the Special… · p. 19

B.1 Haar Measure on SU(n)

• Haar measure supplies an invariant probability measure on SU(n), and the appendix describes sampling from this prior using Algorithm 1 [32].Integrals over SU(n) can be expressed through diagonal matrices under eigendecomposition.Source: §B.1 · p. 19

B.2 Eigendecomposition on SU(n)

• The SU(n) invariant-potential computation requires eigendecomposition and differentiation through it; explicit characteristic-polynomial formulas cover n=2 and n=3, but do not scale beyond n=4.Twice-differentiable complex eigendecomposition in newer PyTorch versions permits extension to higher dimensions.Source: §B.2 · pp. 19–20

B.2.1 Explicit Formula for SU(2)

• For SU(2), the appendix derives eigenvalues from the characteristic polynomial and notes an S3 isomorphism that permits learning a regular flow on S3 instead.For SU(3), it derives the characteristic polynomial and solves the resulting cubic using Cardano’s formula.Source: §B.2.1 · p. 20, §B.2.2 · p. 20

B.2.2 Explicit Formula for SU(3)

• The SU(3) eigenvalue derivation forms the characteristic polynomial, transforms it to a depressed cubic, and applies Cardano’s formula.The cubic roots are selected to satisfy the stated product condition.Source: §B.2.2 · p. 20

• The section provides additional details about the experiments learning invariant densities on SU(n).Source: C Experimental Details for Learning… · p. 21
• The authors exclude ESS because mode-collapsed densities could score near-perfectly while failing to match target geometry.Source: C Experimental Details for Learning… · p. 21
• They instead use distribution-geometry visualizations, noting that Boyda et al. [3] released no code and reported ESS only for some test cases.Source: C Experimental Details for Learning… · p. 21

C.1 Training Details

• Training minimizes KL divergence using uniform SU(n) batches, Adam, and gradients of the batch KL; the implementation also required numerical-stability and learning-rate tuning.The reported Figure 2 run used 300 iterations, batch size 8192, weight decay 0.01, and starting learning rate 0.01.Source: §C.1 · p. 21

C.2 Conjugation-Invariant Target Distributions

• The SU(2) and SU(3) experiments use three instances of a conjugation-invariant target density family parameterized by c_k and β.Table 1 lists the parameter sets for the experiments.Source: §C.2 · p. 21, Table 1 · p. 21

C.2.1 Case for SU(2)

• For SU(2), the target density is expressed using the eigenvalue angles θ and −θ, with tr(U) = 2 cos(θ).Source: §C.2.1 · p. 22

C.2.2 Case for SU(3)

• For SU(3), the eigenvalues are represented by angles θ1, θ2, and −θ1−θ2.Source: §C.2.2 · p. 22

D Learning Continuous Normalizing Flows over Manifolds with Boundary

• The boundary-manifold approach represents a manifold with boundary as a quotient of a larger boundary-free manifold and learns with an invariant potential so the density descends smoothly.The interval [0, 1] is presented as a quotient of S2 by the isotropy group at the north pole; the resulting density is constant along latitude circles.Source: D Learning Continuous Normalizing… · p. 22