Generalization Guarantees for Representation Learning via Data-Dependent Gaussian Mixture Priors
Authors: Milad Sefidgaran
Abdellatif Zaidi
Piotr Krasnowski
Venue: ICLR 2025 Spotlight (2025)
Keywords: Representation learning algorithm
Gaussian-Mixture
regularizer
rate-disotortion
Primary area: learning theory
TL;DR: We derive generalization bounds for the representation learning algorithms and, inspired by the bounds, propose a regularizer with data-dependent Gaussian mixture priors.

Abstract: We establish in-expectation and tail bounds on the generalization error of representation learning type algorithms. The bounds are in terms of the relative entropy between the distribution of the representations extracted from the training and "test'' datasets and a data-dependent symmetric prior, i.e., the Minimum Description Length (MDL) of the latent variables for the training and test datasets. Our bounds are shown to reflect the "structure" and "simplicity'' of the encoder and significantly improve upon the few existing ones for the studied model. We then use our in-expectation bound to devise a suitable data-dependent regularizer; and we investigate thoroughly the important question of the selection of the prior. We propose a systematic approach to simultaneously learning a data-dependent Gaussian mixture prior and using it as a regularizer. Interestingly, we show that a weighted attention mechanism emerges naturally in this procedure. Our experiments show that our approach outperforms the now popular Variational Information Bottleneck (VIB) method as well as the recent Category-Dependent VIB (CDVIB).

Published as a conference paper at ICLR 2025
GENERALIZATION GUARANTEES FOR REPRESENTA- TION LEARNING VIA MIXTURE PRIORS
GAUSSIAN
Milad Sefidgaran', 1, Abdellatif Zaidi 1 Piotr Krasnowski!
Paris Research Center, Huawei Technologies France Université Gustave Eiffel, France (mi i lad.aad.. sef sefidgaren piotr.g krasnowski @huawei. com, abdellatif. wai.iiiiii-e.iee..l...
ABSTRACT
We establish in-expectation and tail bounds on the generalization error of repre- sentation learning type algorithms. The bounds are in terms of the relative entropy between the distribution of the representations extracted from the training and "test" datasets and a data-dependent symmetric prior, i.e., the Minimum Description Length (MDL) of the latent variables for the training and test datasets. Our bounds are shown to reflect the "structure" and "simplicity" of the encoder and signifi- cantly improve upon the few existing ones for the studied model. We then use our in-expectation bound to devise suitable data-dependent regularizer; and we inves- tigate thoroughly the important question of the selection of the prior. We propose a systematic approach to simultaneously learning data dete-deenndent Gaussian mixture prior and using it as a regularizer. Interestingly. we show that a weighted attention mechanism emerges naturally in this procedure. Our experiments show that our approach outperforms the now popular Variational Information Bottleneck (VIB) method as well the recent Bategory-Depeendee VIB (CDVIB).
INTRODUCTION
One major problem in learning theory pertains to how to guarantee that a statistical learning algorithm performs on new, unseen data as well as on the used training data, i.c., it has good generaliza-
tion properties. This key question, which has roots in various scientific disciplines, has been
studied using seemingly unrelated approaches, including compression-basee
(LW86) BEHW87
AGNZI8 BL03
ISA++20 HITW2I BSE=+11 [K HKS19) BHMZ20 HK21
HKSW20 20
CK22 SGRS22 SZ24), information-theoreti (RZ16, XR17 SZ20, EGI20 BZV20 HDMR21
NDHR21 ABT 21 HRVSG2T] ZTT222 [L22]] HDD22,, PAC- Bayes (Se022 LCOT, Cat03] Mau04
GLLM09 TSI3: BGLRI6 TIWSI7 DRI7 IBSSIg, RKSST20 NDR20 NHD 20 VGHM21).
and intrinsic dimension-bact ($SDE20) BLGS2I HSKM22 LW$22) approaches.
In practice, common approach advocates the usage of two-part, encoder-decoder, model, often referred to as representation learning. In this approach, the encoder part of the model shoots for extracting "minimal" representation of the input (i.e., small generalization error), whereas the decoder part shoots for small empirical risk. One popular approach is the information bottleneck (IB), which was first introduced in (TPBO0) and then extended in various ways SSTTO) AFDM17 EAZI8 KTW19) Fis20;; RRGTS20) KASK22). The IB principle is mainly based on the assumption that Shannon's mutual information (MI) between the input and the representation is good indicator of the generalization error. However, this assumed relationship has been refuted in several works
KTVKI8 RG119 AG19 GK DKSV20 LLS+23) SZKK3).. As shown in these works, the few existing theoretical MI-based generalization bounds (e.g., (VPV18) KDJH23) become vacuous in most reasonable setups. Also, in practice, no consistent relation between the generalization error and the MI has been observed experimentally so far. Rather, recent works (BL03) GK19; SZK23) have shown that the generalization error of representation learning algorithms is related to the minimumm description length (MDL) of the latent variable and the so-called geometric compression. 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

---

Published as a conference paper at ICLR 2025
X PUIX,W
U
PyJu,wa Y
Encoder
Decoder
'car'
Figure 1: Studied representation learning setup.
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 classic MI-based bounds (VPV18) KDJH23) which are known to be vacuous when the latent vectors are deterministic functions of the inputs.
In this work, we study the problem of representation learning depicted in Fig. from a generalization error perspective. Then, use the obtained generalization bound to design and discuss various choices of generalization-inpaa inspired regularizers using data-dependent Gaussian mixture priors. To the best knowledge of the authors, generalization error bounds that account suitably for the encoder- decoder structure of the representation learning problem of Fig. are very scarce; and, in fact, with the exception of (SZK23), non-vacuous bounds for these settings have been reported far.
Contributions: Our main contributions in this work are summarized as follows.
We establish in-expectation and tail bounds the generalization error of the representation learning algorithms. Our bounds are expressed in terms of the relative entropy between the distribution of the representations extracted from training and "test" datasets and a data-dependent symmetric prior Q, i.e. the Minimum Description Length (MDL(Q)) of the latent variables for training and test datasets- (Bounds that depend on MDL Q) are arguably better bounds because they capture the structure and simplicity of the encoders in sharp contrast with IB-based approaches (BLO3)). However, our bounds are shown be possibly tighter than those of (SZK23). For instance, while the bounds of (SZK23) are of the order of /MDL((Q) /n. where n designates the size of the used training dataset, ours is approximately of the order of MDL(Q)/r nr. for the realizable setup.
We propose systematic approach to finding suitable "data-dependent" prior that we then use to construct judiciously a regularizer during training (based on our newly established bounds). Specifically, first, observe that if the latent variables are generated according to Gaussian distribution, then the prior Q that minimizes the empirical MDL(() term is a Gaussian mix- ture distribution. Then, using this and the known fact that Gaussian mixture distributions can approximate sufficiently well any arbitrary distribution when the number of mixture components is large enough (DH83, GBCI6 NND 22), we propose two methods for simultaneously finding a Gaussian mixture prior and using it as a regularizer along the optimization iterations. The methods are coined 'lossless Gaussian mixture prior" and "lossy Gaussian mixture prior", respectively. In essence, the procedure consists of finding the underlying "structure" of the latent variables in the form of a Gaussian mixture prior; and, simultancously. steers the latent variables to best fit with this found structure. Interestingly, in the lossy version of the approach, which is shown to generally yield better performance, the components of the Gaussian mixture are updated using a mechanism that is similar the self-attention mechanism. In particular, the components updated according to the extent they each "attend" to the latent variables statistically.
We validate the advantages of our generalization aware regularizer in practice through experiments using various datasets (CIFARIO, CIFAR100, INTEL. and USPS) and encoder architectures (CNN4 and ResNet1 8). In particular, we show that our approach outperforms the popular VIB of (AFDM17) and the recent Detegory-Deeendent VIB of (SZK23). The reader is referred to Scction 5 and Appendix B for details on the datasets, models, and experiments.
We emphasize once more that our approach here, which measures complexity using MDL of the involved latent variables, has two appealing features: (i) it yields generalization bounds that only depend on the encoder part of representation type statistical learning algorithms, and (ii) the employed lossy compression enables the yielded bounds to only take finite values, i.e., not vacuous, in reasonable setups, by opposition to the MI bounds of (VPVI8) KDJH233. The described approach and results must be contrasted with classes of prior-art bounds that measure complexity differently. The first class
2

---

Published as a conference paper at ICLR 2025
of bounds involves the complexity of the hypothesis (model) space and includes, e.g., MI-based, PAC- Bayes, and some of the compression- based bounds (e.g. (AGNZI8)). Such bounds mostly involve "data- "ddaaa-iddependemn" priors on the model; and seldom use "data-dependent" priors (DR 8 PORSTS21 see (Alq21] Section 3.3) for a detailed review. Generalization bounds that use model complexity do not seem be amenable to using them for regularization since, in practice, one has only single instance of the posterior. The second class of bounds are intrinsic dimension-based bounds that measure the complexity of the model along the optimization trajectories. Although in this approach multiple instances of the posterior are available, measuring the trajectory complexity of large models not practical. The third class of bounds uses prediction complexity such as with f-CMI (HRVSG21: (HD22)- see also the related (BL03) SZK23). In such bounds, typically the complexity appears both the loss function and the regularizer; and this is generally not reasonable in practice.
Notations. We denote the random variables and their realizations by upper and lower case letters and use Calligraphy fonts to refer to their support set e.g., X. and X. The distribution of X is denoted by Px which for simplicity, is assumed to be a probability mass function for a random variable with discrete support set and to be probability density function otherwise. With this assumption, the Kull- Leib-ee (KL) between two distributions P and Q is defined as DxL (Pl Q) = [pllo((//1] if P & Q, and oo otherwise. Lastly, we denote the set (1,..., n1,, n No, by [n].
PROBLEM SETUP
We consider C-class classification setup, as described below.
Data. We assume that the input data Z, which take value according to an unknown distribution 4., is composed of two parts Z =XX, Y), where (i) X represents the feature of the input data, taking values in the feature space X, and (ii) Y e ) represents the label ranging from I to C, i.e. y' = [C]. We denote the underlying distribution of X and Y by px and jy, respectively, and their joint distribution by == MxlY py := μXPYЛXX
Training dataset. To learn model, we assume the availability of training dataset S (Z,...., } ~ 1 =: Ps. composed of n i.i.d samples Z; (X,, Y.) of the input data. In our analysis. we often use test dataset (known also as ghost dataset (SZ20)) S (Zi..... Z,} u =: Pss where Z = (X,, Y?). To simplify the notation, we denote the features and labels of
6on S and 5" by X == X px Y = Y ny respectively.
Rn and Y' == ym ~ by &n X' = ~ 1'
80n
Encoder-decoder model. We assume that the model (hypothesis) is composed of two parts: an encoder and a decoder part. The encoder w. We takes as input the feature and generates as output the representation or the latent variable U U. possibly stochastically. For simplicity, we assume that 21 = Rd for some d E N The decoder wa E Wa takes U as input and outputs an estimate ý of the true label Y. The overall model is denoted by w := (We, E W = W. x Wa. The setup is shown in Fig.
Learning algorithm. We consider a general stochastic learning framework in which the learning algorithm A: -n - W has access to a training dataset S and uses it to choose a model (or hypothesis) A(S)= = W E W, where W = (We, Wa). The distribution induced by the learning algorithm A is denoted by Pwis= Pw... W. Also, the joint distribution of (S, W) is denoted by Ps.w and the marginal distribution of W under this distribution is denoted by Pw: Furthermore, we denote the induced conditional distribution of the latent variable U given the encoder and the input by Pupx,w.-
Finally, we denote the conditional distribution of the model's prediction Y, conditioned on the decoder and the latent variable, by Pyju.w. It is easy see that Pyxx.w EUi x, [ [..].... Lastly and as a general rule, we use the following shorthand notation
PU.
zejm) (PulIX,W. Pujx(w.)
(I)
Similar notation is used to shorten products of distributions, egg, Puix,w, and P
'We, however, make an exception for the input data, whose distribution is denoted by n,, as it is common in thcoretical papers, e.g. (XR17 BZV20 [222.
3

---

Published as a conference paper at ICLR 2025
Risks. The quality of model w is assessed by the below 0-1 loss function &: x W {0,1):
{(2, w) = Ry Prle. (wwpy) BU-Po Poia.w. Py-Pr10..,
In learning theory, the ultimate goal is to find model that minimizes the population risk, defined as [(w) = Ez-nl w)]. However, since the underlying distribution şi is unknown, only the empiri- cal risk, defined as [(s, w) >= 1 Lie(m) {(2;, w), is accessible and can be minimized. Therefore, a central question in learning theory and this paper to control the difference between these two risks, known generalization error:
gen w) = [(w) - L(s, w).
(3)
In our results, for simplicity, we also use the following shorthand notations:
L(y.9) = 1 Eiering (0ttts) 2(y',))= = 1 -ie(m)) 1 (06+1)
(4)
Note that
[(8, w) By P fly, 8)). ((s',w) B po [(ys' [())
(5)
Symmetric prior. Our results are stated in terms of the KL-divergence between a posterior (e.g.. PU,U'XXX'',, and prior Q that needs to satisfy some symmetry property.
Definition (Symmetric prior). conditional prior 2 y22 _22 We) said be symmetric if QUU?n) Y 2n X22,, ,We) is invariant under all permutations :: [2n] - [2n] for which vi: Y;=Y=()--
GENERALIZATION BOUNDS FOR REPRESENTATION LEARNING ALGORITHMS
In this section, we establish novel in-expectation and tail bounds on the generalization error of representation learning algorithms for the setup of Fig.
3.1 IN-EXPECTATION BOUND
Define the function hp: [0, [0, [0, 2] as
+ x2)
hD(x1, = 2hv
h2(x1)))))))))))))(<)))) ho(22),
where hz(2) - (o), - (1 - x) log, [1 x) is the binary Shannon entropy function. It is
easy that (2 22) 2 equals the Jensen Saasen-Shannon divergence between two binary Bernoulli
distributions with parameters in E [0, 1] and 22 E [0, 1]. Also, let the function lic: [0. 1] x [0, 1] x
R+ R+ be defined
(21 = max h1l + +e) (2(2122) + ho(En - -ho(Tzre2)) d
(6)
where F142 = min(x1, x2), Z1v2 = max(#1, x2), and the maximization in
6 is over
de 0. min (c, 31v2 .2 )
(7)
Hereafter we sometimes use the handy notation
hy 8 (E) = ho (C(, s')),
(8)
Now, state in-expectation generalization bound for representation learning algorithms. Theorem 1. Consider C-class classification problem and learning algorithm A: n W that induces the joint distribution (S' S, W. U, U² U, Y,Y) Ps Ps,w PU.u X.X.W. w. Py YU.U...
4

---

Published as a conference paper at ICLR 2025
Then, for any symmetric conditional distribution QUU.U Y, X,X' We) and 2 10, we have
B.s. v hn (2(Y, YY,, (x) MDL(Q) + log(n) E 8 S /y
<
(flor - fiix--xxli)]
(9)
where by and py: are empirical distributions Yf and Y', respectively, and
MDL(Q) == D...s.....D.. (PU.UKUPPUUU..X. [Q] [Q]]-
(10)
The proof of Theorem which appears in Appendix G. consists of two main proof steps,
moment generation function (MGF). Specifically, we use the Donsker-Varadhan's 's lemma (DV75 Lemma 2.1) to change the distribution of the latent variables from Pu, PU W. to Q. This change in measure results in penalty term equal to MDL Q Let f be given by n times the difference of hD and the term on the right- sidd-e (HS())) We apply the Donsker- -arkddann change of measure on the function f. in sharp contrast with related proofs in MI-based bounds literature (XR17 SZ20 Alq2I). The second step consists of bounding the MGF of naf. For every label c E [C], let B, denote the set of those samples of S and S' that have label c. By construction, any arbitrary reshuffling of the latent variables associated with the samples in the set B& preserves the labels. In addition, such reshuffling does not change the value of the symmetric prior Q. The rest of the proof consists of judiciously bounding the MGF of nff under the uniform distribution induced by such reshuffles.
change of measure argument followed by the computation
(RHS) of (9) i.e., f = n(hp
It is easy to see that the left hand side (LHS) of (9) is related to the expected generalization error. For instance, since by (SZK23 Lemma 1) the function h hp(a1, 22) is convex in both arguments, hn(21, e F1, and hn(2 11, x2)
> (x) x2) 2 for I1, 22 [0, 1J. one has that
Es,w Bs.w[gen(s (S, w)] Ess. y [bn Y), }))]
and
Es,w [gen(S, w)j² (b L(Y, Y))]
for the "realizable" and "unrealizable" cases, respectively.
Several remarks now in order. First, note that the generalization gap bound of Theorem does not depend on the classification head; it only depends on the encoder part! In particular, this offers theoretical justification of the intuition that in representation-tyee neural architectures the main goal of the encoder part is to seek a good generalization capability whereas the main goal of the decoder part is seek minimize the empirical risk. Also, it allows the design of regularizers that depend only on the encoder, namely the complexity of the latent variables, as we will elaborate on thoroughly in the next section. (2) The dominant term of the RHS of ( is MDL(Q))rn /n. This can be seen by noticing that the total variation term pp- px l1 is of the order /C/n n as shown in (BK12 Theorem 2); and, hence, the residual
Bemp_diff == E Ey NYY (1 1-11--xxll)]
(11)
is small for large n (see below for additional numerical justification of this statement). (3) The term MDL(Q), as given by (10). expresses the average (w.r.t. data and training stochasticity) of KL-divergence terms of the form DKL (PIIQ) where P is the distribution of the representation in the training samples n and the test samples n conditioned on the features of the 2n examples for given encoder, while Q is a fixed symmetric prior distribution for representations given 2n samples for the given encoder. As stated in Definition] T Q is symmetric for any permutation n: and, in a sense, this means that Q induces distribution on (U, U') conditionally given (Y,Y', Y', X, X', W.) that is invariant under all permutations that preserve the labels of training and ghost samples. (4) The minimum description length of the representations arguably reflects the encoder's "structure" and "simplicity"" (SZK23). In contrast, mutual information (MI) type bounds and regularizers, used, e.g., in the now popular IB method, are known to fall short of doing so [G LL)+ In fact, as mentioned in these works, most existing theoretical basba generalization bounds (e.g., (VPV18 KDIH23) ) become vacuous in reasonable setups. In addition, no consistent
5

---

Published as a conference paper at ICLR 2025
Trasenn 1. -
0.6
+. 0.0050
Bounds
0.4
Thearn
0.0025
0.2
0.0008.0
0.2
0.4
00.0
0.05
0.10
0.15
0.20
0.25
Generalization error
MDL/
Figure 2: Values of hc ((yy ((y f):e) for various values the generalization error for the CIFARIO 10 dataset.
Figure Comparison of the generalization bounds of Theorem | (for various values of C(S. W)) and (SZK23 Theorem 4) for the Cl- FAR10 dataset.
relation between the generalization error and MI has been reported experimentally so far. Therefore, MDL is better indicator of the generalization error than the mutual information used in the IB principle.
As we already mentioned, the total variation jey - pY- is of the order v CC/n (BK12 Theorem 2); and for this reason, the second term on the RHS (9) negligible in practice. Figure shows the values of the term inside the expectation of Bemp_iin as given by (TI) for the CIFART 10 dataset for various values the generalization error. The values are obtained for empirical risk 0.05 and
| set to be of the order CC/n. As it is visible from the figure, the term inside the expectation of Bemp_diff is the order of magnitude smaller than the generalization error. This illustrates that even for settings with moderate dataset size such as CIFAR, the generalization bound of Theorem is mainly dominated by MDL (Q)nn.
As stated in the Introduction section, generalization bounds for the representation learning setup of Fig. [ are rather scarce; and, to the best of our knowledge, the only non-vacuous existing in- expectation bound was provided recently in (SZK23) Theorem 4). This bound states that
2MDL( Q) + +
Es w www.genn W)]
(12)
where C' is the number of classes.
Investigating 9 and 12 it is easy that, order-wise, while the bound of (ZZK23 Theorem 4)
evolves O MDL Q)/n our bound of Theorem is tighter comparatively and it evolves approximately as O(MDL(Q)) Q)/n) for realizable setups with large n (i.e., for most settings in practice).
ii Figure the evolution of both bounds function of MDL(Q)/n /n for the CIFAR10 dataset and for different values of the empirical risk. important to emphasize that, in doing so, we account for the contribution of all terms of the RHS of (9) including the residual Bemp diff which is then not neglected. As is clearly visible from the figure, our bound of Theorem [ is tighter comparatively. Also, the advantage over (12) becomes larger for smaller values of the empirical risk and larger values of MDL(Q)/n.
3.2 TAIL BOUND
The following theorem provides a probability tail bound on the generalization error of the representa- tion learning setup of Fig.
Theorem 2. Consider the setup of Theorem [add consider some symmetric conditional distribution QUU,U'Y Y'. X. X', W.). Then, for any ò 2 0 and for 10, with probability at least 1 over choices of (S, St W), it holds that
DEL (P log(n/8)
hp (ISS,,)

---

Published as a conference paper at ICLR 2025
+ Ey X 1y X
(iim--x.l.))
(13)
where py and py, are empirical distributions of Y and Y', respectively.
The proof of Theorem appears in Appendix G.2
3.3 LOSSY GENERALIZATION BOUNDS
The bounds of the previous section can be regarded as lossless versions of ones that more general, and which we refer to as lossy bounds. The lossy bounds are rather easy extensions of the corresponding lossless versions, but they have the advantage of being guaranteed stay non- vacuous even when the encoder is set to be deterministic. Also, such bounds are useful to explain the empirically observed geometrical compression phenomenon (Gei2T). For comparison, MI-based bounds, such as Xu-Raginsky (XR17) are known to suffer both shortcomings (HRGT+23 Liv23). The aforementioned shortcomings have been shown that can be addressed using the lossy approach (SGRS22 SZ24). For the sake of brevity, in the rest of this section we only illustrate how the bound 112 can be extended to corresponding lossy one. Let W. We be any quantized model defined by P w. that satisfy the distortion criterion Eps,w w Pwa.s
gen W) gen(S, w)]
where W- (We, Wa). Then, we get
2 2MDL(()) C++ +,
Bs,w Bs.w[gen(SS W)]
(14)
n
where now MDL(Q) considered for the quantized encoder, i.e.
MDL(Q) = E s.st w. DKL (Pu.uIX.X wr http:xxxxxxw///I,UUU |Q
(15)
REGULARIZATION USING DATA- DEPENEETN GAUSSIAN MIXTURE PRIORS
Theorems D and 2 essentially mean that if for a given learning algorithm the minimum description length MDL(Q) small, then the algorithm is guaranteed generalize well. Hence, it is natural to use the term MDL(Q)) Q) as suitable regularizer. The question the choice of the prior Q is pivotal for this. In this section, we propose effective method to simultaneously find date-depe Q and use it to build a suitable regularizer term along the optimization iterations.
We assume that for given input the encoder outputs the mean Mr E R and standard deviation O₂ E IRd. Also, we assume that the latent variable is distributed according to a multivariate Gaussian distribution with diagonal covariance matrix, iee. U ~ TNN(uz, diag( o2)) where diag(o?)
denotes d x d diagonal matrix with diagonal elements o2 With this assumption, have
PU.UIXXXWW.
Oiejm] (M dizg(eiag(o (Z)) (mzz, ding(o?)))
In our approach, we model the prior Q as a suitable Gaussian mixture, with the mixture coefficients chosen judiciously in manner that is training-datt dependent and along the optimization iterations. The rationale for this choice two-fold: (i) The Gaussian mixture distribution is known to possibly approximate well enough any arbitrary distribution provided that the number of mixture components sufficiently large ( GBC16) (see also (NNL* 22 Theorem I)); and (ii) given distributions (pe)ee[M] the distribution q that minimizes Eie[N] DKL (P:ll) is 9 N Eme[N] Pt. Thus, if all distributions pt Gaussian, the minimizer is Gaussian mixture.
Let, for c E [C], Qe denote the data-dependent Gaussian mixture prior Qe for label c. Also, is easy to see that this prior satisfies the
let Q(U, U' S", W.) = Ihe (U,)) symmetry property of Definition In what follows, we explain how the priors {Qe) are chosen and updated along the optimization iterations. As it will become clearer, our method is somewhat reminiscent of the expectation-maximizatoon (EM) algorithm for finding Gaussian mixture priors that maximize the log- likelhiooh, but with notable major differences: (i) our case the prior must be learned along the optimization iterations with the underlying distribution of the latent variables
7

---

Published as a conference paper at ICLR 2025
possibly changing at every iteration. (ii) The Gaussian mixture prior is intended to be used in regularizer term, not to maximize the log- log-likelhhoo;; and, hence, the approach must be adapted accordingly. (iii) Unlike the usual scenario where the goal is to find appropriate Gaussian mixture given a set of points, here we are given a set of distributions ie., N (Ja,, diag((2)) that generate such points. (iv) The found prior must satisfy (at least partially certain "symmetry" properties.
4.1 LOSSLESS GAUSSIAN MIXTURE PRIOR
For each label e [C], we let the prior Qe to be defined
Qc = - Ocemm Qe,m.
(16)
over Rd where aemm E [o, 1]. Eme(M]) Oc.m 1 for each c e [C], and where {Qamlem are multivariate Gaussian distributions with diagonal covariance matrix:
Qc.m= (pe,m diag diag(oe..m) (m).)).
[M [C].
With the above prior choice, the regularizer term simplifies as EiE DKL (Ru(Pu,Ix.w. Qx,). How- ever, since the KL-divergence between a Gaussian and a Gaussian mixture distributions does not have closed-fome expression, we estimate it using a slightly adapted method from (H007). Our estimate an average of the upper and lower bounds of the KL-divergence, denoted as Dvar and D prod Please refer Appendix [ for more details on this estimation. For better readability, present the approximation of the KT-divergence by its upper bound Dvar in the main part of this paper and we refer the reader to Appendix C for the approach using (Deur /2
Finally, similar to (AFDMI7: SZKK3)), we consider only the part of the upper bound MDL(()) corresponding to the training dataset S. simply because the test dataset S' is not available during the training phase. With this assumption and for a mini-batch B = {ā1,..., 26} S S, the regularizer term is equal to
Regularizer( == DKL (PUg)Xn. /
(17)
where the indices B indicate the restriction to the set B. For better exposition, we will drop the notation dependence on B in the rest of this section. Now, we are ready explain how the Gaussian mixtures are initialized, updated, and used as a regularizer simultaneously and along the optimization iterations. what follows, the superscript (t) denotes the optimization iteration E N
Initialization. First, we initialize the priors as Q! by initializing dc,m and the parameters pea (0) Je,m of the components Q& for ce [C], m E [M]] similar to the method of initializing the centers in k-means++ (Art77). The reader is referred to Appendix C.1 for further details.
Update of the priors. Let the mini-batch picked at iteration / be ( (=(")
2:. =bb (()). By dropping
the dependence on (t) for better readability, the regularizer 17 at iteration (t), can be written as
Regularizer Q) EEE DEL (Prrlenwe
Em
a Me
1a. (U:))
Q(t)
(e)
Σ Σ
(DKL (DKL(PU |Q(0) mi.m
(U.)) - log (0f0m//amm) (18)
[e[b]
me [M]]
where the last step holds for any choices of 7s.m > such that Σ Eme[M] [M] 7tm for every [b].
To see why the step (a) holds, we refer the reader to Appendix F to see how the variational bound Dear is derived.
Now, to update the components of the priors, first (similar to 'E'-step) note that the coefficients 7i that minimizes the above upper bound are equal to
anymm DEL (Pogtreime m.
7i.m
(1) Em' m.
DKI. (Pogleqiee (0()) (())
i M.me m. e [M].
(19)
WWilee the bounds of Theorems and 2 require the prior Q to satisfy the exact symmetry of Definition i it can be shown that these bounds still hold (with small penalty) if such exact symmetry requirement is relaxed partially. The reader is referred to Appendix B where formal results and their proofs are provided for the case of "almost symmetric" priors.
8

---

Published as a conference paper at ICLR 2025
Let Yi.c. Yym if Ys and Niccm 0 otherwise. Next, (similar Mf-tep) treat Ni as constants, and find the parameters H..m= c,m= a that minimizes the upper bound (18). by simply taking the partial derivatives and equating them to zero. Simple calculations show that the closed-form solutions are
Me.m
Ziepy Bo
C.m.i
Σ
Tiem
(Cais + (Maxn (t) m)mm))))
be,m
be.m
Ocmm
=be,m/ /be
De,m
E
Lajoa 7i,
be E bem--
(20)
where E [d] denotes the index of the coordinate in IRd and o (m a) Finally, to
reduce the dependence of the prior on the dataset and to partially preserve the symmetry property, let
1)++) =(1 - 1/1 mp)pe "/ m + +31-11)
(t+1)2 a.m
2 (1 (1-11))0") 12 c.m. 1220 32 (1+1)
Pemm (I+1) =(1 - 1/3 T3)) c.m. + +730 c.m. a c,m
(t)
(21)
(++11 where 7h, f2. 7s e [0, 1] are some fixed coefficients and 3,+ e [2], are iid. multivariate
Gaussian random variables distributed as (N( c/ Ia) Here 0₂ = (O..... 0) R and c)
cft+1)
Rt are some fixed constants.
Regularizer. Finally, using
[9). the upper bound (88 that we use as regularizer can be simplified
Σ ke[ log Emee[My
DRy (Pucinewe 10( k.mm)
(22)
4.2 LOSSY GAUSSIAN MIXTURE PRIOR
The lossy case is explained in Appendix C.3 when the KL-divergence estimate (Dpost + D considered. Similar to Section 4.1 it can be shown that if only Dvar considered for the KI.- divergence estimate, then the regularizer term becomes equal to
Σ [Elu] log Emetar Be.m
DICL, (Postenieel194 1Q (0))..)))
(23)
where Los, (Pu 110m is defined as
DKL N (N(Wun
Vdl Hos
Ld / Homm va (1 + +DKI (N diag(o"+ (N diag(o?.. c. + +e))
2
(24)
where E = (e....,,,, E) E Rd and e R is fixed hyperparameter.
Furthermore the components are updated according (1I. where N.c. P and a* defined as before, but o* c.m be.m Eiely 7h.mm. 2 and Y. is equal to
DKL.LOW ( (2)))
o(t) Po.me
(Pap (1) vil
me
Ti
Emve[My
(t) m
DKL,doney (Piglanee (e)m!
Em'e(Ar ((t)) w.m.
Ja
where ((()
m Ja
(lggl m /a ) t 2 (2a(1) mi In cases where the
agg.m
means of the components are normalized and the variances are fixed, (l0) (e)
The parameters Tt.m measure the contribution of the component m in Qy, generating the latent variable U,. One can observe similarity between how these parameters are chosen in our approach and the attention mechanism, with the difference that here we are considering a weighted version of this mechanism, and without key and query matrices since do not consider projections other spaces. Intuitively, to measure the contribution of each component, we measure how much that component "attend" to U..

---

Published as a conference paper at ICLR 2025
VIB
0.68
CDVIB
GM-MDL
Accuracy
0.66
no reg.
0.64
0.62
0.60
1e-5
1e-4
1e-3
1e-2
1e-1
Regularization parameter B
Figure 4: Test performance of the CNN-based encoder trained on CIFAR10 using standard VIB (AFDM17) regularization, Category-dependen VIB (CDVIB) (SZK23) regularization, and our proposed Gaussian Mixture MDL (GM-MDL) regularization.
5 EXPERIMENTS
In this section, present the results of our simulations. The reader is referred to Appendix additional details, including used datasets, models, and training hyperparameters.
E for
For the experiments, considered the lossy regularizer approach with Gaussian mixture prior and the KL-divergence estimate of (Dpuut + v/ as detailed in Appendix C.3 this section, we refer regularizer Gaussian mixture MDL (GM-MDL). To verify the practical benefits of the introduced regularizer, we conducted several experiments considering different datasets and encoder architectures as summarized below and detailed in Appendix
Datasets: CIFAR10, CIFAR100, INTEL, and USPS image classification,
Encoder architectures: CNN4 and ResNet18.
To compare our approach with the previous literature, in addition to the no-regularizer case, we also considered the Variational Information Bottleneck (VIB) of (AFDMI7) and the Category-dependen VIB (CDVIB) of (SZK23).
The results presented in Fig. and Table [ show the practical advantages of proposed approach. All experiments are run independently 5 times and the reported values and plots are the average over 5 runs. In Fig 4 plotted the performance of different regularizers function of the trade-off regularization parameter B. Table I we reported the best achieved average test accuracy for each regularizer.
Table 1:
Test performance of representation learning models with different encoder architectures, and trained on selected datasets using VIB (AFDMI7), Category-dependent VIB (CDVIB) (SZK23), and our proposed Gaussian Mixture MDL (GM-MDL).
Encoder
Dataset
no reg.
VIB
CDVIB
GM-MDL
CNN4
CIFARI0
0.612
0.626
0.649
0.681
2
CNN4
USPS
0.948
0.952
0.955
0.963
3
CNN4
INTEL
0.756
0.759
0.763
0.776
4
ResNet18
CIFAR10
0.824
0.829
0.835
0.848
5
ResNet18
CIFARIO0 100
0.454
0.458
0.463
0.497
10

---

Published as a conference paper at ICLR 2025
REFERENCES
[ABT*21] Gholamali Aminian, Yuheng Bu, Laura Toni, Miguel Rodrigues, and Gregory Worell. An exact characterization of the generalization error for the gibbs algorithm. Advances in Neural Information Processing Systems, 34:8106-81 1100-11188 8118, 2021.
[AFDM17] Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep varia- tional information bottleneck. In International Conference on Learning Representa- tions, 2017.
[AG19] Rana Ali Amjad and Bernhard C Geiger. Learning representations for neural network- based classification using the information bottleneck principle. IEEE transactions pattern analysis and machine intelligence, 42(9):2225-2233, 2019.
[AGNZ18] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. In International Conference on Machine Learning, pages 254-263. PMLR, 2018.
[Alq21] Pierre Alquier. User-friendly introduction to pac-bayes bounds. arXiv preprint arXiv:2110.11216. 2021.
[ARS17] Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep neural networks. arXiv preprint arXiv:/711.08856 2017.
[Art0]] David Arthur. K-means++: The advantages if careful seeding. In Proc. Eighteenth Annual ACM-SIAM Symposium on Discrete Algorithms, 2007, pages 1027-1035, 2007.
[BEHW87] Anselm Blumer, Andrzej Ehrenfeucht, David Haussler, and Manfred K Warmuth. Occam's 's razor. Information processing letters, 24(6):377-380, 1987.
[BGLR16] Luc Bégin, Pascal Germain, François Laviolette, and Jean-Francis Roy. Pac-bayesian bounds based on the rényi divergence. In Artificial Intelligence and Statistics, pages 435 444.444.B PMLR, 2016.
[BHMZ20] Olivier Bousquet. Steve Hanneke, Shay Moran, and Nikita Zhivotovskiy. Proper learning, helly number, and an optimal svm bound. In Conference on Learning Theory, pages 582-609. PMLR, 2020.
[BK12] Daniel Berend and Aryeh Kontorovich. On the convergence of the empirical distribu- tion. arXiv preprint 12Xiv:1005.6711, 2012.
[BL03] Avrim Blum and John Langford. Pac-mdl bounds. In Learning Theory and Kernel Machines: 16th Annual Conference on Learning Theory and 7th Kernel Workshop. COLT/Kernel 2003, Washington, DC, USA, August 24-27, 2003. Proceedings, pages Springer, 2003.
[BLGŞ21] Tolga Birdal, Aaron Lou, Leonidas Guibas, and Umut Şimşekli. Intrinsic dimension, persistent homology and generalization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[Bro03] Paul Bromiley. Products and convolutions of gaussian probability density functions. Tina- Vision Memo, 3(4):1, 2003.
[BSE+21] Melih Barsbey, Milad Sefidgaran, Murat A Erdogdu, Gaël Richard, and Umut Şimşekli. Heavy tails SGD and compressibility of overparametrized neural networks. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
[BZV20] Yuheng Bu, Shaofeng Zou, and Venugopal V. Veeravalli. Tightening mutual information-basel bounds on generalization error. IEEE Journal on Selected Ar- eas Information Theory, 1(1):121-130 May 2020.
[Cat03] Olivier Catoni. A pac-bayesian approach to adaptive classification. preprint, 840. 2003.
11

---

Published as a conference paper at ICLR 2025
[CK22] Dan Tsir Cohen and Aryeh Kontorovich. Learning with metric losses. In Conference on Learning Theory. pages 662-700. PMLR, 2022.
[CRHC18] David M Chan, Roshan Rao, Forrest Huang, and John F Canny. t-sne-cuda: Gpu- accelerated t-sne and its applications to modern data. In 2018 30th International Symposium on Computer Architecture and High Performance Computing (SBAC-PAD), pages 330-338. IEEE, 2018.
[DFH* 15] Cynthia Dwork, Vitaly Feldman, Moritz Hardt, Toni Pitassi, Omer Reingold, and Aaron Roth. Generalization in adaptive data analysis and holdout reuse. Advances in neural information processing systems, 28, 2015.
[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.
[KSV20] Yann Dubois, Douwe Kiela, David Schwab, and Ramakrishna Vedantam. Learning optimal representations with the decodable information bottleneck. Advances in Neural Information Processing Systems, 33:1863:1111 18841-1668,,0 2020.
[DR + 14) Cynthia Dwork, Aaron Roth, et al. The algorithmic foundations of differential privacy. Foundations and Trends@ Theoretical Computer Science, 9(3-4)2211-40) 2014.
[DR17] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint PrrXiv:1703..1 1008, 2017.
[DR18] Gintare Karolina Dziugaite and Daniel M Roy. Data-dependent pac-bayes priors via differential privacy. Advances in neural information processing systems, 31, 2018.
[DTK12] J-L Durrieu, J-Ph Thiran, and Finnian Kelly. Lower and upper bounds for approxi- mation the bollback-leible divergence between gaussian mixture models. In 2012 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1933-4836. leee, 2012.
[DV75] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time, Communications on pure and applied muthematics, 28(1):1-47, 1975.
[Dwo06] Cynthia Dwork. Differential privacy. In International colloquium on automata, languages, and programming, pages 1-12. Springer, 2006.
[EAZ18] Iñaki Estella Aguerri and Abdellatif Zaidi. Distributed information bottleneck method for discrete and gaussian sources. In International Zurich Seminar on Information and Commanication (IZS 2018). Proceedings, pages 35-39. ETH Zurich, 2018.
[EGI20] Amedeo Roberto Esposito, Michael Gastpar, and Ibrahim Issa. Generalization error bounds via Rényi-, fddierrgnnes and maximal leakage, 2020.
[Fis20] Ian Fischer. The conditional entropy bottleneck. Entropy, 22(9):999, 2020.
[GB10] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feed- forward neural networks. In Proceedings the thirteenth international conference artificial intelligence and statistics, pages Proceedings, 2010.
JMLR Workshop and Conference
[GBC16] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning, 2016.
[Gei21] Bernhard C Geiger. On information plane analyses of neural network classifiers-a review, IEEE Transactions on Neural Networks and Learning Systems, 2021.
[GK19] Bernhard Geiger and Tobias Koch. On the information dimension of stochastic processes. IEFE transactions information theory, (110))6996-65188 2019.
12

---

Published as a conference paper at ICLR 2025
[GLLM09] Pascal Germain, Alexandre Lacasse, François Laviolette, and Mario Marchand. Pac- bayesian learning of linear classifiers. In Proceedings of the 26th Annual International Conference Machine Learning. pages 353-360. 2009.
[HD22] Fredrik Hellström and Giuseppe Durisi. A new family of generalization bounds using samplewise evaluated emi. Advances Neural Information Processing Systems, 35:10108 10111108-011,,,,1111111111011111 2022.
[HDMR21] Mahdi Haghifam, Gintare Karolina Dziugaite, Shay Moran, and Daniel M. Roy. Towards unified information-theorrtt framework for generalization. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
[HJTW21] Daniel Hsu, Ziwei Ji, Matus Telgarsky, and Lan Wang. Generalization bounds via distillation. In International Conference on Learning Representations, 2021.
[HK19] Steve Hanneke and Aryeh Kontorovich, A sharp lower bound for agnostic learning with sample compression schemes. In Algorithmic Learning Theory, pages 489-505. PMLR, 2019.
[HK21] Steve Hanneke and Aryeh Kontorovich, Stable sample compression schemes: New applications and optimal svm margin bound. In Algorithmic Learning Theory, pages 697-721. PMLR, 2021.
[HKS19] Steve Hanneke, Aryeh Kontorovich, and Menachem Sadigurschi. Sample compression for real-valued learners. In Algorithmic Learning Theory, pages 466-488. PMLR, 2019.
[HKSW20] Steve Hanneke, Arych Kontorovich, Sivan Sabato, and Roi Weiss. Universal bayes consistency in metric spaces. In 2020 Information Theory and Applications Workshop (ITA), pages 1-33. IEEE, 2020.
[H007] John R Hershey and Peder A Olsen. Approximating the kullback leibler divergence be- tween gaussian mixture models. In 2007 IEEE International Conference on Acoustics, Speech and Signal Processing-ICASSSPP volume 4, pages IV-317. IEFE, 2007.
[HRGT+23] Mahdi Haghifam, Borja Kodríguez-Gálver Ragnar Thobaben, Mikael Skoglund, Daniel M Roy, and Gintare Karolina Dziugaite. Limitations of information-theoreee generalization bounds for gradient descent methods in stochastic convex optimization. In International Conference on Algorithmic Learning Theory. pages 663-706. PMLR. 2023.
[HRVSG21] Hrayr Harutyunyan, Maxim Raginsky, Greg Ver Steeg, and Aram Galstyan. information-thoretti generalization bounds for black-box learning algorithms. Ad- vances in Neural Information Processing Systems, 34, 2021.
[HSKM22] Liam Hodgkinson, Umut Simsekli, Rajiv Khanna, and Michael Mahoney. General- ization bounds using lower tail exponents in stochastic optimizers. In International Conference on Machine Learning, pages $774-8795, PMLR, 2022.
[Hu194] Jonathan J. Hull. A database for handwritten text recognition research. JEEE Transac- lions on patlern analysis and machine intelligence, 16(5):550-554, 1994.
[KASK22] Michael Kleinman, Alessandro Achille, Stefano Soatto, and Jonathan Kao. Gaes- korner common information variational autoencoder. arXiv preprint arXiv:2205.12238 2022.
[KB15) Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[KDJH23] Kenji Kawaguchi, Zhun Deng, Xu Ji, and Jiaoyang Huang. How does information bottleneck help deep learning? In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 16099. PMLR, 23-29 Jul 2023.
13

---

Published as a conference paper at ICLR 2025
[K]+ Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. Toronto, ON, Canada, 2009.
[MMN 16] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: (00::1609.44400 04836, 2016.
[KTVK18] Artemy Kolchinsky, Brendan D Tracey, and Steven Van Kuyk. Caveats for information bottleneck in deterministic scenarios. arXiv preprint arXiv: 07593, 2018.
[KTW19] Artemy Kolchinsky, Brendan D Tracey, and David H Wolpert. Nonlinear information bottleneck. Entropy, 21(12):1181, 2019.
[KW14] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. ICLR, 2014.
[LC01] John Langford and Rich Caruana. (not) bounding the error. Advances in Neural Information Processing Systems, 14, 2001.
[Liv23] Roi Livni. Information theoretic lower bounds for information theoretic upper bounds. Advances in Neural Information Processing Systems, 36, 2023.
[3] Yilin Lyu, Xin Liu, Mingyang Song, Xinyue Wang, Yaxin Peng. Tieyong Zeng, and Liping Jing. Recognizable information bottleneck. arXiv preprint arXiv:2304.1 14641168,, 2023.
[LN22] Gábor Lugosi and Gergely Neu. Generalization bounds via convex analysis. In Conference on Learning Theory, pages 3524-3546. PMIR, 2022.
[LW86] Nick Littlestone and Manfred Warmuth. Relating data compression and learnability, Citeseer, 1986.
(LW$22) Soon Hoe Lim, Yijun Wan, and Umut Şimşekli. Chaotic regularization and heavy- tailed limits for deterministic gradient descent. arXiv preprint arXiv:2205.11661, 2022.
[Mau04] Andreas Maurer. A note on the pac bayesian theorem. arXiv preprint cs/0411099, 2004.
[NBS18] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-hayesian ap- proach to spectralyynnormaiized margin bounds for neural networks, 2018.
[NDHR21] Gergely Neu, Gintare Karolina Dziugaite, Mahdi Haghifam, and Daniel M. Roy. Information-theoretie generalization bounds for stochastic gradient descent, 2021.
[NDR20] Jeffrey Negrea, Gintare Karolina Dziugaite, and Daniel Roy. In defense of uniform convergence: Generalization via derandomization with an application to interpolating predictors. In International Conference on Machine Learning, pages 7263-7272 PMLR, 2020.
20] Jeffrey Negrea, Mahdi Haghifam, Gintare Karolina Dziugaite, Ashish Khisti, and Daniel M. Roy. information-theoree generalization bounds for sgld via data- dependent estimates, 2020.
[NND+22] Tan Nguyen, Tam Nguyen, Hai Do. Khai Nguyen, Vishwanath Saragadam, Minh Pham, Khuong Duy Nguyen, Nhat Ho, and Stanley Osher. Improving transformer with an admixture of attention heads. Advances in neural information processing systems, 35:27937-27952 2022.
INNL+22] Tam Minh Nguyen, Tan Minh Nguyen, Dung DD Le, Duy Khuong Nguyen, Viet-Anh Tran, Richard Baraniuk, Nhat Ho, and Stanley Osher. Improving transformers with probabilistic attention keys. In International Conference on Machine Learning, pages 16595-16621. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2025
[PGM* 19] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high tegh-eerfrmaace deep learning library. Advances in neural information processing systems, 32, 2019.
[PORSTS21] María Pérez-Ortiz, Omar Rivasplata, John Shawe- Taylor, and Csaba Szepesvári. Tighter risk certificates for neural networks. Journal of Machine Learning Research, 22(227):1-40, 2021.
[RG19] Borja Rodriguez Galvez. The information bottleneck: Connections to other problems, learning and exploration of the ib curve, 2019.
[RGTS20] Borja Rodríguez Gálvez, Ragnar Thobaben, and Mikael Skoglund. The convex information bottleneck lagrangian. Entropy, 22(1):988, 2020.
[RKSST20] Omar Rivasplata, IIja Kuzborskij, Csaba Szepesvári, and John Shawe- Shawe-Taylor. Pac- bayes analysis beyond the usual bounds. Advances in Neural Information Processing Systems, 33:1 66833-1885,, 2020.
[RZ16] Daniel Russo and James Zou. Controlling bias in adaptive data analysis using infor- mation theory. In Arthur Gretton and Christian C. Robert, editors, Proceedings the 19th International Conference on Artificial Intelligence and Statistics, volume 51 of Proceedings of Machine Learning Research, pages 1232-1240, Cadiz, Spain, 09-11 May 2016. PMLR.
[SAM 20] Taiji Suzuki, Hiroshi Abe, Tomoya Murata, Shingo Horiuchi, Kotaro Ito, Tokuma Wachi, So Hirai, Masatoshi Yukishima, and Tomoaki Nishimura. Spectral pruning: Compressing deep neural networks via spectral analysis and its generalization error. International Joint Conference on Artificial Intelligence, pages 2839, 2020.
[See02] Matthias Seeger. Pac-bayesian generalisation error bounds for gaussian process classification. Journal of machine learning research, 3(Oct):233 2002.
[SGRS22] Milad Sefidgaran, Amin Gohari, Gael Richard, and Umut Simsekli. Rate-distortion theoretic generalization bounds for stochastic learning algorithms. In Conference on Learning Theory, pages 4416 4463. PMLR, 2022.
[$SDE20) Umut Şimşekli, Ozan Sener, George Deligiannidis, and Murat A Erdogdu. Hausdorff dimension, heavy tails, and generalization neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 5138-5151, Curran Associates, Inc., 2020.
[SSTI0] Ohad Shamir, Sivan Sabato, and Naftali Tishby. Learning and generalization with the information bottleneck. Theoretical Computer Science, 41 411(2993))2266-2211 2010.
[SZ20] Thomas Steinke and Lydia Zakynthinou. Reasoning about generalization via condi- tional mutual information. In Jacob Abemnethy and Shivani Agarwal, editors, Proceed- ings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3437-3452. PMLR, 09-12 Jul 2020.
[SZ24] Milad Sefidgaran and Abdellatif Zaidi. Data-dependent generalization bounds via variable-size compressibility. IEEE Transactions on Information Theory, 2024.
[SZK23] Milad Sefidgaran, Abdellatif Zaidi, and Piotr Krasnowski. Minimum description length and generalization guarantees for representation learning. In Thirty-eevetth Conference on Neural Information Processing Systems (NeurIPS), 2023.
[TWS117] Niklas Thiemann, Christian Igel, Olivier Wintenberger, and Yevgeny Seldin. A strongly quasiconvex pac-bayesian bound. In International Conference on Algorithmic Learning Theory, pages 466-492. PMLR, 2017.
[TPB00] Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004057, 2000.
15

---

Published as a conference paper at ICLR 2025
[TS13] Ilya O Tolstikhin and Yevgeny Seldin. Pac--baes-emmiircall-eee bernsseein inequality. Advances in Neural Information Processing Systems, 26, 2013.
[VGHM21] Paul Viallard, Pascal Germain, Amaury Habrard, and Emilie Morvant. A general frame- work for the disintegration of pac-bayesian bounds. arXiv preprint arXiv:2102.08669) 2021.
[VPV18] Matí Vera, Pablo Piantanida, and Leonardo Rey Vega. The role of the information bottleneck in representation learning. In 2018 IEEE International Symposium on Information Theory (ISIT), pages 1580-1584, 2018.
[XR17] Aolin Xu and Maxim Raginsky. Information-theorede analysis of generalization capability of learning algorithms. Advances in Neural Information Processing Systems, 30, 2017.
[YLL12] Minn-Shen Yang, Chien-Yo Yo Lai, and Chih-Ying Lin. robust em clustering algorithm for gaussian mixture models. Pattern Recognition, 45(11):3950-3661 2012.
[ZT1.22] Ruida Zhou, Chao Tian, and Tie Liu. Individually conditional individual mutual information bound on generalization error. IEEE Transactions on Information Theory, 68(5):3304-3316, 2022.
16

---

Published as a conference paper at ICLR 2025
Appendices
The appendices are organized as follows:
In Appendix we provide the intuition behind the lossy generalization bounds and we present an extension Theorem to lossy compression settings.
In Appendix we show how the established generalization bounds of this work can be extended to cases where the prior violates the symmetry condition.
In Appendix C explain in detail our approach finding the Gaussian mixture priors and how to use them in a regularizer term. This subsection is further divided into three parts, describing
initialization method in Appendix C.I
the lossless approach in Appendix C.2
and the lossy approach in Appendix C.3
In Appendix D
we discuss the potential future directions.
Appendix explains the details of our experiments.
Appendix contains the used approximation method for the KL divergence between Gaussian distribution and Gaussian mixture distribution, and also between two Gaussian mixture distributions.
Finally, the deferred proofs are presented in Appendix G
A INTUITION BEHIND LOSSY GENERALIZATION BOUNDS
The bounds Theorems D and for the deterministic encoders may become vacuous due to the KL- divergence term, and the bounds cannot explain the empirically observed geometrical compression phenomenon (Gei2l). These issues can be addressed using the lossy compressibility approach, as opposed to the lossless compressibility approach considered in previous sections. To provide better intuition for these approaches, we first briefly explain their counterparts in information theory, i.e., lossless and lossy source compression.
Consider a discrete source V ~ Pv and assume that we have n i.i.d. realizations Vi...., Vn of this source. Then, for sufficiently large values of n. the classical lossless source coding result in information theory states that this sequence can be described by approximately (V) bits, where (VV the Shannon entropy function. Thus, intuitively, H (V) ) the complexity of the source V. Now, suppose that V is no longer discrete. Then Vi,..., V,, can no longer be described by any finite number of bits. However, if we consider some "vector quantization" instead, sufficiently close vector can be described by finite number of bits. This concept is called lossy compression. The amount of closeness is called the distortion, and the minimum number of needed bits (per sample) to describe the source within given distortion level is given by the rate-distortion function.
Similar to (SZK23 Section 2.2.1 and Appendix C.1.2), we borrow such concepts to capture the "lossy complexity" of the latent variables in order to avoid non-vacuous bounds, which can also explain the geometrical compression phenomenon GGe2 SZK23). This is achieved by considering the compressibility of "quantized" latent variables derived using the "distorted" encoders W.. Note that W. is distorted only for the regularization term to measure the lossy compressibility (rate-distortion), and the undistorted latent variables are passed the decoder. This is different from approaches that simply add noise to the output of the encoder and pass it to the decoder.
Finally, we show how derive similar lossy bounds (14 in terms of the function hp. We first define the inverse of the function fop as follows. For any [0, 2] and x2 E [0, 1]. let
(y( sup(za e [0, 1]: h (2( & v)-
(25)
Let We E We be any quantized model defined by Pw Pw.is+ that satisfy the distortion criterion Epg Paris gen((S, W) - gen(S, (W) & e, where V = (W., Wa). Then, using Theorem for the
17

---

Published as a conference paper at ICLR 2025
quantized model, have
Ess W. 8 hD (C(Y', 8), 8)){ $
MDL(Q) + log(n) Ex (Ix.r. (fller ((lir--xrli)) A(W Q)
n
(26)
Next, using the Jensen inequality, we have hn [Ew (C()), s., [É(S, wO1]] < E.. (nD (c(X' Y L(x, 81) ((x,
(27)
Combining the above two inequalities yields
hD (F (
AA(, Q)
(28)
Finally, we have
Es.w (ww.genSs < ir gen(S, w]] +
[EwTC((T)) s. [C(S, W)] +c
<nD (min(,, Q) Q [2(S, w1)) E s. W)] + (29)
In particular, for negligible values of s.w [É(S, h)) min (2, Es,
[2(8,
(2 A(W, Q)) +log()) which gives
Bg.n (W) $ MDL.(Q) + log(n)
+ c.
B GENERALIZATION BOUNDS VIA NON-SYMMETRIC PRIORS
In this section, we discuss how the bounds of Theorems I and 2 can be extended to settings in which the requirement of symmetry is relaxed partially. We focus "diffrrentillly private" and "partially symmetric" data-dependent priors.
B. DIFFERENTIALLY PRIVATE DATA ATTA-DEPENDENT PRIORS
One way extend the results to include the partially symmetric data-dependent priors is by leveraging the differential privacy tools (Dwo06) DR 14: DFH 15 (R18). The reader is referred to (Alq2i Section 3.3) for more differentially private priors.
Recall that given the dataset 5 we train model W using the learning algorithm A(.), ie.., W = A(S). Now, assume that by having the dataset S and the trained model W A(S) choose the prior QS.W using a potentially stochastic mechanism T x W - 2, where Q denotes the space of all conditional distributions of U, U' given (Y, Y"). that is "strongly" symmetric. To state the definition of strongly symmetric prior, we first recall the notations of Uz. U = and Y,, Y for any permutation π: [2n] [2m]. Let 222n
(Y, Y'). Then, we define Y, and Y as
Y, =YYD) Yz(njs
Y (m( (=()n)-
(30)
The variables Uz and U', are defined in similar manner.
Definition 2 (Strongly symmetric prior). A conditional distribution Q of U, U' given (Y, Y') is strongly symmetric, for every (U, U', Y, Y') and every permutation :: [2m] + [2m] that preserves the labeling (i.e.. Y, Y and Y' = Y') we have
QUY.Y QUU,, U.
(31)
18

---

Published as a conference paper at ICLR 2025
Note that any strongly symmetric prior satisfies the symmetry condition of Definition In addition, the per-label Gaussian-mixture prior of Section meets the strongly symmetric condition. To show this, recall that for any [C]] the Gaussian mixture prior for label is denoted by Q.. Given these per-label priors, the prior Q is defined as
Q(U, U", W.) -Q Y')
II (v)( d=(n]
It is immediate to see that this prior is strongly symmetric under any permutation that preserves the labeling.
Next, we define the notion of learning the prior in differentially private manner. For simplicity, we consider the case where A(S) can be written as a deterministic function (S, V), where represents all the stochasticity in the learning algorithm that is independent of S. An example of such learning algorithm is the Stochastic Gradient Descent (SGD) algorithm.
Definition (Differentially private prior). We say T S x W - Q is e--iifeerentially private if for any fixed V, and all datasets S and S₁ that are different in only coordinate and for all measurable subsets B = Q, we have
P ((Q ((A5)) EB) & c < (S e B).
(32)
where A(S) g(S. V) and A(S1) V).
Now, we state our tail-bound result for E--differentially private prior.
Proposition 1. Consider the setup of Theorem and suppose the prior QS.A(S) is chosen using
E--differentially private mechanism T: S x W Q. Then, for any 3 0 and for n e 10, with probability least 8 over choices of (S, S', W), it holds that
DEL (Pu,UIXX,,, w. Q (.((s) ) log(2n/0)
hD (B)
n
NYYY [glixx - ixxl11)
+ Ey X
+ 228 V log(4/6)
(33)
where py and px: are empirical distributions of Y and Y', respectively.
The proof stated Appendix G.3
an extension of Theorem 2 usingg DFH+15 Theorems 18&19 19)
and (DR18 Theorem 4.2).
B.2 PARTIALLY SYMMETRIC
PRIORS
In this section, show alternative way extend our generalization bound results by defining the partially symmetric priors.
Definition 4 (Partially symmetric prior). The prior Q is (c, 8) partially symmetric for the learning
algorithm A: nn W, if with probability at least over choices of (S', S, U')
Ps Ps.w. Q
(34) where this should hold for any permutation TY,y' (which could potentially depend on Y, Y') that satisfies the labeling.
VTYYY:
Y & e Q U' U' Y,Y
Note that the partially symmetric prior can potentially depend on (5, W).
Propasition 2. Consider the setup of Theorem Then, for any (c, 8) -partiall symmetric conditional
distribution Q and for n > 10, we have
MDL(Q) Q) + log (Se + ne=)
E Egs. ho ((YY, ( X')
T MY (flir (
E Ex
(35)
19

---

Published as a conference paper at ICLR 2025
where py and pyr are empirical distributions of Y and Y', respectively and MDL(Q) is defined in
70
This result is proved in Appendix G.4
C
GAUSSIAN MIXTURE PRIOR APPROXIMATION AND REGULARIZATION
In this section, we explain in detail our approach to finding appropriate data-dependent Gaussian mixture prior and how to use it in a regularizer term along the optimization trajectories. The section subdivided into three parts: the first part explains how initialize the components of the Gaussian mixture prior, and the other two parts explain the lossless and lossy versions of our approach.
Recall that we are considering a regularizer term equal to
Regularizer( Q) := DxI. (U
(36)
where the indices B indicate the restriction the set B. However, for the sake of simplicity, will drop the dependence on B in the rest of this section. Also, in the following, the superscript (t) is used denote the optimization iteration t N
We choose a Gaussian mixture prior Q in lossless and lossy ways. In both approaches, we initialize (0) (0) three sets of parameters oeem- peem, and (0) for [C] and m E [M]. similarly. We will explain
this first.
C.I INITIALIZATION OF THE COMPONENTS
(0) We let Ceem 1/M, for e e [C] and m [M]. The standard deviation values dem are randomly chosen from the distribution (0,TTa)
(0) The means of the components μeem are initialized in a way that the centers are initialized in the k-means++ method (Ar07). More specifically, they are initialized as follows.
1. The model's encoder We is initialized.
2. A mini-batch Z [Zr...., Zz}, with large mini- size b > b, of the training data is selected. Let X and Y be the set of features and labels of this mini-batch.
For simplicity, we denote by X. mini-batch with label e [C]. Note that EmElcy be =
S X the subset of features of the
Using the initialized encoder, compute the corresponding parameters of the latent spaces for this mini-batch. Denote their mean vector as μ. {pe,i Me For each e [C]. we let uod be equal one the elements in p.. uniformly.
3. For 2 & & M. take mini-batel Z, with per-label features and latent variable means X, and Mtc Then, for all c [C]. we compute the below distances:
di)n,c(t) min m/emm
Me,
He,m!
(0)
2
ie [bc]-
Then, we randomly sample index 1 from the set [bc] according to weighted probability distribution, where the index i a weight proportional to dimin.(()). We let jemn be equal to Me.
C.2 LOSSLESS GAUSSIAN MIXTURE PRIOR
We start with the lossless version, which is easier to explain. Based on our observations in the experiments, the final population accuracy achieved when using the lossless regularizer is better than when using VIB (AFDM17) or CDVIB (SZK23) but worse than when using the lossy version, explained in Appendix C.3
Update of the priors. Suppose the mini-batch picked at iteration t (B()) {2(1) ..... 2 (0) We drop the dependence of the samples (I) for better readability. Then, the regularizer (36). at iteration
20

---

Published as a conference paper at ICLR 2025
(t), can be written as
Regularizer( Q) E
DKL (Pu (U)).
(37)
We propose upper and lower bounds this term. The upper bound is already presented in (18) denoted as vaa::
Regularizer Q)
D
Σ [b]
Σ [M]
Yim DEL (Prslass, Uy.m (Ui)) loog 1
(agi.
(38)
7m
se
me
The upper bound holds for all choices of 7i.m 3 such that [me[M] 7i.m for any [6]. As
explained in Section the coefficients "i that minimize the above upper bound and thus make tighter are equal to
(1) ayemm DEL (Pogleeiie D 1Q
i [b], e [M].
(39)
Σ
(1) "i..me
Dxc (Pogleswee 100 m))m)))))
m'e[M]
Ni
if Vi-
Denote 7.
0,
otherwise..
Next, we establish an estimated lower bound on the regularizer as
Regularizer( Q)
Σ
log
(2nTe) II [I 22.3)
log
Σ
ge ti.
2
jE[d]
i=[b]
m==
Σ
2 log (2TTed II je[d] 22.5) + log Σ
(1)
b=[6]
m=
prod:
(40)
where
(Pecsl Sgmma
Lictol 2(02 to Beim.
1..m := EV EV SSj-Poo Popglasie
Q10)
(a)
2T 22.0 + (e) 2)
mi
(1 Sh
Eye[4]
2o!1)
1,m
(41)
2 (1) [Ieea
2
yimmi
where the step (a) is derived from (Bro03). The reader is referred to Appendix for details on how this upper bound is derived.
has already been observed in (H07)) for the case of two Gaussian mixture distributions that the Ki--divergence is better estimated by considering the average of the product lower bound and the variational upper bound. We then consider the following estimate as the regularizer term
D
D prod
Regularizer( Q)
Dest:>
(42)
where Dvar and D prod are defined in 38 and (40). respectively.
Next, treat Rim constants and find the parameters p o.mo a that minimize Dea by solving the following equations
2Dest
DDest
DDest 0,
0,
0,
ppe.m.
do do
6ae.m
21

---

Published as a conference paper at ICLR 2025
with the constraint that Em aemm for each E [C]. The above equations have the following optimal solutions Vemm and a, and amm=
be.m E [E[b]] Wamm/aa
eem
Enetol (Memmoos Peemm(aaed )²)
c.m.j
be,m
a c,m
-b.m/bes
be
be
Σ 2
be.
[e(()] N.a.
(43)
where
7. = The. + B....
2
Ti.m
ifc= yh:
Emmretaa Dimm
Bre.m
otherwise.
(t) (Pred Muromma
Esataj
Ti.m ==
(t) D5.m
2e (8)
(44)
Note that E [d] denotes the index of the coordinate in I&d and a. reduce the dependence of the prior the dataset, choose the updates
(o mm o a. Finally,
Pemm =(1 - m) mp... + 3(-11)
(t+1) o.m = (1 (. 120 cm + 32
(t+1)
(1+1) =(1 (13))2 (
(8) + 7300 c.m.
(45)
c,m
where 7/1 72, 73 E [0, 1] are some fixed coefficients and 3/ 1((+1) j E [2]. are i.d. multivariate
Gaussian random variables distributed as N(Oa, 55t (()+)) Ia). Here 0₂2 = (o...., 0) e Rd and c) ((++1)
Bt are some fixed constants.
Regularizer. Finally, the regularizer estimation (42) can be simplified
Regularizer Q)
1
Σ
log (>metan 1 DKC. (Poglen 1Q 1922)m))
2
#=(b]
1 Σ
log (2TT)) (I jejg 22nd) + log
Σ mm
M
a (1)
km)
(46)
2
C.3 LOSSY GAUSSIAN MIXTURE PRIOR
Now, we proceed with the lossy version of the regularizer. For this, we consider the MDL of the "perturbed" latent variable while passing the unperturbed latent variable to the decoder. Fix some C R and let = (c....,c) R...
For the regularizer, we first consider the perturbed U
Ú- + Ż (ux Z1) Zz =: Úz Ü2
(47)
where 2, Zı, and Z2 are independent multi-varian random variables, drawn from the distributions
N Os. vd/4 La + diag (e)) .N (0.,. / 2//41a) and A (O.,, diag (ož., + +)), respectively. Conse- quently, Ún N Nuxx, v d/4 T2) is independent from Ú² A000,, diag(o <)), given (X, W.).
22

---

Published as a conference paper at ICLR 2025
For each label e [C], we consider two Gaussian mixture priors Qe, and Qe,2 for Ú1 and Us. respectively, as follows:
Qc,1 Eme[ary ac,m Qemm,1,
(48)
Qa,2 [mmEMn) Oc.m Qemm.2
(49)
over R., where Oc.. [0, 11, [mE[M]] Oc,m for each [C], and where (Q a and 2) are multivariate Gaussian distributions with a diagonal covariance matrix:
Qe,m.1 =N -N(Ha.m V /d/AL).
Qc,m. =N diag(o (o m E)).
Note that the Gaussian mixture priors Qe,1 and Qe have the parameters of ac.m Now, let the prior Qe be the distortion of U Ü1 + U2 when Ün Qe,1 and U2 Qa.2.
Now, for the variation upper bound D for the regularizer, we first consider the inequality
DRL Pojr. 1Qv.) DDL (N vd/4Ta)lQu.)) DE (W(Os, e))lQuu.?
=:DLLL pe 19u.)
(50)
Using the same arguments as in the lossless version but for DKILLossy ( 1011,,,,QQ11 instead of DKL. (e 1Qu.). we derive the following upper bound, denoted as D
Regularizer(Q)) & Dvar = Σ
Σ
R.m (eLLttossyPEtpp... w.. sim ( (U:)) - log Of
1
Ti
ze [M[M]
(51)
which is minimized for
(t) Oueim
Det. toney (Pugleg, (..
7.m
i [b]. m [M].
(52)
Losay (eu wo (Q(1) m.
Σ
m'e m'e[M]
cf V/i
Denote Tr..
0,
otherwise..
For the lower bound, we apply similar lower bound as the lossless case. This (estimated) lower bound, denoted by Dpod, is equal to
M
D prod == Σ
2
log Teva) + log
Σ
(1)
Ne.m
(53)
KE(
where
(4)
f.
2vd
(54)
(2 vdyd
We then consider the following estimate as the regularizer term
D D Proud
Regularizer Q)
Destr
(55)
2
where D and D prod defined in 51 and 53 respectively.
Next, similar to the lossless case, we treat Ts as constants and find the parameters fem T.mo a that minimize Dest by solving the following equations
@Dest 0,
ÖDest
DDest O, daa.m
0,
Ôpe.m.
do da
23

---

Published as a conference paper at ICLR 2025
with the constraint that 2m Ocm 1 for each [C]. The exact closed-fome solutions ue and 0 cam and o e.m.
are equal to
1
pe,
H.e,maa
be,m
1
0 c.m.j
Die(o)
Morma.... zif1
De.m
*
be.m/ Bee
be.m
> fic.m
bo
Σ meeeM [M] bom
b
> SE(b Ti.cm
br.m
[e([b] Fam.m
(56)
where
n.a. == Pe + Ba.em
2 27s.em + Br.c.
3
Di.m
ifc yo: otherwise.
Bi..m
0,
(8)
7i
(57)
Note that j E [d] denotes the index of the coordinate in IRd and a (o mm t oemma)) Finally, to reduce the dependence of the prior on the dataset, we choose the updates
1++1) =(1 - m)) + Mp + +3 (L-11 Mem (I+1) =(1 1/3 (t)) 7,
(t+1)2 a.m = (1-12)0 (() 2
c.m. + 120 c,m + 32 (1+1)
(58)
where 7/1 72, 7ss e [0. 1] are some fixed coefficients and 3/ (((+1), e [2]. are i.id. multivariate Gaussian random variables distributed as NNO0,, cf I2) Here cf(t e R are some fixed constants.
Regularizer: Finally, the regularizer estimation (55) can be simplified as
Regularizer( Q)
Σ
log EmELM (1) m e tonay (Peyles we sr
(Q(9)
is [b]
d
log (evvd) log
Σ a (oonttm))) mm
(59)
2
[b]
D FUTURE DIRECTIONS
In this work, we have established generalization bounds in terms of the minimum description length (MDL) of the latent variables. These bounds particularly suitable for encoder-decoder architectures since they depend only on the encoder part of the model. The bounds improve the state-of-the-aar results from (MDL(Q)) to MDL(Q)) Q)/n some cases.
Inspired by established bounds, propose systematic approach to finding data-dependent prior and using it as a regularizer. The approach consists of first finding the underlying "structure" " of the latent variable space, modeling it as Gaussian mixture and then steering the latent variables in
24

---

Published as a conference paper at ICLR 2025
order to fit that mixture model. Conducted various datasets and with various encoder architectures, reported experiments show promising results.
Our work opens up the door for several interesting future work directions, which we summarize hereafter.
1. In the main body of this work, have established generalization bounds in terms of symmetric priors. However, the proposed practical approach for the design of the prior slightly violates the symmetry condition. While it is not uncommon (sometimes preferred to stretch the technical assumptions for practical designs little, in Appendix [ we resolve the tension by showing that small deviations from the required technical symmetry only yield small penalty in the bound. In the context of this paper, this result could be made more precise by studying the exact deviation of the proposed Gaussian mixture prior from the required symmetry and the caused deterioration of the bound.
The introduced 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. This is major advantage of our approach. In addition, our approach relatively easy to implement. Nevertheless, similar many other regularizers, this comes the expense of some additional computational overhead. Possible means of reducing that overhead include: (i) using the regularizer only in the first K epochs (which, generally, are the most critical (KMN+ ARS17)) and (ii) applying the regularizer in suitable lower-dimensional space, e.g., after proper projection of the latent vector onto that space.
3. In Section we have shown how a weighted attention mechanism emerges naturally in the process of finding the data-dependent Gaussian mixture prior. This may particularly interesting: and is worth further exploration especially when our approach is applied to self-attention layers.
Section 4 proper selection of the number of components of the Gaussian mixture (M) should depend, among other factors, on the dimension d of the problem and the number of hidden "subpopulations"" the latent vector (which itself depends on the used encoder!). Thus, suitable values of M seem difficult to obtain beforehand; and, instead, one resort simply treating it as a hyper-parameter. One approach to circumventing this could be to explore the "structure" of the training data using some common dimensionality reduction and unsupervised clustering techniques, such as the t-SNE of (CRHCI8) or the method of (YLL12).
5. Finally, we mention that in this work, we focused primarily the application to classification tasks. However, the approach and results of this paper can be extended to other setups, such as semi-supervised and transfer learning settings.
E DETAILS OF THE EXPERIMENTS
This section provides additional details about the experiments that were conducted. The code used in the experiments is available at httttps://githhi http://ottttttaaaasst///uuaaaaa Mixture Priors for Representation Learning
E. DATASETS
In experiments, we used the following image classification datasets:
CIFAR10 (KH+09) a dataset of 60,000 labeled images of dimension 32 x 32 x 3 representing 10 different classes animals and vehicles.
CIFAR100 (KH+09) dataset of 60,000 labeled images of dimension 32 32 x 3 representing 100 different classes.
USPS (ul94 dataset of 9,298 labeled images of dimension 16 x 16 x representing 10 classes of handwritten digits.
https:: http///w.ctttt.cc ntu. edu http://ww...aaa.......ooooooooa/ htt.:///iiionttommooooo//oseatee hit.eettttss/mmlllss.sss..ss html #usps
25

---

Published as a conference paper at ICLR 2025
INTEL dataset of over 24,000 labeled images of dimension 150 150 3 representing 6 classes of different landscapes ('buildings', 'forest', "glacier', 'mountain', 'sea', "street").
All images normalized before feeding them the encoder.
E.2 ARCHITECTURE DETAILS
The experiments were conducted using two types of encoder models: a custom convolutional encoder and a pre-trained ResNet18 followed by linear layer (more specifically, the model "ResNet18 Weights. IMAGENETI VI" in PyTorch). The architecture of the CNN-based en- This custom encoder a concatenation of four convolutional layers
coder can be found in Table 2 and two linear layers. We apply max-pooling and LeakyReLLU activation function with a negative slope coefficient set to 0.1. The encoders take re-scaled images as input and generate parameters Paz and variance o² of the latent variable of dimension m = 64. Latent samples are produced using the reparameterization trick introduced by (KW14). Subsequently, the generated latent samples are fed into a decoder with a single linear layer and softmax activation function. The decoder's 's output is soft class prediction.
Our tested encoders were complex enough to make them similar to "a universal function approxima- tor", in line with (DKSV20). Conversely, we employ a straightforward decoder akin to (AFDMI7) to minimize the unwanted regularization caused by highly complex decoder. This approach allows us to emphasize the advantages of our regularizer in terms of generalization performance. However, note that the used ResNet18 model is already pre-trained using various regularization and data augmentation techniques. Therefore, the effect of a new regularizer is naturally less visible.
Table The architecture of the convolutional encoder used in the experiments. The convolutional layers parameterized respectively by the number of input channels, the number of output channels, and the filter size. The linear layers are defined by their input and output sizes.
Encoder
Encoder cont'd dd
Encoder cont' d
Number
Layer
Number
Layer
Number
Layer
Conv2D(3,8,5)
6
10nv2D(16,16.33
Lea)
2
CDn,2D(3,,855
7
LeakyReLU(O..)
12
Linear(256,128)
Leak)ReI.U(0.1)
8
MaxPool(2,2)
Decoder
4
MaxPool(2,2)
9
Flatten
Linear(64,.00
16)))))88,1633)
10
Linear(N,256)
2
Softmax
E.3 IMPLEMENTATION AND TRAINING DETAILS
The PyTorch library (GGM 19 and a GPU Tesla P100 with CUDA 11.0 were utilized train our prediction model. We employed the PyTorc Xavier initialization scheme (GB10) to initialize all weights, except biases set zero. For optimization, we used the Adam optimizer (KB15) with parameters 31 = 0.5 and Bg = 0.999, an initial learning rate of 10 40-4, an exponential decay of 0.97, and batch size of 128.
We trained the encoder and decoder models for 200 epochs five times independently for each considered regularization loss and for each value of the regularization parameter ranging between zero and one. The training was done using conventional cross-entropy loss for image category classification at the decoder's output, and regularization of the encoder's output based on either the standard VIB, the Category-dependent VIB, or our Gaussian mixture objective functions. For the Gaussian mixture objective function, we selected M=20 priors for each class category. The Gaussian mixture priors were initialized using the approaches in C.I The priors were updated after each training iteration using the procedure in C with a moving average coefficient 771 1e-2. 2 for the priors' means Pe.m- 1/2 5c-4 for the priors' variances o² and 7/s 1e-2 for the mixture weights O,,m- Following the approach outlined in (AFDM17), we generated one latent sample per
image during training and 12 samples during testing.
26

---

Published as a conference paper at ICLR 2025
F KL-DIVERGENCE ESTIMATION
In this section, we first recall the KL-divergence estimation of two Gaussian mixture distributions developed in HOO7 Then, adapt these approaches to the case where the KL-divergence estimation of a Gaussian distribution and a Gaussian mixture distribution is considered.
F.1 KL-DIVERGENCE ESTIMATION OF TWo GAUSSIAN MIXTURE DISTRIBUTIONS
In this section, we recall the results of (HO07 DTK12 We give the results only for the case where the covariance matrices of the Gaussian components are diagonal, for simplicity and because only diagonal covariance matrices are considered in our work. However, the results hold for the general form of the covariance matrix.
Consider two Gaussian mixture distributions P and Q, defined as
ByP
Q
a.Qi.
where a, 3, > 0. jjery B, = 1, and Eie [M] a4 = 1. In addition, each component is a multivariate
Gaussian distribution with diagonal covariance matrices.
Pj NN diag(o
Q: - -N diag (.)))
F.1.1 PRODUCT OF GAUSSIAN APPROXIMATION
In this approximation, DKI(PIQ) approximated as (HO07):
Eve(M E, [Py] Eie[M] P;EP,, [Q,]
D Prod( (IQ)=
Σ
By log
(60)
je[N]
Note that this approximation is generally neither an upper bound nor lower bound.
F.1.2 VARIATIONAL APPROXIMATION
In this approximation, DkI(PIQ) approximated as (H007):
Eve[my 33e Dwn (PrlP)) Eie(M age (x (Q:)
D (PIQ)
Σ
By log
(61)
JE[N]
Note that this approximation is again not an upper or lower bound general.
F.1.3 AVERAGE OF TWO APPROXIMATIONS
has been shown (H007 DTK12), that the average of the product and variational approximation provides a better estimate of the KL-divergence between two Gaussian prior distributions.
Dpreal P|[)) + Dvar (l[Q)
Dest Pl(Q)
(62)
2
F.2 KL-DIVERGENCE ESTIMATION BETWEEN A GAUSSIAN AND A GAUSSIAN MIXTURE DISTRIBUTION
In this section, adapt the approaches of (H007) for the setup where P is d-dimennsionl Gaussian distribution with diagonal covariance matrix and Q is Gaussian mixture of M of d-dimensional Gaussians with diagonal covariance matrix.
27

---

Published as a conference paper at ICLR 2025
Formally, let
P = [M, digg
and Q be Gaussian mixture
Q Qin
where 0; 0. Lie[M] tz and
Qi M diag(o (.,))
F.2.1 PRODUCT OF GAUSSIAN BOUND
Denoting Lp(f)= Ep[log(f) (f), we have DKI(PIQ) Lp(P) Lp(Q). Note that
Lp(P) (PP) - log (2TTe) II je(dj ).)))
where h(-) is the differential entropy. Next, to bound Lp(Q), using the idea of (H077, we have
Lp(Q) - -Ep
log Σ a.Q
& log
Σ [.EP((.)) = log Σ aatt)
(63)
where
Ep(Q.] (z))
(64)
is the normalization constant of the product of the Gaussians (refer to (DTК12 Appendix A)). Note that by choice of the diagonal covariance matrices, these constants can be written as the product m coordinate-wis. constants.
Thus, we have
M Σ agtt)
DRI(PIQ) =
1
log (22xe))
II jeld]
273)
log
(opra(PP|Q)
(65)
2
Note that, unlike the KL divergence estimation of two Gaussian mixture priors, here the product of Gaussian approaches provides a lower bound.
F.2.2 VARIATIONAL BOUND
Fix some > 0, i E [M] such that Σ, = 1. Then,
M Lp(Q)= =Pp log (aQi))
= EPp log Σ
M
Qi 7/
)
am7
(aiQ. ]
Σ
E.p
log
%
[M
Σ (Lr((Q.)) log(as/77)) /7))
(66)
Maximizing this lower bound with respect to gives
a,e Dxc
%/s EvelM ave -xxe-Px(Pl))))) Dxc(P)Q0))
(67)
28

---

Published as a conference paper at ICLR 2025
Using this choice, 666
simplifies
Lr( =Ep log
a, ai log Σ
(68)
ie[M]
Hence, overall
D ELp(P) log
Σ
(
(69)
iE[M
=- -olog
Σ [E[M]
oge Pxz(PIQi)
(70)
(PIQ).
(71)
Note that again, unlike the KI. divergence estimation of two Gaussian mixture priors, where the variation approach provides only an approximation, this approach provides an upper bound.
F.2.3 AVERAGE OF TWO APPROXIMATIONS
Finally, to estimate the KL-divergence between Gaussian distribution and Gaussian mixture distribution, we consider the average of the product of the Gaussian lower bound and the variational upper bound.
Dea (PllQ) Prooal P(|Q) Dear Pear(P])) (PIIQ)
(72)
2
G PROOFS
In this section, we present the deferred proofs.
G.1 PROOF OF THEOREM
Fix some symmetric conditional prior Q(U, Y', X, X, We will show that
E... hD |n
- ((10x--x-l.)
- i
MDL(Q) log(n)
n
(73)
where py and frr empirical distributions of Y and Y', respectively,
MDL((Q) := Es.s.... [DL (Pu.uix...w. [Q)]
(74)
and
W) ~ Ps Psw.sPPux.PPPPx. Px P Prqu... (S.S', U,U Ý, W)
Denote
P1 A==Ps. Ps Pupx.w PU... PrU...PPyu..a.
Pa = Ps.w Ps. NU P.I.IW Pr.
Y, Y, & ()) =ho LYY', L'), L(Y Ny (2ixx px-ll)).
Next, similar
(e.g. (RI77 SZK233 and PAC- MAAC-Baye--bal based approaches
(e.g. (Alq2l RKSST20) we use Donsker- Varadhan's inequality to change the measure from P1 to Pz. The cost of such change Dxr(Pi MDL(Q). We apply Donsker- Varadhan on
Donsker-Varadhan's
function nf. Concretely, we have
Ess. (Y, 8')
DKL log Ep. (x((xxx:888))
=MDL(Q) + log Ep (.)(Y.r))),))))
29

---

Published as a conference paper at ICLR 2025
Hence, it remains to show that
Brg (x)) n.
(75)
Let Qy Y Y be the conditional distribution of (Y, Y') given (Y, Y"), under the joint distribution P2. It can be easily verified that Q satisfies the symmetry property since Q is symmetric (as defined in Definition 1. For better clarity, re-state the symmetry property of Q and define some notations that will be used in the rest of the proof.
Let 2 = (Y, Y') and 8 == (Y, Y'). For given permutation म: [2n] - [2n]. the permuted vectors Y and y are defined as
2 Ya(i):- -............. Ya(2n)
2 YA(I).... " Ya(zn)-
(76)
Furthermore, under the permutation म, denote the first coordinates of Y and Y by
Y; :- Yx)).-... Ya(M):
Y y (())..... R(():
(77)
respectively, and the next n coordinates Y and Y by
Y = Y=(++1) Y=()n),
Y (((++1).... P
(78)
respectively. By Q being symmetric, we mean that Qy & remains invariant under all permutations such that Y, Ya() for all i e [2m]. In other words, all permutations such that Y = Y2 and Y' = Y
Hence, we can write
Ery [ [x)((x))))8))
= (
(79)
where Y. Y', Y.
1 22n QU
Fix some Y and Y'. Without loss of generality and for simplicity, assume that and Y' are ordered, in the sense that for [R]. Y, Yr, and (YR+I..... Yn
Y 2, where
R 2 by -xrll1-
Otherwise, easy see that the following analysis holds by proper (potentially non-identical) re-orderings of Y and Y' and corresponding predictions Y (according to the way Y is re-ordered) and Y' (according to the way Y' is re-ordered), such that and Y' coincidence in all first coordinates and do not have any overlap in the remaining
R coordinates.
Furthermore, for r [n:] let J, (r, n + r) Bern ()) be a uniform binary random variable and detine Jf as its complement, i.e., J, u Jf {r, n r). Define the mapping कR == [2n] + [2n] as following: For
[R]. TR(r)= = Jr and ( + n) = J.. For [R + 1,n]. TR(r) = Note that лH depends Y' =Y TR* where Y, and Y are defined in (77) and (78). respectively. Hence, 1øx - px =
and Ta(n + +) n
(Y Y") and under πR. Y = Y na and
jv.. - pr. prral To simplify the notations, in what follows we denote the coordinates of Y by
Y Y (Vaa...... Y Yaan);
and the coordinates of Y! by nn
Y KR
(Y VaRdi-.... Y ).
30

---

Published as a conference paper at ICLR 2025
Note that by (77 and (8)). have Y Fnn
Y (n(i) and Y F Y mi(t+n) for i E [n], where y =n(()) defined in 76 Similar notations are used for the prediction vectors, i.e.
RR == (Y Yamm)
Y R = TR. TR.
With these notations, for fixed ordered Y and Y' we have
Ey Y ( (Xxx Ey Y By... JR. Bem yaR
=Ey Y Ji....n Ja--eeneeeeeee you e " ane Y Yoze RR
(80)
where the first step follows due the symmetric property of Q and the second step follows since Y YTR and Y'
Y "na
Now, consider another mapping π [2n] [2n] such that и is identical to gH for the indices in the range [1 R] o I
+ 1 n R). i... for r [R],
( (r()) = Jr,
a n) = TR(L + n) = J
Furthermore, for the indices in the range in [R + 1: n: U [n + R+ 1: 2m]. æ is defined as follows: for [R + 1,
*(r) = Jys л(n + +r) Jr, where previously defined, J, fr,n n + r] Bern(3) ( is uniform binary random variable and J, its complement. Denote
Jhtn Jnt1,...,
With the above definitions, we have e" (Yaan Y Y 8 R)
ndo 1 Fn 148 2a " + En 10 (2
=Ejg Bemm m)
ns(x=n 8 x " -nnn 2im1 4 4. 218 188 (Y
x E
(a)
nhn
1 Ei
18
En 18
& E EyRyy 1 ~ Rem je
e
(81)
where (a) holds due to the following Lemma, shown in Appendix G.5
Lemma 1. The below relation holds:
F(YTRo Y nn ARY
&
Y xn & hp
Σ
02 YY
Σ l (
(82)
Hence, for fixed ordered Y and Y', combining 80 and yields
E E [ [=
aho + Enm 119% trg Dim M (ea
Y. J.... Benr + @@n
n.
(83)
where the last step is derived by using (SZK23 Proof of Theorme 3). As mentioned before, is easy to see that the above analysis holds for non-ordered Y and Y', by simply considering proper (potentially non-identical) re-orderings of Y and Y' and corresponding predictions Y (according to the way Y is re-ordered) and Y' (according the way Y' is re-ordered), such that Y and Y' coincidence in all first R coordinates and do not have any overlap in the remaining n R coordinates.
Combining
79
80
and
83 shows 75 which completes the proof.
31

---

Published as a conference paper at ICLR 2025
G.2 PROOF OF THEOREM! 2
First note that by convexity of the function hp (SZKK33 Lemma 1)), we have
hD (É(s", W), (()))) & Eg Y hD Y`) ((), 80), Hence, it suffices to show that with probability at least 1 - over choices (S, 5", W),
(84)
Dxu (RU(PUUUXXXVV. Q log(n/6)
Ey EY.
NDD (2(Y)) (.),)
+I x / Gflēr px pxll)]] (85)
Similar to the proof of Theorem define
Pt = PuIx P PrI.. PriUU
Pj = QU.UUIXXXXXXY.WW Priu.. Prru.a
= (", ((), ,Y, (lror -rrli)
Using Donsker-Varadhan's 's inequality, we have
nL y X
<DKi + log Bry (
= =KL (DUU(PUUU/xx..:10)) + log Ery ( (
(86)
Hence,
P B (x
log(n/6)
n
(0) PP log Epg
( (xx((xx))xx))))) log(n/6)
=P Ery en (Y.( n/6)
Es, W. Ep (e
1
8,
(87)
where
(a) follows by 166.
(b) is derived using the Markov inequality,
and (c) is shown in (75).
This completes the proof.
G.3 PROOF OF PROPOSITION
To state the proof, first, we need to recall the notion of /3-approximate max-information; as previously defined in (DRT8) Definition 3.2) and (DR18, Definition 3.2). Here, we state the definition adapted to our setup. For ease of notation, denote
ev = (S.A(S)) = (S))(S.V)).
(88)
Definition 5. Let O. Then, define the 3-max-information between S and Qev(S), denoted by 1 as the minimal value l such that for all product events E and all fixed V, we have
(((s. (())) E EF.. <<'(((S, E F: + B.
(89)
where 5 is independent dataset with the same distribution S.
32

---

Published as a conference paper at ICLR 2025
Fix some 8' > O, which will be made explicit in the following, Now, "similar" to the proof of (DR88
Theorem 4.2), for any Q E Q, define
R(Q) (5. (S,S',W)) w): hD ( (I,)s,), ((S.5)
(90)
where
W, DKu(PU,UXxxx.W. () Q) log(n/8")
+ x bY.Y (1 lixx - 1xx
(91)
Fix some 3 > For every fixed and V. by Definition we know that
PP (S,,, E R(Q"())))S', v) c (S. W,S') (( + 8.
where S is independent of (e(S), 8*). Hence,
(S, Ps.w.s W5, e R ((Q(v())) E Pe (IS {((Qv((5)) +
22 + 3,
(92)
where (a) is derived since by Theorem 2 know that P(R(Q)) 8 for every independent of S and 5" Recall that strong symmetry implies symmetry.
Let B 6/2 and 1 8' 8/2. Equivalently,
8/2
8'
2
With these choices, with probability 8 over choices of (S,S", W), we have
DKL (PU (c((9) + log(2n/8) + 15/2
n
y Ii (llie-brll)) ir
(93)
The final result follows by (DFH 15 Theorem 20), where they showed that
n log(4/0)
1
22 2² Ep
This completes the proof.
G.4 PROOF OF PROPOSITION
Recall the following notations in the proof of Theorem
D
P = Ps,w Ps Pujx.w PU w. PYJU. Priu.
Pa = Ps.w Ps. QU.UUXUXU.YYYYU.. Pr Priju,
VY h [ Y") L(Y - - (flir ivell1)
Using the identical steps as in the proof Theorem we have
E.s. { (xx Y, Y', &, x1) <MDL(Q)) + log Ep, (xx((xxx
Hence, it remains to show that
Ep. ( $ 6e + ne..
(94)
33

---

Published as a conference paper at ICLR 2025
Let Ily, Y denote the set of all permutations that preserve the labeling. Denote the size of this set as N. = N. Then, the prior
Q(U, U Y' X.X
1 N
Σ
QUU,.
(95)
relly,y'
T Furthermore, by Definition we have with probability at least 1 over choices of (S', We, U, U') Psy Ps.w. Q
symmetric in the sense of Definition
QUU Y Y, x,X, < QU Y x, x W.).
(96)
Hence, since f( Y, $ we have that
E e ( < 6E + tep (
(97)
where
P3 =P Ps
P QU P U, Wew.sPrTUrr Pu..
The result now follows since Q is symmetric and hence identical to the proof of Theorem Eps nf (x..
we have
<
(98)
This completes the proof.
G.5 PROOF LEMMA II
For ease of notations, for i E [m]. denote
4i
18
n1
1
(Re + mn
(8%
Consider similar notations for the mapping define li, and h. Furthermore, denote
A (liAR L(x)=
Σ (Li,a l,xx,,
i=R+1
AC = (Ce,n ((,x)=
Σ
(G,rn - (
i- R+1
It is easy to verify that Al
A and
A
(n R)
_210x px
(99)
n
With these notations,
5(x (Y
Y an
V
Y
-ho
(5 Σ li
NY.
(21xx -
Tuo
aR
1=1
(a)
hD
L
AC Σ
linn
AE
=hp
(100)
which completes the proof, assuming the step (a) holds.
then remains show the step (a). To show this step, it is sufficient to prove that for every
F1, [0, 1]. ? RR,, and E IR such that (x1 + c), (x2 c) e [] and |c| & ?, the below
inequality holds:
hp(21,42) - 10(01, < h + e, e). Without loss of generality, assume that x & Ø2. We show the above inequality for different ranges of e, separately.
(101)
34

---

Published as a conference paper at ICLR 2025
If < 0, then since by (SZK23) Lemma 1), hp 22) is decreasing in the real-value range of
a [0, x2] and (x(itrr is increasing in the real-value range of E [x1, 1], we have
hp(21, 22) ha 6, c) <
(11.
where the last inequality follows using the fact that byg is non-negative.
I 22 31, then by letting é' = (I2 - x1) - & 0, we have
hp(21, x2) hp(a1 +,42 =ho(z1, x2) hD (22 - c + E)
ho(()1 - (01 + + e))
<
2 (21;
where (a) is deduced by the symmetry of hn and steps (b) and (c) are deduced similar to the case
E < above.
ef (0, 21)/2]. then we have
hp(a1, hD( 6,22 - =10 + +) ho(02 - hh(x1)) ho hx2))
where the last step follows by definition of the function lic, and since e belongs to the below interval:
[0, é] n [0, (œ1v2 -1122/2]
(102)
If ((12-x1)/22 (x/2 x x1)] then by letting e' (w2 xi) e, we have e' [0, ( x1 )/2]
and
x2( - hD( 2, - =h6(01 + E') ho(02 - - -ho(21) (x(22)
<he(
where the last step follows by definition of the function he, and since e belongs to the below interval:
[0, [0, (2iv2 x122)/22..
(103)
Note that e' & ē, since e' E [0, (I2 x )/2] and [ x1)/2, 2, (02 x1)]. Hence, e' e, and
by assumption & é.
This completes the proof of the lemma.
35


Reviews:
[{"review_id": "l9iyc7REB6", "reviewer": "Reviewer_c35E", "summary": "This paper studies the setting of representation learning, where the task is to jointly learn a good representation and a good predictor. PAC-like generalization bounds are provided, which crucially depend on the complexity of the latent representation. Based on these theoretical insights, the authors then suggest a new regularizer term, which shows promising results in some applications.", "questions": "I have some concrete questions:\n- What is the \"joint\" (line 1188, Appendix C.3) learning procedure that you followed in the experiments? Can you provide more details on the loss used, etc?\n- How do the proof techniques compare to the ones used by SZK23?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Novel generalization bounds for representation learning: The paper derives generalization guarantees for representation learning. To my knowledge, this setting has not been studied in depth previously. I found the remarks about these bounds interesting, particularly the fact that they do not depend on the encoder component.\n- Novel regularization term based on the theoretical analysis: As is often the case with generalization bounds, they can directly inspire regularization terms or other modifications to the learning pipeline. The authors propose a new regularization term for representation learning. In experiments with image classification datasets, the introduced method is shown to compete favorably with other recently introduced methods for this task.", "weaknesses": "In general, I found the presentation of some results in the paper to be a bit lacking, while the discussion of related work seems somewhat insufficient. The authors discuss prior work in the introduction, but the discussion is high-level and mainly directed at 'experts' in the field. An additional (perhaps small) section covering related work would have served this paper well. Additionally, the main results of the paper (Theorem 1 and Theorem 2) are presented without much (if any) discussion on the techniques used in their proofs. Another example is in line 305, where the authors mention and emphasize the 'geometrical compression phenomenon' but do not describe it further. In various places, the manuscript would benefit from additional polishing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kKdc83UNLg", "reviewer": "Reviewer_6GSR", "summary": "This paper provides generalization guarantees for representation learning algorithms by deriving in-expectation and tail bounds on generalization error. The authors develop these bounds based on the MDL principle, using a data-dependent Gaussian mixture prior as a regularizer. By establishing bounds related to the relative entropy between representations from training and test datasets, the method aims to leverage the simplicity and structure of the encoder for improved generalization.", "questions": "1. What might be the potential limitations of the current method? What could be the future work?\n\n2. Are there any constraints on the applicability in real-world scenarios introduced by relying on a Gaussian mixture prior?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper provides non-vacuous generalization bounds for representation learning.\n\n2. The proposed regularizer has been validated in many image datasets.\n\n3. Theoretical results have been rigorously presented and supported.", "weaknesses": "1. The Gaussian mixture prior introduces additional computational overhead, especially during training, which may make the approach less practical for very large datasets.\n\n2. Although the emergence of a weighted attention mechanism is highlighted, the paper could benefit from a more detailed analysis of this component and its role in generalization.\n\n3. The empirical validation is limited to standard classification datasets. Testing on real-world tasks beyond classification (e.g., transfer learning, semi-supervised learning) would strengthen the claims of generalization benefits.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BfjDmbjyNR", "reviewer": "Reviewer_ce5p", "summary": "The paper studies the generalization error of representation learning, in which we have the encoder-decoder model. The authors establish a new generalization bound based on the minimum description length (MDL) of a symmetric prior and the induced representation distribution of the encoder. The paper explains how to choose the prior using a mixture of Gaussians, and connects the generalization bound to the regularizer of the optimization problem, in which the prior and the representation distribution are learned jointly. The authors examine their method in simple datasets and model architectures, and in both lossy and lossless scenarios.", "questions": "1. In section 3.3, the bound in Eq 14 seems to be similar to [1] with $\\sqrt{\\text{MDL}(Q)/n}$ rather than $\\text{MDL}(Q)/n$. Do the authors have any intuition/proof on why this change happens in the lossy setting?\n\n2. Is there any intuition behind $h_C$ in Eq 6?\n\n3. In section 4.1, M denotes the number of mixtures. Do the authors suggest any systematic way of choosing M? Or is it a hyperparameter that needs to be tuned?\n\nSome typos also can be addressed:\n- Line 342, an addition | in KL\n- Eq. 17, the index of X should be batch size and not beta", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is well-written and the motivation is clear. The authors improved the generalization bound of the encoder-decoder representation learning from $\\sqrt{\\text{MDL}(Q)/n}$ to $\\text{MDL}(Q)/n$. The proposed symmetric prior using mixtures of Gaussians is relatively practical and easy to implement. The idea of adding the generalization bound as a regularizer and learning the prior and the induced distribution of the encoder together is interesting.", "weaknesses": "While the authors have improved on the previous generalization bound, the technical work and the idea of regularization are heavily based on [1]. However, I think the mixture of Gussians is a nice addition.\n\nThe lossy generalization bound in Section 3.3 seems a bit incomplete to me. While the authors explain what lossy means, I could not follow how it results in Eq 14 and 15. I also checked the appendix, but could not find anything related to this section. I would appreciate it if the authors could elaborate on this section more, or add a theorem and the proof either in the main text or the appendix. \n\nI also think there are some parts that the author can improve the writing or the intuition of their work (see questions).\n\n[1] Minimum Description Length and Generalization Guarantees for Representation Learning, Sefidgaran et al. 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C78GoWolpx", "reviewer": "Reviewer_w8tc", "summary": "In this paper, the generalization performance of representation learning algorithms (where the models consist of an encoder and a decoder) is considered. Specifically, bounds based on the compressibility of the latent variables of the models are derived. On the basis of these bounds, regularizers using priors with a Gaussian mixture model are devised. These regularizers are numerically shown to provide improved performance compared to prior work for some tasks.", "questions": "1. Can you clarify the data-dependence of the prior, and how this relates to the assumptions under which Theorem 1 is derived?\n\n2. For Theorem 1, the prior to be symmetric (equivalent to the exchangeable priors of Audibert (“A better variance control for PAC-Bayesian classification”, 2004) and Catoni (“PAC-Bayesian Supervised Classification: The Thermodynamics of Statistical Learning”, 2007). In their work, it is often sufficient to consider “almost exchangeable priors”, where permutations are restricted to only swap the $i$th element of the training sample with the $i$th sample in the ghost sample. Would a similar weaker requirement work for your results?\n\n3. Line 266: “when the empirical risk is set to 0.05” What does this mean exactly? Does it mean that you train until the empirical risk reaches 0.05 or below and then stops?\n\n4. In Theorem 1, the bound is provided in terms of a scaled Jensen-Shannon divergence between two Bernoullis. This is reminiscent of the Maurer-Langford-Seeger (MLS) bound in PAC-Bayes, where the corresponding KL divergence is used in the LHS. There are results (Foong et al, “How Tight Can PAC-Bayes be in the Small Data Regime?”, Thm. 4, 2021) indicating that the MLS bound is the tightest possible bound (in some sense) up to the log term. What is the relation between Thm. 1 in the present paper and such preceding bounds? Is it possible to instead use the binary KL in the LHS? (Such bounds also have a fast-rate behavior for zero training loss). I understand that the present paper considers a slightly different setup.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The general problem under consideration is interesting, and the twin goals of explaining generalization and finding ways to improve it are valuable. Interesting connections are drawn between, e.g., the presented bounds and classical compressibility, and between the form of the prior updates and attention. The way in which the regularizer is constructed is, as far as I am aware, new and yields promising results. The numerical experiments compare to several prior alternatives from the literature.", "weaknesses": "The presentation is not always clear. For instance, one of the key concerns about the validity of the data-dependent prior, as far as I understood, is to guarantee that it is symmetric with respect to the training data and a ghost sample. However, the procedure seems to only depend on the training data, and it was unclear to me whether symmetry was preserved. Noise was added at some stages to “partially” preserve symmetry, and allusions were made to prior work which allows for conditions to be only partially fulfilled (e.g. differential privacy), but the results were discussed for exact symmetry.\n\nThe relation between prior literature and the presented generalization bounds and data-dependent priors could be discussed to a greater extent.\n\nMinor:\n\n— “date-dependent” in abstract\n\n— Mix of numbering and words in some lists (“(1) First” in line 218, Lines 350-356)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fafpQylcEI", "reviewer": "Reviewer_9bKp", "summary": "This paper provides a new generalization bound for representation learning in multi-class classification tasks. The error bound incorporates the recent notion of minimum description length (MDL), which empirically has proven to be more useful than mutual information-based bounds. In particular, the in-expectation bound and tail bound only depend on the encoder part. This has implications that in order to achieve better generalization bound, one can propose a regularizer based on the MDL. This is further studied by approximations using a Gaussian mixture whose mean and variance are data-dependent, and an update scheme for the prior is provided. Several numerical experiments verify the theory.", "questions": "1. Typo in line 22 of the abstract, \"date-dependent\" should be \"data-dependent\".\n\n2. In line 49, the author should make it clear what \"MI\" stands for before writing \"MI-based\".\n\n3. In line 86, the author should cite earlier references about the universal approximation with Gaussian mixture, for example, some articles in JRSSB.\n\n4. How do you determine the number of modes in the Gaussian mixture prior?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is clearly written and well-organized.\n\n2. The derived bound is tighter than the best bound with MDL.\n\n3. This work provides a quantitative explanation that the encoder plays the role of generalization, as reflected in the bounds in Theorem 1 and 2.\n\n4. Based on the regularization using MDL, a practical and explicit optimization scheme for the prior is provided using Gaussian mixture models.\n\n5. The theory is verified on a few datasets.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)