Paper: Lifting Architectural Constraints of Injective Flows (ICLR-2024)

Reviewer comment (Novelty):
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).

Author response:
Thank you for your constructive feedback. Let us answer to your questions, and we hope that they also soften the mentioned weaknesses:

1. FIF are fast, perform well and are simple to set up:
    1. FIF vs diffusion model: Diffusion models sample slow, which makes them expensive in inference.
    2. FIF vs VAE: VAEs are limited by (i) their variational family, and (ii) by learning a distribution in the embedding space, which is limiting if the data actually lies on a manifold (cf. manifold overfitting). While (i) can be addressed by more expressive likelihoods/posteriors, this comes with additional hyperparameters. Our model only has the architecture and beta as hyperparameters.
    3. FIF vs GAN: FIF is trained via maximum likelihood, so it guarantees against mode collapse.
    4. FIF vs latent flow/diffusion/…, e.g. Stable Diffusion: FIF trains manifold and distribution jointly, while post-hoc learning the latent distribution is a two-step procedure.

Reviewer comment (Novelty):
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).

Author response:
3. We do not simply repeat the problem mentioned in Brehmer & Cranmer, this seems to be a misunderstanding. Let us give a different perspective on the well-behaved loss in section 4.2. We need to address *two* pathologies, each with a different solution. The first was described by Brehmer & Cranmer and appears already for a linear model, and the second one is new and only applies to nonlinear models. You correctly noted that our fix to the second problem does not help against the first problem. So let us quickly sum up the two pathologies and the fix to each.
    1. **Degenerate subspaces** (identified by Brehmer & Cranmer): If the data is only supported on a low-dimensional subspace, then there is a projection of the data with $-\infty$ log-likelihood. This possible projection can be eliminated by adding a tiny bit of noise to the data. In Figure 3, we demonstrate that a high enough reconstruction weight $\beta$ is sufficient to avoid learning degenerate projections. We also present an analytic result on the linear case in Appendix C.
    2. **Strong curvature** (new): If the learnt manifold curves strongly, then data on the outside of the curve is concentrated by projecting to the manifold (see Figure 2). This decreases the entropy of $p(\hat x)$, lowering the lower bound of the manifold log-likelihood, causing divergence. We fix this by computing the encoder Jacobian off-manifold, which punishes spurious curvature (see discussion in appendix B.2). This is also cheaper because it avoids one forward pass through the model.

Reviewer comment (Evaluation):
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).

Author response:
5. We decided to use the Pythae benchmark because it allows a systematic comparison to a large collection of generative autoencoders. However, it puts a compute limit on training by fixing the architecture and the number of training epochs. The purpose of the benchmark is to fairly compare methods, not to achieve SOTA results. This explains the sample quality, see Figure 11 in their paper for samples from competitor models: https://arxiv.org/pdf/2206.08309.pdf

Reviewer comment (Novelty):
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?

Author response:
1. FIF are fast, perform well and are simple to set up:
    1. FIF vs diffusion model: Diffusion models sample slow, which makes them expensive in inference.
    2. FIF vs VAE: VAEs are limited by (i) their variational family, and (ii) by learning a distribution in the embedding space, which is limiting if the data actually lies on a manifold (cf. manifold overfitting). While (i) can be addressed by more expressive likelihoods/posteriors, this comes with additional hyperparameters. Our model only has the architecture and beta as hyperparameters.
    3. FIF vs GAN: FIF is trained via maximum likelihood, so it guarantees against mode collapse.
    4. FIF vs latent flow/diffusion/…, e.g. Stable Diffusion: FIF trains manifold and distribution jointly, while post-hoc learning the latent distribution is a two-step procedure.

Reviewer comment (Theory):
Is the density (not its gradient) of the model tractable? That would extend use cases substantially.

Author response:
2. The density is tractable by computing the Jacobian of the transformation using as many autograd calls as there are latent dimensions. While this may not be computationally feasible for very large latent dimensions, we find that for all problems considered it is still fine. However, there are unbiased estimators for the change of variables which require less autograd evaluations [Chen et al. 2019: Residual Flows for Invertible Generative Modeling].

Reviewer comment (Experiments):
Do you have an explanation for why the problem discussed in Sec. 4.2 does not affect the experiments?

Author response:
> How relevant is this problem, though? I was under the impression that the manifold hypothesis / zero reconstruction error is quite central in motivating this work. In your experiments, what reconstruction errors did you achieve?

In our experiments, this problem seems to be important: If we do not apply our fix, training is unstable and we find substantially worse results. Thanks to the idea of reviewer SnSS, we performed an ablation during this rebuttal that compared the performance with and without the fix. The table below takes the numbers from our Table 1 and adds an additional row which shows the performance on the tabular data of a free-form architecture if the loss is evaluated on-manifold (i.e. without our fix):

| Experiment                               | power         | gas         | hepmass      | miniboone     |
|------------------------------------------|---------------|-------------|--------------|---------------|
| RF: injective flow via RF loss with CG   | 0.08 ± 0.02   | 0.11 ± 0.02 | 0.78 ± 0.19  | 1.00 ± 0.05   |
| FIF: free-form net, off manifold (fixed) | 0.04 ± 0.01   | 0.28 ± 0.03 | 0.54 + 0.03  | 0.60 ± 0.02   |
| free-form net, on manifold (no fix)      | 19.54 ± 20.81 | 7.48 ± 5.40 | 29.03 ± 5.42 | 77.23 ± 16.55 |

We therefore find that this fix was crucial for the successful training of free-form architectures as injective flows.

Reviewer comment (Experiments):
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.

Author response:
5. We decided to use the Pythae benchmark because it allows a systematic comparison to a large collection of generative autoencoders. However, it puts a compute limit on training by fixing the architecture and the number of training epochs. The purpose of the benchmark is to fairly compare methods, not to achieve SOTA results. This explains the sample quality, see Figure 11 in their paper for samples from competitor models: https://arxiv.org/pdf/2206.08309.pdf

Reviewer comment (Experiments):
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.

Author response:
# Ablation results:

We perform the requested ablations to compare the on- and off-manifold variants of our loss using a free-form and an injective flow architecture. The injective flow experiments were performed with both FIF and RF hyperparameters. The results, along with the original results from our submission (top two rows):

| Experiment                               | power         | gas         | hepmass      | miniboone     |
|------------------------------------------|---------------|-------------|--------------|---------------|
| RF: injective flow via RF loss with CG   | 0.08 ± 0.02   | 0.11 ± 0.02 | 0.78 ± 0.19  | 1.00 ± 0.05   |
| FIF: free-form net, off manifold         | 0.04 ± 0.01   | 0.28 ± 0.03 | 0.54 + 0.03  | 0.60 ± 0.02   |
|------------------------------------------|---------------|-------------|--------------|---------------|
| free-form net, on manifold               | 19.54 ± 20.81 | 7.48 ± 5.40 | 29.03 ± 5.42 | 77.23 ± 16.55 |
| injective flow, off manifold, FIF hparams| **0.11 ± 0.06**   | **0.45 ± 0.09** | **1.30 ± 0.14**  | **1.55 ± 0.04**   |
| injective flow, off manifold, RF hparams | 0.98 ± 0.69   | 6.16 ± 4.20 | 2.02 ± 0.74  | 1.80 ± 0.10   |
| injective flow, on manifold, FIF hparams | 3.71 ± 2.19   | 0.40 ± 0.22 | **0.71 ± 0.05**  | 3.13 ± 0.42   |
| injective flow, on manifold, RF hparams  | **0.33 ± 0.22**   | **0.33 ± 0.17** | 0.82 ± 0.07  | **1.84 ± 0.11**   |

Note: bold entries show the best hyperparameters (out of FIF or RF) for each experiment.

## Conclusion:

The ablation shows that (a) the free-form network is unstable for the on-manifold variant of our loss, and (b) the injective flow is more stable against evaluating our loss on manifold: 
For the injective architecture, on- and off-manifold settings are both stable, indicating that there is a stabilizing inductive bias of the architecture. Note that they do not perform as well as the original networks, but this is probably due to the lack of tuning hyperparameters.

## Details:
- FIF hparams means the same training hyperparameters as in our paper
- RF hparams means the same training hyperparameters (including a warm up in the NLL loss) as detailed in the appendix of the rectangular flows paper, except that a fixed number of epochs (150) was used rather than using early stopping
- Injective flow means the RNVP-based architecture detailed in the appendix of RFs
- Hepmass is trained with LR=1e-4 instead of LR=3e-4 when using FIF hparams
- Injective flow with FIF hparams have 4 instead of 5 runs each

Reviewer comment (Theory):
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.

Author response:
### Why $x$ instead of $\hat x$?

**No pathological behavior in rectangular flows?**

From a theoretic point of view, rectangular flows should also suffer from pathologically strong curvature since their conjugate gradient does evaluate the likelihood term on-manifold.


**Using $\hat x$ is an inappropriate way to enforce consistency**

It is true that if $f$ and $g$ are consistent then $f(x) = f(\hat x)$, but this is not true for the Jacobians: $f’(x) \neq f’(\hat x)$ in general (see fig. 4 and section B.2 in the appendix for a discussion). Hence substituting $x$ with $\hat x$ is not just another way of enforcing consistency, it gives a different numerical result even in consistent models, and helps fight against degenerate high-curvature solutions.

**Computational benefit**

You are correct that there is a computational benefit to using $x$ instead of $\hat x$ (one less forward pass). We will add a mention of this in section 4.2.

Reviewer comment (Experiments):
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\)?

[1] Diagnosing and Enhancing VAE Models, Dai & Wipf, ICLR 2019

[2] Diagnosing and Fixing Manifold Overfitting in Deep Generative Models, Loaiza-Ganem et al., 2022

Author response:
### More flexible p(z)

We do find it beneficial in practice to have a few residual blocks which work with the latent dimension. This can be interpreted as being the same structure of an injective flow, which has a latent space flow. Hence, you can view our model as having an expressive latent distribution.

We hope that this addresses your questions and we look forward to further discussion.

Reviewer comment (Theory):
Without architectural constraints, how to guarantee that $det [g′(z)^T g′(z)] > 0$?

Author response:
### How to guarantee a positive determinant?

The reconstruction loss is the only direct loss on the decoder. Reconstruction loss cannot be minimized unless the decoder is injective. Therefore if the reconstruction loss is well-optimized, the determinant is very likely (although not guaranteed) to be positive.

Reviewer comment (Theory):
In 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?

Author response:
### Assumption underlying $f(\hat x) = f(x)$

By definition, $\hat x = g(f(x))$. In the section you are referring to (4.2) we are assuming that $f$ and $g$ are pseudoinverses for the purpose of the derivation, then later drop that assumption. Under the assumption, $f(\hat x) = f(x)$. Sorry for the confusion, we will revise the text to make this assumption more clear.

### Eq. (12)

To clarify your confusion: it is not the case that $p_\text{data}(x) = \hat p_\text{data}(\hat x)$. We replace $p_\text{data}$ with $\hat p_\text{data}$ in eq. (12) because the quantity in the expectation is invariant to projection, so we can replace the original data distribution with the projected distribution.

Reviewer comment (Theory):
After 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?

Author response:
### Relation to maximum likelihood (ML)

You are right that the final objective is not exactly maximum likelihood. The differences are:
- We introduce a reconstruction loss which ensures that high-entropy directions are encoded in the latent space (they would be ignored with a ML loss only).
- We evaluate the encoder Jacobian off-manifold to ensure that the solution does not become degenerate due to excessive curvature (see section 4.2 “Towards a well-behaved loss”).
Given that the two modifications make joint maximum likelihood and manifold learning tractable, we argue that it is justified to denote this as maximum likelihood training. Note that maximum likelihood training alone would not lead to sensible solutions when using bottleneck architectures, so our modifications are necessary to learn ML and manifold simultaneously.

Reviewer comment (Novelty):
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, 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.

Author response:
To clarify the advantages of our model compared to the VAE and its variants: They are limited by (i) their variational family, and (ii) by learning a distribution in the embedding space, which may be limiting if the data actually lies on a manifold (cf. manifold overfitting: Dai and Wipf, 2019 and Loaiza-Ganem et al. 2022). While (i) can be addressed by more expressive likelihoods/posteriors, this comes with additional hyperparameters. Our model only has the architecture and beta as hyperparameters.

> Your response focuses on a comparison to VAE and its variants - but how about non-variational autoencoders? 

Here are some points comparing our model to non-VAE autoencoders included in the benchmark:
- Vanilla autoencoders: not generative models; our model can generate samples
- Adversarial AE and VAEGAN: require an additional adversarial model and potentially unstable adversarial training and mode collapse; our model requires only encoder and decoder, training is based on maximum likelihood so no mode collapse
- Regularized AE: requires two-step training to learn the latent distribution; our model uses a standard normal latent and is trained in one step
- Wasserstein AEs: Regularize the latent space via MMD and not ML. MMD suffers from the curse of dimensionality as can be seen by its diminishing performance going from MNIST to CIFAR.

Are there any other specific autoencoder models you have in mind for a comparison?

Reviewer comment (Presentation):
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.

Author response:
Adding a limitations section is a good idea which we will add for the camera ready version. We will address: The Pythae benchmark is limited because it does not scale training to competitive image generation, the FID and IS metrics are limited, and we observed that architectural changes can significantly improve results. Also, while the empirical results are good, theoretically showing whether the final loss actually learns the exact distribution on the manifold is still open; in particular when the reconstruction loss is not at optimality. Finally, the reconstruction weight is left open as a hyperparameter to tune for each new data set.

Reviewer comment (Evaluation):
In section 5.3, it is suggested that Inception Scores are a reliable measure of diversity, although it is unclear if that's actually a modern viewpoint. Furthermore, the Inception Scores generally seem just in-line with other methods, or worse at times. There is also confusion about why the best Inception Scores are not bolded in Table 3.

Author response:
We included Inception Scores because they are part of the Pythae benchmark. We included the Inception scores because they indicate diversity. We do not bold them since FID is more suited at comparing sample quality to real data.

> And is this [Inception scores indicate diversity] truly still a valid viewpoint considering the advancements made to attempt to quantify diversity of generated datasets?

We base this viewpoint on a subjective comparison of the generated CelebA images against Inception scores for the models in the Pythae benchmark, see fig. 11 and table 7 in [1]. For example an IS of 1 (the min possible value) corresponds to a single blank image being generated and an IS of 2.8 (RAE-L2) corresponds to very diverse images. We are very interested in metrics with better performance and would greatly appreciate some pointers to relevant literature.

Reviewer comment (Presentation):
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. A suggestion is to make the requisite definitions in Section 5.2 first and then use acronyms or reduced descriptions in Section 5.3 as appropriate.

Author response:
Thank you for this suggestion. We agree that some definitions are out of order and it makes sense to make all the definitions in section 5.2.

Reviewer comment (Theory):
The paragraph on Page 6 starting with "Unfortunately" is not sufficiently convincing. Figure 2 does not necessarily prove that reconstruction error is insufficient, as it is conceivable that reconstruction becomes increasingly more difficult if the entropy becomes negative infinity, which should therefore regularize the solution on the left to some extent such that it does not fall on the degenerate, negative-infinite entropy solution.

Author response:
To elaborate on why reconstruction loss alone is not sufficient: imagine in fig. 2 the limit where the learned manifold has infinitely high curvature. In this case, all data will be projected to the same point, so the reconstruction error will take on some finite value. In contrast, the NLL will be negative infinity. Hence we can see the problem: this solution has an infinitely small loss, for any (finite) $\beta$ we choose. Therefore reconstruction error alone is insufficient.

> Is it worth including a sentence on this [why reconstruction is not enough] in the main paper?

Certainly. We will add this to the discussion in section 4.2.

Reviewer comment (Theory):
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.

Author response:
You are right that the convergence speed is exponential in the number of iterations, but the convergence behavior might be different with dimension. We will weaken the wording in the introduction and point to our experiment, where we find that wall clock time is longer with CG than with our method.

Reviewer comment (Theory):
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, this is not fully proven:
- The reconstruction error is not checked in the paper  
- One of the other changes to the loss function may be responsible for the training stability  
- Trade-offs between optimizing the reconstruction error and reconstruction the likelihood contribution may prevent the reconstruction loss from being fully optimized (cf. Figure 3)

Author response:
We always monitored reconstruction loss on validation data during training and found that it was always low in our best-performing models. For example, Figure 1 shows reconstructions of CelebA validation data. Regarding training, we find that when reconstruction loss is high (due to too low $\beta$) training becomes unstable. We do find that reconstruction loss never goes to zero, which it can’t given the abundance of details in face images.

Reviewer comment (Evaluation):
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, so supporting the FID numbers with actual generated samples would be useful.

Author response:
The aim of the Pythae benchmark is to compare models on a given computational budget, and the evaluation metrics are FID and IS. We agree that there are flaws to these metrics, but they are currently the only feasible way to compare against a large number of competing generative autoencoders. See Figure 1 for samples from our improved architecture. See Figure 11 in the Pythae paper for samples from competitor models: https://arxiv.org/pdf/2206.08309.pdf

> Can more samples be included in the experiments section or the appendix?

We have updated our submission pdf by adding a comparison of samples from all models considered in the Pythae benchmark, see Figure 11 at the end of the appendix.

Reviewer comment (Evaluation):
Why is the training time speedup inconsistent in Table 1? It doesn't seem to scale with dimension in any predictable way.

Author response:
The inconsistency in training time in table 1 is due to the unpredictability of rectangular flow training, which terminates when no further improvement is made. Training time can vary widely between runs. Please see table 5 on page 30 in the appendix for a fuller picture.

Reviewer comment (Theory):
What is the definition of entropy as in e.g. (13) for a distribution that is supported on a manifold?

Author response:
If we restrict integration to the manifold, the entropy has a sensible definition. Suppose that $\hat p$ is normalized on the manifold: $\int \hat p(\hat x) d \hat x = 1$ where integration is only over the manifold. Then the entropy of $\hat p$ is $-\int \hat p(\hat x) \log \hat p(\hat x) d \hat x$, again integrating only over the manifold. Note that the KL divergence in eq. (13) is defined similarly, with integration restricted to the manifold. There is a typo in eq. (13): the expectation should be over $\hat p_\text{data}(\hat x)$, not $p_\text{data}(\hat x)$.