SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tianyang Hu<sup>1</sup>, Fei Chen<sup>1</sup>, Haonan Wang<sup>2</sup>, Jiawei Li<sup>1</sup>, Wenjia Wang<sup>3</sup>, Jiacheng Sun<sup>1*</sup>, Zhenguo Li<sup>1</sup>
<sup>1</sup> Huawei Noah's Ark Lab, <sup>2</sup> National University of Singapore
$^{3}$  Hong Kong University of Science and Technology (Guangzhou)
In generative modeling, numerous successful approaches leverage a low-dimensional 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.
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 \times 256$  to  $32 \times 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.
*Corresponding to: Jiacheng Sun (sunjiacheng1@huawei.com)
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).
Notations. For a function  $f: \Omega \to \mathbb{R}$ , let  $\|f\|_p = \left( \int_{\Omega} |f(x)|^p dx \right)^{1/p}$ . For a vector  $x$ ,  $\|x\|_p$  denotes its  $l_p$ -norm.  $L_p$  and  $\ell_p$  are used to distinguish function norms and vector norms. For two positive sequences  $\{a_n\}_{n \in \mathbb{N}}$  and  $\{b_n\}_{n \in \mathbb{N}}$ , we write  $a_n \lesssim b_n$  if there exists a constant  $C > 0$  such that  $a_n \leq C b_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 \circ 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 \in \mathcal{G}$  from random noises  $z$  to the data  $x$ . At the population level, its objective function can be expressed in general as
$$
\inf  _ {g \in \mathcal {G}} D _ {h} \left(P _ {x}, P _ {g (z)}\right), z \sim P _ {z}, \tag {2.1}
$$
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],
${}^{2}$  Calculated using reconstructed data. More details can be found in Section 6.2
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 \in \mathbb{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<sup>3n</sup> 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 Wasserstein-  $p$  distance and  $f$ -divergence can all be naturally extended to such settings, either by embedding  $P_{z}$  to  $\mathbb{R}^{d}$  or projecting  $P_{x}$  to  $\mathbb{R}^{dz}$  and taking the infimum over all orthogonal linear projections. Specifically, for any distance between distributions  $D(\cdot ,\cdot)$  in the Wasserstein-  $p$  or  $f$ -divergence family, define its generalized version to be
$$
D ^ {+} (P _ {x}, P _ {z}) := \inf  _ {P _ {\widehat {x}} \in \Phi^ {+} (P _ {z}, d)} D (P _ {x}, P _ {\widehat {x}}),
$$
where  $\Phi^{+}(P_{z},d)\coloneqq \{P_{\widehat{x}}:\mathbf{A}\widehat{x} +\mathbf{b}\sim P_{z},\mathbf{A}\in \mathbb{R}^{d_{z}\times d},\mathbf{b}\in \mathbb{R}^{d_{z}},\mathbf{A}\mathbf{A}^{\top} = \mathbb{I}_{d_{z}}\}$  is the family of distributions in  $\mathbb{R}^d$  embedded from  $P_z$  by a linear orthogonal transformations from  $\mathbb{R}^{d_z}\to \mathbb{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})\leq \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  $\mathcal{M}(\Omega)$  denote the set of all Borel probability measures on  $\Omega$ . For any distance between distributions  $D(\cdot, \cdot)$  in the data space  $\mathcal{M}(\mathbb{R}^d)$ , define the generalized version associated with a generator  $g \in \mathcal{G}$  mapping from  $\mathbb{R}^{d_z}$  to  $\mathbb{R}^d$  as:
$$
D ^ {\mathcal {G}} \left(P _ {z}, P _ {x}\right) := \inf  _ {g \in \mathcal {G}} D \left(P _ {g (z)}, P _ {x}\right). \tag {3.1}
$$
<sup>3</sup>We use the word "distance" quite generally, including not only well-defined metrics, but also divergences.
Compared with  $D^{+}$  in [9],  $D^{\mathcal{G}}$  substitutes the linear orthogonal mappings to a general function space. Although  $D^{\mathcal{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^{\mathcal{G}}$  to be nontrivial, the function space  $\mathcal{G}$  cannot be too large, as it is known that one-dimensional distributions can approximate higher dimensional distributions arbitrarily well by neural network transformations in Wasserstein-p distance [86]. On the other hand,  $\mathcal{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  $\mathcal{G}$  to be a family of neural networks with bounded complexity, i.e.,  $\mathcal{G}_c = \{g\in \mathcal{G}:C(g)\leq c\}$ , where  $c > 0$  and  $C(\cdot)$  is some complexity measure to be specified.  $C(\cdot)$  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\in \mathcal{G}}$ . A toy example is given in Appendix A.1 to illustrate the global minimizers of  $g\in \mathcal{G}$ .
Remark 3.2 (Complexity scaling law). In the toy case of Appendix A.1, the relationship between  $D^{\mathcal{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^{\mathcal{G}}(P_z,P_x)$  in (3.1) is zero if and only if there exists  $g^{*}\in \mathcal{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_{p}^{\mathcal{G}}(p_{x},p_{z})\leq W_{q}^{\mathcal{G}}(p_{x},p_{z})$ , for  $p\le q$ .
Equipped with  $D^{\mathcal{G}}$  to compare different latent distributions, we can characterize the ideal  $P_{z}$  as the one that minimizes  $D^{\mathcal{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<sup>4</sup>. 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^{\mathcal{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^{\mathcal{G}}$  serves as a well-defined measurement of closeness between  $p_z$  and  $p_x$ , the latent  $p_z$  that minimizes  $D^{\mathcal{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^{\mathcal{G}}(P_z(\lambda z), P_x)$  is non-increasing with  $\lambda$ . 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
Now that 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 \in \mathcal{F}: \mathbb{R}^{d} \to \mathbb{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  $\pmb{x}$  and the reconstructed  $g \circ f(\pmb{x})$ , i.e.,
$$
\inf  _ {f \in \mathcal {F}, g \in \mathcal {G}} D (P _ {x}, P _ {g \circ f (\boldsymbol {x})}) = \inf  _ {f \in \mathcal {F}} \left(\inf  _ {g \in \mathcal {G}} D (P _ {x}, P _ {g \circ f (\boldsymbol {x})})\right) = \inf  _ {f \in \mathcal {F}} D ^ {\mathcal {G}} (P _ {x}, P _ {f (\boldsymbol {x})}).
$$
4More mathematical ground for this argument can be found in Appendix A.3.
Looking at the encoder and the decoder together, we can see that optimizing  $D(P_{g \circ 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
$$
f ^ {*} = \underset {f \in \mathcal {F}} {\operatorname {a r g m i n}} D ^ {\mathcal {G}} \left(P _ {x}, P _ {f (x)}\right). \tag {3.2}
$$
Notice that  $P_{f}^{*}$  not only depends on  $P_{x}$  and  $D$ , but also  $\mathcal{F}$  and  $\mathcal{G}$ . The change-of-scale problem in Remark 3.6 is mostly taken care of by the inherent boundedness of  $\mathcal{F}$ . More discussion on the choice of  $\mathcal{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 well-separated 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}$ ,  $\mathcal{F}$  can also serve an auxiliary role in defining a more general  $P^{*}$ . To this end, in parallel to  $D^{\mathcal{G}}$ , we can also define  $D^{\mathcal{F}}(P_{z}, P_{x}) := \inf_{f \in \mathcal{F}} D(P_{z}, P_{f(x)})$ . As  $D^{\mathcal{G}}$  is to  $D^{+}$ ,  $D^{\mathcal{F}}$  generalizes the  $D^{-}$  in [9]. Combining  $D^{\mathcal{G}}$  and  $D^{\mathcal{F}}$ , we get a generalized measure of distance between the latent and data for Autoencoders as  $D^{\mathfrak{x}}(P_{z}, P_{x}) = D^{\mathcal{G}}(P_{x}, P_{z}) + D^{\mathcal{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
$$
P _ {z} ^ {*} = \underset {P _ {z}} {\operatorname {a r g m i n}} D ^ {\mathfrak {x}} \left(P _ {z}, P _ {x}\right). \tag {3.3}
$$
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^{\mathfrak{e}}(P_z,P_x)$  can be significantly larger for complicated data.
Remark 3.8.  $D^{\mathcal{G}}$  and  $D^{\mathcal{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,  $\mathcal{G}$  and  $\mathcal{F}$  should be chosen carefully in order for the induced  $D^{\mathcal{G}}$  and  $D^{\mathcal{F}}$  to be meaningful.
4.1 Necessity of relatively weak decoders
Given a total budget $^5$  of  $\mathcal{G}$  and  $\mathcal{F}$ , intuitively,  $C(\mathcal{G}) = C(\mathcal{F})$  tend to result in the best approximation of  $\mathcal{G} \circ \mathcal{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} \in \mathbb{R}^{d \times d_{z}}$  and  $W_{d} \in \mathbb{R}^{d_{z} \times 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.
The budget and  $C(\cdot)$  here refer to the general concept of size (width, depth, etc.) of the network.
Figure 1: Overview of the DAE 2-stage training. DAE addresses the trade-off between generation quality and the latent quality where the encoder is only updated in the first stage with an auxiliary weaker decoder and then frozen in the second stage while the actual decoder is being trained.
Theorem 4.1. Let the singular values of  $W_{d}$  be  $\lambda_{1} \geq \ldots \geq \lambda_{d_{z}}$ . If  $\| W_{e} \| < 1 / \lambda_{d_{z}}$ , the linear Autoencoder cannot recover the principal components, and the reconstruction error is sub-optimal.