Isometric Manifold Learning for Injective Flows

• The paper introduces an isometric regularization approach that preserves local distances in injective normalizing flows to enable tractable density estimation.
• It employs a dual-stage architecture combining isometric autoencoders with latent normalizing flows to efficiently model the intrinsic geometry of data manifolds.
• Empirical results confirm that this method improves reconstruction accuracy and density estimation on both synthetic and real datasets while adapting to curved manifolds.

Isometric manifold learning for injective @@@@1@@@@ addresses the geometric and probabilistic representation of high-dimensional data that concentrate near low-dimensional manifolds. The central goal is to utilize parameterized injective mappings (often realized by neural networks) to embed latent codes into the ambient data space, capturing the underlying manifold while preserving local geometric structure (isometry), and supporting density estimation directly on the manifold. This approach unifies principles from Riemannian geometry, normalizing flows, and autoencoding architectures, and has motivated a series of specialized formulations, algorithms, and regularization strategies.

1. Mathematical Framework for Injective and Isometric Flows

Let $Z \cong \mathbb{R}^d$ denote a low-dimensional latent space and $X \subset \mathbb{R}^D$ the ambient data space with $d < D$. An injective normalizing flow is a smooth, injective mapping

$f:Z \longrightarrow X,\qquad x = f(z)$

whose image $\mathcal{M} = f(Z)\subset X$ forms a $d$-dimensional immersed submanifold representing the learned data manifold. When $f$ is parameterized by neural networks (as in standard flows), the Jacobian $J_f(z) \in \mathbb{R}^{D \times d}$ pulls back the ambient Euclidean metric to

$g(z) = J_f(z)^T J_f(z) \in \mathbb{R}^{d\times d}$

There is a unique preimage $z = f^{-1}(x)$ for $x \in \mathcal{M}$, and the manifold inherits a Riemannian volume element

$$d\mathrm{Vol}_\mathcal{M}(x) = \sqrt{\det \big( J_f(z)^T J_f(z) \big) }\,dz$$

The induced probability density, confined to the manifold, is

$$p_\mathcal{M}(x) = p_Z(z) \big| \det( J_f(z)^T J_f(z) ) \big|^{-1/2} \qquad \text{with } z=f^{-1}(x)$$

and is normalized on $\mathcal{M}$ (Brehmer et al., 2020).

A perfect isometric embedding satisfies $J_f(z)^T J_f(z) = I_d$ for all $z$, preserving first-order distances on the manifold (Cramer et al., 2022).

2. Isometry Regularization and Iso-Riemannian Geometry

Exact global isometry is unattainable for general nonlinear data manifolds embedded in high-dimensional ambient space. Instead, approximate isometry is enforced by regularization terms such as

$$R_{\rm iso}[f] = \frac1N \sum_{i=1}^N \| J_f(z_i)^T J_f(z_i) - I_d \|_F^2, \qquad z_i \sim p_Z$$

or, in autoencoder formulations,

$$\mathcal{L}_{\rm iso}(\theta_F) = \mathbb{E}_{z\sim p_Z} \mathbb{E}_{u\in S^{d-1}} \left( \| J_f(z) u \|_2 - 1 \right)^2$$

where $S^{d-1}$ denotes the unit sphere in $\mathbb{R}^d$ (Cramer et al., 2022, Brehmer et al., 2020). Additional “pseudo-isometry” penalties can be imposed on the encoder for autoencoder-based schemes.

The framework of iso-Riemannian geometry (Diepeveen et al., 12 May 2025) further formalizes these ideas: the pullback metric $g_Z(z)$ is controlled either explicitly by penalizing deviation from identity, or indirectly by constraining the architecture (constant-determinant blocks, bounded derivatives) and by post-hoc isometrization—reparameterizing geodesics so that their ambient speed is constant.

3. Model Construction and Algorithmic Designs

Modern isometric-injective manifold flows comprise two principal design patterns:

Decoupled Autoencoder + Flow composition:

• Stage 1: An isometric autoencoder (I-AE) $(g, f)$ maps $x\mapsto z\mapsto x'$, trained for low reconstruction loss and isometry regularization, with $g:\mathcal{X}\to\mathcal{Z}$ and $f:\mathcal{Z}\to\mathcal{X}$.
• Stage 2: A $d$-dimensional standard normalizing flow $T$ is trained on the latent codes $z=g(x)$, yielding latent density $p_Z(z)$ (Cramer et al., 2022).

The tractable composite density on the manifold is then

$$p_X(x) \approx p_V(T^{-1}(g(x))) |\det J_T(T^{-1}(g(x)))|^{-1}$$

with the isometric property guaranteeing $|\det J_f| \approx 1$ (thus the ambient Jacobian determinant need not be modeled) (Cramer et al., 2022).

Principal Manifold Flows (PF, iPF):

Rather than explicitly learning a single invariant chart, PF/iPF methods seek to align latent-space partitions (“contours”) with principal directions of the data manifold by orthogonality constraints on sub-Jacobians, with empirical objectives

$I_{s,t} = L_{s\cup t} - L_s - L_t \geq 0$

that vanish only for mutually orthogonal latent blocks. Injective variants (iPF) adopt surrogates requiring only block-wise Jacobian determinants:

$$L_{\rm iPF}(\theta) = \sum_{z=g(x), x\in D} \Big[ -\log p_Z(z) + \frac12 \sum_{k\in \mathcal{P}} \log|J_k^T J_k| \Big]$$

enabling scalable, regularized density estimation on variable-rank manifolds (Cunningham et al., 2022).

Alternating (“M/D”) training:

M-flows (injective manifold flows) alternate between manifold-learning (minimizing reconstruction error via $f$) and density-learning (maximizing likelihood via $h$), decoupling geometry from density estimation for improved stability and tractability (Brehmer et al., 2020).

4. Implementation, Regularization, and Practical Considerations

Architectures employ invertible coupling layers and invertible linear transformations; regularity is enforced by

• restricting the network’s activation functions (e.g., finite polynomials in $\tanh$ (Diepeveen et al., 12 May 2025)),
• volume-control measures (constant-determinant blocks, e.g., $|\det Df| = 1$ per block), and
• Jacobian norm or Frobenius norm regularization terms (Diepeveen et al., 12 May 2025, Cramer et al., 2022).

For numerical stability, strategies such as LU-decomposed linear layers, actnorm or spectral-norm constraints, and rational-quadratic splines are used within coupling blocks (Brehmer et al., 2020). Post-hoc isometrization, e.g., reparameterizing geodesics with constant-speed in the ambient space, further corrects for metric distortions (Diepeveen et al., 12 May 2025).

Key aspects are summarized below:

Approach	Isometry Enforcement	Density Tractability
I-AE + Flow (Cramer et al., 2022)	Explicit loss ($\mathcal{L}_{\rm iso}$, $\mathcal{L}_{\rm piso}$)	Exact via isometric mapping
M-Flow (Brehmer et al., 2020)	Jacobian regularizer $R_{\rm iso}$	Change-of-variables
PF/iPF (Cunningham et al., 2022)	Orthogonality penalty $I_{s,t}$	Block determinants
Iso-NF (Diepeveen et al., 12 May 2025)	Architecture + post-hoc isometrization	Standard flow loss

These strategies offer a spectrum of trade-offs: strict isometry improves density tractability and training stability at the cost of expressivity on highly curved manifolds (Cramer et al., 2022), while looser isometry with post-hoc adjustment preserves greater flexibility (Diepeveen et al., 12 May 2025).

5. Density Estimation on Manifolds and the Change-of-Variables Formula

The induced density on the learned manifold $\mathcal{M}$ follows

$$p_\mathcal{M}(x) = p_Z(z) \left| \det (J_f(z)^T J_f(z)) \right|^{-1/2}, \quad z=f^{-1}(x)$$

For decoupled I-AE + flow, with the isometry $J_f^T J_f \approx I_d$, this reduces to the base flow’s density in latent space, facilitating density estimation as if the ambient curvature were not present (Cramer et al., 2022). In M-flows and iPFs, the normalization involves the Riemannian volume element or blockwise Jacobians, enabling density estimation and sampling directly on the manifold (Brehmer et al., 2020, Cunningham et al., 2022).

PF/iPF allows local principal manifolds (variable-rank contours) and adapts density estimation to data with non-uniform intrinsic dimensionality, leveraging volume-stretch criteria for selecting relevant latent directions at each point (Cunningham et al., 2022).

6. Empirical Results and Applications

Empirical studies demonstrate the superiority of isometry-regularized and isometric-injective flows in representing data with intrinsic manifold structure, dimensionality reduction, and generative modeling:

• Synthetic Manifolds: On S-curve or circles in $\mathbb{R}^2$ or $\mathbb{R}^3$, isometric-injective flows recover the correct manifold and density, while classical flows allocate probability mass off-manifold (Cramer et al., 2022, Brehmer et al., 2020).
• Complex Image Data: On MNIST (latent $d=16$), I-AE + RealNVP achieves lower reconstruction error and FID (46.6) than PCA-based or full-ambient NFs, and matches or exceeds W-GAN (Cramer et al., 2022). PF/iPF on MNIST align learned contour tangents with principal directions; standard injective flows collapse contours (Cunningham et al., 2022).
• Curved/Nonflat Manifolds: Strict isometry on highly curved manifolds (e.g., sphere) degrades reconstruction quality, reflecting the theoretical impossibility of global nonlinear isometries. Relaxation or trade-offs are necessary in such cases (Cramer et al., 2022).
• Physical and Particle Data: M-flows outperform standard ambient flows and PIE on LHC data for likelihood-free inference and OOD detection (Brehmer et al., 2020).
• Manifold Interpolation and PCA: Post-hoc isometrization in Iso-NF models restores constant-speed geodesics and improves low-rank approximations at the barycenter across toy and real datasets (Diepeveen et al., 12 May 2025).

7. Comparative Assessment and Open Directions

Isometric manifold learning for injective flows provides a principled, modular, and tractable methodology for modeling data with low-dimensional manifold structure, incorporating explicit geometric regularization to ensure faithful embeddings and tractable flow-based likelihoods. Key advantages include:

• Exact density evaluation under isometric (or nearly isometric) embeddings without intractable Jacobian determinants (Cramer et al., 2022).
• Decoupling of geometry learning and density estimation, yielding more stable and interpretable training (Brehmer et al., 2020, Cramer et al., 2022).
• Adaptation to local manifold geometry (PF/iPF) and variable local rank (Cunningham et al., 2022).
• Scalability and expressivity via layered invertible architectures with regularity constraints (Diepeveen et al., 12 May 2025).

A known limitation is reduced expressivity for strictly isometric flows on highly curved manifolds; further, perfect global isometry is unattainable except in special cases (e.g., flat submanifolds). For these, relaxed penalties or post-hoc geometric adjustments (iso-geodesics) provide practical alternatives (Diepeveen et al., 12 May 2025).

Ongoing developments focus on balancing regularity versus expressivity, scalable enforcement of iso-Riemannian constraints, and leveraging learned geometry for downstream tasks including clustering, interpolation, and scientific inference.