Technical Summary: Radial Compensation in Chart-Based Generative Models

1. Problem Statement

Generative models on Riemannian manifolds (e.g., Variational Autoencoders, Normalizing Flows, and Continuous Normalizing Flows) often rely on a "chart-based" approach. This involves sampling from a Euclidean base distribution in the tangent space $T_pM$ and mapping the sample to the manifold $M$ via a chart (e.g., the exponential map).

The paper identifies a fundamental semantic distortion in this standard pipeline:

1. Radial Distortion: Standard isotropic Euclidean bases (like a Gaussian) do not preserve the intended distribution of the geodesic radius $R = d(p, q)$ on the manifold. For instance, an isotropic Gaussian in the tangent space maps to a $\chi_n$ distribution of radii on the manifold, which rarely matches the practitioner's intended prior (e.g., HalfNormal, Gamma).
2. Chart Dependence: Non-geodesic charts (such as Lambert equal-area maps) introduce additional distortions where the same tangent-space scale corresponds to different geodesic radii depending on the chart choice and manifold curvature.
3. Hidden Modeling Decisions: Consequently, the choice of chart and base distribution silently fixes a radial law on the manifold. To compensate for this hidden distortion, the model must learn to adjust its scale or curvature parameters, leading to unstable training, large Jacobian terms, and difficult-to-interpret learned curvatures.

2. Methodology: Radial Compensation (RC)

The authors propose Radial Compensation (RC), a method to decouple the statistical meaning of the model from the numerical properties of the chart.

Core Concept

Instead of starting with a Euclidean base and letting the chart distort it, RC starts with a user-specified target law for the geodesic radius $R \sim p_{R,\theta}$ on the manifold. It then constructs a specific isotropic base distribution in the tangent space such that, after mapping through any scalar-Jacobian azimuthal chart, the resulting manifold density depends on the sample only through the geodesic radius $R$.

Mathematical Formulation

For a scalar-Jacobian azimuthal chart $T$ with radial Jacobian $J_T(r)$ and radius map $R_T(r) = d(p, T(ru))$, the RC base density $f_\theta(x)$ on the tangent space $T_pM \cong \mathbb{R}^n$ is defined as:
$$f_\theta(x) = \phi_\theta(R_T(\|x\|)) J_T(\|x\|)$$
where $\phi_\theta(R)$ is the target manifold density derived from the user-specified radial law $p_{R,\theta}$ and the manifold's volume element (specifically, $\phi_\theta(R) = p_{R,\theta}(R) / S_\kappa(R)$, where $S_\kappa$ is the area of the geodesic sphere).

This construction ensures:

• Geodesic-Radial Likelihoods: The pushforward density on the manifold is $\rho_\theta(q) = \phi_\theta(d(p, q))$.
• Chart Invariance: The realized manifold density and the radial Fisher information are invariant to the choice of scalar-Jacobian chart.
• Uniqueness: Within the class of isotropic bases and scalar-Jacobian azimuthal charts, RC is the unique construction that simultaneously preserves geodesic-radial likelihoods, chart-invariant Fisher information, and tangent-space isotropy.

Balanced Exponential Charts (bExp)

To address numerical conditioning without altering the statistical meaning, the authors introduce the Balanced Exponential family ($bExp_\alpha$). This family interpolates between:

• $\alpha = 0$: Equal-area Lambert charts (volume-preserving, $|\det| \equiv 1$).
• $\alpha = 1$: The standard exponential map (geodesic-exact radii).

Under RC, varying $\alpha$ acts purely as a numerical preconditioner. It tunes the conditioning of the likelihood optimization (reducing the variance of the chart term in the gradient) without changing the realized manifold density. Theoretically, the variance of the chart term scales as $O(\alpha^2)$.

3. Key Contributions

1. Radial Compensation (RC): A theoretical framework and practical recipe to pre-warp the tangent-space base distribution so that the geodesic radius on the manifold follows a user-specified law, regardless of the chart used.
2. Theoretical Guarantees:
   • Theorem 1: Proves that under RC, the Fisher information and KL divergence on the manifold reduce exactly to their one-dimensional radial counterparts, independent of the chart.
   • Theorem 2 & 3: Establishes that RC is the unique solution within the isotropic scalar-Jacobian class to achieve geodesic-radial likelihoods and chart-invariant Fisher information.
   • Theorem 4 & 12: Characterizes the $bExp_\alpha$ family and proves that chart-term variance scales quadratically with $\alpha$, allowing for numerical tuning without statistical cost.
3. Balanced Exponential Charts: A new family of charts that allows users to trade off between volume distortion and geodesic fidelity purely for numerical stability.

4. Experimental Results

The authors evaluate RC on synthetic manifolds ($S^2, H^2$) and standard image benchmarks (MNIST, Fashion-MNIST, Omniglot, CIFAR-10) using Manifold VAEs and Continuous Normalizing Flows (CNFs).

• Semantic Calibration: RC successfully recovers the intended geodesic radius law. On $S^2$ and $H^2$, RC reduces the KL divergence between the realized and target radius distributions from $\approx 1.48$ and $0.28$ (for raw exponential wrapping) to $\approx 10^{-3}$.
• Numerical Stability:
  • In high-dimensional latent CNFs ($d_z \in \{32, 64, 128\}$), raw exponential wrapping leads to unstable training (radius blow-ups, objective excursions). RC variants remain stable under the same solver settings.
  • Varying $\alpha$ in $bExp_\alpha$ reduces chart-term variance significantly (e.g., from 1.016 to 0.167 on CIFAR-10) and decreases the mean number of function evaluations (NFEs) for ODE solvers, indicating reduced stiffness.
• Curvature Interpretability: In mixed-curvature VAEs, RC prevents "curvature absorption," where the model learns incorrect curvatures to compensate for chart distortions. RC yields learned curvatures that are less chart-compensatory and improves the ELBO by approximately 7 nats (primarily through a lower KL term).
• Application Checks: RC maintains or improves performance in hyperbolic WordNet flows and protein orientation models on $S^3$, confirming that the method generalizes beyond latent space modeling to data with intrinsic manifold structure.

5. Significance and Claims

The paper claims that RC fundamentally changes how chart-based generative models should be constructed:

• Decoupling Semantics and Numerics: RC separates the statistical meaning of the model (the law of the geodesic radius) from its numerical conditioning (the choice of chart). The chart becomes a "numerical preconditioner" rather than a hidden modeling decision.
• Stability and Interpretability: By removing the need for the model to compensate for chart-induced distortions, RC leads to more stable training, cleaner curvature estimates, and models where the learned parameters have a direct, interpretable geometric meaning.
• Drop-in Replacement: RC requires no changes to the architecture, decoder, solver, or generative objective; it is a modification solely to the base distribution, making it a practical "drop-in" improvement for existing manifold flow and VAE tooling.

The authors acknowledge limitations: the exact invariance theory requires a tractable polar volume factor and scalar-Jacobian azimuthal charts. They note that for per-example analytic Gaussians (non-isotropic), the standard wrapped exponential remains the appropriate tool, and scaling RC to manifolds without closed-form polar volumes is a future direction.