SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Consider a generator family  $G$  with bounded complexity and we want to have the generated distribution close to the data at level  $\epsilon$  measured by  $D$ , i.e.,  $D(P_x, P_{g(z)}) \leq \epsilon$  for some  $P_z$  and  $g \in G$ .
For a certain (non-degenerate)  $P_{z}$ , there exists  $c > 0$  such that  $D^{G_c}(P_z,P_x) = \epsilon$ , that is, by using this latent  $P_{z}$ , we need at least complexity  $c$  to achieve the  $\epsilon$  goal. If  $P_{z}$  is not ideal, by definition we know that there exists  $P_{z}^{*}$  that  $D^{G_c}(P_z^*,P_x) = \epsilon^* < \epsilon$ . By Assumption A.2, we know there exists  $c' < c$  such that  $D^{G_{c'}}(P_z^*,P_x) = \epsilon$ . 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^{\mathcal{G}}$ .
$D^{\mathcal{G}}$  is to  $D^{\mathcal{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^{\mathcal{G}}(P_z,P_x)\leq D^{\mathcal{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
$$
D _ {K L} \left(P _ {x}, P _ {\widehat {x}}\right) = \sum_ {k = 1} ^ {m} \int_ {X _ {k}} p _ {x} (x) \log \left(\frac {p _ {x} (x)}{p _ {\widehat {x}} (x)}\right) d x.
$$
To minimize  $D_{KL}(P_x, P_{\widehat{x}})$ , clearly  $p_{\widehat{x}}(x) = 0$  implies  $p_x(x) = 0$ . As  $p_{\widehat{x}}(x)$  increases,  $D_{KL}(P_x, P_{\widehat{x}})$  decreases. Since  $P_{\widehat{x}}$  can be selected according to our will, if  $p_x(x) = 0$  but  $p_{\widehat{x}}(x) \neq 0$ , it is possible to build  $\widetilde{p_{\widehat{x}}}$  such that  $\widetilde{p_{\widehat{x}}} (x) \geq p_{\widehat{x}}(x)$  for all  $x$ , and  $1 = \sum_{k = 1}^{m}\int_{X_k}\widetilde{p_{\widehat{x}}} (x)dx > \sum_{k = 1}^{m}\int_{X_k}p_{\widehat{x}}(x)dx$ , which implies there exists a positive measure such that  $\widetilde{p_{\widehat{x}}} (x) > p_{\widehat{x}}(x)$ . This implies  $p_{\widehat{x}}(x)$  cannot minimize  $D_{KL}(P_x, P_{\widehat{x}})$ . Thus, we must have if  $p_x(x) = 0$ , then  $p_{\widehat{x}}(x) = 0$ . This further implies that  $z$  must have at least  $m$  disjoint supports. If  $z$  has less than  $m$  disjoint supports, then  $\widehat{x}$  must have less than  $m$  disjoint supports, because the linear map is continuous. Then there must exist  $p_x(x) = 0$  but  $p_{\widehat{x}}(x) \neq 0$ .
In [9], it has been proved that  $\inf D_{KL}(P_x, P_{\widehat{x}}) = \inf D_{KL}(P_z, P_{\widehat{z}})$ . Following the same logic, the number of supports of  $\widehat{z}$  must be smaller than the number of supports of  $x$ , which implies the number of supports of  $\widehat{z}$  must be smaller than  $m$ . Combining these results with the previous one, we must have the number of supports of  $\widehat{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:
$$
\min  _ {W \in \mathbb {R} ^ {m \times n}} \| X - W W ^ {T} X \| _ {F} ^ {2}, \tag {A.1}
$$
which corresponds to Principle Component Analysis (PCA), and
$$
\min  _ {W _ {1}, W _ {2} \in \mathbb {R} ^ {m \times n}} \| X - W _ {2} W _ {1} X \| _ {F} ^ {2}, \tag {A.2}
$$
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
$$
L _ {1} = \left\| X - W ^ {*} \left(W ^ {*}\right) ^ {T} X \right\| _ {F} ^ {2}, L _ {2} = \left\| X - W _ {2} ^ {*} W _ {1} ^ {*} X \right\| _ {F} ^ {2}. \tag {A.3}
$$
Then,  $L_{1} = L_{2}$
Proof. Clearly,  $L_{2} \leq L_{1}$ , because we can set  $W_{1}^{T} = W_{2} = W^{*}$ , which leads to
$$
L _ {2} = \left\| X - W _ {2} ^ {*} W _ {1} ^ {*} X \right\| _ {F} ^ {2} \leq \left\| X - W ^ {*} \left(W ^ {*}\right) ^ {T} X \right\| _ {F} ^ {2} = L _ {1}.
$$
Therefore, it suffices to show  $L_{1} \leq L_{2}$ .
For a fixed  $W_{2}$ , By [7], the optimal solution with respect to  $W_{1}$  is
$$
W _ {1} = \left(W _ {2} ^ {T} W _ {2}\right) ^ {- 1} W _ {2} ^ {T}.
$$
$$
L _ {2} = \left\| X - W _ {2} ^ {*} \left(\left(W _ {2} ^ {*}\right) ^ {T} W _ {2} ^ {*}\right) ^ {- 1} \left(W _ {2} ^ {*}\right) ^ {T} X \right\| _ {F} ^ {2}.
$$
Let the singular value decomposition of  $W_2^*$  be
$$
W _ {2} ^ {*} = U \Sigma V,
$$
where  $U \in \mathbb{R}^{n \times n}$  and  $V \in \mathbb{R}^{m \times m}$  are orthogonal matrices, and  $\Sigma \in \mathbb{R}^{n \times m}$  is a diagonal matrix with singular values  $\lambda_1 \geq \ldots \geq \lambda_m$ . Therefore,
$$
\begin{array}{l} W _ {2} ^ {*} \left(\left(W _ {2} ^ {*}\right) ^ {T} W _ {2} ^ {*}\right) ^ {- 1} \left(W _ {2} ^ {*}\right) ^ {T} X = U \Sigma V \left(V ^ {T} \Sigma^ {T} U ^ {T} U \Sigma V\right) ^ {- 1} V ^ {T} \Sigma^ {T} U ^ {T} X \\ = U \Sigma (\Sigma^ {T} \Sigma) ^ {- 1} \Sigma^ {T} U ^ {T} X. \\ \end{array}
$$
Note that  $\Sigma^T\Sigma$  is a diagonal matrix. Therefore, we can set
$$
\widetilde {W} = U \Sigma (\Sigma^ {T} \Sigma) ^ {- 1 / 2}.
$$
Clearly,
$$
\left(\widetilde {W}\right) ^ {T} \widetilde {W} = \left(\Sigma^ {T} \Sigma\right) ^ {- 1 / 2} \Sigma^ {T} U ^ {T} U \Sigma \left(\Sigma^ {T} \Sigma\right) ^ {- 1 / 2} = I _ {m},
$$
$$
\begin{array}{l} L _ {2} = \left\| X - W _ {2} ^ {*} \left(\left(W _ {2} ^ {*}\right) ^ {T} W _ {2} ^ {*}\right) ^ {- 1} \left(W _ {2} ^ {*}\right) ^ {T} X \right\| _ {F} ^ {2} \\ = \| X - U \Sigma (\Sigma^ {T} \Sigma) ^ {- 1} \Sigma^ {T} U ^ {T} X \| _ {F} ^ {2} \\ = \left\| X - U \Sigma_ {1} U ^ {T} X \right\| _ {F} ^ {2} \\ = \| X - U _ {m} U _ {m} ^ {T} X \| _ {F} ^ {2} \\ \geq \| X - U _ {m} ^ {*} \left(U _ {m} ^ {*}\right) ^ {T} X \| _ {F} ^ {2} \\ = L _ {1}, \\ \end{array}
$$
where  $\Sigma_{1} = \mathrm{diag}(1,\dots,1,0,\dots,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
$$
\begin{array}{l} \left\| W _ {1} \right\| _ {2} = \left\| \left(W _ {2} ^ {T} W _ {2}\right) ^ {- 1} W _ {2} ^ {T} \right\| _ {2} = \left\| \left(\Sigma^ {T} \Sigma\right) ^ {- 1} \Sigma^ {T} \right\| _ {2} = \lambda_ {m} ^ {- 1}, \\ \left\| W _ {2} \right\| _ {2} = \left\| \Sigma \right\| _ {2} = \lambda_ {1}. \tag {A.4} \\ \end{array}
$$
Thus, as long as  $\lambda_1\lambda_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  $\pmb{x}$  is encoded by an encoder  $f$  into latent representation  $z = f(\pmb{x})$ . A quantization step is applied to obtain  $z_{q} = \tau(z)$ , which contains nearest neighbors of entries of  $z$  from a learnable codebook. Then a decoder  $g$  reconstructs  $\widehat{\pmb{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  $256 \times 256$ .
We use the official VQGAN implementation 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 \times 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 \times 10^{-6}$ , the batch
<sup>9</sup>https://github.com/CompVis/taming-transformers
Dataset	Training	Validation	Total
CelebaHQ	25k	5k	30k
FFHQ	60k	10k	70k
FacesHQ	85k	15k	100k
Table 5: Number of images in the face datasets.
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 \times 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  $\mathcal{G}^{\mathcal{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.
(a) Auxiliary  $g_{aux}$
(b) Dropout
Figure 3: Training reconstruction loss trends for two implementations of DAE-VQGAN. (a) Halving the channels: while the reconstruction loss of our first stage  $(f + g_{aux})$  converges to a higher value than the baseline  $(f + g)$ , our second stage  $(f + g_{aux} / g)$  converges to a lower level with much faster convergence speed; (b) Dropout: the switch between the 2 stages of DAE occurs at halfway (40 epoch, around step  $53\mathrm{k}$ ), where we can see that our method trails behind before the switch and surpasses the baseline after the switch and converges to a lower level of the reconstruction error.
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} \in \mathcal{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.
(a) Encoder
(b) Decoder
Figure 4: Lipschitz complexity of encoder and decoder during training. Our DAE second stage  $f + g_{aux} / g$  utilizes the frozen encoder  $f$  from the first stage  $f + g_{aux}$  at epoch 40 (around step 53k, where the blue line and the orange line intersect). Since  $g \circ f$  is trained to be the identity map, the trends for  $C_{lip}(g)$  and  $C_{lip}(f)$  are approximately reciprocal and the ideal values for them are 1. Compared with the baseline  $f + g$  (blue line), our DAE stage one  $f + g_{aux}$  (orange line) has a significantly lower  $C_{lip}(f)$  that is closer to 1. Since the second stage  $f + g_{aux} / g$  does not update the encoder, the green line stays flat.
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:
$$
C _ {L i p} (f) = \mathbf {E} _ {x _ {1}, x _ {2} \sim P _ {x}} \left[ \frac {\| \tau \left(f \left(x _ {1}\right)\right) - \tau \left(f \left(x _ {2}\right)\right) \| _ {2}}{\| x _ {1} - x _ {2} \| _ {2}} \right], \tag {B.1}
$$
$$
C _ {L i p} (g) = \mathbf {E} _ {x _ {1}, x _ {2} \sim P _ {x}} \left[ \frac {\| g (\tau (f (x _ {1})) - g (\tau (f (x _ {2}))) \| _ {2}}{\| \tau (f (x _ {1})) - \tau (f (x _ {2})) \| _ {2}} \right].
$$
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 \circ 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  $\ell_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
Figure 5: Lipschitz complexity with perceptual similarity (LPIPS complexity) of encoder and decoder. The overall trends are similar to those in  $C_{lip}$  shown in Figure 4.
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.
Figure 6: (a) Histogram of cosine similarity of codes. (b) Top 20 eigenvalues of the cosine distance matrix. (c) Number of appearances of codes in sorted order.
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 \times 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.