Original: PAC-B AYES I NFORMATION B OTTLENECK
Revised: PAC-B AYES I NFORMATION B OTTLENECK

Original: Anonymous authors
Revised: UIUC Tencent

Original: Paper under double-blind review
Revised: Tencent

Original: A BSTRACT
Revised: A BSTRACT

Original: Information bottleneck (IB) depicts a trade-off between the accuracy and conciseness ofencoded representations. IB theory has succeeded in explaining the objective and behavior of neural networks (NNs) as well as learning better representations.
Revised: 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.

Original: However, there are still critics of theuniversality of IB, e.g., the transition from ﬁtting to compression usually fades away, representation compression  is not causally related to generalization, and IB is trivial in deterministic cases.
Revised: 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.

Original: In this work, we build a new IB based on the trade-off betweenthe accuracy and  complexity of learned weights of NNs. We argue that this new IB represents a more solid connection to the objective of NNs since the information stored in weights (IIW) bounds their PAC-Bayes generalization capability, hence we name it as PAC-Bayes IB (PIB).
Revised: In this paper, we propose an algorithm for the efﬁcient approximation of IIW. Then, we build an IIW-based information bottleneck on the trade-off between accuracy and information complexity of NNs, namely PIB.

Original: On IIW, we can  identify the  phase transition phenomenon in general cases andsolidify the causality between  compression and  generalization.
Revised: From PIB, we can empirically identify the ﬁtting to compressing phase transition during NNs’ training and the concrete connection between the IIW compression and the generalization.

Original: We then derive a tractable solution of PIB and design a stochastic inference algorithm  by Markov chain Monte Carlo sampling. We empirically verify our claims through extensive experiments. We also substantiate the superiority of the proposed algorithm on training NNs.
Revised: Besides, we verify that IIW is able to explain NNs in broad cases, e.g., varying batch sizes, overparameterization, and noisy labels. Moreover, we propose an MCMC-based algorithm to sample from the optimal weight posterior characterized by PIB, which fulﬁlls the potential of IIW in enhancing NNs in practice.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: Information bottleneck (IB) (Tishby et al., 2000) was a promising candidate to reveal the principle of NNs through the lens of information stored in encoded representations of inputs.
Revised: Information bottleneck (IB) (Tishby et al., 2000) was a promising candidate to reveal the principle of NNs through the lens of information stored in encoded representations of inputs.

Original: Drawn from the conception of representation minimality and sufﬁciency in information theory, IB describes the objective of NNs as a trade-off, where NNs are abstracted by a Markov chain Y ↔ X ↔ T , as
Revised: Drawn from the conception of representation minimality and sufﬁciency in information theory, IB describes the objective of NNs as a trade-off, where NNs are abstracted by a Markov chain Y ↔ X ↔ T , as

Original: [Equation]
Revised: [Equation]

Original: where I ( T ; X ) and I ( T ; Y ) are the mutual information of representation T towards inputs X and labels Y , respectively.
Revised: where I ( T ; X ) and I ( T ; Y ) are the mutual information of representation T towards inputs X and labels Y , respectively.

Original: IB theory claimed (Tishby & Zaslavsky, 2015) and then empirically corroborated (Shwartz-Ziv & Tishby, 2017) that NNs trained by plain cross entropy loss and SGD all confront an initial ﬁtting phase and a subsequent compression phase.
Revised: IB theory claimed (Tishby & Zaslavsky, 2015) and then empirically corroborated (Shwartz-Ziv & Tishby, 2017) that NNs trained by plain cross entropy loss and SGD all confront an initial ﬁtting phase and a subsequent compression phase.

Original: It also implied that the representation compression is a causal effect to the good generalization capability of NNs.
Revised: It also implied that the representation compression is a causal effect to the good generalization capability of NNs.

Original: IB theorem points out the importance of representation compression and ignites a series of follow-ups to propose new learning algorithms on IB to explicitly take compression into account (Burgess et al., 2018;
Revised: IB theorem points out the importance of representation compression and ignites a series of follow-ups to propose new learning algorithms on IB to explicitly take compression into account (Burgess et al., 2018;

Original: Achille & Soatto, 2018b; Dai et al., 2018; Li & Eisner, 2019; Achille & Soatto, 2018a; Kolchinsky et al., 2019; Wu et al., 2020; Pan et al., 2020; Goyal et al., 2019; Wang et al., 2019; 2020a).
Revised: Achille & Soatto, 2018b; Dai et al., 2018; Li & Eisner, 2019; Achille & Soatto, 2018a; Kolchinsky et al., 2019; Wu et al., 2020; Pan et al., 2020; Goyal et al., 2019; Wang et al., 2019; 2020a).

Original: However, recent critics challenged the universality of the above claims.
Revised: However, recent critics challenged the universality of the above claims.

Original: At ﬁrst, Saxe et al.
Revised: At ﬁrst, Saxe et al.

Original: (2019) argued that the representation compression phase only appears when double-sided saturating nonlinearities like tanh and sigmoid are deployed.
Revised: (2019) argued that the representation compression phase only appears when double-sided saturating nonlinearities like tanh and sigmoid are deployed.

Original: The boundary between the two phases fades away with other nonlinearities, e.g., ReLU .
Revised: The boundary between the two phases fades away with other nonlinearities, e.g., ReLU .

Original: Second, the claimed causality between compression and generalization was also questioned (Saxe et al., 2019; Goldfeld et al., 2019), i.e., sometime networks that do not compress still generalize well, and vice versa.
Revised: Second, the claimed causality between compression and generalization was also questioned (Saxe et al., 2019; Goldfeld et al., 2019), i.e., sometime networks that do not compress still generalize well, and vice versa.

Original: To alleviate this issue, Goldfeld et al.
Revised: To alleviate this issue, Goldfeld et al.

Original: (2019) proposed that the clustering of hidden representations concurrently occurs with good generalization ability.
Revised: (2019) proposed that the clustering of hidden representations concurrently occurs with good generalization ability.

Original: However, this new proposal still lacks solid theoretical guarantee.
Revised: However, this new proposal still lacks solid theoretical guarantee.

Original: Third, mutual information becomes trivial in deterministic cases (Shwartz-Ziv & Alemi, 2020).
Revised: Third, mutual information becomes trivial in deterministic cases (Shwartz-Ziv & Alemi, 2020).

Original: Other problems encountered in deterministic cases were pointed out by Kolchinsky et al.
Revised: Other problems encountered in deterministic cases were pointed out by Kolchinsky et al.

Original: 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
Revised: 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

Original: IB or be contradictory to the high performance of NNs.
Revised: IB or be contradictory to the high performance of NNs.

Original: Motivated by these developments, we focus on the following questions:
Revised: Motivated by these developments, we focus on the following questions:

Original: • Does a universal two-phase training behavior of NNs exist in practice?
Revised: • Does a universal two-phase training behavior of NNs exist in practice?

Original: If this claim is invalid with the previous information measure I ( T ; X ) , can we achieve this two-phase property through another information-theoretic perspective?
Revised: If this claim is invalid with the previous information measure I ( T ; X ) , can we achieve this two-phase property through another information-theoretic perspective?

Original: • As I ( T ; X ) was unable to fully explain NN’s generalization, can we ﬁnd another measure with theoretical generalization guarantee?
Revised: • As I ( T ; X ) was unable to fully explain NN’s generalization, can we ﬁnd another measure with theoretical generalization guarantee?

Original: Also, how do we leverage it to amend the IB theory for deep neural network?
Revised: Also, how do we leverage it to amend the IB theory for deep neural network?

Original: • How do we utilize our new IB for efﬁcient training and inference of NNs in practice?
Revised: • How do we utilize our new IB for efﬁcient training and inference of NNs in practice?

Original: In this work, we propose to handle the above questions through the lens of information stored in weights (IIW), i.e., I ( w ; S ) where S =
Revised: In this work, we propose to handle the above questions through the lens of information stored in weights (IIW), i.e., I ( w ; S ) where S =

Original: { X i , Y i } ni =1 is a ﬁnite-sample dataset.
Revised: { X i , Y i } ni =1 is a ﬁnite-sample dataset.

Original: Our main contributions are four-fold: (1) we propose a new information bottleneck under the umbrella of PACBayes generalization guarantee, namely P AC-Bayes I nformation B ottleneck (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 posterior speciﬁed by PIB; and (4) we demonstrate that our new information measure covers the wide ground of NN’s behavior.
Revised: Our main contributions are four-fold: (1) we propose a new information bottleneck under the umbrella of PACBayes generalization guarantee, namely P AC-Bayes I nformation B ottleneck (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 posterior speciﬁed by PIB; and (4) we demonstrate that our new information measure covers the wide ground of NN’s behavior.

Original: Crediting to the plug-and-play modularity of SGD/SGLD, we can adapt any existing NN to a PAC-Bayes IB augmented NN seamlessly.
Revised: Crediting to the plug-and-play modularity of SGD/SGLD, we can adapt any existing NN to a PAC-Bayes IB augmented NN seamlessly.

Original: Demo code is available in supplementary materials.

Original:  A N EW B OTTLENECK WITH PAC-B AYES G UARANTEE
Revised:  A N EW B OTTLENECK WITH PAC-B AYES G UARANTEE

Original: In this section, we present the preliminaries of PAC-Bayes theory and then introduce our new information bottleneck.
Revised: In this section, we present the preliminaries of PAC-Bayes theory and then introduce our new information bottleneck.

Original: A loss function (cid:96) ( f w ( X ) , Y ) is a measure of the degree of prediction accuracy f w ( X ) compared with the ground-truth label Y .
Revised: A loss function (cid:96) ( f w ( X ) , Y ) is a measure of the degree of prediction accuracy f w ( X ) compared with the ground-truth label Y .

Original: Given the ground-truth joint distribution p ( X, Y ) , the expected true risk (out-of-sample risk) is taken on expectation as
Revised: Given the ground-truth joint distribution p ( X, Y ) , the expected true risk (out-of-sample risk) is taken on expectation as

Original: [Equation]
Revised: [Equation]

Original: Note that we take an additional expectation over p ( w | S ) because we are evaluating risk of the learned posterior instead of a speciﬁc value of parameter w .
Revised: Note that we take an additional expectation over p ( w | S ) because we are evaluating risk of the learned posterior instead of a speciﬁc value of parameter w .

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: w ) p ( S | [Equation] ) p
Revised: =

Original: [Equation]
Revised: [Equation]

Original: The PAC-Bayes bounds hold even if prior p ( w ) is incorrect and posterior p ( w | S ) is arbitrarily chosen.
Revised: The PAC-Bayes bounds hold even if prior p ( w ) is incorrect and posterior p ( w | S ) is arbitrarily chosen.

Original: In practice, we only own ﬁnite samples S .
Revised: In practice, we only own ﬁnite samples S .

Original: This gives rise to the empirical risk as
Revised: This gives rise to the empirical risk as

Original: [Equation]
Revised: [Equation]

Original: With the above Eqs.
Revised: With the above Eqs.

Original: (2) and (3) at hand, the generalization gap of the learned posterior p ( w | S ) in out-of-sample test is ∆ L ( w ) (cid:44) L ( w ) − L S ( w ) .
Revised: (2) and (3) at hand, the generalization gap of the learned posterior p ( w | S ) in out-of-sample test is ∆ L ( w ) (cid:44) L ( w ) − L S ( w ) .

Original: Xu & Raginsky (2017) proposed a PAC-Bayes bound based on the information contained in weights
Revised: Xu & Raginsky (2017) proposed a PAC-Bayes bound based on the information contained in weights

Original: I ( w ; S ) that
Revised: I ( w ; S ) that

Original: [Equation]
Revised: [Equation]

Original: when (cid:96) ( f w ( X ) , Y ) is σ -sub-Gaussian.
Revised: when (cid:96) ( f w ( X ) , Y ) is σ -sub-Gaussian.

Original: 1 A series of follow-ups tightened this bound and veriﬁed it is an effective measure of generalization capability of learning algorithms (Mou et al., 2018;
Revised: 1 A series of follow-ups tightened this bound and veriﬁed it is an effective measure of generalization capability of learning algorithms (Mou et al., 2018;

Original: Negrea et al., 2019; Pensia et al., 2018; Zhang et al., 2018).
Revised: Negrea et al., 2019; Pensia et al., 2018; Zhang et al., 2018).

Original: Therefore, it is natural to build a new information bottleneck grounded on this PAC-Bayes generalization measure, namely the PAC-Bayes information bottleneck (PIB), as
Revised: Therefore, it is natural to build a new information bottleneck grounded on this PAC-Bayes generalization measure, namely the PAC-Bayes information bottleneck (PIB), as

Original: [Equation]
Revised: [Equation]

Original: For classiﬁcation term, the loss term L S ( w ) becomes the cross-entropy between the prediction p ( Y | X, w ) and the label p ( Y | X ) , hence PIB in Eq.
Revised: For classiﬁcation term, the loss term L S ( w ) becomes the cross-entropy between the prediction p ( Y | X, w ) and the label p ( Y | X ) , hence PIB in Eq.

Original: (5) is equivalent to
Revised: (5) is equivalent to

Original: [Equation]
Revised: [Equation]

Original: which demonstrates a trade-off between maximizing the sufﬁciency (the information of label Y contained in w ) and minimizing the minimality of learned parameters w (the information of dataset
Revised: which demonstrates a trade-off between maximizing the sufﬁciency (the information of label Y contained in w ) and minimizing the minimality of learned parameters w (the information of dataset

Original: S contained in w ).
Revised: S contained in w ).

Original: Unlike previous IB based on representations , our PIB is built on weights that are not directly inﬂuenced by inputs and selected activation functions.
Revised: Unlike previous IB based on representations , our PIB is built on weights that are not directly inﬂuenced by inputs and selected activation functions.

Original: Likewise, the trade-off described by PIB objective is more reasonable since its compression term is explicitly correlated to generalization of NNs.
Revised: Likewise, the trade-off described by PIB objective is more reasonable since its compression term is explicitly correlated to generalization of NNs.

Original:  E STIMATING I NFORMATION S TORED IN W EIGHTS
Revised:  E STIMATING I NFORMATION S TORED IN W EIGHTS

Original: In this section, we present a new notion of IIW I ( w ; S ) built on the Fisher information matrix that relates to the ﬂatness of the Riemannian manifold of loss landscape.
Revised: In this section, we present a new notion of IIW I ( w ; S ) built on the Fisher information matrix that relates to the ﬂatness of the Riemannian manifold of loss landscape.

Original: Unlike Hessian eigenvalues of loss functions used for identifying ﬂat local minima and generalization but can be made arbitrarily large (Dinh et al., 2017), this notion is invariant to re-parameterization of NNs (Liang et al., 2019).
Revised: Unlike Hessian eigenvalues of loss functions used for identifying ﬂat local minima and generalization but can be made arbitrarily large (Dinh et al., 2017), this notion is invariant to re-parameterization of NNs (Liang et al., 2019).

Original: Also, our measure is invariant to the choice of activation functions because it is not directly inﬂuenced by input X like I ( T ; X ) .
Revised: Also, our measure is invariant to the choice of activation functions because it is not directly inﬂuenced by input X like I ( T ; X ) .

Original: We leverage it to monitor the information trajectory of NNs trained by SGD and cross entropy loss and verify it is capable of reproducing the two-phase transition for varying activations (e.g., ReLU , linear , tanh , and sigmoid ) in §5.1.
Revised: We leverage it to monitor the information trajectory of NNs trained by SGD and cross entropy loss and verify it is capable of reproducing the two-phase transition for varying activations (e.g., ReLU , linear , tanh , and sigmoid ) in §5.1.

Original: 3. C LOSED - FORM S OLUTION WITH G AUSSIAN A SSUMPTION
Revised: 3. C LOSED - FORM S OLUTION WITH G AUSSIAN A SSUMPTION

Original: By deriving a new information bottleneck PIB, we can look into how IIW I ( w ; S ) and L S ( w ) evolve during the learning process of NNs optimized by SGD.
Revised: By deriving a new information bottleneck PIB, we can look into how IIW I ( w ; S ) and L S ( w ) evolve during the learning process of NNs optimized by SGD.

Original: Now the key challenge ahead is how to estimate the IIW I ( w ; S ) , as
Revised: Now the key challenge ahead is how to estimate the IIW I ( w ; S ) , as

Original: [Equation]
Revised: [Equation]

Original: is the expectation of Kullback-Leibler (KL) divergence between p ( w | S ) and p ( w ) over the distribution of dataset p ( S ) .
Revised: is the expectation of Kullback-Leibler (KL) divergence between p ( w | S ) and p ( w ) over the distribution of dataset p ( S ) .

Original: And, p ( w ) is the marginal distribution of p ( w | S ) as p ( w ) (cid:44) E p ( S )
Revised: And, p ( w ) is the marginal distribution of p ( w | S ) as p ( w ) (cid:44) E p ( S )

Original: [ p ( w | S )] .
Revised: [ p ( w | S )] .

Original: When we assume both p ( w ) =
Revised: When we assume both p ( w ) =

Original: N ( w | θ 0 , Σ 0 ) and p ( w | S ) =
Revised: N ( w | θ 0 , Σ 0 ) and p ( w | S ) =

Original: N ( w | θ S , Σ S ) are Gaussian distribu- tions, the KL divergence term in Eq.
Revised: N ( w | θ S , Σ S ) are Gaussian distribu- tions, the KL divergence term in Eq.

Original: (7) has closed-form solution as
Revised: (7) has closed-form solution as

Original: [Equation]
Revised: [Equation]

Original: Here det A and tr ( A ) are the determinant and trace of matrix A , respectively; D is the dimension of parameter w and is a constant for a speciﬁc NN architecture; θ S are the yielded weights after
Revised: Here det A and tr ( A ) are the determinant and trace of matrix A , respectively; D is the dimension of parameter w and is a constant for a speciﬁc NN architecture; θ S are the yielded weights after

Original: SGD converges on the dataset S .
Revised: SGD converges on the dataset S .

Original: If the covariances of prior and posterior are proportional, 2 the logarithmic and trace terms in Eq.
Revised: If the covariances of prior and posterior are proportional, 2 the logarithmic and trace terms in Eq.

Original: (8) both become constant.
Revised: (8) both become constant.

Original: Therefore, the mutual information term is proportional to the quadratic term as
Revised: Therefore, the mutual information term is proportional to the quadratic term as

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: In the next section, we will see how to set prior covariance Σ 0 .
Revised: In the next section, we will see how to set prior covariance Σ 0 .

Original: 3. B OOTSTRAP C OVARIANCE OF O RACLE P RIOR
Revised: 3. B OOTSTRAP C OVARIANCE OF O RACLE P RIOR

Original: Since the computation of exact oracle prior Σ 0 needs the knowledge of p ( S ) 3 , we propose to approximate it by bootstrapping from S as
Revised: Since the computation of exact oracle prior Σ 0 needs the knowledge of p ( S ) 3 , we propose to approximate it by bootstrapping from S as

Original: [Equation]
Revised: [Equation]

Original: Σ 0 is called oracle prior because it minimizes the term I ( w ; S ) = E p ( S ) [ KL ( p (w | S ) (cid:107) p ( W )] when p ( w ) =
Revised: where S k is a bootstrap sample obtained by re-sampling from the ﬁnite data S , and S k ∼ p ( S ) is still a valid sample following p ( S ) .

Original: E p ( S )

Original: [ p ( w | S )] (Dziugaite et al., 2021).

Original: Data: Total number of samples n , batch size B , number of mini-batches in one epoch T 0 , number of information estimation iterations T 1 , learning rate η , moving average hyperparameters ρ and K , a sequence of gradients set ∇L = ∅

Original: Result:

Original: Calculated approximate information (cid:101) I ( w ; S )

Original: [-1 Pretrain the model by vanilla SGD to obtain the prior mean θ 0 ; 2 for t=1: T 0 do ∇ L t ← ∇ θ 1 B (cid:80) b (cid:96) b (ˆ θ t − 1 ) , ˆ θ t ← ˆ θ t − 1 − η ∇ L t ; /* Vanilla SGD */ ∇L ← ∇L (cid:83) {∇ L t } ; /* Store gradients */ ¯ θ t ← (cid:113) ρ ¯ θ 2 t − 1 + 1 − ρK (cid:80) K − 1 k =
-]
Revised: [-1 Pretrain the model by vanilla SGD to obtain the prior mean θ 0 ; 2 for t=1: T 0 do ∇ L t ← ∇ θ 1 B (cid:80) b (cid:96) b (ˆ θ t − 1 ) , ˆ θ t ← ˆ θ t − 1 − η ∇ L t ; /* Vanilla SGD */ ∇L ← ∇L (cid:83) {∇ L t } ; /* Store gradients */ ¯ θ t ← (cid:113) ρ ¯ θ 2 t − 1 + 1 − ρK (cid:80) K − 1 k =
-]

Original: ˆ θ 2 t − k ; /

Original: * Moving average */ 6 end 7 ∆ θ ← ¯ θ T 0 − θ 0 , ∆ F 0 ← 0 ; 8 for t=1: T 1 do ∆ F t ← ∆ F t − 1 + (∆ θ (cid:62) ∇ L t ) 2 ; /

Original: * Storage-friendly computation */ 10 end 11 (cid:101) I ( w ; S ) ← nT 1 ∆ F T 1 ; where S k is a bootstrap sample obtained by re-sampling from the ﬁnite data S , and S k ∼ p ( S ) is still a valid sample following p ( S ) .

Original: Now we are closer to the solution but the above term is still troublesome to calculate.
Revised: Now we are closer to the solution but the above term is still troublesome to calculate.

Original: For getting { θ S k } Kk =1 , we need to optimize on a series of ( K times) bootstrapping datasets { S k } Kk =1 via SGD until it converges, which is prohibitive in deep learning practices.
Revised: For getting { θ S k } Kk =1 , we need to optimize on a series of ( K times) bootstrapping datasets { S k } Kk =1 via SGD until it converges, which is prohibitive in deep learning practices.

Original: Therefore, we propose to approximate the difference θ S − θ 0 by inﬂuence functions drawn from robust statistics literature (Cook & Weisberg, 1982; Koh & Liang, 2017; Wang et al., 2020b).
Revised: Therefore, we propose to approximate the difference θ S − θ 0 by inﬂuence functions drawn from robust statistics literature (Cook & Weisberg, 1982; Koh & Liang, 2017; Wang et al., 2020c;b).

Original: Lemma 1 (Inﬂuence Function (Cook & Weisberg, 1982)) .
Revised: Lemma 1 (Inﬂuence Function (Cook & Weisberg, 1982)) .

Original: Given a dataset S
Revised: Given a dataset S

Original: = {
Revised: = {

Original: ( X i , Y i ) } ni =1 and the parameter ˆ θ S (cid:44) argmin θ L S ( θ ) = argmin θ 1 n (cid:80) ni
Revised: ( X i , Y i ) } ni =1 and the parameter ˆ θ S (cid:44) argmin θ L S ( θ ) = argmin θ 1 n (cid:80) ni

Original: =1 (cid:96) i ( θ ) 4 that optimizes the empirical loss function, if we drop sample ( X j , Y j ) in S to get a jackknife sample S \ j and retrain our model, the new parameters are ˆ θ S \ j = argmin θ L S \ j ( θ ) = argmin
Revised: =1 (cid:96) i ( θ ) 4 that optimizes the empirical loss function, if we drop sample ( X j , Y j ) in S to get a jackknife sample S \ j and retrain our model, the new parameters are ˆ θ S \ j = argmin θ L S \ j ( θ ) = argmin

Original: θ 1 n [Equation] ni =1 (cid:96) i ( θ ) − 1 n (cid:96) j ( θ ) .
Revised: θ 1 n (cid:80) n i =1 (cid:96) i ( θ ) − 1 n (cid:96) j ( θ ) .

Original: The approximation of parameter difference ˆ θ S \ j − ˆ θ S is deﬁned by inﬂuence function ψ , as [Equation] , where ψ j [Equation]
Revised: The approximation of parameter difference ˆ θ S \ j − ˆ θ S is deﬁned by inﬂuence function ψ , as ˆ θ S \ j − ˆ θ S (cid:39) − 1 n ψ j , where ψ j = − H − 1ˆ θ S ∇ θ (cid:96) j

Original: − [Equation] [Equation] and H ˆ θ S (cid:44) 1 n (cid:80) ni
Revised: (ˆ θ S ) ∈ R D , (11) and H ˆ θ S (cid:44) 1 n (cid:80) ni

Original: =1 ∇ 2 θ (cid:96) i
Revised: =1 ∇ 2 θ (cid:96) i

Original: (ˆ θ S ) ∈
Revised: (ˆ θ S ) ∈

Original: R D
Revised: R D

Original: × D is Hessian matrix.
Revised: × D is Hessian matrix.

Original: The application of inﬂuence functions can be further extended to the case when the loss function is perturbed by a vector ξ = ( ξ 1 , ξ 2 , . . . , ξ n ) (
Revised: The application of inﬂuence functions can be further extended to the case when the loss function is perturbed by a vector ξ = ( ξ 1 , ξ 2 , . . . , ξ n ) (

Original: cid:62) ∈ R n as ˆ θ S, ξ = argmin θ 1 n (cid:80) ni
Revised: cid:62) ∈ R n as ˆ θ S, ξ = argmin θ 1 n (cid:80) ni

Original: =1 ξ i (cid:96) i ( θ ) .
Revised: =1 ξ i (cid:96) i ( θ ) .

Original: In this scenario, the parameter difference can be approximated by
Revised: In this scenario, the parameter difference can be approximated by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where Ψ =
Revised: where Ψ =

Original: ( ψ 1 , ψ 2 , . . .
Revised: ( ψ 1 , ψ 2 , . . .

Original: , ψ n ) (cid:62) ∈ R n × D is a combination of all inﬂuence functions ψ ; 1 = (1 , 1 , . . .
Revised: , ψ n ) (cid:62) ∈ R n × D is a combination of all inﬂuence functions ψ ; 1 = (1 , 1 , . . .

Original: , 1) (cid:62) is an n -dimensional all-one vector.
Revised: , 1) (cid:62) is an n -dimensional all-one vector.

Original: Lemma 1 gives rise to the following lemma on approximation of oracle prior covariance in Eq.
Revised: Lemma 1 gives rise to the following lemma on approximation of oracle prior covariance in Eq.

Original: Lemma 2 (Approximation of Oracle Prior Covariance) .
Revised: Lemma 2 (Approximation of Oracle Prior Covariance) .

Original: Given the deﬁnition of inﬂuence functions (Lemma 1) and Poisson bootstrapping (Lemma A.2), the covariance matrix of the oracle prior can be approximated by[Equation] [Equation] [Equation]
Revised: Given the deﬁnition of inﬂuence functions (Lemma 1) and Poisson bootstrapping (Lemma A.2), the covariance matrix of the oracle prior can be approximated by

Original: sequence of weights { w t } t

Original: [Equation]

Original: ˆ k

Original: [Equation]

Original: p

Original: [Equation]
Revised: [Equation]

Original: w

Original: [Equation]

Original: [Equation] Gaussian noise [Equation] where F ˆ θ is Fisher information matrix (FIM); we omit the subscript S of ˆ θ S and ˆ θ S, ξ for notation conciseness, and ξ k is the bootstrap resampling weight in the k -th experiment.
Revised: where F ˆ θ is Fisher information matrix (FIM); we omit the subscript S of ˆ θ S and ˆ θ S, ξ for notation conciseness, and ξ k is the bootstrap resampling weight in the k -th experiment.

Original: Please refer to Appendix A for the proof of this lemma.
Revised: Please refer to Appendix A for the proof of this lemma.

Original: 3. E FFICIENT I NFORMATION E STIMATION A LGORITHM
Revised: 3. E FFICIENT I NFORMATION E STIMATION A LGORITHM

Original: After the approximation of oracle prior covariance, we are now able to rewrite the IIW term I ( w ; S ) in Eq.
Revised: After the approximation of oracle prior covariance, we are now able to rewrite the IIW term I ( w ; S ) in Eq.

Original: (9) to
Revised: (9) to

Original: [Equation]
Revised: [Equation]

Revised: Data: Total number of samples n , batch size B , number of mini-batches in one epoch T 0 , number of information estimation iterations T 1 , learning rate η , moving average hyperparameters ρ and K , a sequence of gradients set ∇L = ∅

Revised: Result:

Revised: Calculated approximate information (cid:101) I ( w ; S )

Revised: 1 Pretrain the model by vanilla SGD to obtain the prior mean θ 0 ; 2 for t=1: T 0 do

Revised: [Table]

Revised: ∇

Revised: [Table]

Revised: ← ∇

Revised: [Table]

Revised: (cid:96)

Revised: [Table]

Revised: [Table]

Revised: t

Revised: [Table]

Revised: [Table]

Revised: L t ; /* Vanilla SGD */

Revised: [Table]

Revised: ∇L

Revised: [Table]

Revised: ∇L (cid:83) {

Revised: [Table]

Revised: ; /* Store gradients */

Revised: [Table]

Revised: ; /* Moving average */

Revised: [Table]

Revised: end

Revised: [Table]

Revised: [Table]

Revised: do

Revised: [Table]

Revised: ; /* Storage-friendly computation */

Revised: [Table]

Revised: end

Revised: [Table]

Revised: Data: Total number of samples n ,

Revised: [Table]

Revised: temperature

Revised: [Table]

Revised: Result

Revised: [Table]

Revised: of

Revised: [Table]

Revised: following

Revised: [Table]

Revised: w

Revised: [Table]

Revised: [Table]

Revised: /* Mini-batch gradient of energy function */

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: ; / * SGLD by gradient plus isotropic Gaussian noise *

Revised: [Equation]

Revised: * Learning rate & temperature decay

Revised: [Equation]

Revised: [Equation]

Revised: [Table]

Revised: β

Revised: [Table]

Revised: [Equation]

Revised: t

Revised: [Equation]

Revised: t

Revised: [Equation]

Revised: [Equation]

Revised: 5 until the weight sequence { w t } t ≥ ˆ k becomes stable

Original: We deﬁne the approximate information by (cid:101) I ( w ; S ) where we approximate the expectation
Revised: We deﬁne the approximate information by (cid:101) I ( w ; S ) where we approximate the expectation

Original: E p ( S )
Revised: E p ( S )

Original: [ θ (cid:62) S F ˆ θ θ S ] by ¯ θ S = (cid:113) 1 K (cid:80)
Revised: [ θ (cid:62) S F ˆ θ θ S ] by ¯ θ S = (cid:113) 1 K (cid:80)

Original: Kk
Revised: Kk

Original: =1 ˆ θ 2 k = (cid:18)(cid:113) 1 K (cid:80)
Revised: =1 ˆ θ 2 k = (cid:18)(cid:113) 1 K (cid:80)

Original: Kk
Revised: Kk

Original: =1 ˆ θ 21 ,k , . . .
Revised: =1 ˆ θ 21 ,k , . . .

Original: , (cid:113) 1 K (cid:80)
Revised: , (cid:113) 1 K (cid:80)

Original: Kk
Revised: Kk

Original: =

Revised: =

Original: ˆ θ 2 D,k
Revised: ˆ θ 2 D,k

Original: (cid:19) (cid:62) .
Revised: (cid:19) (cid:62) .

Original: In Eq.
Revised: In Eq.

Original: (14), the information consists of two major components: ∆ θ = ¯ θ S − θ 0 ∈ R D and F ˆ θ ∈ R D × D , which can easily cause out-of-memory error due to the high-dimensional matrix product operations.
Revised: (14), the information consists of two major components: ∆ θ = ¯ θ S − θ 0 ∈ R D and F ˆ θ ∈ R D × D , which can easily cause out-of-memory error due to the high-dimensional matrix product operations.

Original: We therefore hack into FIM to get
Revised: We therefore hack into FIM to get

Original: [Equation]
Revised: [Equation]

Original:  such that the high dimensional matrix vector product reduces to vector inner product.
Revised: (15) such that the high dimensional matrix vector product reduces to vector inner product.

Original: We encapsulate the algorithm for estimating IIW during vanilla SGD by Algorithm 1.
Revised: We encapsulate the algorithm for estimating IIW during vanilla SGD by Algorithm 1.

Original:  B AYESIAN I NFERENCE FOR THE O PTIMAL P OSTERIOR
Revised:  B AYESIAN I NFERENCE FOR THE O PTIMAL P OSTERIOR

Original: Recall that we designed a new bottleneck on the expected generalization gap drawn from PAC-Bayes theory in §2, and then derived an approximation of the IIW in §3.
Revised: Recall that we designed a new bottleneck on the expected generalization gap drawn from PAC-Bayes theory in §2, and then derived an approximation of the IIW in §3.

Original: The two components of PACBayes IB in Eq.
Revised: The two components of PACBayes IB in Eq.

Original: (6) are hence tractable as a learning objective.
Revised: (6) are hence tractable as a learning objective.

Original: We give the following lemma on utilizing it for inference.
Revised: We give the following lemma on utilizing it for inference.

Revised: [Figure]

Original: Lemma 3 (Optimal Posterior for PAC-Bayes Information Bottleneck) .
Revised: Lemma 3 (Optimal Posterior for PAC-Bayes Information Bottleneck) .

Original: Given an observed dataset
Revised: Given an observed dataset

Original: S ∗ , the optimal posterior p ( w | S ∗ ) of PAC-Bayes IB in Eq.
Revised: S ∗ , the optimal posterior p ( w | S ∗ ) of PAC-Bayes IB in Eq.

Original: (5) should satisfy the following form that
Revised: (5) should satisfy the following form that

Original: [Equation]
Revised: [Equation]

Original: where U S ∗ ( w ) is the energy function deﬁned as U S ∗ ( w )
Revised: where U S ∗ ( w ) is the energy function deﬁned as U S ∗ ( w )

Original: =
Revised: =

Original: ˆ L S ∗ ( w ) − β log p ( w ) , and Z ( S ) is the normalizing constant.
Revised: ˆ L S ∗ ( w ) − β log p ( w ) , and Z ( S ) is the normalizing constant.

Original: Please refer to Appendix B for the proof.
Revised: Please refer to Appendix B for the proof.

Original: The reason why we write the posterior in terms of an exponential form is that it is a typical Gibbs distribution (Kittel, 2004)
Revised: The reason why we write the posterior in terms of an exponential form is that it is a typical Gibbs distribution (Kittel, 2004)

Original: (also called Boltzmann distribution) with energy function U S ∗ ( w ) and temperature β (the same β of PIB appears in Eq.
Revised: (also called Boltzmann distribution) with energy function U S ∗ ( w ) and temperature β (the same β of PIB appears in Eq.

Original: Crediting to this formula, we can adopt Markov chain Monte Carlo (MCMC) for rather efﬁcient
Revised: Crediting to this formula, we can adopt Markov chain Monte Carlo (MCMC) for rather efﬁcient

Original: Bayesian inference.
Revised: Bayesian inference.

Original: Speciﬁcally, we propose to use stochastic gradient Langevin dynamics (SGLD)
Revised: Speciﬁcally, we propose to use stochastic gradient Langevin dynamics (SGLD)

Original: (Welling & Teh, 2011) that has been proved efﬁcient and effective in large-scale posterior inference.
Revised: (Welling & Teh, 2011) that has been proved efﬁcient and effective in large-scale posterior inference.

Original: SGLD can be realized by a simple adaption of SGD as
Revised: SGLD can be realized by a simple adaption of SGD as

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where η k is step size, ε k ∼ N ( ε | 0 , I D ) is a standard Gaussian noise vector, and g k is an unbiased estimate of energy function gradient ∇ U ( w k ) .
Revised: where η k is step size, ε k ∼ N ( ε | 0 , I D ) is a standard Gaussian noise vector, and g k is an unbiased estimate of energy function gradient ∇ U ( w k ) .

Original: SGLD can be viewed as a discrete Langevin diffusion described by stochastic differential equation (Raginsky et al., 2017; Borkar & Mitter, 1999): d w ( t )
Revised: SGLD can be viewed as a discrete Langevin diffusion described by stochastic differential equation (Raginsky et al., 2017; Borkar & Mitter, 1999): d w ( t )

Original: = −∇ U ( w ( t ))
Revised: = −∇ U ( w ( t ))

Original: dt + √ 2 βdB ( t ) , where { B ( t ) } t ≥ 0 is the standard Brownian motion in R D .
Revised: dt + √ 2 βdB ( t ) , where { B ( t ) } t ≥ 0 is the standard Brownian motion in R D .

Original: The Gibbs distribution π ( w ) ∝ exp( − 1 β U ( w )) is the unique invariant distribution of the Langevin diffusion.
Revised: The Gibbs distribution π ( w ) ∝ exp( − 1 β U ( w )) is the unique invariant distribution of the Langevin diffusion.

Original: And, distribution of w t converges rapidly to π ( w ) when t → ∞ with sufﬁciently small β (Chiang et al., 1987).
Revised: And, distribution of w t converges rapidly to π ( w ) when t → ∞ with sufﬁciently small β (Chiang et al., 1987).

Original: Similarly for SGLD in Eq.
Revised: Similarly for SGLD in Eq.

Original: (17), under the conditions that (cid:80) ∞
Revised: (17), under the conditions that (cid:80) ∞

Original: t η t → ∞ and (cid:80) ∞ t η 2 t → 0 , and an annealing temperature β , the sequence of { w k } k ≥ ˆ k converges to Gibbs distribution with sufﬁciently large ˆ k .
Revised: t η t → ∞ and (cid:80) ∞ t η 2 t → 0 , and an annealing temperature β , the sequence of { w k } k ≥ ˆ k converges to Gibbs distribution with sufﬁciently large ˆ k .

Original: As we assume the oracle prior p ( w ) 
Revised: As we assume the oracle prior p ( w ) =

Revised: N ( w | θ 0 , Σ 0 ) , log p ( w ) satisﬁes

Original: [Equation]
Revised: [Equation]

Original: The inference of the optimal posterior is then summarized by Algorithm 2.
Revised: The inference of the optimal posterior is then summarized by Algorithm 2.

Original: φ η ( · ) and φ β ( · ) are learning rate decay and temperature annealing functions, e.g., cosine decay, respectively.
Revised: φ η ( · ) and φ β ( · ) are learning rate decay and temperature annealing functions, e.g., cosine decay, respectively.

Original: Our SGLD based algorithm leverages the advantage of MCMC such that it is capable of sampling from the optimal posterior even for very complex NNs.
Revised: Our SGLD based algorithm leverages the advantage of MCMC such that it is capable of sampling from the optimal posterior even for very complex NNs.

Original: Also, it does need to know the groundtruth distribution of S while still theoretically allows global convergence avoiding local minima.
Revised: Also, it does need to know the groundtruth distribution of S while still theoretically allows global convergence avoiding local minima.

Original: It can be realized with a minimal adaptation of common auto-differentiation packages, e.g., PyTorch (Paszke et al., 2019), by injecting isotropic noise in the SGD updates.
Revised: It can be realized with a minimal adaptation of common auto-differentiation packages, e.g., PyTorch (Paszke et al., 2019), by injecting isotropic noise in the SGD updates.

Original: Please refer to Appendix C for the details of computation for PIB object.
Revised: Please refer to Appendix C for the details of computation for PIB object.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we aim to verify the intepretability of the proposed notion of IIW by Eq.
Revised: In this section, we aim to verify the intepretability of the proposed notion of IIW by Eq.

Original: We monitor the information trajectory when training NNs with plain cross entropy loss and SGD for the sake of activation functions (§5.1), architecture (§5.2), noise ratio (§5.3), and batch size (§5.4).
Revised: We monitor the information trajectory when training NNs with plain cross entropy loss and SGD for the sake of activation functions (§5.1), architecture (§5.2), noise ratio (§5.3), and batch size (§5.4).

Original: We also substantiate the superiority of optimal Gibbs posterior inference based on the proposed
Revised: We also substantiate the superiority of optimal Gibbs posterior inference based on the proposed

Original: Algorithm 2, where PIB instead of plain cross entropy is used as the objective function (§5.5).
Revised: Algorithm 2, where PIB instead of plain cross entropy is used as the objective function (§5.5).

Original: We conclude the empirical observations in §5.6 at last.
Revised: We conclude the empirical observations in §5.6 at last.

Original: Please refer to Appendix D for general experimental setups about the used datasets and NNs.
Revised: Please refer to Appendix D for general experimental setups about the used datasets and NNs.

Original: 5. I NFORMATION WITH D IFFERENT A CTIVATION F UNCTIONS
Revised: 5. I NFORMATION WITH D IFFERENT A CTIVATION F UNCTIONS

Revised: [Figure]

Revised: All follow the general trend of ﬁtting to compression phase transition.

Revised: And, deeper layers can accelerate both the ﬁtting and compressing phases.

Original: We train a 2-layer MLP (784-512-10) with plain cross-entropy loss by Adam on the MNIST dataset, meanwhile monitor the trajectory of the IIW I ( w ; S ) .
Revised: We train a 2-layer MLP (784-512-10) with plain cross-entropy loss by Adam on the MNIST dataset, meanwhile monitor the trajectory of the IIW I ( w ; S ) .

Original: Results are illustrated in Fig.
Revised: Results are illustrated in Fig.

Original: 1 where different activation functions ( linear , tanh , ReLU , and sigmoid ) are testiﬁed.
Revised: 1 where different activation functions ( linear , tanh , ReLU , and sigmoid ) are testiﬁed.

Original: We identify that there is a clear boundary between ﬁtting and compression phases for all of them.
Revised: We identify that there is a clear boundary between ﬁtting and compression phases for all of them.

Original: For example, for the linear activation function on the ﬁrst column, the IIW I ( w ; S ) surges within the ﬁrst several iterations then drops slowly during the next iterations.
Revised: For example, for the linear activation function on the ﬁrst column, the IIW I ( w ; S ) surges within the ﬁrst several iterations then drops slowly during the next iterations.

Original: Simultaneously, we could see that the training loss reduces sharply at the initial stage, then keeps decreasing during the information compression.
Revised: Simultaneously, we could see that the training loss reduces sharply at the initial stage, then keeps decreasing during the information compression.

Original: At the last period of compression, we witness the information ﬂuctuates near zero with the recurrence of memorization phenomenon (IIW starts to increase).
Revised: At the last period of compression, we witness the information ﬂuctuates near zero with the recurrence of memorization phenomenon (IIW starts to increase).

Original: This implies that further training is causing over-ﬁtting.
Revised: This implies that further training is causing over-ﬁtting.

Original: IIW shows great universality on representing the information compression of NNs.
Revised: IIW shows great universality on representing the information compression of NNs.

Original: 5. I NFORMATION WITH D EEPER AND W IDER A RCHITECTURE
Revised: 5. I NFORMATION WITH D EEPER AND W IDER A RCHITECTURE

Original: [Figure]

Original: [Figure]

Original: All follow the general trend of ﬁtting to compression phase transition.

Original: And, deeper layers can accelerate both the ﬁtting and compressing phases.

Original: Having identiﬁed the phase transition of the 2layer MLP corresponding to IIW I ( w ; S ) , we test it under more settings: different architectures and different batch sizes.
Revised: Having identiﬁed the phase transition of the 2-layer MLP corresponding to IIW I ( w ; S ) , we test it under more settings: different architectures and different batch sizes.

Original: For the architecture setting, we design MLPs from 1 to 4 layers.
Revised: For the architecture setting, we design MLPs from 1 to 4 layers.

Original: Results are shown in Fig.
Revised: Results are shown in Fig.

Original: The ﬁrst and the last ﬁgures show the information trajectory of the 1-layer/4-layer version of MLP-Large (78410/784-512-100-80-10) where clear two-phase transitions happen in all these MLPs.
Revised: The ﬁrst and the last ﬁgures show the information trajectory of the 1-layer/4-layer version of MLP-Large (784-10/784-512-100-80-10) where clear two-phase transitions happen in all these MLPs.

Original: The 1-layer MLP is actually a softmax regression model.
Revised: The 1-layer MLP is actually a softmax regression model.

Original: It can be identiﬁed that this model ﬁts and compresses very slowly w.r.t.
Revised: It can be identiﬁed that this model ﬁts and compresses very slowly w.r.t.

Original: IIW compared with deeper NNs.
Revised: IIW compared with deeper NNs.

Original: 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.
Revised: 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.

Original: Furthermore, when we add more layers, the ﬁtting phase becomes shorter.
Revised: Furthermore, when we add more layers, the ﬁtting phase becomes shorter.

Original: Speciﬁcally, we observe the incidence of overﬁtting at the end of the 4-layer MLP training as IIW starts to increase.
Revised: Speciﬁcally, we observe the incidence of overﬁtting at the end of the 4-layer MLP training as IIW starts to increase.

Original: We also examine how IIW explains the generalization w.r.t.
Revised: We also examine how IIW explains the generalization w.r.t.

Original: the number of hidden units, a.k.a.
Revised: the number of hidden units, a.k.a.

Original: the width of NNs, by Fig.
Revised: the width of NNs, by Fig.

Original: We train a 2-layer MLP without any regularization on MNIST.
Revised: We train a 2-layer MLP without any regularization on MNIST.

Original: The left panel shows the training and test errors for this experiment.
Revised: The left panel shows the training and test errors for this experiment.

Original: Notably, the difference of test and train acc can be seen an indicator of the generalization gap in Eq.
Revised: Notably, the difference of test and train acc can be seen an indicator of the generalization gap in Eq.

Original: IIW should be aligned to this gap by deﬁnition.
Revised: IIW should be aligned to this gap by deﬁnition.

Original: While 32 units are (nearly) enough to interpolate the training set, more hidden units still achieve better generalization performance, which illustrates the effect of overparameterization.
Revised: While 32 units are (nearly) enough to interpolate the training set, more hidden units still achieve better generalization performance, which illustrates the effect of overparameterization.

Original: In this scenario, weights (cid:96) 2 -norm keeps increasing with more units while IIW decays, similar to the test error.
Revised: In this scenario, weights (cid:96) 2 -norm keeps increasing with more units while IIW decays, similar to the test error.

Original: We identify that more hidden units do not render much increase of IIW, which is contrast to the intuition that wider NNs always have larger information complexity.
Revised: We identify that more hidden units do not render much increase of IIW, which is contrast [Figure] [Figure] [Figure]

Revised: [Figure]

Revised: to the intuition that wider NNs always have larger information complexity.

Original: More importantly, we ﬁnd IIW is consistent to the generalization gap on each width.
Revised: More importantly we ﬁnd IIW is consistent to the generalization gap on each width

Original: 5. R ANDOM L ABELS VS . T RUE L ABELS
Revised: 5. R ANDOM L ABELS VS . T RUE L ABELS

Original: According to the PAC-Bayes theorem, the IIW is a promising measure to explain/predict the generalization capability of NNs.
Revised: According to the PAC-Bayes theorem the IIW is a promising measure to explain/predict the generalization capability of NNs.

Original: NNs are often over-parameterized thus can perfectly ﬁt even the random labels, obviously without any generalization capability.
Revised: NNs are often over-parameterized thus can perfectly ﬁt even the random labels, obviously without any generalization capability.

Original: For example, 2-layer MLP has 15 , 728 , 640 parameters that are much larger than the sample number of CIFAR-10 (50,000).
Revised: For example, 2-layer MLP has 15 , 728 , 640 parameters that are much larger than the sample number of CIFAR-10 (50,000).

Original: That causes the number of parameters an unreliable measure of NN complexity in overparameterization[Figure] [Figure] settings (Neyshabur et al., 2015).
Revised: That causes the number of parameters an unreliable measure of NN complexity in overparameterization settings (Neyshabur et al., 2015).

Original: Alternatively, (cid:96) 2 -norm is often used as a complexity measure to be imposed on regularizing model training in practices.[Figure]
Revised: Alternatively, (cid:96) 2 -norm is often used as a complexity measure to be imposed on regularizing model training in practices.

Original: We investigate the model trained with different levels of label corruption, as shown by the left panel of Fig. 4.
Revised: We investigate the model trained with different levels of label corruption, as shown by the left panel of Fig. 4.

Original: We train a 2-layer MLP on CIFAR-10 and ﬁnd that the increasing noise causes sharp test acc decay while train acc does not change much.
Revised: We train a 2-layer MLP on CIFAR-10 and ﬁnd that the increasing noise causes sharp test acc decay while train acc does not change much.

Original: Meanwhile, IIW keeps growing with the fall of test acc and expansion of generalization gap.
Revised: Meanwhile, IIW keeps growing with the fall of test acc and expansion of generalization gap.

Original: This demonstrates IIW’s potential in identifying the noise degree in datasets or the mismatch between the test and train data distributions.
Revised: This demonstrates IIW’s potential in identifying the noise degree in datasets or the mismatch between the test and train data distributions.

Original: We further build a random-label CIFAR-10, results are on the right of Fig.
Revised: We further build a random-label CIFAR-10, results are on the right of Fig.

Original: 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.
Revised: 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.

Original: This phenomenon is also captured by IIW.
Revised: This phenomenon is also captured by IIW.

Original: 5. I NFORMATION IN WEIGHTS W . R . T . B ATCH S IZE
Revised: 5. I NFORMATION IN WEIGHTS W . R . T . B ATCH S IZE

Original: [Figure]

Original: PIB.

Original: It is seen that PIB reaches the least IIW thus obtaining the best generalization ability.

Original: We also consider how batch size inﬂuences IIW and generalization gap.
Revised: We also consider how batch size inﬂuences IIW and generalization gap.

Original: Recent efforts on bounding I ( w ; S ) of iterative algorithms (e.g., SGD and SGLD)
Revised: Recent efforts on bounding I ( w ; S ) of iterative algorithms (e.g., SGD and SGLD)

Original: (Mou et al., 2018; Pensia et al., 2018) imply that the variance of gradients is a crucial factor.
Revised: (Mou et al., 2018; Pensia et al., 2018) imply that the variance of gradients is a crucial factor.

Original: For the ultimate case where batch size equals  full sample size, the gradient variance is zero, and the model is prone to over-ﬁtting grounded on empirical observations.
Revised: For the ultimate case where batch size equals [Table] full sample size , the gradient variance is zero, and the model is prone to over-ﬁtting grounded on empirical observations.

Original: On the other hand, when batch size equals one, the variance becomes tremendously large.
Revised: When batch size equals one, the variance becomes tremendously large.

Original: We conjecture there is an optimal batch size that reaches the minimum generalization gap, in other word, the minimum IIW.
Revised: We conjecture there is an optimal batch size that reaches the minimum generalization gap, in other word, the minimum IIW.

Original: This conjecture is raised on our empirical ﬁndings, displayed in Fig.where we test IIW on model with varying batch size.
Revised: This conjecture is raised on our empirical ﬁndings, displayed in Fig.where we test IIW on model with varying batch size.

Original: Each model is updated with the same total number of iterations and the same learning rate.
Revised: Each model is updated with the same total number of iterations and the same learning rate.

Original: 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).
Revised: 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).

Original: On the left panel, the modelalso reaches the minimum IIW when batch size is 16.
Revised: On the left, the model reaches the minimum IIW when batch size is 16.

Original: [Table]

Original: 5. B AYESIAN I NFERENCE WITH V ARYING E NERGY F UNCTIONS
Revised: 5. B AYESIAN I NFERENCE WITH V ARYING E NERGY F UNCTIONS

Original: To conﬁrm the superiority of the proposed algorithm for PAC-Bayes IB in § [Table] we compare it with vanilla SGD and two widely used regularizations: (cid:96) 2 -norm and dropout.
Revised: To conﬁrm the superiority of the proposed PIB in §4, we compare it with vanilla SGD and two widely used regularizations: (cid:96) 2 -norm and dropout.

Original: We train a large VGG network (Simonyan & Zisserman, 2014) on four open datasets: CIFAR10/100 (Krizhevsky et al., 2009), STL10 (Coates et al., 2011) [Table] and SVHN (Netzer et al., 2011), as shown in Table 
Revised: We train a large VGG network (Simonyan & Zisserman, 2014) on four open datasets: CIFAR10/100 (Krizhevsky et al., 2009), STL10 (Coates et al., 2011), and SVHN (Netzer et al., 2011), as shown in Table 1, where we ﬁnd PIB consistantly outperforms the baselines.

Original: [Table]

Original: where we ﬁnd PIB consistantly outperforms the baselines.

Original: 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.
Revised: 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.

Original: This is veriﬁed by Fig.
Revised: This is veriﬁed by Fig.

Original: 6 where PIB helps restrict IIW in the lowest level.
Revised: 6 where PIB helps restrict IIW in the lowest level.

Original: Please refer to Appendix D for experimental setups.
Revised: Please refer to Appendix D for experimental setups.

Original: 5. S UMMARY OF E XPERIMENTS
Revised: 5. S UMMARY OF E XPERIMENTS

Original: We made the following observations from our experiments:
Revised: We made the following observations from our experiments:

Original: We can clearly identify the ﬁtting-compression phase transition during training through our new information measure, i.e., information stored in weights (IIW).
Revised: We can clearly identify the ﬁtting-compression phase transition during training through our new information measure, i e., information stored in weights (IIW).

Original: Unlike the representationbased information measure I ( X ; Z ) , IIW applies to various activation functions including ReLU , sigmoid , tanh , and linear .
Revised: Unlike the representationbased information measure I ( X ; Z ) , IIW applies to various activation functions including ReLU , sigmoid , tanh , and linear .

Original: We further identify that the phase transition applies to deeper and wider architecture.
Revised: We further identify that the phase transition applies to deeper and wider architecture.

Original: More importantly, deeper model is proved to reach faster ﬁtting and compression than shallow models.
Revised: More importantly, deeper model is proved to reach faster ﬁtting and compression than shallow models.

Original: Unlike (cid:96) 2 -norm of weights that rise together with wider models, IIW better illustrates the true model complexity and its generalization gap.
Revised: Unlike (cid:96) 2 -norm of weights that rise together with wider models, IIW better illustrates the true model complexity and its generalization gap.

Original: IIW can explain the performance drop w.r.t.
Revised: IIW can explain the performance drop w.r.t.

Original: the degree of label noise.
Revised: the degree of label noise.

Original: Also, IIW can even identify the generalization gap for models learned from random labels.
Revised: Also, IIW can even identify the generalization gap for models learned from random labels.

Original: There might exist an optimal batch size for the minimum generalization gap, which is empirically demonstrated by our experiments.
Revised: There might exist an optimal batch size for the minimum generalization gap, which is empirically demonstrated by our experiments.

Original: Adopting SGD based on the energy function derived from PAC-Bayes IB enables good inference to the optimal posterior of NNs.
Revised: Adopting SGD based on the energy function derived from PAC-Bayes IB enables good inference to the optimal posterior of NNs.

Original: This works for practical large networks in the literature.
Revised: This works for practical large networks in the literature.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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 regularization.
Revised: 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 regularization.

Original: Empirical results show the universality of our information measure on explaining NNs, which sheds light on understanding NNs through information bottleneck.
Revised: Empirical results show the universality of our information measure on explaining NNs, which sheds light on understanding NNs through information bottleneck.

Original: We aim to further investigate its performance and develop this into practical NNs for production in the future.
Revised: We aim to further investigate its performance and develop this into practical NNs for production in the future.

Original: A P ROOF OF L EMMA A. 
Revised: A P ROOF OF L EMMA A.

Original: Before the proof of Lemma A.3, we need to introduce a lemma by Martens (2020) as:
Revised: Before the proof of Lemma A.3, we need to introduce a lemma by Martens (2020) as:

Original: Lemma A.1 (Approximation of Hessian Matrix in NNs (Martens, 2020)) .
Revised: Lemma A.1 (Approximation of Hessian Matrix in NNs (Martens, 2020)) .

Original: The Hessian matrix of
Revised: The Hessian matrix of

Original: NNs on a local minimum ˆ θ can be decomposed based on Fisher information matrix as
Revised: NNs on a local minimum ˆ θ can be decomposed based on Fisher information matrix as

Original: [Equation]
Revised: [Equation]

Original: where C is the total number of classes, ˆ y is the output (or prediction) of the network given input x , and H [ f ] c is the Hessian of the c -th component of ˆ y .
Revised: where C is the total number of classes, ˆ y is the output (or prediction) of the network given input x , and H [ f ] c is the Hessian of the c -th component of ˆ y .

Original: Speciﬁcally, for a well-trained NN, we could have 
Revised: Speciﬁcally, for a well-trained NN, we could have ∇ ˆ y (cid:96) i

Original: [Equation]{+(ˆ θ ) (cid:39) 
+}
Revised: {+(ˆ θ ) (cid:39) 
+}

Revised: thus H ˆ θ (cid:39) F ˆ θ .

Original: We further introduce a lemma of Poisson bootstrapping as
Revised: We further introduce a lemma of Poisson bootstrapping as

Original: Lemma A.2 (Poisson Bootstrapping (Efron, 1992; Chamandy et al., 2012)) .
Revised: Lemma A.2 (Poisson Bootstrapping (Efron, 1992; Chamandy et al., 2012)) .

Original: Given an inﬁnite number of samples, the bootstrap resampling weight ξ has the property that lim n →∞ Binomial (cid:0) n, 1 n (cid:1) =
Revised: Given an inﬁnite number of samples, the bootstrap resampling weight ξ has the property that lim n →∞ Binomial (cid:0) n, 1 n (cid:1) =

Original: Poisson (1) .
Revised: Poisson (1) .

Original: This approximation becomes precise in practice when n ≥ 100 .
Revised: This approximation becomes precise in practice when n ≥ 100 .

Original: Also, we know E [ ξ i ] = 

Revised: Also, we know E [ ξ i ] =

Original: and Var [ ξ i ] = 1 by the deﬁnition of Poisson distribution when n is large enough.
Revised: and Var [ ξ i ] = 1 by the deﬁnition of Poisson distribution when n is large enough.

Original: When bootstrap resampling from dataset S , each individual sample Z i =
Revised: When bootstrap resampling from dataset S , each individual sample Z i =

Original: ( X i , Y i ) has a probability of 1 n being picked, causing the weight ξ i to follow a binomial distribution as ξ i ∼ Binomial (cid:0) n, 1 n (cid:1) .
Revised: ( X i , Y i ) has a probability of 1 n being picked, causing the weight ξ i to follow a binomial distribution as ξ i ∼ Binomial (cid:0) n, 1 n (cid:1) .

Original: As a result, all weights ξ follow a multinomial distribution as ξ ∼ Multinomial (cid:0) n, 1 n , 1 n , . . .
Revised: As a result, all weights ξ follow a multinomial distribution as ξ ∼ Multinomial (cid:0) n, 1 n , 1 n , . . .

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

Original: [Equation]
Revised: [Equation]

Original: with the total number of samples constrained to be n .
Revised: with the total number of samples constrained to be n .

Original: When it comes to big data, i.e., n is prohibitively large, this multinomial resampling thus becomes rather slow.
Revised: When it comes to big data, i.e., n is prohibitively large, this multinomial resampling thus becomes rather slow.

Original: Based on Lemma A.2, we can now start to prove our lemma of approximating oracle prior covariance.
Revised: Based on Lemma A.2, we can now start to prove our lemma of approximating oracle prior covariance.

Original: Lemma A.3 (Approximation of Oracle Prior Covariance) .
Revised: Lemma A.3 (Approximation of Oracle Prior Covariance) .

Original: Given the deﬁnition of inﬂuence functions (Lemma 1) and Poisson bootstrapping (Lemma A.2), the covariance matrix of the oracle prior can be approximated by
Revised: Given the deﬁnition of inﬂuence functions (Lemma 1) and Poisson bootstrapping (Lemma A.2), the covariance matrix of the oracle prior can be approximated by

Original: [Equation]
Revised: [Equation]

Original: where F ˆ θ is Fisher information matrix (FIM); we omit the subscript S of ˆ θ S and ˆ θ S, ξ for notation conciseness, and ξ k is the bootstrap resampling weight in the k -th experiment.
Revised: where F ˆ θ is Fisher information matrix (FIM); we omit the subscript S of ˆ θ S and ˆ θ S, ξ for notation conciseness, and ξ k is the bootstrap resampling weight in the k -th experiment.

Original: Proof.
Revised: Proof.

Original: Recall that in the k -th bootstrap resampling process, the loss function is reweighted by ξ k = ( ξ k, 1 , ξ k, 2 , . . . , ξ k,n ) (cid:62) .
Revised: Recall that in the k -th bootstrap resampling process, the loss function is reweighted by ξ k = ( ξ k, 1 , ξ k, 2 , . . . , ξ k,n ) (cid:62) .

Original: Also, we have an inﬂuence matrix Ψ
Revised: Also, we have an inﬂuence matrix Ψ

Original: = ( ψ 1 , ψ 2 , . . .
Revised: = ( ψ 1 , ψ 2 , . . .

Original: , ψ n ) (cid:62) ∈ R n × D .
Revised: , ψ n ) (cid:62) ∈ R n × D .

Original: The original risk minimizer on the full dataset S is
Revised: The original risk minimizer on the full dataset S is

Original: [Equation]
Revised: [Equation]

Original: and the reweighted empirical risk minimizer (after bootstrapping) is deﬁned by
Revised: and the reweighted empirical risk minimizer (after bootstrapping) is deﬁned by

Original: [Equation]
Revised: [Equation]

Original: where we omit the subscript k for the sake of conciseness.
Revised: where we omit the subscript k for the sake of conciseness.

Original: Given the deﬁnition of inﬂuence function from Lemma 1, the difference between the two risk minizers above, ˆ θ S and ˆ θ S, ξ , can be written as
Revised: Given the deﬁnition of inﬂuence function from Lemma 1, the difference between the two risk minizers above, ˆ θ S and ˆ θ S, ξ , can be written as

Original: [Equation]
Revised: [Equation]

Original: As a result
Revised: As a result

Original: [Equation]
Revised: [Equation]

Original: the oracle prior can be transformed as
Revised: the oracle prior can be transformed as

Original: [Equation]
Revised: [Equation]

Original: Furthermore, based on the deﬁnition of inﬂuence function, we know Ψ (cid:62) 1 = 1 (cid:62) Ψ = 0 .
Revised: Furthermore, based on the deﬁnition of inﬂuence function, we know Ψ (cid:62) 1 = 1 (cid:62) Ψ = 0 .

Original: The term in Eq.
Revised: The term in Eq.

Original: (A.8) can be further writen to
Revised: (A.8) can be further writen to

Original: [Equation]
Revised: [Equation]

Original: From Lemma A
Revised: From Lemma A

Original: [Equation]
Revised: [Equation]

Original: 2 we know E [ ξ i ] = 1 and Var [ ξ i ] = 1 when n ≥ 100 .
Revised: 2 we know E [ ξ i ] = 1 and Var [ ξ i ] = 1 when n ≥ 100 .

Original: We also know that
Revised: We also know that

Original: [Equation]
Revised: [Equation]

Original: This gives rise to the ﬁnal solution that
Revised: This gives rise to the ﬁnal solution that

Original: [Equation]
Revised: [Equation]

Original: I n in Eq.
Revised: I n in Eq.

Original: (A.10) is an identity matrix with size n × n ; Eq.
Revised: (A.10) is an identity matrix with size n × n ; Eq.

Original: (A.11) is true by re-applying the property of inﬂuence functions
Revised: (A.11) is true by re-applying the property of inﬂuence functions

Original: that Ψ (cid:62) 1 = 1 (cid:62) Ψ
Revised: that Ψ (cid:62) 1 = 1 (cid:62) Ψ

Original: = 0 ; and Eq.
Revised: = 0 ; and Eq.

Original: (A.15) is achieved by the result from Lemma
Revised: (A.15) is achieved by the result from Lemma

Original: A.1.
Revised: A.1.

Original: Summarizing all the above results, the oracle prior covariance can be approximated by
Revised: Summarizing all the above results, the oracle prior covariance can be approximated by

Original: [Equation]
Revised: [Equation]

Original: B P ROOF OF L EMMA
Revised: B P ROOF OF L EMMA

Original: Lemma B.1 (Optimal Posterior for PAC-Bayes Information Bottleneck) .
Revised: Lemma B.1 (Optimal Posterior for PAC-Bayes Information Bottleneck) .

Original: Given an observed dataset
Revised: Given an observed dataset

Original: S ∗ , the optimal posterior p ( w | S ∗ ) of PAC-Bayes IB in Eq.
Revised: S ∗ , the optimal posterior p ( w | S ∗ ) of PAC-Bayes IB in Eq.

Original: (5) should satisfy the following form that
Revised: (5) should satisfy the following form that

Original: [Equation]
Revised: [Equation]

Original: where U S ∗ ( w ) is the energy function deﬁned by
Revised: where U S ∗ ( w ) is the energy function deﬁned by

Original: [Equation]
Revised: [Equation]

Original: and Z ( S ) is the normalizing constant.
Revised: and Z ( S ) is the normalizing constant.

Original: Proof.
Revised: Proof.

Original: Recap that the PAC-Bayes information bottleneck in Eq.
Revised: Recap that the PAC-Bayes information bottleneck in Eq.

Original: (5) is
Revised: (5) is

Original: [Equation]
Revised: [Equation]

Original: Given an observed dataset S ∗ , our object of interest is to ﬁnd the optimal posterior p ( w | S ∗ ) that minimizes the L PIB .
Revised: Given an observed dataset S ∗ , our object of interest is to ﬁnd the optimal posterior p ( w | S ∗ ) that minimizes the L PIB .

Original: Consider a constraint of posterior distribution that
Revised: Consider a constraint of posterior distribution that

Original: [Equation]
Revised: [Equation]

Original: we can formulate the problem by
Revised: we can formulate the problem by

Original: [Equation]
Revised: [Equation]

Original: A Lagrangian can hence be built to relax the above optimization problem by
Revised: A Lagrangian can hence be built to relax the above optimization problem by

Original: [Equation]
Revised: [Equation]

Original: with α = { α S |∀ S ∼ p
Revised: with α = { α S |∀ S ∼ p

Original: ( X, Y ) ⊗ n } corresponding to Lagrange multipliers; we denote the empirical risk by ˆ L S ( w ) = 1 n (cid:80) ni
Revised: ( X, Y ) ⊗ n } corresponding to Lagrange multipliers; we denote the empirical risk by ˆ L S ( w ) = 1 n (cid:80) ni

Original: =1 (cid:96) i ( w ) .
Revised: =1 (cid:96) i ( w ) .

Original: Differentiating (cid:101)
Revised: Differentiating (cid:101)

Original: L PIB w.r.t.
Revised: L PIB w.r.t.

Original: p ( w | S ∗ ) results in
Revised: p ( w | S ∗ ) results in

Original: [Equation]
Revised: [Equation]

Original: Setting ∇ p ( w | S ∗ ) (cid:101)
Revised: Setting ∇ p ( w | S ∗ ) (cid:101)

Original: L PIB = 0 and solving for p ( w | S ∗ ) yields
Revised: L PIB = 0 and solving for p ( w | S ∗ ) yields

Original: [Equation]
Revised: [Equation]

Original: The second exponential term exp (cid:110) − 1 − α S ∗ β (cid:111) is the partition function that normalizes the posterior distribution.
Revised: The second exponential term exp (cid:110) − 1 − α S ∗ β (cid:111) is the partition function that normalizes the posterior distribution.

Original: Denoting the normalization term as Z ( S ) , we hence obtain the optimal posterior solution as
Revised: Denoting the normalization term as Z ( S ) , we hence obtain the optimal posterior solution as

Original: [Equation]
Revised: [Equation]

Original: C C OMPUTATION FOR PIB OBJECT
Revised: C C OMPUTATION FOR PIB OBJECT

Original: Our Alg. 2 presents the training process based on the proposed PIB objective function.
Revised: Our Alg. 2 presents the training process based on the proposed PIB objective function.

Original: It differs from vanilla SGD on two aspects: (1) the objective function (also called energy function U ) consists of the negative log-likelihood plus a regularization term log p ( w ) (line 2); (2) the parameter w gets update by the gradient of the energy function plus an isotropic Gaussian noise (line 3).
Revised: It differs from vanilla SGD on two aspects: (1) the objective function (also called energy function U ) consists of the negative log-likelihood plus a regularization term log p ( w ) (line 2); (2) the parameter w gets update by the gradient of the energy function plus an isotropic Gaussian noise (line 3).

Original: Since (2) has no additional computational cost, the major change is the regularization term log p ( w ) in (1), described in Eq.
Revised: Since (2) has no additional computational cost, the major change is the regularization term log p ( w ) in (1), described in Eq.

Original: (18) as
Revised: (18) as

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst term is just matrix-vector product based on the approximation of Σ 0 in Eq.
Revised: The ﬁrst term is just matrix-vector product based on the approximation of Σ 0 in Eq.

Original: And the second term can be written by
Revised: And the second term can be written by

Original: [Equation]
Revised: [Equation]

Original: where λ i is the eigenvalue of Σ 0 , which can be obtained by efﬁcient eigen-decomposition techniques.
Revised: where λ i is the eigenvalue of Σ 0 , which can be obtained by efﬁcient eigen-decomposition techniques.

Original: D E XPERIMENTAL P ROTOCOL
Revised: D E XPERIMENTAL P ROTOCOL

Original: All experiments are conducted on MNIST (LeCun et al., 1998) or CIFAR-10 (Krizhevsky et al., 2009).
Revised: All experiments are conducted on MNIST (LeCun et al., 1998) or CIFAR-10 (Krizhevsky et al., 2009).

Original: We design two multi-layer perceptron (MLPs): MLP-Small and MLP-Large, where MLPSmall is a two-layer NN, i.e., 784(3072)-512-10, and MLP-Large is a ﬁve-layer NN, i.e., 784(3072)- 100-80-60-40-10.
Revised: We design two multi-layer perceptron (MLPs): MLP-Small and MLP-Large, where MLPSmall is a two-layer NN, i.e., 784(3072)-512-10, and MLP-Large is a ﬁve-layer NN, i.e., 784(3072)- 100-80-60-40-10.

Original: The number of input units is 784 with permutation MNIST inputs or 3072 with permutation CIFAR-10 inputs.
Revised: The number of input units is 784 with permutation MNIST inputs or 3072 with permutation CIFAR-10 inputs.

Original: For the performance comparison in §5.5, we use a reduced version of VGG-Net where two blocks are cut due to the memory constraint.
Revised: For the performance comparison in §5.5, we use a reduced version of VGG-Net where two blocks are cut due to the memory constraint.

Original: In the general setting, we pick
Revised: In the general setting, we pick

Original: Adam optimizer (Kingma & Ba, 2014) to accelerate the convergence of NN training.
Revised: Adam optimizer (Kingma & Ba, 2014) to accelerate the convergence of NN training.

Original: We use one
Revised: We use one

Original: RTX 3070 GPU for all experiments.
Revised: RTX 3070 GPU for all experiments.

Original: Speciﬁcally for the Bayesian inference experiment, the batch size is picked within { 8 , 16 , 32 , 64 , 128 , 256 , 512 } ; learning rate is in { 1 e − 4 , 1 e − 3 , 1 e − 2 , 1 e − 1 } ; weight decay of (cid:96) 2norm is in { 1 e − 3 , 1 e − 4 , 1 e − 5 , 1 e − 6 } ; noise scale of SGLD is in { 1 e − 4 , 1 e − 6 , 1 e − 8 , 1 e − 10 } ; β of
Revised: Speciﬁcally for the Bayesian inference experiment, the batch size is picked within { 8 , 16 , 32 , 64 , 128 , 256 , 512 } ; learning rate is in { 1 e − 4 , 1 e − 3 , 1 e − 2 , 1 e − 1 } ; weight decay of (cid:96) 2norm is in { 1 e − 3 , 1 e − 4 , 1 e − 5 , 1 e − 6 } ; noise scale of SGLD is in { 1 e − 4 , 1 e − 6 , 1 e − 8 , 1 e − 10 } ; β of

Original: PAC-Bayes IB is in { 1 e − 1 , 1 e − 2 , 1 e − 3 } ; and the dropout rate is ﬁxed as 0 .
Revised: PAC-Bayes IB is in { 1 e − 1 , 1 e − 2 , 1 e − 3 } ; and the dropout rate is ﬁxed as 0 .

Original: 1 for the dropout regularization.
Revised: 1 for the dropout regularization.