PAC-Bayes Information Bottleneck
Authors: Zifeng Wang
Shao-Lun Huang
Ercan Engin Kuruoglu
Jimeng Sun
Xi Chen
Yefeng Zheng
Venue: ICLR 2022 Spotlight (2022)
Keywords: information bottleneck
representation learning
generalization

Abstract: Understanding the source of the superior generalization ability of NNs remains one of the most important problems in ML research. There have been a series of theoretical works trying to derive non-vacuous bounds for NNs. Recently, the compression of information stored in weights (IIW) is proved to play a key role in NNs generalization based on the PAC-Bayes theorem. However, no solution of IIW has ever been provided, which builds a barrier for further investigation of the IIW's property and its potential in practical deep learning. In this paper, we propose an algorithm for the efficient approximation of IIW. Then, we build an IIW-based information bottleneck on the trade-off between accuracy and information complexity of NNs, namely PIB. From PIB, we can empirically identify the fitting to compressing phase transition during NNs' training and the concrete connection between the IIW compression and the generalization. Besides, we verify that IIW is able to explain NNs in broad cases, e.g., varying batch sizes, over-parameterization, and noisy labels. Moreover, we propose an MCMC-based algorithm to sample from the optimal weight posterior characterized by PIB, which fulfills the potential of IIW in enhancing NNs in practice.

Published as a conference paper at ICLR 2022
PAC-BAYES INFORMATION BOTTLENECK
Zifeng Wang" UIUC
Shao-Lun Huang
Ercan E. Kuruoglu Tsinghua University
Tsinghua University
Jimeng Sun UTUC
Xi Chen Tencent
Yefeng Zheng
Tencent
ABSTRACT
Understanding the source of the superior generalization ability of NNs remains one of the most important problems in ML research. There have been series of theoretical works trying to derive non-vacuous bounds for NNs. Recently, the compression of information stored in weights (IIW) is proved to play key role in NNs generalization based on the PAC-Bayes theorem. However, no solution IIW has ever been provided, which builds barrier for further investigation of the IIW's property and its potential in practical deep learning. In this paper, we propose an algorithm for the efficient approximation of IIW. Then, we build an IIW-based information bottleneck on the trade-off between accuracy and informa- tion complexity of NNs, namely PIB. From PIB, we can empirically identify the fitting to compressing phase transition during NNs' training and the concrete con- nection between the IIW compression and the generalization. Besides, we verify that IIW is able explain NNs in broad cases, e.g., varying batch sizes, over- parameterization, and noisy labels. Moreover, we propose an MCMC- bbsed al- gorithm to sample from the optimal weight posterior characterized by PIB, which fulfills the potential of IIW in enhancing NNs in practice.
INTRODUCTION
Understanding the behavior of neural networks (NNs) learned through stochastic gradient descent (SGD) is a prerequisite to revealing the source of NN's generalization in practice. Information bot- tleneck (IB) (Tishby et al., 2000) the lens of information stored in encoded representations of inputs. Drawn from the conception of representation minimality and sufficiency in information theory, IB describes the objective of NNs as a trade-off, where NNs are abstracted by Markov chain Y + X ++ T, as
promising candidate reveal the principle of NNs through
max I(T:Y) - (X)
where X() and YY;Y) are the mutual information of representation T towards inputs X and labels Y, respectively. IB theory claimed (Tishby & Zaslavsky, 2015) and then empirically cor- roborated (Shwartz-7iv & Tishby, 2017) that NNs trained by plain cross entropy loss and SGD all confront an initial fitting phase and subsequent compression phase. It also implied that the repre- sentation compression a causal effect to the good generalization capability of NNs. IB theorem points out the importance of representation compression and ignites a series of follow-ups to pro- pose new learning algorithms on IB explicitly take compression into account (Burgess al., 2018; Achille & Soatto, 2018b; Dai et al., 2018: Li & Eisner, 2019; Achille & Soatto, 201 8a; Kolchinsky et al., 2019; Wu et al., 2020; Pan al., 2020; Goyal et al., 2019; Wang et al., 2019; 2020a).
However, recent critics challenged the universality of the above claims. At first, Saxe et al. (2019) argued that the representation compression phase only appears when double-sided saturating nonlin- earities like tanh and si sigmoid are deployed. The boundary between the two phases fades away with other nonlinearities, e.g., ReLU Second, the claimed causality between compression and gen- eralization was also questioned (Saxe et al., 2019; Goldfeld al., 2019), i.e., sometime networks that do not compress still generalize well, and vice versa. To alleviate this issue, Goldfeld et al. (2019)
Coorreponndence fengw281 noiss edu

---

Published as a conference paper at ICLR 2022
proposed that the clustering of hidden representations concurrently occurs with good generalization ability. However, this new proposal still lacks solid theoretical guarantee. Third, mutual information becomes trivial in deterministic cases (Shwartz-2a) Ziv & Alemi, 2020). Other problems encountered in deterministic cases were pointed out by Kolchinsky et al. (2018). Although several techniques (Shwartz-Ziv & Tishby, 2017; Goldfeld et al., 2019), e-g., binning and adding noise, are adopted to make stochastic approximation for the information term, they might either violate the principle of IB or be contradictory to the high performance of NNs. Motivated by these developments, we focus on the following questions:
Does universal two-phase training behavior of NNs exist in practice? If this claim invalid with the previous information measure I(T; X), can we achieve this two-phase property through another infornation-theoreiic perspective?
As (T; X) was unable to fully explain NN's generalization, can we find another measure with theoretical generalization guarantee? Also, how do we leverage it to amend the IB theory for deep neural network?
How do we utilize our new IB for efficient training and inference of NNs in practice?
In this work, we propose to handle the above questions through the lens of information stored in weights (IIW), i.e., [(w; S) where S = {_x,, Yi);=1 is a finite-sample dataset. Our main contri- butions are four-fold: (1) we propose a new information bottleneck under the umbrella of PAC- Bayes generalization guarantee, namely PAC-Bayes Information Bottleneck (PIB); (2) we derive an approximation of the intractable IIW; (3) we design a Bayesian inference algorithm grounded on stochastic gradient Langevin dynamics (SGLD) for sampling from the optimal weight pos- terior specified by PIB; and (4) we demonstrate that our new information measure covers the wide ground of NN's behavior. Crediting to the plug-and-play modularity of SGD/SGLD, we can adapt any existing NN to a PAC-Bayes IB augmented NN seamlessly. Demo code is at https: /tt::://gttttt htw.hhh..coom/..e /ttp:///aaaaaa/////////////
2 A NEW BOTTLENECK WITH PAC-BAYES GUARANTEE
In this section, we present the preliminaries of PAC-Bayes theory and then introduce our new infor- mation bottleneck. A loss function (((((); Y) is measure of the degree of prediction accuracy fw (X) compared with the ground-truth label Y. Given the ground-truth joint distribution p(X, Y), the expected true risk (out-of-sang) -faample-- risk) is taken on expectation as
L L(w)
Y).
(2)
Note that we take an additional expectation over p(w p(w|S) (S) because we are evaluating risk of the learned posterior instead of specific value of parameter w. In addition, we call p (W[S) posterior here for convenience while it is not the Bayesian posterior that is computed through Bayes theorem
p(w|S) p(w) The PAC-Bayes bounds hold even if prior p(w) is incorrect and posterior
p(5)
p( (S) is arbitrarily chosen.
In practice, we only own finite samples S. This gives rise to the empirical risk as
Ls(w)= (p(w)s)
e(f (x.))
With the above Eqs. (2) and (3) at hand, the generalization gap of the learned posterior p(w]S) in out-of-sample test is (Lw A L(w) - [s(w). Xu & Raginsky (2017) proposed a PAC-Bayes bound based on the information contained in weights 1(w; S) that
200 T(W: S),
20"
Epts) (T(() - L. [.((w)] <
n
when {(f"(X),Y) Y) is o-sub-Gaussian. series of follow-ups tightened this bound and verified is an effective measure of generalization capability of learning algorithms (Mou al., 2018; Negrea
AA casy way to fulfill this condition is to clip the lost function to E [0. a] hence satisfies ş g-sub-Gaussian (Philippe, 2015; Xu & Raginsky, 2017).
2

---

Published as a conference paper at ICLR 2022
al., 2019; Pensia et al., 2018: Zhang al., 2018). Therefore, it natural to build new information bottleneck grounded on this PAC-Bayes generalization measure, namely the -CCyB information bottleneck (PIB),
min LPB Ls(w) (I: p(w[S)
(5)
For classification term, the loss term Lg(w) becomes the cross-entropy between the prediction P(Y|X, w) and the label p(Y)X), hence PIB Eq. (5) equivalent to
max I [(w; YX,,) (); p(w (S)
(6)
which demonstrates trade-off between maximizing the sufficiency (the information of label Y con- tained in w) and minimizing the minimality of learned parameters w (the information of dataset S contained in w). Unlike previous IB based on representations, are not directly influenced by inputs and selected activation functions. Likewise, the trade-off de- scribed by PIB objective is more reasonable since its compression term is explicitly correlated to generalization of NNs.
PIB is built on weights that
ESTIMATING INFORMATION STORED IN WEIGHTS
In this section, present new notion of IIW I(w; S) built on the Fisher information matrix that relates to the flatness of the Riemannian manifold of loss landscape. Unlike Hessian eigenvalues of loss functions used for identifying flat local minima and generalization but be made arbitrarily large (Dinh et al., 2017) 7), this notion is invariant to re-parameterization of NNs (Liang et al., 2019). Also, our measure invariant to the choice of activation functions because enced by input X like / (T;; We leverage to monitor the information trajectory of NNs trained by SGD and cross entropy loss and verify it capable of reproducing the two-phase transition for varying activations (e.g., ReLU, linear, tanh, and s sigmoin i d) in $5.1.
directly influ-
3.1 CLOSED-FORM SOLUTION WITH GAUSSIAN ASSUMPTION
By deriving a new information bottleneck PIB, we can look into how IIW [(w; and Ls(w) evolve during the learning process of NNs optimized by SGD. Now the key challenge ahead is how to estimate the IIW T(w: S), as
1(w; S) = Ep()) (LL || p(w))]
the expectation of Kullback-Leibbe (KL) divergence between p(w S) and p(w) over the distribu- tion of dataset p(S). And, p(w) is the marginal distribution p (|[) as pl p(w) A Rpp(s [p[w]S))
When we assume both p(w) = (w [f,, Σg) and p(w[S) -N(w w 0s, Es) are Gaussian distribu- tions, the KL. divergence term in Eq. (7) has closed-form solution as
KL(p(w(S) [[w()) 2 log det Es
+ (0s 0,) T ' - 0o) + tr (En Es) (E)'Es)
(8)
det En
Here del A and tr(A) are the determinant and trace of matrix A, respectively; D is the dimension of parameter and is constant for specific NN architecture; Os are the yielded weights after SGD converges on the dataset S. the covariances of prior and posterior are proportional, the logarithmic and trace terms in Eq. (8) both become constant. Therefore, the mutual information term is proportional to the quadratic term as
I I(w: S) xx Ep(s) [(Ss 0,) 105(00000)) (0s 0b)] = Ep(s) [ 0 Σ 0
In the next section, will see how prior covariance Co.
Assuming the same covariance for the Gaussian randomization of posterior and prior is common practice of building PAC-Bayes bound for simplification, (Dziugaite & Roy, 2018; Rivasplata al., 2018).
3

---

Published as a conference paper at ICLR 2022
3.2 BOOTSTRAP COVARIANCE OF ORACLE PRIOR
Since the computation of exact oracle prior Eo needs the knowledge of p(S) 3, we propose to approximate it by bootstrupping from S as
1 5(0s Oss 0s)(Os. Oz) K
Eo pS) [00g (o)(0s 0o)
(10)
where Sy is bootstrap sample obtained by re-sampling from the finite data S, and Sk p(S) still valid sample following P(S). Now we are closer to the solution but the above term is still troublesome calculate. For getting [0s,)-m we need optimize on a series of (K times) bootstrapping datasets (S via SGD until it converges, which is prohibitive in deep learning practices. Therefore, we propose to approximate the difference - 0g by influence ceffuencefunccion functions drawn from robust statistics literature (Cook & Weisberg, 1982; Koh & Liang, 2017: Wang al., 2020c;b).
Lemma (Influence Function (Cook & Weisberg, 1982)), Given dataset S = {(X,, Yo) and the parameter 0₂ - argming Lg (0) argming n EE= (0)² that optimizes the empirical loss function, if we drop sample (Xj, Y2) in S to get jackknife sample Sig and retrain our model, the new parameters are 0s,, argming Lss, (0) argming 1. mation of parameter difference Os Og is defined by influence function ub, as
n. Li= (0) (0), The approxi-
05, - 0g ~ bb where ub; - (05) ER.
(11)
and H A
2 A. (0s) E DRD is Hessian matrix.
The application of influence functions be further extended to the case when the loss function perturbed by vector f ,., (x) Rn 0g,ç argming Ei=1 (.)( In this scenario, the parameter difference can be approximated by
ŌS.E 0g
(&: 1) 1:
LT T (E 1),
(12)
1 is a combination of all influence functions eb; 1 (1, 1,,,,...11 1) an n-dimennional all-one vector. Lemma gives rise to the following lemma on approximation of oracle prior covariance in Eq. (10):
where v = (7/1, 222..................... )
Lemma 2 (Approximation of Oracle Prior Covariance). Given the definition of influence functions (Lemma 1) and Poisson bootstrapping (Lemma A.2), the covariance matrix of the oracle prior can be approximated by
K
Eg = Ep(s) [(0s (0s - 0o) -0)
1
Σ 0gt
(0g - 0)
1 -H n
FA
FF
K
(13)
where Fg is Fisher information matrix (FIM); we omit the subscript S of Os and Ôs.e for notation conciseness, and th is the bootstrap resampling weight the k-tth experiment.
Please refer to Appendix A for the proof of this lemma.
3.3 EFFICIENT INFORMATION ESTIMATION ALGORITHM
After the approximation oracle prior covariance, we are now able to rewrite the IIW term ((w;S) in Eq. (9) to
T(w; S) x nEp(s) [00g 0p) (0 0n)] n00 0n) 0T) [F
0p) T(w; S). (14)
"2 is called oracle prior because it minimizes the term I(w; S) = Epc9) (p((p())S) || p p(W) when p(w)= Ep(s) [P(W]S]] (Dziugaite et al., 2021).
*Note (0) is not the expected empirical risk (s(w) in Eq. (3): instead, is the deterministic empir- ical risk that only relates to the mean parameter 0. We also denote (X,)) Yi) by (a()) for the notation conciseness.
4

---

Published as a conference paper at ICLR 2022
Algorithm 1: Efficient approximate information estimation of I(w;
Data: Total number of samples nt, batch size B, number of mini-batches in one epoch To.
number of information estimation iterations T1, learning rate 7, moving average
hyperparameters and K, sequence of gradients set V C = 0
Result: Calculated approximate information (w:5)
Pretrain the model by vanilla SGD to obtain the prior mean 0p:
for /= do
L Ve B E, (v(00-1). O₂ 0 V
/* Vanilla SGD */
VL + vLU(VL);
/ Store gradients */
0, + V p0% + Le
Ek=0 0 K
/+ Moving average */
end
TT. for 1=1:7 do
A0
tg.
AF 0;
AF,- f AF 1-1 (AO VL,)²
/* Shorage-friendly computation */
end
n AFT; 7(w; 5) - Ti
Algorithm 2: Optimal Gibbs posterior inference SGLD.
Data: Total number samples n. batch size B, learning rate temperature / Result: A sequence of weights {We)r2k following p(w(S*) repeat
/x Mini-batch gradient of energy function
+ (- E. logg((/|X,,, w:-1) - B logr(w;-1)) SGLD by gradient plus isotropic Gauss ian noise sse
W(elo, In), + Wi-1 Th- I---US- 7Us-(W.-1)) + 11.
/x Learning rate & temperature decay
n (n(n-1)) + 6g(
until the weight sequence (wiheg becomes stable;
We define the approximate information by T(W;S) where approximate the expectation
Ep(s) 03 by 05 = ~ Σ
(V x
Σ
Σ
R
02 D.k
In Eq.
(14), the information consists of two major components: A = 0s - E R and F DXD which can easily cause out-of-memory error due the sich-dimensional matrix product operations. We therefore hack into FIM to get
T Σ Vo
7(w:5)= = TAAT
TaOT
AO
Σ AA vo)
(15)
such that the high dimensional matrix vector product reduces to vector inner product. We encapsu- late the algorithm for estimating IIW during vanilla SGD by Algorithm I.
BAYESIAN INFERENCE FOR THE OPTIMAL POSTERIOR
Recall that we designed a new bottleneck on the expected generalization gap drawn from PAC- PAC-Bayes theory in $2, and then derived approximation of the IIW in $3. The two components of PAC- Bayes IB in Eq. (6) are hence tractable as a learning objective. We give the following lemma on utilizing it for inference.
Thhe quadratic mean is closer to the true value than arithmetic mean because of the quadratic term within the expectation function.
5

---

Published as a conference paper at ICLR 2022
ITW linar MLP
IW of tark MLP
BW relk MLP
By af's gread MLP
1.0900
00:00:
layer layer
loyer
62001
Fiiing
100005
layer 2
(10003
layer
5,000
Compressing
11000
layer
000s
21035
layor
In spoch
3.000
10
0:030
001
3.5
75
spoch
<poch
<pcc7
2.5
2.55
06
loss
train.
2.0
150
test
test
0.7
07
Inflection point
Lo
1.3
3.00
10
000
3.5 5:11
7.5
000
spoch
opoch
opoch
<pcc7
IIW (up). loss and accuracy (down) of different activation functions (linear, tanh, ReLU, and s sigmoi d) NNs. There is clear boundary between the initial fitting and the compres- sion phases identified by IIW. Meanwhile, the train loss encounters the inflection point that keeps decreasing with slower slope. Note that the learning rate is set small (1e-4) except for s sigmoion i-NN for the better display of two phases.
Figure 1:
Lemma 3 (Optimal Posterior for PAC-Bayes Information Bottleneck). Given an observed dataset S", the optimal posterior p( of PAC-Bayes IB in Eq. (5) should satisfy the following form that
1 p( (w) exp { Ls- (is-(w))
{
1 Z(S)
exp 1-1
(
p(w(s*)=
(16)
Z(S)
where Ug- (w) is the energy function defined Us- (w) Ls- (w) -3 p(w), and Z(S) is the normalizing constant.
Please refer to Appendix for the proof. The reason why we write the posterior in terms of an exponential form is that it is typical Gibbs distribution (Kittel, 2004) (also called Boltzmann dis- tribution) with energy function Us- (w) and temperature / (the same 3 of PIB appears in Eq. (5)). Crediting to this formula, we adopt Markov chain Monte Carlo (MCMC) for rather efficient Bayesian inference. Specifically, we propose to use stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011) that has been proved efficient and effective in large-scale posterior inference. SGLD can be realized by simple adaption of SGD as
Wk+1 = Wk TRgs V 2m
(17)
where 7 is step size, E/ ~ O, is a standard Gaussian noise vector, and gs is an unbi- ased estimate of energy function gradient (WW).. SGLD can be viewed diffusion described by stochastic differential equation (Raginsky et al., 2017; Borkar & Mitter, 1999): dw(()
discrete Langevin
(UU(((t)) + v23dB(t), where (B(t))e2o is the standard Brownian mo- tion in RR.. The Gibbs distribution IT <(w) x exp(-) (ww)) is the unique invariant distribution of the Langevin diffusion. And, distribution of w, converges rapidly to x(w) when t co with sufficiently small 3 (Chiang et al., 1987). Similarly for SGLD in Eq. (17), under the conditions that D: Ttt + oo and Er T/2 + O, and an annealing temperature B, the sequence of {wehazk converges to Gibbs distribution with sufficiently large k.
As we assume the oracle prior p( w) = -N(w 0o,, Eo). logp(w) satisfies
llogp(w) x (w 0,) Σ, (w- 0o) + log(det Eo).
(18)
The inference of the optimal posterior is then summarized by Algorithm 2. (.((.) and Φs()) are learning rate decay and temperature annealing functions, e.g., cosine decay, respectively. Our SGLD based algorithm leverages the advantage of MCMC such that it is capable sampling from the op- timal posterior even for very complex NNs. Also, it does need to know the groundtruth distribution of S while still theoretically allows global convergence avoiding local minima. It can be realized with a minimal adaptation of common diteeefifrentiiiionnn packages, e.g., PyTorch (Paszke et al., 2019), by injecting isotropic noise in the SGD updates. Please refer to Appendix for the details of computation for PIB object.

---

Published as a conference paper at ICLR 2022
5 EXPERIMENTS
In this section, we aim to verify the intepretability of the proposed notion of HIW by Eq. (15). We monitor the information trajectory when training NNs with plain eross entropy loss and SGD for the sake of activation functions ($5.1), architecture ($5.2), noise ratio ($5.3), and batch size ($5.4).
We also substantiate the superiority of optimal Gibbs posterior inference based on the proposed Algorithm 2, where PIB instead of plain cross entropy is used as the objective function ($5.5). We conclude the empirical observations in $5.6 at last. Please refer to Appendix D for general experimental setups about the used datasets and NNs.
5.1 INFORMATION WITH DIFFERENT ACTIVATION FUNCTIONS
We train 2-layer MLP (784-512-10) with plain cross-entropy loss by Adam on the MNIST dataset, meanwhile monitor the trajec- tory of the HI S(. Results are illustrated in Fig. I where different actívation functions (linear, tanhh, ReLU, and sigmoidi d) are testified. We identify that there a clear bound- ary between fitting and compression phases for all of them. For example, for the linear ac- tivation function on the first column, the IIW I ((w;S) surges within the first several itera- tions then drops slowly during the next itera- tions. Simultancously, we could see that the training loss reduces sharply at the initial stage, then keeps decreasing during the information compression. At the last period of compres- sion, we witness the information fluctuates near zero with the recurrence of memorization phe- nomenon (IIW starts to increase). This implies that further training is causing over-fitting IIW shows great universality on representing the infor- mation compression of NNs.
IIW of 1 layer MLP
of 2-leper
layer
Tager
0.0
1000 iteration
2001
iteration
IIW of 3-layer MLP
IIW 4-layer MLLP
layer
layer
itenation
iteration
Figure 2: Information compression with varying number of layers (1, 2, 3, and 4) for ReLU MLPs. All follow the general trend of fitting to compres- sion phase transition. And, deeper layers can ac- celerate both the fitting and compressing phases.
5.2 INFORMATION WITH DEEPER AND WIDER ARCHITECTURE
Having identified the phase transition of the 2-layer MLP corresponding to IIW 1(w; S), we test it under more settings: different architectures and different batch sizes. For the architecture setting, design MLPs from 1 to 4 layers. Results are shown in Fig. 2. The first and the last figures show the information trajectory of the 1-layer/4-layer version of MLP-Large (784-10784-512-0004 where clear two-phase transitions happen in all these MLPs.
The 1-layer MLP is actually a softmax regression model. It can be identified that this model fits and compresses very slowly w.r.t. IIW compared with deeper NNs. This phenomenon demonstrates that deep models have overwhelmingly learning capacity than shallow models, because deep layers can not only boost the memorization of data but also urges the model to compress the redundant information to gain better generalization ability. Furthermore, when we add more layers, the fitting phase becomes shorter. Specifically, we observe the incidence of overfitting at the end of the 4-yer MIP training as IIW starts to increase.
We also examine how IIW explains the generalization w.r.t. the number of hidden units, a.k.a. the width of NNs, by Fig. 3. We train a 2-layer MLP without any regularization on MNIST. The left panel shows the training and test errors for this experiment. Notably, the difference of test and train acc can be seen indicator of the generalization gap in Eq. (4). IIW should be aligned to this gap by definition. While 32 units are (nearly) enough to interpolate the training set, more hidden units still achieve better generalization performance, which illustrates the effect of overparameteriizatoo. In this scenario, weights 4 no-m keeps increasing with more units while IIW decays, similar to the test error. We identify that more hidden units do not render much increase of IIW, which is contrast
7

---

Published as a conference paper at ICLR 2022
2.0
Trar/lest Wisth
IIM612
INT COTSupEed LI0SE
LEAA d sdgi 0003
:30
Cereriitaaiio $30 oZ
1.4
IWW
2.0
nom
12
220
number hidden units
$320
$352
number hidden units
05 00
Figure 3: Left: Training and test accuracy w.r.t. # units; Right: Complexity measure (IIW and l2-no (--nom) w.r.t. # units. Blue dash line shows the gap between train/test acc (general- ization gap). We find f--norm keeps increas- ing with more hidden units. Instead, IIW keeps pace with the generalization gap: the larger the
Figure 4: Left: IIW, train, and test accuracy when noise ratio in labels changes. IIW rises when noise ratio grows. Right: IIW with vary- ing size of random-label data. Test acc keeps constant while train acc decays. Hence, more data causes lower IIW because of the shrinking gap between the train and test accuracy.
gap, the larger the IIW.
Nost and to babch size
HIW batch sioe
vanilla 12
0.0075
5 8as
0.0000
dropout
PIB
Train 80c
6.000
1601
0.6
256
756
number of batch sire
NUMOO batch sire
spoch
Figure 5: Left: Training and test accuracy W.r.t. # batch size; Right: IIW w.r.t. # batch
Figure 6: The tracked IIW of the VGG net dur- ing the training by four ways: vanilla, f--nrrm regularization. dropout, and PIB training. We can identify that: first, all of them still follow a cotrin-compressing paradigm specified by IIW; second, vanilla VGG reaches the largest IIW far above the others; third, PIB regularizes IIW di- rectly thus yielding the smallest IIW.
size. We find IIW keeps pace with the gener- alization gap: the larger the gap, the larger the IIW. From IIW we can specify that 16 is the best which reaches the least generalization gap without the need of having the model tested.
to the intuition that wider NNs always have larger information complexity. More importantly, find IIW is consistent to the generalization gap on each width.
5.3 RANDOM LABELS VS. TRUE LABELS
According to the PAC-Bayes theorem, the IIW is promising measure to explain/predict the gen- eralization capability of NNs. NNs are often werrpparaeeterized thus can perfectly fit even the random labels, obviously without any generalization capability. For example, 2-layer MLP has 15, 728, 640 parameters that are much larger than the sample number of CIFAR-10 (50,000). That causes the number of parameters an unreliable measure of NN complexity in overparameterization settings (Neyshabur al., 2015). Alternativlly, momm is often used as complexity measure to be imposed on regularizing model training in practices.
We investigate the model trained with different levels of label corruption, shown by the left panel of Fig. 4. We train a 2-layer MIP on CIFAR-10 and find that the increasing noise causes sharp test acc decay while train acc does not change much. Mcanwhile, IIW keeps growing with the fall of test acc and expansion of generalization gap. This demonstrates IIW's potential in identifying the noise degree in datasets or the mismatch between the test and train data distributions.
We further build random- CIFAR- 10, results are on the right of Fig. 4. Although the model can still nearly interpolate the train data, with the rise of random-label data, the model keeps 10% test acc but has less train acc, which renders larger generalization gap. This phenomenon is also captured by IIW.
5.4 INFORMATION IN WEIGHTS W.R.T. BATCH SIZE
We also consider how batch size influences IIW and generalization gap. Recent efforts on bounding I [(w; S) of iterative algorithms (c.g.. SGD and SGLD) (Mou et al., 2018; Pensia et al., 2018) imply that the variance of gradients is crucial factor. For the ultimate case where batch size equals
8

---

Published as a conference paper at ICLR 2022
Table 1: Test performance of the proposed PIB algorithm compared with two other common regu- larization techniques: f2-norm and dropout, on VGG-net (Simonyan & Zisserman, 2014). The 95% confidence intervals are shown in parentheses. Best values are in bold.
Test ACC (%)
CIFAR10
CIFARI00
STL10
SVHN
vanilla SGD
77.03(0.57)
52.07(0.44)
54.31(0.65)
93.57(0.67)
SGD++2-norm
77.1300.53)
50.84(0.71)
55.30(0.68)
93.60(0.68)
SGD+dropout
78.95(0.60)
52.34(0.66)
56.35(0.78)
93.61(0.76)
SGD++IB
80.19(0.42)
56,47(00.62)
58.83(0.75)
93.88(0.88)
full sample size, the gradient variance zero, and the model is prone to over-fitting grounded on empirical observations. When batch size equals one, the variance becomes tremendously large.
We conjecture there is an optimal batch size that reaches the minimum generalization gap, in other word, the minimum IIW. This conjecture is raised our empirical findings, displayed in Fig. where we test IIW on model with varying batch size. Each model is updated with the same total number of iterations and the same learning rate. We identify that when batch size is 16, the model reaches the best test acc and the least generalization gap, which means this optimal batch size should fall into (4,16) or (16, 64). On the left, the model reaches the minimum IIW when batch size is 16.
5.5 BAYESIAN INFERENCE WITH VARYING ENERGY FUNCTIONS
To confirm the superiority of the proposed PIB in $4, we compare it with vanilla SGD and two widely used regularizations: f--norm and dropout. We train a large VGG network (Simonyan & Zisserman, 2014) on four open datasets: CIFAR 10/100 (Krizhevsky et al.. 2009), STLIO (Coates et al., 20, and SVHN (Netzer et al., 2011), as shown Table I, where we find PIB consistantly outperforms the baselines. We credit the improvement to the explicit consideration of information regularization during the training, which forces the model to forget the training dataset to regularize the generalization gap. This is verified by Fig. 6 where PIB helps restrict IIW in the lowest level. Please refer to Appendix for experimental setups.
5.6 SUMMARY OF EXPERIMENTS
We made the following observations from our experiments:
1. We clearly identify the fitting- compresoomresioonn phase transition during training through new information measure, i.c., information stored in weights (IIW). Unlike the representation- based information measure T(X; Z), IIW applies to various activation functions including ReLU, a sigmoid, tanh, and linear.
2. We further identify that the phase transition applies to deeper and wider architecture. More importantly, deeper model is proved to reach faster fitting and compression than shallow models. Unlike f2-norm of weights that rise together with wider models, IIW better illustrates the true model complexity and its generalization gap.
4. IIW explain the performance drop w.r.t. the degree of label noise. Also, IIW can identify the generalization gap for models learned from random labels.
5. There might exist an optimal batch size for the minimum generalization gap, which is empirically demonstrated by our experiments.
Adopting SGD based on the energy function derived from PAC-Bayes enables good inference to the optimal posterior of NNs. This works for practical large networks in the literature.
CONCLUSION
In this paper, we proposed PAC-Bayes information bottleneck and the corresponding algorithm for measuring information stored in weights of NNs and training NNs with information principled reg- ularization. Empirical results show the universality of our information measure explaining NNs, which sheds light on understanding NNs through information bottleneck. We aim to further investi- gate its performance and develop this into practical NNs for production in the future.

---

Published as a conference paper at ICLR 2022
REFERENCES
Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. The Journal of Machine Learning Research, 19(1):1947-1980, 201
Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy
computation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(12):2897-22 2905, 2018b.
Vivek S Borkar and Sanjoy K Mitter. A strong approximation thcorem for stochastic recursive algorithms. Journal of Optimization Theory and Applications, 100(3):499-513 1999.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 8-VAE. arXiv preprint arXiv: 1844:1990,,3399,, 2018.
Nicholas Chamandy, Omkar Muralidharan, Amir Najmi, and Siddartha Naidu. Estimating uncertainty for massive data streams. Technical report, Google, 2012.
Tzuu-Shuh Chiang, Chii-Ruey Hwang, and Shuenn Jyi Sheu. Diffusion for global optimization in R", SIAM Journal Control and Optimization, 25(3):737-733, 1987.
Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In International Conference Artificial Intelligence and Statistics, pp. 215-223, 2011.
R Dennis Cook and Sanford Weisberg. Residuals and influence in regression. New York: Chapman and Hall, 1982.
Bin Dai, Chen Zhu, Baining Guo, and David Wipf. Compressing neural networks using the variational infor- mation bottleneck. International Conference on Machine Learning, pp. 1135 1144. PMLR, 2018.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Data-dependent PAC-Bayes priors differential privacy. In Advances in Neural Information Processing Systems, pp. 8440 8450, 2018.
Gintare Karolina Dziugaite, Kyle Hsu, Wascem Gharbich, Gabriel Arpino, and Daniel Roy. On the role of data in PAC-Bayes. In International Conference on Artificial Intelligence and Statistics, pp. 604-612. PMLR, 2021.
Bradley Efron. Bootstrap methods: another look at the jackknife. In Breakthroughs in Statistics, pp. 569-593. Springer, 1992.
Ziv Goldfeld, Ewout van den Berg, Kristjan Greenewald, Igor Melnyk, Nam Nguyen, Brian Kingsbury, and Yury Polyanskiy. Estimating information flow in deep neural networks. In International Conference on Machine Learning, 2019.
Anirudh Goyal, Riashat Islam, DJ Strouse, Zafarali Ahmed, Hugo Larochelle, Matthew Botvinick, Yoshua Ben- gio, and Sergey Levine. InfoBot: transfer and exploration via the information bottleneck.. In International Conference on Learning Representations, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization.
arXiv preprint
arXiv:1412.69800 2014.
Charles Kittel. Elementary statistical physics. Courier Corporation, 2004.
Pang Wei Koh and Percy Liang. Understanding black box predictions via influence functions. In International Conference on Machine Learning, pp. 1885-1894 PMLR, 2017.
Artemy Kolchinsky, Brendan D Tracey, and Steven Van Kuyk. Caveats for information bottleneck in determin- istic scenarios. In International Conference on Learning Representations, 2018.
Artemy Kolchinsky, Brendan D Tracey, and David H Wolpert. Nonlinear information bottleneck. Entropy, 21 (12):1181, 2019.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, (6(11))2278-283 1998.
10

---

Published as a conference paper at ICLR 2022
Xiang Lisa Li and Jason Eisner. Specializing word embeddings (for parsing) by information bottleneck. In Conference on Empirical Methods Natural Language Processing, pp. 2744-2754, 2019.
Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-Rao metric, geometry, and complexity of neural networks. In International Conference Artificial Intelligence and Statistics, pp- 988-896. PMLR, 2019.
James Martens. New insights and perspectives the natural gradient method. Journal of Machine Learning Research, 21:1-76, 2020.
Wenlong Mou, Liwei Wang, Xiyu Zhai, and Kai Zheng. Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints. In Conference Learning Theory, pp. 605-638 PMLR, 2018.
Jeffrey Negrea, Mahdi Haghifam, Gintare Karolina Dziugaite, Ashish Khisti, and Daniel M Roy. Information- thcoretic generalization bounds for SGLD via data-dependent estimates. arXiv preprint 11.00111 2019.
Yuval Netzer, Tao Wang. Adam Coates, Alessandro Bissacco. Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. In International Conference on Learning Representations Workshop, 2015.
Ziqi Pan, Li Niu, Jianfu Zhang, and Liqing Zhang. Disentangled information bottleneck. arXiv preprint 20X2:2001..03372, 2020.
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, hegh-performance deep learning library. In Advances in Neural Information Processing Systems, pp. 8024-8035, 2019.
Ankit Pensia, Varun Jog, and Po-Ling Loh. Generalization error bounds for noisy, iterative algorithms. In IEEF International Symposium on Information Theory, pp. 546-550. IEEE, 2018.
Rigollet Philippe. High- high-Diimessooal Statistics, chapter 1. Massachusetts Institute of Technology: MIT Open- CourseWare, 2015.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient Langevin dynamics: a nonasymptotic analysis. In Conference on Learning Theory, pp. 1674-1703. PMLR, 2017.
Omar Rivasplata, Emilio Herrado-Heránder John Shawe- Taylor, Shiliang Sun, and Csaha Szepesvári. PAC- Bayes bounds for stable algorithms with instance-dependent priors. In Advances in Neural Information Processing Systems, pp. 9234-9244, 2018.
Andrew M Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan D Tracey, and David D Cox. On the information bottleneck theory of deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2019(12): (10112)::222000 140200, 2019.
Ravid Shwartz-Ziv and Alexander A Alemi. Information in infinite ensembles of inie neural net- works, In Symposium on Advances in Approximate Bayesian Inference, pp. 1-17. PMLR, 2020.
Ravid Shwartz-ZZi and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint / 000/0, 2017."
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 40Xiv: 1556, 2014.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In IEEE Informa- tion Theory Workshop, pp. 1-5. IEFF, 2015.
Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004057. 2000.
Qi Wang, Claire Boudreau, Qixing Luo, Pang-Ning Tan, and Jiayu Zhou. Deep multi-view information bottle- neck. In Proceedings of the SIAM International Conference on Data Mining, pp. 37-45. SIAM, 2019.
Zifeng Wang, Xi Chen, Rui Wen, Shao-Lun Huang, Erean E. Kuruoglu, and Yefeng Zheng. Information theoretic counterfactual learning from missing-not--t-raadoo feedback. In Advances in Neural Information Processing Systems, 2020a.
11

---

Published as a conference paper at ICLR 2022
Zifeng Wang, Rui Wen, Xi Chen, Shao-Lun Huang, Ningyu Zhang. and Yeleng Zheng. Finding influential instances for distantly supervised relation extraction. arXiv preprint wxivv2209.098841 2020b.
Zifeng Wang, Hong Zhu, Zhenhua Dong. Xiugiang He, and Shao-Lun Huang. Less is better: Unweighted
data subsampling via influence function. AAAI Conference on Artificial Intelligence, volume 34, pp. 6340-6347, 2020c.
Max Welling Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference on Machine Learning, pp. 681 188, 2011.
Tailin Wu, Hongyu Ren, Pan Li, and Jure Leskovec. Graph information bottleneck. arXiv:2010. 12111, 2020.
arXiv preprint
Aolin Maxim Raginsky. Information-theoreee theoretic analysis of generalization capability of learning algo- rithms. In Advances in Neural Information Processing Systems, pp. 2521-2530, 2017.
Jingwei Zhang, Tongliang Liu, and Dacheng Tao. theforettion--heoretiic view deep learning. arXiv preprint arxiv: 1804 09060, 2018.
12

---

Published as a conference paper at ICLR 2022
A
PROOF OF LEMMA A.3
Before the proof of Lemma A.3, we need to introduce lemma by Martens (2020) as:
Lemma A.1 (Approximation of Hessian Matrix in NNs (Martens, 2020)). The Hessian matrix of NNs on a local minimum A can be decomposed based on Fisher information matrix as
1 n C TI
H = F
(A.1)
where C is the total number of classes, ý is the output (or prediction) of the network given input x, and fl. the Hessian of the e-th component of y. Specifically, for well-trained NN, we could
have V (0) = 0 thus g ~ Fø.
We further introduce a lemma of Poisson bootstrapping as
Lemma A.2 (Poisson Bootstrapping (Efron, 1992: Chamandy et al., 2012)). Given an infinite num- ber of samples, the bootstrap resampling weight has the property that limg Binomial (n, +) Poisson(1). This approximation becomes precise in practice when n 2 100. Also, we know [[g.] 1 and Varls. 1 by the definition of Poisson distribution when large enough.
When bootstrap resampling from dataset S, each individual sample Z; (X,, Y;) has a probability of | being picked, causing the weight & to follow binomial distribution & ~ Binomial (n, 4 As a result, all weights & follow a multinomial distribution as & ~ Multinomial (n, ni m.-... n) with the total number of samples constrained to be n. When it comes to big data, i.e., is pro- hibitively large, this multinomial resampling thus becomes rather slow. Based on Lemma A.2, we can now start to prove our lemma of approximating oracle prior covariance.
π
Lemma A.3 (Approximation of Oracle Prior Covariance). Given the definition of influence functions (Lemma 1) and Poisson bootstrapping (Lemma A.2), the covariance matrix of the oracle prior can be approximated by
K
Σo Pps) [00s (0) - 0] =
1 (0=
1 K
Σ 0cz
-0)
-0)
1
H
F-
(A.2)
where Fğ is Fisher information matrix (FIM); we omit the subscript S of 0s and os.. for notation conciseness, and th is the bootstrap resampling weight in the k-th experiment.
Proof. Recall that in the k-th bootstrap resampling process, the loss function is reweighted by &
1 ......... .....) T. Also, we have an influence matrix d = (151, ,2,...,, 3xx) E R The original risk minimizer on the full dataset S is
1 n
i g e argmin
Σ (0)) TI
(A.3)
and the reweighted empirical risk minimizer (after bootstrapping) is defined by
Ôs.c A
argmin - (,)
(A.4)
0
2=1
where we omit the subscript A for the sake of conciseness. Given the definition of influence function from Lemma 1, the difference between the two risk minizers above, 0g and A can be written as
(E+ 1)b. = 1 T (E 1).
05.5 0g n
(A.5)
As result, the oracle prior can be transformed as
Σn ~ Ep(s)
(Ōs.s (s) 0s)T]
(A.6)
(LVI (E (1) ( (LVT(E (5-11
~ Ep((s)
(A.7)
nス Ep(s) [VT(E 1)( TV..
(A.8)
13

---

Published as a conference paper at ICLR 2022
Furthermore, based on the definition of influence function, know v/ V in Eq. (A.8) can be further writen to
-T
0. The term
Eg n² P()) [ (E 1)(5 - 1) v] = n²
T
E (P) 55 JV.
(A.9)
From Lemma A.2 know E[E.] 1 and Varlgi] when > 100. We also know that
Pp()) pt))) = Ep(s) [4,] Var[Ei] 13 2,
i i J.
(p()))iE)) {
i=
This gives rise to the final solution that
V E
"
v (11 +I,) v
(A.10)
(A.11)
(A.12)
=nH
V Vel(()]oo (0)
H
(A.13)
FFgHT
(A.14)
nF
(A.15)
I, Eq. (A.10) is an identity matrix with size n n: Eq. (A.11) true by re-applying the property of influence functions that v 1 = 1 TT O; and Eq. (A.15) is achieved by the result from Lemma A. Summarizing all the above results, the oracle prior covariance can be approximated by
nF Eo
F-1
(A.16)
n
72
0
B PROOF OF LEMMA
Lemma B.1 (Optimal Posterior for PAC-Bayes Information Bottleneck). Given an observed dataset S", the optimal posterior p( //w S*) of PAC-Bayes IB in Eq. (5) should satisfy the following form that
(-bls-(()) = Z(S) exp {-1
{-3 (Us- ((s-(w)) {w))]
P(w[S*)= = Z(S) w exp
(B.1)
where Us- (w) the energy function defined by
Us. (w) Lss (w) logg(ww)
(B.2)
and ZSS) is the normalizing constant.
Proof. Recap that the PAC-Bayes information bottleneck in Eq. (5) is
min Lrm Ls(w) + (w: S).
(B.3)
p (w
Given an observed dataset S", our object of interest is to find the optimal posterior p(w[S*) that minimizes the LpIB. Consider constraint of posterior distribution that
f p(w Sidw 1, AS p(,,))(X Y)
(B.4)
can formulate the problem by
min LPBB T.s(w) (3((w: (W: S): p(w[S)
(B.5)
s.t. / p(w]S)dw=
14

---

Published as a conference paper at ICLR 2022
A Lagrangian can hence be built to relax the above optimization problem by
min LPm [.s(w) + (T(:::)) f as (o(w).) 1) dwds p(w (8)
f p(w]s) [ss(w] dw 3 (lopp((W)) - log p()p dwds
(B.6)
+ 1 as (p(w)() dwds,
with a fassvs P(X,)) } corresponding to Lagrange multipliers; we denote the empirical risk by Ls(w) Ei= (:(w).
Diffcrentiating EPIB w.r.t p(w| results in
V Viwss-> Cru = Ls- (w) + 3 po (S*) - 3 log p(w) + 3 as
(B.7)
Setting V yp(w]s+) Spm and solving for p(w/5*) yields
logp(w(s-) -3 L.s- (w) -logp(() as
(B.8)
p(w(S*) p(w) exp LS- (w) exp
1
0 -1 1-1- 3
The second exponential term cxp {11 aga is the partition function that normalizes the poste- rior distribution. Denoting the normalization term as Z(S), we hence obtain the optimal posterior solution as
(-glss(())
p(w(S*) [())p(())))) exp
7(S)
(B.9)
Ls- (w) Blg plog
Z(S) exp
C COMPUTATION FOR PIB OBJECT
Our Alg. presents the training process based on the proposed PIB objective function. It differs from vanilla SGD two aspects: (1) the objective function (also called energy function U) consists of the negative log-likelihood plus a regularization term (og((w) (line 2); (2) the parameter w gets update by the gradient of the energy function plus an isotropic Gaussian noise (line 3). Since (2) has no additional computational cost, the major change is the regularization term log l(w((w) (1), described in Eq. (18) as
log p w) x (w 00)
0,) + log(det Eo). ' -o
(C.1)
The first term is just matrix vector product based the approximation of Ep in Eq. (13). And the second term can be written by
D
log(det Eo) =
Σ log A4
(C.2)
where A is the eigenvalue of Ea. which can be obtained by efficient eigen-decomposition tech- niques.
D EXPERIMENTAL PROTOCOL
All experiments are conducted on MNIST (LeCun et al., 1998) or CIFAR-10 (Krizhevsky et al., 2009). We design two multi-layer perceptron (MLPs): MLP-Small and MLP-Large, where MLP- Small is a two-layer NN, ie.. 784(3072)-512-10, and MLP-Large a five- NN, ie., 784(3072)- 100-80-60-40 -0-0000-00--0000 10. The number of input units is 784 with permutation MNIST inputs or 3072 with
15

---

Published as a conference paper at ICLR 2022
permutation CIFAR-10 inputs. For the performance comparison in $5.5, reduced version of VGG-Net where two blocks are cut due to the memory constraint. In the general setting, we pick Adam optimizer (Kingma & Ba, 2014) to accelerate the convergence of NN training. We use one RTX 3070 GPU for all experiments.
Specifically for the Bayesian inference experiment, the batch size is picked within {8, 16, 32, 64, 128, 256, 512); learning rate in {le-1, le ,,,,,,,__ 1e=";; weight decay of 2-- norm is in {le leee. lee :ee=); }; noise scale of SGLD is in {le llee le 1e-10;; of PAC-Bayes IB is in (le-1,11 1e 1,,1, le 1..): and the dropout rate is fixed as 0.1 for the dropout regular- ization.
16


Reviews:
[{"review_id": "lNhZYBsPjkm", "reviewer": "Reviewer_F2VH", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths:\n- The paper proposes an exciting general principle of deep learning. As far as I know, the contributions here are novel and will be of high interest to the community. \n- The authors build on previous work by showing how their IB objective addresses the shortcomings of previous work in this area.\n- It is very well written. This is a highly-technical paper, and the details are presented in a careful and thoughtful way.\n- The experiments are well done and the results support the conclusions. Specifically, this objective is motivated by a PAC-bound (the tightness of which is not clear) and various approximations are used to estimate I(w;S) (the accuracy of these are not immediately clear). The experiments address these issues by showing that the motivation and the approximations are reasonable. \n\nWeaknesses:\n- The limitations of this method are not discussed clearly. For example, the paper provides an algorithm for sampling from the weight posterior p(w|S), but how does this compare computationally to standard training of a neural network, or a estimating the posterior in a Bayesian Neural Network? \n- There are some minor grammatical and spelling typos throughout, e.g. \"infection point\".\n", "summary_of_the_paper": "This paper proposes a new version of the Information Bottleneck objective for training neural networks. This is in part motivated by previously-derived PAC-Bayes bounds on the generalization error that are proportional to the square root of the mutual information of the weights and the training dataset: I(w;S). Thus this new information bottleneck objective attempts to minimize both the empirical risk and this mutual information.The paper derives a computationally-tractable algorithm for estimating I(w;S), then this algorithm is used to show that this quantity is inversely correlated with generalization loss on a variety of neural network architectures. ", "summary_of_the_review": "An excellent paper with exciting ideas, clear presentation, and technical depth. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "ilmFyAucarL", "reviewer": "Reviewer_HBHz", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strength:\n\nThis paper addresses an interesting and important problem. The proposed PIB is novel. The experiments show that the proposed \\tilde{I}(w;S) correlates with the generalization gap, and helps improving the performance.\n\nWeakness:\nIn order to make the computation of I(w;S) tractable, the authors make several important assumptions. It would strengthen the paper a lot if the paper discuss and perform experiments to show if the assumptions are valid, in the experiments the authors run.\n\nFurthermore, in section 5.5, can the authors show the generalization gap (together with the train and test acc), with different regularization? Ideally we should see that with PIB as the objective, the generalization gap is much smaller than the other methods. With this, we can then be confident that the improvement is due to reduced generalization gap instead of better training.", "summary_of_the_paper": "This paper introduces the PAC-Bayes Information Bottleneck (PIB). Starting from the generalization bound Eq. 4 which shows that the generalization gap is upper bounded by a function of I(w;S), the authors proposes PIB which has an additional regularization term of \\beta I(w;S). Since the computation of I(w;S) is intractable, the authors then make several assumptions to simplify its computation, arriving at an estimate of I(w;S) by \\tilde{I}(w;S) (eq. 15), and use SGLD to compute it in practice. Experiments show that (1) there is a two-phase transition in SGD training as indicated by \\tilde{I}(w;S); (2) \\tilde{I}(w;S) seems to correlate with the generalization gap, under different variations of experiment hyperparameters: number of hidden layers, noise ratio, and number of random label data; (3) it improves performance compared to l2 and dropout regularization. ", "summary_of_the_review": "In summary, this paper is novel, but the experiment should be strengthened as detailed in the main review.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "K-lOGqto8Hu", "reviewer": "Reviewer_gEEc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "The paper is generally well written and treats an interesting and timely topic. The idea to limit the information about the sample that is contained in the weights is not new (the authors cite several works that bound the generalization error via this information), but this is the first time that I have seen a corresponding cost function implemented in practice. There are, however, a few issues that are not perfectly clear to me:\n\n- The authors cite the literature stating that the generalization gap is limited by I(S;W) if the loss is sigma-sub-Gaussian. Does this hold for the negative log-likelihood in (6)? Also, in (6) is S a random variable or not? (4) requires that I(S;W) is computed as an expectation over p(S), while the log-likelihood in (6) is an expectation over P(w|S), i.e., not over p(S) but over a concrete S. How can this be understood? \n- Connected to this, is it safe to call the resulting cost function an information bottleneck cost function? I assume that this is better called an IIW-regularization rather than an IB cost. The IB cost is a very specific formulation that combines a mutual information cost with a mutual information utility, whereas here we have a general cost with an additional mutual information cost as regularization term.\n- The authors correctly claim that I(X;T) becomes trivial if the network is deterministic. More precisely, this mutual information becomes infinite in many of these cases (see \"Learning Representations for Neural Network-Based Classification Using the Information Bottleneck Principle\" by Amjad and Geiger). I believe that this result carries over to I(S;W) being infinite for deterministic learning algorithms. This may not hold for all learning algorithms, but certainly for some. My own gut feeling suggests that I(S;W) is infinite for SGD with finitely many epochs (e.g., by the fact that there are only combinatorially many options to shuffle the batches), but that it is finite for SGLD, where noise is added to the weights. It is therefore not clear to me in which settings the analysis in Section 3 is a valid approximation. In other words, in which settings is the assumption that p(w|S) is Gaussian valid? Does it only hold for SGLD?\n- Connected to the point above: In which cases is the assumption that p(w) is Gaussian a valid approximation? \n- Can this Gaussian assumption about p(w) be used to bound I(S;W) from above? (E.g., for a Gaussian learning algorithm, can it be shown that the term I(S;W) is maximized if W becomes Gaussian as well? This would be at least intuitive from a channel coding perspective, where a Gaussian channel input is known to maximize the mutual information through a Gaussian channel, and which is then known to produce a Gaussian channel output.)\n- In Algorithms 1, pls. compare line 9 with your equation (15). In (15), you sum over squared inner products. In line 9 and 11, you square over the resulting sum of inner products. Is this difference intended, and if so, how can it be explained? Also, do we have $T_0 \\ge T_1$ in Algorithm 1?\n- In Fig. 1, why is the mutual information I(W;S) evaluated for different layers? What is the exact meaning of splitting the IIW between layers in terms of the generalization bound? I was assuming that the generalization bounds all consider the entire set of weights, and that the proposed PIB should do so as well. \n- Also in Fig. 1, the discussion of the inflection point is not fully clear. \n- In Section 5.1, it is claimed that the variance of the information explodes. Can this be made more precise (e.g., by writing down the mathematical symbol for this variance)? Furthermore, this is not shown in the figures, if I remember correctly.\n- In all figures, why is the mutual information I(S;W) so small? These numbers do not seem right. I would assume that it is necessary to \"learn\" more than 10⁻2 bits/nats to successfully solve a classification problem. In other words, while the general trend of IIW seems to be correct, I am not convinced of the correctness of the absolute numbers. Can you provide some intuition about these small numbers? Is this connected with the proportionality symbol in (14)? (But going from (8) to (9) it seems to be that additive constants are dropped, not multiplicative constants.)\n\nFor the sake of clarity, I would prefer that footnote 3 is in the main text. Also, in some instances the notation and terminology is not clear. E.g., is S sampled iid in (4)? Why is the \"oracle prior\" called an oracle? How exactly is the bootstrapping resampling weight \\zeta_k defined? Why is the temperature $\\beta$ called the annealing temperature just before (18)? At the end of Section 5.2 you write that the l2-norm keeps increasing -- the norm of what?", "summary_of_the_paper": "The authors propose a formulation of the information bottleneck problem, replacing the mutual information between input X and latent representation Z via the mutual information between the sample S and the weights W obtained from the sample. They derive closed-form solutions for this mutual information in the Gaussian setting and propose an SGLD scheme to optimize the objective. Using this objective and optimization algorithm, the authors investigate several interesting scenarios, including different activation functions and noisy labels.", "summary_of_the_review": "A very interesting paper, dealing with an interesting and timely topic. Unfortunately, the paper is not perfectly clear throughout all sections.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0COwuVECf47", "reviewer": "Reviewer_CyLM", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "**Strengths**\nTo my knowledge this paper has significant technical and empirical novelty.\nThe authors do a good job of summarizing previous work and differentiating their contributions.\nThis is not my area of expertise, but the derivation of PIB, the estimator for $I(\\mathbf{w}; S)$, the proposed optimal posterior all look novel and correct.\nThe experiments are well done, thorough, and support the main claims of the paper.\n\n**Weaknesses**\nThe main weaknesses of this paper are language and clarity.\nI would recommend a thorough Grammarly or perhaps external advice.\nThe graphs should report means and standard error intervals over multiple random seeds.\n\n**Some specifics**\n- Some technical terms are used before definition (e.g., phase transition in abstract)\n- In the abstract IB cannot explain, maybe IB theory can, as you use in the intro.\n- IIW and $I(\\mathbf{w}; S)$ are redundant, would recommend just using $I(\\mathbf{w}; S)$\n- \"Third, mutual information becomes trivial in deterministic cases.\" Please elaborate / cite.\n- \"(2) we derive a solution to the intractable...,\" can something intractable have a solution? maybe approximation is better.\n- \"optimal posterior of PIB,\" does PIB have a posterior or is the posterior over the weights?\n- Figure 2. IIW only shows compression phase, can the loss also be included in these plots?", "summary_of_the_paper": "The authors propose new interpretation of the Information Bottleneck (IB), dubbed PAC-Bayes Information Bottleneck (PIB).\n\nWhere the IB is defined wrt the mutual information between feature representations $T$ and either inputs $X$ or targets $Y$, PIB is defined wrt the empirical risk over the dataset $S = \\\\{X_i, Y_i\\\\}_{i=1}^n$ and the mutual information between model parameters and $S$, $I(\\mathbf{w}; S)$, or information stored in weights.\n\nThe authors show that PIB is a bound on generalization error.\n\nThe authors derive a tractable estimator for $I(\\mathbf{w}; S)$.\n\nThe authors present an approximate inference method for p(\\mathbf{w} \\mid S)  that utilizes the proposed PIB.\n\nThe authors show that PIB reflects the hypothesized two-phase fitting and compression modes of neural networks across different activation functions, network depth, and network width.\n\nThey show that $I(\\mathbf{w}; S)$ yields a good estimator of the generalization error that is robust to label noise.\n\nThey show that their inference method improves generalization across several benchmark datasets.", "summary_of_the_review": "To my knowledge this paper demonstrates significant technical and empirical novelty.\nI believe the main weaknesses can be addressed prior to publication.\nTherefore I recommend acceptance.\nHowever, I am not an expert on this topic, so my confidence is only a 2.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper revisits the information bottleneck principle, but in terms of the compression inherent in the weights of a neural network, rather than the representation. This gives the resulting IB principle a PAC-Bayes flavor. The key contribution is a generalization bound based on optimizing the objective dictated by this principle, which is then tractably approximated and experimentally verified. Reviews raise concerns about assumptions made to achieve the tractable version and a public discussion debates whether this is truly a PAC-Bayes bound. The authors address these adequately. Another concern is whether improvements in experiments can be ascribed to the new objective. Authors add new experiments in support of this. Additional concerns about the clarity of certain aspects of the paper were or were promised to be addressed by the authors. Overall, the perspective of this paper, its technical contributions, and experimental evaluations appear to be worthwhile to share with the community, as they advance the applicability of the information bottleneck principle.
Decision: Accept (Spotlight)