Year: 2020

Abstract:
Flow-based generative models typically define a latent space with dimensionality identical to the observational space. In many problems, however, the data does not populate the full ambient data-space that they natively reside in, but rather inhabit a lower-dimensional manifold. In such scenarios, flow-based models are unable to represent data structures exactly as their density will always have support off the data manifold, potentially resulting in degradation of model performance. In addition, the requirement for equal latent and data space dimensionality can unnecessarily increase model complexity for contemporary flow models. Towards addressing these problems, we propose to learn a manifold prior that affords benefits to both the tasks of sample generation and representation quality. An auxiliary product of our approach is that we are able to identify the intrinsic dimension of the data distribution.

heading: 1 INTRODUCTION
text: Normalizing flows (Rezende and Mohamed, 2015; Kobyzev et al., 2020) have shown considerable potential for the tasks of modelling and inferring expressive distributions through the learning of well-specified probabilistic models. Contemporary flow-based approaches define a latent space with dimensionality identical to the data space, typically by parameterizing a complex model pX(x|θ) using an invertible neural network fθ. Samples drawn from an initial, simple distribution pZ(z) (e.g. Gaussian) can be mapped to a complex distribution as x = fθ(z). The process results in a tractable density that inhabits the full data space. However, contemporary flow models may make for an inappropriate choice to represent data that resides in a lower-dimensional manifold and thus does not populate the full ambient space. In such cases, the estimated model will necessarily have mass lying off the data manifold, which may result in under-fitting and poor generation qualities. Furthermore, principal objectives such as Maximum Likelihood Estimation (MLE) and Kullback-Leibler (KL) divergence minimization are ill-defined, bringing additional challenges for model training.
In this work, we propose a principled strategy to model a data distribution that lies on a continuous manifold and we additionally identify the intrinsic dimension of the data manifold. Specifically, by using the connection between MLE and KL divergence minimization in Z space, we can address the important problem of ill-defined KL divergence under typical flow based assumptions.
Flow models are based on the idea of “change of variable”. Assume a random variable Z with distribution PZ and probability density pZ(z). We can transform Z to get a random variable X: X = f(Z), where f : RD → RD is an invertible function with inverse f−1 = g. Suppose X has distribution PX and density function pX(x), then log pX(x) will have the following form
log pX(x) = log pZ(g(x)) + log ∣∣∣∣det(∂g∂x )∣∣∣∣ , (1)
where log ∣∣∣det( ∂g∂x)∣∣∣ is the log determinant of the Jacobian matrix. We call f (or g) a volumepreserving function if the log determinant is equal to 0.
Training of flow models typically makes use of MLE. We denote Xd as the random variable of the data with distribution Pd and density pd(x). In addition to the well-known connection between MLE and minimization of the KL divergence KL(pd(x)||pX(x)) in X space (see Appendix A for detail), MLE is also (approximately) equivalent to minimizing the KL divergence in Z space, this is due to the KL divergence is invariant under invertible transformation (Yeung, 2008; Papamakarios et al.,
2019). Specifically, we define ZQ : ZQ = g(Xd) with distribution QZ and density function q(z), the KL divergence in Z space can be written as
KL(q(z)||p(z)) = ∫ q(z) log q(z)dz − ∫ q(z) log p(z)dz (2)
= − ∫ pd(x) ( log pZ (g(x)) + log ∣∣∣∣det(∂g∂x )∣∣∣∣) dx+ const., (3)
The full derivation can be found in Appendix A. Since we can only access samples x1, x2, . . . , xN from pd(x), we approximate the integral by Monte Carlo sampling
KL(q(z)||p(z)) ≈ − 1 N N∑ n=1 log pX(xn) + const.. (4)
We thus highlight the connection between MLE and KL divergence minimization, in Z space, for flow based models. The prior distribution p(z) is usually chosen to be a D-dimensional Gaussian distribution. However, if the data distribution Pd is singular, for example a measure on a low dimensional manifold, the induced latent distribution QZ is also singular. In this case, the KL divergence in equation 2 is typically not well-defined under the considered flow based model assumptions. This issue brings both theoretical and practical challenges that we will discuss in the following section.

heading: 2 FLOW MODELS WITH MANIFOLD DATA
text: We assume a data sample x ∼ Pd to be a D dimensional vector x ∈ RD and define the ambient dimensionality of Pd, denoted by Amdim(Pd), to be D. However for many datasets of interest, e.g. natural images, the data distribution Pd is commonly believed to be supported on a lower dimensional manifold (Beymer and Poggio, 1996). We assume the dimensionality of the manifold to beK where K < D, and define the intrinsic dimension of Pd, denoted by Indim(Pd), to be the dimension of this manifold. Figure 1a provides an example of this setting where Pd is a 1D distribution in 2D space. Specifically, each data sample x ∼ Pd is a 2D vector x = {x1, x2} where x1 ∼ N (0, 1) and x2 = sin(2x1). Therefore, this example results in Amdim(Pd) = 2 and Indim(Pd) = 1.
In flow-based models, function f is constructed such that it is both bijective and differentiable. When the prior PZ is a distribution whose support is RD (e.g. Multivariate Gaussian distribution), the marginal distribution PX will also have support RD and Amdim(PX) = Indim(PX) = D. When the support of the data distribution lies on aK-dimensional manifold andK < D, Pd and PX are constrained to have different support. That is, the intrinsic dimensions of PX and Pd are always different; Indim(PX) 6= Indim(Pd). In this case it is impossible to learn a model distribution PX identical to the data distribution Pd. Nevertheless, flow-based models have shown strong empirical success in real-world problem domains such as the ability to generate high quality and realistic images (Kingma and Dhariwal, 2018). Towards investigating the cause and explaining this disparity between theory and practice, we employ a toy example to provide intuition for the effects and consequences resulting from model and data distributions that possess differing intrinsic dimension.
Consider the toy dataset introduced previously; a 1D distribution lying in a 2D space (Figure 1a). The prior density p(z) is a standard 2D Gaussian p(z) = N (0, IZ) and the function f is a nonvolume preserving flow with two coupling layers (see Appendix C.1). In Figure 1b we plot samples from the flow model; the sample x is generated by first sampling a 2D datapoint z ∼ N (0, IZ) and then letting x = f(z). Figure 1c shows samples from the prior distributions PZ and QZ . QZ is defined as the transformation of Pd using the bijective function g, such that QZ is constrained to support a 1D manifold in 2D space, and Indim(QZ) = Indim(Pd) = 1. Training of QZ to match PZ (which has intrinsic dimension 2), can be seen in Figure 1c to result in curling up of the manifold in the latent space, contorting it towards satisfying a distribution that has intrinsic dimension 2. This ill-behaved phenomenon causes several potential problems for contemporary flow models:
1. Poor sample quality. Figure 1b shows examples where incorrect assumptions in turn result in the model generating bad samples.
2. Low quality data representations. The discussed characteristic that results in “curling up” of the latent space may cause degradations of the representation quality.
3. Inefficient use of network capacity. Neural network capacity is spent on contorting the distribution QZ to satisfy imposed dimensionality constraints.
A natural solution to the problem of intrinsic dimension mismatch is to select a prior distribution PZ with the same dimensionality as the intrinsic dimension of the data distribution such that: Indim(PZ) = Indim(Pd). However, since we do not know Indim(Pd) explicitly, one option involves to instead learn it from the data distribution. In the following section, we will introduce a parameterization approach that enables us to learn Indim(Pd).

heading: 3 LEARNING A MANIFOLD PRIOR
text: Consider a data vector x ∈ RD, then a flow based model prior PZ is usually given by a Ddimensional Gaussian distribution or alternative simple distribution that is also absolutely continuous (a.c.) in RD. Therefore, the intrinsic dimension Indim(PZ) = D. To allow a prior to have intrinsic dimension strictly less than D, we let PZ have the ‘generalized density’1 p(z) with the form
p(z) = N (0, AAT ), (5)
where z ∈ RD and A is a D×D lower triangular matrix with D(D+1)2 parameters, such that AA T is constrained to be a positive semi-definite matrix. When AAT has full rank D, then PZ is a (nondegenerate) multivariate Gaussian on RD. When Rank(AAT ) = K and K < D, then PZ will degenerate to a Gaussian supported on a K-dimensional manifold, such that the intrinsic dimension Indim(PZ) = K. Figure 2 illustrates a sketch of this scenario. In practice, we initialize A to be an identity matrix, so AAT will also be an identity matrix and PZ is initialized as a standard Gaussian.
When Rank(AAT ) < D, the degenerate covariance AAT is no longer invertible and we are unable to evaluate the density value of p(z) for a given random vector z. Furthermore, when the data distribution Pd is supported on a K-dimensional manifold, QZ will also be supported on a Kdimensional manifold and no longer has a valid density function. Using equation 2 to train the flow model then becomes impossible as the KL divergence between PZ and QZ is not well defined2. Recent work by Zhang et al. (2020) proposed a new family of divergence to address this problem. In the following section we briefly review the key concepts pertaining to this family of divergences.
1We use the generalized density to include the case that AAT is not full rank. 2The KL divergence KL(Q||P) is well defined when Q and P have valid densities and their densities have
the same support (Ali and Silvey, 1966).

heading: 4 FIX THE ILL-DEFINED KL DIVERGENCE
text: random Let ZQ and ZP be two random variables with distribution QZ and PZ . The KL divergence between QZ and PZ is not defined if QZ or PZ does not have valid density function. Let K be an a.c. random variable that is independent of ZQ and ZP and has density pK , We define ZP̃ = ZP +K;ZQ̃ = ZQ +K with distributions P̃Z and Q̃Z respectively. Then P̃Z and Q̃Z are a.c. with density functions
q(z̃) = ∫ z pK(z̃− z)dQZ p(z̃) = ∫ z pK(z̃− z)dPZ . (6)
We can thus define the spread KL divergence between QZ and PZ as the KL divergence between Q̃Z and P̃Z as:
K̃L(QZ ||PZ) ≡ KL(Q̃Z ||P̃Z) ≡ KL (q(z̃)||p(z̃)) . (7)
In this work we letK be a Gaussian with diagonal covariance σ2ZI to satisfy the sufficient conditions such that K̃L is a valid divergence (see Zhang et al. (2020) for details) and has the properties:
K̃L(QZ ||PZ) ≥ 0, K̃L(QZ ||PZ) = 0⇔ QZ = PZ . (8)
Since QZ and PZ are transformed from Pd and PX using an invertible function g, we have
QZ = PZ ⇔ Pd = PX . (9)
Therefore, the spread KL divergence can be used to train flow based models with a manifold prior in order to fit a dataset that lies on a lower-dimensional manifold.

heading: 5 IDENTIFIABILITY OF INTRINSIC DIMENSION
text: A byproduct of our model is that the intrinsic dimension of the data manifold can be identified. Section 4 shows that when K̃L(QZ ||PZ) = 0 ⇔ PX = Pd, the supports of PX and Pd will also have the same intrinsic dimension: Indim(PX) = Indim(Pd). The flow function g, and its inverse g = f−1, are bijective and continuous so f is a diffeomorphism (Kobyzev et al., 2020). Due to the invariance of dimension property of diffeomorphisms (Lee, 2013, Theorem 2.17), the manifold that supports PZ will have the same dimension as the manifold that supports PX . Therefore, we have
Indim(PZ) = Indim(PX) = Indim(Pd). (10)
Since the intrinsic dimension of PZ is equal to the rank of the matrix AAT , we can calculate Rank(AAT ) by counting the number of non-zero eigenvalues of the matrix AAT . This allows for identification of the intrinsic dimension of the data distribution as
Indim(Pd) = Rank(AAT ). (11)
We have shown that we can identify the intrinsic dimension of Pd using the spread KL divergence. In the next section, we will discuss how to estimate the spread KL divergence in practice.

heading: 6 ESTIMATION OF THE SPREAD KL DIVERGENCE
text: Our goal is to minimize the spread KL divergence between QZ and PZ . Using our definition of spread divergence (equation 7), this is equivalent to minimizing
KL (q(z̃)||p(z̃)) = ∫ q(z̃) log q(z̃)dz̃︸ ︷︷ ︸
Term 1
− ∫ q(z̃) log p(z̃)dz̃︸ ︷︷ ︸
Term 2
. (12)
where q(z̃) and p(z̃) are defined in equation 6. We separate the objective into two terms and now discuss the estimation for each of them.
Term 1: We use H(·) to denote the differential entropy. Term 1 is the negative entropy −H(ZQ̃). For a volume preserving g and Xd is a.c., the entropy H(ZQ) = H(Xd) and is independent of the
model parameters and can be ignored during training. However, the entropy H(ZQ̃) = H(ZQ +K) will still depend on g, see Appendix B.1 for an example. We claim that when the variance of K is small, the dependency between the H(ZQ̃) and volume preserving function g is weak, thus we can approximate equation 12 by leaving out term 1 and will not affect the training.
To build intuitions, we first assume Xd is a.c., so ZQ = g(Xd) is also a.c.. Using standard entropic properties (Kontoyiannis and Madiman, 2014), we can pose the following relationship
H(ZQ) ≤ H(ZQ +K) = H(ZQ) + I(ZQ +K,K), (13)
where I(·, ·) denotes the mutual information. Since ZQ is independent of function g and if σ2Z → 0, then I(ZQ +K,K)→ 0 (see Appendix B.2 for a proof), the contribution of the I(ZQ +K,K) term, with respect to training g, becomes negligible in the case of small σ2Z .
Unfortunately, equation 13 is no longer valid when Pd lies on a manifold since ZQ will be a singular random variable and the differential entropy H(ZQ) is not defined. In Appendix B.3, we show that leaving out the entropy H(ZQ̃) corresponds to minimizing an upper bound of the spread KL divergence. To further find out the contribution of the negative entropy term, we compare between leaving out −H(ZQ̃) and approximating −H(ZQ̃) during training. In Appendix B.4, we discuss the approximation technique and give the empirical evidence which shows that the ignoring −H(ZQ̃) will not affect the training of our model. Therefore, we make use of volume preserving g and small variance σ2Z = 1× 10−4 in our experiments. In contrast to other volume preserving flows, that utilize a fixed prior PZ , our method affords additional flexibility by way of allowing for changes to the ‘volume’ of the prior towards matching the distribution of the target data. In this way, our decision to employ volume-preserving flow functions does not limit the expressive power or abilities of the model, in principle. Popular non-volume preserving flow structures, e.g. affine coupling flow, may also easily be normalized to become volume preserving, thus further extending the applicability of our approach (see Appendix C.1 for an example).
Term 2: The noisy prior p(z̃) is defined to be a degenerate Gaussian N (0, AAT ), convolved with a Gaussian noise N (0, σ2ZI), and has a closed form density
p(z̃) = N (0, AAT + σ2ZI). (14)
Therefore, the log density log p(z̃) is well defined, we can approximate term 2 by Monte Carlo∫ q(z̃) log p(z̃)dz̃ ≈ 1
N N∑ n=1 log p(z̃n), (15)
where q(z̃) = ∫ p(z̃|z)dQZ . To sample from q(z̃), we first get a data sample x ∼ Pd, use function g to get z = g(x) (so z is a sample of QZ) and finally sample z̃ ∼ p(z̃|z).

heading: 7 EXPERIMENTS
text: To demonstrate the effectiveness of our approach, Sections 7.1–7.4 report experiments on four datasets; toy 2D and 3D data, the fading square dataset and a synthesized adaptation of MNIST. We use the Adam optimizer (Kingma and Ba, 2014) in all our experiments. Our flow networks are consisted of incompressible affine coupling layer introduced by Sorrenson et al. (2020); Dinh et al. (2016). We compare a volume preserving flow with a learnable prior (our method) to a non-volume preserving flow with fixed prior, so both models have the ability to adapt their ‘volume’ to fit the target distribution and retains fair comparisons. See Appendix C.1 for a detailed discussion of the incompressible affine coupling layer and the network structures we use for all the experiments.

heading: 7.1 2D TOY DATA
text: We firstly verify our method using the toy dataset described in Section 2 and Figure 1a. The flow function has two coupling layers. We train our model using learning rate 3×10−4 and batch size 100 for 10k iterations. Figure 3 shows the samples from the model, the learned prior and the eigenvalues of AAT . We observe the sample quality is better than that in Figure 1b and the prior has learned
a degenerate Gaussian with Indim(PZ) = 1, which matches Indim(Pd). In Appendix C.2, we show that our model can not only learn the manifold support of target distribution but also capture the ‘density’ allocation on the manifold.

heading: 7.2 S-CURVE DATASET
text: We fit our model to the S-curve dataset shown in Figure 4a. The data distribution lies on a 2D manifold in a 3D space, therefore Indim(Pd) = 2. Specific network structure and training details can be found in Appendix C.1. After training, our model learns a nonlinear function g to transform Pd to QZ , where the latter lies on a 2D linear subspace in 3D space (see Figure 4b). Following this, a linear dimensionality reduction can be conducted to generate 2D data representations, we now briefly outline a general procedure for this.
For QZ with Amdim(QZ) = D and Indim(QZ) = K, we first find the eigenvectors e1, . . . , eD of AAT , sorted by their eigenvalues. When Rank(AAT ) = K ≤ D, there exist K eigenvectors with positive eigenvalues. We select the first K eigenvectors and form the matrix E = [e1, . . . , eK ] with dimension D×K. We then transform each data sample x ∈ RD into Z space: z = g(x), such that z ∈ RD. Afterwards, a linear projection is carried out zproj = zE to obtain the lower dimensional representation zproj ∈ RK . This procedure can be seen as a nonlinear dimensionality reduction, where the nonlinear component is solved using the learned function g.
We plot the resulting representations in Figure 4b. The colormap indicates correspondence between the data in 3D space and the representation in 2D space. We observe that our method can successfully (1) identify that the intrinsic dimension of the data is two and (2) project the data into a 2D space that faithfully preserves the structure of the original data distribution. We also compare the sample generation quality with a flow that has a fixed Gaussian prior, see Appendix C.3 for details.

heading: 7.3 FADING SQUARE DATASET
text: The fading square dataset (Rubenstein et al., 2018) was proposed in order to diagnose model behavior when data distribution and model possess differing intrinsic dimension and therefore constitutes a further relevant test bed for our current work. The dataset consists of 32×32 pixel images with 6×6 grey squares on a black background. The grey scale values are sampled from a uniform distribution with range [0, 1], so Indim(Pd) = 1. Figure 5a shows the data samples. We fit our model to the dataset, the network structure and the training details can be found in Appendix C.
Figure 5b shows samples from our trained model. Figure 5d shows the first 20 eigenvalues of the AAT (ranked from high to low), we can see only one eigenvalue is larger than zero and the others have converged to zero. This illustrates that we have successfully identified the intrinsic dimension of Pd. We further carry out the dimensionality reduction process that was introduced in Section 7.2; the latent representation z is projected onto a 1D line and we visualize the correspondence between the projected representation and the data in Figure 5e. Pixel grey-scale values can be observed to decay as the 1D representation is traversed from left to right, indicating that our representations are consistent with the properties of the original data distribution. In contrast, we find that the traditional flow model, with a fixed 1024D Gaussian p(z), fails to learn such a data distribution, see Figure 5c.

heading: 7.4 MNIST DATA
text: We further investigate training of our model using images of digits. However, for datasets like MNIST (LeCun, 1998), the true intrinsic dimension is unknown. In order to verify the correctness of our model’s ability to identify the intrinsic data dimension, we construct synthetic datasets by first fitting an implicit model pθ(x) = ∫ δ(x− g(z)p(z))dz to the MNIST dataset, and then use samples from the trained model x ∼ pθ(x) as training data. The intrinsic dimension of the training dataset is the same as the dimension of the latent variable z in the implicit model Indim(Pd) = dim(z). We construct two datasets with dim(z) = 5 and dim(z) = 10 such that Indim(Pd) = 5 and Indim(Pd) = 10, respectively. Further details on the implicit model, flow network structure, training method and samples from the learned models are found in Appendix C.
In contrast to the fading square dataset, we find that in order to successfully train the model (i.e. such that valid image samples are generated), it is necessary to add small Gaussian noise to the training data. This trick is commonly used in the training of flow based models for image data (Sorrenson et al., 2020). We note that adding Gaussian noise breaks the assumption that the data lies on a manifold and, alternatively, the intrinsic dimension of the training data distribution will be equal to its ambient dimension. Towards alleviating this undesired effect, we firstly add Gaussian noise with standard deviation σx = 0.05 and anneal σx after 2000k iterations with a factor of 0.9 every 10k
iterations. However, to help retain successful model training, we disallow the annealing procedure to reach a state where zero noise is added due to the outlined model behavior observed when considering this image dataset. Experimentally, we cap a lower-bound Gaussian noise level of 0.01 and leave further investigation of this phenomenon to future work.
In Figure 6a and 6b, we plot the first 20 eigenvalues of the AAT (ranked from high to low) after training on two synthetic MNIST datasets with intrinsic dimension 5 and 10. It can be observed that 5 and 10 eigenvalues are significantly larger than the remaining values, respectively. It can thus be concluded that the intrinsic dimension of the two datasets are 5 and 10. Remaining non-zero eigenvalues can be attributed to the small Gaussian noise added to the training data.
For the original MNIST dataset, it was shown that digits have different intrinsic dimension (Costa and Hero, 2006). This suggests the distribution of MNIST may lie on several, disconnected manifolds with differing intrinsic dimension. Although our model assumes that Pd lies on one continuous manifold, it is interesting to investigate the case when this model assumption is not fulfilled. We thus fit our model to the original MNIST data and plot the eigenvalues in Figure 6c. In contrast to Figures 6a and 6b, the gap between eigenvalues can be seen to be less pronounced, with no obvious step change. However the values suggest that the intrinsic dimension of MNIST is between 11 and 14. This result is consistent with previous estimations stating that the intrinsic dimension of MNIST is between 12 and 14 (Facco et al., 2017; Hein and Audibert, 2005; Costa and Hero, 2006). Recent work Cornish et al. (2019) discusses fitting flow models to a Pd that lies on disconnected components, by introducing a mixing prior. Such techniques may be easily combined with our method towards constructing more powerful flow models; a promising direction for future work.

heading: 8 RELATED WORK
text: Classic latent variable generative models assume that data distributions lie around a low-dimensional manifold, for example the Variational Auto-Encoder (Kingma and Welling, 2013) or, recently introduced, Noisy Injective Flows (Cunningham et al., 2020). Such methods typically assume that observational noise is not degenerated (e.g. a fixed Gaussian distribution), therefore the model distribution is absolutely continuous and maximum likelihood learning is thus well defined. However, common distributions such as natural images usually don’t have Gaussian observational noise (Zhao et al., 2017). Therefore, in this work, we focus on modeling distributions that lie on a low-dimensional manifold.
The study of manifold learning for nonlinear dimensionality reduction Cayton (2005) and intrinsic dimension estimation Camastra and Staiano (2016) is a rich field with an extensive set of tools. However, most methods commonly do not model data density on the manifold and are thus not used for the same purpose as the models introduced here. There are however a number of recent works that introduced normalizing flows on manifolds that we now highlight and relate to our approach.
Several works define flows on manifolds with prescribed charts. Gemici et al. (2016) generalized flows from Euclidean spaces to Riemannian manifolds by proposing to map points from the manifold M to RK , apply a normalizing flow in this space and then map back toM. The technique has been
further extended to Tori and Spheres (Rezende et al., 2020). These methods require knowledge of the intrinsic dimension K and a parameterization of the coordinate chart of the data manifold.
Without providing a chart mapping a priori, M-flow (Brehmer and Cranmer, 2020) recently proposed an algorithm that learns the chart mapping and distribution density simultaneously. However, their method still requires that the dimensionality of the manifold is known. They propose that the dimensionality can be learned either by a brute-force solution or through a trainable variance in the density function. The brute-force solution is clearly infeasible for data embedded in extremely high dimensional space, as is often encountered in deep learning tasks. Use of a trainable variance is natural and similar to our approach. However, as discussed at the beginning of this paper, without carefully handling the KL or MLE term in the objective, a vanishing variance parameter will result in wild behavior of the optimization process since these terms are not well defined.
While the GIN model considered in Sorrenson et al. (2020) could recover the low dimensional generating latent variables following their identifiability theorem, the assumptions therein require knowledge of an auxiliary variable, e.g. the label, which is not required in our model. Behind this superficial difference is the essential discrepancy between the concept of informative dimensions and intrinsic dimensions. The GIN model discovers the latent variables that are informative in a given context, defined by the auxiliary variable u instead of the true intrinsic dimensions. In their synthetic example, the ten dimensional data is a nonlinear transformation of ten dimensional latent variables where two out of ten are correlated with the labels of the data and the other eight are not. In this example, there are two informative dimensions, but there are ten intrinsic dimensions. Nevertheless, our method for intrinsic dimension discovery can be used together with informative dimension discovery methods to discover finer structures of data.

heading: 9 CONCLUSION
text: We presented a principled strategy to learn the data distribution that lies on a manifold and identify its intrinsic dimension. We fix the ill-defined KL divergence and show, across multiple datasets, the resulting benefits for flow based models under both sample generation and representation quality. There remain a number of open questions and interesting directions for future work. Namely; further exploration of the effects of the entropy term in the case of non-volume preserving networks and, additionally, investigation of the phenomenon concerning the apparent necessity of noise addition in cases pertaining to complex real-world distributions e.g. image data.

heading: A MAXIMUM LIKELIHOOD ESTIMATION AND KL DIVERGENCE
text: Given data x1, x2, . . . , xN sampled independently from the true data distribution Pd, with density function pd(x), we want to fit the model density p(x)3 to the data. A popular choice to achieve this involves minimization of the KL divergence where:
KL(pd(x)||p(x)) = ∫ pd(x) log pd(x)dx− ∫ pd(x) log p(x)dx (16)
= − ∫ pd(x) ( log pZ (g(x)) + log ∣∣∣∣det(∂g∂x )∣∣∣∣) dx+ const.. (17)
Since we can only access samples from pd(x), we approximate the integral by Monte Carlo sampling
KL(pd(x)||p(x)) ≈ − 1
N N∑ n=1 log p(xn) + const.. (18)
Therefore, minimizing the KL divergence between the data distribution and the model is (approximately) equivalent to Maximum Likelihood Estimation (MLE).
When p(x) is a flow-based model with invertible flow function f : Z → X, g = f−1, minimizing the KL divergence in X space is equivalent to minimizing the KL divergence in the Z space. We let Xd be the random variable of data distribution and define ZQ : ZQ = g(Xd) with density q(z), so q(z) can be represented as
q(z) = ∫ δ(z − g(x))pd(x)dx. (19)
Let p(z) be the density of the prior distribution PZ , the KL divergence in Z space can be written as KL(q(z)||p(z)) = ∫ q(z) log q(z)dz︸ ︷︷ ︸
Term 1
− ∫ q(z) log p(z)dz︸ ︷︷ ︸
Term 2
. (20)
Term 1: using the properties of transformation of random variable (Papoulis and Pillai, 2002, pp. 660), the negative entropy can be written as∫
q(z) log q(z)dz = ∫ pd(x) log pd(x)dx︸ ︷︷ ︸
const.
− ∫ pd(x) log ∣∣∣∣det(∂g∂x )∣∣∣∣ dx. (21)
Term2: the cross entropy can be written as∫ q(z) log p(z)dz = ∫ ∫ δ(z − g(x))pd(x) log p(z)dzdx (22)
= ∫ pd(x) log p(g(x))dx. (23)
Therefore, the KL divergence in Z space is equivalent to the KL divergence in X space
KL(q(z)||p(z)) = KL(pd(x)||p(x)). (24) We thus build the connection between MLE and minimizing the KL divergence in Z space.

heading: B ENTROPY

heading: B.1 AN EXAMPLE
text: Assume a 2D Gaussian random variable X with covariance [ 1 0 0 1 ] , g is a volume preserving flow
with parameter θ. Let Z1 = gθ1(Z) be a Gaussian with covariance [ 2 0 0 12 ] and Z2 = gθ2(X) be
3For simplicity, we use notation p(x) to represent the model pX(x) unless otherwise specified.
a Gaussian with covariance [ 3 0 0 13 ] . Therefore the entropy H(Z1) = H(Z2) = H(X) and doesn’t
not depend on θ. Let K be an Gaussian with zero mean and covariance [ 1 0 0 1 ] , so Z1 + K is a
Gaussian with covariance [ 3 0 0 32 ] and Z2 + K is a Gaussian with covariance [ 4 0 0 43 ] . Therefore H(Z1 +K) 6= H(Z2 +K) and H(gθ(X) +K) depends on θ. A similar example can be constructed when X is not absolutely continuous.

heading: B.2 Z IS AN ABSOLUTELY CONTINUOUS RANDOM VARIABLE
text: For two mutually independent absolutely continuous random variable Z and K, the mutual information between Z +K and K is
I(Z +K,K) = H(Z +K)−H(Z +K|K) (25) = H(Z +K)−H(Z|K) (26) = H(Z +K)−H(Z). (27)
The last equality holds because Z andK are independent. Since mutual information I(Z+K,K) ≥ 0, we have
H(Z) ≤ H(Z +K) = H(Z) + I(Z +K,K). (28)
AssumeK has a Gaussian distribution with 0 mean and variance σ2Z 4. When σ2Z = 0,K degenerates to a delta function, so Z +K = Z and I(Z +K,K) = H(Z +K)−H(Z) = 0, (29)
this is because the mutual information between an a.c. random variable and a singular random variable is still well defined, see (Yeung, 2008, Theorem 10.33). Assume K1, K2 are Gaussian random variables with 0 mean and variances σ21 and σ 2 2 respectively. Without loss of generality, we assume σ21 > σ 2 2 and σ 2 1 = σ 2 2 + σ 2 δ , and let Kδ be the random variable of a Gaussian that has 0 mean and variance σ2δ such that K1 = K2 +Kδ . By the data-processing inequality, we have I(Z +K2,K2) ≤ I(Z +K2 +Kδ,K2 +Kδ) = I(Z +K1,K1). (30)
Therefore, I(Z+K,K) is a monotonically decreasing function when σ2Z decreases and when σ 2 Z → 0, I(Z +K,K)→ 0.

heading: B.3 UPPER BOUND OF THE SPREAD KL DIVERGENCE
text: In this section, we show that leaving out the entropy term H(ZQ̃) in equation 12 is equivalent to minimizing an upper bound of the spread KL divergence.
For singular random variable ZQ = g(Xd) and absolutely continuous random variable K that are independent, we have
H(ZQ +K)−H(K) = H(ZQ +K)−H(ZQ +K|ZQ) (31) = I(ZQ +K,ZQ) ≥ 0. (32)
The second equation is from the definition of Mutual Information (MI); the MI between an a.c. random variable and a singular random variable is well defined and always positive, see (Yeung, 2008, Theorem 10.33) for a proof.
Therefore, we can construct an upper bound of the spread KL objective in equation 12 KL(q̃||p̃) = ∫ q(z̃) log q(z̃)dz̃︸ ︷︷ ︸ −H(ZQ+K) − ∫ q(z̃) log p(z̃)dz̃ (33)
≤ −H(K)︸ ︷︷ ︸ const.
− ∫ q(z̃) log p(z̃)dz̃. (34)
Therefore, ignoring the negative entropy term during training is equivalent to minimizing an upper bound of the spread KL objective.
4The extension to higher dimensions is straightforward.

heading: B.4 EMPIRICAL EVIDENCE FOR IGNORING THE NEGATIVE ENTROPY
text: In this section, we first introduce the approximation technique to compute the negative entropy term, and then discuss the contribution of this term for the training. The negative entropy of random variable ZQ̃ is
−H(ZQ̃) = ∫ q(z̃) log q(z̃)dz̃, (35)
where
q(z̃) = ∫ z pK(z̃− z)dQZ , (36)
and pK is the density of a Gaussian with diagonal covariance σ2ZI . We first approximate q̃(z̃) by a mixture of Gaussians
q(z̃) ≈ 1 N N∑ n=1 N (z̃; zn, σ2ZI) ≡ q̂N (z̃) (37)
where zn is the nth sample from distribution Qz by first sampling xn ∼ Pd and letting zn = g(xn). We denote the random variable of this Gaussian mixture as ẐNQ̃ and approximate
−H(ZQ̃) ≈ −H(Ẑ N Q̃ ). (38)
However, the entropy of a Gaussian mixture distribution does not have a closed form, so we further conduct a first order Taylor expansion approximation (Huber et al., 2008)
−H(ẐNQ̃ ) ≈ 1
N N∑ n=1 log q̂N (z̃ = zn), (39)
this approximation is accurate for small σ2Z . Finally we have our approximation
−H(ZQ̃) ≈ 1
N N∑ n=1 log q̂N (z̃ = zn). (40)
To evaluate the contribution of the negative entropy, we train our flow model on both low dimensional data (Toy datasets: 2D) and high dimensional data (Fading square dataset: 1024D) by optimization that uses two objectives: (1) ignoring the negative entropy term in equation 7 and (2) approximating the negative entropy term in equation 7 using the approximation discussed above. During training, we keep tracking of the value of the entropy H(ZQ̃) (using the approximation value) in both objectives. We let N equal the batch size when approximating the entropy. Additional training details remain consistent with those described in Appendix C.
To evaluate the contribution of the negative entropy, we train our flow model on both low dimensional data (Toy datasets: 2D) and high dimensional data (Fading square dataset: 1024D) by optimization that uses two objectives: (1) ignoring the negative entropy term in equation 7 and (2) approximating the negative entropy term in equation 7 using the approximation discussed above. During training, we keep tracking of the value of the entropy H(ZQ̃) (using the approximation value) in both objectives. We let N equal the batch size when approximating the entropy. Additional training details remain consistent with those described in Appendix C.
In Figure 7, we plot the (approximated) entropy value H(ZQ̃) during training for both experiments. We find the difference between having the (approximated) negative entropy term and ignoring the negative entropy to be negligible. We leave theoretical investigation on the effects of leaving out the entropy term during training to future work.

heading: C EXPERIMENTS

heading: C.1 NETWORK ARCHITECTURE
text: The flow network we use consists of incompressible affine coupling layers (Sorrenson et al., 2020; Dinh et al., 2016). Each coupling layer splits a D-dimensional input x into two parts x1:d and xd+1:D. The output of the coupling layer is
y1:d = x1:d (41) yd+1:D = xd+1:D exp(s(x1:d)) + t(x1:d), (42)
where s : Rd → RD−d and t : Rd → RD−d are scale and translation functions parameterized by neural networks, is the element-wise product. The log-determinant of the Jacobian of a coupling layer is the sum of the scaling function ∑ j s(x1:d)j . To make the coupling transform volume preserving, we normalize the output of the scale function, so the i-th dimension of the output is
s̃(x1:d)i = s(x1:d)i − 1 D − d ∑ j s(x1:d)j , (43)
and the log-determinant of the Jacobian is ∑ i s̃(x1:d)i = 0. We compare a volume preserving flow with a learnable prior (normalized s(·)) to a non-volume preserving flow with fixed prior (unnormalized s(·)). In this way both models have the ability to adapt their ‘volume’ to fit the target distribution, retaining comparison fairness.
In our affine coupling layer, the scale function s and the translation function t have two types of structure: fully connected net and convolution net. Each fully connected network is a 4-layer neural network with hidden-size 24 and Leaky ReLU with slope 0.2. Each convolution net is a 3-layer convolutional neural network with hidden channel size 16, kernel size 3×3 and padding size 1. The activation is Leaky ReLU with slope 0.2. The downsampling decreases the image width and height by a factor of 2 and increases the number of channels by 4 in a checkerboard-like fashion (Sorrenson et al., 2020; Jacobsen et al., 2018). When multiple convolutional nets are connected together, we randomly permute the channels of the output of each network except the final one. In Table 1, 2, 3, 4, we show the network structures for our four main paper experiments.

heading: C.2 TOY DATASET
text: We also construct a second dataset and train a flow model using the same training procedure discussed in Section 7.1. Figure 8a shows the samples from the data distribution Pd, each data point is a 2D vector x = [x1, x2] where x1 ∼ N (0, 1) and x2 = x1, so Indim(Pd) = 1. Figure 8f shows that the prior PZ has learned the true intrinsic dimension Indim(PZ) = Indim(Pd) = 1. We compare to samples drawn from a flow model that uses a fixed 2D Gaussian prior, with results shown in Figure 8. We can observe, for this simple dataset, flow with a fix Gaussian prior can generate reasonable samples, but the ‘curling up’ behavior, discussed in the main paper, remains highly evident in the Z space, see Figure 8c.
We also plot the ‘density allocation’ on the manifold for the two toy datasets. For example, for the data generation process x = [x1, x2] where x1 ∼ p = N (0, 1) and x2 = x1, we use the density value p(x = x1) to indicate the ‘density allocation’ on the 1D manifold. To plot the ‘density allocation’ of our learned model, we first sample xs uniformly from the support of the data distribution, the subscript ‘s’ here means that they only contain the information of the support. Specifically, since xs = [x s 1, x s 2], we sample x s 1 ∼ p = U(−3σ, 3σ) (σ is the standard deviation of N (0, 1), U is the uniform distribution) and let xs2 = x s 1 or x s 2 = sin(x s 1), depending on which dataset is used. We use the projection procedure that was described in Section 7.2 to obtain the projected samples zprojs , so zprojs ∈ RIndim(Pd). We also project the learned prior PZ to RIndim(Pd) by constructing P proj Z as a Gaussian with zero mean and a diagonal covariance contains the non-zeros eigenvalues of AAT . Therefore PprojZ is a.s. in RIndim(Pd) and we denote its density function as pproj(z). We then use the density value pproj(z = zprojs ) to indicate the ‘density allocation’ at the location of xs on the manifold support. In Figure 9, we compare our model with the ground truth and find that we can successfully capture the ‘density allocation’ on the manifold.

heading: C.3 S-CURVE DATASET
text: To fit our model to the data, we use the Adam optimizer with learning rate 5× 10−4 and batch size 500 and train the model for 200k iterations. We anneal the learning rate with a factor of 0.9 every 10k iterations.
We compare our method with a traditional normalizing flow with a fixed 3D Gaussian prior. Both models have the same network architecture and training procedure. Figure 10a and 10b show the samples form our model and the traditional flow with a fixed 3D Gaussian prior. We can observe more samples lying outwith the true data distribution in Figure 10b than in Figure 10a. We can conclude that our model has better sample quality considering this S-Curve dataset. We also compare the latent representation for both models, see Figure 10c and 10d. We can see the representation distribution QZ captures the structure of the data distribution well whereas the distribution QZ in Figure 10d is unable to do so.

heading: C.4 FADING SQUARE DATASET
text: To fit the data, we train our model for 20k iterations with batch-size 100 using the Adam optimizer. The learning rate is initialized to 5×10−4 and decays with a factor of 0.9 at every 1k iterations. We additionally use an L2 weight decay with factor 0.1.

heading: C.5 MNIST DATASET
text: C.5.1 IMPLICIT DATA GENERATION MODEL
To fit an implicit model to the MNIST dataset, we first train a Variational Auto-Encoder (VAE) (Kingma and Welling (2013)) with Gaussian prior p(z)=N (0, I). The encoder is q(z|x) = N (µθ(x),Σθ(x)) where Σ is a diagonal matrix. Both µθ and Σθ are parameterized by a 3-layer
feed-forward neural network with a ReLU activation and the size of the two hidden outputs are 400 and 200. We use a Gaussian decoder p(x|z) = N (gθ(z), σ2xI) with fixed variance σx = 0.3. The gθ is parameterized by a 3-layer feed-forward neural network with hidden layer sizes 200 and 400. The activation of the hidden output uses a ReLU and we utilize a Sigmoid function in the final layer output to constrain the output between 0 and 1. The training objective is to maximize the lower bound of the log likelihood
log p(x) ≥ ∫ q(z|x) log p(x|z)dz −KL(q(z|x)||p(z)),
see Kingma and Welling (2013) for further details. We use an Adam optimizer with learning rate 1 × 10−4 and batch size 100 to train the model for 100 epochs. After training, we sample from the model by first taking a sample z ∼ p(z) and letting x = gθ(z). This is equivalent to taking a sample from an implicit model pθ(x) = ∫ δ(x − gθ(z))d(z)dz, see Tolstikhin et al. (2017) for further discussion regarding this implicit model construction. In Figure 11, we plot samples from the trained implicit model with dim(z) = 5, dim(z) = 10 and the original MNIST data.

heading: C.5.2 FLOW MODEL TRAINING
text: We train our flow models to fit the synthetic MNIST dataset with intrinsic dimensions 5 and 10 and the original MNIST dataset. In all models, we use the Adam optimizer with learning rate 5 × 10−4 and batch size 100. We train the model for 3000k iterations and, following the initial 1000k iterations, the learning rate decays every 10k iterations by a factor 0.9. In Figure 12, we plot the samples from our models trained on three different training datasets. In Figure 13, we also plot the samples from traditional flow models that trained on these three datasets using the same experiment settings , we found the samples from our models are sharper than that from traditional flow models.

Review summaries:
This paper proposes a new method to train flow models on data from low dimensional manifolds embedded in high dimensional ambient spaces. The basic idea is based on minimizing the KL divergence in the latent space, which is equivalent to maximizing expected log-likelihood over the data distribution. Since the KL between two low-dimensional distributions are often undefined, authors propose to use spread divergence as a surrogate. Experiments demonstrate that the authors' technique can successfully model the distributions on low-dimensional manifolds without knowing the manifold beforehand. In particular, the proposed method can recover the intrinsic dimensionality of data manifolds.

This paper proposes a modification to the latent distribution of a flow model, replacing the commonly used full-rank Normal with a low-rank one which has the form N(0, AA^T). To train this degenerate model, the spread divergence from Zhang et al. (2020) is used and some approximation is made (e.g., ignoring the first entropy term). Experiments on toy data (with 1d intrinsic dimension) and mnist digits generated from a GAN (with 5-10 intrinsic dimension) demonstrated that the proposed model can identify the intrinsic dimension through the rank of AA^T.

This paper proposes a new method of training flow models, instead of minimizing KL divergence in the data space X, the paper proposes to minimize the KL divergence in the latent space Z. However, the problem of this is dimension mismatch, so the KL divergence is ill-defined. The proposed solution is to add noise, such that the KL divergence can be well defined again. In addition, this paper proposes learning a low dimensional prior p(Z) by a Gaussian distribution with low rank covariance matrix.