Injective Probability Flow RAE

• The paper introduces an injective probability flow RAE that bridges RAEs, VAEs, and normalizing flows by relaxing invertibility and incorporating explicit Jacobian-based regularization.
• It employs differentiable lower bounds and stochastic Jacobian-vector products for scalable training when the latent space is lower-dimensional than the ambient space.
• Empirical evaluations on datasets like MNIST, CIFAR-10, and CelebA demonstrate improved sample quality and computational efficiency with lower Fréchet Inception Distance scores.

An injective probability flow regularized autoencoder (RAE) is a generative modeling paradigm that extends flow-based and autoencoding models to settings where the latent space dimension is lower than the ambient data space, relaxing the standard invertibility constraint to injectivity. This framework introduces new objectives and training procedures derived from lower bounds on the induced probability density, resulting in scalable, tractable models with explicit Jacobian-based regularization. The injective probability flow RAE forms a bridge between regularized autoencoders, variational autoencoders (VAEs), and normalizing flows, enabling flexible manifold learning, efficient density estimation, and high-quality sample generation, notably for domains where the intrinsic data dimension is significantly less than the ambient dimensionality (Kumar et al., 2020).

1. Injective Probability Flow: Foundations and Mathematical Formulation

Central to the injective probability flow RAE is the relaxation of the usual bijectivity demand in flow-based models. Standard normalizing flows require a smooth, invertible map $g: \mathbb{R}^d \rightarrow \mathbb{R}^D$ with $d=D$ and tractable Jacobian determinants to permit both efficient forward sampling and exact likelihood computation via the change-of-variables formula:

$\ln p_x(x) = \ln p_z(h(x)) - \ln |\det J_g(h(x))|,$

where $h = g^{-1}$ and $p_z$ is a tractable prior over latents.

The injective formulation instead posits $d \ll D$ and only requires that $g$ is injective. For an infinitesimal volume $dz$ at $z$, the push-forward support lies on the $d$-dimensional manifold $g(Z) \subset \mathbb{R}^D$. The density on the manifold is given by

$$\ln p_x(x) = \ln p_z(z) - \frac{1}{2} \ln \det[J_g(z)^\top J_g(z)],$$

with $x = g(z)$. To avoid needing $g^{-1}$ at training time, an encoder $h: X \rightarrow Z$ is introduced, yielding the tractable surrogate

$$\ln p_x(x) = \ln p_z(h(x)) - \frac{1}{2} \ln \det[J_g(h(x))^\top J_g(h(x))],$$

subject to $x = g(h(x))$. The change-of-variables term now involves the locally linear volume expansion under $g$ rather than a full determinant as in the bijective case (Kumar et al., 2020).

2. Derivation of Training Objectives and Jacobian Regularization

Direct evaluation of the manifold determinant term is computationally intractable, motivating the derivation of differentiable lower bounds. By bounding $\ln s_i^2$ (the log singular values squared of $J_g$) using Jensen's inequality with a scalar $\lambda > 0$,

$$\ln s_i^2 \leq \frac{s_i^2}{\lambda} + \ln \lambda - 1,$$

and summing yields lower bounds on $\ln p_x(x)$ parameterized by $\lambda$:

• For fixed $\lambda$ ("squared-Frobenius"): the regularizer becomes proportional to $\|J_g(h(x))\|_F^2$.
• When $\lambda$ is optimized analytically ("log-Frobenius"): tighter bounds with explicit $\ln \|J_g(h(x))\|^2_F$ regularization.

These bounds motivate training objectives of the form:

$$\min_{h,g} \mathbb{E}_{x,v}\left[ \frac{1}{2\sigma^2} \|h(x)\|^2 + \mu \|x - g(h(x))\|^2 + R_\text{Jacobian} + \mu_\text{in} [\cdots] \right],$$

where $R_\text{Jacobian}$ involves either a log-Frobenius or squared-Frobenius norm of the decoder Jacobian, approximated via Hutchinson's stochastic trace estimator ($\|Jv\|^2$), and $\mu, \mu_\text{in}$ modulate reconstruction and injectivity penalties (Kumar et al., 2020).

3. Relationship to Regularized Autoencoders, VAEs, and Flow Models

The injective probability flow RAE objective generalizes and subsumes standard regularized autoencoders. In the limit of perfect reconstruction and fixed-latent prior penalty, the resulting loss is exactly a regularized autoencoder with a Frobenius-norm penalty on the decoder Jacobian. The injective model's log-likelihood lower bound becomes the RAE objective, thus providing a probabilistic interpretation and theoretical grounding for RAE regularizers.

In comparison to VAEs, the injective RAE employs an explicit deterministic encoder $h(x)$, eschews variational sampling, and acts directly on the latent prior with a Jacobian penalty, avoiding issues with the variational posterior and known problems such as variance collapse. Bijective flows, requiring equal input and latent dimension and exact log-determinant computations, are replaced by the injective flow's manifold-based approach and stochastic Jacobian-vector products, offering dimensionality reduction and significant computational savings (Kumar et al., 2020).

4. Architectural Considerations and Computational Complexity

Injective probability flow RAEs use standard convolutional encoder-decoder networks with batch normalization and ELU activations; latent dimension $d$ is selected to reflect intrinsic dataset structure (e.g., $d=32$ for MNIST, $d=128$ for CIFAR-10 and CelebA). Unlike bijective flows—whose determinant computation and invertibility checks scale as $\mathcal{O}(D^2)$ to $\mathcal{O}(D^3)$—injective flows require only $\mathcal{O}(Dd)$ per sample for Jacobian-vector products and $\mathcal{O}(D)$ for reconstruction, yielding significant efficiency improvements particularly when $d \ll D$ (Kumar et al., 2020).

A related line of work exploits isometric regularization—enforcing that $J_G(z)^\top J_G(z) = I_d$ through a dedicated penalty—in "isometric autoencoder + normalizing flow" constructions (Cramer et al., 2022). This further simplifies density estimation, decoupling manifold learning from latent density modeling and eliminating determinants from the likelihood computation, simplifying both optimization and hyperparameter selection.

5. Empirical Evaluation and Performance Characteristics

Experiments on MNIST, CIFAR-10, and CelebA demonstrate that injective probability flow RAEs achieve superior performance over standard AEs, VAEs, and several RAE variants when measured by the Fréchet Inception Distance (FID) on both reconstructions and sample generation. On high-dimensional datasets (CelebA, CIFAR-10), injective models outperform baselines by wide FID margins (10–20 points), while on MNIST their performance is competitive with or slightly below the best $\beta$-VAE configurations.

Qualitative analysis reveals that injective flows produce reconstructions with sharper and more detailed features compared to VAEs, albeit sometimes at the cost of minor artifacts. This supports the claim that Jacobian regularization and dimensional reduction within the injective framework preserves geometrical richness on learned manifolds and enhances generative fidelity (Kumar et al., 2020).

6. Connections to Isometric Manifold Learning and Injective Normalizing Flows

Subsequent research develops injective normalizing flow models that further emphasize isometric embedding, wherein the decoder's Jacobian is strongly regularized to be close to orthonormal ($J_G(z)^\top J_G(z) \approx I_d$). The density is decoupled as $p_X(x) = p_Z(E(x)) \cdot |\det \nabla_x E(x)|$, and if the isometry holds exactly, the Jacobian term drops out, leaving only a latent-space flow density.

Contrasting with prior injective flows such as M-Flow and Trumpets—which require costly determinant calculations per example—these models split training into two stages: RAE-based manifold fitting and normalizing flow fitting in latent space. This separation avoids difficult joint optimization and eliminates reconstruction-likelihood trade-offs, leading to tractable, interpretable models and efficient sampling on manifolds in high-dimensional ambient spaces (Cramer et al., 2022).

7. Impact, Limitations, and Future Prospects

The injective probability flow RAE enables scalable manifold learning, tractable density estimation, and efficient sampling in generative modeling when the true data support is of lower intrinsic dimension. Its explicit connection between RAEs and probabilistic manifold learning unifies several strands of generative modeling theory and offers practical training benefits. Limitations include reliance on the smoothness and injectivity of $g$, computational cost for large $D$ and $d$ (though reduced compared to bijective flows), and the need to estimate or regularize the Jacobian spectrum robustly.

Future advances may further tighten the connection between manifold geometry and tractable density estimation, develop improved regularization or architectural schemes for high-complexity data, and expand applications in domains where structured manifold support is characteristic, such as molecular data, 3D scenes, and complex image manifolds.

---

Key References:

• "Regularized Autoencoders via Relaxed Injective Probability Flow" (Kumar et al., 2020)
• "Nonlinear Isometric Manifold Learning for Injective Normalizing Flows" (Cramer et al., 2022)