Title: Complexity Matters: Rethinking the Latent Space for Generative Modeling
Venue: NeurIPS 2023
Accepted: yes
Keywords: generative model

latent space

distance between distributions

generative adversarial network

vqgan
Abstract: In generative modeling, numerous successful approaches leverage a low-dimensional latent space, e.g., Stable Diffusion models the latent space induced by an encoder and generates images through a paired decoder. Although the selection of the latent space is empirically pivotal, determining the optimal choice and the process of identifying it remain unclear. In this study, we aim to shed light on this under-explored topic by rethinking the latent space from the perspective of model complexity. Our investigation starts with the classic generative adversarial networks (GANs). Inspired by the GAN training objective, we propose a novel "distance" between the latent and data distributions, whose minimization coincides with that of the generator complexity. The minimizer of this distance is characterized as the optimal data-dependent latent that most effectively capitalizes on the generator's capacity. Then, we consider parameterizing such a latent distribution by an encoder network and propose a two-stage training strategy called Decoupled Autoencoder (DAE), where the encoder is only updated in the first stage with an auxiliary decoder and then frozen in the second stage while the actual decoder is being trained. DAE can improve the latent distribution and as a result, improve the generative performance. Our theoretical analyses are corroborated by comprehensive experiments on various models such as VQGAN and Diffusion Transformer, where our modifications yield significant improvements in sample quality with decreased model complexity.

Complexity Matters: Rethinking the Latent Space for Generative Modeling

Abstract

In generative modeling, numerous successful approaches leverage a lowdimensional latent space, e.g., Stable Diffusion [68] models the latent space induced by an encoder and generates images through a paired decoder. Although the selection of the latent space is empirically pivotal, determining the optimal choice and the process of identifying it remain unclear. In this study, we aim to shed light on this under-explored topic by rethinking the latent space from the perspective of model complexity. Our investigation starts with the classic generative adversarial networks (GANs). Inspired by the GAN training objective, we propose a novel "distance" between the latent and data distributions, whose minimization coincides with that of the generator complexity. The minimizer of this distance is characterized as the optimal data-dependent latent that most effectively capitalizes on the generator's capacity. Then, we consider parameterizing such a latent distribution by an encoder network and propose a two-stage training strategy called Decoupled Autoencoder (DAE), where the encoder is only updated in the first stage with an auxiliary decoder and then frozen in the second stage while the actual decoder is being trained. DAE can improve the latent distribution and as a result, improve the generative performance. Our theoretical analyses are corroborated by comprehensive experiments on various models such as VQGAN [21] and Diffusion Transformer [60], where our modifications yield significant improvements in sample quality with decreased model complexity.

1 Introduction

In the past decade, deep generative models have achieved great success across various domains such as images, audio, videos, and graphs [32, 59, 33, 34]. The advances are epitomized by recent breakthroughs in text-driven image generation [69, 66, 65, 68]. Behind the empirical success are fruitful developments of a wide variety of deep generative models, among which, many can be associated with a latent space, e.g., generative adversarial networks (GANs) [23, 6, 40], Variational Autoencoders (VAEs) [44], and latent diffusion models [79, 68, 60]. It is commonly believed that for structured data such as images, the intrinsic dimension is much lower than the actual dimension [4, 62]. By utilizing a proper low-dimensional latent space, generative modeling can be more efficient [79, 50, 78]. Take text-to-image generation as an example, state-of-the-art models such as DALL-E [65], Parti [88], Stable Diffusion [68] all utilized discrete Autoencoders to alleviate the computation cost by downsampling images, e.g., from 256 × 256 to 32 × 32.

In practice, the choice of the latent plays a critical role. Stable Diffusion [68] employs the Autoencoder from VQGAN [21], which is an improvement over VQVAE [80] by incorporating adversarial training.

The discrete Autoencoder was later substituted to continuous ones regularized with Kullback-Leibler (KL) divergence. ViT-VQGAN [87] further boosted VQGAN by using vision transformers [20] as the backbone and it was adopted by Parti [88]. Despite the methodological advances, our theoretical understanding of the latent space is still limited. Important questions such as what constitutes a good latent space, what is the optimal choice, and how it depends on data remain elusive.

Ideally, the latent distribution P z should preserve the important information about the data distribution P x as much as possible so that the required effort from the model learning is minimal. Finding such a data-dependent latent can be viewed as a self-supervised learning (SSL) task. Can we borrow insights from the vast literature on SSL to understand and improve the latent? Unfortunately, existing methods [57, 14, 25, 28, 29] are designed for discriminative tasks and not suited for generative modeling. To illustrate, consider the classic GANs [24, 76] where the latent distribution is usually pre-defined. When modeling the CIFAR-10 dataset [45], if we substitute the DCGAN's latent [63] from standard Gaussian to the features learned from SimCLR [14] with the same dimension, the Inception Score (IS) drops from 5.68 to 3.93 2 . Therefore, new theoretical insights are in dire need to elucidate the ideal P z and uncover its connection to P x .

In this work, we aim to provide a new understanding of the latent space in generative modeling from the angles of SSL and minimizing model complexity, where we first formalize the problem for GAN models and then generalize it to Autoencoders.

Drawing inspiration from the training objective of GANs, a novel distance between distributions in different dimensions can be defined to measure the closeness between the latent and the data (Section 3.1). As is typically emphasized in learning theory [82, 83, 35, 61], the complexity of the generator is an important factor in this formulation, where a latent is deemed closer to the data if a generator can achieve the same (better) performance with lower (the same) complexity. The latent closest to the data in a GAN-induced distance is characterized as the optimal choice for that GAN, which leads to the simplest map between the latent and the data.

With our formulation of the optimal latent distribution, the immediate question is how to estimate it, which gives rise to a new SSL task for generative modeling. Naturally, we consider utilizing an encoder network to parameterize the latent, which uncovers the popular Autoencoder architecture with the paired generator as a decoder (Section 3.2). In the learning process, we further emphasize the importance of a relatively weak decoder and the trade-off between the informativeness of the latent and the capability of the decoder (Section 4.1). To this end, we propose a 2-stage training scheme called Decoupled Autoencoder (DAE), where the encoder is only updated in the first stage with an auxiliary decoder and then frozen in the second stage while the actual decoder is being trained (Section 4.2). Our theoretical analyses are corroborated by comprehensive experiments on both synthetic data and real image data with various models such as DCGAN [63], VAEGAN [47, 90], VQGAN [21], and Diffusion Transformer (DiT) [60], where our modifications yield significant improvements in image quality with decreased model complexity (Section 6).

2 Preliminary

Notations. For a function f : Ω → R , let ∥ f ∥ p = ( ∫ Ω | f ( x ) | p d x ) 1 /p . For a vector x , ∥ x ∥ p denotes its l p -norm. L p and ℓ p are used to distinguish function norms and vector norms. For two positive sequences { a n } n ∈ N and { b n } n ∈ N , we write a n ≲ b n if there exists a constant C > 0 such that a n ≤ Cb n for all sufficiently large n . We use P x and p x to denote the probability distribution and density of the random variable x . g ◦ f ( x ) := g ( f ( x )) denotes composition of functions.

Generative adversarial networks are a type of implicit generative models that aim to learn transformations g ∈ G from random noises z to the data x . At the population level, its objective function can be expressed in general as

where P z is usually pre-determined to be some simple distribution such as (mixture) Gaussian or uniform, and D h is realized by the adversarially-trained discriminator h and when optimized properly, can give rise to various well-established distances, e.g., Jensen-Shannon divergence [24],

maximum mean discrepancy [48], f -divergence [58], Wasserstein distance [5], etc. Exploiting the flexible architectures of various neural networks as generators [41, 42], GAN enjoys state-of-the-art performance in generating various data types [41, 85, 74, 89] with notable advantage in the sampling speed, especially compared with diffusion probabilistic models (DPMs) [72, 32]. However, GAN models often suffer from mode collapse, failing to keep the data diversity intact. Besides the usual suspect of the min-max adversarial training, the poor choice of latent noise distribution is also to blame [22].

Self-supervised learning. The goal of SSL is to learn a (low-dimensional) feature representation f ( x ) from unlabeled data that is suited for various discriminative downstream tasks, e.g., classification, detection, etc. There are mainly three categories, i.e., pretext task-based [57], contrastive methods such as SimCLR [14], MoCo [28], BYOL [25], SimSiam [15], and generative-based such as Masked Autoencoder (MAE) [29]. The learned feature map f ( x ) defines a distribution over the feature space z ∈ R d z , denoted as P f . Ideally, P f should preserve the important information about P x as much as possible. In this sense, representation learning in general can be seen as preserving certain "distances 3 " between distributions in different dimensions [36, 8].

Distance between distributions in different dimensions. To quantitatively measure the closeness, various similarity measures between P f and P x are defined. Since d z < d , typical distances between distributions will not work due to the different metric spaces z and x reside in. Fortunately, there are two existing tools we can turn to. The first is Gromov-Wasserstein distance [56, 71], which circumvents the metric spaces mismatch by comparing pairwise joint distributions. Typical manifold learning methods such as stochastic neighbor embedding can be thought of as a special case [31, 81, 36]. Second, as in [9], popular distances such as Wassersteinp distance and f -divergence can all be naturally extended to such settings, either by embedding P z to R d or projecting P x to R d z and taking the infimum over all orthogonal linear projections. Specifically, for any distance between distributions D ( · , · ) in the Wassersteinp or f -divergence family, define its generalized version to be

where Φ + ( P z , d ) := { P ̂ x : A ̂ x + b ∼ P z , A ∈ R d z × d , b ∈ R d z , AA ⊤ = I d z } is the family of distributions in R d embedded from P z by a linear orthogonal transformations from R d z → R d . Similarly, we can define D -( P x , P z ) and [9] showed that D + ( P x , P z ) = D -( P x , P z ) .

Contrastive learning has deep connections with manifold learning [8] in terms of preserving pairwise similarity, which can be viewed as a special case of preserving the Gromov-Wasserstein distance [36]. In contrast, as we will demonstrate in this work, the optimal latent for generative modeling such as GANs is more closely related to D + ( P x , P z ) .

3 Optimal latent distribution for generative models

Recall the typical training objective of GANs (2.1) and denote its empirical version as L n ( g | P z ) , where n is the sample size. For any given P z , with limited capacity on the generator, min g L n ( g | P z ) reflects the compatibility or closeness between P z and P x . If P z 1 is a better latent than P z 2 , intuitively min g L n ( g | P z 1 ) ≤ min g L n ( g | P z 2 ) . This can serve as an empirical and relative evaluation of the quality of GAN latents. Inspired by this, we can generalize the distance proposed in [9] to construct a more meaningful version specifically for implicit generative models such as GANs.

3.1 GAN-induced distance

Let M (Ω) denote the set of all Borel probability measures on Ω . For any distance between distributions D ( · , · ) in the data space M ( R d ) , define the generalized version associated with a generator g ∈ G mapping from R d z to R d as:

Compared with D + in [9], D G substitutes the linear orthogonal mappings to a general function space. Although D G is not a valid metric as it is defined between different probability measure spaces, it can serve as a viable measurement of the closeness between P z and P x .

In order for D G to be nontrivial, the function space G cannot be too large, as it is known that onedimensional distributions can approximate higher dimensional distributions arbitrarily well by neural network transformations in Wassersteinp distance [86]. On the other hand, G cannot be too small, or it may not be able to adequately extract the information from the latent distribution. In this work, we consider G to be a family of neural networks with bounded complexity, i.e., G c = { g ∈ G : C ( g ) ≤ c } , where c > 0 and C ( · ) is some complexity measure to be specified. C ( · ) can be as specific as the Lipschitz constant, or as general as the size (width, depth, etc.) of the network.

Remark 3.1 . It is worth emphasizing that in (3.1), we do not consider the optimization problem but instead, focus on the existence or the global minimum of inf g ∈G . A toy example is given in Appendix A.1 to illustrate the global minimizers of g ∈ G .

Remark 3.2 (Complexity scaling law) . In the toy case of Appendix A.1, the relationship between D G c and c is an embodiment of the popular scaling law phenomenon [39, 91] with respect to the model complexity. It's worth noting that the requirement on the complexity can be drastically decreased if we consider multiple steps of generation, rather than single-step push-forward transformations. From this perspective, one of DPM's main advantages may just be the increased complexity due to the recursive sampling process.

Directly following the definition in (3.1), we have the following propositions.

Proposition 3.3. Let D be a p -Wasserstein metric, a Jensen-Shannon divergence, a total variation, or an f -divergence. Then, D G ( P z , P x ) in (3.1) is zero if and only if there exists g ∗ ∈ G such that P g ∗ ( z ) = P x .

Proposition 3.4. For two distributional distances, generalizing them across different dimensions via (3.1) does not change their relative relationships, e.g., if D is Wasserstein distance, we still have W G p ( p x , p z ) ≤ W G q ( p x , p z ) , for p ≤ q .

Equipped with D G to compare different latent distributions, we can characterize the ideal P z as the one that minimizes D G ( P z , P x ) , i.e., the latent that requires the lowest generator capacity for L ( g | P z ) to be below a certain threshold 4 . Such a latent depends on both the data and the GAN training algorithm and can be seen as an SSL task, targeting specifically minimizing the required capacity of the corresponding generator, thus easing the training of the generator.

Remark 3.5 (Dimensional collapse) . In contrast to the common dimensional collapse phenomenon [38] found in contrastive learning where the learned feature distribution is only supported on a low-dimensional subspace, minimizing D G ( P z , P x ) with respect to P z encourages it to maintain the intrinsic dimension of the data so that D ( P g ( z ) , P x ) is not ill-posed. The dimensional collapsed latent contributes to the poor performance of SimCLR+DCGAN mentioned in Section 1.

Remark 3.6 (Change-of-scale problem) . Although D G serves as a well-defined measurement of closeness between p z and p x , the latent p z that minimizes D G ( p z , p x ) might be ill-posed. If the complexity is not invariant to scaling, e.g., L p norm, Lipschitz constant, etc., D G ( P z ( λz ) , P x ) is non-increasing with λ . In the next section, we impose extra constraints to circumvent this problem, and the optimal P ∗ z is characterized therein.

3.2 Learning the latent by an encoder

Nowthat we have characterized the ideal latent distribution given the generator family and the data, the next question is how to find such a good latent. Casting it as an estimation problem, the first decision to make is how to parameterize P z . We could adopt Gaussian mixtures as in [84]. Nonetheless, a more universal and powerful way is to employ an encoder network f ∈ F : R d → R d z , similar to SSL methods. With P z parameterized by an encoder, the overall structure is an Autoencoder. In the remaining part, we also refer to the generator as the decoder and use the two words interchangeably.

At the population level, we want to match the distributions of x and the reconstructed g ◦ f ( x ) , i.e.,

Looking at the encoder and the decoder together, we can see that optimizing D ( P g ◦ f , P x ) amounts to learning an encoder f that minimizes D G . This strategy is employed by VQGAN [21], where the encoder and decoder are optimized together with the discretized latent code. Inspired by this observation, we can characterize the optimal latent distribution P ∗ z for a given P x from the perspective of minimizing the distance between distributions in different dimensions by defining P f ∗ as

Notice that P ∗ f not only depends on P x and D , but also F and G . The change-of-scale problem in Remark 3.6 is mostly taken care of by the inherent boundedness of F . More discussion on the choice of G can be found in Section 4.1.

Remark 3.7 (Classification) . As the solution of a new SSL task, f ∗ in (3.2) will preserve wellseparated clusters of P x , though not as discriminative as the features learned from existing SSL methods. See Proposition A.3 in the appendix for more details. We also conducted experiments on simulated data evaluating the classification accuracy in Section 6.1.

With the encoder, (3.2) is not the only option to characterize the optimal latent distribution. Besides directly parameterizing P z , F can also serve an auxiliary role in defining a more general P ∗ z . To this end, in parallel to D G , we can also define D F ( P z , P x ) := inf f ∈F D ( P z , P f ( x ) ) . As D G is to D + , D F generalizes the D -in [9]. Combining D G and D F , we get a generalized measure of distance between the latent and data for Autoencoders as D æ ( P z , P x ) = D G ( P x , P z ) + D F ( P x , P z ) , which resembles the condition number of matrices, but for latent distributions. The optimal latent distribution can also be characterized by

If d z = d , then identity mapping, i.e., P ∗ z = P x , obviously solves (3.2) and (3.3). If like DPMs or VAEs where the latent distributions are chosen to be data-agnostic, e.g., standard Gaussian, D æ ( P z , P x ) can be significantly larger for complicated data.

Remark 3.8 . D G and D F can be zero, especially when P x is supported on a low-dimensional space [9]. As a result, the defined f ∗ or P ∗ z might not be unique. We provide a toy example in Appendix A.2 to illustrate how P ∗ z minimizes the complexity.

4 Improving the latent with Decoupled AutoEncoder

In practice, the encoder and decoder family should be as powerful as possible. However, as stated before, G and F should be chosen carefully in order for the induced D G and D F to be meaningful.

4.1 Necessity of relatively weak decoders

Given a total budget 5 of G and F , intuitively, C ( G ) = C ( F ) tend to result in the best approximation of G ◦ F to the identity mapping. This is reflected in practice where the encoder-decoder pair is often designed with symmetric architecture, with the same number of blocks and parameters, e.g., up-sampling vs. down-sampling, convolution vs. deconvolution, etc. However, when it comes to the quality of the encoded latent, only targeting the reconstruction error may not be the ideal strategy. As we will demonstrate in Section 4.2, a better way is to first target a good latent distribution, and then do reconstruction using the learned latent. Consider two symmetric cases: (larger encoder, smaller decoder) vs (smaller encoder, larger decoder). The reconstruction ability may not be significantly different, although both are sub-optimal, the former tends to produce better latent distribution. See Section 6.1 for empirical evidence in the toy case comparing the two cases.

For a concrete demonstration, consider the simple linear case where the encoder and the decoder are two matrices, denoted as W e ∈ R d × d z and W d ∈ R d z × d respectively. Take the matrix 2 -norm as a measurement of complexity, which is equal to the largest singular value. A similar setting is also considered in [37]. In the linear case, the quality of the latent is equivalent to the reconstruction error and the optimal encoder should recover the principal components. The next theorem states that the Autoencoder could fail if the encoder is not large enough.

Stage 1: Update Encoder with Auxiliary Decoder

9A

fEF

Real Samples

Z

Generated Samples

gEg

Stage 2: Freeze Encoder, Update Actual Decoder

Theorem 4.1. Let the singular values of W d be λ 1 ≥ . . . ≥ λ d z . If ∥ W e ∥ < 1 /λ d z , the linear Autoencoder cannot recover the principal components, and the reconstruction error is sub-optimal.

The linear case provides insights into neural networks as Autoencoders, indicating that the encoder should be relatively strong to better learn the latent distribution. In the more practical nonlinear cases, things become more complicated. The intuition that (larger encoder, smaller decoder) is preferred to (smaller encoder, larger decoder) for learning a more informative latent is supported by the well-known posterior collapse phenomenon in the VAE literature [1, 27, 67, 54], where the encoded distribution becomes completely uninformative if the decoder is too powerful. However, for implicit generative models such as GANs, the generator is obviously the key component and it should be as powerful as possible. There exists a trade-off between generation quality and the quality of the latent. Using a decoder too powerful can hinder the learning process of a good latent while a decoder too small will hurt the generation performance. In practice, whether a decoder family is weak or not can be summarized in an oversimplified way by its size, i.e., a smaller decoder is weaker. Remark 4.2 . We could consider pairing a sufficiently large G with an even larger F . Unfortunately, such a combination may not be a good solution. Besides the added computational cost, an encoderdecoder pair too powerful can drive D G + D F too close to zero, so that it may get overwhelmed by random noise. The signal-to-noise ratio may be too low for effective optimization.

4.2 Decoupled AutoEncoder

To address the trade-off between generation quality and the quality of the latent, we can introduce an auxiliary decoder G A , potentially much smaller/weaker, to substitute G for training the encoder. Once the encoder is trained, we can then freeze it and pair it with G for the decoder training. The key idea is to decouple the learning process of the encoder and the decoder in a 2-stage training scheme we call Decoupled AutoEncoder ( DAE ). Figure 1 illustrates the method overview where the two stages of DAE can be summarized as:

DAE Stage 1: Train ̂ f ∈ F with a small decoder G A to learn a good latent.

DAE Stage 2: Freeze the trained encoder ̂ f . Then train the regular decoder g ∈ G to ensure good generation performance.

By the 2-stage training, one with G A and one with G , without much hyperparameter tuning, DAE can learn a better latent distribution and as a result, achieve better reconstruction and generative modeling. Remark 4.3 (Realizing G A ) . The essence of the first stage is to employ an auxiliary decoder that is relatively weak compared with the encoder. Straightforwardly, we can introduce an extra decoder family. For an easier way of implementation, we can also use various forms of regularization to realize G A . For instance, with a balanced encoder-decoder pair, we can apply strong Dropout [73] to the decoder in the first stage to make it weaker. Both methods are experimented in Section 6.

DAE can be applied in a plug-and-play style to any existing Autoencoder training pipelines with the decoupled 2-stage modification. We demonstrate in Section 6 that compared to various baselines,

e.g., V AE, VQGAN, etc., our DAE-modified versions are better in terms of both reconstruction error and quality of the latent. For instance, we refer to Figure 3 in the appendix where we experimented with VQGAN [21]. In terms of the reconstruction error, Stage 1 is worse than the baseline, which is to be expected since a weaker decoder is used. However, paired with the Stage 1 encoder, the decoder (the same size as the baseline) can achieve better reconstruction performance.

4.3 Sampling from the latent distribution

In most Autoencoder-based generative models, e.g., VAE, adversarial Autoencoder [55], Wasserstein Autoencoder [76], etc., the latent distributions are all modeled by Gaussian distributions. DPMs [72, 32] can also be seen as a type of Autoencoder-based method with the forward noising process as the encoder and the denoising process as the decoder, with the latent distribution being Gaussian. As discussed in Remark 3.2, the decoding process of DPM enjoys drastically increased capacity due to the multi-step generative process.

In our work, with P z parameterized by the encoder, how to generate new samples from it requires extra modeling of the latent. Otherwise, we can only do reconstruction. There are two mainstream methods to sample from the latent distribution. The first is inspired by natural language processing. In VQGAN[21] and Parti [88], the discretized latent codes are autoregressively modeled by transformers as a sequence, similar to GPT [64] while masked image modeling methods [11, 12] resembles BERT [19]. The second is by diffusion process. Latent space diffusion models [68, 60] can be seen as using DPMs to model the latent space of a powerful Autoencoder. We mainly adopt the VQGAN formulation in our numerical experiments for training the Autoencoders. To validate the effectiveness of our DAE, we consider modeling the latent by both transformers and diffusion models.

5 Related works

BourGAN [84] proposed to model the latent distribution as Gaussian mixtures as a remedy for the mode collapse problem of GANs. Instance-conditioned GAN [10] proposed to partition the data manifold into a mixture of overlapping neighborhoods described by a data point and its nearest neighbors to improve the quality of unconditional generation. There is also a line of work that considers post-sampling latent space mining of GANs by exploiting the trained discriminator [75, 13]. [49] further proposed adding an implicit latent transform before the mapping function to improve the latent from its initial distribution in a bi-level optimization fashion. Though these works targeted improving the latent distribution of GANs, they did not characterize the optimal choice nor employ an encoder to model the latent. In comparison, we approach the problem from the perspective of preserving the distance between distributions in different dimensions and propose a 2-stage training scheme involving an encoder.

There are GANs that incorporate an Autoencoder structure. VAEGAN [47, 90] combines a VAE by collapsing the decoder and the generator into one. [2] utilized a standard autoencoder to embed data into the latent space to address the mismatch of the continuous neural network and the discontinuous optimal transport (OT) map. [3] later proposed to utilize semi-discrete OT maps to sample from the latent and train GAN models. VQGAN [87] uses vector-quantized latent codes and transformers to model the tokenized latent distribution. ViT-VQGAN [87] later made improvements over VQGAN. Besides using vision transformers [20] as the backbone, it also introduced an l 2 regularization term to the latent code. However, the aforementioned works did not characterize the optimal latent either and did not point out the necessity of the decoupled 2-stage training. Our DAE approach can also be applied to them for further improvement.

In the literature of V AE, [77] introduced new priors consisting of mixture distribution with components given by variational posteriors conditioned on learnable pseudo-inputs. [17] proposed a two-stage learning scheme to amend the problems associated with the Gaussian encoder/decoder assumptions. [1] stated that obtaining a good ELBO is not enough for good representation learning and proposed to incorporate mutual information and rate-distortion curve to achieve a better balance between the informativeness of the latent and the reconstruction quality. In comparison, our characterization of the optimal latent is from the perspective of matching distributions and minimizing complexities and is not restricted to Gaussian assumptions. It is also worth emphasizing that while the mutual information, closely related to V AEs, can also serve as a "distance" between distributions in different dimensions, it will not work for our benefit since we are considering deterministic encoders and

decoders. In the deterministic case, the conditional entropy would be zero and the mutual information would always be the entropy of the latent. Hence uniformly distributed latent is always ideal, which does not offer new insights.

Masked Autoencoder (MAE) [29] is a generative-based SSL method with the training objective of reconstructing masked images. In MAE, a high mask ratio and a relatively weak decoder have been verified as important factors for good discriminative performance, which is also corroborated by our perspective of preserving distributions in different dimensions. MAEs are designed for discriminative tasks and their generation capability is yet to be explored.

6 Experiments

In this section, we conduct empirical evaluations of our proposed DAE training scheme on a variety of datasets and generative models, from toy Gaussian mixture data to DCGAN on CIFAR-10, to VQGAN and DiT on larger datasets. The detailed experiment settings can be found in Appendix B.

6.1 Toy examples

To showcase the effectiveness of the proposed 2-stage training, we consider a toy Gaussian mixture case with 8 clusters, whose centers are evenly placed on the unit circle, and the variances are 0.25. All samples in R 2 are then projected to R 10 via a linear orthogonal mapping. We employ the vanilla VAE [44] for demonstration with d z = 2 . In our implementation, both the encoder and the decoder are 3-layer fully connected networks with different hidden dimensions to control the size. Giving each cluster a label, we evaluate (1) the nearest neighbor classification accuracy from the 2D latent; (2) the nearest neighbor classification accuracy from the 10D reconstruction; and compare different configurations of the encoder and decoder pair. The average classification accuracy is reported in Table 1, which corroborates our analysis in Section 4, i.e., a relatively weak decoder can learn a better latent and the following 2-stage training can further improve reconstruction.

(64, 128)

Latent Acc (%)

80.6 (7.0)

(128, 128)

Rec Acc (%)

92.2 (6.1)

6.2 GAN on CIFAR-10

As a more comprehensive proof-of-concept, we experiment with DCGAN [63] on the CIFAR-10 [45] dataset. Specifically, we consider a vanilla baseline where the generator has 5 convolutional layers and the discriminator has 4 convolutional layers, followed by a linear classification layer. The latent for DCGANis standard Gaussian. To corroborate our analysis in Section 4, we adapted the DCGAN with varying latent space configurations: (1) We utilize a latent space defined by a CIFAR-10 pre-trained SimCLR model [16], denoting this adaptation as DCGAN-SimCLR. (2) We consider VAEGAN with a learnable latent space. The structures of the decoder and discriminator mirrored those of our baseline DCGAN while the encoder is implemented with a five-layer CNN. (3) We modified VAEGAN to DAE-VAEGAN, where the width of convolutional layers in the decoder is halved in the first stage and then set back to its original structure in the second stage.

We evaluated the aforementioned four methods using the metrics of Inception Score (IS) [70] and Frechlet Inception Distance [30]. With the exception of DCGAN, which generates images from random noises, the other methods first encode an image into latent space, and then generate an image based on that encoding. To this end, we omit the latent space modeling and calculate their metrics from reconstructed images. The results are listed in Table 2, with two take-home messages: (1) Traditional SSL methods may not help with the latent for generative modeling; (2) Decoupling the encoder and decoder training can improve the overall performance of generative modeling.

(128, 64)

87.8

(5.8)

DAE (128, 128)

98.0

(1.3)

6.3 VQGAN

VQGAN [21] consists of two major components. The first is a discrete Autoencoder, where the latent is quantized by looking up nearest neighbors from a learnable codebook. After the encoder and decoder are trained (objective contains adversarial loss), the second component, a transformer model, learns to generate discrete codes autoregressively in the latent space, which are then mapped to images by the learned decoder.

To demonstrate the effectiveness of our DAE, we introduce the decoupled training strategy to the first component of the standard VQGAN. In the first stage, we implement the relatively weak auxiliary decoder by applying Dropout [73] to the decoder with ratio p , and train encoder and the dropped decoder jointly 6 . Then in the second stage, we freeze the encoder and train the full decoder without Dropout. Training of the transformer component is the same as VQGAN. The modified model is denoted as DAE-VQGAN.

Setup. We evaluate DAE-VQGAN by comparing with the baseline VQGAN on the FacesHQ dataset, which is a combination of two face datasets CelebAHQ [51] and FFHQ [43], with 85k training images and 15k validation images in total. In all experiments, the input image size is 256x256. We use the official VQGAN implementation 7 and model architectures for FacesHQ. The Autoencoder is dubbed VQ-f16. The training process of DAE-VQGAN is divided equally into two stages, where in the first stage we apply 2D Dropout (channel-wise) to the decoder with ratio p = 0 . 5 .

Table 4: Class-conditional image generation on ImageNet 256 × 256 using the DiT-L/4 model with different Autoencoders. *: Our reproduction is based on the official DiT implementation, which is consistent with the results reported in Fig.5 of [60].

Improved image reconstruction and generation. To quantitatively evaluate the encoder-decoder component, we train the model on FacesHQ and compute the reconstruction FID over the full training and validation splits. We further train a transformer separately for CelebAHQ and FFHQ with the same architecture as in [21], and compute the generation FID with 50k generated samples against the training split of the respective dataset. As shown in Table 3, DAE-VQGAN achieves significantly improved FID compared to VQGAN, indicating that a stronger encoder-decoder component can be unlocked for VQGAN by our DAE approach. We show qualitative generation results in Figure 2. We see that models trained by DAE-VQGAN can reconstruct and generate images with high fidelity. Figure 3 in the Appendix shows the reconstruction losses during training. It is worth noting that while our DAE-modified VQGAN has higher reconstruction losses in the first stage, it catches up fast in the second stage and converges to a lower level than the baseline. This is to be expected as the first stage of DAE focuses on learning a better latent, which can ease the training of the decoder in the second stage and result in better performance overall.

Decreased Complexity. Minimizing the model complexity is a key motivation for our DAE. To verify our claims, we evaluate the Lipschitz complexity ( C Lip defined in (B.1)) of the encoder f and decoder g in VQGAN. Since g ◦ f is trained to be approximately the identity map, the trends of C Lip ( f ) and C Lip ( g ) are expected to be the opposite, with the ideal number being close to 1 for both if we want the overall complexity C Lip ( f ) + C Lip ( g ) to be minimized. We observed in Figure 4 that the encoder complexity of VQGAN and DAE-VQGAN are 1.40 and 0.88, respectively, with an increasing trend in terms of training steps; the decoder complexity of VQGAN and DAE-VQGAN

(a)

(b)

are 0.73 and 1.16, respectively, with a decreasing trend in terms of training steps. Compared with the baseline, our DAE-modified version produces better latent distribution with a significantly lower complexity while achieving a lower reconstruction error. See Appendix B.2 for more details.

6.4 Diffusion Transformers (DiT)

The encoder in VQGAN is discrete with a codebook. In Appendix B.3, we look into the codebook and investigate how different training strategies affect the behavior of the learned latent. Our DAE can also apply to continuous Autoencoders that are usually regularized by KL divergence. Such KL-regularized AEs are adopted by latent diffusion models such as DiT [60]. Following similar settings as VQGAN, we also modified DiT to DAE-DiT by employing the decoupled training scheme and experimented with the larger ImageNet dataset [18]. The details can be found in Appendix B.4.

For evaluation, we train DiT on the ImageNet dataset with 256 × 256 resolution following the official implementation 8 . Instead of the largest DiT-XL/2 model (675M Params), we select the DiT-L/4 model (458M Params) for higher training efficiency. The corresponding Autoencoder is dubbed KL-f8. AdamW [53] optimizer is employed with a constant learning rate of 10 -4 and a weight decay of 3 × 10 -2 . The batch size is 1024, and the number of epochs is 120. The trained model with Exponentially Moving Average (EMA) is then used to generate 50k images of the 1000 categories equally via a 250-step DDPM sampling [32] and a classifier-free guidance scale of 4. Table 4 compares the image generation quality where we can see that our DAE-DiT achieves significant improvement, similar to what we observed for DAE-VQGAN.

7 Discussion

This work investigates the ideal latent distribution for generative models. We introduce a novel distance between distributions to characterize the ideal P z that minimizes the required model complexity. Practically, we propose a two-stage training scheme called DAE that achieves practical improvements on various models such as GAN, VQGAN, and DiT. Since many of the most powerful generative models are associated with a latent space, the impact of such investigations is potentially very high.

Nevertheless, there are many limitations of our work that call for further research along this line. First, our formulation of the optimal latent distribution is based on D G , D F and D æ , which serve mainly illustrative purposes. The proposed distances cannot be effectively calculated. Second, the effectiveness of DAE is not proven mathematically and the resulting latent is not guaranteed to be closer to P ∗ z . Our work could be further strengthened if the aforementioned limitations can be addressed. It would also be an interesting direction to explore how our method affects latent space disentanglement [52, 26].

References

Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing a broken elbo. In International conference on machine learning , pages 159-168. PMLR, 2018.

Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. Ae-ot: A new generative model based on extended semi-discrete optimal transport. ICLR 2020 , 2019.

Dongsheng An, Yang Guo, Min Zhang, Xin Qi, Na Lei, and Xianfang Gu. Ae-ot-gan: Training gans from data specific latent distribution. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXVI 16 , pages 548-564. Springer, 2020.

Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan. Intrinsic dimension of data representations in deep neural networks. Advances in Neural Information Processing Systems , 32, 2019.

Martín Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning , 2017.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning , pages 214-223. PMLR, 2017.

Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks , 2(1):53-58, 1989.

Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods. arXiv preprint arXiv:2205.11508 , 2022.

Yuhang Cai and Lek-Heng Lim. Distances between probability distributions of different dimensions. arXiv preprint arXiv:2011.00629 , 2020.

Arantxa Casanova, Marlene Careil, Jakob Verbeek, Michal Drozdzal, and Adriana Romero Soriano. Instance-conditioned gan. Advances in Neural Information Processing Systems , 34: 27517-27529, 2021.

Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11315-11325, 2022.

Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming-Hsuan Yang, Kevin Murphy, William T Freeman, Michael Rubinstein, et al. Muse: Text-to-image generation via masked generative transformers. arXiv preprint arXiv:2301.00704 , 2023.

Tong Che, Ruixiang Zhang, Jascha Sohl-Dickstein, Hugo Larochelle, Liam Paull, Yuan Cao, and Yoshua Bengio. Your gan is secretly an energy-based model and you should use discriminator driven latent sampling. Advances in Neural Information Processing Systems , 33:12275-12287, 2020.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15750-15758, 2021.

Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04297 , 2020.

Bin Dai and David Wipf. Diagnosing and enhancing vae models. arXiv preprint arXiv:1903.05789 , 2019.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12873-12883, 2021.

Ruili Feng, Zhouchen Lin, Jiapeng Zhu, Deli Zhao, Jingren Zhou, and Zheng-Jun Zha. Uncertainty principles of encoding gans. In International Conference on Machine Learning , pages 3240-3251. PMLR, 2021.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems , 27, 2014.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems , 27, 2014.

Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07733 , 2020.

Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris. Ganspace: Discovering interpretable gan controls. Advances in neural information processing systems , 33:9841-9850, 2020.

Junxian He, Daniel Spokoyny, Graham Neubig, and Taylor Berg-Kirkpatrick. Lagging inference networks and posterior collapse in variational autoencoders. arXiv preprint arXiv:1901.05534 , 2019.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 9729-9738, 2020.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. arXiv preprint arXiv:2111.06377 , 2021.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Geoffrey Hinton and Sam T Roweis. Stochastic neighbor embedding. In NIPS , volume 15, pages 833-840. Citeseer, 2002.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458 , 2022.

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning , pages 8867-8887. PMLR, 2022.

Tianyang Hu, Wenjia Wang, Cong Lin, and Guang Cheng. Regularization matters: A nonparametric perspective on overparametrized neural network. In International Conference on Artificial Intelligence and Statistics , pages 829-837. PMLR, 2021.

Tianyang Hu, Zhili Liu, Fengwei Zhou, Wenjia Wang, and Weiran Huang. Your contrastive learning is secretly doing stochastic neighbor embedding. In International Conference on Learning Representations , 2023.

Wenlong Ji, Zhun Deng, Ryumei Nakada, James Zou, and Linjun Zhang. The power of contrast for feature learning: A theoretical analysis. arXiv preprint arXiv:2110.02473 , 2021.

Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:2110.09348 , 2021.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4401-4410, 2019.

Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 8110-8119, 2020.

Tero Karras, Miika Aittala, Samuli Laine, Erik Härkönen, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Alias-free generative adversarial networks. Advances in Neural Information Processing Systems , 34:852-863, 2021.

Vahid Kazemi and Josephine Sullivan. One millisecond face alignment with an ensemble of regression trees. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1867-1874, 2014.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Alex Krizhevsky. Learning multiple layers of features from tiny images. University of Toronto , 2009.

Alina Kuznetsova, Hassan Rom, Neil Alldrin, Jasper Uijlings, Ivan Krasin, Jordi Pont-Tuset, Shahab Kamali, Stefan Popov, Matteo Malloci, Alexander Kolesnikov, et al. The open images dataset v4: Unified image classification, object detection, and visual relationship detection at scale. International Journal of Computer Vision , 128(7):1956-1981, 2020.

Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Autoencoding beyond pixels using a learned similarity metric. In International conference on machine learning , pages 1558-1566. PMLR, 2016.

Chun-Liang Li, Wei-Cheng Chang, Yu Cheng, Yiming Yang, and Barnabás Póczos. Mmd gan: Towards deeper understanding of moment matching network. Advances in neural information processing systems , 30, 2017.

Yang Li, Yichuan Mo, Liangliang Shi, and Junchi Yan. Improving generative adversarial networks via adversarial learning in latent space. Advances in Neural Information Processing Systems , 35:8868-8881, 2022.

Tengyuan Liang. How well generative adversarial networks learn distributions. The Journal of Machine Learning Research , 22(1):10366-10406, 2021.

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.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning , pages 41144124. PMLR, 2019.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! a linear vae perspective on posterior collapse. Advances in Neural Information Processing Systems , 32, 2019.

Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders. arXiv preprint arXiv:1511.05644 , 2015.

Facundo Mémoli. Gromov-wasserstein distances and the metric approach to object matching. Foundations of computational mathematics , 11(4):417-487, 2011.

Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invariant representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 6707-6717, 2020.

Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. In NIPS , 2016.

Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv:1609.03499 , 2016.

William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748 , 2022.

Paweł Piwek, Adam Klukowski, and Tianyang Hu. Exact count of boundary pieces of relu classifiers: Towards the proper complexity measure for classification. In Conference on Uncertainty in Artificial Intelligence , pages 1673-1683. PMLR, 2023.

Phillip Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. arXiv preprint arXiv:2104.08894 , 2021.

Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434 , 2015.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning , pages 8821-8831. PMLR, 2021.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Ali Razavi, Aäron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with delta-vaes. arXiv preprint arXiv:1901.03416 , 2019.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487 , 2022.

Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems , 29, 2016.

Antoine Salmona, Julie Delon, and Agnès Desolneux. Gromov-wasserstein distances between gaussian distributions. arXiv preprint arXiv:2104.07970 , 2021.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2020.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research , 15(1):1929-1958, 2014.

Sandeep Subramanian, Sai Rajeswar, Francis Dutil, Chris Pal, and Aaron C. Courville. Adversarial generation of natural language. In Rep4NLP@ACL , 2017.

Akinori Tanaka. Discriminator optimal transport. Advances in Neural Information Processing Systems , 32, 2019.

Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein autoencoders. arXiv preprint arXiv:1711.01558 , 2017.

Jakub Tomczak and Max Welling. Vae with a vampprior. In International Conference on Artificial Intelligence and Statistics , pages 1214-1223. PMLR, 2018.

Ananya Uppal, Shashank Singh, and Barnabás Póczos. Nonparametric density estimation & convergence rates for gans under besov ipm losses. Advances in neural information processing systems , 32, 2019.

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems , 34:11287-11302, 2021.

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 30, 2017.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research , 9(11), 2008.

Vladimir Vapnik. The nature of statistical learning theory . Springer science & business media, 1999.

Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems , 32, 2019.

Chang Xiao, Peilin Zhong, and Changxi Zheng. Bourgan: Generative networks with metric embeddings. Advances in Neural Information Processing Systems , 31, 2018.

Ryuichi Yamamoto, Eunwoo Song, and Jae-Min Kim. Parallel wavegan: A fast waveform generation model based on generative adversarial networks with multi-resolution spectrogram. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 6199-6203. IEEE, 2020.

Yunfei Yang, Zhen Li, and Yang Wang. On the capacity of deep generative networks for approximating distributions. Neural Networks , 145:144-154, 2022.

Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved vqgan. arXiv preprint arXiv:2110.04627 , 2021.

Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al. Scaling autoregressive models for content-rich text-to-image generation. arXiv preprint arXiv:2206.10789 , 2022.

Lantao Yu, Weinan Zhang, Jun Wang, and Yong Yu. Seqgan: Sequence generative adversarial nets with policy gradient. In AAAI Conference on Artificial Intelligence , 2017.

Xianwen Yu, Xiaoning Zhang, Yang Cao, and Min Xia. Vaegan: A collaborative filtering framework based on adversarial variational autoencoders. In IJCAI , pages 4206-4212, 2019.

Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12104-12113, 2022.

Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 586-595, 2018.

A Technical details

A.1 Toy example 1

To illustrate the interplay between D G and G , we devise a one-dimensional toy example where x ∼ N (0 , σ 2 ) , z ∼ U (0 , 1) , g : R → R is monotonically increasing and C ( g ) is the Lipschitz constant upper bounded by c . In this case, the optimal transformation is g ∗ = P -1 x , whose derivative around 0 or 1 is exponentially dependent on µ/σ . Due to the Lipschitz constraint, by the change of distribution formula, we have

From the perspective of regression, the optimal ℓ 2 -projection of g ∗ to G c is

Let z c = √ 2 σ 2 log( c/ 2 π ) and it is clear that ∇ g ∗ ( z c ) = c , or equivalently, p x ( z c ) = 1 /c . As a result, the push-forward distribution is a truncation of P x supported on | x | ≤ z c + c · P x ( -z c ) , i.e., P ̂ g = P x if p x ≥ 1 /c and 1 /c otherwise, which is equivalent to minimizing the total variation (TV) between P g and P x . In this oversimplified example, it is easy to see the scaling between c and D G c , i.e., TV G c ( P z , P x ) ≲ P x ( -z c ) , which is polynomial with c in the Gaussian case.

A.2 Toy example 2

To illustrate the characterization of P ∗ z and how it minimizes the complexity, consider a toy example where d z = 1 , x ∼ N (0 , Σ d ) and both G and F are linear functions with bounded Lipschitz constant c > 1 . Denote the singular values of Σ d as λ 1 ≥ . . . ≥ λ d . In this case, we have C ( P ∗ z ) = 0 and P ∗ z = N (0 , σ 2 ) with σ ∈ ( √ λ n , √ λ 1 ) .

Lemma A.1 (Example 6.1 in [9]) . If D is Wasserstein-2 metric,

where λ 1 ≥ · · · ≥ λ d are the singular values of Σ .

As a direct corollary of the above lemma, we have C ( P ∗ z ) = 0 and P ∗ z = N (0 , σ ) with σ ∈ ( √ λ n , √ λ 1 ) .

A.3 Discussion on P ∗ z and C ( g )

The ideal P ∗ z is characterized as the one that minimizes D G ( P z , P x ) . The link between its definition and the capacity of the generator can be made mathematically sound with an extra assumption.

Assumption A.2. For any P z and P x , D G c ( P z , P x ) is continuous and monotonically decreasing with c.

Assumption A.2 is not provable for general D ( · , · ) and C ( · ) . Nevertheless, we give proof that the ideal P ∗ z that minimizes D G ( P z , P x ) will give rise to the minimal complexity generator if it holds.

Proof. First, recall that G c := { g ∈ G : C ( g ) ≤ c } where C ( g ) is defined as some complexity measurement. It is easy to see that for any c 2 ≥ c 1 > 0 , G c 1 ⊂ G c 2 . Therefore, D G c 2 ( P z , P x ) ≤ D G c 1 ( P z , P x ) .

Consider a generator family G with bounded complexity and we target to have the generated distribution close to the data at level ϵ measured by D , i.e., D ( P x , P g ( z ) ) ≤ ϵ for some P z and g ∈ G .

For a certain (non-degenerate) P z , there exists c > 0 such that D G c ( P z , P x ) = ϵ , that is, by using this latent P z , we need at least complexity c to achieve the ϵ goal. If P z is not ideal, by definition we know that there exists P ∗ z that D G c ( P ∗ z , P x ) = ϵ ∗ < ϵ . By Assumption A.2, we know there exists c ′ < c such that D G c ′ ( P ∗ z , P x ) = ϵ . This means that the goal can be achieved using a lower complexity generator if the latent is ideal.

A.4 Proof of claims in Section 3

Proposition 3.3 and 3.4 directly follow the definition of D G .

D G is to D F as D + is to D -. Due to the enlarged function space, the equality D + = D -no longer holds. Directly following Lemma 2.1 in [9], we have the following inequality: D G ( P z , P x ) ≤ D F ( P z , P x ) where D can be a p -Wasserstein metric or an f -divergence.

Proposition A.3 (Classification guarantee for clustered data) . Suppose P x has m disjoint supports. Then the corresponding P ∗ z must preserve all the clusters.

Proof. Let D KL be the KL divergence (for example). Suppose P x has m disjoint supports. Then

To minimize D KL ( P x , P ̂ x ) , clearly p ̂ x ( x ) = 0 implies p x ( x ) = 0 . As p ̂ x ( x ) increases, D KL ( P x , P ̂ x ) decreases. Since P ̂ z can be selected according to our will, if p x ( x ) = 0 but p ̂ x ( x ) = 0 , it is possible to build ˜ p ̂ z such that ˜ p ̂ x ( x ) ≥ p ̂ x ( x ) for all x , and 1 = ∑ m k =1 ∫ X k ˜ p ̂ x ( x ) dx > ∑ m k =1 ∫ X k p ̂ x ( x ) dx , which implies there exists a positive measure such that ˜ p ̂ x ( x ) > p ̂ x ( x ) . This implies p ̂ x ( x ) cannot minimize D KL ( P x , P ̂ x ) . Thus, we must have if p x ( x ) = 0 , then p ̂ x ( x ) = 0 . This further implies that z must have at least m disjoint supports. If z has less than m disjoint supports, then ̂ x must have less than m disjoint supports, because the linear map is continuous. Then there must exists p x ( x ) = 0 but p ̂ x ( x ) = 0 .

̸

In [9], it has been proved that inf D KL ( P x , P ̂ x ) = inf D KL ( P z , P ̂ z ) . Following the same logic, the number of supports of ̂ z must be smaller than the number of supports of x , which implies the number of supports of ̂ z must be smaller than m . Combining these results with the previous one, we must have the number of supports of ̂ z is m .

A.5 Proof of claims in Section 4

Before we prove Theorem 4.1, we introduce the following lemma on linear auto-encoder.

Lemma A.4. Consider two optimization problems:

which corresponds to Principle Component Analysis (PCA), and

which is a linear auto-encoder. Let W ∗ be the solution to PCA and W ∗ 1 , W ∗ 2 be the solution to linear auto-encoder, and

Then, L 1 = L 2 .

Proof. Clearly, L 2 ≤ L 1 , because we can set W T 1 = W 2 = W ∗ , which leads to

Therefore, it suffices to show L 1 ≤ L 2 .

For a fixed W 2 , By [7], the optimal solution with respect to W 1 is

Therefore,

̸

Let the singular value decomposition of W ∗ 2 be

where U ∈ R n × n and V ∈ R m × m are orthogonal matrices, and Σ ∈ R n × m is a diagonal matrix with singular values λ 1 ≥ . . . ≥ λ m . Therefore,

Note that Σ T Σ is a diagonal matrix. Therefore, we can set

Clearly, which implies

where Σ 1 = diag(1 , ..., 1 , 0 , ..., 0) is a rank m diagonal matrix. This finishes the proof.

With the above lemma, Theorem 4.1 is a simple corollary.

Proof of Theorem 4.1. Recall the proof of Lemma A.4 where the linear encoder is W 1 and the linear decoder is W 2 . Then, in order to realize the PCA solution, we have

Thus, as long as λ 1 λ m < 1 , the encoder is larger in terms of 2-norm.

B More on Experiments

B.1 Experiment details for DAE-VQGAN

Let's recap the VQGAN basics with more details. VQGAN [21] consists of two major components. In the first component, an input image x is encoded by an encoder f into latent representation z = f ( x ) . A quantization step is applied to obtain z q = τ ( z ) , which contains nearest neighbors of entries of z from a learnable codebook. Then a decoder g reconstructs ̂ x = g ( z q ) from the codes. After the encoder and decoder are trained (with adversarial loss), the second component, a transformer model, learns to generate discrete codes in the latent space, which are then mapped to images by the learned decoder g .

We evaluate our DAE modifications to VQGAN on the FacesHQ dataset, which is a combination of two face datasets CelebAHQ and FFHQ, with 85k training images and 15k validation images in total (Table 5). In all experiments, the input image size is 256x256.

We use the official VQGAN implementation 9 and model architectures for FacesHQ, where both encoder f and decoder g have 128 convolution channels. The spatial dimension of the latent representation is 16 × 16 . See Table 7 and Table 8 in [21] for details. All experiments are run on eight V100 GPUs. For training the encoder and decoder, the learning rate is 4 . 5 × 10 -6 , the batch

size is 8 on each GPU (total batch size 64), and the number of training epochs is 80. For training the transformer the learning rate is 2 × 10 -6 and the batch size is 12 on each GPU. The FIDs in Table 3 are obtained by selecting the best results from three independent runs, for both VQGAN and DAE-VQGAN.

To implement our 2-stage DAE-VQGAN, we consider two ways to realize the auxiliary decoder family G A during the first stage of training the encoder. Similar to the settings in Section 6.1 and 6.2, we first utilize an extra auxiliary decoder g aux with half the number of channels as the actual decoder. Then, we experimented with decoder Dropout [73], as mentioned in Remark 4.3.

0

50k

100k

150k

200k

250k

step

0.3

0.4

0.5

0.6

0.7

0.8

reconstruction loss

f

+

g

f

+

gaux

f

+

gaux

/

g

(a) Auxiliary

g

aux

0

20k

40k

60k

80k

100k

step

0.3

0.4

0.5

0.6

0.7

0.8

0.9

reconstruction loss

VQGAN

DAE-VQGAN

(b) Dropout

Halving the channels. We use the same f and g architectures as the VQGAN baseline, and g aux is the same as g except for the number of channels set to 64. To distinguish, we use f + g to represent the models from the baseline VQGAN, and use f + g aux and f + g aux /g to denote the first stage and second stage of DAE-VQGAN, respectively. For the DAE training, we jointly train encoder f and the auxiliary decoder g aux ∈ G A in the first stage (first 40 epochs). Then in the second stage (last 40 epochs), we replace g aux with g , and train g from scratch with f (and the trained codebook) fixed. For a more thorough investigation, we extended the run for both VQGAN and DAE-VQGAN to 200 epochs and the reconstruction error trends are shown in Figure 3(a), where we can see that our DAE-VQGAN achieves significantly better reconstruction. For the actual evaluation on the FacesHQ dataset, we stick to the 80 epoch training scheme for both the baseline and our DAE-VQGAN, with the switch point between the two stages at epoch 40.

Dropout. In the first stage, we implement the relatively weak auxiliary decoder by applying 2D Dropout (channel-wise) to the decoder. In particular, we replace the Dropout layer in each ResNet Block in the decoder of the original VQGAN with 2D Dropout and set ratio p = 0 . 5 , and train the encoder and the dropped decoder jointly. Then in the second stage, we freeze the encoder and train the full decoder without Dropout. There are in total 80 training epochs for both the baseline and DAE-VQGAN (the first 40 epochs belong to the first DAE stage and the second DAE stage takes up the next 40 epochs). Figure 3(b) shows the trends of the reconstruction loss during training, where

we can see that our DAE-VQGAN achieves significantly better reconstruction as well, despite the fact that the encoder f is only updated in the first stage.

0

50k

100k

150k

200k

250k

step

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Lipschitz complexity

f

+

g

f

+

gaux

f

+

gaux

/

g

(a) Encoder

0

50k

100k

150k

200k

250k

step

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Lipschitz complexity

f

+

g

f

+

gaux

f

+

gaux

/

g

(b) Decoder

B.2 Complexity of encoder and decoder

One benefit of our DAE is the ability to better exploit the capacity/complexity of the encoder and decoder. To this end, we investigate the Lipschitz complexity C Lip of encoder f and decoder g defined as:

Without loss of generality, we choose the DAE-VQGAN implemented by introducing an auxiliary g aux for demonstration. The conclusions for Dropout are similar. Figure 4 shows the Lipschitz complexity of the encoder and decoder during 200 epochs of training, where f + g represents the baseline VQGAN, f + g aux represents the first stage of our method, and f + g aux /g represents the second stage of our method where encoder f is initialized with weights from the 40-th epoch (around training step 53k) of the first stage training. Since g ◦ f is trained to be approximately the identity map, the trends of the encoder and decoder are the opposite. As stated in Section 3, the Lipschitz complexity should be ideally close to 1 for both the encoder and decoder. Compared with the baseline, our DAE-VQGAN produces a better latent distribution with significantly lower complexity overall (closer to 1) while achieving smaller reconstruction errors. We also compute a variant of C Lip where the ℓ 2 norms of difference between x 1 and x 2 , and that between their reconstructions, are replaced with perceptual similarity (LPIPS) [92]. This LPIPS complexity variant shows similar trends as C Lip (Figure 5).

B.3 VQ codebook

We look into the codebook and investigate how different training strategies affect the behavior of the learned latent. For both VQGAN and DAE-VQGAN, the codebook size is 1024. For the codebook itself, we compute the pairwise cosine similarity among the learned codes and show the histogram of the similarity scores in Figure 6(a). For VQGAN, the cosine similarities are much more concentrated near 1 and -1, indicating that a large portion of learned codes may be highly similar. In comparison, codes learned by DAE-VQGAN are more scattered in the latent space and therefore can be more efficient. The expressiveness of our codebook is more powerful in approximating the ideal latent

50k

100k

150k

200k

250k

step

400

450

500

550

600

650

LPIPS complexity

f

+

g

f

+

gaux

f

+

gaux

/

g

(a) Encoder

50k

100k

150k

200k

250k

step

0.0016

0.0018

0.0020

0.0022

0.0024

LPIPS complexity

f

+

g

f

+

gaux

f

+

gaux

/

g

(b) Decoder

distribution than the original one. We further plot the eigenvalues of the cosine distance matrix for learned codes in Figure 6(b), where we can clearly see that the eigenvalue decay for the baseline is significantly faster, indicating more severe dimensional collapse.

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

cosine similarity

0

20000

40000

60000

count

VQGAN

DAE-VQGAN

(a)

0

5

10

15

index

10

4

10

3

10

2

10

1

10 0

10 1

10 2

10 3

eigenvalue

VQGAN

DAE-VQGAN

(b)

0

200

400

600

800

1000

code

0

2000

4000

6000

8000

10000

12000

count

VQGAN

DAE-VQGAN

(c)

For the latent expressed using the codebook, we count and sort the number of appearances of codes over the training set of CelebAHQ (Figure 6(c)). We see that DAE-VQGAN can learn less redundant codes and hence their appearances are more evenly distributed. Interestingly, all codes are utilized in the reconstruction process for both VQGAN and DAE-VQGAN, perhaps due to the limited size of the codebook.

B.4 Experiment details for DAE-DiT

We note that while the auxiliary decoder g aux approach presents clear trends of complexity changing, it requires manually changing the network between the two stages, and hence is less convenient to implement than the Dropout approach in practice. Therefore we adopt Dropout as the default approach in our DAE-DiT experiments.

Similar to DAE-VQGAN, in the first stage of DAE-DiT, we implement the relatively weak auxiliary decoder by applying Dropout[73] to g with ratio p = 0 . 2 , and train encoder f and the dropped decoder jointly. Then in the second stage, we freeze encoder f and train the full decoder g without Dropout. We train on the OpenImages [46] dataset and adopt the KL-f8 model as in [68]. The encoder and decoder are trained with learning rate 4 . 5 × 10 -6 and batch size 64 .

For the diffusion training, we adjust the scale factor from the default 0.18215 to 0.8 such that the variance of features from the encoder is close to 1. Following the same setup as the official implementation, the EMA rate is 0.9999 and the classifier-free guidance scale is 4. The computation during both the training and the testing process was executed using FP16 precision except for the attention module for computational stability.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work investigates what constitutes a good latent space for generative models, and proposes a new training paradigm for generative models – DAE. Simply put, with DAE generative models are trained as Autoencoder in two stages. First, a relatively weak decoder is employed, whose purpose is to aid the encoder in learning meaningful representations. Second, the weak decoder is replaced with the “actual” decoder and training is continued. The second decoder is the one eventually evaluated as a generative model.\n\nSTRENGTHS:\n1. The paper is very well written, easy and interesting to follow.\n2. The questions of what constitutes a good latent space and how is one constructed are often overlooked. Most works follow longstanding paradigms of using predetermined distributions (e.g., Gaussian) or whatever is learned from an autoencoder. Posing these questions and formulating the setting, on its own, is impactful.\n3. The proposed method DAE is empirically proven effective, and since it is rather simple to implement, I conjecture it might have a strong impact on future generative modeling works.\n\nWEAKNESSES:\nWhile the mathematical formulation in the paper looks sound to me, I think it does not benefit the paper, and significant portions of it could be moved into the appendix. As the authors acknowledge in the Discussion section, the formulation“ serves mainly illustrative purposes” and is “not proven mathematically”. I don’t see it as an issue, as much of Generative Modeling research (and ML in general) is empiric in nature. The proposed method, DAE, could be introduced as an empirically-supported design, while some of the mathematical formulation could be described to serve as intuition (e.g., Theorem 4.1) . However, dedicating over three pages to it seems excessive to me. In my opinion, most readers would greatly benefit if the experiments in Appendix B were present in the main paper instead of some of the mathematical formulation.\n\nQUESTIONS:\nCould the authors please elaborate on “there are other algorithms to implement the\ndecoupled 2-stage training, and the evaluations in this work are not comprehensive”? Why would the authors not include such evaluations in the paper if they find them relevant?\n\nLIMITATIONS:\nAuthors addressed limitations. One of them (non comprehensive evaluation) raises questions. I would appreciate a clarification on that.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an approach with theoretical analysis to explore a more suitable latent distribution for generation. For this purpose, this paper proposes a novel distance between the latent and target distributions and tries to minimize it to obtain the optimal data-dependent latent distribution. In practice, a two-stage training strategy called Decoupled Autoencoder (DAE) is introduced to leverage a superior latent distribution to improve the generative performance. The experiments on VQGAN and DiT show the effectiveness of the proposed method.\n\nSTRENGTHS:\n1.The paper is well-structured and easy to follow.\n\n2.The methodology is supported by theoretical analysis.\n\nWEAKNESSES:\n1.The proposed methodology is somehow similar to AE-OT-GAN [1], which also learns a latent distribution via an autoencoder and then utilizes the learned latent to train GAN. The authors should explicitly highlight the advantages of their approach over [1].\n\n[1] AE-OT-GAN: Training GANs from Data Specific Latent Distribution. ECCV 2020. \n\n2.Some important related works are missed, particularly the line of works that explore improved latent sampling in GANs, e.g., [1] [2] [3] [4]. In special, AdvLatGAN [2] seeks to adjust and discover a more suitable latent distribution via adversarial learning in both the post-training sampling and the training, and it has already introduced the concept of optimal latent distribution noted $p_z^{op}$ in its theory.\n\n[2] Improving Generative Adversarial Networks via Adversarial Learning in Latent Space. NeurIPS 2022.\n\n[3] Discriminator optimal transport. NeurIPS 2019.\n\n[4] Your gan is secretly an energy-based model and you should use discriminator-driven latent sampling. NeurIPS 2020.\n\n3.The experimental results are insufficient. The comparison basically focuses on the performance gain beyond the vanilla backbones. Authors should compare the other efforts on improving the latent sampling/distribution, e.g., the previously mentioned works.\n\nQUESTIONS:\n1.As discussed in [2], the natural image distribution lie upon multiple disjoint manifolds, which potentially requires a discontinuous latent distribution (Prop. 3.2). Does the proposed method adhere to this requirement?\n\n2.In Fig. 2 (a), actually it is hard to discern noticeable performance improvements in the comparison.\n\n3.I am curious why this method can be applied to diffusion models, for which the Gaussian is the stable distribution for noise addition. If the latent distribution is altered, how does the mechanism of diffusion models still function?\n\nPlease also see the discussion of weaknesses. If the issues are resolved, I would be inclined to reassess and potentially raise my rating.\n\nLIMITATIONS:\nPlease see above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper first proposes a new framework to analyze latent spaces in the context of generative models.\nThis framework takes inspiration from prior results about GANs, which allowed interpreting the min-max training objective as computing a distance between distributions to be minimized, to define a similar distance between the latent distribution and the data distribution.\nFrom their analysis, they derive a simple two-step training for auto-encoders to learn better latents and reconstruction, in which they first train the encoder with a weak decoder to extract good latents, and then train a larger decoder to get better reconstructions.\n\nThey perform experiments in a simple toy case, and then with commonly used models such as GANs, VQGAN, and DiTs.\n\nSTRENGTHS:\n- The paper is easy to follow\n- It proposes a novel view on latent codes, including an explanation of some different properties between SSL and generative latent codes, and a theoretical framework to describe why a powerful encoder/decoder pair can't learn a good latent code.\n- The proposed practical solution is simple and their experiments show that it can improve performances in a variety of generative settings.\n\nWEAKNESSES:\n- While the fresh view on latent codes is interesting, it doesn't provide any theoretical guarantees.\nAs far as I understand, the main conclusion is that in order to obtain a good latent code the encoder and decoder should not be too powerful but is not able to give an indication about how the correct balance. Note that in the context of Variational Autoencoders, a similar conclusion had been reached before: that a too powerful decoder is a good explanation for the phenomenon called *posterior collapse* that describes a state in which the latent code is completely uninformative [1].\nIt is good that it is formally extended to other auto-encoders, but unsurprising.\n\n- Because of this lack of quantification of optimal complexity, it is quite unclear whether the positive results obtained in VQGAN and DiT settings are actually related to the theoretical conclusions or not. It could very well be completely unrelated and just confirmation bias.\n\n- Related work present different generative models with loose links to the proposed method, but does not discuss any study about the latent spaces of generative models. In addition to links with posterior collapse in VAEs [1], I would also be very interested to read what the authors have to say about sparsity and disentanglement properties of VAE [2], and PCA directions in GAN space [3] among other things.\nMoreover, I have to point out that contrary to the paragraph in related works, Mask Autoencoders have been explored for their generative capabilities [4].\n\n\n--- \n[1] Fixing a Broken ELBO. ICML 2018.\nAlexander A. Alemi, Ben Poole, Ian Fischer, Joshua V. Dillon, Rif A. Saurous, Kevin Murphy.\n\n[2] Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations. ICML 2019.\nFrancesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, Olivier Bachem\n\n[3] GANSpace: Discovering Interpretable GAN Controls. NeurIPS 2020.\nErik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, Sylvain Paris\n\n[4] MaskGIT: Masked Generative Image Transformer. CVPR 2022.\nHuiwen Chang, Han Zhang, Lu Jiang, Ce Liu, William T. Freeman\n\nQUESTIONS:\nOne minor thing that was not clear to me is Remark 3.5. Could the authors explicit why scaling invariance would be a problem in this context?\n\nIn any case, despite its limitations, I still find the paper interesting for both its theoretical and practical contributions.\nI trust the authors will fix the related work section, and my current rating anticipates that the authors will demonstrate their willingness to do so.\n\n-----\nPost-rebuttal update:\nHaving taken into consideration the rebuttal, discussions and comments from the other reviewers, I updated my score from WA to Accept.\n\nLIMITATIONS:\nThe authors are very open regarding the limitations of their submission.\n\nPotential negative societal impacts of generative models are not. Widely acknowledge ones include amplification of biases, potential misuse for propagating fake information, and concerns about lack of attribution and copyright infringements.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an asymmetric training scheme for auto-encoders that double as image generator. Based on analytical insights that the decoder should have less capacity than the encoder for the encoder to capture correctly the data distribution, they propose a first training cycle where a strong encoder and a weak decoder are trained jointly. This produces a latent distribution able to better capture the data distribution. Then, in a second stage crucial for the end application, the encoder is frozen and a strong decoder is trained using the \"good\" latent distribution. Experiments are carried out on faces datasets with VQGAN and class conditional imagenet with a diffusion model and show the proposed training scheme is promising.\n\nSTRENGTHS:\nThe strengths of this paper are :\n- It is a simple method, yet effective. As such, it should be easy to reproduce. The improvements are maybe questionable (no error bars, test against the training set - as is traditional in this domain), but given that they are reported with widely different architecture and on different dataset, there is more confidence than usual that it actually works.\n- The analysis part is really nice. The linear example in particular sheds some light as to why under a constrained budget, the complexity of the encoder should exceed that of the decoder to properly minimize the projected distance (divergence) to the data distribution.\n\nWEAKNESSES:\nThe main weakness is that most of the paper is about the analytical part, which is a bit handwavy (acknowledged in the conclusion), and lacks a good structure. The reading could be improved by having a clear outline of the work such that one does not wonder where the text is headed to after each section.\n\nQUESTIONS:\nl.146 the link between the ideal $P_z$ and a low capacity generator is tenuous at best. There is no formal proof of that statement as far as I understand the paper. I can see a structural risk argument (ie, among all generator that equally minimize $D(p_x, p_z)$ chose the lowest complexity one), but it is more a principle and not a formal proof. Correct?\nl. 190 $D(P_z, P_x)$ will be significantly larger in case of arbitrarily chosen $P_z$. This you don't know. If $P_x$ and $P_z$ have the same topology, then there exists a continuous mapping between the two that can be perfectly approximated by a sufficiently large neural network. In practice, maybe $D$ will be larger, but not necessarily in theory.\nl. 198 I do not agree that $C(F) = C(G)$ is intuitive. For optimal reconstruction, we have $f = g^{-1}$ (assuming all that is required) and the Lipschitz constant of $g^{-1}$ can be arbitrarily large before that of $g$. For example, with $g(x) = x^2$ and $f(x) = \\sqrt{x}$ on $(\\epsilon, 1)$, the Lipschitz constant of $f$ can be made arbitrarily large. And vice versa. I think this depends whether it is easier (ie, the mapping has lower Lipschitz) to map from $x$ to $z$ or from $z$ to $x$. But what does it depend on? The size of $z$?\nl. 312 this is a bit cheating. You should learn an encoder that inverse the (frozen) Gaussian space and evaluate on reconstruction for DCGAN.\nl 354 Why is the Lipschitz of the encoder lower than that of the decoder for the proposed method? Isn't it completely contradictory to all the story of the paper?\nA2. Fix notation errors.\n\nLIMITATIONS:\nThe conclusion is a list of limitations that is very honest.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}