Manifold Learning with Normalizing Flows: Towards Regularity, Expressivity and Iso-Riemannian Geometry

This paper addresses distortions and modeling errors in multi-modal data by proposing a framework that isometrizes learned Riemannian structures and balances the regularity and expressivity of diffeomorphism parametrizations within the context of manifold learning with normalizing flows.

The Gist

Imagine you are trying to navigate a complex, hilly landscape where the "ground" isn't flat, but a twisted, folded sheet of paper floating in 3D space. This is how modern AI often views data: high-dimensional information (like a photo of a cat) actually lives on a much simpler, lower-dimensional "manifold" (the hidden rules that make a cat a cat).

The paper you shared is about building better maps for this landscape. The authors, Willem Diepeveen and Deanna Needell, are tackling two main problems that happen when AI tries to learn these maps: distortion and wobbly roads.

Here is the breakdown using simple analogies:

1. The Problem: The "Rubber Sheet" Distortion

Imagine you have a rubber sheet representing your data. You want to draw a line (a geodesic) between two points, say, a picture of a "2" and a picture of a "6" in the MNIST dataset.

• The Old Way (The Stretchy Map): In previous methods, the AI learned a map where the rubber sheet stretches and shrinks unevenly. If you walk from the "2" to the "6" at a constant speed on the map, you might spend 90% of your time walking through a "forest" of empty space (low-density areas) and only 10% walking through the "city" where the actual data lives.

  • The Result: If you try to interpolate (create a smooth transition) between the two, the AI generates weird, blurry images that look like nothing because it spent too much time in the empty zones. It's like trying to drive from New York to London by spending 9 hours driving through the middle of the Atlantic Ocean just because the map says that's the shortest path.

• The Solution: "Iso-Riemannian" Geometry (The Speedometer Fix): The authors propose a trick called Isometrization. Think of this as putting a speedometer on your car that forces you to travel at a constant speed relative to the actual terrain, not the distorted map.

  • They mathematically re-wire the map so that every step you take covers the same amount of "real" distance.
  • The Result: Now, when you travel from "2" to "6," you spend your time exactly where the data is. The transition is smooth, logical, and makes sense. It's like switching from a distorted Mercator projection (where Greenland looks huge) to a map that accurately represents travel time.

2. The Problem: The "Wobbly Bridge" (Irregularity)

Now, imagine the AI is trying to build a bridge between two islands (two clusters of data, like cats and dogs).

• The Old Way (The Wild Architect): To make the bridge look fancy and handle complex shapes, the AI uses very flexible, "expressive" building blocks. However, these blocks are so flexible that they twist and turn wildly in the empty space between the islands.

  • The Result: The bridge might connect the islands, but it takes a bizarre, looping path that doesn't make sense. If you try to cross it, you might fall off or end up in a weird place. This is bad for fairness and interpretation because the AI's "logic" is chaotic in the gaps between data.

• The Solution: The "Steady Architect" (Regular Flows): The authors suggest using a specific type of building block that is regular (smooth and predictable) but still expressive (capable of handling complex shapes).

  • They combine simple, stable linear layers with smart, bounded non-linearities. Think of it as using a flexible hose that can bend, but has a rigid internal skeleton so it doesn't kink or twist uncontrollably.
  • The Result: The bridge takes the most direct, natural path between the islands. It avoids the weird loops and ensures that the transition between data types is smooth and logical.

3. The Grand Finale: Putting It All Together

The paper's main discovery is that you need both solutions to get the best results.

• Analogy: Imagine you are building a roller coaster.
  • Isometrization ensures the train moves at a constant, safe speed so passengers don't get thrown around (fixing the distortion).
  • Regular Flows ensure the track is built smoothly without sudden, dangerous spikes or loops (fixing the wobbly architecture).

When the authors combined these two techniques, they found that their AI could:

1. Interpolate better: Creating smooth, realistic transitions between data points (like morphing a "2" into a "6" without it turning into a blob).
2. Reduce dimensions better: Compressing complex data into simpler forms without losing important details or introducing errors.
3. Be more fair: Ensuring that the AI treats all parts of the data equally, rather than making huge errors for data points that are far apart.

Summary

In the world of machine learning, data is often a twisted, high-dimensional shape. Previous methods tried to flatten this shape but ended up stretching it weirdly or building unstable bridges between data points.

This paper says: "Let's fix the map so distances are fair (Isometrization), and let's build the bridges with steady, reliable materials (Regular Flows)." By doing both, they create a system that understands data more naturally, leading to better, more interpretable, and more reliable AI.

Technical Summary

1. Problem Statement

Modern machine learning often relies on the Manifold Hypothesis, which posits that high-dimensional data lies near a low-dimensional, non-linear manifold. While learning Riemannian geometry on these manifolds improves tasks like clustering and interpolation, current approaches face two critical challenges when dealing with multi-modal data:

1. Distortions from Non-Isometry: Existing methods using "pullback geometry" (deriving a metric from a diffeomorphism $\phi$) often fail to maintain local $\ell_2$-isometry on the data support. This leads to geodesics that do not travel at constant $\ell_2$-speed. Consequently, interpolation spends disproportionate time in low-density regions, creating "distorted" paths that misrepresent the data structure and degrade interpretability.
2. Regularity vs. Expressivity Trade-off: To model complex, multi-modal data, normalizing flows (NFs) typically use highly expressive, non-volume-preserving architectures (e.g., affine couplings, spline flows). However, these architectures often lack regularity (smoothness/stability), leading to incorrect geodesic transitions between modes and unstable dimensionality reduction. Conversely, volume-preserving flows (e.g., additive flows) are regular but often lack the expressivity needed for complex manifolds.

The core question is: Can we construct a framework that balances regularity and expressivity while ensuring the learned geometry is "isometrized" to prevent distortions in downstream tasks?

2. Methodology

The authors propose a two-pronged approach to address these issues: Iso-Riemannian Geometry (a post-processing correction) and Regular Normalizing Flows (a parametric improvement).

A. Iso-Riemannian Geometry (Theoretical Framework)

The authors introduce a systematic method to "isometrize" any learned Riemannian structure without altering the underlying manifold topology. This involves reparametrizing geodesics to ensure constant $\ell_2$-speed.

• Iso-Geodesics: Instead of the standard geodesic $\gamma_{x,y}(t)$, they define an iso-geodesic $\gamma^{iso}_{x,y}(t)$ by applying a time-change $\tau_{x,y}(t)$. This mapping ensures that the curve traverses the path at a constant speed relative to the Euclidean metric, regardless of the intrinsic curvature.
• Iso-Manifold Mappings: Based on the iso-geodesic, they define:
  • Iso-Logarithm ($\log^{iso}$): A rescaled logarithmic map where the $\ell_2$-length equals the intrinsic arc length.
  • Iso-Exponential ($\exp^{iso}$): The inverse of the iso-logarithm.
  • Iso-Parallel Transport: A transport operator that preserves $\ell_2$-lengths along the geodesic.
• Iso-Distances: A distance metric derived from the arc length, ensuring consistency with the isometrized mappings.
• Application to Dimension Reduction: They propose Algorithm 2, an isometrized version of Riemannian low-rank approximation. By projecting data onto the tangent space using $\log^{iso}$ and reconstructing via $\exp^{iso}$, the global approximation error is minimized more effectively than standard methods, as it decouples the tangent space approximation from the distortion caused by non-constant geodesic speeds.

B. Regular Normalizing Flow Parametrization

To address the modeling errors in multi-modal settings, the authors propose a specific architecture for the diffeomorphism $\phi$ used to generate the pullback metric.

• Architecture Design: Instead of standard expressive flows (like affine couplings) or purely volume-preserving flows, they propose a composition of invertible linear layers and regularized additive coupling layers.
  • Linear Layers: Use learnable normalization followed by orthogonal matrices (via Householder decomposition) or 1x1 convolutions. These ensure the determinant is constant (though not necessarily 1), relaxing the strict volume-preserving constraint while maintaining stability.
  • Non-Linearities: Use additive coupling with neural networks employing bounded activation functions (specifically a sum of $\tanh$ functions with learnable coefficients). This ensures the derivatives are bounded, promoting regularity.
• Training Objective: They simplify the training loss to the standard Negative Log-Likelihood (NLL) with weight decay, removing the complex regularization terms required by previous methods (like [7]) to enforce local isometry. The architecture itself enforces the necessary regularity.

3. Key Contributions

1. Iso-Riemannian Geometry: A formal framework to reparametrize geodesics and manifold mappings to have constant $\ell_2$-speed. This eliminates distortions in interpolation and dimension reduction caused by non-isometric pullback metrics.
2. Regular yet Expressive Flows: A novel diffeomorphism architecture that combines the stability of linear/orthogonal layers with the expressivity of bounded non-linearities, specifically tailored for multi-modal data.
3. Simplified Training: Demonstrating that with the proposed regular architecture, complex regularization terms are unnecessary; standard NF training suffices to learn valid pullback geometries.
4. Synergistic Performance: Showing that combining Iso-Riemannian geometry with Regular Flows yields superior results compared to using either technique in isolation.

4. Results

The authors validate their approach on synthetic data (bimodal distributions, hemispheres) and real-world data (MNIST).

• Interpolation (Geodesics):
  • Standard/Expressive Flows: Geodesics often take unnatural paths between modes (e.g., entering a mode from the "side" rather than the top) and exhibit variable speeds, causing interpolation artifacts.
  • Regular Flows + Iso-Riemannian: Geodesics follow the natural high-density path between modes. Iso-geodesics ensure uniform sampling along the path, improving interpretability.
• Dimension Reduction (Rank-$r$ Approximation):
  • Standard Methods: Show significant distortion in tangent space projections, leading to poor reconstruction (high RMSE).
  • Proposed Method: Significantly reduces the Low-Rank Relative RMSE. For example, on modeled pullback data, the error dropped from 0.1741 (non-isometrized) to 0.0606 (isometrized).
• Multi-Modal Data (MNIST):
  • On MNIST, the regularized flow successfully learned the transition between digits (e.g., '2' to '6').
  • While the improvement in reconstruction error was smaller for MNIST compared to synthetic data, the geodesic interpolation remained significantly more faithful to the data manifold, and the "distortion" in the tangent space was eliminated.
• Distance from Barycentre: Experiments show that the benefits of isometrization become increasingly significant as data points move further from the Riemannian barycentre, where non-linear distortions typically accumulate.

5. Significance

This work bridges a critical gap between generative modeling (which prioritizes expressivity and non-volume-preserving flows) and geometric data analysis (which requires stability, regularity, and isometry).

• Interpretability: By ensuring geodesics travel at constant speed and follow high-density regions, the method provides a more intuitive and faithful representation of "intermediate" data points, crucial for applications in fairness and explainable AI.
• Scalability: The proposed framework is scalable to high-dimensional data and avoids the computational bottlenecks of previous metric-learning approaches.
• Theoretical Unification: It establishes that "Iso-Riemannian Geometry" is a robust generalization of Riemannian data analysis, allowing standard algorithms (like PCA on manifolds) to be applied to complex, learned geometries without suffering from metric-induced distortions.

In summary, the paper provides a principled path forward for manifold learning: use regular normalizing flows to learn the manifold structure and iso-Riemannian corrections to ensure the geometric operations performed on that manifold are stable and interpretable.