Isometric Manifold Learning for Injective Normalizing Flows

Eike Cramera,b, Felix Rauha,c,d, Alexander Mitsose,a,b, Raúl Temponef,g,
Manuel Dahmena,∗

  * a

Forschungszentrum Jülich GmbH, Institute of Energy and Climate Research,
Energy Systems Engineering (IEK-10), Jülich 52425, Germany

  * b

RWTH Aachen University, Process Systems Engineering (AVT.SVT), Aachen 52074,
Germany

  * c

RWTH Aachen University, Aachen 52062, Germany

  * d

KU Leuven, Research Center for Operations Management, 3000 Leuven, Belgium

  * e

JARA-ENERGY, Jülich 52425, Germany

  * f

RWTH Aachen University, Chair of Mathematics for Uncertainty Quantification,
Aachen 52062, Germany

  * g

King Abdullah University of Science and Technology (KAUST), Computer,
Electrical, and Mathematical Sciences & Engineering Division (CEMSE), Saudi
Arabia

To model manifold data using normalizing flows, we employ isometric
autoencoders to design embeddings with explicit inverses that do not distort
the probability distribution. Using isometries separates manifold learning and
density estimation and enables training of both parts to high accuracy. Thus,
model selection and tuning are simplified compared to existing injective
normalizing flows. Applied to data sets on (approximately) flat manifolds, the
combined approach generates high-quality data.
Keywords: image generation; machine learning; multivariate statistics;
nonparametric statistics .

Isometric Manifold Learning for Injective Normalizing Flows

Eike Cramera,b, Felix Rauha,c,d, Alexander Mitsose,a,b, Raúl Temponef,g,
Manuel Dahmena,∗

  * a

Forschungszentrum Jülich GmbH, Institute of Energy and Climate Research,
Energy Systems Engineering (IEK-10), Jülich 52425, Germany

  * b

RWTH Aachen University, Process Systems Engineering (AVT.SVT), Aachen 52074,
Germany

  * c

RWTH Aachen University, Aachen 52062, Germany

  * d

KU Leuven, Research Center for Operations Management, 3000 Leuven, Belgium

  * e

JARA-ENERGY, Jülich 52425, Germany

  * f

RWTH Aachen University, Chair of Mathematics for Uncertainty Quantification,
Aachen 52062, Germany

  * g

King Abdullah University of Science and Technology (KAUST), Computer,
Electrical, and Mathematical Sciences & Engineering Division (CEMSE), Saudi
Arabia

To model manifold data using normalizing flows, we employ isometric
autoencoders to design embeddings with explicit inverses that do not distort
the probability distribution. Using isometries separates manifold learning and
density estimation and enables training of both parts to high accuracy. Thus,
model selection and tuning are simplified compared to existing injective
normalizing flows. Applied to data sets on (approximately) flat manifolds, the
combined approach generates high-quality data.
Keywords: image generation; machine learning; multivariate statistics;
nonparametric statistics

##  1 Introduction

Normalizing flows are deep generative models (DGM) that represent the
probability distribution of high-dimensional data sets as a change of
variables of a multivariate Gaussian [1, 2]. Using the inverse of this
transformation, normalizing flows can compute the probability density
functions (PDFs) explicitly, thus enabling training via the statistically
consistent and asymptotically efficient [3] likelihood maximization [4, 5].

Despite the advantage of likelihood maximization, normalizing flows often
result in poor fits of the target probability distributions [6, 7]. The poor
fits often result from the structural setup of the normalizing flow that is
ill-posed to describe probability distributions over lower-dimensional
manifolds [6, 7]. As the probability distributions of many high-dimensional
data sets reside on lower-dimensional manifolds [8, 9], the inability of
normalizing flows to fit such probability distributions rescinds the
advantages gained from the direct likelihood maximization. Advanced
normalizing flow models like Glow [10] and FFJORD [11] can partially
compensate for this issue by relying on modern model architectures.

Approaches to training normalizing flows on manifold data substitute the
bijective transformation with an injective transformation, i.e., a pseudo
invertible transformation that maps the data to a lower-dimensional Gaussian
[12, 6]. Such injective normalizing flows then combine manifold learning with
density estimation in a single model and, thus, a simultaneous training
approach. Some of the published approaches assume a dimensionality-reducing
map to be known and available a priori [12, 13]. Other works use compositions
of manifold learning models and normalizing flows that are trained
simultaneously, e.g., the $\bm{\mathcal{M}}$-Flow [6], Noisy Injective Flows
[14], piecewise injective flows called Trumpets [15], and neural manifold
ordinary differential equations [16]. Approaches without dimensionality
reduction use inflation with noise to overcome the manifold structure [17, 18,
19].

Most injective normalizing flows do not support exact PDF tractability and
instead require PDF approximations [6], approximations of the inverse
transformation [14], or approximations of the Jacobian determinant terms [20].
Notable exceptions are the conformal embedding flow [21] and the principal
manifold flow [22] that maintain tractable PDFs. Most injective normalizing
flows are designed to perform manifold learning and density estimation using
the same transformation. Hence, likelihood maximization and manifold learning
have to be trained either simultaneously [21] or iteratively [6]. Both the
simultaneous and the iterative training lead to a difficult tuning problem as
the loss function has to balance two uneven objectives: The likelihood
maximization, which is unbounded [3], and the reconstruction loss of the
manifold learning that has a theoretical global optimum of zero.

In this work, we propose a combination of isometric embeddings and normalizing
flows that rescinds the issue of the combined loss function. We show
theoretically that the PDF described by an injective normalizing flow is
invariant to isometric embeddings. Thus, isometric embeddings allow us to
separate the embedding and the normalizing flow training into two distinct
problems, and to train both model parts to the highest attainable accuracy.
Meanwhile, the composition still allows for explicit sampling as well as
explicit and efficient PDF computation.

This work presents a generalization of our previous work on principal
component analysis (PCA)-based dimensionality reduction [23] to nonlinear
embeddings. In particular, we use the isometric autoencoder (I-AE) [24] that
enforces the isometry via regularization. A family of strictly isometric
embeddings would be limited to learning flat manifolds [25] and, the
regularization approach can balance any trade-offs between the embedding and
the learned density [26]. Notably, the proposed method does not aim to improve
on the expressivity of other injective flows like the $\bm{\mathcal{M}}$-Flow
[6]. In contrast, our approach simplifies model tuning and training.

The remainder of this paper is organized as follows: In Section 2, we briefly
review the basic principles of normalizing flows, present the general concept
of injective normalizing flows, and explicate the difficulties associated with
the combination of manifold learning and density estimation. Then, we show how
isometric embeddings allow for a seamless separation of the two learning
problems. Section 3 reviews practical designs of isometric embeddings. In
Section 4, we apply the injective normalizing flows to generate samples of the
artificial S-curve data set and the MNIST image data set of handwritten
digits. Furthermore, Section 4 highlights the limitation of isometric
embeddings by the example of a distribution on a spherical surface. Finally,
Section 5 concludes our work.

##  2 Isometric embedding normalizing flows

This Section first reviews the basic principles of normalizing flows in
Section 2.1, and then states the general loss function for an injective
normalizing flow with combined manifold learning in Section 2.2. Section 2.3
proposes a separation of manifold learning and density estimation by using
isometric embeddings. Finally, Section 2.4 discusses the limitations of
isometric embeddings.

###  2.1 Normalizing flows

Normalizing flows model the probability distribution of a $D$-dimensional
multivariate random variable $X$ through a change of variables of a
multivariate standard Gaussian [2, 5]. The change of variables utilizes a
diffeomorphism, i.e., a bijective and differentiable transformation,
$T(\cdot):\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$. Thus, each point
$\mathbf{x}$ in the data probability distribution is uniquely mapped to a
point $\mathbf{z}$ in the Gaussian and vice versa. The links are given by the
forward $T(\cdot)$ and inverse $T^{-1}(\cdot)$ transformation, respectively.
Exploiting the diffeomorphism, the PDF $p_{X}(\mathbf{x})$ can be calculated
explicitly using the change of variables formula [5]:

$p_{X}(\mathbf{x})=\phi\left(\mathbf{z}\right)\left|\det\mathbf{J}_{T}(\mathbf{z})\right|^{-1}$
(1)

Here, $\mathbf{J}_{T}$ is the Jacobian of the transformation $T$, and
$\phi(\mathbf{z})$ is the Gaussian PDF. Given a data set $\mathcal{X}$ and a
trainable diffeomorphism $T_{\bm{\theta}}$ with parameters $\bm{\theta}$, the
logarithm of Equation (1) can be used to train the normalizing flow via direct
log-likelihood maximization [5].

###  2.2 Injective normalizing flows

Given a $d$-dimensional manifold $\mathcal{M}\subset\mathbb{R}^{D}$ with
$d<D$, let $\psi(\cdot):\mathbb{R}^{d}\rightarrow\mathcal{M}$ be a decoder
function that maps the data from the latent space $\mathbf{z}$ to the data
space $\mathbf{x}$ and let
$\psi^{-1}(\cdot):\mathbb{R}^{D}\rightarrow\mathbb{R}^{d}$ be an encoder
function:

$\displaystyle\mathbf{x}$ $\displaystyle=\psi(\mathbf{z})$
$\displaystyle\mathbf{z}$ $\displaystyle=\psi^{-1}(\mathbf{x})$

Then, $\psi^{-1}$ is a pseudo-inverse of $\psi$, i.e., an exact inverse of
$\psi$ for $\mathbf{x}\in\mathcal{M}$. Since using encoder and decoder
functions results in non-square Jacobians, the change of variables has to be
set up in a generalized form [12]:

$\displaystyle
p_{X}(\mathbf{x})=\phi(\mathbf{z})\left|\det\left(\mathbf{J}_{\psi}(\mathbf{z})^{T}\mathbf{J}_{\psi}(\mathbf{z})\right)\right|^{-0.5}$
(2)

Note that the generalized change of variables formula in Equation (2) is
typically prohibitively expensive to compute during training.

If the injective normalizing flow is used for both density estimation and
manifold learning, the reconstruction loss has to be considered in addition to
the likelihood maximization. The combined loss function then is the
expectation of the two parts over the distribution of the random variable $X$
[6]:

$\mathcal{L}=\mathbb{E}_{X}\left[-\log\leavevmode\nobreak\
p_{X}(\mathbf{x})+\beta\leavevmode\nobreak\
||\mathbf{x}-\psi(\psi^{-1}(\mathbf{x}))||_{2}\right]$ (3)

In Equation (3), $||\cdot||_{2}$ is the Euclidean norm and $\beta$ is a
scaling hyperparameter to balance the two loss function components. In
general, injective normalizing flows require manifold learning and likelihood
maximization to be performed either simultaneously [15, 21] or iteratively
[6]. However, this balancing is particularly difficult, as the reconstruction
loss $||\mathbf{x}-\psi(\psi^{-1}(\mathbf{x}))||_{2}$ is bounded from below by
zero, while the likelihood loss $-\log\leavevmode\nobreak\ p_{X}(\mathbf{x})$
is theoretically unbounded. Furthermore, the absolute values of the likelihood
vary by orders of magnitude between data sets, which makes tuning the
hyperparameter $\beta$ difficult.

###  2.3 Isometric embeddings

As an alternative to performing both manifold learning and density estimation
with a single model, we propose a separation of the two objectives into a
composition of two distinct learning problems. By leveraging isometric
embeddings, we can design dimensionality reduction schemes without distorting
the probability distribution [24]. The full transformation $\psi$ then
consists of an injective transformation
$f(\cdot):\mathbb{R}^{d}\rightarrow\mathcal{M}$ with pseudo-inverse
$g(\cdot):\mathcal{M}\rightarrow\mathbb{R}^{d}$ and a diffeomorphism
$T(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ in the lower-dimensional
latent space:

$\displaystyle\mathbf{x}$ $\displaystyle=\psi(\mathbf{z})=f\circ
T(\mathbf{z})$ $\displaystyle\mathbf{z}$
$\displaystyle=\psi^{-1}(\mathbf{x})=T^{-1}\circ g(\mathbf{x})$

Applying the chain rule of differentiation to Equation (2), the change of
variables of the composition is given by [6]:

$\displaystyle p_{X}(\mathbf{x})=$ $\displaystyle\leavevmode\nobreak\
\phi(\mathbf{z})\left|\det\mathbf{J}_{T}(\mathbf{z})\right|^{-1}$ (4)
$\displaystyle\left|\det\left(\mathbf{J}_{f}\left(T(\mathbf{z})\right)^{T}\mathbf{J}_{f}\left(T(\mathbf{z})\right)\right)\right|^{-0.5}$

Options on how to mitigate the computational cost of Equation (4) typically
use approximate expressions for
$\det\left(\mathbf{J}_{f}\left(T(\mathbf{z})\right)^{T}\mathbf{J}_{f}\left(T(\mathbf{z})\right)\right)$
[6].

In our previous work [23], we used PCA encodings and highlighted that the PDF
in Equation (4) is invariant to a PCA encoding. By definition, Equation (4) is
invariant to any other isometric embedding as well. Mathematically, this can
be expressed as

$\mathbf{J}_{f}(\mathbf{z})^{T}\mathbf{J}_{f}(\mathbf{z})=\mathbf{I}_{d}\quad\forall\mathbf{z}\in\mathbb{R}^{d},$

where $\mathbf{I}_{d}$ is the $d$-dimensional identity matrix. As the
determinant of the identity matrix equals one, the PDF in Equation (4) is
invariant to the isometric embedding. Therefore, the PDF for any isometric
embedding can be simplified to:

$p_{X}(\mathbf{x})=\phi(\mathbf{z})\left|\det\mathbf{J}_{T}(\mathbf{z})\right|^{-1}$
(5)

A visual representation of the two-step composition is shown in Figure 1.

Figure 1:  Injective normalizing flow designed as a composition of a
diffeomorphism $T:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ and an isometric
embedding $f:\mathbb{R}^{d}\rightarrow\mathcal{M}$ with pseudoinverse
$g:\mathcal{M}\rightarrow\mathbb{R}^{d}$.

In conclusion, the isometric embedding can be trained independently from the
diffeomorphic normalizing flow without any adjustments to the normalizing flow
training, which avoids the need to balance the objectives of likelihood
maximization and reconstruction loss described in Equation (3) and allows us
to train both models to high accuracy.

###  2.4 Limitations of isometric embeddings

Restricting the manifold learning problem to strict isometries limits the
expressivity of the embeddings. The Nash embedding theorem [27] states that
isometric embeddings can encode $d$-dimensional flat Riemannian manifolds
embedded in $D$-dimensional ambient space arbitrarily well if $D\geq d+1$
[24]. Thus, the decoder $f$ can be designed to be isometric and the encoder
$g$ to be a pseudoinverse of $f$ without loss of information in the embedding,
if and only if the considered manifold is flat [25]. We refer to [26] for a
more detailed discussion of the trade-offs between the accuracy and the
isometry of the learned embedding for non-flat manifolds.

##  3 Isometric embeddings

In this Section, we review two possible isometric embeddings and discuss the
selection of the latent dimensionality. Note that the use of PCA as affine
isometric embedding for normalizing flows was already discussed in our
previous work [23].

###  3.1 Affine isometric embedding with PCA

For PCA, both the encoder $g_{\text{PCA}}$ and the decoder $f_{\text{PCA}}$
are affine transformations [28]:

$\displaystyle\mathbf{x}$
$\displaystyle=f_{\text{PCA}}(\mathbf{z})=\mathbf{V}_{d}\mathbf{z}+\bm{\mu}_{\mathcal{X}}$
$\displaystyle\mathbf{z}$
$\displaystyle=g_{\text{PCA}}(\mathbf{x})=\mathbf{V}_{d}^{T}\left(\mathbf{x}-\bm{\mu}_{\mathcal{X}}\right)$

Here, $\bm{\mu}_{\mathcal{X}}$ is the empirical mean of the dataset
$\mathcal{X}$ and $\mathbf{V}_{d}\in\mathbb{R}^{D\times d}$ are the first $d$
column vectors of the matrix of right singular vectors
$\mathbf{V}\in\mathbb{R}^{D\times D}$ that stems from the singular value
decomposition (SVD) of the empirical covariance matrix. Because PCA is an
affine transformation, $\mathbf{V}_{d}$ is the constant Jacobian of the
decoding transformation $f_{\text{PCA}}$, hence:
$\mathbf{J}_{f_{\text{PCA}}}=\mathbf{V}_{d}=\text{const}$. By design of the
SVD, $\mathbf{V}_{d}$ has orthonormal column vectors and thus:

$\mathbf{J}_{f_{\text{PCA}}}^{T}\mathbf{J}_{f_{\text{PCA}}}=\mathbf{V}_{d}^{T}\mathbf{V}_{d}=\mathbf{I}_{d}$

For PCA, the latent dimensionality $d$ equals the number of considered
principal components. We base the selection of $d$ on the explained variance
ratio, i.e., the singular values $\sigma_{k}$ of the empirical covariance
matrix normalized by their total sum. The sum over the explained variance
ratio of the considered principal components then is the cumulative explained
variance (CEV):

$\text{CEV}=\frac{\sum_{k=1}^{d}\sigma_{k}}{\sum_{i=1}^{D}\sigma_{i}}100\%$

In practice, we select $d$ such that the CEV is greater or equal to the
desired percentage, e.g., 99%.

###  3.2 Nonlinear isometric embeddings with I-AE

As a possible nonlinear isometric embedding, we utilize the isometric
autoencoder (I-AE) proposed by Gropp et al. [24]. Notable alternatives are the
regularized autoencoders in [26] and the local conformal autoencoders in [29].

The I-AE uses a nonlinear autoencoder and specifically enforces the isometry
during training via penalization. Both encoder $g_{\text{I-AE}}$ and decoder
$f_{\text{I-AE}}$ are deep neural networks
$\textbf{G}(\mathbf{x};\bm{\theta}_{\textbf{G}})$ and
$\textbf{F}(\mathbf{z};\bm{\theta}_{\textbf{F}})$, respectively:

$\displaystyle\mathbf{x}$
$\displaystyle=f_{\text{I-AE}}(\mathbf{z})=\textbf{F}(\mathbf{z};\bm{\theta}_{\textbf{F}})$
$\displaystyle\mathbf{z}$
$\displaystyle=g_{\text{I-AE}}(\mathbf{x})=\textbf{G}(\mathbf{x};\bm{\theta}_{\textbf{G}})$

Here, $\bm{\theta}_{\textbf{F}}$ and $\bm{\theta}_{\textbf{G}}$ are the
parameters of decoder and encoder, respectively. To minimize the loss of
information, the loss function penalizes the mean-squared-error in the
reconstruction:

$\mathcal{L}_{\text{AE}}(\bm{\theta}_{\textbf{G}},\bm{\theta}_{\textbf{F}})=\mathbb{E}_{X}\left[||\mathbf{x}-\textbf{F}\left(\textbf{G}(\mathbf{x})\right)||_{2}\right]$
(6)

Furthermore, the training has to promote the isometry of the decoder:

$\mathbf{J}_{\textbf{F}}(\mathbf{z})^{T}\mathbf{J}_{\textbf{F}}(\mathbf{z})=\mathbf{I}_{d}\quad\forall\mathbf{z}\in\mathbb{R}^{d}$
(7)

Computing the expression in Equation (7) requires expensive matrix
multiplications of the Jacobian and its transpose. To avoid these, Gropp et
al. [24] propose enforcing the isometry condition implicitly by penalizing the
change in length of a vector to the surface of the unit ball
$\mathbf{u}\in\mathbb{S}^{d-1}=\\{\mathbf{y}\in\mathbb{R}^{d}:||\mathbf{y}||_{2}=1\\}$,
i.e., a unit vector, after multiplication with the Jacobian:

$\mathcal{L}_{\text{iso}}(\bm{\theta}_{\textbf{F}})=\mathbb{E}_{Z,U}\left[\left(||\mathbf{J}_{\textbf{F}}(\mathbf{z})\mathbf{u}||_{2}-1\right)^{2}\right]$
(8)

Here, $\mathbb{E}_{Z,U}$ is the expectation over the latent space
$\mathbf{z}\in\mathbb{R}^{d}$ and the space of all unit vectors
$\mathbf{u}\in\mathbb{S}^{d-1}$. In addition to the isometry loss of the
decoder, the isometry of the encoder (pseudoinverse) has to be promoted as
well:

$\mathcal{L}_{\text{piso}}(\bm{\theta}_{\textbf{G}})=\mathbb{E}_{X,U}\left[\left(||\mathbf{u}^{T}\mathbf{J}_{\textbf{G}}(\mathbf{x})||_{2}-1\right)^{2}\right]$
(9)

Here, $\mathbb{E}_{X,U}$ is the expectation over the manifold
$\mathbf{x}\in\mathcal{M}$ and the space of all unit vectors
$\mathbf{u}\in\mathbb{S}^{d-1}$. Finally, the three loss functions are
combined to a total loss using weights $\lambda_{\text{iso}}$ and
$\lambda_{\text{piso}}$ for $\mathcal{L}_{\text{iso}}$ and
$\mathcal{L}_{\text{piso}}$, respectively:

$\displaystyle\mathcal{L}(\bm{\theta}_{\textbf{G}},\bm{\theta}_{\textbf{F}})=$
$\displaystyle\mathcal{L}_{\text{AE}}(\bm{\theta}_{\textbf{G}},\bm{\theta}_{\textbf{F}})$
(10)
$\displaystyle+\lambda_{\text{iso}}\mathcal{L}_{\text{iso}}(\bm{\theta}_{\textbf{F}})+\lambda_{\text{piso}}\mathcal{L}_{\text{piso}}(\bm{\theta}_{\textbf{G}})$

Unlike PCA, the I-AE does not come with a recipe to estimate the latent
dimensionality $d$, i.e., $d$ must be determined via empirical studies [24] or
via advanced machine noise inflation techniques [30]. For more details on the
implementation, selection of the weights $\lambda_{\text{iso}}$ and
$\lambda_{\text{piso}}$, and the computation of the loss functions, we refer
to the original I-AE publication [24].

##  4 Numerical experiments

We apply the proposed combination of isometric embeddings and normalizing
flows to learn and sample from the probability distributions of three
different data sets. Specifically, we analyze the generated samples from
standard full-space normalizing flows and isometric embedding flows compared
to the real data. The three data sets are the synthetic S-Curve data set
provided by the machine learning library scikit-learn [31], a two-dimensional
Gaussian distribution on a spherical surface, and the MNIST data set with
images of handwritten digits [32]. For all normalizing flow models, we employ
the commonly used RealNVP affine coupling layer architecture [4] as the
invertible neural network. All models are implemented in TensorFlow, version
2.5.0 [33].

###  4.1 S-Curve

The S-Curve data set describes a two-dimensional, nonlinear, flat manifold
embedded in three-dimensional Euclidean space (see left of Figure 2).
Following the training methodology derived in Section 2, we first train the
I-AE using Equation (10) and then train a two-dimensional RealNVP model.

Figure 2:  Encoding and decoding of three-dimensional S-Curve data (1000
samples) [31] using I-AE [24] with
$\lambda=\lambda_{\textbf{iso}}=\lambda_{\textbf{piso}}=1.0$. Test data
$\mathbf{x}\in\mathcal{X}_{\text{test}}$ (left); reconstruction of test data
$f(g(\mathbf{x}))\leavevmode\nobreak\
\forall\mathbf{x}\in\mathcal{X}_{\text{test}}$ (center); two-dimensional
latent space of I-AE $g(\mathbf{x})\leavevmode\nobreak\
\forall\mathbf{x}\in\mathcal{X}_{\text{test}}$ (right). Color coding is based
on the position of test data.

Figure 2 shows the reconstruction (center) and the latent space (right) of the
trained I-AE for a test set $\mathcal{X}_{\text{test}}$ of 1000 samples. The
results in Figure 2 clearly show that the encoding leads to an unfolding of
the three-dimensional S-shape to a two-dimensional rectangle and that the full
S-Curve is recovered by the decoder. Table 1 shows the final training and test
losses for the total and the three constituents of the loss function.

Table 1:  Final training loss and test loss for I-AE training on the S-Curve
data set [31]. The constituents of the loss function are defined in Equations
(10), (6), (8), and (9).

| $\mathcal{L}$ | $\mathcal{L}_{\text{AE}}$ | $\mathcal{L}_{\text{iso}}$ | $\mathcal{L}_{\text{piso}}$
---|---|---|---|---
Training | 2.21$\times 10^{-4}$ | 1.55$\times 10^{-4}$ | 4.78$\times 10^{-5}$ | 1.87$\times 10^{-5}$
Test | 2.11$\times 10^{-4}$ | 1.54$\times 10^{-4}$ | 3.81$\times 10^{-5}$ | 1.89$\times 10^{-5}$

To verify that the embedding is in fact isometric, we compute the Jacobian
determinant of the decoder. For the 1000-sample test set
$\mathcal{X}_{\text{test}}$, the expectation of the Jacobian determinant is:

$\mathbb{E}_{\mathcal{X}_{\text{test}}}\left[\left|\det\mathbf{J}_{\textbf{F}}(\textbf{G}(\mathbf{x}))^{T}\mathbf{J}_{\textbf{F}}(\textbf{G}(\mathbf{x}))\right|^{-0.5}\right]=1.0096$

Thus, the PDF in Equation (5) only deviates by a factor of 1.0096 from an
exact isometry.

Figure 3:  1000 generated S-Curve samples [31]. Test data
$\mathbf{x}\in\mathcal{X}_{\text{test}}$ (left); samples from RealNVP in
reduced space with I-AE embedding (center); samples from RealNVP in full space
(right). The colors are based on the z-axis of the respective data points.

Figure 3 shows sampled data of the normalizing flow with I-AE embedding
(center) in comparison to a full-space normalizing flow (right). The results
show that the manifold shape is correctly recovered by the injective flow,
whereas the full-space normalizing flow fails to reconstruct the shape of the
S-Curve distribution. Note that we do not show the absolute PDF values of the
injective (Equation (5)) and full space normalizing flow (Equation (1)), as
they describe the PDFs of two systems with different dimensionality and are,
therefore, not comparable.

The results in Figure 2 and Table 1 show that the I-AE learns embeddings with
very high accuracy. Furthermore, the learned embedding is almost isometric,
which justifies its usage in an injective normalizing flow as proposed in
Section 2. This finding is confirmed by the generated data in Figure 3 from
the combination of I-AE and RealNVP, which reconstructs the true distribution
with only a few outliers.

###  4.2 Spherical surface

We highlight the trade-off between obtaining isometric embeddings and low
reconstruction losses via a generated Gaussian distribution on a spherical
surface. The spherical surface data is sampled from a two-dimensional Gaussian
distribution with a random covariance matrix and projected to partially cover
a three-dimensional spherical surface. The two-dimensional Gaussian is
parameterized by a zero mean vector and a randomly sampled covariance matrix:

$\mathbf{z}\sim\mathcal{N}\left(\mathbf{z};\begin{bmatrix}0.0\\\
0.0\end{bmatrix},\begin{bmatrix}0.548&0.602\\\
0.602&0.544\end{bmatrix}\right)$

The projection onto the spherical surface is done via stereographic
projection:

$\begin{bmatrix}x_{1}\\\ x_{2}\\\
x_{3}\end{bmatrix}=\frac{1}{1+z_{1}^{2}+z_{2}^{2}}\begin{bmatrix}2z_{1}\\\
2z_{2}\\\ -1+z_{1}^{2}+z_{2}^{2}\end{bmatrix}$

Figure 4:  Pareto curve of reconstruction loss (MSE-loss, Equation (6)) and
isometry loss (ISO-loss, Equation (8)). The parameter
$\lambda=\lambda_{\textbf{iso}}=\lambda_{\textbf{piso}}$ is the weighting
parameter between reconstruction loss and isometry loss (see Equation (10)).
Dots are validation losses for different values of $\lambda$ and dashed lines
are linear interpolations of the Pareto curve for illustration purposes.

Figure 4 shows the Pareto front of the reconstruction loss (Equation (6)) and
the isometry loss (Equation (8)) for different values of the weighting factor
$\lambda=\lambda_{\text{iso}}=\lambda_{\text{piso}}$ and a test set
$\mathcal{X}_{\text{test}}$ of 1000 samples. The Pareto front shows that
neither of the two loss functions can be improved without worsening the other
indicating a trade-off between the two objectives.

Figure 5:  Gaussian distribution on a spherical surface. Ground truth (left),
I-AE trained with $\lambda=10^{-3}$ (center), I-AE trained with $\lambda=10$
(right). The parameter
$\lambda=\lambda_{\textbf{iso}}=\lambda_{\textbf{piso}}$ is the weighting
factor between reconstruction loss and isometry loss (see Equation (10).

Figure 5 shows samples of the true distribution (left), a composition of
RealNVP, and an I-AE with low weighting factor ($\lambda=0.01$) on the
isometry (center) and with high weighting factor ($\lambda=100$) on the
isometry (right), respectively. The I-AE with $\lambda=0.01$ learns the shape
of the spherical surface well but fails to learn the probability density on
the surface. The I-AE with $\lambda=100$ fails to learn the surface and, thus,
the composition fails to learn the distribution.

Table 2: Average residual of surface equation $R=|-1+\sum_{i=1}^{3}x_{i}^{2}|$ for 1000 samples from compositions of RealNVP [4] and I-AE [24] for different $\lambda=\lambda_{\textbf{iso}}=\lambda_{\textbf{piso}}$.  $\lambda$ | 0.01 | 0.1 | 0.5 | 1 | 5 | 10 | 100
---|---|---|---|---|---|---|---
$R$ | 0.03 | 0.05 | 0.11 | 0.12 | 0.19 | 0.29 | 0.29

Table 2 shows the average residuals of the surface equation for sampled data
using the isometric manifold normalizing flow with different values of
$\lambda$. The trend shows that the representation of the spherical surface
becomes worse with increasing $\lambda$. In fact, the average residual
increases by an order of magnitude by changing $\lambda$ from 0.01 to 100.

The different extremes in Figure 5 and the residuals in Table 2 show the
inability of the I-AE to deal with the curvature of the spherical surface and
underline that isometries cannot transform data from a lower-dimensional
latent space to a non-flat manifold. Thus, the trade-off between representing
the true shape of the manifold and finding accurate probability density
functions must be considered in applications to real data.

###  4.3 MNIST

Finally, we consider the MNIST data set of handwritten digits [32] with
$28\times 28$ pixel resolution. The MNIST manifold is nonlinear and not
strictly flat [26]. Lee et al. [26] discuss the trade-off between the isometry
and the accuracy of the encoding of the MNIST data set. For both the I-AE and
PCA, we use latent dimensionalities of 16 based on the original I-AE paper
[24].

Figure 6:  I-AE and PCA reconstruction after encoding of MNIST data set of
handwritten digits [32]. Row 1: Samples from the test set, row 2:
Reconstruction after encoding with I-AE, row 3: Reconstruction after encoding
with PCA. Latent dimensionalities for I-AE and PCA are 16 based on the
original I-AE paper [24].

Figure 6 shows a random selection of images from the MNIST test set (first
row) as well as the reconstruction of these images after encoding using I-AE
(second row) and PCA (third row), respectively. For all test images in Figure
6, the I-AE is able to preserve the images well. The PCA reconstruction shows
a loss of features, blurred digits, and a significant amount of noise around
the digits.

Table 3: Reconstruction loss of PCA and I-AE for MNIST test data set [32]. The table shows the mean-squared-error (MSE) and the mean-absolute-error (MAE) with data scaled to $[0,1]$. | I-AE | PCA
---|---|---
MSE | 8.28$\times 10^{-3}$ | 2.69$\times 10^{-2}$
MAE | 3.20$\times 10^{-2}$ | 8.81$\times 10^{-2}$

In addition to the visual comparison in Figure 6, Table 3 lists the mean-
squared-error (MSE) and mean-absolute-error (MAE) for both PCA and I-AE on the
MNIST test set. With respect to both metrics, the I-AE achieves lower values
compared to the PCA. Notably, the MSE for the I-AE encoding is low despite the
MNIST manifold being not strictly flat. Thus, the nonlinear I-AE shows a
better ability to encode the data compared to the affine PCA and overall good
performance despite the isometry regularization.

Figure 7:  Generated images of MNIST data set of handwritten digits [32]. Row
1: Random samples of the I-AE and RealNVP composition, row 2: Random samples
of the PCA and RealNVP composition, row 3: random samples from a full-space
RealNVP, row 4: random samples from a Wasserstein GAN [34]. Latent
dimensionalities for I-AE and PCA are 16.

We train different RealNVP models on samples encoded using I-AE and PCA as
well as the full images. Figure 7 shows samples from the combined isometric
embedding normalizing flows in comparison to a full $28\times 28$-dimensional
normalizing flow. For reference, we also include images generated using a
Wasserstein-GAN (W-GAN) [34], which is often the method of choice for image
generation [35]. Note that the W-GAN was significantly more difficult to train
and fine-tune compared to the isometric embedding normalizing flows. Hence,
the presented W-GAN results are the best of the 20 trained models. Meanwhile,
both isometric normalizing flow compositions converge to the same result in
almost every iteration.

The generated images from the composition of I-AE and RealNVP (row 1) appear
like handwritten digits from the original MNIST data set. The images generated
from the RealNVP with PCA embedding (row 2) resemble handwritten digits to
some extent. However, as observed in Figure 6 the images are blurred, and the
digits are difficult to read or ambiguous. The full-space normalizing flow-
generated images (row 3) are very noisy and do not resemble handwritten digits
at all, i.e., the full-space RealNVP was unable to learn the probability
distribution of the MNIST data set. The W-GAN-generated images (row 4) show
lines with high contrast that resemble digits to some extent. However, these
digits are often incomplete or ambiguous. Apparently, the discriminator model
fails to identify at least some of the images in row 4 as fake.

In addition to the visual analysis (cf. Figure 7), we compute the inception
score (IS) [36] and the Fréchet inception distance (FID) [37] for each variant
of the normalizing flow and the W-GAN. Both IS and FID are quantitative
measures that analyze both the quality as well as the diversity of the
generated images. While the IS considers generated samples only, the FID also
includes information from real data. In general, the IS is a positively
oriented score, i.e., high scores indicate high-quality data sets and the FID
is a negatively oriented score. We use a classifier model with $>$99.3%
accuracy. As recommended by Salimans et al. [36], we generated 50000 samples
from each normalizing flow variant to compute both IS and FID. For the IS, we
batch the data set and compute the empirical mean and standard deviation.
Table 4 lists IS and FID for the generated image data sets.

Table 4:  Inception score (IS) [36] and Fréchet inception distance (FID) [37] for MNIST test set $\mathcal{X}_{\text{Test}}$ and generated images from I-AE and PCA compositions with RealNVP, full-space RealNVP (FSNF), and Wasserstein-GAN (W-GAN). The number of generated images is 50000, and the mean and standard deviation for the IS are computed over 10 batches. A high IS and low FID indicate good scores.  | IS | FID
---|---|---
$\mathcal{X}_{\text{Test}}$ | 11.39 $\pm$ 0.39 | -
I-AE | 12.10 $\pm$ 0.35 | 46.60
PCA | 7.04 $\pm$ 0.25 | 724.51
FSNF | 2.55 $\pm$ 0.08 | 1011.60
W-GAN | 12.15 $\pm$ 0.51 | 95.10

The results in Table 4 highlight that the image data set generated using the
I-AE embedding scores as high as the test set with real images with respect to
the IS and also shows the lowest FID of all models. The W-GAN scores an IS as
good as the I-AE but shows a higher FID than the I-AE. The results suggest
that the W-GAN images are diverse and can be classified well as indicated by
the IS, but do not resemble the actual data as well as the I-AE normalizing
flow as highlighted by the FID. The IS also correctly indicates a lower
quality of the images generated using the PCA embedding as well as the
exceptionally low quality of the FSNF-generated images. Thus, the quantitative
results in Table 4 confirm our conclusions from the visual assessment in
Figure 7. Note that despite the isometry regularization, the FID score for the
I-AE images is competitive with that of, e.g., Ross and Cresswell [21], who
report an FID score of 38.5.

In conclusion, the isometric dimensionality reduction vastly increases the
quality of the generated images and outperforms the W-GAN reference despite
the MNIST manifold not being strictly flat. In fact, the I-AE learns highly
accurate encodings of the MNIST data set despite a latent dimensionality as
low as 16. Our analysis of Figure 7 and Table 4, thus, shows that the
composition of I-AE with RealNVP can build highly descriptive and accurate
generative models for data that resides on nonlinear manifolds. As expected,
the affine PCA shows a distinctively lower quality, which, however, is still
far better than the full-space normalizing flow images, which are completely
unusable.

##  5 Conclusion

This work proposes a two-parted approach that separates injective normalizing
flows into an isometric manifold-learning part and a lower-dimensional density
estimation part. The probability density function (PDF) described by the
change of variables formula is invariant to isometries, which leads to two
major benefits in training: First, no expensive Jacobian determinant
computations result from the embedding. Second, embedding functions and
normalizing flow models can be trained separately, which avoids the difficult
balancing of the unbounded likelihood maximization and the bounded
reconstruction loss. Thus, both model parts can be trained to the highest
attainable accuracy.

We apply the isometric embeddings to the artificial S-Curve data set, a
projection of a Gaussian to a spherical surface, and the MNIST data set with
images of handwritten digits. For the S-Curve and the MNIST data sets, the
isometric embeddings excel at modeling the probability distributions and
sample highly realistic data compared to the standard full-space normalizing
flow. Meanwhile, the spherical surface data set highlights the limitations of
isometric embeddings, which are unable to describe curved manifolds. The high-
quality MNIST samples highlight that the composition of I-AE and RealNVP
performs well, despite the trade-off between a strict isometry and a minimal
reconstruction loss. In practice, the complexity of the embedding should be
increased progressively, i.e., starting with PCA and progressively
transitioning to higher complexity if the embedding performs poorly.

In conclusion, the proposed combination of isometric embeddings and
normalizing flows presents an effective approach to handling manifold data and
takes full advantage of the direct log-likelihood maximization and explicit
sampling of normalizing flows. For established injective normalizing flow
models, the typical combined learning approach is difficult to tune. Thus, the
separation of embedding and density estimation is of great benefit as both
model parts can be trained to high accuracy without trade-offs between the two
loss functions. In short, the proposed separation makes training the injective
normalizing flow easy.

## Acknowledgements

This work was performed as part of the Helmholtz School for Data Science in
Life, Earth and Energy (HDS-LEE) and received funding from the Helmholtz
Association of German Research Centres. This publication is based upon work
supported by the King Abdullah University of Science and Technology (KAUST)
Office of Sponsored Research (OSR) under Award No. OSR-2019-CRG8-4033, and the
Alexander von Humboldt Foundation.

## Nomenclature

Description | Symbol
---|---
Latent dimension | $d$
Data dimensionalty | $D$
Decoder | $\mathbf{f}(\mathbf{z})$
Encoder | $\mathbf{g}(\mathbf{x})$
Jacobian | $\mathbf{J}$
Data manifold | $\mathcal{M}\subset\mathbb{R}^{D}$
Probability density function | $p_{X}(\mathbf{x})$
Surface of $d$-dimensional unit ball | $\mathbb{S}^{d-1}$
Diffeomorphism | $T(\cdot)$
Samples of multivariate random variables |  $\mathbf{x}$,$\mathbf{z}$
Random variable | $X$
Data set | $\mathcal{X}$
Sample mean vector of $\mathcal{X}$ | $\bm{\mu}_{\mathcal{X}}$
PDF of standard Gaussian | $\phi(\mathbf{z})$
Parameter vector | $\bm{\theta}$

## Bibliography

  * [1] Esteban G Tabak and Eric Vanden-Eijnden.  Density estimation by dual ascent of the log-likelihood.  Communications in Mathematical Sciences, 8(1):217–233, 2010.
  * [2] Esteban G Tabak and Cristina V Turner.  A family of nonparametric density estimation algorithms.  Communications on Pure and Applied Mathematics, 66(2):145–164, 2013\.
  * [3] Richard J Rossi.  Mathematical statistics: An introduction to likelihood based inference.  John Wiley & Sons, 2018.
  * [4] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using Real NVP.  In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
  * [5] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  Journal of Machine Learning Research, 22(57):1–64, 2021.
  * [6] Johann Brehmer and Kyle Cranmer.  Flows for simultaneous manifold learning and density estimation.  Advances in Neural Information Processing Systems, 33:442–453, 2020\.
  * [7] Jens Behrmann, Paul Vicol, Kuan-Chieh Wang, Roger Grosse, and Joern-Henrik Jacobsen.  Understanding and mitigating exploding inverses in invertible neural networks.  In Arindam Banerjee and Kenji Fukumizu, editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 1792–1800. PMLR, 13–15 Apr 2021.
  * [8] Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan.  Testing the manifold hypothesis.  Journal of the American Mathematical Society, 29(4):983–1049, 2016\.
  * [9] Madeleine Udell and Alex Townsend.  Why are big data matrices approximately low rank?  SIAM Journal on Mathematics of Data Science, 1(1):144–160, 2019\.
  * [10] Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [11] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud.  FFJORD: free-form continuous dynamics for scalable reversible generative models.  CoRR, abs/1810.01367, 2018.
  * [12] Mevlana C. Gemici, Danilo Rezende, and Shakir Mohamed.  Normalizing flows on Riemannian manifolds.  arXiv preprint arXiv:1611.02304, 2016.
  * [13] Danilo Jimenez Rezende, George Papamakarios, Sébastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer.  Normalizing flows on tori and spheres.  In International Conference on Machine Learning, pages 8083–8092. PMLR, 2020.
  * [14] Edmond Cunningham, Renos Zabounidis, Abhinav Agrawal, Ina Fiterau, and Daniel Sheldon.  Normalizing flows across dimensions.  CoRR, abs/2006.13070, 2020.
  * [15] Konik Kothari, AmirEhsan Khorashadizadeh, Maarten V. de Hoop, and Ivan Dokmanic.  Trumpets: Injective flows for inference and inverse problems.  In Cassio P. de Campos, Marloes H. Maathuis, and Erik Quaeghebeur, editors, Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence, UAI 2021, Virtual Event, 27-30 July 2021, volume 161 of Proceedings of Machine Learning Research, pages 1269–1278. AUAI Press, 2021.
  * [16] Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa.  Neural manifold ordinary differential equations.  In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.
  * [17] Christian Horvat and Jean-Pascal Pfister.  Density estimation on low-dimensional manifolds: an inflation-deflation approach.  arXiv preprint arXiv:2105.12152, 2021.
  * [18] Christian Horvat and Jean-Pascal Pfister.  Denoising normalizing flow.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, virtual, 2021\.
  * [19] Janis Postels, Martin Danelljan, Luc Van Gool, and Federico Tombari.  ManiFlow: Implicitly representing manifolds with normalizing flows.  arXiv preprint arXiv:2208.08932, 2022.
  * [20] Anthony L Caterini, Gabriel Loaiza-Ganem, Geoff Pleiss, and John P Cunningham.  Rectangular flows for manifold learning.  Advances in Neural Information Processing Systems, 34, 2021.
  * [21] Brendan Leigh Ross and Jesse C Cresswell.  Conformal embedding flows: Tractable density estimation on learned manifolds.  In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, pages 1–11. ICML, 2021.
  * [22] Edmond Cunningham, Adam D Cobb, and Susmit Jha.  Principal component flows.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 4492–4519. PMLR, 17–23 Jul 2022.
  * [23] Eike Cramer, Alexander Mitsos, Raúl Tempone, and Manuel Dahmen.  Principal component density estimation for scenario generation using normalizing flows.  Data-Centric Engineering, 3:e7, 2022.
  * [24] Amos Gropp, Matan Atzmon, and Yaron Lipman.  Isometric autoencoders.  arXiv preprint arXiv:2006.09289, 2020.
  * [25] Franz Kamber and Philippe Tondeur.  Flat Manifolds, volume 1.  Springer Berlin, Heidelberg, 1 edition, 1968.
  * [26] Yonghyeon Lee, Sangwoong Yoon, MinJun Son, and Frank C. Park.  Regularized autoencoders for isometric representation learning.  In International Conference on Learning Representations, virtual, 2022.
  * [27] John Nash.  The imbedding problem for Riemannian manifolds.  Annals of Mathematics, pages 20–63, 1956.
  * [28] Karl Pearson.  On lines and planes of closest fit to systems of points in space.  The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 2(11):559–572, 1901.
  * [29] Erez Peterfreund, Ofir Lindenbaum, Felix Dietrich, Tom Bertalan, Matan Gavish, Ioannis G Kevrekidis, and Ronald R Coifman.  Local conformal autoencoder for standardized data coordinates.  Proceedings of the National Academy of Sciences, 117(49):30918–30927, 2020.
  * [30] Christian Horvat and Jean-Pascal Pfister.  Intrinsic dimensionality estimation using normalizing flows.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  * [31] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay.  Scikit-learn: Machine learning in Python.  Journal of Machine Learning Research, 12:2825–2830, 2011.
  * [32] Yann LeCun and Corinna Cortes.  MNIST handwritten digit database, 2010.  http://yann.lecun.com/exdb/mnist/, accessed on 10-10-2021.
  * [33] Martín Abadi and Ashish Agarwal.  TensorFlow: Large-scale machine learning on heterogeneous systems.  https://www.tensorflow.org/, 2015.  Accessed on 08-08-2021.
  * [34] Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein generative adversarial networks.  In International Conference on Machine Learning, pages 214–223, Sydney, Australia, 2017. PMLR.
  * [35] Sam Bond-Taylor, Adam Leach, Yang Long, and Chris G. Willcocks.  Deep generative modelling: A comparative review of vaes, gans, normalizing flows, energy-based and autoregressive models.  IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1–1, 2021.
  * [36] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.  Improved techniques for training GANs.  In Proceedings of the 30th International Conference on Neural Information Processing Systems, pages 2234–2242. Curran Associates Inc., Red Hook, NY, USA, 2016.
  * [37] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  GANs trained by a two time-scale update rule converge to a local nash equilibrium.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.