SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Applied to supervised learning, Theorem 3.1 bounds the expected risk of a Gibbs posterior  $q$  which, upon receiving a previously unseen datapoint  $\mathbf{x} \sim \mu$ , samples a predictor  $h$  dependent on  $\mathbf{x}$ , and uses it to make a prediction. Note that the family  $\mathcal{E}$  from Assumption 1 does not appear in the bound, which has nice consequences in practice. Indeed one may pick a loss function  $\ell$  that fits the problem, and then find a family  $\mathcal{E}$  for which the continuity assumption is satisfied with constant  $K$  that is as small as possible.
Note also that, in the tradition of PAC-Bayesian bounds, Theorem 3.1 does not make any assumptions on the nature of the elements of  $\mathcal{H}$  (e.g.  $\mathcal{H}$  could be a class of functions, a set of neural network's parameters, etc). Therefore, the theorem is very general and could be applied to different domains and models. In the following sections, we will use a specific kind of hypothesis class  $\mathcal{H} = \mathcal{Z}$ , in order to capture the VAE's latent space.
4 Generalization bounds for the Reconstruction Loss
For the remainder of this work,  $\| \cdot \|$  denotes the  $L_{2}$  norm, and we assume the instance space  $\mathcal{X}$  is Euclidean, and the latent space  $\mathcal{Z} = \mathbb{R}^{d_{\mathcal{Z}}}$ , where  $d_{\mathcal{Z}} > 0$ . Both  $\mathcal{X}$  and  $\mathcal{Z}$  are equipped with the Euclidean distance as the underlying metric. Therefore, if  $\mathbf{x}, \mathbf{x}' \in \mathcal{X}$ ,  $d(\mathbf{x}, \mathbf{x}') = \| \mathbf{x} - \mathbf{x}' \|$ .
The following assumption states that the encoder and decoder networks have finite Lipschitz norms.
Assumption 2. The encoder and decoder are Lipschitz-continuous w.r.t. their inputs, meaning there exist real numbers  $K_{\phi}, K_{\theta} > 0$  such that for any  $\mathbf{x}_1, \mathbf{x}_2 \in \mathcal{X}$  and  $\mathbf{z}_1, \mathbf{z}_2 \in \mathcal{Z}$ ,
$$
\left\| Q _ {\phi} \left(\mathbf {x} _ {1}\right) - Q _ {\phi} \left(\mathbf {x} _ {2}\right) \right\| \leq K _ {\phi} \| \mathbf {x} _ {1} - \mathbf {x} _ {2} \| \tag {4}
$$
$$
\left\| g _ {\theta} \left(\mathbf {z} _ {1}\right) - g _ {\theta} \left(\mathbf {z} _ {2}\right) \right\| \leq K _ {\theta} \left\| \mathbf {z} _ {1} - \mathbf {z} _ {2} \right\|. \tag {5}
$$
Recall the definition of  $Q_{\phi}$  from Equation (2). Note that in practice, one can estimate the Lipschitz constant of trained networks (Fazlyab et al., 2019; Latorre et al., 2020) or train the VAE with preset Lipschitz constants (Barrett et al., 2022).
Moreover, we define the reconstruction loss  $\ell_{\mathrm{rec}}^{\theta}$  with the  $L_{2}$  norm, instead of the squared  $L_{2}$  norm, which enables us to exploit the properties of a metric. We discuss this choice in Section 6. In order to be consistent with the PAC-Bayesian framework, we define the loss function as follows:  $\ell_{\mathrm{rec}}^{\theta}:\mathcal{Z}\times \mathcal{X}\to [0,\infty)$ ,
$$
\ell_ {\operatorname {r e c}} ^ {\theta} (\mathbf {z}, \mathbf {x}) = \| \mathbf {x} - g _ {\theta} (\mathbf {z}) \|. \tag {6}
$$
Our goal is to apply the general bound of Theorem 3.1 to the VAE model. But first, since Theorem 3.1 requires Assumption 1 to be satisfied, we start by showing that if the encoder and decoder networks have finite Lipschitz norms, then Assumption 1 holds.
Proposition 4.1. Consider a VAE with parameters  $\phi$  and  $\theta$  and let  $K_{\phi}, K_{\theta} \in \mathbb{R}$  be the Lipschitz norms of the encoder and decoder respectively. Then the variational distribution  $q_{\phi}(\mathbf{z}|\mathbf{x})$  satisfies Assumption 1, with  $\mathcal{E} = \{f: \mathcal{Z} \to \mathbb{R} \text{ s.t. } \| f \|_{Lip} \leq K_{\theta}\}$ ,  $\ell = \ell_{rec}^{\theta}$ , and  $K = K_{\phi}K_{\theta}$ .
Proof idea. The proof of Proposition 4.1 is in Appendix C, we provide a brief summary here. To prove the first part of Assumption 1, we first notice that if  $\mathcal{E}$  is the set of real-valued  $K_{\theta}$ -Lipschitz continuous functions, then  $d_{\mathcal{E}}$  is a scaling of the Wasserstein distance. In addition, since  $W_{1} \leq W_{2}$ , using the closed form of the Wasserstein-2 distance between Gaussian distributions, one can show that  $d_{\mathcal{E}}(q_{\phi}(\cdot|\mathbf{x}_1), q_{\phi}(\cdot|\mathbf{x}_2)) \leq K_{\phi}K_{\theta}\| \mathbf{x}_1 - \mathbf{x}_2\|$ . Finally, the second part of the assumption is a consequence of the definition of the loss function and the Lipschitz continuity of the decoder.
Proposition 4.1 tells us that Assumption 1 holds for VAEs. Consequently, we can utilize our general bound of Theorem 3.1 to obtain generalization guarantees. This leads to the following general PAC-Bayesian bound for the VAE's reconstruction loss.
Theorem 4.2. Let  $\mathcal{X}$  be the instance space,  $\mu \in \mathcal{M}_{+}^{1}(\mathcal{X})$  the data-generating distribution,  $\mathcal{Z}$  the latent space,  $p(\mathbf{z})\in \mathcal{M}_+^1 (\mathcal{Z})$  the prior distribution on the latent space,  $\theta$  the decoder's parameters,
$\delta \in (0,1), \lambda > 0$  be real numbers. With probability at least  $1 - \delta$  over the random draw of  $S \sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
\begin{array}{l} \underset {\mathbf {x} \sim \mu} {\mathbb {E}} \underset {q _ {\phi} (\mathbf {z} | \mathbf {x})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x}) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) + \frac {1}{\lambda} \left[ \sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) | | p (\mathbf {z})) + \right. \\ \frac {\lambda K _ {\phi} K _ {\theta}}{n} \sum_ {i = 1} ^ {n} \underset {\mathbf {x} \sim \mu} {\mathbb {E}} d (\mathbf {x}, \mathbf {x} _ {i}) + \log \frac {1}{\delta} + n \log \underset {\mathbf {z} \sim p (\mathbf {z})} {\mathbb {E}} \underset {\mathbf {x} \sim \mu} {\mathbb {E}} e ^ {\frac {\lambda}{n} \left(\underset {\text {r e c}} {\mathbb {E}} _ {\mathbf {x} ^ {\prime} \sim \mu} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} ^ {\prime}) - \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x})\right)} \Bigg ], \\ \end{array}
$$
where  $K_{\phi}$  and  $K_{\theta}$  are the Lipschitz norms of the encoder and the decoder (see (4) and (5)) and  $\mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})}$  is a shorthand for  $\mathbb{E}_{\mathbf{z}\sim q_{\phi}(\mathbf{z}|\mathbf{x})}$ .
Note that the choice of the hyperparameter  $\beta$  in the VAE's optimization objective (3) correlates with the choice of the hyperparameter  $\lambda$  in Theorem 4.2 (e.g.  $\lambda = n$  corresponds to  $\beta = 1$ ). Note also that the encoder and decoder are not treated the same way in Theorem 4.2. Indeed, the inequality holds for a given decoder, but uniformly for any encoder. We discuss this subtle difference and its practical consequences in Section 6.
Theorem 4.2 can be seen as a general framework. In order to obtain a useful upper bound, one needs to bound the average distance and the exponential moment on the right-hand side. In the sections below, we provide upper bounds for these terms under various assumptions on the instance space.
4.1 Reconstruction Guarantees for Bounded Instance Spaces
In the following theorem, we provide a special case of Theorem 4.2 when the instance space's diameter  $\Delta \stackrel{\mathrm{def}}{=} \sup_{\mathbf{x}, \mathbf{x}' \in \mathcal{X}} d(\mathbf{x}, \mathbf{x}')$  is finite (see Section C.2 for the proof).
Theorem 4.3. Let  $\mathcal{X}$  be the instance space,  $\Delta < \infty$  its diameter,  $\mu \in \mathcal{M}_{+}^{1}(\mathcal{X})$  the data-generating distribution,  $\mathcal{Z}$  the latent space,  $p(\mathbf{z}) \in \mathcal{M}_{+}^{1}(\mathcal{Z})$  the prior on the latent space,  $\theta$  the decoder's parameters,  $\delta \in (0,1)$ ,  $\lambda > 0$  be real numbers. With probability at least  $1 - \delta$  over the random draw of  $S \sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
\begin{array}{l} \underset {\mathbf {x} \sim \mu} {\mathbb {E}} \underset {q _ {\phi} (\mathbf {z} | \mathbf {x})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x}) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) \right\} + \frac {1}{\lambda} \left(\sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) \| p (\mathbf {z})) + \right. \\ \lambda K _ {\phi} K _ {\theta} \Delta + \log \frac {1}{\delta} + \frac {\lambda^ {2} \Delta^ {2}}{8 n} \Bigg). \\ \end{array}
$$
The left-hand side of this inequality is the expected reconstruction loss for samples  $\mathbf{x} \sim \mu$ , while the right-hand side is the empirical reconstruction and KL losses, plus an additional term depending on the Lipschitz constants of the VAE and the model's diameter.
Note that for real-life datasets, the diameter of the instance space might be very large and non-representative of the structure and complexity of the data. Indeed, it is common to scale image datasets in order to utilize a specific architecture (Radford et al., 2016). In the following section, we provide a special case of Theorem 4.2 under the manifold hypothesis on the data-generating process.
4.2 Reconstruction Guarantees under the Manifold Assumption
The manifold assumption (Fodor, 2002; Narayanan and Mitter, 2010; Fefferman et al., 2016) states that most high-dimensional datasets encountered in practice lie close to low-dimensional manifolds. This assumption is exploited by latent variable generative models such as GANs and VAEs, which approximate high-dimensional datasets using transformations of distributions on a low-dimensional space. The works of Schreuder et al. (2021) and Mbacke et al. (2023) provide generalization bounds for GANs, by assuming that the data-generating distribution is a smooth transformation of the uniform distribution on  $[0,1]^{d^*}$ , where  $d^*$  is the intrinsic dimension. However, since the standard VAE calls for a standard Gaussian prior, in the following theorem, we assume  $\mu$  is a smooth transformation of the standard Gaussian distribution  $p^*$  on  $\mathbb{R}^{d^*}$ . We consider the case when  $p^*$  is the uniform distribution on  $[0,1]^{d^*}$  in the supplementary material.
Theorem 4.4. Let  $\mathcal{X}$  be the instance space,  $\mu \in \mathcal{M}_{+}^{1}(\mathcal{X})$  the data-generating distribution,  $\mathcal{Z}$  the latent space,  $p(\mathbf{z})\in \mathcal{M}_{+}^{1}(\mathcal{Z})$  the prior distribution on the latent space,  $\theta$  the decoder's parameters,  $\delta \in (0,1),\lambda >0,a > 0$  real numbers. Assume the data-generating distribution  $\mu = g^{*}\sharp p^{*}$ , where  $p^*$  is the standard Gaussian distribution on  $\mathbb{R}^{d^*}$  and  $g^{*}\in \mathrm{Lip}_{K_{*}}(\mathbb{R}^{d^{*}},\mathcal{X})$ . With probability at least  $1 - \delta -\frac{nd^{*}}{2} e^{-a^{2} / 2}$  over the random draw of  $S\sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
\begin{array}{l} \underset {\mathbf {x} \sim \mu} {\mathbb {E}} \underset {q _ {\phi} (\mathbf {z} | \mathbf {x})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x}) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) \right\} + \frac {1}{\lambda} \left(\sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) | | p (\mathbf {z})) + \right. \\ \lambda K _ {\phi} K _ {\theta} K _ {*} \sqrt {(1 + a ^ {2}) d ^ {*}} + \log \frac {1}{\delta} + \frac {\lambda^ {2} K _ {*} ^ {2}}{2 n} \Bigg). \\ \end{array}
$$
Let us clarify the role of the new parameter  $a > 0$ . Each training sample  $\mathbf{x}_i \in S$  can be expressed as  $\mathbf{x}_i = g^*(\mathbf{w}_i)$ , where  $\mathbf{w}_i \sim p^*$ . Since  $p^*$  is the standard Gaussian distribution on  $\mathbb{R}^{d^*}$ , all samples  $\mathbf{w}_i$  will be inside a hypercube  $[-a, a]^{d^*}$ , with high probability. This uncertainty is reflected in the lowered confidence (from  $1 - \delta$  in Theorem 4.2 to  $1 - \delta - \frac{nd^*}{2} e^{-a^2 / 2}$  in Theorem 4.4), and can be controlled by choosing a large enough value of  $a$ . The proof of Theorem 4.4 is in the supplementary material (Section C.3), we provide a short summary below.
Proof idea. The proof starts with Theorem 4.2, and uses the assumptions of Theorem 4.4 to obtain upper bounds on the exponential moment and the average distance. To derive the upper bound on the exponential moment, we observe that the function  $\mathbf{z} \mapsto \ell_{\mathrm{rec}}^{\theta}(\mathbf{z}, \mathbf{x})$  is  $K_*$ -Lipschitz continuous, then we use a dimension-free upper bound on the MGF of Lipschitz-continuous functions of Gaussian random variables. Furthermore, we obtain the upper bound on the average distance  $\frac{1}{n} \sum_{i=1}^{n} \mathbb{E}_{\mathbf{x} \sim \mu} \| \mathbf{x} - \mathbf{x}_i \|$ , by using Holder's inequality and the expectation of a non-central  $\chi^2$  distribution. Then, we upper-bound the probability that  $\mathbf{w}_i \in [-a, a]^{d^*}$  for all  $1 \leq i \leq n$  using the error function and Bernoulli's inequality. Finally, we use the union bound to update the overall confidence.
5 Generalization Bounds for Regeneration and Generation
Let  $\hat{\mu}_{\phi,\theta}$  be the empirical regenerated distribution, meaning
$$
\hat {\mu} _ {\phi , \theta} = \frac {1}{n} \sum_ {i = 1} ^ {n} g _ {\theta} \sharp q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}). \tag {7}
$$
In other words, sampling  $\mathbf{x} \sim \hat{\mu}_{\phi, \theta}$  is done by sampling  $\mathbf{z} \sim q_{\phi}(\mathbf{z}|\mathbf{x}_i)$  where  $i$  is uniformly sampled from  $\{1, \dots, n\}$ , then passing  $\mathbf{z}$  through the decoder:  $\mathbf{x} = g_{\theta}(\mathbf{z})$ . It is therefore the distribution regenerated by the VAE, given the training set  $S = \{\mathbf{x}_1, \dots, \mathbf{x}_n\}$  as input.
In this section, we provide statistical guarantees on the regenerative and generative properties of VAEs. More precisely, we derive upper bounds for the quantities  $W_{1}(\mu ,\hat{\mu}_{\phi ,\theta})$  and  $W_{1}(\mu ,g_{\theta}\sharp p(\mathbf{z}))$ . Note that the average distance term does not appear in the bounds of this section. This is because instead of relying on Theorem 3.1, the results of this section depend upon a preliminary lemma (Lemma B.1), which does not necessitate Assumption 1.
5.1 Regeneration and Generation Guarantees for Bounded Instance Spaces
The following theorem presents our first upper bound on the distance between the input distribution and the empirical regenerated distribution.
Theorem 5.1. Under the definitions and assumptions of Theorem 4.3, we have that with probability at least  $1 - \delta$  over the random draw of  $S\sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
W _ {1} (\mu , \hat {\mu} _ {\phi , \theta}) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) \right\} + \frac {1}{\lambda} \left(\sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) | | p (\mathbf {z})) + \log \frac {1}{\delta} + \frac {\lambda^ {2} \Delta^ {2}}{8 n}\right).
$$
As we can see, the right-hand side of Theorem 5.1 depends on the empirical reconstruction loss and KL-divergence. This guarantees that as the VAE's empirical risk decreases, the regenerated distribution gets closer to the data-generating distribution. The proof of Theorem 5.1 exploits the fact
that the underlying metric on  $\mathcal{X}$  is the Euclidean distance  $d(\mathbf{x},\mathbf{x}^{\prime}) = \| \mathbf{x} - \mathbf{x}^{\prime}\|$ , which is also used to define the reconstruction loss  $\ell_{\mathrm{rec}}^{\theta}$  (see Equation 6). The full proof can be found in Appendix D.
The following theorem provides an upper bound of the distance between the input distribution and the VAE's generated distribution.
Theorem 5.2. Under the definitions and assumptions of Theorem 4.3, we have that with probability at least  $1 - \delta$  over the random draw of  $S\sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
\begin{array}{l} W _ {1} (\mu , g _ {\theta} \sharp p (\mathbf {z})) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) \right\} + \frac {1}{\lambda} \left(\sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) | | p (\mathbf {z})) + \right. \\ \log \frac {1}{\bar {\delta}} + \frac {\lambda^ {2} \Delta^ {2}}{8 n}) + \frac {K _ {\theta}}{n} \sum_ {i = 1} ^ {n} \sqrt {\| \mu_ {\phi} (\mathbf {x} _ {i}) \| ^ {2} + \left\| \sigma_ {\phi} (\mathbf {x} _ {i}) - \vec {1} \right\| ^ {2}}, \\ \end{array}
$$
where  $\vec{1} \in \mathbb{R}^{d_z}$  denotes the vector whose entries are all 1.
The right-hand side of Theorem 5.2 is equal to the right-hand side of Theorem 5.1, plus an additional term depending on the Wasserstein-2 distance  $W_{2}(q_{\phi}(\mathbf{z}|\mathbf{x}_{i}),p(\mathbf{z}))$ , which is used in the proof because of its closed form for Gaussian distributions. Hence, the right-hand side of Theorem 5.2 augments the VAE's optimization objective with  $W_{2}(q_{\phi}(\mathbf{z}|\mathbf{x}_{i}),p(\mathbf{z}))$ , suggesting that a good generative performance may require the latent codes to be even closer to the prior. This is consistent with the findings of Zhao et al. (2019), who showed that in order to improve generative performance, the latent codes need to be much closer to the prior, which may disrupt the balance between reconstruction loss and KL-loss.
5.2 Regeneration and Generation Guarantees under the Manifold Assumption
Similar to what we did in Section 4.2, we assume that the data-generating distribution is a smooth transformation of the standard Gaussian distribution on  $\mathbb{R}^{d^*}$ , where  $d^*$  is the intrinsic dimension of the dataset. This yields the following results.
Theorem 5.3. Under the definitions and assumptions of Theorem 4.4, with probability at least  $1 - \delta$  over the random draw of  $S \sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
W _ {1} (\mu , \hat {\mu} _ {\phi , \theta}) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) \right\} + \frac {1}{\lambda} \left(\sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) | | p (\mathbf {z})) + \log \frac {1}{\delta} + \frac {\lambda^ {2} K _ {*} ^ {2}}{2 n}\right).
$$
Note that the intrinsic and extrinsic dimensions do not explicitly appear in this inequality, although they may affect the reconstruction and KL loss.
We now present our last result, an upper bound on the Wasserstein distance between the input distribution and the VAE's generated distribution, under the manifold assumption.
Theorem 5.4. Under the definitions and assumptions of Theorem 4.4, with probability at least  $1 - \delta$  over the random draw of  $S \sim \mu^{\otimes n}$ , the following holds for any posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$ :
$$
\begin{array}{l} W _ {1} (\mu , g _ {\theta} \sharp p (\mathbf {z})) \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\underset {q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i})} {\mathbb {E}} \ell_ {r e c} ^ {\theta} (\mathbf {z}, \mathbf {x} _ {i}) \right\} + \frac {1}{\lambda} \left(\sum_ {i = 1} ^ {n} \mathrm {K L} (q _ {\phi} (\mathbf {z} | \mathbf {x} _ {i}) | | p (\mathbf {z})) + \right. \\ \log \frac {1}{\bar {\delta}} + \frac {\lambda^ {2} K _ {*} ^ {2}}{2 n}) + \frac {K _ {\theta}}{n} \sum_ {i = 1} ^ {n} \sqrt {\| \mu_ {\phi} (\mathbf {x} _ {i}) \| ^ {2} + \left\| \sigma_ {\phi} (\mathbf {x} _ {i}) - \vec {1} \right\| ^ {2}}, \\ \end{array}
$$
where  $\vec{1} \in \mathbb{R}^{d_{\mathcal{Z}}}$  denotes the vector whose entries are all 1.
Theorem 5.2 and Theorem 5.4 show that by minimizing the VAE's objective, one is also minimizing the Wasserstein distance between the input distribution and the VAE's generated distribution.
From the upper bounds given by Theorems 5.1 and 5.3, one can deduce rates of convergence of  $O(n^{-1/2})$  (when  $\lambda \approx \sqrt{n}$ ) for the empirical regenerated distribution. Note that  $\lambda \approx n$  leads to the much faster rate of  $n^{-1}$ , but then the bounds do not converge to the empirical risk, but to a larger positive number, dependent on the input distribution. Similarly, Theorems 5.2 and 5.4 provide rates of convergence of  $O(n^{-1/2})$  for the VAE's generated distribution.
The different treatments of  $\theta$  and  $\phi$ . The bounds we've presented in this work hold for a given decoder  $\theta$ , but uniformly for all encoders. In practice, this means that the risk certificate has to be computed using samples different from the ones used to train the VAE. This is different from the usual PAC-Bayesian trick (Germain et al., 2009; Parrado-Hernandez et al., 2012; Pérez-Ortiz et al., 2021, see also Remark F.2) of splitting the training set to learn the prior, then training the model on the whole training set, because the decoder and encoder are jointly optimized. Instead, one has to make sure that the model is only trained on samples distinct from the ones used to compute the bound. The same method would be necessary when computing the risk certificates given by the recent PAC-Bayesian bounds of Rivasplata et al. (2020) and Haddouche and Guedj (2022), since those bounds are not uniformly valid for any posterior.
The reconstruction loss. In our bounds, the reconstruction loss is the  $L_{2}$  norm (RMSE), instead of the squared  $L_{2}$  norm (MSE). In practice, one can still optimize a VAE with the MSE (or any other reconstruction loss, e.g. the cross entropy loss), and then compute the bounds using the RMSE. However, if the reconstruction loss is not the RMSE, then the optima of the chosen optimization objective might differ from the ones minimizing the right-hand side of the bounds. Therefore, if the goal is to minimize the bounds, one should utilize the RMSE as the reconstruction loss.
Conclusion. It is common, when applying PAC-Bayesian theory to new problems, to add additional stochasticity in order to account for the PAC-Bayesian distributions on the hypothesis class. For instance, Mbacke et al. (2023) added distributions on the parameters of a WGAN's generator, in order to perform a PAC-Bayesian analysis. However, because of the seamless integration of the PAC-Bayesian and VAE frameworks, such modification to the original problem has been avoided in this work. We matched the prior and posterior distributions on the VAE's latent space to the PAC-Bayesian prior and posterior, which allowed us to recover the VAE's optimization objective. We provide preliminary experiments on synthetic datasets in the supplementary material.
This work is a humble contribution to the theoretical understanding of VAEs. We developed novel PAC-Bayesian bounds suited to the analysis of VAEs and provided generalizations bounds for the VAE's reconstruction loss. In addition, we also derived upper bounds on the Wasserstein distance between the input distribution and the VAE's generative model's distribution. These bounds depend on the VAE's empirical optimization objective and the data-generating process. By integrating the VAE and PAC-Bayesian frameworks, we hope to establish PAC-Bayesian theory as a prime tool for the theoretical analysis of VAEs.
This research is supported by the Canada CIFAR AI Chair Program, and the NSERC Discovery grant RGPIN-2020- 07223. F. Clerc is funded by IVADO through the DEEL Project CRDPJ 537462 18 and by a grant from NSERC.