Matching Normalizing Flows and Probability Paths on Manifolds

Motivation

• Manifold generative modeling addresses data from spheres, protein angles, and robot motion, extending beyond Euclidean applications.The paper frames manifold modeling as important across several scientific domains.Source: §1 · p. 1

Existing methods

• CNFs offer greater expressive power than manifold NFs, but likelihood training requires solving an ODE to evaluate model densities.CNFs parameterize diffeomorphisms through neural-network vector fields rather than restricted compositions of invertible transformations.Source: §1 · p. 1

CNFM

• PPD compares a target probability path with the CNF-generated path through log mass conservation, avoiding model-density evaluation during training.The divergence depends on the vector field, its first derivatives, and the target path’s log density.Source: §1 · p. 2
• CNFM accepts flexible target paths and applies to manifolds and higher-dimensional data, while PPD bounds standard divergences for different values of ℓ.The paper states bounds for total variation at ℓ=1, α-divergence for 1<ℓ<∞, and reversed KL for ℓ=∞.Source: §1 · p. 2

Experiments

• Experiments report state-of-the-art sample quality and likelihoods on low-dimensional manifold datasets, faster optimization than prior CNF training, and scaling to moderately high-dimensional manifolds.The introduction specifically contrasts this scalability with previous manifold generative models.Source: §1 · p. 2

Manifold setting

• A smooth Riemannian manifold supplies the metric, volume form, tangent spaces, gradients, and divergence used to define manifold probability models.The paper restricts attention to strictly positive, smooth densities and bounded smooth time-dependent vector fields.Source: §2 · p. 2

CNFs

• A time-dependent tangent vector field defines a one-parameter family of diffeomorphisms through an ODE.This family is the geometric flow underlying the CNF.Source: §2 · p. 2
• In generative modeling, a CNF pushes forward or pulls back probability densities and events on the manifold.The notation distinguishes pushed-forward densities from pulled-back densities under the flow.Source: §2 · p. 2

Probability paths

• A probability path is a smooth time-indexed family of positive densities normalized over the manifold.The paper denotes the collection of such paths by P(M).Source: §2 · p. 2

CNF visualization

• Figure 1 depicts the forward CNF transporting noise to data on a sphere and the reverse-time flow transporting data back to noise.The forward sequence runs from t=0 to t=1, while the reverse path is shown separately.Source: Figure 1 · p. 3, §2 · p. 3

3. Matching CNF and target probability

• CNF Matching optimizes a divergence d(p, q) between a target path p and the CNF-generated path q, with q produced from prior p0 by a learnable vector field.The target path begins from a simple prior and approximates the data distribution at t=1.Source: §3 · p. 3

3.2. Probability path divergence

• PPD measures how far log p and the CNF vector field fail to satisfy the Logarithmic Mass Conservation formula, defining a divergence between target and generated density paths.The LMC formula is necessary and sufficient for a vector field to generate the target path, while its logarithmic form reduces high-dimensional numerical issues.Source: §3.1 · p. 3, §3.2 · p. 4
• PPD evaluation and parameter derivatives do not require access to q or the flow map, sidestepping ODE solving during training; it also bounds corresponding f-divergences of pT and qT for every T.The bound covers standard divergences through the f-divergence family, including total variation, KL, reverse KL, and α-divergences.Source: §3 · p. 3, §3.2 · p. 4
• For ℓ=1, the path divergence is equivalent to the total variation norm of the density ratio along flow trajectories, while increasing ℓ approaches the −log t limit.Except for ℓ=1, the associated f functions are differentiable and share the same derivative at 1.Source: §3.2 · p. 4

3.3. Target paths

• Valid target paths must connect a simple prior to an approximation of pdata, support efficient sampling at each time, and provide or approximate derivatives of log pt.These requirements do not require an SDE or Fokker–Planck PDE characterization; constructions include unimodal and uniform priors suited to different manifold types.Source: §3.3 · p. 5
• Target paths can use unimodal priors for non-compact manifolds, uniform priors for compact manifolds, and product-manifold kernels whose concentration follows componentwise concentration.The product construction applies to manifold-valued states such as combinations of positions and rotations.Source: §3.3 · pp. 5–6

Relations to existing CNF models

• CNF likelihood training requires solving ODEs for trajectories and log densities, whereas PPD minimization avoids ODE solves during training.The contrast follows from evaluating log q along characteristics in standard CNF training versus directly minimizing PPD.Source: §4.1 · p. 6
• Moser Flow also avoids training-time ODE solves and applies to manifolds, but its positivity loss, density-based formulation, and fixed path hinder high-dimensional scaling and flexibility.Moser Flow uses convex combinations of prior and model densities, while the paper’s approach is formulated in log-probabilities and supports broader target paths.Source: §4.1 · pp. 6–7

Relations to diffusion and score based generative models

• CNFM does not require a target path generated by a known diffusion or SDE, and its reverse process is obtained by solving the learned ODE backward in time.This makes target-path selection more flexible than in diffusion, score, and SDE models with restricted forward processes.Source: §4.2 · p. 7
• On manifolds, diffusion approaches may require numerical solutions because heat-equation kernels are rarely known, whereas CNFM can use arbitrary paths such as closed-form vMF paths on the sphere.The paper also notes that ODE sampling has higher asymptotic convergence rates than SDE sampling and supports widely used higher-order solvers.Source: §4.2 · p. 7

5.1. Toy densities on R2 and S2

• CNFM learns highly similar model and ground-truth densities on 2D Euclidean and spherical toy datasets.The experiments use PPD with ℓ=1 and visualize data samples, generated samples, and learned densities; spherical densities are shown with Mollweide projection.Source: §5.1 · p. 8

5.2. Earth and climate dataset

• CNFM improves state of the art by a large margin in Earth and Climate negative log likelihood, with Moser Flow as the runner-up.The comparison uses a six-layer, 512-neuron MLP, PPD order ℓ=2, and κ1=55K selected from three tested values.Source: §5.2 · p. 8, Table 1 · p. 8

5.3. Higher dimensional spheres

• The higher-dimensional experiment constructs increasingly complex densities r_k on S^15 for k=2,3,4 and compares CNFM with vMF mixtures and a stereographic-projection manifold CNF.The datasets contain 45K training samples and 5K test samples, while Table 2 reports NLL scores for CNFM and baselines.Source: §5.3 · p. 8, Table 2 · p. 8
• CNFM and FFJORD use comparable architectures and run for about 4K seconds, while epoch-time differences increase on the log scale for larger architectures.The comparison reports timing behavior for both the higher-dimensional experiment and larger architecture types.Source: §5.3 · p. 9, Figure 6 · p. 9

Robotics data

• The robotics experiment models Walker and Cheetah on R^3×(S^1)^6 and Humanoid on R^3×(S^1)^8×(S^3)^6, generating uncurated samples from trained CNFMs.The data are generated from 50 simulated trajectories, with 17.5K randomized samples for each robot type.Source: §5.4 · p. 9, Figure 7 · p. 9

• Dimensions above 100, including image-scale data, require further work for CNFM.Source: §6 · p. 9
• The trained CNFM produces noise-to-data paths on product manifolds representing Cheetah, Walker, and Humanoid robot state spaces.Source: Figure 8 · p. 9
• The main scaling challenges are biased gradient estimates from stochastic divergence approximation and conditioning from time-dependent PPD scales.Source: §6 · p. 9
• On MNIST with d = 784 and batch size 128, CNFM training remains stable, but generation quality does not match state-of-the-art CNF models.Source: §6 · p. 9

• CNFM minimizes Probability Path Divergence without sampling model densities, making CNF training easier to apply on manifolds.Source: §7 · p. 9
• PPD upper bounds standard divergences and supports a flexible family of target paths on manifolds.Source: §7 · p. 9
• Empirically, CNFM improves training time, scales to moderately dimensional manifolds, and produces state-of-the-art samplings and log likelihoods.Source: §7 · p. 9

• The mass-conservation formula applies because the vector field is locally Lipschitz and bounded, with the stated integrability condition.Source: A. Proof of Theorem 1 · p. 11
• Assuming p_t > 0 permits division by p_t, converting the mass-conservation equation into a density-normalized relation.Source: A. Proof of Theorem 1 · p. 11
• Using div(fv) = ⟨∇_x f, v⟩ + f div(v) expands the divergence term into directional-derivative and vector-field-divergence components.Source: A. Proof of Theorem 1 · p. 11
• The identity ∂_t log p_t = ∂_t p_t / p_t establishes equivalence between the density equation and the logarithmic mass-conservation equation.Source: A. Proof of Theorem 1 · p. 11

• A time-dependent vector field defines a two-parameter diffeomorphism family through an ordinary differential equation, with the CNF map given by φ_t = Φ_t,0.Source: B. Proof of Lemma 1 · p. 11
• The theorem supplies the relevant flow identity for every time t and point x on the manifold.Source: B. Proof of Lemma 1 · p. 11
• Subtracting the theorem’s identity in the loss yields the expression used in the lemma’s derivation.Source: B. Proof of Lemma 1 · p. 11
• The composition Φ_s,t(φ_t(x)) = Φ_s,0(x) = φ_s(x) follows from the two-parameter flow property, while the distributional step uses φ_t(x) ∼ q_t when x ∼ p_0.Source: B. Proof of Lemma 1 · pp. 11–12

Proof details

• Theorem 2’s proof uses monotone convergence after establishing an increasing sequence converging to −log(t).The argument also invokes inequalities justified through convexity, tangency, triangle inequality, and Hölder inequality.Source: C. Additional details for the proof of… · pp. 12–13

Numerical implementation

• The implementation stabilizes von Mises–Fisher normalization and derivatives using scaled Bessel functions and approximations for high dimensions.The derivative approximation averages lower and upper bounds proposed by Ruiz-Antolín & Segura (2016), following Oh et al. (2019).Source: D. Numerically stable derivative of… · p. 13

Manifold implementation

• The appendix represents Euclidean manifold vector fields with MLPs and computes PPD using the standard inner product, gradient, and divergence.For spherical manifolds, the construction uses an MLP-induced tangent field and the induced Euclidean metric.Source: E. Vector field representation · p. 14
• On product manifolds, the tangent field combines component MLPs, while inner products, gradients, and divergence decompose across manifold factors.The target path is defined using the stated kernel construction, extending similarly to products of more than two manifolds.Source: E. Vector field representation · p. 14

R2 and S2 datasets

• R2 experiments use 3-layer, 256-hidden-unit MLPs with Adam, batch size 1000, learning rate 1e −4, σ1 = 0.01, and ℓ= 1.Learning rates and σ1 values were searched over the listed alternatives.Source: §F.1 · p. 14
• S2 experiments use 6-layer, 512-hidden-unit MLPs with Adam, batch size 1000, learning rate 1e −4, κ = 5000, and ℓ= 1.The passage specifies the training configuration but not a separate result comparison.Source: §F.1 · p. 14

S15 datasets

• S15 experiments vary k across 2, 3, and 4 using 3-layer, 64-hidden-unit MLPs, Adam, batch size 7000, κ = 5K, and ℓ= 2.Learning rates were searched over 1e −3, 1e −4, and 1e −5.Source: §F.3 · p. 14

Robotics datasets

• Robotics experiments use 6-layer, 512-hidden-unit MLPs with Adam, learning rate 1e −4, batch size 1000, κ = 55K, and ℓ= 1.The κ search includes 5K and 55K; CNFM and S-FFJORD use approximate divergence via the Hutchinson estimator.Source: §F.4 · p. 14, §F.3 · p. 14

Robotics visualizations

• Figures 9 and 10 show additional uncurated samples and noise-to-data paths for Cheetah, Walker, and Humanoid.The visualizations use trained CNFs on product manifolds representing each robot’s state space.Source: G. Extra Experimental Results · p. 14, Figure 9 · p. 15, Figure 10 · p. 15