Generalization bounds via distillation

1 OVERVIEW AND MAIN RESULTS . Generalization bounds are statistical tools which take as input various measurements of a predictor on training data , and output a performance estimate for unseen data — that is , they estimate how well the predictor generalizes to unseen data . Despite extensive development spanning many decades ( Anthony & Bartlett , 1999 ) , there is growing concern that these bounds are not only disastrously loose ( Dziugaite & Roy , 2017 ) , but worse that they do not correlate with the underlying phenomena ( Jiang et al. , 2019b ) , and even that the basic method of proof is doomed ( Zhang et al. , 2016 ; Nagarajan & Kolter , 2019 ) . As an explicit demonstration of the looseness of these bounds , Figure 1 calculates bounds for a standard ResNet architecture achieving test errors of respectively 0.008 and 0.067 on mnist and cifar10 ; the observed generalization gap is 10−1 , while standard generalization techniques upper bound it with 1015 . Contrary to this dilemma , there is evidence that these networks can often be compressed or distilled into simpler networks , while still preserving their output values and low test error . Meanwhile , these simpler networks exhibit vastly better generalization bounds : again referring to Figure 1 , those same networks from before can be distilled with hardly any change to their outputs , while their bounds reduce by a factor of roughly 1010 . Distillation is widely studied ( Buciluŭ et al. , 2006 ; Hinton et al. , 2015 ) , but usually the original network is discarded and only the final distilled network is preserved . The purpose of this work is to carry the good generalization bounds of the distilled network back to the original network ; in a sense , the explicit simplicity of the distilled network is used as a witness to implicit simplicity of the original network . The main contributions are as follows . • The main theoretical contribution is a generalization bound for the original , undistilled network which scales primarily with the generalization properties of its distillation , assuming that wellbehaved data augmentation is used to measure the distillation distance . An abstract version of this bound is stated in Lemma 1.1 , along with a sufficient data augmentation technique in Lemma 1.2 . A concrete version of the bound , suitable to handle the ResNet architecture in Figure 1 , is described in Theorem 1.3 . Handling sophisticated architectures with only minor proof alterations is another contribution of this work , and is described alongside Theorem 1.3 . This abstract and concrete analysis is sketched in Section 3 , with full proofs deferred to appendices . • Rather than using an assumption on the distillation process ( e.g. , the aforementioned “ wellbehaved data augmentation ” ) , this work also gives a direct uniform convergence analysis , culminating in Theorem 1.4 . This is presented partially as an open problem or cautionary tale , as its proof is vastly more sophisticated than that of Theorem 1.3 , but ultimately results in a much looser analysis . This analysis is sketched in Section 3 , with full proofs deferred to appendices . • While this work is primarily theoretical , it is motivated by Figure 1 and related experiments : Figures 2 to 4 demonstrate that not only does distillation improve generalization upper bounds , but moreover it makes them sufficiently tight to capture intrinsic properties of the predictors , for example removing the usual bad dependence on width in these bounds ( cf . Figure 3 ) . These experiments are detailed in Section 2 . 1.1 AN ABSTRACT BOUND VIA DATA AUGMENTATION . This subsection describes the basic distillation setup and the core abstract bound based on data augmentation , culminating in Lemmas 1.1 and 1.2 ; a concrete bound follows in Section 1.2 . Given a multi-class predictor f : Rd → Rk , distillation finds another predictor g : Rd → Rk which is simpler , but close in distillation distance Φγ , m , meaning the softmax outputs φγ are close on average over a set of points ( zi ) mi=1 : Φγ , m ( f , g ) : = 1 m m∑ i=1 ∥∥φγ ( f ( zi ) ) − φγ ( g ( zi ) ) ∥∥1 , where φγ ( f ( z ) ) ∝ exp ( f ( z ) /γ ) . ( 1.1 ) The quantity γ > 0 is sometimes called a temperature ( Hinton et al. , 2015 ) . Decreasing γ increases sensitivity near the decision boundary ; in this way , it is naturally related to the concept of margins in generalization theory , as detailed in Appendix B. due to these connections , the use of softmax is beneficial in this work , though not completely standard in the literature ( Buciluŭ et al. , 2006 ) . We can now outline Figure 1 and the associated empirical phenomenon which motivates this work . ( Please see Section 2 for further details on these experiments . ) Consider a predictor f which has good test error but bad generalization bounds ; by treating the distillation distance Φγ , m ( f , g ) as an objective function and increasingly regularizing g , we obtain a sequence of predictors ( g0 , . . . , gt ) , where g0 = f , which trade off between distillation distance and predictor complexity . The curves in Figure 1 are produced in exactly this way , and demonstrate that there are predictors nearly identical to the original f which have vastly smaller generalization bounds . Our goal here is to show that this is enough to imply that f in turn must also have good generalization bounds , despite its apparent complexity . To sketch the idea , by a bit of algebra ( cf . Lemma A.2 ) , we can upper bound error probabilities with expected distillation distances and errors : Prx , y [ arg max y′ f ( x ) y′ 6= y ] ≤ 2Ex ∥∥φγ ( f ( x ) ) − φγ ( g ( x ) ) ∥∥1 + 2Ex , y ( 1− φγ ( g ( x ) ) y ) . The next step is to convert these expected errors into quantities over the training set . The last term is already in a form we want : it depends only on g , so we can apply uniform convergence with the low complexity of g. ( Measured over the training set , this term is the distillation error in Figure 1 . ) The expected distillation distance term is problematic , however . Here are two approaches . 1 . We can directly apply uniform convergence ; for instance , this approach was followed by Suzuki et al . ( 2019 ) , and a more direct approach is followed here to prove Theorem 1.4 . Unfortunately , it is unclear how this technique can avoid paying significantly for the high complexity of f . 2 . The idea in this subsection is to somehow trade off computation for the high statistical cost of the complexity of f . Specifically , notice that Φγ , m ( f , g ) only relies upon the marginal distribution of the inputs x , and not their labels . This subsection will pay computation to estimate Φγ , m with extra samples via data augmentation , offsetting the high complexity of f . We can now set up and state our main distillation bound . Suppose we have a training set ( ( xi , yi ) ) ni=1 drawn from some measure µ , with marginal distribution µX on the inputs x . Suppose we also have ( zi ) m i=1 drawn from a data augmentation distribution νn , the subscript referring to the fact that it depends on ( xi ) ni=1 . Our analysis works when ‖dµX/dνn‖∞ , the ratio between the two densities , is finite . If it is large , then one can tighten the bound by sampling more from νn , which is a computational burden ; explicit bounds on this term will be given shortly in Lemma 1.2 . Lemma 1.1 . Let temperature parameter γ > 0 be given , along with sets of multiclass predictors F and G. Then with probability at least 1 − 2δ over an iid draw of data ( ( xi , yi ) ) ni=1 from µ and ( zi ) m i=1 from νn , every f ∈ F and g ∈ G satisfy Pr [ arg max y′ f ( x ) y′ 6= y ] ≤ 2 ∥∥∥∥dµXdνn ∥∥∥∥ ∞ Φγ , m ( f , g ) + 2 n n∑ i=1 ( 1− φγ ( g ( xi ) ) yi ) + Õ ( k3/2 γ ∥∥∥∥dµXdνn ∥∥∥∥ ∞ ( Radm ( F ) + Radm ( G ) ) + √ k γ Radn ( G ) ) + 6 √ ln ( 1/δ ) 2n ( 1 + ∥∥∥∥dµXdνn ∥∥∥∥ ∞ √ n m ) , where Rademacher complexities Radn and Radm are defined in Section 1.4 . A key point is that the Rademacher complexity Radm ( F ) of the complicated functions F has a subscript “ m ” , which explicitly introduces a factor 1/m in the complexity definition ( cf . Section 1.4 ) . As such , sampling more from the data augmentation measure can mitigate this term , and leave the complexity of the distillation class G as the dominant term . Of course , this also requires ‖dµX/dνn‖∞ to be reasonable . As follows is one data augmentation scheme ( and assumption on marginal distribution µX ) which ensures this . Lemma 1.2 . Let ( xi ) ni=1 be a data sample drawn iid from µX , and suppose the corresponding density p is supported on [ 0 , 1 ] d and is Hölder continuous , meaning |p ( x ) − p ( x′ ) | ≤ Cα‖x− x′‖α for some Cα ≥ 0 , α ∈ [ 0 , 1 ] . Define a data augmentation measure νn via the following sampling procedure . • With probability 1/2 , sample z uniformly within [ 0 , 1 ] d. • Otherwise , select a data index i ∈ [ n ] uniformly , and sample z from a Gaussian centered at xi , and having covariance σ2I where σ : = n−1/ ( 2α+d ) . Then with probability at least 1− 1/n over the draw of ( xi ) ni=1 , ∥∥∥∥dµXdνn ∥∥∥∥ ∞ = 4 +O ( √ lnn nα/ ( 2α+d ) ) . Though the idea is not pursued here , there are other ways to control ‖dµX/dνn‖∞ , for instance via an independent sample of unlabeled data ; Lemma 1.1 is agnostic to these choices . 1.2 A CONCRETE BOUND FOR COMPUTATION GRAPHS . This subsection gives an explicit complexity bound which starts from Lemma 1.1 , but bounds ‖dµX/dνn‖∞ via Lemma 1.2 , and also includes an upper bound on Rademacher complexity which can handle the ResNet , as in Figure 1 . A side contribution of this work is the formalism to easily handle these architectures , detailed as follows . Canonical computation graphs are a way to write down feedforward networks which include dense linear layers , convolutional layers , skip connections , and multivariate gates , to name a few , all while allowing the analysis to look roughly like a regular dense network . The construction applies directly to batches : given an input batch X ∈ Rn×d , the output Xi of layer i is defined inductively as XT0 : = X T , XTi : = σi ( [ WiΠiDi|〉Fi ] XTi−1 ) = σi ( [ WiΠiDiX T i−1 FiX T i−1 ] ) , where : σi is a multivariate-to-multivariate ρi-Lipschitz function ( measured over minibatches on either side with Frobenius norm ) ; Fi is a fixed matrix , for instance an identity mapping as in a residual network ’ s skip connection ; Di is a fixed diagonal matrix selecting certain coordinates , for instance the non-skip part in a residual network ; Πi is a Frobenius norm projection of a full minibatch ; Wi is a weight matrix , the trainable parameters ; [ WiΠiDi|〉Fi ] denotes row-wise concatenation of WiΠiDi and Fi . As a simple example of this architecture , a multi-layer skip connection can be modeled by including identity mappings in all relevant fixed matrices Fi , and also including identity mappings in the corresponding coordinates of the multivariate gates σi . As a second example , note how to model convolution layers : each layer outputs a matrix whose rows correspond to examples , but nothing prevents the batch size from changes between layers ; in particular , the multivariate activation before a convolution layer can reshape its output to have each row correspond to a patch of an input image , whereby the convolution filter is now a regular dense weight matrix . A fixed computation graph architecture G ( ~ρ , ~b , ~r , ~s ) has associated hyperparameters ( ~ρ , ~b , ~r , ~s ) , described as follows . ~ρ is the set of Lipschitz constants for each ( multivariate ) gate , as described before . ri is a norm bound ‖W Ti ‖2,1 ≤ ri ( sum of the ‖ · ‖2-norms of the rows ) , bi √ n ( where n is the input batch size ) is the radius of the Frobenius norm ball which Πi is projecting onto , and si is the operator norm of X 7→ [ WiΠiDiXT|〉FiXT ] . While the definition is intricate , it can not only model basic residual networks , but it is sensitive enough to be able to have si = 1 and ri = 0 when residual blocks are fully zeroed out , an effect which indeed occurs during distillation . Theorem 1.3 . Let temperature parameter γ > 0 be given , along with multiclass predictors F , and a computation graph architecture G. Then with probability at least 1− 2δ over an iid draw of data ( ( xi , yi ) ) n i=1 from µ and ( zi ) n i=1 from νn , every f ∈ F satisfies Pr [ arg max y′ f ( x ) y′ 6= y ] ≤ inf ( ~b , ~r , ~s ) ≥1 g∈G ( ~ρ , ~b , ~r , ~s ) 2 [ ∥∥∥∥dµXdνn ∥∥∥∥ ∞ Φγ , m ( f , g ) + 2 n n∑ i=1 ( 1− φγ ( g ( xi ) ) yi + Õ ( k3/2 γ ∥∥∥∥dµXdνn ∥∥∥∥ ∞ Radm ( F ) ) + 6 √ ln ( 1/δ ) 2n ( 1 + ∥∥∥∥dµXdνn ∥∥∥∥ ∞ √ n m ) + Õ ( √ k γ √ n ( 1 + k ∥∥∥∥dµXdνn ∥∥∥∥ ∞ √ n m ) ( ∑ i [ ribiρi L∏ l=i+1 slρl ] 2/3 ) 3/2 ) ] . Under the conditions of Lemma 1.2 , ignoring an additional failure probability 1/n , then ‖dµXdνn ‖∞ = 4 +O ( √ lnn nα/ ( 2α+d ) ) . A proof sketch of this bound appears in Section 3 , with full details deferred to appendices . The proof is a simplification of the covering number argument from ( Bartlett et al. , 2017a ) ; for another computation graph formalism designed to work with the covering number arguments from ( Bartlett et al. , 2017a ) , see the generalization bounds due to Wei & Ma ( 2019 ) .

Summary: The paper overall is of good quality. The story of the work is well-written which makes the contributions easier to digest. One suggestion would be to comment a bit more on the relevance of the margin distribution for readers that are unfamiliar with it, for instance, in Figure 1, the term margin distribution is thrown without explaining why one should look into it.