PAC-Bayes Compression Bounds So Tight That They Can Explain Generalization
Anonymous Author(s)
AffiliationAddress email
Abstract
While there has been progress in developing non-vacuous generalization boundsfor deep neural networks, these bounds tend to be uninformative about why deeplearning works.
In this paper, we develop a compression approach based on quantizing neural network parameters in a linear subspace, profoundly improving onprevious results to provide state-of-the-art generalization bounds on a variety oftasks, including transfer learning.
We use these tight bounds to better understandthe role of model size, equivariance and the implicit biases of optimization, forgeneralization in deep learning.
Notably, we find large models can be compressedto a much greater extent than previously known, encapsulating Occam’s razor.
Introduction
There has been great effort to understand why deep learning models generalize so well on commondatasets, despite the fact that these models have many more parameters than the number of trainingdatapoints and can fit random labels [70].
These observations are not explainable through classicalstatistical learning theory such as VC dimension or Rademacher complexity which focus on uniformconvergence over the hypothesis class [51].
The PAC-Bayes framework, in contrast, gives a convenient way of constructing bounds where the generalization gap depends on the deep learning model(or models) that are actually found by training rather than the hypothesis set as a whole.
Using thisframework, many different potential explanations have been proposed drawing on properties of adeep learning model that are induced by the training dataset, such as low spectral norm [55], noisestability [2], flat minima [30], derandomization [53], robustness, and compression [2, 71].
In this work, we show that neural networks are substantially more compressible than previouslybelieved when paired with structured training datasets.
Constructing tighter generalization boundsthan have been previously achieved, we show that this compression alone is sufficient to explainmuch of the generalization of neural networks.
Contributions of this Paper
We develop a new approach for training compressed neural networks that adapts the com- pressed size to the difficulty of the problem.
We train in selected random subspaces of the parameters and perform learned quantization.
Consequently, we achieve extremely low compressed sizes for neural networks at a given accuracy level, which is a requirement for our tight bounds.
Using a prior encoding Occam’s razor and our compression scheme, we construct the best generalization bounds to date on image datasets, both with data-dependent and dataindependent priors.
We also show how transfer learning improves compression and thusour generalization bounds, explaining the practical performance benefits of pre-training.
[Table]
Our bounds
[Table]
Through the lens of compression and our bounds, we are able to explain why deep learningmodels generalize on structured datasets like CIFAR-10, but not when structure is brokensuch as by shuffling the pixels or shuffling the labels, the benefits of equivariant models,and whether the implicit regularization of SGD is necessary for generalization.
Related Work
In this section, we delineate recent techniques developed for achieving non-vacuous PAC-Bayesbounds for neural networks.
Optimizing the PAC-Bayes bound
Dziugaite and Roy [18] obtain the first non-vacuous generalization bounds for deep stochastic neural networks on binary MNIST.
To do so, the authorsconstruct a relaxation of the Langford and Seeger [40] bound and optimize it with the objective offinding a posterior distribution that covers a large volume of low-loss solutions around a local minimum obtained using SGD.
Similarly, Rivasplata et al.
[58] extend the optimization to other boundsby developing novel relaxations of other PAC-Bayes bounds based on previous work from Blundellet al.
Compressing Models
It is well-established that neural networks are robust to small perturbations [29, 30, 39, 38, 36, 55, 8].
The work of Arora et al.
[2] leverages this connection to establish acompression-based approach that uses noise stability and the ability to reconstruct weight matriceswith random projections is used to study generalization of neural networks.
Following the insighton compressibility, Zhou et al.
[71] develops a PAC-Bayes bound that uses the representation of acompressed model in bits, and they also add noise stability through their use of Gaussian posteriorsand Gaussian mixture priors.
Furthermore, they achieve the smallest model representations throughpruning and quantization following ideas of Han et al.
[27] and Cheng et al.
Data-Dependent Priors
Dziugaite et al.
[19] demonstrated that for linear PAC-Bayes boundssuch as Thiemann et al.
[62], a tighter bound can be achieved by choosing the prior distribution to bedata-dependent.
That is, the prior is trained to concentrate around low loss regions on held-out data.
More precisely, the authors show that the optimal data-dependent prior is the conditional expectationof the posterior given a subset of the training data.
They approximate this data-dependent prior bysolving a variational problem over Gaussian distributions.
They evaluate the bounds for SGD-trainednetworks on data-dependent priors obtaining tight bounds on MNIST, Fashion MNIST, and CIFAR10.
In a similar vein, P´erez-Ortiz et al.
[57] combine the data-dependent priors approach Dziugaiteet al.
[19] with the PAC-Bayes with Backprop (PBB) approach from Rivasplata et al.
[58] to obtainstate-of-the-art PAC-Bayes non-vacuous bounds for MNIST and CIFAR-10 with data-dependentpriors.
Table 1 summarizes the PAC-Bayes bounds obtained by these works alongside our bounds.
Our upper bound on the KL term closely follows Zhou et al.
[71], but with key improvements:(i) we train in lower dimensional subspaces using intrinsic dimensionality and FiLM subspaceswhich proves to be more effective and adaptable than pruning, (ii) we consider a more effectivequantization scheme with variable length codes and quantization aware training, and (iii) we takeadvantage of the increased compression provided by transfer learning and data dependent priors.
As a result, we obtain state-of-the-art non-vacuous bounds for MNIST, Fashion MNIST, CIFAR-10, and
CIFAR-100 in both the from scratch and transfer learning setting as shown in Table 1.
Motivatedby the tightness of these bounds, we seek to better understand the role of equivariance, stochasticoptimization, and data dependent bounds, for explaining generalization.
Ding et al.
[16] is another closely related work which optimizes the PAC-Bayes bounds to predict thetransferability of various upstream models to a downstream task.
In contrast, we aim to explain whyneural networks generalize through compression.
This is different than predicting transferabilityas there are several factors other than compression (which we demonstrate is sufficient, but notnecessary for generalization) that could make a model perform well on a downstream task.
PAC-Bayes Bounds
PAC-Bayes bounds are fundamentally an expression of Occam’s razor: simpler descriptions of thedata generalize better.
As an illustration, consider the finite hypothesis generalization bound.
Letˆ R ( h )
= 1 n (cid:80) ni
=1 ℓ ( h ( x i ) , y i ) be the empirical risk of a hypothesis h from a finite hypothesis set
H for the 0 - 1 loss ℓ , and R ( h ) =
E [ ˆ R ( h )] its population risk.
With probability at least 1 − δ , thepopulation risk for a hypothesis h using n data samples satisfies
[Equation]
bounded by the empirical risk and a complexity term log |H| counting the number of bits needed tospecify any h ∈ H .
But what if we don’t believe that each hypothesis is equally likely?
If we consider that a given priordistribution P over H concentrates around the optimal hypothesis, then we can construct a variablelength code that uses fewer bits for the more likely hypothesis.
With the best possible code for P ,a given hypothesis h will take log 2 1 P ( h ) bits to represent.
Using this prior may allow us to payfor a smaller complexity term even if the hypothesis set is large, so long as the hypotheses that areconsistent with the data are also likely under the prior.
The number of bits payed can further be reduced from log 2 1 P ( h ) to KL ( Q ∥ P ) by considering adistribution of “good” solutions Q .
If we don’t care which element of Q we arrive at, we can gainsome bits back from this insensitivity (which could be used to code a separate message).
The averagenumber of bits to code a sample from Q using the prior P is the cross entropy H ( Q, P ) , and onegets on average the entropy H ( Q ) bits back from being agnostic about which sample h ∼ Q to use,yielding the KL-divergence between Q and P : H ( Q, P ) − H ( Q ) =
KL ( Q ∥ P ) .
With these improvements of the finite hypothesis bound: replacing log |H| with KL ( Q ∥ P ) , andsampling a hypothesis h ∈ H , one arrives (with minor bookkeeping) at the PAC-Bayes boundintroduced in McAllester [49].
With probability at least 1 − δ ,
[Equation]
Many refinements of Eq.
(2) have been developed [40, 48, 7, 62] but retain the same character.
Thelower the ratio of the KL to the number of data points n , the tighter the gap between empiricaland expected risk.
In this work, we use the tighter Catoni [7] variant of the PAC-Bayes bound (see
Appendix I for details).
Universal Prior.
Connecting back to Occam’s razor, one can define a prior that explicitly penalizesthe minimum compressed length of the hypothesis, also known as the universal prior [60]: P ( h ) =
2 − K ( h ) /Z , where K is the prefix Kolmogorov complexity [33] of h (the length of the shortestprogram that produces h and also delimits itself), and Z ≤ 1 .
1 Using a point mass posterior on asingle hypothesis h ∗ , KL ( Q ∥ P ) from the bound reduces to
[Equation]
where l ( h ) is the length of a given program that reproduces h not including the delimiter, which wecan use when upper bounding the KL.
For convenience, we can condition on using the same methodfor compression and decompression for all elements of the prior.
If we can reduce the compressedlength l ( h ) of the hypotheses we find that fits the training data, then we can improve the tightness ofthe PAC-Bayes bound.
Model Compression.
Model compression aims to find a nearly equivalent model that can be expressed fewer bits and is useful for deployments on mobile devices or for improving inference timeon specialized hardware [9, 10].
For computing PAC-Bayes generalization bounds, however, onlythe model size matters.
Therefore, we can employ compression methods which may otherwise beunfavorable in practice due to worse computational requirements.
Pruning and quantization areamong the most widely used methods for model compression.
In this work, we rely on quantization(Section 4.2) to achieving tighter generalization bounds.
Tighter Generalization Bounds via Adaptive Subspace Compression
Training a neural network involves taking many gradient steps in a high-dimensional space R D .
Although D may be large, the loss landscape has been, in some sense, found to be simpler thantypically believed [14, 26, 23].
Analogous to the notion of intrinsic dimensionality more generally,
Li et al.
[44] attempt to find the lowest dimensional subspace in which the network can be trainedand still fit the training data.
The weights of a neural network θ ∈ R D are parametrized in terms ofan initialization θ 0 and a projection w ∈ R d to a lower dimensional subspace through a fixed matrix
P ∈ R D
×
d ,
[Equation]
To facilitate favorable conditioning during optimization, P is chosen to be approximately orthonormal P ⊤ P
≈ I d × d and for scalability, Li et al.
[44] utilize random normal matrices of the form
P ∼ N (0 , 1) D × d / √ D , and their sparse approximations [45, 42].
In its original form, intrinsic dimensionality (ID) is used as a scientific tool to measure the complexity of the learning task.
Unlike other methods like pruning, ID has the benefit of scaling dramaticallyup or down in size with difficulty of the problem.
While ID was not used with quantization as a fullfledged model compression method, we find that it can be.
We also highlight that concepts similar tothe intrinsic dimensionality of a model have been explored in the context of model compression forestimating bounds as seen in Arora et al.
For our work, the ability to find the intrinsic dimensiond ≪ D has profound implications for the compressibility of the models, and therefore our ability toconstruct generalization bounds.
As demonstrated by Zhou et al.
[71], the compressibility of neuralnetwork models has a direct connection to generalization and immediately allows us to computenon-vacuous PAC-Bayes bounds for transfer learning.
Building upon ID, we describe our approachnext.
4.1 Finding Better Random Subspaces to Train In
To further improve upon the scalability and effectiveness of the projections P used by Li et al.
[44],we introduce three new methods.
Kronecker Sum Projector :
Using the Kronecker product ⊗ , we construct the matrix
P ⊕ =
( 1 ⊗ R 1 + R 2 ⊗ 1 ) / √ 2 D where R 1 , R 2 ∼ N (0 , 1) √ D × d and 1 is the vector of all onesin R √ D .
From the fact that R 1 ⊥⊥ R 2 and the entries are standard normal, one can verify that
P ⊤⊕ P ⊕ =
I
d × d + O (1 / √ D ) .
Kronecker Product Projector : Alternatively, we can form the matrix P ⊗ = Q 1 ⊗ Q 2 / √ D withthe smaller Q 1 , Q 2 ∼ N (0 , 1) √ D × √ d , and again this matrix is approximately orthogonal: P ⊤⊗ P ⊗ =( Q ⊤ 1 Q 1 / √ D ) ⊗ ( Q ⊤ 2 Q 2 / √ D )
=
I ⊗ I
+ O ( D − 1 / 4 )
=
I d × d + O ( D − 1 / 4 ) .
2 a perfect square, we concatenate a dense random matrix to pad out the difference between D , d , and a perfect square.
As ( √ D + 1)
= D + 2 √ D
+ 1 , we have that the size of
1.5 d ×100.
T r a i n
A cc u r a c y
[Figure]
1.5 d ×10
R un t i m e ( G P U
H o u r s )
[Figure]
FiLM
[Figure]
Kronecker
[Figure]
3 d ×10 4 0.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
i
[Figure]
The matrix vector multiply w (cid:55)→ Pw for both the above projectors can performed in time
O ( d √ D ) and O ( √ dD ) respectively, rather than the O ( dD ) that is required by the dense randommatrix.
Figure 1 demonstrates the runtime speedup and training performance improvement in comparison to the methods using by Li et al.
Notably, the Kronecker-structured projections retainthe fidelity of the dense random matrix while being orders of magnitude faster than the alternativeoperators when scaling to larger values of d .
FiLM projector : It is well known that the BatchNorm parameters have an outsized effect on thedownstream task performance relative to their size, and this fact has been used in Featurewise independent Linear Modulation (FiLM) [56, 17] for efficient control of neural networks in many different settings.
Several authors have explored performing finetuning for transfer learning solely onthese parameters and the final linear layer [35].
Drawing on these observations, we construct aprojection matrix P FiLM where only columns corresponding with BatchNorm or head parametersare non-zero and sampled from N (0 , 1) d / √ D , which we also show in Figure 1.
While the FiLMprojector is highly effective for transfer learning (shown in Figure 1 left), the performance saturatesquickly when training from scratch.
For this reason, when training from scratch we employ the sum
P FiLM+ ⊗ =
( P FiLM + P ⊗ ) / √ 2 , which outperforms the two projectors individually as shown in
Figure 1 right.
4.2 Quantization Scheme and Training
Through quantization, the average number of bits used per parameter can be substantially reduced.
When optimizing purely for model size rather than efficiency on specialized hardware, one canchoose non-linearly spaced quantization levels which are learned, and use variable length codingschemes as shown in Han et al.
On the other side, the straight through estimator has beencentral to learning weights in binary neural networks [31].
We combine these ideas to simultaneouslyoptimize the quantized weights and the quantization levels for maximum compression.
Given the full precision weights w =
[ w 1 , . .
w d ] ∈ R d and a vector c =
[ c 1 , . . .
c L ] ∈ R L of Lquantization levels, we construct the quantized vector ˆ w =
[ ˆ w 1 , . .
, ˆ w d ] such that ˆ w i =
c q ( i ) whereq ( i ) =
arg min k | w i − c k | .
The quantization levels c are learned alongside w , where the gradientsare defined using the straight through estimator [3, 68]:
[Equation]
and
[Equation]
We initialize c with uniform spacing between the minimum and maximum values in parametervector w , which performs better than k-means initialization, as also observed in Choi et al.
To further compress the network, we use a variable length code in the form of arithmetic coding[47], which takes advantage of the fact that certain quantization levels are more likely than others.
Given probabilities p k (empirical fractions) for cluster c k , arithmetic coding of w takes at most
KL ( Q
P
← GET KL ( ˆ w
P )
▷ ▷
9: function TRAIN QUANTIZE
▷
Initialize c ← GET CLUSTERS
11: for i = 1 to quant epochs do 12: c ← c − ρ ∇ c L ( w, c ) w ← w − ρ ∇ w L ( w, c ) 13: end for 14: return ˆ w 15: end function 16: function GET KL ( ( ˆ w, P ) )
17: c , count ← GET UNIQUE VALS COUNTS ( ˆ w ) 18: message size ← DO ARITHMETIC ENCODING ( ˆ w , c , count ) 19: message size ← message size + hyperparam search ▷
(Section 4.3) 20: return message size + 2 × log ( message size ) 21: end function ⌈ d × H ( p ) ⌉ + 2 bits, where H ( p ) is the entropy H ( p ) =
− (cid:80) k
p k log 2 p k .
For a small number ofquantization levels, we find it makes a small improvement in the compression over more common
Huffman coding.
In summary, we use ⌈ d × H ( p ) ⌉ + 2 bits for coding the quantized weights ˆ w , 16 L bits for thecodebook c (represented in half precision), and additional L × ⌈ log 2 d ⌉ bits for representing theprobabilities p k , arriving at l ( w ) ≤ ⌈ d × H ( p )
+ L × (16 + ⌈ log 2 d ⌉ ) + 2 .
4.3 Subspace Dimension Optimization and Hyperparameters in the Universal Prior
The smaller the chosen intrinsic dimension d , the more similar θ is to the initialization θ 0 in Eq.
Consequently, that value of θ is more likely under the universal prior given the shorter descriptionlength.
Note that in this prior, we condition on the random seed used to generate θ 0 and P .
Aswe optimize over different parameters such as the subspace dimension d = 1 , .., D , and possiblyother hyperparameters such as the learning rate, or number of quantization levels L , we must encodethese into our prior and thus pay a penalty for optimizing over them.
We can accomodate this verysimply by considering the hypothesis h as not just specifying the weights, but also specifying thesehyperparameters:
h = ( θ, d, L, lr) , and therefore using the universal prior P ( h ) = 2 − K ( h ) /Z
wepay additional bits for each of these quantities: K ( h ) ≤ K ( θ | d, L ) + K ( d ) + K ( L ) + K (lr) .
Ifwe optimize over a fixed number H of distinct values known in advance for a given hyperparametersuch as L , then we can code L using this information in only log 2 ( H ) bits.
In general, we can alsobound the dimensionalities searched over by the maximum D so that K ( d ) ≤ ⌈ log 2 D ⌉ in any case.
For transfer learning, we replace θ 0 with a learned initialization θ D that is found using the pretrainingtask and data D .
With the ID compression, the universal prior P ( h |D )2 − K ( h |D ) /Z will place higherlikelihood on solutions θ that are close to the pretraining solution θ D .
Empirical Non-Vacuous Bounds
5.1 Non-Vacuous PAC-Bayes Bounds
To compute our bounds on benchmark datasets, we use the joint approach of projected subspacetraining, parameter quantization, and optimizing the subspace dimension with respect to the bounds
[Table]
as summarized in Algorithm 1.
We also optimize other bound hyperparameters such as the learningrate of the quantization-aware training.
We offer more details about bound optimization in Appendix D.
Our results for data-independent priors, i.e., priors that do are not pre-trained on a subsetof the training data, are shown in Table 2.
We derive the first non-vacuous bounds on FashionMNIST, CIFAR-10, and CIFAR-100 without data-dependent priors.
These results are of significanceas we argue in Section 5.2 that using data-dependent priors are not explanatory about the learning process.
In particular, we improve over the compression bound results obtained by Zhou et al.[71] on MNIST from < 46% to 11 .
We also dramatically improve their compression resultsas we reduce the compressed size for the best bound by 94% bringing it down from 6 .
KB to0 .
38 KB for MNIST with LeNet5.
The tightness of our SOTA subspace compression bounds allowsus to improve the understanding of several deep learning phenomena as discussed in Section 6.
See
Appendix A for additional results.
5.2 Data-Dependent Priors
So far, we demonstrated the strength of our bounds on data-independent priors, where we considerably improve on the state-of-the-art.
However, a number of recent papers have considereddata-dependent priors as a way of achieving tighter bounds [57, 19].
In this setup, the training data
D = { ( x i , y i ) } ni =1 is partitioned into two parts, D a and D b , with length n − m and m .
The firstdataset is used to construct a data dependent prior P ( h |D a ) , and then the bound is formed over theremaining part of the process: the adaptation of the prior P ( h |D a ) to the posterior Q ( h ) using thedata D b .
The empirical risk is computed over D b only.
Intuitively, using dataset D a it is possible toconstruct a much tighter prior over the possible neural network solutions.
In our setting, similar to transfer learning, we use the prior P D a ( θ ) = 2 − K
( θ | θ D a ) /Z
where forcompression we use θ = θ D a + Pw , and θ D a is the solution found by training the model (withoutrandom projections) on the data D a rather than initializing randomly.
With these data-dependentpriors, we achieve the best bounds in Table 2.
However, our adaptive approach exposes a significant downside of data-dependent priors.
To theextent that PAC-Bayes bounds can be used for explanation, in this case they only yield insights intothe procedure used to adapt P D a ( θ ) to Q using D b , any learning that is done in finding P D a ( θ ) isnot constrained or explained by the bound (we have no information on the compressibility of θ D a ).
Given the ability to adapt the size of the KL to the difficulty of the problem, it is possible to squeezeall of the learning into P D a ( θ ) and none in this adaption to Q .
This happens as the KL → 0 ,which we find happens empirically (or very nearly so) for many splits of the data where n − m islarge.
Setting Q ( θ ) =
[ θ = θ D a ] , the KL has only the contribution from the optimization over d : 251 KL ( Q || P D a ) ≤ log D we find that the bound is nothing more than a variant of the simple Hoeffdingbound where D b is the validation set R ( θ D a ) ≤
ˆ R D b ( θ D a ) + (cid:113) log( Dm/δ )
+2 2 m −.
This phenomenon can be seen in Appendix C, where optimization over d = 1 , ..., D leads to verysmall values of the KL when n − m is a large fraction of n , and hence highly simple solutionsapproaching the Hoeffding bound or confidence interval.
This demonstrates that data-dependentpriors based PAC-Bayes bounds are not particularly meaningful, and attempts to optimize these bounds will lead inevitably to the statistical bounds formed from the validation set such as a simple
Hoeffding bound or confidence interval.
While these bounds are correct, they are only informativeabout the difference in performance of the pretrained model on validation data and test data.
Note that this is in contrast with the data-independent bounds that constrain and explain generalization for the entirety of the learning process.
Similarly, our transfer learning bounds meaningfullyconstrain what happens in the finetuning on the downstream task, but they do not constrain the priordetermined from the upstream task.
5.3 Non-Vacuous PAC-Bayes Bounds for Transfer Learning
By directly interpreting PAC-Bayes bounds through the lens of compression, we immediately see thebenefits of employing an upstream dataset for transfer learning.
In fact, transfer learning allows us toconstrain the prior around parameters we know are consistent with the upstream task, and hopefullythe downstream task as well, reducing the KL-divergence between the prior and the posterior andleading to even tighter bounds as we show in Table 2.
We obtain unequivocally tighter boundsusing transfer learning for data-independent priors, providing a theoretical certification that transferlearning can improve generalization.
Our PAC-Bayes transfer learning approach also indicates thattransfer learning can boost generalization whenever codings optimized on a pre-training task aremore efficient for encoding a downstream posterior than an a priori guess made before seeing data.
Incontrast, downstream tasks which differ greatly from the upstream task may only be consistent withmodels that are not compressible under the learned prior, a scenario which may describe negativetransfer.
As pre-training becomes the default in deep learning, our bounds open the door to providinglearning guarantees for state-of-the-art techniques.
See Appendix B for more details.
Understanding Generalization through PAC-Bayes Bounds
Unlike the classical statistical learning theory viewpoint of uniform convergence, where the focus ison the properties of the hypothesis class as a whole, the ability of a given network to be compressedis highly dependent on the actual dataset that the model was trained on.
In this respect, the PACBayes perspective shows us that the ability to compress and therefore generalize is not a function ofthe hypothesis space but rather of individual functions: there are many networks in that space thatfit corrupted data and provably cannot be compressed (see Appendix A.1).
But these models are notthe ones that we find when training on real world datasets, as real world datasets have a tremendousamount of structure.
If they did not, it would be pointless trying to learn from them as famouslyargued by Hume [32] and the No Free Lunch theorem [65, 25].
When training on these structureddatasets that have low Kolmogorov complexity, the compressibility of the dataset is reflected in thecompressiblity of the model.
Breaking the structure in the dataset also breaks the structure in themodel and leads to worse generalization bounds for shuffled labels, shuffled pixels, and equivariancemore generally.
[Figure]
MLPs vs CNNs: It is well known that convolutional neural networks (CNNs) generalize muchbetter than standard multilayer perceptrons (MLPs) even when controlling for the number of parameters.
We find that this fact is reflected in the improved compressibility of CNNs and our generalization bounds, shown in Figure 2 (left) on the image dataset CIFAR-10 using the same number ofparameters for the two models.
The CNN is able find a more structured and lower description lengthexplanation of the training data.
Shuffled Pixels:
However, when this image structure is broken by shuffling the pixels we find that
CNNs are no better at generalizing than MLPs, and this is reflected in the compressed sizes.
The
CNNs become substantially less compressible when trained on this dataset, and hence our boundsshow them generalizing worse than MLPs, see Figure 2 (center).
Not using this image structure (andinvariant to a global permutation of the features), MLPs do not suffer when this structure is brokensince they never used it in the first place.
Shuffled Labels: When the structure of the dataset is entirely broken by shuffling the labels, thecompressibility of the models (both for CNNs and MLPs) are lost.
Regardless of the subspacedimension used, our generalization bounds are all at 100% error as shown in Figure 2 (right).
Usinglow subspace dimensions it is not possible to fit the training data, and when using a large enoughdimension to fit the data the compressed size of the model is larger than the training data and hencethe generalization bounds are vacuous.
Equivariance: Designing models which are equivariant to certain symmetry transformations hasbeen a guiding principle for the development of data efficient neural networks in many domains[12, 13, 63, 64, 22, 34].
While intuitively it is clear that respecting dataset symmetries severelyimproves generalization, relatively little has been proven for neural networks [21, 72, 4, 20].
Wecompress and evaluate rotationally equivariant and non-equivariant Wide ResNets [64, 69] trainedon MNIST and a rotated version of MNIST.
As shown in Appendix E, the rotationally equivariantmodels are more compressible and provably generalize better than their non equivariant counterpartswhen paired with a dataset that also has the rotational symmetry.
Discussion
In this work, we have constructed a new method for compressing deep learning models which ishighly adaptive to the level of structure in the model and the training dataset.
According to the
Occam’s prior which considers shorter compressed length models to be more likely, we constructstate-of-the-art generalization bounds across a variety of settings.
Through our bounds and theviewpoint of Occam’s razor and compressibility, we show how generalization relates to the structurein the dataset and the structure in the model, and we are able to explain aspects of neural networkgeneralization for natural image datasets, shuffled pixels, shuffled labels, equivariant models, andfor stochastic training.
Limitations We note that despite the power of this compression scheme and our bounds, there arecertain aspects of deep learning generalization that remain unexplained.
Our compression boundsstill show a preference towards models with a smaller number of base parameters as shown in Appendix G, despite the fact that larger models tend to generalize better.
While we have achievedbetter model compression than previous works, it is unlikely that we are close to theoretical limits.
We hope that with a further improved compression scheme (possibly using nonlinear parameter projections), we might find that larger deep learning model solutions are compressible to smaller sizesthan smaller models.
There is also a question of whether models not trained in a special way can berelated and bounded by the performance of the compressed models, perhaps leveraging ideas from
Nagarajan and Kolter [50] and others investigating this from a theoretical lens, and this would behelpful for applying our findings more broadly.
Additionally, while our bounds show that the compressibility of our models implies generalization, we have no statements in the reverse direction.
Webelieve however that the compression approach and Occam’s razor have yet untapped explanatorypower in deep learning.
Checklist
For all authors...
(a)
Do
the main
claims
made in the abstract
and introduction accurately reflect
the
pa
per’s contributions and
scope
[Yes
We
provide evidence for the claims made in the
abstract
and introduction in Section 4, Section 5.1, and Section
(b) Did you describe the limitations of your work?
[Yes]
We
describe
limitations in Sec-
tion 7.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Ap-
pendix I.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix I.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix D.
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix D.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] See Appendix J.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
Our code is submitted in the supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data
you’re using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifi-
able information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N
/A]
Supplementary Material for
PAC-Bayes Compression Bounds So Tight That They Can Explain Generalization
Appendix Outline
The appendix is organized as follows.
• In Appendix A, we report results for additional bounds for SVHN and ImageNet.
We also
report the compression size corresponding to our best bound values and compare it to the
compression size obtained through standard pruning.
Furthermore, in Appendix A.1 we
prove why models cannot both be compressible and fit random labels.
• In Appendix B, we show how our PAC-Bayes bound benefit from transfer learning.
• In Appendix C, we discuss data-dependent priors and their effect on the subspace dimen-
sion optimization.
• In Appendix D, we detail our experimental setup including models, datasets, and hyperpa-
rameter settings for training and bound computation.
• In Appendix E, we provide a compression perspective to why equivariant models may bemore desirable for generalization.
• In Appendix F, we further discuss how our through our PAC-Bayes compression bounds,we provide evidence that SGD is not necessary for generalization.
• In Appendix G, we ablate the model size and show how it impacts our bounds and compressibility, we identify the best performing size of models for our bounds.• In Appendix H, we present our observations on double descent and their preditability fromour PAC-Bayes bounds.
• In Appendix I, we expand our theoretical discussion and emphasize conceptual differences
between our method and previous ones in the literature.• In Appendix J we provide licensing information on the datasets we use.• Lastly, in Appendix K we discuss stochasticity and double descent through the PAC-Bayeslens.
A Additional Results
In addition to the results reported in Table 2, we report the best bounds for SVHN and ImageNet-1kas well as the corresponding compressed size in Tables 3 and 4.
In Table 3 we show how compressingthe model via intrinsic dimension (ID) yields better results than standard pruning.
In this table,we basically run our method but substitute ID with pruning and then proceed by quantizing theremaining weights and encoding them through arithmetic encoding.
When pruning we used thestandard iterative procedure following Han et al.
[27], for the MNIST model we pruned 98.8%of the weights, for the FMNIST model 97.0% of the weights, for the SVHN model 98.8% of theweights and for both the CIFAR-10 and CIFAR-100 models we pruned 52.1% of the weights andstopped there as the accuracy dropped significantly if we kept pruning.
Error bars on our bounds: We re-run the bounds computation for 10 times and observe that thevalues are consistent.
On average, we obtain ± 0 .
5% variation in our bounds for models trained fromscratch and ± 0 .
1% variation for transfer learning models.
[Table] [Table]
A.1 Models that can fit random labels cannot be compressed
Our ability to construct nonvacuous generalization bounds rests on the ability to construct modelswhich both fit the training data and are highly compressible.
However, when the structure in thedataset has been completely destroyed by shuffling the labels, then we do not find that our modelsare compressible (shown in Figure 2 right).
This is not just an empirical fact, but one that can beproven apriori: models which fit random labels cannot be compressed.
While this result is a trivialconsequence of complexity theory, we present an argument here for illustration.
Almost all random datasets are incompressible
When sampling labels uniformly at random, almost all datasets are not substantially compressible.
Given a dataset D = { x i , y i } ni =1 (where we are only considering the labels y i , and conditioningon the inputs x i ), and denoting |D| as the length of the string of labels, the probability that a givendataset can be compressed to size |D|− c is less than 2 − c +1 .
To see this, one must consider that thereare only (cid:80) |D|− c
i = 2 i ≤ 2 |D|− c +1 programs of length ≤ |D| − c (fewer still when restricting to selfdelimiting programs), and there are 2 |D| possible datasets.
Therefore averaging over all randomlylabeled datasets the fraction which are compressible to less than or equal to |D| − c bits is at most2 |D|− c +1 / 2 |D| = 2 − c +1 .
A compressible model which fits the data is a compression of the dataset
Let prior P that includes a specification of the model architecture, and the model h which outputsprobabilities for each of the outcomes: p ( y = k
| x i )
=
h ( x i ) k .
We can decompose the (prefix)
Kolmogorov complexity of the dataset (given the prior) as
[Equation]
The term K ( D | h, P ) can be interpreted as a model fit term and upper bounded by the totalnegative log likelihood simply using the model probabilities as a distribution to encode the labels:
K ( D | h
P ) ≤ − (cid:80) i log 2 h ( x i )
y
+ 1 = NLL
( D | h )
+ 1 .
Using the fact that almost all random datasets are incompressible, and choosing c = 1 + log 2 (1 /δ ) ,we have that with probability at least 1 − δ over all randomly sampled datasets
K ( D | P ) > |D| − log 2 (1 /δ )
− 1 .
Plugging into Eq.
(5) and rearranging, we have with probability
[Equation] [Equation]
In Figure 5 we plot the quantity K ( h | P ) + NLL
( D | h ) which represents the compressed sizeof the dataset achieved by our model (related to the minimum description length principle).
We seethat the value is considerably lower than the size of the dataset |D| , emphasizing that real machinelearning datasets such as CIFAR-10 have a very low Kolmogorov complexity and are very unlikethose with random labels.
B Transfer Learning Bounds
We show the expanded results both with and without transfer learning in Table 3.
When finetuningfrom ImageNet we use the larger EfficientNet-B0 models rather than the small convnet.
Despite thefact that the model is significantly larger than the convnet or resnet models that we use to achieve thebest bounds for from scratch training, the difference between the finetuned and pretrained models ishighly compressible.
C Data Dependent Priors
We observe that when using data dependent priors, our optimization over the subspace dimension(and the complexity of the model used to fit the data when measured against the prior) favors verylow dimensions and low KL values which we show empirically in Figure 3.
This phenomenon makessense since a large fraction of the data fitting is moved into fitting a good prior, and particularly whenthe dataset fraction used to train the prior is large.
When the prior is already fitted on the data, thefinal solution can have a very low complexity with respect to that prior and is encouraged to do so.
[Figure]
D Experimental Details
In this section we provide experimental details to reproduce our results.
D.1 Model Training Details
We use a standard small convolutional architecture for our experiments, which we find producesbetter bounds than its ResNet counterparts.
The architecture is detailed in Table 5, and we usek = 16 for experiments, but this value is ablated in Figure 5.
Stochastic training: All models were trained for 500 epochs using the Adam optimizer with learning rate 0 .
001 , except for ImageNet which was trained for 80 epochs with SGD using learning rateof 0 .
05 and weight decay of 0 . 00002 .
The model architectures for each dataset are listed below:
• MNIST [43] ( + SVHN [54] Transfer): LeNet-5 [43].
• FashionMNIST [67] ( + CIFAR-10 [37] Transfer): ResNet20 [28].
• SVHN: ConvNet (Table 5).
• SVHN + ImageNet Transfer: EfficientNet-B0 [61].
• CIFAR-10: ConvNet.
• CIFAR-10 + ImageNet Transfer: EfficientNet-B0 [61].
• CIFAR-100 [37]: ConvNet.
• CIFAR-100 + ImageNet [15] Transfer: EfficientNet-B0.
[Table]
Full-batch training: We train all models for 3000 epochs, use learning rates equal to 0 .
1 (MNIST+ LeNet-5 and CIFAR-10 + ResNet-18) and 0 .
5 (CIFAR-10 + ConvNet), and a cosine learningrate scheduler that we warm-up for 10 epochs.
We also clip the full gradient to have an L 2 -norm ofat most 0 .
before performing parameter updates in each epoch [24].
Transfer Learning All previous training details remain the same, except that θ 0 from Eq.
(3) isinitialized from a pre-trained checkpoint instead of a random initialization.
As typically done inliterature, the final classification layer is replaced with a randomly initialized fully-connected layerto account for the number of classes in the downstream task.
D.2 Bound Hyperparameter Optimization
As explained in Section 4.3, we optimize the bound hyperparameters by considering that the hypothesis of interest h specifies the hyperparameters in addition to the weights.
Therefore, we paybits back for the combination of hyperparemeters that we select.
For example, if we are doing a gridsearch over 2 values of the quantization-aware training learning rate, 2 values of the intrinsic dimensionality values, 2 values of the quantization levels, and use k-means by default, then the number ofbits that we pay is log 2 (2 × 2 × 2)
=
3 bits.
Optimizing PAC-Bayes bounds for data-independent priors: Our PAC-Bayesian subspace compression bounds for data-independent priors have 4 hyperparameters that we list here-under alongside the possible values that we consider for each hyperparameter:
• The learning rate for the quantization-aware training, possible values: { 0 . 001 , 0 . 003 , 0 .
The intrinsic dimensionality, possible values: { 0 , 1000 , 2500 , 3000 , 3500 , 4000 , 5000 , 7500 ,8000 , 10000 , 12000 , 15000 , 20000 , 25000 , 50000 , 100000 , 250000 , 500000 } , except forthe ImageNet transfer learning which was conducted on the more limited range:{ 500 , 1000 , 2000 , 3000 , 4000 , 6000 , 8000 }• The number of quantization levels, possible values: { 0 , 7 , 11 , 30 , 50 } .
• The quantization initialization, possible values: { uniform , k-means } .
[Table]
Note that we only use a subset of these hyperparameter values for some datasets, depending on thedataset size and other considerations.
For all bound computations, we use arithmetic encoding and30 epochs of quantization-aware training.
In Table 6, we summarize the hyperparameters corresponding to the data-independent bounds thatwe report in Table 2.
[Table]
Optimizing PAC-Bayes bounds for data-dependent priors:
In addition to the hyperparameterslisted above, we also tune the hyperparameter corresponding to the subset of the training dataset thatwe use to train the prior on.
We consider the following values for the subset training
In Table 7, we summarize the the hyperparameters corresponding to the data-dependent bounds thatwe report in Table 2.
The best bounds are obtained for intrinsic dimensionality equal to 0 , thereforeno quantization is performed.
[Table]
D.3 Computational Infrastructure & Resources
Our computational hardware involved a mix of NVIDIA GeForce RTX 2080 Ti (12GB), NVIDIA
TITAN RTX (24GB), NVIDIA V100 (32GB), and NVIDIA RTX8000 (48GB).
The experimentswere managed via W&B [5].
The total computational cost of all experiments (including the onesthat do not appear in this work) amounts to ≈ 8000 GPU hours.
E Equivariance
We conduct a simple experiment to evaluate the extent to which model equivariance has on thecompressibility of deep learning models and the tightness of our generalization bounds.
We usethe rotationally equivariant C 8 WideResNet model from Weiler and Cesa [64] which has an 8 -foldrotational symmetry, and we also use a non equivariant version of this model.
The equivariant modelhas a depth of 10 and a widen factor of 4 yielding 1 .
451 M parameters.
We control for the number of parameters by adjusting the widen factor of the non equivariant model to 4 .
67 yielding 1 . 447 Mparameters.
[Figure]
We evaluate these models both on MNIST and the RotMNIST dataset [41] consisting of 12K trainingexamples of rotated MNIST digits.
As shown in Figure 4 (a), when paired with the rotationallysymmetric RotMNIST dataset, the rotationally equivariant model achieves better bounds and is morecompressible than it’s non equivariant counterpart despite having the same number of parameters.
However, when this symmetry of the dataset is removed by considering standard MNIST, we seethat the benefits of equivariance to the generalization bound and compressibility vs the WRN modeldissapear.
F Full-Batch vs. Stochastic Training (SGD)
To further expand on the results that we present in Section 6, we study the impact of hyperparameters, namely the weight decay and the architecture choice, on the bounds obtained throughfull-batch (F-B) training.
Table 8 summarizes these results and we provide the training and boundcomputation details in Appendix D. Our PAC-Bayes subspace compression bounds provide similartheoretical guarantees for both full-batch and stochastic training, suggesting that the implicit biasesof SGD are not necessary to guarantee good generalization.
Moreover, we see that the results areconsistent for different configurations, which result in comparable bounds overall.
Transfer learning using full-batch training: In addition to results reported in Section 6, we perform full batch training for transfer learning from SVHN to MNIST using LeNet-5 and the sameexperimental setup described in Appendix D. Our best PAC-Bayes subspace compression boundsfor SVHN to MNIST transfer are 8 .
7% and 9 . 0% for full-batch and SGD training, respectively.
This finding provides further evidence that good generalization of neural networks, and the successof transfer learning in particular, does not necessarily require stochasticity or additional flatnessinducing procedures to be achieved.
Finally, we note that we optimize over the same set of hyperparameters for the bound computationfor both full-batch and stochastic training.
G Model Size vs. Compressibility
We perform an ablation to determine how the size of the model affects our generalization bounds.
Using the fixed model architecture Table 5, we vary the width k from 4 to 192 .
Using our subspacecompression scheme, we find that the compression ratio of the model does increase with model size,however the total compressed size still increases slowly making our bounds less strong for larger
[Table]
models.
For this paper, we find the sweet spot k
= 16 is just above the point with equal number ofparameters and data points.
We note that this finding leaves room for an improved compression scheme and generalizationbounds which are able to explain why even larger models still generalize better.
Curiously, whenplotting to the total compressed dataset size ( K ( h | P ) + NLL ) using the model as a compressionscheme, we find that the MDL principle which favors shorter description lengths of the data actuallyprefers larger models than our PAC-Bayes generalization bounds selects.
[Figure]
a compression scheme.
While the raw labels have size 20 . 3 KB (shown by the black line), the best model compresses the labels down to 8 .
6 KB.
Curiously, the compressed dataset size and hence the MDL principle favors larger models than our generalization bounds.
H Double Descent
Under select conditions, we are able to reproduce the double descent phenomenon in our generalization bounds.
In Figure 6 (right), we show that our bound exhibits a double descent that is similarto what we see in terms of the test error Figure 6 (left).
The results we show in Figure 6 (right) areobtained for a fixed intrinsic dimensionality that is equal to 35000 , but we observed that this middledescent consistently appears in our bounds plots for a given (fixed) intrinsic dimensionality wherewe select the best bound for each base width.
However, we expect that extending the plot out tolarger model widths the bound gets worse again as explained in Appendix G.
 
[Figure]
     7 H V W ( UU R U   /D V W ( SR F K  /DEHO1RLVH  0RGH 6FUDWFK 7UDQVIHU 
[Figure]
  
 
 ( UU  % RXQG 

I PAC-Bayes Bounds
[Figure]
I.1 Catoni PAC-Bayes Bound
In our case, since neural networks achieve low training error, we focus on a bound like Catoni [7]which becomes tighter when KL ( Q, P ) is large.
This is the same bound used in Zhou et al.
Theorem 1 (Catoni [7]) .
Given a 0 - 1 loss ℓ , a fixed α > 1 and a confidence level δ ∈ (0 , 1) then
[Equation]
holds with probability higher than 1 − δ and where
[Equation]
I.2 Variable Length Encoding and Robustness Adjustment
In Zhou et al.
[71], the authors assume a fixed length encoding for the weights.
Given that thedistribution over quantization levels is highly nonuniform, using a variable length encoding (suchas Huffman encoding or arithmetic encoding) can represent the same information using fewer bits.
While this choice gives significant benefits, it means that we cannot immediately make use of robustness adjustment from Zhou et al.
[71], where the robustness adjustment comes from consideringneighboring models that result from perturbing the weights slightly.
Revisiting the prior derivation in Zhou et al.
[71], we show why the method used for bounding the
KL does not transfer over to variable length encodings.
In Zhou et al.
[71], the prior used is
[Equation]
where S denotes the encoding of the position of the pruned weights, C denotes the codebook, Q thecodebook value that the weight take, d the number of nonzero weights, L the number of clusters, ˆ wthe quantized weight and τ 2 the prior variance.
Note that ˆ w changes depending on S, Q, C and alsonote that the fixed-length encoding can be seen in how we sum over d ⌈ log L ⌉ options.
This prior isa mixture of Gaussians centered at the quantized values.
With this choice of prior Zhou et al.
[71]and setting the posterior to be also Gaussian centered at a quantized value, one can upper bound the
KL with a computationally tractable term involving the sum over dimensions.
Crucially, for their decomposition they use the fact that the size of the encoding | Q | is d ⌈ log L ⌉ , which is independentof the coding Q and only dependent on the codebook C .
Therefore they are able to upper bound
the
KL.
KL
[Equation]
due to the independence of the fixed length encoding to that of the values that each quantized valuetakes, see appendix in Zhou et al.
This independence is broken for variable length encoding asthe cluster centers and the values that each weight can take are interlinked.
Thus, we cannot expressand satisfactorily approximate the first high-dimensional KL term as a sum of one-dimensionalelements that can be estimated through quadrature or Monte Carlo.
J Licences
MNIST 3 [43] is made available under the terms of the Creative Commons Attribution-Share Alike3.0 license.
FashionMNIST 4 [67] is available under the MIT license.
CIFAR-10 and CIFAR-100[37] are made available under the MIT license.
ImageNet 5 [15] is the copyright of Stanford Vision
Lab, Stanford University, Princeton University.
SVHN 6 [54] is released for non-commericial useonly.
K Stochasticity and Double Descent through the PAC-Bayes Lens
Is Stochasticity Necessary for Generalization?
It is widely hypothesized that the implicit biasesof SGD help find solutions which generalize better.
For example, Arora et al.
[1] argue that thereis no regularizing function with a gradient that replicates the benefits of gradient noise.
Wu et al.[66], Smith et al.
[59], and Li et al.
[46] advocate that gradient noise is necessary to achieve stateof-the-art performance.
In contrast, recent work by Geiping et al.
[24] shows that full-batch gradientdescent can match state-of-the-art performance.
We train ResNet-18 and LeNet5 models on CIFAR-10 and MNIST, respectively, using full-batchand stochastic gradient descent with different values of the intrinsic dimensionality.
We provide thetraining details in Appendix F, but it is worth noting that, unlike Geiping et al.
[24], we do not useflatness inducing regularization for full batch training.
For MNIST with LeNet5, the best generalization bounds that we obtain are 11 .
55% and 11 .
20% using stochastic gradient descent (SGD) andfull batch training respectively.
The best generalization bounds that we obtain for CIFAR-10 with
ResNet-18 are 74 .
68% and 76 .
25% using SGD and full batch training, respectively.
These closetheoretical guarantees on the generalization error for both SGD and full batch training suggest thatthe implicit biases of SGD may be helpful, but are not at all crucial to understanding why neuralnetworks generalize well.
We expand this analysis in Appendix F.
Double Descent [52]: In Appendix H, we show that our bounds are tight and fine-grained enoughto predict the double descent phenomenon, and we discuss the significance of this result.