Generalization Guarantees for Representation Learning via Data-Dependent Gaussian Mixture Priors

We thank the reviewer for the relevant comments.

1. Thank you. Modified.

2. In the revised version, we made it clearer by writing "mutual information (MI)" in the second paragraph of the introduction.

3. We agree with the reviewer and have added [DH83] and [GBC16]. If the reviewer recommends any other related references, we would be thankful to share them with us.

   [DH83] SR Dalal and WJ Hall. Approximating priors by mixtures of natural conjugate priors. Journal of the Royal Statistical Society: Series B (Methodological), 45(2):278–286, 1983.

   [GBC16] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning, 2016.

4. The suitable choice of the value of the parameter $M$ depends on several factors, especially the dimension $d$ of the latent variables and the "sparsity" or number of "subpopulations" in the latent variables (which depends on the complexity of the data and the complexity of the encoder). Therefore, in general, it is not possible to know precisely an appropriate value of $M$ in advance and should be hence considered as a hyperparameter.

   As future work, it would be interesting to investigate how common dimensionality reduction and unsupervised clustering techniques, such as the t-SNE dimensionality reduction and clustering using a Gaussian mixture, could be used to obtain a good "initial guess" for the choice of $M$.

We sincerely thank the reviewer for the relevant comments and suggestions that helped improve our paper.

Regarding the literature: We have added a brief discussion of how our approach compares to previous related work.

Questions:

1. We thank the reviewer for bringing up this discussion. Generally, it is a common practice to design practical approaches that are "inspired" by the theoretical results by deviating slightly from the required assumptions. Furthermore, to make the proposed approach more founded, in the initial version (footnote on page 7) we briefly mentioned how the established bounds can be extended to cases where the symmetry is slightly violated using differentially private tools, e.g. ones similar to [DR18]. In fact, there are various ways to do such an extension and we believe a separate dedicated work is required to study them. However, to address the reviewer's concern, in Appendix~B of the revised versions we have detailed the explanations in the footnote and proposed two preliminary methods to extend the results to allow a slight violation of the symmetry condition (with a slight penalty). The first one is by using the differentially private prior tools (as was mentioned in the footnote on page 7 of the original manuscript) and the second one is by a new definition of a "partially symmetric prior". As a future work, it would be interesting to quantify more precisely the amount of "violation" from the symmetry of the prior in our proposed approach. Here, we would like to emphasize that in particular in the few first epochs, which have been reported in many works to be critical for the rest of the training procedure, such violation is intuitively small since the model has not yet "memorized" the training dataset, and both adding noise and passing over different batches has the effect of "forgetting" which reduces the dependence of the prior on the next batch.

2. Thanks for such a keen remark. First, we would like to elaborate a bit onto why, and how, the type of symmetry that is needed to derive bounds in terms of latent variable complexity is different from those that appear, e.g., in the mentioned exchangeable priors of Audibert. In those works, as mentioned by the reviewer, we swap some training samples with test samples. This will result in a new training dataset and since the decoder part also depends on the training dataset, this kind of shuffling yields bounds that involve, among other terms, one that accounts for the "decoder" or "prediction" complexity. This contrasts with our work here which seeks establishing bounds on the generalization that depend on the encoder complexity only (via the latent variable complexity), not the decoder complexity. Hence, in the type of symmetry that we choose we shuffle only the assigned latent variables (under label preservation) while keeping the training and test dataset as before. This is the reason for which in our definition of the needed symmetry for our results we require $\mathbf{Q}(U^{2n}_{\pi}|Y^{2n},X^{2n},W_e)$ to be invariant under all label-preserving permutations $\pi\colon [2n]\mapsto [2n]$. Here, we swap the latent variables but keep training and test datasets unchanged. Such analyses are in general more delicate.

   That being said, in this work, we consider maximal shuffling of latent variables $\mathbf{U}$ and $\mathbf{U}'$ for training and test datasets; that preserves the labeling (i.e. we shuffle two latent variables if they are related to samples that have the same label). Instead, as can be concluded naturally by the comment of the reviewer, we could only consider shuffling one latent variable in the training dataset with only one latent variable in the test dataset that has the same label as the training sample. Such an approach would give the bound established in [SZK23, Theorem 4] with order of intuitively $\sqrt{\text{MDL}(\mathbf{Q})/n}$. The maximal shuffling allows to improve the bound to order $\text{MDL}(\mathbf{Q})/n$ in some cases.

Questions

1. On the lossy compressibility and Equations (14) and (15): First, note that, by assumption, we have $\mathbb{E}\left[\text{gen}(S,W)-\text{gen}(S,\hat{W})\right] \leq \epsilon$, or equivalently $\mathbb{E}\left[\text{gen}(S,W)\right] \leq \mathbb{E}\left[\text{gen}(S,\hat{W})\right]+\epsilon$. Therefore, it is sufficient to upper bound $\mathbb{E}\left[\text{gen}(S,\hat{W})\right]$ for the "quantized model" $\hat{W}=(\hat{W}_e,W_d)$. This is easily done by invoking [SZK23, Theorem~4]; and it yields (14) and (15).

   Note that this simple trick has both practical and conceptual implications. Conceptually, it measures "lossy compressibility". To give an example, let's consider source compression for simplicity. While discrete random variables with a finite support set can be expressed (or "stored") using a finite number of bits, a "continuous" random variable requires an infinite number of bits for lossless description. However, if we allow the random variable to be recovered with some precision or distortion, we can express it using a finite number of bits, which is characterized by the rate-distortion function in the source coding literature. More technically, in our setup, while $\text{MDL}(\mathbf{Q})$ considered for $W$ may be infinite for deterministic encoders with continuous output support, $\text{MDL}(\mathbf{Q})$ considered for $\hat{W}$ is finite.

   As already discussed in Section 3.3, we have stated the lossy compressibility result by invoking [SZK23, Theorem 4] mainly for reasons of simplicity (and this has yielded a generalization bound of the order of $\sqrt{\text{MDL}(\mathbf{Q})/n}$). However, taking into account the suggestion of the reviewer, we have added in the revised version (see Appendix~A) a lossy compressibility result as well as a detailed proof of it. Observe that the new result is expressed in terms of the function $h_D$ and has a better dependence on the number of training samples $n$, i.e., $\text{MDL}(\mathbf{Q})/n$ instead of the $\sqrt{\text{MDL}(\mathbf{Q})/n}$.

2. On the intuition about the function $h_C$: First, let us give an intuition about $h_D(x_1,x_2)$. Fix some large integer $m$. Suppose that we are given $2m$ balls where $[m(x_1+x_2)]$ of them are red and the rest of them are blue -- $[\cdot]$ denotes the integer part. We take $m$ balls randomly and without replacement from these $2m$ balls. Let $P_m$ be the probability that $[m x_1]$ of the picked balls are red. Then, $h_D(x_1,x_2) = - \lim_{m\to \infty} \frac{1}{m} \log(P_m)$. In other words, $h_D(x_1,x_2)$ is the "asymptotic error exponent" of the event of having $[m x_1]$ red ball. To put this in the context of our work, the red balls can be thought of as samples with a loss value of one. Thus, in a simple scenario, if we have a superset containing samples from the training and test datasets that have a total error of $n(x_1+x_2)$, then $h_D(x_1,x_2)$ is the asymptotic error exponent of the probability that the empirical risk is equal to $x_1$, if we choose $m$ samples randomly from the superset as the training dataset.

   Now, regarding the question of the reviewer; assuming $x_1 \leq x_2$, we can re-write the $h_C$ function as h\_C(x\_1,x\_2,\epsilon) = \max\_{\epsilon'}\left\\{ h\_D(x\_1,x\_2)-h\_D(x\_1+\epsilon',x\_2-\epsilon')\right\\},

   where $\epsilon' \in [0,\min(\epsilon,(x_1-x_2)/2)]$. For a given $\epsilon'$, the term $h_D(x_1+\epsilon',x_2-\epsilon')$ captures the asymptotic error exponent of picking $m(x_1+\epsilon')$ red balls in the above scenario (from $2m$ balls with total $m(x_1+x_2)$ red balls). Hence, $h_C(x_1,x_2,\epsilon)$ can be seen as the maximum change of the asymptotic error exponent of the probability of picking $m(x_1+\epsilon')$ red balls (w.r.t. prob. of picking $mx_1$ red balls) in the range $\epsilon' \in [0,\min(\epsilon,(x_1-x_2)/2)]$. It is worth mentioning that whenever $x_2\leq 1/2$, then the maximum is achieved for $\epsilon'=\min(\epsilon,(x_1-x_2)/2)$.

3. On the selection of the value of parameter $M$: This depends on several factors, especially the dimension $d$ of the latent variables and the "sparsity" or number of "subpopulations" in the latent variables (which depends on the complexity of the data and the complexity of the encoder). Therefore, in general, it is not possible to know precisely an appropriate value of $M$ in advance and should be hence considered as a hyperparameter.

   As future work, it would be interesting to investigate how common dimensionality reduction and unsupervised clustering techniques, such as the t-SNE dimensionality reduction and clustering using a Gaussian mixture, could be used to obtain a good "initial guess" for the choice of $M$.

We thank the reviewer for the careful reading of the paper and the comments, which are relevant.

Regarding the difference with [1]: Indeed, as mentioned by the reviewer, the design of prior as a mixture of Gaussians is a significant addition, comparatively. However, we kindly attract the attention of the reviewer on that, in addition to this one, there are other important differences with the approach of [1], especially in the proof techniques of our main results which are required to get our tighter generalization bounds, with a better dependence on the number of training samples $n$ (namely, $\text{MDL}(\mathbf{Q})/n$ instead of the $\sqrt{\text{MDL}(\mathbf{Q})/n}$ of [1]). For instance, as we now elaborate after Theorem 1 and our response to Reviewer C35E (see the part "Regarding the proof techniques" of the rebuttal to that reviewer), unlike [1] we do not find another dataset having the same labels as the training dataset, as this introduces the penalty of order $\sqrt{\frac{C}{n}}$, and we analyze directly the terms (such as computing moment-generating functions) for $S$ and $S'$, which have potentially unequal numbers of samples per label. Secondly, for each label $c\in[C]$, we consider all samples in $S$ and $S'$ that have the label $c$. Denoting this set as $\mathcal{B}_c$, our proof then relies on the complete reshuffling of the latent variables corresponding to the samples in the set $\mathcal{B}_c$. Analyzing the complete reshuffling of the latent variables, especially taking into account the unequal number of samples per label, makes the analysis more involved (for example applying Donsker-Varadhan on $nf$ as defined in the proof of Theorem 1, rather than $nh_D$, and bounding the moment-generating function of $nf$ rather than $nh_D$). But, it is precisely these two steps that allow us to obtain bounds that behave as $\text{MDL}(\mathbf{Q})/n$ in some cases.

We thank the reviewer for the relevant comments and suggestions.

Weaknesses:

1. First, we emphasize that in our approach, the regularizer is used only during the training phase. Thus, there is no overhead in the inference phase. Moreover, similar to VIB, our regularizer depends on the dimension of the latent variable, rather than on the dimension of the model or the input data, which are often much larger. Moreover, our approach is rather easy to implement, as acknowledged for example by Reviewer ce5p.

   However, we agree with the reviewer that, similar to many ones that introduce/use regularizers, the gains of the approach come at the expense of a slight addition in the computational cost. Possible approaches towards reducing further that overhead could include: (i) using the regularizer only in the first $K$ epochs (which, usually, are the most critical ones for convergence) and (ii) first projecting the latent space onto a lower dimensional one and then applying the regularizer in that lower-dimensional space. We now elaborate more on this in the newly added "Future Directions" section in Appendix D.

2. Regarding the connection with the attention mechanism, we agree with the reviewer that further investigations along this direction are very interesting. In particular, in combination with point 1 above, we expect that when our approach is employed in the models that use the self-attention mechanism some of the computations performed therein could be used to update the GM prior and thus partially reduce the computational overhead.

3. Please note that our theoretical results of Theorem 1 and Theorem 2 are established for the classification setup and naturally the experiments are performed only in this setup (but on various datasets, as already noted by the reviewer). In fact, we agree that it would be interesting to extend both the theoretical and experimental results of this work to use in transfer learning and semi-supervised learning, which is beyond the scope of the current work. For example, in transfer learning, can the learned priors be used for the transfer learning task? Or in semi-supervised learning, can unlabeled data be used to find the GM prior more accurately? In the revised version, we have discussed such aspects in the newly added "Future Directions" section.

Questions:

1. We discussed some limitations (e.g. computational complexity and being limited to classification) and future directions (e.g. computational efficiency, considering our approach in the self-attention layer, and extending to transfer learning and semi-supervised learning setups) in the revised version.

2. As discussed in the paper, the Gaussian mixture is chosen for two reasons: (i) the Gaussian mixture distribution is known to possibly approximate any arbitrary distribution well enough, provided the number of mixture components is sufficiently large [NND+22]; and (ii) given distributions $\\{p_{i}\\}\_{i\in[N]}$, the distribution $q$ that minimizes $\sum\_{i\in[N]} D\_{KL}(p\_i\||q)$ is $q=\frac{1}{N}\sum\_{i\in[N]}p\_i$. Thus, if all $p\_{i}$ distributions are Gaussian, the minimizer is a Gaussian mixture.

   In many of the variational encoders, the posterior is indeed Gaussian, and thus it is optimal to consider GM. Moreover, in deterministic encoders, by using the lossy approach, we can consider the posterior to be a Gaussian with the mean being equal to the deterministic encoder output and the covariance being equal to $\varepsilon\mathrm{I}\_d$, where $\varepsilon$ is the lossy compressibility hyperparameter. Thus, in most situations, the choice of GM is a natural one.

We thank the reviewer for the interest in the work and relevant comments.

Regarding the related works: We have added a brief discussion of how our approach compares to previous related work.

Regarding the proof techniques: As suggested by the reviewer, we now add in the revised version, right after Theorem 1, a sketch of its proof with emphasis on the main elements of that proof. In what follows, we focus in particular on the reviewer's question about the difference between our proof techniques and those in [SZK23].

While, at a high level, the proof techniques used in our work and those of [SZK23] can be viewed as being both based on the idea of "shuffling" of the latent variables associated with the training and test samples that have the same label, ours are more involved and the analysis is deeper, comparatively. This explains why the resulting bounds are tighter than those of [SZK23]. In order to show this, for convenience, we briefly recall the main ideas of the proof technique of [SZK23]. In the approach of [SZK23], given the datasets $S$ and $S'$, the authors (implicitly) find another dataset $\tilde{S}$ such that: (i) the vectors of labels $\tilde{\mathbf{Y}}$ (associated to $\tilde{S}$) and $\mathbf{Y}$ (associated to $S$) are equal up to some suitable re-ordering, (ii) the found dataset $\tilde{S}$ and the ghost dataset $S'$ differ in no more than $\sqrt{\frac{C}{n}}$ samples (since $\mathbf{Y}$ and $\mathbf{Y}'$ are drawn independently, with some suitable reordering of $\mathbf{Y}'$ they differ in about $\sqrt{\frac{C}{n}}$ indices). This first proof step introduces a penalty of order $\sqrt{\frac{C}{n}}$ into the bound. Thus, with such an approach, obtaining bounds with better dependence on $n$ than $1/\sqrt{n}$ seems elusive. Next, assuming that $\tilde{Y}$ (after suitable reordering) is equal to $Y$, the rest of the proof uses arguments similar to those of the conditional mutual information (CMI) approach; and, for instance, it relies on shuffling the latent variables corresponding to the $i$th samples of the training and test datasets. This second proof step yields bounds of the order of $1/\sqrt{n}$.

Based on the above, it appears clearer that in order to obtain new bounds with better dependence on $n$ than those of [SZK23] one needs different, 'finer', proof techniques in both aforementioned steps. This is precisely what we do. First, we do not find another dataset having the same labels as the training dataset, as this introduces the penalty of order $\sqrt{\frac{C}{n}}$, and we analyze directly the terms (such as computing moment-generating functions) for $S$ and $S'$, which have potentially unequal numbers of samples per label. Secondly, for each label $c\in[C]$, we consider all samples in $S$ and $S'$ that have the label $c$. Denote this set by $\mathcal{B}_c$. Our proof then relies on the complete reshuffling of the latent variables corresponding to the samples in the set $\mathcal{B}_c$. Analyzing the complete reshuffling of the latent variables, especially taking into account the unequal number of samples per label, makes the analysis more involved (for example applying Donsker-Varadhan on $nf$ as defined in the proof of Theorem 1, rather than $nh_D$, and bounding the moment-generating function of $nf$ rather than $nh_D$). But, it is precisely these two steps that allow us to obtain bounds that behave as $\text{MDL}(\mathbf{Q})/n$ in some cases.

Regarding the geometrical compression: In the revised version, we have added a short explanation in the introduction section of what is meant by "Geometric compression", a term already used in related works such as [AG19, GK19]. More precisely, "Geometric compression occurs when latent vectors are designed so as to concentrate around a limited number of representatives which form centroid vectors of associated clusters [AG19, GK19]. In such settings, inputs can be mapped to the centroids of the clusters that are closest to their associated latent vectors (i.e., lossy compression); and this yields non-vacuous bounds at the expense of only a small (distortion) penalty. The benefit of this lossy compression approach can be appreciated when opposed to classic MI-based bounds [VPV18, KDJH23] which are known to be vacuous when the latent vectors are deterministic functions of the inputs."

Regarding the experiments: By joint training, we mean joint training of the encoder and decoder parts of the model. This was stated as such in order to emphasize the dependence of both encoder and decoder parts on the dataset. In the revised version, we removed this to avoid any confusion.

The considered loss function is the cross entropy function, as often considered in classification tasks. This is now stated more clearly in the documents. Also, further details about the experiments are added in Appendix E.

We have now changed it. Thank you for bringing our attention on that.