Lifting Architectural Constraints of Injective Flows
Authors: Peter Sorrenson
Felix Draxler
Armand Rousselot
Sander Hummerich
Lea Zimmermann
Ullrich Koethe
Venue: ICLR 2024 poster (2024)
Keywords: normalizing flows
injective flows
manifold learning
maximum likelihood
generative model
auto encoder
Primary area: generative models

Abstract: Normalizing Flows explicitly maximize a full-dimensional likelihood on the training data. However, real data is typically only supported on a lower-dimensional manifold leading the model to expend significant compute on modeling noise. Injective Flows fix this by jointly learning a manifold and the distribution on it. So far, they have been limited by restrictive architectures and/or high computational cost. We lift both constraints by a new efficient estimator for the maximum likelihood loss, compatible with free-form bottleneck architectures. We further show that naively learning both the data manifold and the distribution on it can lead to divergent solutions, and use this insight to motivate a stable maximum likelihood training objective. We perform extensive experiments on toy, tabular and image data, demonstrating the competitive performance of the resulting model.

Published as a conference paper at ICLR 2024
LIFTING ARCHITECTURAL CONSTRAINTS OF INJECTIVE FLOWS
Peter Sorrenson; Felix Draxler,' Armand Rousselot, Sander Hummerich, Lea Zimmermann & Ullrich Köthe Computer Vision and Learning Lab Heidelberg University firstname Lastname iir. uni wwweeeeieelleee de
ABSTRACT
Normalizing Flows explicitly maximize full-dimensional likelihood on the train- ing data. However, real data typically only supported on a lower-dimensional manifold leading the model to expend significant compute on modeling noise. In- jective Flows fix this by jointly learning manifold and the distribution on it. So far, they have been limited by restrictive architectures and/or high computational cost. We lift both constraints by new efficient estimator for the maximum like- lihood loss, compatible with free-form bottleneck architectures. We further show that naively learning both the data manifold and the distribution on it can lead to divergent solutions, and use this insight to motivate stable maximum likelihood training objective. We perform extensive experiments on toy, tabular and image data, demonstrating the competitive performance of the resulting model.
INTRODUCTION
Generative modeling is one the most important tasks in machine learning, having numerous ap- plications across vision (Rombach al. [2022), language modeling (Brown al.) 2020), science (Ariizzone et al.] 2018 Radey et al.] 2022) and beyond. One of the best-motivated approaches to generative modeling is maximum likelihood training, due its favorable statistical properties (Hastie et al. 2009). In the continuous setting, exact maximum likelihood training is most com- monly achieved by normalizing flows (Rezende & Mohamed 2015 Dinh et al. 2015 Kobyzev et al. 2021) which parameterize an exactly invertible function with tractable change of variables (log- (log-determinant term). This generally introduces a trade-off between model expressivity and com- putational cost, where the cheapest networks to train and sample from, such as coupling block architectures, require very specifically constructed functions which may limit expressivity (Draxler) et al. 2022). In addition, normalizing flows preserve the dimensionality of the inputs, requiring latent space of the same dimension as the data space.
Due to the manifold hypothesis (Bengio et al. 2013), which suggests that realistic data lies low-dimensional manifold embedded into a high-dimensions data space, it is more efficient to only model distributions a now-dinensional manifold and regard deviations from the manifold as un- informative noise. Prior works such as Caterini et al. (0021): Brehmer & Cranmer (2020) have restricted normalizing flows to low-dimensional manifolds via specially-construccee bottleneck chitectures (known as "invertible autoencoders" (Teng & Choromanska 2019) or "injective flows" (Kothari et al., (02T)) where encoder and decoder share parameters. These injective flows are opti- mized by some version of maximum likelihood training. This is not an ideal design decision, as the restrictive architectures used such models were originally designed for tractable change of vari- ables calculations in normalizing flows, but such calculations are not possible in the presence of bottleneck (Brehmer & Cranmer 2020).. As a result, we propose to drop the restrictive constructions (such as coupling blocks), instead use an unconstrained encoder and decoder, and introduce new technique to get around calculating the change of variables. This greatly simplifies the design of the model and makes it more expressive.
"Equal contribution.

---

Published as a conference paper at ICLR 2024
ENLL
Dec
8
Gauile
Guarin
Lecon
Figure 1: Free-form injective flow (FIF) training and inference. (Left) We combine a recon- struction loss Lreoon. with a novel maximum likelihood loss EULL to obtain an injective flow without architectural constraints. (Right) We generate novel samples by decoding standard normal latent samples with our best-performing models on CelebA and MNIST. The reconstructions shown are on CelebA validation data, the samples uncurated samples from our models.
We build on the unbiased maximum likelihood estimator used by rectangular flows (Caterini et al., 2021] to approximate the gradient of the change variables term. We simplify the estimator con- siderably by replacing iterative conjugate gradient with an efficient single-step estimator. This is fast: a batch can be processed about twice the time (or less) as an autoencoder trained on recon- struction loss only. In addition, we make novel observation about injective flows: naively training with maximum likelihood ill-defined due to the possibility diverging curvature the decod- ing function. To fix this problem, we propose modification to maximum likelihood estimator which counteracts the possibility of diverging curvature. We call our model the free-form injective flow (FIF).
To summarize, we make the following contributions:
We introduce an efficient maximum likelihood estimator for free-form injective flows and use it to train an unconstrained injective flow for the first time (section 4.1.
We identify pathological behavior in the naive application of maximum likelihood training in the presence of bottleneck, and offer solution avoid this behavior while maintaining computational efficiency (section 4.2).
We outperform previous injective flows and demonstrate competitive performance to gen- erative autoencoders toy, tabular and image data (section (5. We provide code to im- plement our model and reproduce our results at hhttps:///e htt.tt.ttuub...o cww/ciicciiiiio... www.vislearr//////ooccm.mmm FFF
2 RELATED WORK
Injective flows jointly learn manifold and maximize likelihood on that manifold. The latter requires estimating the Jacobian determinant the transformation to calculate the change variables. Ef- ficient computation of this determinant traditionally imposed two major restrictions flow architectures: Firstly, the latent space has to match dimension with the data space, ruling bottleneck architectures. Secondly, normalizing flows restricted certain functional forms, such as coupling and autoregressive blocks. Below we outline the existing approaches to overcome these problems and how our solution compares.
normalizing
Lover-dimensional latent spaces One set of methods attempts to use full-dimensions normal- izing flows, with some additional regularization or architectural constraints such that a subspace of the latent space corresponds the manifold. One strategy adds noise to the data to make it full-dimensional distribution then denoises to the manifold (Horvat & Pfister [021]] Loaiza-Gannoe
et al. 2022). Another restricts the non-manifold latent dimensions have small variance (Beitler
et al.
2021 Silvestri et al. 2023 Zhang et al. 2023).
Other methods sidestep the problem by making training into two-step procedure. First, toencoder is trained on reconstruction loss, then a normalizing flow is trained to learn the resulting latent distribution. In this line of work, Brehmer & Cranmer (2020) and Kothari et al. (202T) use
2

---

Published as a conference paper at ICLR 2024
an injective flow, while Böhhm & Seljak (2022) use unconstrained networks as autoencoder. Ghosh et al. (2020) additionally regularize the decoder.
Conformal embedding flows (Ross & Cresswell 2021 ensure the decomposition of the determinant into the contribution from each block by further restricting the architecture to exclusively conformal transformations. Cramer et al. (2022) uses isometric autoencoder such that the change of variables trivial. However, the resulting transformations are quite restrictive and cannot represent arbitrary manifolds.
The most similar work to ours is the rectangular flow (Caterini et al. 2021 which estimates the gradient of the log-determinant via iterative, unbiased estimator. The resulting method is quite slow to train, and uses injective flows, which are restrictive.
Unconstrained normalizing flow architectures Several works attempt to reduce the constraints imposed by typical normalizing flow architectures, allowing the use of free-form networks. How- ever, all of these methods only apply to full-dimensional architectures. FFJORD (Grathwohl et al., 2019) is a type of continuous normalizing flow (Chen al.] 2018b) which estimates the change of variables stochastically. Residual flows (Behrmann et al. 2019] Chen et al. 2019) make resid- ual networks invertible, but require expensive iterative estimators to train via maximum likelihood. Self-normalizing flows (Keller et al. 2021)) and relative gradient optimization (Gresele et al.) (020)) estimate maximum likelihood gradients for the matrices used in neural networks, but restrict the ar- chitecture to use exclusively square weight matrices without skip connections. In parallel work, we present the extension of the present paper to full-dimensional normalizing flows based on free-form neural networks, called free-form flows (Drxxler et al. 2024)..
Approximating maximum likelihood Many methods optimize some bound on the full- dimensional maximum likelihood, notably the variational autoencoder (Kingma & Welling, 2014) and its variants. [Cunningham et al. (2020) also optimizes variational lower bound to the likelihood. Other methods fit into the injective flow framework by jointly optimizing a reconstruction loss and some approximation
maximum likelihood on the manifold: Kumar et al. (2020); /Zhang et al. (020) approximate the log-determinant of the Jacobian by its Frobenius norm. The entropic AE (Ghose al. 2020) maximizes the entropy of the latent distribution by nearest-neighhor estimator while constraining its variance, resulting in a Gaussian latent space. In addition, there are other ways to regularize the latent space of an autoencoder which are not based on maximum likelihood, e.g. adversarial methods (Makhzani et al. 2015).
In contrast the above, our approach jointly learns the manifold and maximizes the likelihood on it with an unconstrained architecture, which easily accommodates lower-dimensional latent space.
BACKGROUND
Notation Let RD + Rd be an encoder which compresses data to a latent space and decoder g Rd + RR which decompresses the latent representation. A full-dimensional model has d = D while bottleneck model has d D. If o Rd Rd is the identity, then we call f and consistent. For example, the forward and inverse function of a normalizing flow are consistent as p-1=g.
Injective flows Injective flows (Brehmer Cranmer 2020), also called invertible autoencoders
Teng & Choromanskal] 2019). adapt invertible normalizing flow architectures to bottleneck setting. They parameterize and g as the composition of two invertible functions, w defined in TR. and defined in IR with slicing/padding operation in between:
f= h- slice w= and g pad h,
where sli slice(x) selects the first d elements of and pad(z) concatenates D -d zeros to the end of z. Since slice and pad are consistent, so too are f and g.
Injective flows typically minimize reconstruction loss (to learn a manifold which spans the data) alongside maximizing the likelihood of the data on that manifold, performing joint manifold and maximum likelihood training (alternatively, they trained via two-step training, section 2
3

---

Published as a conference paper at ICLR 2024
Change of variables across dimensions The maximum likelihood objective resulting from the change of variables theorem, used to train normalizing flow models, is only well-defined when map- ping between spaces of equal dimension. A result from differential geometry (Krantz & Parks 2008) allows us to generalize the change of variables theorem to non-equal-dimenion transforma- tions through the formula:
px - pz((())) (det
(f(x))))
(2)
where ff and g are consistent and primes denote derivatives: (((())) is the Jacobian matrix of evaluated at f(x). Note that, since px derived as the pushforward of the latent distribution pz by 9. the formula is valid only for z which lie on the decoder manifold (see appendix [ for more details). Unlike the full-dimensional case, there no architecture known which can represent arbitrary manifolds and at the same time allows efficient exact computation of eq. (see section
Rectangular flows Minimizing the negative logarithm of eq. 2 and adding Lagrange multiplier to restrict the distance of data points from the decoder manifold results in the following per-sample loss term:
(RF(I) = - log pz()) +
2 log det (x)) + B x11?,
where
g(=), and / is hyperparameter.
The log-determinant term is the difficult one to optimize. Fortunately, its gradient with respect to the parameters 0 of the decoder can be estimated tractably by the following construction (Caterini) et al. 2021). Note that = ge but the 0 subscript is dropped to avoid clutter. The relevant quantity (with J (z)):
(JTJ)-
a 0 (TT)))
log det(JT
(4)
00,
00,
The trace can be estimated (Hutchinson
1989
Girard
1989 by:
0 1 00, 2
K
8 J
log de J)
J)++ ~
(J Jee
(5)
2K
00,-
k=
where the 6 K samples from distribution where FlecT] II typically either Rademacher or standard normal. Now, we take steps to compute the above without constructing the entire net- work Jacobian J. First, note that (+) can be written a Jacobian-vector product ii Ju and a vector-Jacobian product J 01 = (n) J) each readily available via automatic differentiation. For (+*), Caterini et al. (202T) propose employ the iterative conjugate gradient method: Write &g (JTJ)
CG(J ) where CG(A; b) denotes the conjugate gradient solution to Ar = b. Conjugate gradient is clever choice here, since again only requires computing terms of the form Ju using autodiff. The parameter derivative can be made to act only on the rightmost Jacobian terms by applying the st stop-gradien ent operation to the output of the conjugate gradient method, which returns its input, but has zero gradient. The final surrogate for the fog-determinant term is therefore:
2K Σ stop-gradient (ce (J JiEN) 7) JJJERR JER.
1
which replaces the log-determinant term in the loss. This computation yields the same gradient as the original loss in eq. (3), even though it has different value. In the following, we present signif- icantly improved version of this estimator based on the insight that the encoder is an (approximate) inverse of the decoder.
4 FREE-FORM INJECTIVE FLOW (FIF)
Our modification rectangular flows is threefold: first. use unconstrained autoencoder archi- tecture (no restrictively parameterized invertible functions); second, we introduce a more computa- tionally efficient surrogate estimator; third, modify the surrogate avoid pathological behavior
4

---

Published as a conference paper at ICLR 2024
related to manifolds with high curvature. Our per-sample loss function is:
K
1 K
((x) - - logg
& (xx)
stop.gradient (d((=)ex) + B x1?
(7)
with fxx). Note the negative sign before the surrogate term, which comes from sending the Log-determinan gradient to the encoder rather than the decoder. We will derive and motivate this formulation of the loss in sections and
4.1 SIMPLIFYING THE SURROGATE
We considerably simplify the optimization of rectangular flows by a new surrogate for the log- determinant term in eq. ( which uses the Jacobian of the encoder as an approximation for the inverse Jacobian of the decoder. This allows the surrogate to be computed in a single pass, avoiding costly conjugate gradient iterations.
We do this by expanding the derivative in eq.
4b:
2 00 (JT))
J 8 303"
2 ir (J
Ir
(8)
where /²= (JT))\\) is the Moore-Penrooe inverse of J. The full derivation is in appendix B To see the advantage of this formulation consider that for encoder and decoder g optimal with respect to the reconstruction loss, f'(2) g' (x) (see appendix B.1 Using the stop.gradient
operation, this leads the following surrogate loss term:
stop gradient (f)) g((2)a
(9)
K
equivalently, using the encoder Jacobian in place of the decoder Jacobian (see appendix B):
K
Σ
[f()sstooo gradient (g'(2)ex).
(10)
K
Each term of the sum can be computed from just two vector-accoian/Jacoiinn wwwpaiiiinnaccoba---eeeettteteee vector products obtained from automatic differentiation. This significant improvement on the iterative conjugate gradient method needed in the original formulation of rectangular flows which requires up to (dd+1) such products to ensure convergence (Caterini et al. 202I). Compared to reconstruction loss only, measure 1.5x to 2x the wall clock time, independent of the latent dimension. Note that the surrogate is only accurate f and are (at least approximately) optimal with respect to the reconstruction loss. We observe stable training in practice, validating this assumption.
4.2 PROBLEMS WITH MAXIMUM LIKELIHOOD IN THE PRESENCE OF A BOTTLENECK
Rectangular flows are trained with a combination of reconstruction and likelihood term. We might ask what happens if only train with the likelihood term, making an analogy to normalizing flows. In this case our loss would be:
1 log pz)) + N log det [g'(2) g')). T')))
(xLL(x)
(11)
Unfortunately, optimizing this loss can lead us to learn degenerate decoder manifold, an issue raised in Brehhmer & Cranmer (2020). Here we expand on their argument. First consider that if and g are consistent, then (5) f(g(f(x)) fxx) and the per-sample loss is invariant to projections: LNLL(E) (x), since LNLL is function only of f(x). This means that we can write our loss as:
LNLL = EE (CNLL
(12)
where Paama (.) is the probability density of the projection of the training data onto the decoder manifold. Now consider that the negative log- likelihooood loss is one part of a KL divergence, and KL divergences are always non-negative:
KT. (T. (6) (llp((() (fttaa F.p (2) (ogg.o
po()
(13)
5

---

Published as a conference paper at ICLR 2024
Naive gradient computation
Corrected gradient computation
[=0
t-2
==
Figure 2: Naive training of autoencoders with negative log-likelihood (NLL,, section 4.2 leads to pathological solutions (left). Starting with the initialization (t 0. black), gradient steps increase the curvature the learnt manifold projected data is reduced, by moving the points closer to one another. This effect is stronger than the reconstruction loss. We fix this problem by evaluating the volume change off off-manifold (right). This moves the manifold closer to the data and reduces the curvature (t 1, 2, green), until eventually centers the manifold on the data with zero curvature (t = xc, green). Light lines show the set of points which map to the same latent point. Data is projected onto the = 2 manifold.
1, 2, orange). This reduces NLL because the entropy the
As result, the loss lower bounded by the entropy of the data projected onto the manifold:
NLLL = - (+) [log (+) 2 ((daaa(i)).
(14)
Unlike in the entropy depends on the projection learned by the model. Thus, the model could modify the projection such that entropy is low possible. We break this pathology down into two cases:
standard normalizing flow optimization, where the right hand side would be fixed, here
A model manifold which does not align with the data manifold but instead intersects it. For example, [Brhhmer & Cranmer (2020) discuss case where linear model learns to project a data distribution to single point on the manifold, thus reducing its entropy to -oc, the lowest possible value. To the best of our knowledge, this be fixed by adding noise and a reconstruction loss with sufficiently high weight. In appendix C we prove as much for linear models and characterize the solutions, which the same PCA if B 1/2g² where o² is the smallest eigenvalue of the data covariance matrix.
2. A model manifold which concentrates the data by use of high curvature, see fig. (left) This newly identified pathological case only occurs in nonlinear models, henee Brehmer & Cranmer (2020) did not notice this effect in their linear example. Importantly, this is not fixed by adding reconstruction loss.
Most existing injective flows avoid this by a two-stage training, which first learns a projection and then the distribution of the projected data in the latent space. To enable jointly learning manifold and maximum-likelihon density on it, we need to find a fix for the pathology.
Towards rell-behaved loss The term which leads pathological behavior in the likelihood loss is the og-dettrminnat. When using the change of variables with f evaluated at f. all that matters is the change volume from the projected data to the latent space, so the model can decrease the loss by choosing a manifold which concentrates the projected data more tightly (the more possibility it has to expand the data, the lower the loss will be). We can counteract this effect by introducing a factor inversely proportional to the concentration. This can be achieved by the fairly simple modifi- cation evaluating f' in our estimator at z rather than é. Namely, we modify eq. (10) to estimate the gradient of the log-determinant term by:
K T [aaatstopggrrrtttnnt (g((=)e))-
(15)
K
See appendix B.2 for detailed explanation.

---

Published as a conference paper at ICLR 2024
15
15
W
W
8.0
0.0
Loos
-05
005
Ukelhood
Seconctnurtion 19nn.
Figure 3: Learning a noisy 2-D sinusoid with a 1-D latent space for different reconstruction weights 3. Color codes denote the value of the latent variable each location. When the reconstruction term has low weight (left), the autoencoder learns to throw away information about the position along the sinusoid and only retains the orthogonal noise. Only sufficiently high weights (right) result in the desired solution, where the decoder spans the sinusoid manifold. The middle plot shows the tradeoff between reconstruction error and NLL as we transition between these regimes (box plots indicate variability across runs).
In this way, we discourage pathological solutions involving high curvature. In fig. 2 (right) we can see the effect of the modified estimator: the manifold now moves towards the data since the optimization is not dominated by diverging curvature. We note that the modified estimator is also computationally cheaper, since the vector-Jacobian product cf f' (z) can reuse the computational graph generated when computing z.
Along with the results of section 4.1
this leads to the following loss (same eq. 7
[(x)= = ÊNLL (x) + 3 Crecon. (x)
(16)
K
log (z))
(x) stop.gradient (g'(2)() + BB x/|2.
(17)
K
k=1
Phase transition Figure shows that when using this loss, if 3 is large enough, the dominant manifold direction is identified. In appendix 1.1.2) we show similar experiment MNIST.
5 EXPERIMENTS
In this section, we test the empirical performance of the proposed model. First, we show that our model much faster than rectangular flows on tabular data. Second, we show that it outperforms previous SOTA injective flows generating images. Finally, we compare against other generative autoencoders on the Pythae image generation benchmark Chadebee et (2022), achieving the best FID) score in some categories.
Implementation details In implementing the trace estimator, we have to make a number of choices. Briefly, i) we chose to formulate the log-determinant gradient in terms of the encoder rather than decoder as it was more stable in practice, ii) we performed traces in the order f'(x)g as this reduces variance (both orderings f'(x)g( (z) is a x d matrix whereas g (z)) is D. the former is typically easier to estimate), iii) we used mixture of forward- and backward-mon patible with our estimator, and iv) we used orthogonalized Gaussian noise in the trace estimator, to reduce variance. Full justification for these choices is given in appendix D
valid due to the cyclic property of the trace but since
automatic differentiation
this was com-
5.1 TABULAR DATA
We evaluate our method on four of the tabular datasets used by Paaaamakarios et al. (201)). using the same data splits, and make comparison to the published rectangular flow results (Caterini et al.) 2021), see table I We adopt the "FID-like metric" from that work, which computes the Wasserstein- 2 distance between the Gaussian distributions with equal mean and covariance as the test data and the data generated by the model. This is measure of the difference of the means and covariance matrices of the generated and test datasets. We outperform rectangular flows on all datasets except speedup in training time of between 1.5 and 6 times between FIF and
GAS. In addition,

---

Published as a conference paper at ICLR 2024
Table I: Free-form injective flows (FIF) are significantly faster than rectangular flows (RF) with
superior performance in FID-like metric on 3 out of tabular datasets (Papamakarios et al. 2017
Both methods use K 1. The results for RF are taken directly from (Caterini et al. 2021
Method
POWER
GAS
HEPMASS
MINIBOONE
RF (Caterini et al. 2021
0.083
0.015
0.110 + 0.021
0.779 0.191
1.001 0.051
FIF (ours)
0.041 + 0.007
0.281 + 0.031
0.541 + 0.034
0.598 = 0.024
Training Time Speedup
3,9 x
2.2
6.1 X
1.5
a rerun of rectangular flows (using the published code) on the same hardware. Full experimental
details are appendix E.2
We also show an ablation study in table 6 and table in the appendix disentangling the effect of the three individual components of our method: The surrogate of eq. (IO), the fix for high curvature solutions in eq. (5)) and the use of free-form architectures. Comparing to table D we find that the fix to estimate the negative log-likelihood with an off-manifoll encoder Jacobian is crucial for good performance of free-from architectures, as the on-manifold variant diverges (see section 4.2)
This is not the case for the injective architecture based on coupling flows, indicating a stabilizing regularization via the architecture.
5.2 COMPARISON INJECTIVE FLOWS
We compare FIF against previous injective flows on CelebA images (Liu et al. 2015) in table Our models significantly improve the quality of the generated images in terms of the Fréchet inception distance (FID) Heusel et al. (017)) and Inception Score (IS) (Sllimans et al. (0116). The former compares generated samples to a set of reference samples, by computing the Wasserstein-22 distance between two Gaussian distributions fit to later measures diversity by the entropy of the distribution class labels in the generated samples, where the class labels provided by some pre-trained classifier. Samples from this model are depicted in fig.
embedding of the respective sets of samples. The
For fair comparison, train each model the same hardware for equal wall clock time with the code provided by the authors. The architectures of previous works were dominated by the need that most layers invertible and have tractable Jacobian determinant. Our loss in eq. 7 does not impose these constraints on the architecture, and we can use an off-the-shelf convolutional auto-encoder with additional fully-connected layers in the latent space. Details can be found in appendix E.3
5.3 COMPARISON GENERATIVE AUTOENCODERS
As free-form injective flows (FIF) do not require any specific architecture, we expand our compar- ison to the much broader range of generative autoencoders. This is general class of bottleneck architectures that encode the training data to a standard normal distribution, so that the decoder can be used as a generator after training.
Recently, Chadebec et al. (2022) proposed the Pythae benchmark for comparing generative autoen- coders on image generation. They evaluate different training methods using two different architec-
CelebA under equal computational budget. Free-form Injective Flows (FIF) outperform previous work significantly in terms FID.
Table 2: Comparison of injective flows
N sampler
GMM sampler
Model
# parameters
FID |
IS f
FID
IS
DNF Horvat & Pfister 2021
39.4M
55.6 + 0.59
1.9
52.7
0.33
2.0
Trumpet (Kothari el al.
2021
19.
56.2 + 1.39
1.8
47.7 +: 2.24
1.9
FIF (ours)
34.3M
47.3 + 1.39
1.7
37.4
1.35
2.0
8

---

Published as a conference paper at ICLR 2024
Table Pythae benchmark results on CelebA, CollowingChhdebec Chadebec et al. (2022). We train their ar- chitectures (ConvNet and ResNet) with our new training objective, achieving SOTA FID on ResNet. We draw latent samples from standard normal "N" or GMM fit using training data "GMM". Mod- els with multiple variants (indicated in brackets) have been merged to indicate only the best result across variants. We mark the best FID each column in bold and underline the second best.
ConvNet
A
ResNet NV
ConvNet
GMM
RecNet
GMM
FID 4
IS 1
FID
IS
FID
IS
FID
IS
VAE Kingma Welling
1.6
IWAF IWAF: of di 2015
55.7
1.9
67.6
1.6
52.7
1.9
64.1
1.7
VAE-lin NE Rezenide & Mohamed 2015
56.5
1,9
1.6
53,3
62.8
1.7
VAE-IAF Kingna
2016
55.4
1.9
66.2
1.6
53.6
62.7
1.7
B-(TC) VAB Higgins of 2017 7 Chen al. 2018a
55.7
1.8
65.9
1.6
51.7
1.9
59.3
1.7
Kim Minli
InfoVAE (KBFIIOQ) el cl 2017
55.5
1.9
66.4
1.6
52.7
1.9
62.3
1.7
AAB Makhzzmi et 2013
59.9
64.8
1.7
53.9
58.7
1.8
MSSSIM- VAL
20177
124.3
1.3
119.0
1.3
124.3
119.2
1.3
Vanilla AE
327.7
1.0
275.0
2.9
55.4
2.0
57.4
1.8
WAE- (RBE-RM)MM (Tolstikhin
2018
64.6
1.6
51.7
57.7
1.8
VQVAE evan den Ooii et ar. 7
306.9
1.0
140.3
2.2
51.6
2.0
57.9
1.8
RAE (1.2GG et 2020)
168.7
FIF (ours)
56.9
2.1
62.3
1.7
47.3
1.9
55.0
1.8
tures on MNIST (eCunn 2010 (data D 784, latent d 16), CIFARIO (Krizhessky 2009 (D 3072, d 256), and CelebA (Liu et al. 2015) (D 12288, d = 64). All models are trained with the same limited computational budget. The goal of the benchmark to provide fair com- parison of different models, not to achieve SOTA image generation results, as this would require significantly more compute.
As shown in table our model performs strongly the benchmark, achieving SOTA on CelebA in Fréchet Inception Distance (FID) (Heusel et al. 2017) on the ResNet architecture with latent codes sampled from standard normal, and on both architectures when sampling from a Gaussian Mixture Model fit using training data. At the same time, the Inception Scores (IS) (Salimans et al. 2016) high, indicating a high diversity. On the combination where model does not outperform the competitors, FIF still achieves a comparable FID and high Inception Score. FIF also performs strongly the other datasets, see appendix E.4
For each method in the benchmark, ten hyperparameter configurations are trained and the best model according FID reported. For our method, we choose vary 10, 15, 20, 25 and the number of Hutchinson samples K = 1, 2. We find the performance to be robust against these choices, and give all details on the training procedure in appendix E.4
CONCLUSION
This paper offers a computationally efficient solution jointly learning manifold and distri- bution on it, which we call the free-form injective flow (FIF). We i) significantly improve an ex- isting estimator for the gradient of the change of variables across dimensions, ii) note that it can be applied to unconstrained architectures, iii) analyze problems with joint manifold and maximum- likelihood training and offer a solution, and iv) implement and test our model toy, tabular and image datasets. We find that the model is practical and scalable, outperforming comparable injec- tive flows, and showing similar or better performance other autoencoder generative models on the Pythae benchmark.
Several theoretical and practical questions remain for future work: We identified a previously over- looked problem with jointly learning a manifold and maximum likelihood. We propose fix in section 4.2 that provides high-quality models, but further investigation is needed for a thorough un- derstanding. Fitting GMM to the latent space after training improves performance on image data, suggesting that our latent distributions are not perfectly Gaussian. We generally find that architec- tures with fully- conneceddd layers in the latent space have more Gaussian latent distribution, suggesting that larger models suffer less from this problem. We leave potential theoretical or practi- cal improvements future work.

---

Published as a conference paper at ICLR 2024
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, llya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wat- tenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learn- ing on heterogeneous systems, 2015.
Lynton Ardizzone, Jakob Kruse, Carsten Rother, and Ullrich Köthe. Analyzing Inverse Problems with Invertible Neural Networks. In International Conference on Learning Representations, 2018.
Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen. Invertible residual networks. In International Conference on Machine Learning, 2019.
Jan Jetze Beitler, Ivan Sosnovik, and Arnold Smeulders. arXiv:2111.00619, 2021.
PIE: Pseudo-invertible encoder.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 3588):1798-1888, 2013.
Vanessa Böhm and Uroš Seljak. Probabilistic auto-encoder. Transactions on Machine Learning Research. 2022.
Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density estima- tion. Advances in Neural Information Processing Systems, 33:442-453, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamers. Advances in neural information processing systems, 33:1877-19011 2020. 2020.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv:1509.000. 2015.
Anthony I. Caterini, Gabriel Loaiza-Ganem, Geoff Pleiss, and John P Cunningham. Rectangular flows for manifold learning. Advances Neural Information Processing Systems, 2021.
Clément Chadebec, Louis Vincent, and Stéphanie Allassonnière. Pythae: Unifying generative autoencoders in python a benchmarking case. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, 2022.
Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disen- tanglement in variational autoencoders. Advances in Neural Information Processing Systems, 31, 2018a.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018b.
Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and föm--ennil Jacobsen. Residual flows for invertible generative modeling. Advances in Neural Information Processing Systems, 2019.
Eike Cramer, Felix Rauh, Alexander Mitsos, Raúl Tempone, and Manuel Dahmen. Isometric mani- fold learning for injective normalizing flows. arXiv preprint arXiv:2203.0380 03934, 2022.
Edmond Cunningham, Renos Zabounidis, Abhinav Agrawal, Madalina Fiterau, and Daniel Sheldon. Normalizing flows across dimensions. International Conference on Machine Learning, Workshop Track, 2020.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear Independent Components Estimation. In International Conference on Learning Representations, Workshop Track, 2015.
10

---

Published as a conference paper at ICLR 2024
Felix Draxler, Christoph Schnörr, and Ullrich Köthe. Whitening Convergence Rate of Coupling- based Normalizing Flows. In Advances Neural Information Processing Systems, 2022.
Felix Draxler, Peter Sorrenson, Lea Zimmermann, Armand Rousselot, and Ullrich Köthe. Free- form Flows: Make Any Architecture a Normalizing Flow. In Artificial Intelligence and Statistics, 2024.
William Falcon and The PyTorch Lightning team. PyTorch lightning, March 2019.
Amur Ghose, Abdullah Rashwan, and Pascal Poupart. Batch norm with entropic regularization turns deterministic autoencoders into generative models. In Conference on Uncertainty Artificial Intelligence, 2020.
Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. International Conference on Learning Representations, 2020.
A Girard. A fast 'Monte-Carlo cross-validation' procedure for large least squares problems with noisy data. Numerische Mathematik, 56:1-23, 1989.
Gene H Golub and Victor Pereyra. The differentiation of pseudo-inverses and nonlinear least squares problems whose variables separate. SIAM Journal on numerical analysis, 10(2):4134432 1973.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In International Con- ference on Learning Representations, 2019.
Luigi Gresele, Giancarlo Fissore, Adrián Javaloy, Bernhard Schölkopf, and Aapo Hyvarinen. Rel- ative gradient optimization of the jacobian term unsupervised deep learning. In Advances in Neural Information Processing Systems, 2020.
Charles Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Bérard-Marchane Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array program- ming with NumPy. Nature, 585(7825):357-36 2020.
Trevor Hastie, Robert Tibshirani, and Jerome H Friedman. The Elements of Statistical Learning: Data Mining, Inference, and Prediction, volume 2. Springer, 2009.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in Neural Information Processing Systems, 30, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. Beta-vae: Learning basic visual concepts with constrained variational framework. In International Conference Learning Representations, 2017.
Christian Horvat and Jean-Pascal Pfister. Denoising normalizing flow. Advances in Neural Informa- tion Processing Systems, 34:9099-9111 2021.
J. D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering, 9(3): 90-95, 2007.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics-Siman Simulation and Computation, 18(3):1059-107) 1989.
Thomas Keller, Jom WT Peters, Priyank Jaini, Emiel Hoogeboom, Patrick Forré, and Max Welling, Self normalizing flows. In International Conference Machine Learning, 2021.

---

Published as a conference paper at ICLR 2024
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference Machine Learning. 2018.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. Advances in Neural Information Processing Systems, 29, 2016.
Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker. Normalizing Flows: An Introduction and Review of Current Methods. IEEE Transactions Pattern Analysis and Machine Intelli- gence, 43 1):3944-3979, 2021.
Konik Kothari, AmirEhsan Khorashadizadeh, Maarten V. de Hoop, and Ivan Dokmanic. Trumpets: Injective flows for inference and inverse problems. In Conference on Uncertainty Artificial Intelligence, 2021.
Steven G Krantz and Harold R Parks. Geometric Integration Theory. Springer Science & Business Media, 2008.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, Toronto, Ontario, 2009.
Abhishek Kumar, Ben Poole, and Kevin Murphy. Regularized autoencoders via relaxed injective probability flow. In International Conference Artificial Intelligence and Statistics, 2020.
Yann LeCun, Corinna Cortes, and CJ Burges. MNIST handwritten digit database. ATT Labs [On- line]. Available: http://yann.lecmmca wwww..com/xxd///misst..t 2, 2010.
Ziwei Liu, Ping Luo. Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Gabriel Daiza--Ganem, Brendan Leigh Ross, Luhuan Wu, John Patrick Cunningham, Jesse Cress- well, and Anthony L Caterini. Denoising deep generative models. In Advances in Neural Infor- mation Processing Systems, Workshop Track, 2022.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders. In International Conference on Learning Representations, 2015.
Wes McKinney. Data Structures for Statistical Computing in Python. Stéfan van der Walt and Jarrod Millman (eds.), 9th Python in Science Conference, 2010.
Francesco Mezzadri. How to generate random matrices from the classical compact groups. Notices of the American Mathematical Society, 54(5):592-60) May 2007. ISSN 0002-9920.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. Advances in Neural Information Processing Systems, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems, 2019.
Stefan T. Radev, Ulf K. Mertens, Andreas Voss, Lynton Ardizzone, and Ullrich Köthe. BayesFlow: Learning Complex Stochastic Models With Invertible Neural Networks. IEEE Transactions on Neural Networks and Learning Systems, 33(4):1452-1466 April 2022. ISSN 2162-2388. doi: D..110//TNNLL
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Interna- tional Conference on Machine Learning, July 2015.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
12

---

Published as a conference paper at ICLR 2024
Brendan Ross and Jesse Cresswell. Tractable density estimation on learned manifolds with confor- mal embedding flows. Advances Neural Information Processing Systems, 34:26635-26648, 2021.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Gianluigi Silvestri, Daan Roos, and Luca Ambrogioni. Deterministic training of generative autoen- coders using invertible layers. In The Eleventh International Conference Learning Represen- tations, 2023.
Jake Snell, Karl Ridgeway, Renjie Liao, Brett D Roads, Michael Mozer, and Richard Zemel. Learning to generate images with perceptual similarity metrics. In IEEE International Conference on Image Processing, pp. 4277-4281. IEEE, 2017.
Yunfei Teng and Anna Choromanska. Invertible autoencoder for domain adaptation. Computation, 7(2):20, 2019.
The pandas development team. Pandas-dev/pandas:: Pandas, February 2020.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. In International Conference Learning Representations, 2018.
Aaron van den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in Neural Information Processing Systems, 30, 2017.
Mingtian Zhang, Yitong Sun, Chen Zhang, and Steven Medonagh. Spread flows for manifold mod- elling. In International Conference on Artificial Intelligence and Statistics, 2023.
Zijun Zhang, Ruixiang Zhang, Zongpeng Li, Yoshua Bengio, and Liam Paull. Perceptual generative autoencoders. In International Conference on Machine Learning, 2020.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Information maximizing variational autoencoders, arXiv: 1806, 02262, 2017.
13

---

Published as a conference paper at ICLR 2024
A CHANGE OF VARIABLES FORMULA ACROSS DIMENSIONS
The change of variables formula describes how probability densities change as they are mapped through an injective "pushforward" function g. It is instructive to derive this formula when g is an
invertible function. Let pz be base density and px the pushforward density obtained by mapping samples from pz through g. Then we can write
px(x)= = / p(l: =)pz(=)dz
(18)
/ 5(x o(=))pz(2)d=
(19)
/ 8(x - z)pz(J (derd) di
(00)
(z())) |det(g (f(x)))
(21)
using the change of variables = g()), meaning that = f(0) and (=)) = dā with f the inverse of g.
Now suppose that g maps from d to RD with d < D. We can generalize the change of variables
= (2) using x = f()) and (x) g' (2) da = dē where f and are consistent (f o g is the
identity) (see chapter 5 Krantz & Parks 20088 This gives us
(x = / 6(x g(=))p2(2)))))))))
(22)
f 8(xz - (e)) (((()))))))))()))))=))))))))))))
(23)
This expression defines probability density in the full ambient space RD (albeit degenerate distribution) but we cannot easily remove the integral. However, we can convert into an expression resembling the full-dimensional case, but defined only on the image of
(x(x) pz(((z)) det(g' det(g'(f(x)) g' (f(()))))
(24)
Note that this expression only integrates to 1 if we restrict integration to the image of g. As such, it should only be regarded as defining a probability distribution on this manifold, not in the ambient space R D
B DERIVATION OF GRADIENT ESTIMATOR
We expand the derivative in eq.
log dot JJ
2 tr
JJ-
(J J)
(25)
00; 2
00,
2 Ir
1
(JTJ)-1 θ
JT
+ すう
ir
(J
JJT
8
(26)
00,
00;"
1 tr
(JJ 1)
0 01)
2
00,
+ tr
2
002
(27)
(JT
00, 1)
tr
(28)
tr jt 08 J
(29)
where we used the eyclic property of the trace and that tr(AB) tr(AT B Jt Jt= is the Moore-Penrose inverse of J.
(JTJ)-
Now will do an equivalent derivation for the encoder. Observe that, since
(30)
14

---

Published as a conference paper at ICLR 2024
can rewrite the deg-determinant term using the encoder Jacobian:
1 log det(J' J) log det(J
(31)
2
Note the negative sign on the right- hand side.
The derivation for the derivative is very similar to that for the decoder, where take deriva- tive with respect to encoder parameters Φ:
0 1 Depy 2
log det JtJ tr (Jtut))- (
(32)
tr (ut Jsy
Jog (D0,11) JtT)
tr (uIt
(
JTT
(33)
tr ( θ 1 jt
Jt
1
tr ((st)tT)-)))
0
j
0oj""
2
dØj
(34)
tr (Jt(ttttT)-
① j-t
(35)
tr J
(36)
رΦj
where we used the cyclic property of the trace, that tr(AB) = tr(A BT) and that J = (Jl)² t JTT (TT(JttT)=
Recall eq. (9) in section 4.1
which gives the surrogate for the log-determinant term:
1
stop gradient K
d(())
(37)
This is formulated in terms of the Jacobian of the decoder, other words is derived from dJt(JT J). The equivalent term, formulated in terms of the Jacobian of the encoder should be derived from dJt JfT) and is therefore (note the negative sign):
I K
Σ
x (((")stop.gradieee (5 (z)ex)
(38)
We write it in the order ('())g'(z) rather than g' (()) since this reduces the variance of the estimate. See appendix D.I for further details on this point.
B. NOTE ON OPTIMALITY WITH RESPECT TO RECONSTRUCTION LOSS
Our estimator relies on the approximation f' (f())) If and g are consistent (f g the identity) this guarantees that g' ff = I, but not that f'()) is the Moore-Penroe inverse of (x)) sufficient requirement is that f and g are optimal with respect to the reconstruction loss, that is, any variation in the functions would lead to higher reconstruction. With calculus of variations, is possible to show that such ff and g are consistent, and
(a - x) g'(()))
(39)
for all x. By taking the derivative with respect to x and evaluating at some x the image of g (so f 2) we have that
f()) (((T(((()))))))))))))))))))))))) ((((=))) ((f ((((()))
(40)
and hence
- of() o((((=))
(41)
In the remainder of the appendix, given encoder-decoder pair and g which are optimal with respect to the reconstruction loss, we refer to g as the pseudoinverse of f. and f as the pseudoinverse of g-
15

---

Published as a conference paper at ICLR 2024
Figure 4: Representation of ill-defined probability density p(x) x p(i)e-3(é=)) (left and center). Solid black lines denote the manifold, dashed lines are a constant distance from the manifold. The probability density is constant along the manifold. The width of the cyan bands is proportional to -- and represents the probability density along the and off- contours. While the density is reasonable for a flat manifold (left), note that the amount of probability mass associated with region of the manifold (bounded by solid lines) larger at some points off the manifold than on it when the manifold has curvature (center). This behavior can lead to divergent solutions when optimizing for likelihood and should be compensated for. The appropriate compensation factor is the ratio of the volume of small region on the manifold (small blue square embedded in green manifold, right) to the equivalent region off the manifold (large blue square, right). The blue arrows represent an orthonormal frame on the manifold, and the equivalent frame in the off- offifoldl region.
B.2 MODIFIED ESTIMATOR
As stated in 4.2 we modify the log-determinant estimator (eq. (38 ) by replacing ('))) by f'(x). The loss we are trying to optimize (sending gradient from the log- fog-determinant to the encoder) is:
[(x) = - (f())(f())) - log vol (f + Bl =I1?
(42)
with ()))) = (f Consider the probability density ğ implied by interpreting this loss as a negative log-likel;hood:
p(x) p(0)e
#||-=||| -( pz(2) vol (f) -f -1
(43)
where p(ē) is the on-manifold density. Unfortunately, this density is ill-defined and leads to patho- logical behavior (see fig. 4. order to provide correction to this density, we need a term which compensates for the volume increase or decrease of off aff-manifold regions in comparison to the on- manifold region they are projected to. This is depicted in fig. (right). The blue arrows in the on-manifold region will span the same latent-space volume as the blue arrows in the off-manifold region. The change in volume between the depicted on-manifold region and the latent space is vol(f'(#) and vol(f'(z)) between the off-manifold region and the latent space. Combining these facts means
(a) (volume of on-manifold region) vol(f'(z) x (volume of off-manifolle region) (44)
and hence the ratio of the volume of the on-manifold region the off-manifold region is vol(f((z))) vol()))/voo(' v.)) Multiplying ((x) by this factor leads to
vol(f'(x)) = pz(2) vol (f'( -111--111 p(x) (
(45)
and the corresponding negative log likelihood loss
[(x) log))(((x)) log vol (f'xz) B|| 20113
(46)
16

---

Published as a conference paper at ICLR 2024
The surrogate for the log log-determinant term is therefore
"nf""rrastoo.graddeenf((ff (f())))
(47)
In order to maintain computational efficcincy, we approximate (x)x) by (f(())))
tt =ttl///)sttop.prrddnn((((((((
(48)
giving the stated correction.
C LINEAR MODEL TRAINED ON MAXIMUM LIKELIHOOD ALONE
Consider linear model, trained on data with zero mean and covariance E. Let the encoder function be f(z) = Ax and suppose that A has positive singular values, meaning that AA" is positive definite. Let the decoder function g(=) = A"z, z, where At A (AAT)-I, We want to minimize a combination of negative log-likelihood and a reconstruction loss (here we use 1/2o² instead of B as prefactor):
C= LNLL + Lrecon.
(49)
= E
|x
log (eet(AA))))) +
A Ar =x1?
(50)
2
2
20
Ex
Ax] A
log (cc((AAA)))) +
Ex(
(AI - a]
(51)
20
1
tr(AEr (100 T ]AT) 2
1
log det AA
1
tr F. [xx T J( At ATA))
(52)
20²
1
= tr(AEAT) log det det (AAT) -
tr(E(I - A\A))
(53)
2
20
where a full rank d x matrix with d D.
Before solving for the minimum, let's review some matrix calculus identities. It is often convenient to consider A as a function of single variable æ, differentiate with respect to æ, and then choose to be Aij. Then the derivative is
d d'i.j
A E(5)
(54)
where E(5) is matrix of zeros, except for the ij entry which is a one. We can write this as E dnd where d is the Kronecker delta. When evaluating El0) inside trace get the simple expression:
tr(E B) Ed Bu da Bix Bj using Einstein notation. The additional matrix identities we will need are Jacobi's formula for a square invertible matrix B:
(55)
B))
dct(B) = Br B-1
(56)
dx
dr
and hence
d
d dx
log del (3) = tr B
B
(57)
dr
and we will prove the following lemma.
Lemma C.1. Suppose the matrix A depends variable Then we have the the following expression for the derivative of the projection operator At A:
dx (A'A)= - (
[I A'A)+ (IA))+++)+++
Al
d dx A (I- (I-A1A)
(58)
dx
Proof. The following is based on the proof to lemma 4.1 in Golub & Pereyra 1973) 3).
Define the
projection operator PA A A and its complement P3 = I PA. Then, since PAP. PA.
d PAPA dx
dx
PA
=
PA dx
PA + PA
dz
PA
(59)
17

---

Published as a conference paper at ICLR 2024
In addition, since PAAT= AA
( ( dx
PAAT = ( PA
(
(dA)
= (
AT
PA
(60)
dx
dx
dr
and therefore
(d PA Pa dx
PA
d PA dx
A ATT
(
A
AIT - PA
(
A
AlT
= PA
d
A
T
AlT
dx
dr
dx
(61)
By similar steps but using
APA
A, we can derive
d PA
d A
PA
At
P
(62)
dx
dr
Putting it all logether gives
d PA
d
At
A
P
+
A
A
P
(63)
dr
dr
dr
Note that the second term is just the transpose of the first.
Now are ready to find the derivative of the loss and set it to zero
Lemma C.2. The derivative of the loss with respect to A takes the form:
dA C = ΣAT A
02 (1 A'A)EAt)
(64)
Proof. Let's upply the above identities to the first term in the loss:
d 1
1
d
d AΣ
(5(A)
)
tr AAEA
2 tr
A
EAT
(65)
dr
dx
dr
d
tr
EAT
(66)
dr
since the trace is invariant under transposition and hence
d 1 (r(AEA)) ( =(EAA)))))))))))))))))))))))))))
(67)
dAi 2
Applying Jacobi's ' formula to the second term in the loss gives:
d
log det (AAT) tr (AAT)-I ( d ( (AAT)))
(68)
dr 2
tr (AAT)-1
A + A
d
(69)
dx
de
and therefore
d 1
log det (AAT
1
tr ((AAT)-1 (E()) + +E(0))
(70)
dA,.; 2
2
1
tr ATA
(71)
2
(AT(AAT)-I)
(72)
AAe
(73)
where we used the cyclic and transpose properties of the trace and that EUOT E(0))
18

---

Published as a conference paper at ICLR 2024
The final term requires derivative of tr(E At A)), which is equal a derivative of -t(EEttt) We use the formula for the derivative of the projection operator to get
d 1 tr(EAA
1 ل
(Al
tr
(74)
dx 2
2
dx
tr
ΣAt
A
(I
AT
(75)
dr
again using the transpose property of the trace, and therefore
d
trEAAA A) (I(E A'AJEA') ((T A A)EA))
(76)
dAi
Putting the three expressions together, we have that
dA -= [EAT - Al
1 r² A'A)EA!)
(77)
Lemma C.3. The critical points of C satisfy the following properties:
1. A A= with UUT= =I
2. U T commutes with 2
Proof. Using lemma
C.2 the critical points satisfy
ΣAT T - A (T- AL A)EA =
(78)
By multiplying by A from the left we have
AΣA = la
(79)
meaning that AΣ must have orthonormal rows (since UUT T). With this definition, we can write = UΣ- 1,
If we now multiply by A from the right, we get
1 EA
1 A 7
ΣATA - A
(80)
Noting that the second and fourth terms are symmetric (since At A is symmetric), this means that the remaining terms must be symmetric:
1A--1AA A 02 -
EA A
A A
G2 A'
(81)
Since ATA commutes with A A, they are simultaneously diagonalizable, and since they are both symmetric, they share an orthonormal basis of eigenvectors. Clearly A A - A A/o" has the same basis. Since this matrix commutes with Σ, it must share basis with Σ and hence has the same basis as AT A. This means that Σ commutes with A A.
Expanding in terms of U, this means that
ΣATA A =
AT - A AE
(82)
and therefore EUTU = U UE, meaning that Σ and TU commute.
Consider as an example the case where diagonal. U is a projection matrix and in this case must be diagonal due commuting with Σ. As result, it must have exactly d ones and D - d zeros along the diagonal. This means that the rows of U are a basis for the d dimensional axis-aligned subspace corresponding to the nonzero entries. In the case of non-diagonal E, this generalizes to the rows of spanning the same subspace as some subset of d cigenvectors of Σ. This leads to the expression of the loss function in the next theorem.
19

---

Published as a conference paper at ICLR 2024
-0.75
-1.00
-1.25
f(A))
-1.50
-177
-2.00
-2.25
-250
0.0
0.5
1.0
1.5
2.5
3.5
4.0
A
Figure 5: Plot of f(A:) = log A; - X1/o² with = 1, showing maximum value at A = o² and
unbounded behavior either side.
Theorem C.4. Let Σ have the eigen-decoompsition VAVT with diag(A). ()). Let U have the eigen-decomposition V EVV with E = diag(a). Then the minimun of the loss is satisfied by such that
La
ai log A Tz
(83)
minimal, subject to the constraint a4 {0, 13 with Σb a; d.
Proof. Let's note couple of properties. First, we have (UEUT)== U due E commuting with U TÜ, so we can say that f(UXU*) UfEE) for any matrix function f with Taylor series. Next, U on the diagonal. We know that the rank of U is d, hence E has exactly d ones and D d zeros along the diagonal. Therefore we have the constraint a {0, 13 with ΣP a4 d. Next, note that A = U.U.
TÚ U is an orthogonal projection matrix, so a diagonal matrix with ones or zeros
Now we substitute back into the loss in terms of U:
f =
2 ) log det (US + tr(E(T UT)
(84)
20
2 tr(U (((Ulog(E)U)) 202 tr(UEUT) + const.
(85)
where we used that
log del (UE lr troo= tr(U (r(Ulog(2)0T))
(86)
Note that tr(UUT) is constant. Consider that
tr(UEUT) tr(U =(VE t:::(((EDD) a
(87)
The same logic holds for the term with log(E)) Therefore, dropping constant terms, the loss can be written in terms of and X:
C.
log A; =2
(88)
20

---

Published as a conference paper at ICLR 2024
The loss will take different values depending on which elements of a are nonzero. Define f(A,) log A 14--// /a2 The loss will be minimized when the nonzero a; correspond to those values of ff (A,) which are minimal. Clearly (Ae) < 0, so f has only one maximum at X r² and is unbounded below on either side of this maximum (see fig. 5). Consider the two extreme cases:
1. All cigenvalues are smaller than a². The minimal values of f(x,) will occur for the smallest values of A₃ Hence the d smallest cigenvalues of 2 will be selected.
2. All eigenvalues are larger than x², The minimal values of / f(A,)) will occur for the largest values of A₁. Hence the d largest cigenvalues of E will be selected.
In the intermediate regime, there will be a phase transition between these two extremes.
In the first case, the reconstructed manifold will be projection onto the d-dimensional subspace with the lowest variance, exactly the opposite result to PCA. In the second case, the reconstructed manifold will be a projection onto the d-dimensional subspace with the highest variance, exactly the same result as PCA. If we maximize the likelihood on the manifold without any reconstruction loss, corresponding to the o² - do limit, actually learn the lowest entropy manifold. It makes more sense to learn the highest entropy manifold as in PCA. We can ensure this is the case by adding Gaussian noise of variance o² the data, ensuring that the minimum eigenvalue of the covariance matrix is at least o², even if the original data is degenerate.
D IMPLEMENTATION DETAILS
In implementing the trace estimator, we have to make a number of choices, each claborated below. The main reasons for each choice are given first, with more technical details deferred to later in the appendix.
Gradient to encoder or decoder The log- degermeeemnnnt term can be formulated either in terms of the Jacobian of the encoder (see eq. (TOP) or the decoder (see eq. (9) ). As discussed in section 4.2 find that formulating it in terms of the encoder Jacobian leads to more stable training. Since the training also minimizes the squared norm of f() we speculate that having gradient from this term and the log-determinant term both being sent the encoder allows the encoder to more effi- ciently shape the latent space distribution. We note the similarity of this formulation to the standard change-of- varviabbless loss used train normalizing flows. If we instead send the gradient of the log degedttermnnnt to the decoder, the information about how the encoder can change can only reach it via the reconstruction term, which doesn't 't allow the encoder to deviate significantly from being the pseudoinverse of the decoder. A change in the decoder will therefore lead to corresponding change the encoder, but this a less direct process than sending gradient to the encoder directly. In addition, this formulation means the decoder is optimized only to minimize reconstruction loss, meaning that it will likely be an approximate pscudoinverse for the encoder, a condition we require for the accuracy of the surrogate estimator.
Space in which trace is performed Considering eqs. 29 and (36), the central component of the surrogate is trace (estimator). Making use of the cyclic property of the trace, i.c. tr(A B) tr(BAT) for any A, B E TR. Dxd we can choose which expansion of the trace to estimate:
Σ [(ATB), (AT B)
(BA)): ).. 1=1
(89)
The variance of a stochastic trace estimator depends on the noise used but in general is roughly proportional to the squared Frobenius norm of the matrix (see appendix [D.3.. Given two matrices A, B E Dxd with d < D, it is likely that ||A B < |BATT] 1 This statement is not true for all A and B, but is almost always fulfilled when d < D.
Transferred to context: In general the matrices f((x) E Rdd and g'(2) E R Dxd rectangular and can be multiplied together in either the f'(x)g'( (2) E Rdxdd order or (x) E DxD order. This matters for applying the trace since generally d < D.
A more precise statement (proven in appendix D.I is that the entries of A and B are sampled from standard normal distributions, then El| A B [ Ddl versus Fi[] BA [|2.] dd.. For
21

---

Published as a conference paper at ICLR 2024
Table Different possible estimators for the gradient of the log-determinant term.
gradient to encoder
gradient to decoder
trace in data space
[g((z)
30, (x))
0
38, of (xx)) (() tr (f) 38,-
trace in latent space
tr (
a Jey
f g((2)
g (2)
d < D the difference becomes significant. The difference between the two estimators may not be large if the two matrices have special structure, particular if they share a basis. However, since the terms being multiplied in our case a Jacobian matrix and the derivative of another Jacobian matrix with respect to a parameter 0g or Φj. it is unlikely that any such particular structure is present.
As result, when the latent space is smaller than the data space, the preferable estimator is the one that performs the trace in the latent space, meaning that products in the estimator have the order f' ((x)g'(>>) (z) (see table 4). In appendix [.I.1] we experimentally test the convergence of trace estimators with increasing Hutchinson samples, performed in both data and latent space, confirming that convergence is much faster when performing the trace in latent space.
Type of gradient Consider the estimator:
tr (10) 0'(2)) θ
DΦ, E. [
(90)
Ignoring the stop gradient operation for now, this requires computing terms of the form ('(x)g'(z)) In order to avoid calculating full Jacobian matrices, we can implement the calcu- lation using some combination of vector-Jacobian (vjp) or Jacobian-vector (ivp) products, which are efficient to compute with backward mode respectively forward-mode automatic differentiation. Note that we can use the result from one product as the vector for another vip or jvp. For example, can compute c f'(()) 01 (c f'(x)) E R yields vector, vector-Jacobian products.
i g' (z)c via two
This gives us three choices: i) backward mode only (two vip), ii) forward mode only (two jvp) or iii) a mix of both (one jvp and one jvp). We opt to use mixed mode (see appendix D.2 further details).
Trace estimator noise Trace estimators rely on the identity E, Ae] tr(AE,(<e tr(A), meaning that we require only E, [ee" to the variance of the estimator. Among all noise vectors whose entries are sampled independently, Rademacher noise has the lowest variance (Hutchinson) 1989). However, if the entries are sampled from a standard normal distribution and then scaled to have length vd where d the dimension of c, the entries are longer independent and the variance of the estimator is comparable to Rademacher noise (Giardd 1989). When using a single Hutchinson sample, we choose to use scaled Gaussian noise for its low variance, and since it covers more directions than Rademacher noise (covering the hypersphere uniformly, rather than at a fixed 2d points). When we have more than one Hutchinson sample, additionally orthogonalize the vectors this further reduces variance. More details in appendix D33
for the noise variable. The choice comes down mainly
Number of noise samples We choose use between and d noise samples in the trace esti- mator (with d samples we already can calculate the exact trace, so more samples are not necessary). Denote the number of samples by K. We find that general K is enough for good performance, especially if the batch size is sufficiently high.
D.1 VARIANCE OF TRACE ESTIMATOR
Theorem D.I. Let A. B E Dxd where the entries of both matrices are sampled from a standard normal distribution. Then
E [ B =dD
and
E [[IIAA T 1 2
(91)
22

---

Published as a conference paper at ICLR 2024
relstvv grselen actanes
rotttr grodlen atotance
16
batch 5:26
256
512
batch d2ze
$12
10"
14
16
2
10
12
14
16
number of Hutchinson samples
number Hu totchiisonn samples
Figure 6: Relative gradient distance to the exact surrogate gradient as a function of Hutchinson samples for varying batch sizes. (Left) Trace estimation in data space. (Right) Trace estimation in latent space. We estimate the trace using orthogonalized Gaussian noise (see appendix [D.355 (solid lines), but also feature non-orthogonaliizd samples for batch size in latent space (dashed line). Note the different scales on the y-axes and the use of symlog in the right-hand plot. We evaluate the trace estimation on the surrogate estimator with gradient encoder as specified in table for converged model trained on conditional MNIST (d = 16, D = 784).
Proof. Consider first ||A Bl We can write this
2
d
D
| B
A Bkg
(92)
Taking an expectation over this expression, the only nonzero contributions will be from terms where the A and B terms are both quadratic, since if not, the term will be multiplied by E[X] = 0 where X is standard normal. This requires giving
E $1A B
E
[A
Bi
(93)
k
d
D
Σ
Σ
E
[A]a]
E
[
(94)
1. k
d
(95)
since As and BRj are independent and the expectation of the square of a standard normal variable is its variance, i.e. 1. The equivalent expressions for BA D in these expressions.
} can be obtained by swapping d and
D.1.1 EXPERIMENTAL CONFIRMATION
To evaluate the convergence behavior of the trace estimation, which is exact for d and 1) Hutchin- son samples in latent and data space respectively, we compute the relative gradient distance of the resulting surrogate gradient with respect to the exact solution function of Hutchinson samples K:
relative gradient distance ()) Vsurrogate(K) - Vexact|e
(96)
/exact||2
Here, Vsurrogate(K)) denotes the gradient of the surrogate loss term after K Hutchinson samples and Vexact the gradient of the exact surrogate loss term, i.e. after d or D samples.
In fig. 6 one can see a clear decrease in gradient distance and variance when computing the trace in latent space instead of data space. Furthermore, we note that increasing the batch size also contributes to fast and steady convergence, which result of sampling an independent noise sample per batch instance.
D.2 FORWARD/BACKWARR AUTOMATIC DIFFERENTIATION
Two of the basic building blocks automatic differentiation (autodiff) libraries are the vector- Jacobian product (vip) and the Jacobian-vector product (jvp). The vip implemented by
23

---

Published as a conference paper at ICLR 2024
backward- mode autodiff and computes vector multiplied with Jacobian matrix from the left, along with the output of the function being used:
f(x), c f'(x) = vip(f.z.c
(97)
In PyTorch, this is implemented by the torch. aut autograd. functional, .vip function, or by first computing (z)), then using torch. autograd. grad.
The jvp is implemented by forward-mods autodiff and computes vector multiplied by Jacobian matrix from the right, along with the output of the function being used:
f(x), fx(e ivp(f,z,c)
(98)
In PyTorch, this is implemented by the torch. aut autograd. forward ad package.
Our preferred estimator for the log-determinant is the following (see section 4.2)
K
1
c fx)sstp gradient (g (2)cx)
(99)
K
Therefore we need compute terms of the form e f'(x)g' with a srop-gradient opera- tion. The stop.gradient operation is implemented by applying the detach () method to tensor PyTorch. We have the following options. Note that we can use product obtained from vip jvp vector input to a subsequent product.
Backward mode todiff.
This mode uses only vector-Jacobian products, requiring backward-mode au-
1. (((x)
2, 11 = vip(f,z,)) f, Un = 2;
2. 12 = =]]((2)
3. c (2)) uic
Forward mode This mode uses only Jacobian vector products, requiring nenward-mode autodiff.
1. = gz))
i ivp
2. 1/2 = f(x)
2, 1/2 = jvp(f.
3. c (())
Mixed mode This mode vector-Jacobian and one Jacobian-vector product, requiring both forward- and backward mode autodiff.
1. oo = f())
2, v1 vip(f,z,
2. 1/2 g'(z)e
f, U2 jvp(g,z,c)
3. e (=)=)
We prefer using backward mode autodiff where possible, since we find that it is slightly faster than forward mode in PyToorc. However for estimator of choice, most easily implemented. Using backward mode would require a stop-gradiemt operation to be introduced after the second step, but in way that allows gradient to flow back to f'(x). While we believe this is possible if implemented carefully, we did not pursue this option. In mixed mode, we can easily detach the gradient of 02 without affecting the first step of the calculation.
mixed mode, since this is
D.3 PROPERTIES OF TRACE ESTIMATOR NOISE
Hutchinson style trace estimators (Hutchinson. 1989) have the form E. Ae] and equal tr(A) in expectation. A is skew-symmetric (A - A), then (c Ac) = =eTAT A -eT Ae and hence Ac 0 with zero variance. Since any matrix be decomposed into symmetric and skew- symmetric part, the variance the estimator comes only from the symmetric part of A, namely A, (A + A //2. From on, suppose A is symmetric and if not, substitute A, for A.
24

---

Published as a conference paper at ICLR 2024
D.3.1 RADEMACHER NOISE
If the entries of are sampled independently from a distribution with zero mean and unit variance, then the variance of the estimator is minimized by the Rademacher distribution which samples the values -1 and 1 each with probability half. This estimator achieves the following variance for symmetric A (see proposition in Hutchinson (189))):
V.le Ac] 2
Σ ith
A
(100)
D.3.2 GAUSSIAN NOISE
With standard normal noise, the estimator is unbiased, but the variance higher (see again Hutchin- 1989)):
V.lc Ac = Σ = 2||A
(101)
iee. twice the Frobenius norm.
D.3.3 SCALED GAUSSIAN NOISE By contrast
T Ac
E
tr(A)
(102)
where e is standard normal variable in R4. The variance this estimator for symmetric A (see theorem 2.2 in Girard (1989) is:
E Ae 2 (A(A)
Ve
(103)
cTc
d
where (A(A)) denotes the variance of the eigenvalues of A.
We can write this estimator in the "Hutchinson" form by sampling c from standard normal distri- bution, then normalizing it such that length vd. Then we have
Edlc Ac| = tr(A)
(104)
and
2d
VEle Ac
2 (A(A))
(105)
d+2
D.3.4 COMPARISON
When the dimension of A becomes large, the variance of Rademacher and scaled Gaussian esti- mators are comparable. Suppose that the eigenvalues of A have zero mean (e.g. the entries are independent normal samples). Then
do (A(A)) = > tr(A)) = ||||?
(106)
we further assume that all entries of have roughly equal magnitude we have that
Σ 44 ~ ||||
(107)
(
since the sum in the Frobenius norm is dominated by the d(d 1) off-diagonal terms. Similarly,
d²
d 0² (X(A)) do²()(A)) |||]
(108)
meaning that the two estimators have approximately the same variance.
the matrix has special structure, we might choose one estimator over the other. For example, the standard deviation of the eigenvalues of A is small in comparison to the mean eigenvalue, the
25

---

Published as a conference paper at ICLR 2024
scaled Gaussian estimator is preferable and if A is dominated by its diagonal then the Rademacher estimator is preferable. We don't tt expect either type of special structure in our matrices, so we consider the estimators interchangeable. We decided to use scaled Gaussian noise since it produces noise which points in all possible directions in Rd whereas Rademacher noise is restricted to a finite 2d points. We assume that there is no reason prefer this set of 2 directions and therefore sampling from all possible directions is better.
D.3.5 REDUCING VARIANCE WHEN SAMPLING MORE THAN 1 HUTCHINSON SAMPLE
When the number of Hutchinson samples K are greater than 1, it is more favorable to sample the noise vectors in dependent way than independently. Consider the case of d d matrix A with K = d. Then we can get an exact estimate of the trace via
Σ 2 Aq. trQ"" AQ) = tr(A)
(109)
with orthogonal Q and qi the i-th column Q. If the qi were sampled independently, we almost certainly wouldn't it achieve this exact result. We therefore sample our noise vectors as the first K columns of randomly sampled d x d orthogonal matrix and scale each column by vd. We show below that this reduces variance compared with sampling independently and make an experimental comparison (see fig. 6 If the resulting noise vectors are denoted ei, we estimate tr(A) by
K
1 K
(A( =
Ac;
(110)
[=1
This estimator is unbiased:
K Σ
Exg.... [T[( =
E,, 1c, Aej
(111)
K
1 K
tr (AE,, [ccT]]
(112)
K 1=1
1
Σ ir(A)
(113)
K
(r(
(114)
since [e,[c] I for all 6,-
This procedure is equivalent to using scaled Gaussian noise when K = 1 and is what we implement in practice for all values of K. Note that it is not necessary to use K d since already achieve the exact value with d.
A note our sampling strategy: in practice sample by taking the Q matrix of the QR decompo- sition of d x K matrix with entries sampled from standard normal. Since the QR decomposition performs Gram-Schmidt orthogonalization, the Q matrix is uniformly sampled from the group of orthogonal matrices if Q is square (Mezzadri) 2007). The same logic applies to the QR decomposi- tion of non-square matrices, yielding a x K matrix made up of the first K columns of a uniformly sampled orthogonal d x d matrix. Strictly speaking, the QR decomposition is only unique if the R matrix has positive diagonal, and uniqueness required for uniform sampling. Let X QR and define by D the sign of the diagonal of R. D is diagonal with or -1 on the diagonal and 1)2 = T. Uniqueness can be achieved by multiplying Q by D from the right and multiplying R by D from the left: X = QDDR QR. The resulting uniformly sampled orthogonal matrix is QD, meaning the columns of Q are multiplied by either or -1. In our setting, we have terms of the form e, Ac, where the e are the columns of Q, multiplying the columns by -1 has effect the trace estimate. As a result we opt not to multiply by D. Therefore, although we do not sample uniformly from the orthogonal group, the final result is equivalent sampling uniformly.
26

---

Published as a conference paper at ICLR 2024
Variance derivation that
Using the formula for the variance of sum of random variables, have
K
V CR EK [T(A))
Σ
Ca.ee Acs, cJ Acgl
(115)
K²
1
where C' denotes the covariance between two random variables. Note that since the permutation of the columns of a randomly sampled orthogonal matrix is arbitrary (permuting the columns results in another randomly sampled orthogonal matrix of equal probability), all columns equivalent and we only have to distinguish between the cases i and i + j. This leads td
V VEL [Tr(A)]
K2 (Ko K(K 1)c) K (0+ (K-
(116)
with U = Ve, and == Ce..e,le, Ac,, c] Acy] when * j. Each column viewed individually
randomly sampled Gaussian vector, scaled have length /d, hence the value of is equal to the scaled Gaussian noise case above, namely
2d
(117)
2
where (A(A)) denotes the variance of the eigenvalues of A. We also know that the variance of the estimator reduces to zero when K and hence (d 1)c leading
(118)
d-1
Putting all together leads
V Vg..ex [TT(A)] = (v d-1 v)
-
(119)
K
1 d K
(120)
K d d-1
2d - K) k(d- ((d-1)(d+2) o (A(A))
(121)
valid for d > The comparable quantity for independently sampled scaled Gaussian noise is
2d
VA & [G((A)
2 (A(A))
(122)
K(d + 2)
iee. 1/K times the K result. The orthogonalized noise strategy always has lower variance since the ratio between the variances is
d - K
(123)
which even reduces zero when K = d. If dl is large, the difference is not great for small K, which aligns with the fact that randomly sampled directions in Rd are close to orthogonal for large d.
E EXPERIMENTAL DETAILS
E. ROLE OF RECONSTRUCTION WEIGITT
E. TOY DATA
To analyze the model behaviour depending on the reconstruction weight ,9, we train the same archi- tecture on simple sinusoid data set with / varied between 0.01 and 100.
For the generation of data points, we draw positions from a ID standard normal distribution and calculate the respective positions by = sin(Tx/22 22). Then, isotropic Gaussian noise with = 0.1 is added. We train an autoencoder architecture built with four residual blocks and
'This argument is http://///////////t/e//////////a
Findi naainniiariaan weeaaaaiaeee-eo of-the--tttteeeee samplamppe. wammple-men--poeee -f www.ooaa---aaaaaaaaaaaaaaaaaaaaa
acee ace
27

---

Published as a conference paper at ICLR 2024
Noise: 0.1
2
Loss value
0
Neg. Log Likellhaod Reconstruction Loss
0.01
0.05
0.1
0.2
0.4
0.6
0.8
1.0
2.0
4.0
6.0
8.0
10.0
100.0
B
Noise: 0.2
2
Loss value
Neq. Log Likelihood
1
Reconstruction Loss
0
0.01
0.05
0.1
0.2
0.4
0.6
0.8
1.0
2.0
4.0
6.0
8.0
10.0
100.0
B
Nolse:
Neg. Log Likellhood Reconstruction Loss
3
Loss value
2
0.01 0.05
0.1
0.2
0.4
0.6
0.8
4.0
6.0
8.0
10.0
100.0
Figure 7: The position of the transition point depends on the data set. The plots show the tradeoff between reconstruction error and NLL with reconstruction weight 3 (box plots summarize 20 runs per condition). The point at which / becomes sufficiently large (transition point) shifts to lower values with increased Gaussian noise added to the data points.
latent space for 50 epochs with learning rate 0.001 until convergence. Each residual block is made up of a feedforward network with one hidden layer of width 256. For each value of models are trained.
To visualize the dimension of the data that is captured by the model, we project samples from the data distribution to the (ID) latent space and color the data points using the respective latent code as color value. Figure 3 illustrates that low reconstruction weight values result in learning the dimension with the lowest entropy (noise) and higher values required to learn the manifold that spans the sinusoid.
Additionally, repeat the procedure with higher noise levels (a 0.2, 0.3). We observe that the point at which the model transitions from learning the noise to representing the manifold is not fixed, but depends features of the data set such the noise (fig. 7
28

---

Published as a conference paper at ICLR 2024
101
10-
amg.lorvauu
1000
number aininuuaar
6-
reconsti weght
10
0.1
-..
3
-0
300
300
10
12
15
10
10f 10:
LOP
number sngu
Figure 8: (Left) Singular value spectra for varying reconstruction weight. (Right) Number of singular values greater or equal to one as a function reconstruction weight. The error bars show the span of the intersection of the shaded region with the line y 1 in the left hand plot, rounded down to the nearest integer. For each trained model we generate 1024 samples per condition, compute the singular value spectra and average over all samples regardless of condition. The mean spectra and their standard deviation are evaluated across five trained models reconstruction weight.
21 2 8
2
P 8
8
8
8 8
P
6
R
8
7
77
7
7
7
666669666
6
4 666
6
6
5S5 555555555
P 5 5 5
556556555
5
5 6
4
#
4
3 3 33333
3 3 33333333 33
3
3333333 333
1
3 2
3
1 222
2 2
2222222
62
//1
3
0 0000000000 0
0
0 0000000000 0
Figure 9: Conditional MNIST samples for varying reconstruction weight B. For each value of B E [0.1, 1, 10, 100] (from left right), and each condition (rows) we generate samples (columns) temperature T
E.1.2 CONDITIONAL MNIST
To measure the structure of the conditional MNIST dataset learned by the generating model g. compute the full decoder Jacobian matrix by calculating Jacobian-vector products (one per column of the Jacobian). We compute the singular values E diag($1..... sa) of J, which roughly speaking indicate the stretching or shrinking of the latent manifold by Hence, the number of non-vanishing singular values suggest the dimension of the data manifold and the of the log singular values is equal to the change in entropy between the latent space and data space, where a higher entropy indicates that more of the data manifold is spanned by the decoder. We can see this from the formula
[
H H (pz)= - Epz
log det (x g'()))
(124)
with H the differential entropy, and noting that 2 tr log E = log (et gz))))
We train multiple FIF models on conditional MNIST with reconstruction weights ranging from 0.1 to 100, and evaluate their singular value spectra. The models are trained for 400 epochs, which sufficient for convergence. The architecture used is the same as table except that has four times many channels in each convolutional layer.
In fig. iit clear that a higher reconstruction weight gives rise to higher number of non- non-vanisee vanishiig singular values. Hence, the reconstruction weight contributes towards learning structure of the true data manifold. This observed additional structure for higher reconstruction weights reflected an increasing diversity of samples (see fig. 19). Nevertheless, for high reconstruction weights we note the trade-off between sample diversity and properly learned latent distributions, which might result in out of distribution samples.
E.2 TABULAR DATA
We compare to the tabular data experiments in Caterini et al. (202I). using the same datasets and data splits, well as the same latent space dimensions. We train models with roughly the same number
29

---

Published as a conference paper at ICLR 2024
Table Dataset-dependent hyperparameters and average total runtime for tabular data experiments. We compare our runtime against the published rectangular flow (Caterini et al., 2021) runtimes (RNFs-ML (K = model), well as rerunning their code on our hardware.
Hyperparameter
POWER
GAS
HEPMASS
MINIBOONE
Latent dimension
10
21
Training epochs
15
30
85
875
Model
Training time (minutes)
FIF (ours)
38
39
41
49
Rectangular flow (published)
113
75
138
34
Rectangular flow (our hardware)
147
86
249
75
Training time speedup (same hardware)
3.9 X
2.2 x
6.1 x
1.5 x
Table 6: Ablation study on the effeet of each component of our proposed improvement to rectangular flows (RF). By NLL estimator we denote how the loss in equation is approximated. For this experiment we used our reimplementation of RF.
Hyperparameters & Model
NLL estimator
POWER
GAS
HEPMASS
MINIBOONE
FIF & free-form net
off manifold (eq.
15
0.041 + 0.007
0.281 + 0.031
0.541 0.034
0.598 + 0.024
FIF free-form net FIF coupling flow
manifold (eq
19.54 +: 20.81
7.48 5.40
29.03
5.42
77.23 16.55
off manifold (eq.
0.11 1.0160006
0.45 0.09
1.30
1,55 1 0.04
& coupling flow FIF coupling flow coupling flow
off manifold (eq.
0.98 0.69
6.16 4.20
2.02 + 0.74
1,80 + 0.10
manifold (eq.
3.71 2.19
0.40 0.22
0.71 0.05
3.13 + 0.42
manifold (eq.
0.33 0.22
0.33 0.17
0.82 + 0.07
1.84 + 0.11
of parameters. Our main architectural difference that we use an unconstrained autoencoder rather than an injective flow. Our encoder consists of two parts: i) feed-forward network with two hidden layers of dimension 256 and ReLU activations (no normalization layers), which maps from the input dimension the latent dimension ii) ResNet with of dimension 256 and ReLU activations. The ResNet has input and output dimension equal to the latent space dimension. The decoder is the inverse: i) an identical ResNet the encoder (but with separate parameters) followed by ii) feed-forward network with two hidden layers of dimension 256 mapping from the latent space to the data space dimension.
blocks, each with two hidden layers
We use batch size of 512, add isotropic Gaussian noise with standard deviation 0.01, use K = Hutchinson samples and reconstruction weight optimizer with the onecycle LR scheduler with LR of 10-4 (except for HEPMASS which has LR of x 10-4) and weight decay of 10-4, The number epochs was chosen such that all experiments had approximately the same number of training iterations. We ran the model 5 times per dataset. The dataset-dependent parameters and average training times are given in table
for all experiments. We use the Adam
We compare training times against the published rectangular flow training times for their RNFs- ML (K = 1) model, as well as rerunning their code on our hardware (a single RTX 2070 card). We find comparable FID-like scores
rerun (except GAS where could not reproduce the score, see section 5), but our hardware is slower, with runs consistently taking at least 15% longer and more than twice as long MINIBOONE. We find that our model runs in half the time less of the rectangular flow on the same hardware, except for MINIBOONE (about 2/3 the time).
E.3 COMPARISON TO EXISTING INJECTIVE FLOWS
We compare against Trumpets Kothari et al. 202 and Denoising Normalizing Flows (DNF) Hor- vat Pfister (021). as they are the best best-performing injective flows to our knowledge, and report performance on CelebA in table Note that Trumpets default to d 192, DNF to d 512, whereas we are able to reduce the bottleneck dimension to d = 64 (consistent with the Pythae benchmark in appendix E.4
30

---

Published as a conference paper at ICLR 2024
Table 7: Reconstruction losses of FIF with free-form architecture on the POWER, GAS HEP- MASS and MINIBOONE datasets. The reconstruction error is always much higher for on-manifold training compared to off-manifold, demonstrating the instability caused by on-manifold NLL eval- uation in free-form networks. Note: the large standard deviations in on-manifold runs are typically the result of a single large outlier. We remove the largest outlier where applicable ("On Manifold (outliers removed)" row).
POWER
GAS
HEPMASS
MINIBOONE
On manifold
237 = 498
5835 + 13006
119 = 34
300 + 160
On manifold (outliers removed)
14 +: 27
18 +: 31
119 + 34
229 + 23
Off manifold
0.072 + 0.002
0.188 + 0.012
2.569 = 0.098
1.077 + 0.011
Both models differ in the recommended wall clock time, and therefore fix the wall clock time available to each model to five hours on a single NVIDIA A40, Trumpets train the manifold and the distribution it in two sequential steps. To accommodate both steps in the reduced training time, we vary the fraction of the five hours spent in training manifold and distribution and report the best FID among the variants tried. We vary number of manifold epochs as manifold 2, 5, 10, with 10 performing best.
Our free-form injective flows (FIF) are not restricted in their architecture, and we choose an off- the-shelf convolutional autoencoder, followed by total of four fully-connected ResNet blocks, see table 11 The fully-connected blocks are important, as can be seen when comparing to the architecture used in the Pythae benchmark (see appendix E.4). We note that the Pythae benchmark could benefit from a modified architecture, but leave this modification open for future work.
For Trumpets and DNF, we point to the training details provided by the respective authors. For FIF, we choose these training hyperparameters: We train with the Adam optimizer with a LR of 10-3 and weight decay of 10 -4, linearly increase 3 from 20 initialization to 40 at the end of training, a single Hutchinson sample K = and a student-t distribution on the latent space. We set the batch size 256.
We conclude that the Pythae benchmark could benefit from an optimized architecture, as this change probably also improves the other methods. From the data at hand, further conclude that the full potential of FIF has not yet been exploited, and that easy gains can be made by improving the architecture and other hyperparameters.
E,4 PYTHAE BENCHMARK ON GENERATIVE AUTOENCODERS
We compare our method to existing generative autoencoder paradigms using the benchmark from Chadebec et al. (2022). We use the provided open-source pipeline and follow the training setup de- scribed by the authors. For MNIST and CIFAR 10 this means training for 100 epochs with the Adam optimizer at a starting LR of 10-4, reserving the last 10k images of the training sets as validation sets. CelebA trains for 50 epochs with a starting LR of 10-3, All experiments are performed with a batch size of 100 and LR is reduced by half when the loss plateaus for 10 epochs. In accordance with the original benchmark we pick 10 sets of fig. TOp and use the model which achieves the best FID on the validation set as the final model. In table we report the FID and IS of this model on the test set. To complement the metrics, we show samples from all models in fig. from the FID comparison, as the model trains more than double the time required for FIF and goes beyond fitting a transformation of the training data standard normal distribution.
compute their validation FID (see
demonstrating convincing quality. We exclude the VAEGAN
As described in section 5 we use the architectures from Chadebec et al, 2022) for the benchmark, which we replicate in tables and 10
E.5 COMPUTE AND DEPENDENCIES
We used approximately 150 GPU hours for computing the Pythae benchmark, and an additional 800 GPU hours for model exploration and testing. The majority of the experiments were performed on an internal cluster of A 1000s. The majority of compute time was spent image datasets.
31

---

Published as a conference paper at ICLR 2024
ConvNet MNIST
ResNet MNIST
27.5
27.5
GMM, K-1
N. K=2
25.0
25.0
GMM, K=2
validation loss
22.5
22.5
K=1
20.0
GMM, K-1
20.0
N. Km2
17.5
GMM, K-2
17.5
10.0
10.0
5.0
7.5
10.0
12.5
17.5 reconstruutionn loss
15.0
20.0
22.5
25.0
5.0
1.5
10.0
17.5 recorsttuuctoo loss
12.5
15.0
20.0
22.5 25.0
ConvNet CIFAR10
ResNet CIFARIO
160
N. K=1
.V. K=1
GMM, K=1
GMM, K-1
150
K-2
150
K=2
GMM, K-2
GMM, K-2
140
140
120
120
110
110
100
100
15.0 B reconstruction loss
reconstruction
ConvNet CELEBA
ResNet CELEBA
M.
GMM, K-1
GMN, K-1
90
A, K-2
90
N. K-2
GMM,
GMM,
validation
70
70
50
50-
5.0
7.5
10.0
12.5
15.0
17.5
20.0
22.5
25.0
5.0
1.5
10.0
12.5
15.0
17.5
20.0
22.5
25.0
reconstruution
recorstruction less
Figure 10: Validation FID of our 10 benchmark model setups the different datasets and architec- tures used by Pythae. We report results for reconstruction weights 3 = 3, 10, 15, 20, 25 and number of Hutchinson samples K 1 in the solid lines and K in the dashed lines. We show the performance of standard normal sampling (A) and a Gaussian mixture model (GMM). We see that GMM sampling always improves performance. In contrast, the reconstruction weight and number of Hutchinson samples have no noticeable effect on performance, except that increasing 3 improves performance on ConvNet MNIST and ConvNet CIFAR10, and decreases performance on ResNet MNIST (normal sampling only).
We build our code upon the following python libraries: PyTorch (Paszke et al.] 2019). Py PyTorch
Lightning (Falcon & The PyTorch Lightning team 2019). Tensorflow (Abadi et al. (215)) for FID score evaluation, Numpy (Harris et al.] 2020). Matplotlib (Hunter) (007) for plotting and Pandas (McKinney 2010 The pandas development team 2020) for data evaluation.
32

---

Published as a conference paper at ICLR 2024
Table 8: Table taken from Chadebec et al. 2022) with our results added at the top. We report Inception Score (IS) and Fréchet Inception Distance (FID) computed with 10k samples the test set. The best model per dataset and sampler is highlighted in bold, the second best underlined.
ConvNet
ResNet CIFARIO
Sampler
FID
IS
FID
1S
FID
IS
FID
IS
FID
IS
FID
IS
FIF (oun))
GMM
.V
23.8
11.0
2.2
2.2
121.0
90.6
3.0
56.9
47.3
2.1
1.9
19.5
11.7
2.1
2.1
132.6
119.2
2.9
62.3
55.0
1.7
NV
28.5
2.1
241.0
2.2
54.8
1.9
31.3
2.0
181.7
66.6
1.6
235.9
179.7
VAMP
VAMP
61.2
2.0
329.0
56.0
1.9
34.5
2.1
181.9
67.2
29.0
2.1
245.3
2.1
55.7
1.9
32.4
2.0
191.2
67.6
1.6
IWAE
GMM
28.4
2.1
241.2
2.1
52.7
1.9
34.4
2.1
188.
2.4
64.1
1.7
185.5
67.1
GMM
28.4
2.1
237.0
53.3
1.9
33.1
2.1
183.1
62.8
VAE-IAF
AV
27.5
2.1
236.0
2.2
55.4
1.9
30.6
2.0
183.6
2.5
66.2
1.6
GMM
2.1
235.4
62.7
2.1
115.4
2.0
65.9
GMM
9.2
2:2
92.2
51.7
1.9
11.4
2.1
112.6
59.3
CO 1)
96.5
2.3
219.4
3.6
130.8
1.6
109.4
2.7
209.
3.2
110.6
1.5
Dis D-VAE
GMM
192.7
2.3
300.8
1.8
94.0
1.5
98.7
2.1
202.3
2.6
83.2
1,6
2-s sampler
MAP sampler
19188
2.2
300.6
1.8
94.3
1.5
98.3
2.1
200.6
2.6
82.6
B-TC VAE
GMM
N
21.3
11.6
2.2
2.1
116.6
89.3
2.8
51.8
55.7
1.8
1.9
20.7
13.3
2.1
2.0
125.8
106.5
3.4
65.9
59.3
1.6
1.7
VActoVVAE
2.1
236.5
66.4
GMM
26.9
2.1
234.0
2.2
52.4
2.0
32.7
2.1
184
2.5
63.3
1.7
InfoVAE- RBF
GMM
N
26.7
27.5
2.1
2.1
230.4
235.2
2.1
2.2
55.5
52.7
1.9
31.1
32.3
2.1
2.0
182.8
179.5
2.5
66.5
62.8
1.6
1.7
InfoVAE IMQ
GMM
27.7
2.1
2.1
231.9
233.8
2.2
53.7
56.7
1.9
32.8
2.1
180.7
2.6
62.3
1.7
AV
16.8
2.2
139.9
2.6
59.9
19.1
2.1
164.9
64.8
1.7
92.1
MSSSIM-VAE
GMM
27.2
26.7
2.2
2.2
279.7
279.9
1.7
121.3
124.3
1.3
1.3
28.8
28.0
2.1
2.1
253.1
254.2
1.7
119.0
119.2
1.3
VAEGAN
2.2
199.5
2.2
39.7
72.8
2.2
198.7
122.8
(Bol compured)
2.2
188.
84.3
AP
GMM
26.7
9.3
2.2
2.1
201.3
97.3
3.6
327.7
55.4
2.0
1.0
221.8
11.0
2.1
1.3
210.1
120.7
3.4
275.0
57.4
1.8
21.2
2.2
175.1
332.6
21.2
2.1
170.2
69.4
GMM
2.2
97.1
55.0
2.0
11.2
2.1
120.3
58.3
WAE IMQ
GMM
.V
18.9
2.2
2.2
164.4
96.5
2.2
$1.7
64.6
2.0
1.7
20.3
2.1
2.1
119.0
150.7
2.5
67.1
57.7
1.6
VQVAE
GMM
28.2
9.1
2.0
2.2
152.2
95.2
2.0
3.7
306.9
$1.6
2.0
1.0
170.7
10.7
2.1
1.6
130.1
195.7
3.4
1.9
140.3
$7.9
1.8
GMM
9.1
2.2
85.3
3.9
55.2
1.9
11.5
2.1
122.5
58.3
1.8
RAB GP
AV
27.1
2.1
196.8
2.1
86.1
2.4
61.5
2.2
229
2,0
201.9
3.1
123.3
Table 9: ConyNet, neural network architecture used for the convolutional networks, adapted from Chadebec et al. (2022).
MNIST
CIFAR
CELERA
Encuder
28, 28)
(3, 32)
Layer
Convi128, 4. BN, ReLU
Conv(128, 4, 2), BN, ReLU
Conv(128, 4, BN, ReLU
Layer 2
Conv(256, 4, 2), BN, ReLU
Conv(256, 2), BN, ReLU
Conv(266, 4, 2), BN, ReLU
Layer
Convi512, 4, BN, ReL.U
Conv(512, 4, BN, ReLU
Conv(512, 2), BN, ReL.U
Layer
Canv(1024, 2), BN, ReLU
Conv(1024, 4, 2). BN, ReLU
Convt 1024, 4, 2), BN, ReLU
Layer
Lincar(1024, lateot dim)
Lincar(4096, lateot.dim)*
Lincar[16844, dim)*
Decoder
Layer I
Linsar(latent.c dim, 16384) Conv 1(51, 2). BN. ReLU ConvT(256, 2), BN, ReLU
Liocar(latent dim, 65536)
Lincar(latent dim, 65536)
Layer 2
ConVI(512, 4, 2), BN, ReLU ConvT((25, BN, ReLU
ConvT[512, 5, BN, ReLU ConvTi256, 5, BN, ReLU ConVT1128, 5, BN, ReLU T13, 1), Sigmoid
Layer
Layer 4
Conv(1, 3, 2). Sigmoid
Conv(3, 4, 1), Sigmoid
Layer
#Parameters
17.2M
39.4M
33.5M
33

---

Published as a conference paper at ICLR 2024
Table 10: ResNet, neural network architecture used for the residual networks, adapted from Chade-
bec et al. 2022).
MNIST
CIFARIO
CELEBA
Encoder
28, 28)
(3. 32, 32)
64, 64)
Layer
Conv164, 2,
Conv(64, 4, 2)
Conv(64, 2,
Layer
Conv(128,
Conv(128,
Conv(128, 4,
Layer 3
Conv(128, 3, 2)
Conv( 128, 3,
Conv(128 3,
Layer
ResBlock
ResBlock*
Conv(128, 3,
Layer
ResBlock
ResBlock*
ResBlock
Layer 6
(204, dimes
Linear(8192, dale))
ResBlock
Layer
Linear latent_dim))
Decoder
Layer
Lineartlatent diinn 2048)
Lincartlaten.doe 8192)
Linear(attent.imn 2048)
Layer
1112, 2) ResBlock* ResBlock*, ReLU ConvT(64, RcLU ConvT(L 3, Sigmoid
ResBlock
Conv (112,, 3, RexBlack* ResBlock*
Layer
ResBlock*
Layer
ConvT164, 4,
Layer
ConvT(3, 2), Sigmoid
Co28 5, 2).). Sigmoid
Layer 6
Conv 1(64, 5, 2). Sigmoid
Layer
T(3, 2). Sigmoid
#Parameters
0.73M
4.8M
1.6M
"The ResBlocks composed of one C3nv(32, 3, Jollowed Conv 1, with ReLU.
Table 11: FIF-ConvNet, neural network architecture used for comparison to Trumpet and Denoising Normalizing Flow.
MNIST
CELEBA
Encoder
(1, 28, 28)
64, 64)
Layer
Conv(32, 4, 2). BN, ReLU
Conv(128. 2). BN. ReLU
Layer
Conv164, BN, ReLU
Conv(256. 2)2).
Layer 3
Conv (128, 4, 2), BN, ReLU Conv(256,
Conv(512, 2), BN, ReLL
Layer
Conv(T024, 4,
Layer 5
Linear256, lotent.dim))* 4xResBlock(512)
Lineart 16384, latent.dim)" 4xResBlock(256)
Layer 6-9
Decoder
Layer
4tResBlock/512) Lineartlatent. dim, 4096)
4sRe>Block(256) Lincar(lateot dim, 65536)
Layer 5
Layer
ConVTY256, 3, 2), BN, ReLU
ConvT(512, 5, 2), BN, ReLU
Layer
Co1/T(188 3, BN, ReLU
Conv I(256, 5, 2), BN, ReLU
L
Com Com/T(64, 2), BN, ReL.U
ConvT(12 5, 2), BN, ReL.!
Layer
Conv( 2), Sigmoid
ConTT 5, Sigmoid
3.3M
34.3M
*The RosBlockatinner dim) are composed of Lincartlatent dim, inner dim). Linear(inner dim, inner dim), SILU, Lineartinner dim, dim) skip connection.
34

---

Published as a conference paper at ICLR 2024
CELEBA - N
CELEBA GMM
VAE
IWAE
VAE-lin-NE
VAE-IAF
B-VAE
8-TC-VAE
Factor-Vas VAE
InfoVAE- IMQ
Info RBF
AAE
MSSSIM-VAE
VAEGAN
AE
WAF-IMQ
WAE-RBE
VQVAE
RAE-L2
RAE-GP
FIF (ours)
Figure 11: Uncurated samples from the CelebA ConvNet experiments in the PythAE benchmark.
Our model is shown the bottom, samples from the other models have been taken from Chadebec et al. 2022).
35

---

Published as a conference paper at ICLR 2024
R = 0.1
R = 1.0
R = 10.0
R = 100.0
R = 1000.0
2
0
-2
-2
0
-2
-2
0
Figure 12: Possible learnt manifolds of varying curvature 1/R for data supported on a subspace,
where d = 1 and D 2.
B 1
B
4
-2
MSE
NL
-4
Total
10-1
102
10-1
102
10-1
102
10-1
102
Radius R
Radius R
Radius R
Radius R
Figure 13: Weighting the reconstruction loss higher does not lead to stable training. The plots show different reconstruction weights B. all settings, highly curved manifolds (i.e. low radius) achieve the lowest loss.
F DETAILS ON PATHOLOGY INDUCED BY CURVATURE
As described in section gradients from the on- on-manifold loss in eq. 10 cause the learned man- ifold to increase curvature. This is visualized in the main text in fig. where the left plot shows that this loss leads to iver-incraasing curvature. The reason that the entropy of data projected to curved manifold smaller than the entropy of data projected to flat manifold.
Here, we provide intuition for why this happens for synthetic data where d is known and the data could in principle be perfectly reconstructed. Figure 12 depicts projections of the data possible model manifolds of varying curvature x. We parameterize the curvature by varying the radius π 1/K. One observe that for with increasing curvature (i.e. decreasing radius), the data is projected to an increasingly small region. Correspondingly, the entropy H (Poatal (i)) of the projected data becomes arbitrarily negative (just like Gaussian with low standard deviation has arbitrarily negative entropy), lowering the achievable negative log- ligeliiel
Adding reconstruction loss alone does not fix this pathology, which we illustrate in fig. The reconstruction loss saturates for small radii, but the best achievable negative log- ligeeiiodoood (i.e. the entropy of data) continues to decrease with the radius. Thus, even when increasing 3, the minimal possible value of the total loss is still achieved by a spuriously curved manifold.
36


Reviews:
[{"review_id": "uwuJeZSPLr", "reviewer": "Reviewer_KrZV", "summary": "This paper proposes a new technique for training flow-based architectures with manifold structure, entitled the \"Free-form Injective Flow\". This technique is derived by loosening the architectural constraints on previously-proposed injective flows; in particular, the proposed autoencoder is completely unconstrained besides using a pre-specified latent dimensionality $d$. Issues with the loss used for previous injective flows are also identified, and this paper derives a novel loss function to address those issues. This loss is computationally tractable while maintaining stability. Experiments are performed to compare with previous injective flow techniques, and other types of autoencoders.", "questions": "1. Why is the training time speedup inconsistent in Table 1? It doesn't seem to scale with dimension in any predictable way.\n2. What is the definition of entropy as in e.g. (13) for a distribution that is supported on a manifold?\n\n**POST-REBUTTAL**\n\nI'll be upgrading my score after discussion with the authors.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I'll enumerate the strengths below for ease of reference in discussion. These are not listed in order of importance.\n\n1. The paper is generally very well-written and comes up quite polished. I'll outline below:\n    - The introduction is very clean. The motivation for the method is clearly laid out.\n    - The paper is well-situated amongst the related work.\n    - The background is easy to digest.\n    - The path to the final model in Section 4 is laid out well.\n    - The appendix is quite thorough\n2. The autograd / linear-algebra-type derivations were quite well-done. I have personally always appreciated works that don't just blindly apply basic automatic differentiation techniques and instead think deeper about the problem and the requisite gradient estimators.\n3. I like the use of figures. Figure 2 in particular is quite nice for explaining how the technique improves on previous injective flow-based methods. Figure 3 also clearly demonstrates the trade-off between reconstruction and log-likelihood.\n4. I feel well-convinced that this technique is clearly better than previous injective flow techniques, both in terms of representation power and computational tractability.", "weaknesses": "I'll write weaknesses in a list as well. Again this list is not ordered in terms of importance.\n\n1. In the end, this paper could be summarized as simply training autoencoders with a different training loss, with the loss motivated by previous work in injective flows. The novelty and significance of this particular choice of loss over other types of autoencoder losses is not completely clear for a couple of reasons: (i) Table 3 is not convincing, as the best results are still produced by other autoencoders, and (ii) more modern autoencoder architectures are not compared against. Furthermore, this paper does not necessarily maintain all of the benefits of injective flows - mainly, we do not get exact inverses on the projections to the learned manifold. To summarize, I think *some* degree of discussion is warranted on the benefits of using this approach over other generative autoencoders, as the benefits over other injective flows are comparatively very well-documented here. \n2. This paper is missing a dedicated limitations section. This is partially covered by the conclusion, but not completely, and would show some more perspective from the authors considering the weaknesses laid out here.\n3. In section 5.3, it is suggested that Inception Scores are a reliable measure of diversity, although I don't know if that's actually a modern viewpoint. Furthermore, the Inception Scores generally seem just in-line with other methods, or worse at times. I am also confused about why the best Inception Scores are not bolded in Table 3.\n4. It seems like Section 5.2 and 5.3 are out-of-order on how things are defined. For example, the FID acronym is both cited and defined in 5.3, yet referred to in 5.2. Table 2 also requires more of a description -- including what \"IS\" is, and what the two samplers are -- some of which is contained within the caption of Table 3. I would just suggest making the requisite definitions in Section 5.2 first and then using acronyms or reduced descriptions in Section 5.3 as appropriate.\n5. It is discussed twice that traces are performed in the order $f' g'$, and that details are in the Appendix -- however, there is certainly space in the paper to provide a bit more discussion on that. \n6. The paragraph on Page 6 starting with \"Unfortunately\" is not sufficiently convincing: I don't think Fig 2 proves that reconstruction error is insufficient, as you could imagine that reconstruction becomes increasingly more difficult if the entropy becomes negative infinity which should therefore regularize the solution on the left to some extent so that it does not fall on the degenerate, negative-infinite entropy solution.\n7. It is suggested that rectangular flows require $O(d)$ `vjp`s / `jvp`s for convergence, but practically conjugate gradient has exponential convergence and thus much fewer iterations suffice.\n8. It is stated that the surrogate loss is only accurate if $f$ and $g$ are optimal with respect to reconstruction, and then assumed that this is indeed fulfilled by optimizing the reconstruction loss and by the fact that training is stable. However, I don't think this is fully proven:\n    - The reconstruction error is not checked in the paper\n    - One of the other changes to the loss function may be responsible for the training stability\n    - Trade-offs between optimizing the reconstruction error and reconstruction the likelihood contribution may prevent the reconstruction loss from being fully optimized (cf. Figure 3)\n9. There are no samples provided for the generative methods, which suggests that the generation quality may not actually be that good. FID has recently come under more scrutiny as an evaluation metric and so supporting the FID numbers with actual generated samples would be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4jsCWlokYT", "reviewer": "Reviewer_vmge", "summary": "A new Injective Flow named free-form injective ﬂow (FIF) is proposed. FIF is developed based on the rectangular flow (Caterini et al., 2021) and change of variables across dimensions. Different from the rectangular flow, FIF leverages the auto-coding architecture to approximately but efficiently calculate the gradient of maximum likelihood surragate wrt the parameters. The authors also identify pathological behavior in the naive application of maximum likelihood training and propose a fix.", "questions": "Without architectural constraints, how to guarantee that $det [g′(z)^T g′(z)] > 0$?\n\nIn the paragraph before Eq. (12), what are the assumptions underlying $f(\\hat x)=f(x)$? Also, why does Eq. (12) hold true? If $p_{data}(x)=\\hat p_{data}(\\hat x)$, then the right-hand side of Eq. (14) is fixed, right?\n\nAfter adopting the modification in Eq. (16), the final objective in Eq. (18) (or its first two terms) ultimately is not identical to the negative maximum likelihood, right? If so, what are the differences?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The presented techniques are original, with contributions on removing limitations from prior methods.\n\nThe presented techniques are interesting and potentially valuable.", "weaknesses": "The clarity should be significantly improved. For example, many important derivations should be moved to the main manuscript, and important assumptions should be highlighted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fCiovfW2EX", "reviewer": "Reviewer_SnSS", "summary": "This papers builds upon rectangular flows, a method for end-to-end training of injective normalizing flows. Three modifications are proposed: (a) not restricting the architecture with normalizing flows, since the reconstruction error encourages injectivity in it of itself, (b) a more efficient gradient estimator, thus addressing a main limitation of rectangular flows, and (c) a further modification to the gradient estimator, which changes the gradient itself and improves numerical stability.\n\nOverall, the paper is well-written and I believe it makes a significant methodological contribution to the area of injective normalizing flows. That being said, I also believe that ablations are missing to properly identify the sources of empirical improvement that the authors observed.\n\n------------------------------------------------------------------------------\n11/22 UPDATE\n------------------------------------------------------------------------------\n\nThe authors have adequately addressed the points I raised in my original review and I am thus increasing my score.\n\n------------------------------------------------------------------------------\n12/04 UPDATE\n------------------------------------------------------------------------------\n\nI did want to raise an additional point after having discussed the pathology described in section 4.2 with other reviewers. In particular, the bound in eq 14 is a lower bound on the loss, not an upper bound. Thus, if the entropy becomes arbitrarily negative, it does not automatically imply that the loss can become equally negative, as seems to be implied in the discussion following the equation (\"the loss can continue to decrease without bound by reducing the entropy of the projected data\").\n\nWhile I do not believe the point made by the authors is wrong, I do believe that it requires additional arguments: for $\\mathcal{L}_{NLL}$ to also become arbitrarily negative along with the entropy of projected data, the KL divergence in eq. 13 needs to not grow faster than the entropy. I still believe this can happen if the encoder and decoder have unrestricted architectures and collapse to point masses (since then the KL between projected data and the model would be controlled), but I am actually not sure if this can happen when the architectures are constrained to be injective. I also believe this might explain the results from the ablations, where the updated loss did not significantly improve results when using injective architectures.\n\nIf the paper is accepted, I would ask the authors to please flesh out the discussion around this, as I agree with reviewer RfcP that the current exposition in section 4.2 can lead to confusion.", "questions": "7. As you point out in eq 1, injective flows typically have a low-dimensional flow $h$ on the latent space. One could also interpret this architecture as a flexible distribution $p_Z$ on latent space, given by $h$, along with a decoder $w \\circ \\texttt{pad}$; rather than thinking of their composition as the decoder $g$. Throughout the paper you mention making $g$ more expressive, but another interpretation is that you are making $w \\circ \\texttt{pad}$ more expressive, and reducing expressivity on the latent space (instead of a flow, you use a Gaussian or a mixture of Gaussians). Previous research has found benefits of having flexible distributions on latent space (rectangular flows prefer using a flow-based p_Z rather than fixing it as a Gaussian, and other works also recommend using flexible distributions on latent space, e.g. [1, 2]), is there a reason why you do not use more flexible $p_Z$?\n\n[1] Diagnosing and Enhancing VAE Models, Dai & Wipf, ICLR 2019\n\n[2] Diagnosing and Fixing Manifold Overfitting in Deep Generative Models, Loaiza-Ganem et al., 2022", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper has several strengths:\n\n1. It is well written and easy to follow, and I think the authors did a good job of deciding which material to include in the main manuscript and which details to include in the appendix.\n\n2. It is well motivated, as I agree with the authors that the current injective flow literature uses overly restrictive architectures and/or computationally intensive training procedures.\n\n3. The simple observation that, when the encoder $f$ is the left inverse of the decoder $g$, allows to write Jacobians of $g$ as Jacobians of $f$ is elegant, and does result in clear computational gains.\n\n4. Empirical results are good, showing that the proposed method outperforms other injective flows and generative autoencoders.", "weaknesses": "5. In my view, the main weakness of the paper is the lack of ablations. As mentioned, the paper proposes 3 improvements over rectangular flows, and it is unclear how much each of these contributes to the empirical performance of the proposed method. I think table 1 provides a perfect test bed to carry out these ablations: results using the same architecture as rectangular flows should be added to the table, both (a) using the gradient estimator from eq 10, and (b) that from eq 16. Ideally, using eq 10 (and the same architecture as rectangular flows) would simply show a speed up and the same performance compared to rectangular flows, whereas using eq 16 should improve performance but not match the results of FIF with a fully flexible architecture. I would see this as strong empirical evidence backing up the claims in the paper. I will increase my score if these ablations are included.\n\n6. While the authors include a discussion as why $x$ should be used (eq. 16) instead of $\\hat{x}$ (eq. 10), I think there are several relevant points missing from the discussion: (a) why does the pathological behaviour described by the authors not happen in rectangular flows? Is it because the more restrictive architecture implicitly regularizes the curvature? Or is this actually a hidden issue in rectangular flows as well (the above ablation will obviously also help answer this question)?. (b) Since, when $f$ and $g$ are consistent, $f(x)=f(\\hat{x})$, it seems to me like one can attempt to justify both objectives as attempting to maximize log-likelihood subject to perfect reconstructions. In this view, the problem of using $\\hat{x}$ could be seen as an inappropriate way of enforcing the constraint through a penalty term. Could you further discuss? (c) There is also an additional computational benefit to using $x$ instead of $\\hat{x}$, namely one less forward pass is required through the encoder, which I believe should also be mentioned.\n\nFinally, some minor points:\n\n- In the notation paragraph in sec 3, you write $f^{-1} = g$, which I think should be avoided: when $d<D$, $f$ cannot be an invertible function, since you defined its domain as $\\mathbb{R}^D$ (its restriction to a manifold could of course be injective though, I am not saying there's anything fundamentally wrong here, just nitpicking the notation): I think it'd be better to stick to the language of left inverses.\n\n- Missing period at the end of the injective flows paragraph in sec 3.\n\n- Use \\citep instead of \\citet in the first paragraph of appendix E.3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "trLKEErw3s", "reviewer": "Reviewer_RfcP", "summary": "- The authors propose a new kind of injective flow (a normalizing flow with lower-dimensional latent space than the data space).\n- Unlike some previous attempts, they use unconstrained encoders and decoders, and introduce a new estimator for the gradient of the change-of-variables term in the log likelihood.\n- The authors also discuss a training issue for injective flows that was previously pointed out by Brehmer & Cranmer (2020).\n- They demonstrate the method by training generative models on tabular data and CelebA, reporting metrics of generative quality like FID.", "questions": "1. What's the main use case for this injective flow? In what situations do you expect benefits from the manifold structure of this generative model compared to, say, diffusion models or VAEs?\n2. Is the density (not its gradient) of the model tractable? That would extend use cases substantially.\n3. I don't understand the \"fix\" of the pathological behaviour pointed out in Sec. 4.2. Could you expand the discussion of why it would work? Is it guaranteed to work? Consider the toy problem that Brehmer & Cranmer (2020) use to illustrate the same problem (Fig 4 in the arXiv version, 2003.13913). Here the encoder $f(x)$ is linear, thus $f'(\\hat{x}) = f'(x)$ , and the \"fix\" does not change anything.\n4. Do you have an explanation for why the problem discussed in Sec. 4.2 does not affect the experiments?\n5. In the experimental evaluation, are the models converged? How do the results change if the models (in particular the baselines) are trained for longer? I find it hard to believe that none of the VAE methods are able to produce higher-quality samples than what we see in Fig. 1.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The main idea (the new estimator of the gradient) is wonderfully simple, sensible, and efficient.\n2. The paper is clearly structured and well-written.", "weaknesses": "1. Injective flows are academically interesting, but do not have a very clear use case, especially if they do not have a tractable density (see questions below).\n2. The discussion of joint manifold and likelihood training is not novel (see Brehmer & Cranmer, 2020), which the authors are open about. The proposed solution leaves questions open (see below).\n3. I am not yet convinced by the experimental evaluation. Given the quality of the samples in Figure 1, I am surprised by the claim that the method outperforms various VAE methods (see questions below).\n4. Overall, the paper's contributions are quite thin.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)