SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathcal{G}$  with an even larger  $\mathcal{F}$ . Unfortunately, such a combination may not be a good solution. Besides the added computational cost, an encoder-decoder pair too powerful can drive  $D^{\mathcal{G}} + D^{\mathcal{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  $\mathcal{G}^A$ , potentially much smaller/weaker, to substitute  $\mathcal{G}$  for training the encoder. Once the encoder is trained, we can then freeze it and pair it with  $\mathcal{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  $\widehat{f} \in \mathcal{F}$  with a small decoder  $\mathcal{G}^A$  to learn a good latent.
- DAE Stage 2: Freeze the trained encoder  $\widehat{f}$ . Then train the regular decoder  $g \in \mathcal{G}$  to ensure good generation performance.
By the 2-stage training, one with  $\mathcal{G}^A$  and one with  $\mathcal{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  $\mathcal{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  $\mathcal{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., VAE, 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 noisig 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.
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 VAE, [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 VAEs, 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.
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  $\mathbb{R}^2$  are then projected to  $\mathbb{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.
Table 1: The mean (variance) of the classification accuracy w.r.t. difference hidden dimensions of the (encoder, decoder) over 10 replications.
Hidden Dim	(64, 128)	(128, 64)
Latent Acc (%)	80.6 (7.0)	87.8 (5.8)
Hidden Dim	(128, 128)	DAE (128, 128)
Rec Acc (%)	92.2 (6.1)	98.0 (1.3)
Table 2: The performance of DCGAN with different latent distributions.
Method	IS (↑)	FID (↓)
DCGAN (reproduced)	5.68	51.76
DCGAN-SimCLR	3.93	168.23
VAEGAN	5.82	48.11
DAE-VAEGAN	6.16	46.12
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 DCGAN is 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 Frechet 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.
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<sup>6</sup>. 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  $256 \times 256$ . We use the official VQGAN implementation<sup>7</sup> 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 3: Reconstruction FID over FacesHQ training and validation sets, and transformer generation FID on CelebAHQ and FFHQ training sets. †: Evaluated on the publicly available pre-trained model on FacesHQ. *: Our reproduction is based on the official VQGAN implementation.
Method	Reconstruction	Generation
Train	Val	CelebaHQ	FFHQ
VQGAN	4.81†	6.27†	10.2	9.6
VQGAN*	4.23	5.83	9.97	10.44
DAE-VQGAN	2.01	3.82	8.58	8.36
Table 4: Class-conditional image generation on ImageNet  $256\times 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].
Method	FID	sFID	IS
DiT*	35.16	7.33	39.25
DAE-DiT	32.29	6.90	41.71
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\circ 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
<sup>6</sup>We also implemented the auxiliary weak decoder by leveraging an extra decoder family with the number of convolution channels halved. The overall performance is similar and we put the results in Appendix B.1
<sup>7</sup>https://github.com/CompVis/taming-transformers
Figure 2: (a) Top: original image; Middle: reconstruction of VQGAN; Bottom: reconstruction of DAE-VQGAN. (b) Generated samples of DAE-VQGAN.
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 \times 256$  resolution following the official implementation. 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 \times 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.
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^{\mathcal{G}}, D^{\mathcal{F}}$  and  $D^{\mathbf{x}}$ , 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].
$^{8}$ https://github.com/facebookresearch/DiT.
[89] Lantao Yu, Weinan Zhang, Jun Wang, and Yong Yu. Seqgan: Sequence generative adversarial nets with policy gradient. In AAAI Conference on Artificial Intelligence, 2017.
[90] 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.
[91] 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.
[92] 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.1 Toy example 1
To illustrate the interplay between  $D^{\mathcal{G}}$  and  $\mathcal{G}$ , we devise a one-dimensional toy example where  $x \sim N(0, \sigma^2)$ ,  $z \sim U(0, 1)$ ,  $g: \mathbb{R} \to \mathbb{R}$  is monotonically increasing and  $C(g)$  is the Lipschitz constant upper bounded by  $c$ . In this case, the optimal transformation is  $g^{*} = P_{x}^{-1}$ , whose derivative around 0 or 1 is exponentially dependent on  $\mu / \sigma$ . Due to the Lipschitz constraint, by the change of distribution formula, we have
$$
p _ {g (z)} (x) = \mathbb {I} _ {\{g ^ {- 1} (x) \in (0, 1) \}} \cdot \nabla g ^ {- 1} \geq \frac {1}{c}.
$$
From the perspective of regression, the optimal  $\ell_2$ -projection of  $g^*$  to  $\mathcal{G}_c$  is
$$
\widehat {g} (z) = \left\{ \begin{array}{l l} g ^ {*} (z) & \text {i f} \nabla g ^ {*} (z) \leq c, \\ c (z - \frac {1}{2}) & \text {e l s e .} \end{array} \right.
$$
Let  $z_{c} = \sqrt{2\sigma^{2}\log(c / 2\pi)}$  and it is clear that  $\nabla 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|\leq z_{c} + c\cdot P_{x}(-z_{c})$ , i.e.,  $P_{\widehat{g}} = P_{x}$  if  $p_x\geq 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^{\mathcal{G}_c}$ , i.e.,  $TV^{\mathcal{G}_c}(P_z,P_x)\lesssim 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$ ,  $\pmb{x} \sim N(0, \Sigma_{d})$  and both  $\mathcal{G}$  and  $\mathcal{F}$  are linear functions with bounded Lipschitz constant  $c > 1$ . Denote the singular values of  $\Sigma_{d}$  as  $\lambda_{1} \geq \ldots \geq \lambda_{d}$ . In this case, we have  $C(P_{z}^{*}) = 0$  and  $P_{z}^{*} = N(0, \sigma^{2})$  with  $\sigma \in (\sqrt{\lambda_{n}}, \sqrt{\lambda_{1}})$ .
Lemma A.1 (Example 6.1 in [9]). If  $D$  is Wasserstein-2 metric,
$$
D ^ {-} (N (0, \sigma^ {2}), N (0, \Sigma)) = \left\{ \begin{array}{l l} \sqrt {\lambda_ {n}} - \sigma & i f \sigma <   \sqrt {\lambda_ {n}}, \\ 0 & i f \sqrt {\lambda_ {n}} \leq \sigma \leq \sqrt {\lambda_ {1}}, \\ \sigma - \sqrt {\lambda_ {1}} & i f \sigma > \sqrt {\lambda_ {1}}, \end{array} \right.
$$
where  $\lambda_1 \geq \dots \geq \lambda_d$  are the singular values of  $\Sigma$ .
As a direct corollary of the above lemma, we have  $C(P_z^*) = 0$  and  $P_{z}^{*} = N(0,\sigma)$  with  $\sigma \in (\sqrt{\lambda_n},\sqrt{\lambda_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(\cdot, \cdot)$  and  $C(\cdot)$ . 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} \coloneqq \{g \in G : C(g) \leq c\}$  where  $C(g)$  is defined as some complexity measurement. It is easy to see that for any  $c_{2} \geq c_{1} > 0$ ,  $G_{c_{1}} \subset G_{c_{2}}$ . Therefore,  $D^{G_{c_{2}}}(P_{z}, P_{x}) \leq D^{G_{c_{1}}}(P_{z}, P_{x})$ .