GENERALIZATION BOUNDS VIA DISTILLATION
Year: 2021

Abstract:
This paper theoretically investigates the following empirical phenomenon: given a high-complexity network with poor generalization bounds, one can distill it into a network with nearly identical predictions but low complexity and vastly smaller generalization bounds. The main contribution is an analysis showing that the original network inherits this good generalization bound from its distillation, assuming the use of well-behaved data augmentation. This bound is presented both in an abstract and in a concrete form, the latter complemented by a reduction technique to handle modern computation graphs featuring convolutional layers, fullyconnected layers, and skip connections, to name a few. To round out the story, a (looser) classical uniform convergence analysis of compression is also presented, as well as a variety of experiments on cifar10 and mnist demonstrating similar generalization performance between the original network and its 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 10. 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 10. Distillation is widely studied (Buciluŭ 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

heading: 1 OVERVIEW AND MAIN RESULTS
text: 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.

heading: 1.1 AN ABSTRACT BOUND VIA DATA AUGMENTATION
text: 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.

heading: 1.2 A CONCRETE BOUND FOR COMPUTATION GRAPHS
text: 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 cannot 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).

heading: 1.3 A UNIFORM-CONVERGENCE APPROACH TO DISTILLATION
text: In this section, we derive a Rademacher complexity bound on F whose proof internally uses compression; specifically, it first replaces f with a narrower network g, and then uses a covering number bound sensitive to network size to control g. The proof analytically chooses g’s width based on the structure of f and also the provided data, and this data dependence incurs a factor which causes the familiar 1/√n rate to worsen to 1/n1/4 (which appears as ‖X‖F/n3/4). This proof is much more intricate than the proofs coming before, and cannot handle general computation graphs, and also ignores the beneficial structure of the softmax. Theorem 1.4. Let data matrix X ∈ Rn×d be given, and let F denote networks of the form x 7→ σL(WL · · ·σ1(W1x)) with spectral norm ‖Wi‖2 ≤ si, and 1-Lipschitz and 1-homogeneous activations σi, and ‖Wi‖F ≤ Ri and width at most m. Then
Rad(F) = Õ ( ‖X‖F n3/4 [∏ j sj ] [∑ i (Ri/si) 4/5 ]5/4 [∑ i lnRi ]1/4) .
The term Ri/si is the square root of the stable rank of weight matrixWi, and is a desirable quantity in a generalization bound: it scales more mildly with width than terms like ‖W Ti ‖2,1 and ‖W Ti ‖F √ width which often appear (the former appears in Theorem 1.3 and Lemma 3.1). Another stable rank bound was developed by Suzuki et al. (2019), but has an extra mild dependence on width.
As depicted in Figure 2, however, this bound is not fully width-independent. Moreover, we can compare it to Lemma 3.1 throughout distillation, and not only does this bound not capture the power of distillation, but also, eventually its bad dependence on n causes it to lose out to Lemma 3.1.

heading: 1.4 ADDITIONAL NOTATION
text: Given data (zi)ni=1, the Rademacher complexity of univariate functionsH is
Rad(H) := E~ sup h∈H
1
n ∑ i ih(zi), where i i.i.d.∼ Uniform({−1,+1}).
Rademacher complexity is the most common tool in generalization theory (Shalev-Shwartz & BenDavid, 2014), and is incorporated in Lemma 1.1 due to its convenience and wide use. To handle multivariate (multiclass) outputs, the definition is overloaded via the worst case labels as Radn(F) = sup~y∈[k]n Rad({(x, y) 7→ f(x)y : f ∈ F}). This definition is for mathematical convenience, but overall not ideal; Rademacher complexity seems to have difficulty dealing with such geometries.
Regarding norms, ‖ · ‖ = ‖ · ‖F will denote the Frobenius norm, and ‖ · ‖2 will denote spectral norm.

heading: 2 ILLUSTRATIVE EMPIRICAL RESULTS
text: This section describes the experimental setup, and the main experiments: Figure 1 showing progressive distillation, Figure 2 comparing Theorem 1.4, Lemma 3.1 and VC dimension, Figure 3 showing width independence after distillation, and Figure 4 showing the effect of random labels.
Experimental setup. As sketched before, networks were trained in a standard way on either cifar10 or mnist, and then distilled by trading off between complexity and distillation distance Φγ,m. Details are as follows.
1. Training initial network f . In Figures 1 and 2a, the architecture was a ResNet8 based on one used in (Coleman et al., 2017), and achieved test errors 0.067 and 0.008 on cifar10 and mnist, respectively, with no changes to the setup and a modest amount of training; the training algorithm was Adam; this and most other choices followed the scheme in (Coleman et al., 2017) to achieve a competitively low test error on cifar10. In Figures 2b, 3 and 4, a 6-layer fully connected network was used (width 8192 in Figure 2b, widths {64, 256, 1024} in Figure 3, width 256 in Figure 4), and vanilla SGD was used to optimize.
2. Training distillation network g. Given f and a regularization strength λj , each distillation gj was found via approximate minimization of the objective
g 7→ Φγ,m(f, g) + λjComplexity(g). (2.1)
In more detail, first g0 was initialized to f (g and f always used the same architecture) and optimized via eq. (2.1) with λ0 set to roughly risk(f)/Complexity(f), and thereafter gj+1 was initialized to gj and found by optimizing eq. (2.1) with λj+1 := 2λj . The optimization method was the same as the one used to find f . The term Complexity(g) was some computationally reasonable approximation of Lemma 3.1: for Figures 2b, 3 and 4, it was just ∑ i ‖W Ti ‖2,1, but for Figures 1 and 2a, it also included a tractable surrogate for the product of the spectral norms, which greatly helped distillation performance with these deeper architectures. In Figures 2b, 3 and 4, a full regularization sequence was not shown, only a single gj . This was chosen with a simple heuristic: amongst all (gj)j≥1, pick the one whose 10% margin quantile is largest (see the definition and discussion of margins below).
Margin histograms. Figures 2b, 3 and 4 all depict margin histograms, a flexible tool to study the individual predictions of a network on all examples in a training set (see for instance (Schapire & Freund, 2012) for their use studying boosting, and (Bartlett et al., 2017a; Jiang et al., 2019a) for
their use in studying deep networks). Concretely, given a predictor g ∈ G, the prediction on every example is replaced with a real scalar called the normalized margin via
(xi, yi) 7→ g(xi)yi −maxj 6=yi g(xi)j
Radn(G) ,
where Radn(G) is the Rademacher complexity (cf. Section 1.4), and then the histogram of these n scalars is plotted, with the horizontal axis values thus corresponding to normalized margins. By using Rademacher complexity as normalization, these margin distributions can be compared across predictors and even data sets, and give a more fine-grained analysis of the quality of the generalization bound. This normalization choice was first studied in (Bartlett et al., 2017a), where it was also mentioned that this normalization allows one to read off generalization bounds from the plot. Here, it also suggests reasonable values for the softmax temperature γ.
Figure 1: effect of distillation on generalization bounds. This figure was described before; briefly, a highlight is that in the initial phase, training and testing errors hardly change while bounds drop by a factor of nearly 1010. Regarding “generalization measure”, this term appears in studies of quantities which correlate with generalization, but are not necessarily rigorous generalization bounds (Jiang et al., 2019b; Dziugaite et al., 2020); in this specific case, the product of Frobenius norms requires a dense ReLU network (Golowich et al., 2018), and is invalid for the ResNet (e.g., a complicated ResNet with a single identity residual block yields a value 0 by this measure).
Figure 2a: comparison of Theorem 1.4, Lemma 3.1 and VC bounds. Theorem 1.4 was intended to internalize distillation, but as in Figure 2a, clearly a subsequent distillation still greatly reduces the bound. While initially the bound is better than Lemma 3.1 (which does not internalize distillation), eventually the n1/4 factor causes it to lose out. Also note that eventually the bounds beat the VC bound, which has been identified as a surprisingly challenging baseline (Arora et al., 2018).
Figure 3: width independence. Prior work has identified that generalization bounds are quite bad at handling changes in width, even if predictions and test error don’t change much (Nagarajan & Kolter, 2019; Jiang et al., 2019b; Dziugaite et al., 2020). This is captured in Figure 3a, where the margin distributions (see above) with different widths are all very different, despite similar test errors. However, following distillation, the margin histograms in Figure 3b are nearly identical! That is to say: distillation not only decreases loose upper bounds as before, it tightens them to the point where they capture intrinsic properties of the predictors.
Figure 2b: failure of width independence with Theorem 1.4. The bound in Theorem 1.4 was designed to internalize compression, and there was some hope of this due to the stable rank term.
Unfortunately, Figure 2b shows that it doesn’t quite succeed: while the margin histograms are less separated than for the undistilled networks in Figure 3a, they are still visibly separated unlike the post-distillation histograms in Figure 3b.
Figure 4: random labels. A standard sanity check for generalization bounds is whether they can reflect the difficulty of fitting random labels (Zhang et al., 2016). While it has been empirically shown that Rademacher bounds do sharply reflect the presence of random labels (Bartlett et al., 2017a, Figures 2 & 3), the effect is amplified with distillation: even randomizing just 25% shrinks the margin distribution significantly.

heading: 3 ANALYSIS OVERVIEW AND SKETCH OF PROOFS
text: This section sketches all proofs, and provides further context and connections to the literature. Full proof details appear in the appendices.

heading: 3.1 ABSTRACT DATA AUGMENTATION BOUNDS IN SECTION 1.1
text: As mentioned in Section 1.1, the first step of the proof is to apply Lemma A.2 to obtain Prx,y[arg max
y′ f(x)y′ 6= y] ≤ 2Ex ∥∥φγ(f(x))− φγ(g(x))∥∥1 + 2Ex,y (1− φγ(g(x))y) ; this step is similar to how the ramp loss is used with margin-based generalization bounds, a connection which is discussed in Appendix B.
Section 1.1 also mentioned that the last term is easy: φγ is (1/γ)-Lipschitz, and we can peel it off and only pay the Rademacher complexity associated with g ∈ G. With data augmentation, the first term is also easy:
EΦγ,m(f, g) = ∫ ‖φγ(f(z))− φγ(g(z))‖1 dµX (z) = ∫ ‖φγ(f(z))− φγ(g(z))‖1
dµX dνn dνn(z)
≤ ∥∥∥∥dµXdνn ∥∥∥∥ ∞ ∫ ‖φγ(f(z))− φγ(g(z))‖1 dνn(z),
and now we may apply uniform convergence to νn rather than µX . In the appendix, this proof is handled with a bit more generality, allowing arbitrary norms, which may help in certain settings. All together, this leads to a proof of Lemma 1.1.
For the explicit data augmentation estimate in Lemma 1.2, the proof breaks into roughly two cases: low density regions where the uniform sampling gives the bound, and high density regions where the Gaussian sampling gives the bound. In the latter case, the Gaussian sampling in expectation behaves as a kernel density estimate, and the proof invokes a standard bound (Jiang, 2017).

heading: 3.2 CONCRETE DATA AUGMENTATION BOUNDS IN SECTION 1.2
text: The main work in this proof is the following generalization bound for computation graphs, which follows the proof scheme from (Bartlett et al., 2017a), though simplified in various ways, owing mainly to the omission of general matrix norm penalties on weight matrices, and the omission of the reference matrices. The reference matrices were a technique to center the weight norm balls away from the origin; a logical place to center them was at initialization. However, in this distillation setting, it is in fact most natural to center everything at the origin, and apply regularization and shrink to a well-behaved function (rather than shrinking back to the random initialization, which after all defines a complicated function). The proof also features a simplified (2, 1)-norm matrix covering proof (cf. Lemma C.3). Lemma 3.1. Let data X ∈ Rn×d be given. Let computation graph G be given, where Πi projects to Frobenius-norm balls of radius bi √ n, and ‖W Ti ‖2,1 ≤ ri, and ‖[WiΠiDi|〉Fi]‖2 ≤ si, and Lip(σi) ≤ ρi, and all layers have width at most m. Then for every > 0 there exists a covering set M satisfying
sup g∈G min X̂∈M
∥∥∥g(XT)− X̂∥∥∥ ≤ and ln |M| ≤ 24/3n ln(2m2) 2 [∑ i ( ribiρi L∏ l=i+1 slρl )2/3]3 .
Consequently,
Rad(G) ≤ 4 n + 12
√ ln(2m2)
n
[∑ i ( ribiρi L∏ l=i+1 slρl )2/3]3/2 .
From there, the proof of Theorem 1.3 follows via Lemmas 1.1 and 1.2, and many union bounds.

heading: 3.3 DIRECT UNIFORM CONVERGENCE APPROACH IN THEOREM 1.4
text: As mentioned before, the first step of the proof is to sparsify the network, specifically each matrix product. Concretely, given weights Wi of layer i, letting XTi−1 denote the input to this layer, then
WiX T i−1 = m∑ j=1 (Wiej)(Xi−1ej) T.
Written this way, it seems natural that the matrix product should “concentrate”, and that considering all m outer products should not be necessary. Indeed, exactly such an approach has been followed before to analyze randomized matrix multiplication schemes (Sarlos, 2006). As there is no goal of high probability here, the analysis is simpler, and follows from the Maurey lemma (cf. Lemma C.1), as is used in the (2, 1)-norm matrix covering bound in Lemma C.3.
Lemma 3.2. Let a network be given with 1-Lipschitz homogeneous activations σi and weight matrices (W1, . . . ,WL) of maximum width m, along with data matrix X ∈ Rn×d and desired widths (k1, . . . , kL) be given. Then there exists a sparsified network output, recursively defined via
X̂T0 := X T, and X̂Ti := Πiσi(WiMiX T i−1), where Mi := ∑ j∈Si Zjeje T j ‖Aej‖ ,
where Si is a multiset of ki = |Si| indices, Πi denotes projection onto the Frobenius-norm ball of radius ‖X‖F ∏ j≤i ‖Wj‖2, and the scaling term Zj satisfies Zj ≤ ‖Wk‖F √ m/kj , and
‖σL(WL · · ·σ1(W1XT) · · · )− X̂TL‖F ≤ ‖X‖F  L∏ i=1 ‖Wi‖2  L∑ i=1 √ ‖Wi‖2F ki‖Wi‖22 ,
The statement of this lemma is lengthy and detailed because the exact guts of the construction are needed in the subsequent generalization proof. Specifically, now that there are few nodes, a generalization bound sensitive to narrow networks can be applied. On the surface, it seems reasonable to apply a VC bound, but this approach did not yield a rate better than n−1/6, and also had an explicit dependence on the depth of the network, times other terms visible in Theorem 1.4.
Instead, the approach here, aiming for a better dependence on n and also no explicit dependence on network depth, was to produce an∞-norm covering number bound (see (Long & Sedghi, 2019) for a related approach), with some minor adjustments (indeed, the∞-norm parameter covering approach was applied to obtain a Frobenius-norm bound, as in Lemma 3.1). Unfortunately, the magnitudes of weight matrix entries must be controlled for this to work (unlike the VC approach), and this necessitated the detailed form of Lemma 3.2 above.
To close with a few pointers to the literature, as Lemma 3.2 is essentially a pruning bound, it is potentially of independent interest; see for instance the literature on lottery tickets and pruning (Frankle & Carbin, 2019; Frankle et al., 2020; Su et al., 2020). Secondly, there is already one generalization bound in the literature which exhibits spectral norms, due to (Suzuki et al., 2019); unfortunately, it also has an explicit dependence on network width.

heading: ACKNOWLEDGMENTS
text: MT thanks Vaishnavh Nagarajan for helpful discussions and suggestions. ZJ and MT are grateful for support from the NSF under grant IIS-1750051, and from NVIDIA under a GPU grant.

heading: A PROOFS FOR SECTION 1.1
text: The first step is an abstract version of Lemma 1.1 which does not explicitly involve the softmax, just bounded functions.
Lemma A.1. Let classes of bounded functions F and G be given with F 3 f : X → [0, 1]k and G 3 g : X → [0, 1]k. Let conjugate exponents 1/p + 1/q = 1 be given. Then with probability at least 1− 2δ over the draw of ((xi, yi))ni=1 from µ and (zi)mi=1 from νn, for every f ∈ F and g ∈ G,
Ef(x)y ≤ 1
n n∑ i=1 g(xi)yi + 2Radn ({ (x, y) 7→ g(x)y : g ∈ G }) + 3 √ ln(1/δ) 2n
+ ∥∥∥∥dµXdνn ∥∥∥∥ Lq(νn) ( 1 m m∑ i=1 ‖f(zi)− g(zi)‖pp + 3 √ ln(1/δ) 2m
+ 2Radm ({ z 7→ min{1, ‖f(z)− g(z)‖pp} : f ∈ F , g ∈ G }))1/p where
Radm ({ z 7→ min{1, ‖f(z)− g(z)‖pp} : f ∈ F , g ∈ G }) ≤ p
k∑ y′=1 [ Radm({z 7→ f(z)y′ : f ∈ F}) + Radm({z 7→ g(z)y′ : g ∈ G}) ] .
Proof of Lemma A.1. To start, for any f ∈ F and g ∈ G, write
Ef(x)y = E(f(x)− g(x))y + Eg(x)y.
The last term is easiest, and let’s handle it first: by standard Rademacher complexity arguments (Shalev-Shwartz & Ben-David, 2014), with probability at least 1− δ, every g ∈ G satisfies
Eg(x)y ≤ 1
n n∑ i=1 g(xi)yi + 2Radn({(x, y) 7→ g(x)y : g ∈ G}) + 3 √ ln(1/δ) 2n .
For the first term, since f : X → [0, 1]k and g : X → [0, 1]k, by Hölder’s inequality E(f(x)− g(x))y = ∫ min{1, (f(x)− g(x))y}dµ(x, y)
≤ ∫ min{1, ‖f(x)− g(x)‖p} dµ(x, y)
= ∫ min{1, ‖f(x)− g(x)‖p}
dµX dνn (x) dνn(x)
≤ ∥∥∥min{1, ‖f − g‖p}∥∥∥
Lp(νn) ∥∥∥∥dµXdνn ∥∥∥∥ Lq(νn) .
Once again invoking standard Rademacher complexity arguments (Shalev-Shwartz & Ben-David, 2014), with probability at least 1 − δ, every mapping z 7→ min{1, ‖f(z) − g(z)‖pp} where f ∈ F and g ∈ G satisfies∫
min{1, ‖f(z)− g(z)‖pp} dνn(z) ≤ 1
m m∑ i=1 min{1, ‖f(zi)− g(zi)‖pp}+ 3 √ ln(1/δ) 2m
+ 2Radm ({ z 7→ min{1, ‖f(z)− g(z)‖pp} : f ∈ F , g ∈ G }) .
Combining these bounds and unioning the two failure events gives the first bound.
For the final Rademacher complexity estimate, first note r 7→ min{1, r} is 1-Lipschitz and can be peeled off, thus
mRadm ({ z 7→ min{1, ‖f(z)− g(z)‖pp} : f ∈ F , g ∈ G }) ≤ mRadm ({ z 7→ ‖f(z)− g(z)‖pp : f ∈ F , g ∈ G
}) = E sup
f∈F g∈G
m∑ i=1 i‖f(zi)− g(zi)‖pp
≤ k∑
y′=1 E sup f∈F g∈G m∑ i=1 i|f(zi)− g(zi)|py′
= k∑ y′=1 mRadm ({ z 7→ |f(z)− g(z)|py′ : f ∈ F , g ∈ G }) .
Since f and g have range [0, 1]k, then (f − g)y′ has range [−1, 1] for every y′, and since r 7→ |r|p is p-Lipschitz over [−1, 1] (for any p ∈ [1,∞), combining this with the Lipschitz composition rule for Rademacher complexity and also the fact that a Rademacher random vector ∈ {±1}m is distributionally equivalent to its coordinate-wise negation − , then, for every y′ ∈ [k],
Radm({z 7→ |f(z)− g(z)|py′ : f ∈ F , g ∈ G}) ≤ pRadm({z 7→ (f(z)− g(z))y′ : f ∈ F , g ∈ G})
= p
m E sup f∈F sup g∈G m∑ i=1 i(f(zi)− g(zi))y′
= p
m E sup
f∈F m∑ i=1 if(zi)y′ + p m E sup g∈G m∑ i=1 − ig(zi)y′
= pRadm({z 7→ f(z)y′ : f ∈ F}) + pRadm({z 7→ g(z)y′ : g ∈ G}).
To prove Lemma 1.1, it still remains to collect a few convenient properties of the softmax. Lemma A.2. For any v ∈ Rk and y ∈ {1, . . . , k},
2(1− φγ(v))y ≥ 1[y 6= arg max i vi].
Moreover, for any functions F with F 3 f : X → Rk,
Radn ({ (x, y) 7→ φγ(f(x))y : f ∈ F }) = Õ
(√ k
γ Radn(F)
) .
Proof. For the first property, let v ∈ Rk be given, and consider two cases. If y = arg maxi vi, then φγ(v) ∈ [0, 1]k implies 2 ( 1− φγ(v) ) y ≥ 0 = 1[y 6= arg max
i vi].
On the other hand, if y 6= arg maxi vi, then φγ(v)y ≤ 1/2, and 2 ( 1− φγ(v) ) y ≥ 1 = 1[y 6= arg max
i vi].
The second part follows from a multivariate Lipschitz composition lemma for Rademacher complexity due to (Foster & Rakhlin, 2019, Theorem 1); all that remains to prove is that v 7→ φγ(v)y is (1/γ)-Lipschitz with respect to the `∞ norm for any v ∈ Rk and y ∈ [k]. To this end, note that
d
dvy φγ(v)y =
exp(v/γ)y ∑ j 6=y exp(v/γ)j
γ( ∑ j exp(v/γ)j) 2 ,
d
dvi6=y φγ(v)y = − exp(v/γ)y exp(v/γ)i γ( ∑ j exp(v/γ)j) 2 ,
and therefore ∥∥∇φγ(v)y∥∥1 = 2 exp(v/γ)y ∑ j 6=y exp(v/γ)j
γ( ∑ j exp(v/γ)j) 2 ≤ 1 γ ,
and thus, by the mean value theorem, for any u ∈ Rk and v ∈ Rk, there exists z ∈ [u, v] such that∣∣φγ(v)y − φγ(u)y∣∣ = ∣∣∣〈∇φγ(z)y, v − u〉∣∣∣ ≤ ‖v − u‖∞ ·∥∥∇φγ(v)y∥∥1 ≤ 1γ ‖v − u‖∞, and in particular v 7→ φγ(v)/y is (1/γ)-Lipschitz with respect to the `∞ norm. Applying the aforementioned Lipschitz composition rule (Foster & Rakhlin, 2019, Theorem 1),
Radn ({ (x, y) 7→ φγ(f(x))y : f ∈ F }) = Õ
(√ k
γ Radn(F)
) .
Lemma 1.1 now follows by combining Lemmas A.1 and A.2.
Proof of Lemma 1.1. Define ψ := 1 − φγ . The bound follows by instantiating Lemma A.1 with p = 1 and the two function classes
QF := {(x, y) 7→ ψ(f(x)y) : f ∈ F} and QG := {(x, y) 7→ ψ(g(x)y) : g ∈ G},
combining its simplified Rademacher upper bounds with the estimates for Radm(QF ) and Radm(QG) and Radn(QG) from Lemma A.2, and by using Lemma A.2 to lower bound the left hand side with
Eψ(f(x))y = E(1− φγ(f(x))y) ≥ 1 2 1 [ arg max y′ f(x)y′ 6= y ] ,
and lastly noting that
1
m m∑ i=1 ‖ψ(f(zi))− ψ(g(zi))‖1 = 1 m m∑ i=1 ‖1− φγ(f(zi))− 1 + φγ(g(zi))‖1 = Φγ,m(f, g).
To complete the proofs for Section 1.1, it remains to handle the data augmentation error, namely the term ‖dµX/dνn‖∞. This proof uses the following result about Gaussian kernel density estimation. Lemma A.3 (See (Jiang, 2017, Theorem 2 and Remark 8)). Suppose density p is α-Hölder continuous, meaning |p(x) − p(x′)| ≤ Cα‖x − x′‖α for some Cα ≥ 0 and α ∈ [0, 1]. There there exists a constant C ≥ 0, depending on α, Cα, maxx∈Rd p(x), and the dimension, but independent of the sample size, so that with probability at least 1 − 1/n, the Gaussian kernel density estimate with bandwidth σ2I where σ = n−1/(2α+d) satisfies
sup x∈Rd
|p(x)− pn(x)| ≤ C √ ln(n)
n2α/(2α+d) .
The proof of Lemma 1.2 follows.
Proof of Lemma 1.2. The proposed data augmentation measure νn has a density pn,β over [0, 1]d, and it has the form pn,β(x) = β + (1− β)pn(x), where β = 1/2, and pn is the kernel density estimator as described in Lemma A.3, whereby
|pn(x)− p(x)| ≤ n := O
( √ lnn
nα/(2α+d)
) .
The proof proceeds to bound ‖dµX/dνn‖∞ = ‖p/pn,β‖∞ by considering three cases.
• If x 6∈ [0, 1]d, then p(x) = 0 by the assumption on the support of µX , whereas pn,β(x) ≥ pn(x)/2 > 0, thus p(x)/pn,β(x) = 0.
• If x ∈ [0, 1]d and p(x) ≥ 2 n, then pn,β(x) ≥ (1− β)p(x)− n) ≥ n/2, and p(x)
pn,β(x) = 1 + p(x)− pn,β(x) pn,β(x)
≤ 1 + βp(x) pn,β(x) + (1− β)|p(x)− pn(x)| pn,β(x)
≤ 1 + βp(x) (1− β)(p(x)− n) + (1− β) n n/2
≤ 1 + β (1− β)(1− n/p(x)) + 1
≤ 4.
• If x ∈ [0, 1]d and p(x) < 2 n, since pn,β(x) ≥ β = 1/2, then p(x)
pn,β(x) < 2 n β = 4 n.
Combining these cases, ‖dµX/dνn‖∞ = ‖p/pn,β‖∞ ≤ max{4, 4 n} ≤ 4 + 4 n.

heading: B REPLACING SOFTMAX WITH STANDARD MARGIN (RAMP) LOSS
text: The proof of Lemma 1.1 was mostly a reduction to Lemma A.1, which mainly needs bounded functions; for the Rademacher complexity estimates, the Lipschitz property of φγ was used. As such, the softmax can be replaced with the (1/γ)-Lipschitz ramp loss as is standard from margin-based generalization theory (e.g., in a multiclass version as appears in (Bartlett et al., 2017a)). Specifically, defineMγ : Rk → [0, 1]k for any coordinate j as
Mγ(v)j := `γ(vj − arg max y′ 6=j vy′), where `γ(z) :=  1 z ≤ 0, 1− zγ z ∈ (0, γ), 0 z ≥ γ.
We now have 1[arg maxy′ f(x)y′ ] ≤ Mγ(f(x))y without a factor of 2 as in Lemma A.2, and can plug it into the general lemma in Lemma A.1 to obtain the following corollary. Corollary B.1. Let temperature (margin!) 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)ni=1 from νn, every f ∈ F and g ∈ G satisfy
Pr[arg max y′ f(x)y′ 6= y] ≤ ∥∥∥∥dµXdνn ∥∥∥∥ ∞ 1 m m∑ i=1 ‖Mγ(f)−Mγ(g)‖1 + 1 n n∑ i=1 Mγ(g(xi))yi
+ Õ ( k3/2
γ ∥∥∥∥dµXdνn ∥∥∥∥ ∞ ( Radm(F) + Radm(G) ) + √ k γ Radn(G) ) + 3 √ ln(1/δ)
2n
( 1 + ∥∥∥∥dµXdνn ∥∥∥∥ ∞ √ n m ) .
Proof. Overload function composition notation to sets of functions, meaning Mγ ◦ F = { (x, y) 7→ Mγ(f(x))y : f ∈ F } .
First note thatMγ is (2/γ)-Lipschitz with respect to the `∞ norm, and thus, applying the multivariate Lipschitz composition lemma for Rademacher complexity (Foster & Rakhlin, 2019, Theorem 1) just as in the proof for the softmax in Lemma A.2,
Radm(Mγ ◦ F) = Õ
( 2 √ k
γ Radm(F)
) ,
with similar bounds for Radm(Mγ ◦ G) and Radn(Mγ ◦ G). The desired statement now follows by combining these Rademacher complexity bounds with Lemma 1.1 applied toMγ ◦ F andMγ ◦ G, and additionally using 1[arg maxy′ f(x)y′ 6= y] ≤Mγ(f(x))y .

heading: C SAMPLING TOOLS
text: The proofs of Lemma 3.1 and Lemma 3.2 both make heavy use of sampling. Lemma C.1 (Maurey (Pisier, 1980)). Suppose random variable V is almost surely supported on a subset S of some Hilbert space, and let (V1, . . . , Vk) be k iid copies of V . Then there exist (V̂1, . . . , V̂k) ∈ Sk with∥∥∥∥∥∥EV − 1k ∑ i V̂i ∥∥∥∥∥∥ 2
F
≤ E V1,...,Vk ∥∥∥∥∥∥EV − 1k ∑ i Vi ∥∥∥∥∥∥ 2
F
= 1
k
[ E‖V ‖2F − ‖EV ‖2F ] ≤ 1 k E‖V ‖2F ≤ 1 k sup V̂ ∈S ‖V̂ ‖2F .
Proof of Lemma C.1. The first inequality is via the probabilistic method. For the remaining inequalities, by expanding the square multiple times,
E V1,...,Vk ∥∥∥∥∥∥EV − 1k ∑ i Vi ∥∥∥∥∥∥ 2
F
≤ E V1,...,Vk
1
k2 ∑ i ‖EV − Vi‖2F + ∑ i 6=j 〈 EV − Vi,EV − Vj 〉 = 1
k EV1‖V1 − EV ‖ 2 F =
1
k
[ E‖V ‖2F − ‖EV ‖2F ] ≤ 1 k E‖V ‖2F ≤ 1 k sup V̂ ∈S ‖V̂ ‖2F .
A first key application of Lemma C.1 is to sparsify products, as used in Lemma 3.2. Lemma C.2. Let matrices A ∈ Rd×m and B ∈ Rn×m be given, along with sampling budget k. Then there exists a selection (i1, . . . , ik) of indices and a corresponding diagonal sampling matrix M with at most k nonzero entries satisfying
M := ‖A‖2F k k∑ j=1 eije T ij ‖Aeij‖2 and
∥∥ABT −AMBT∥∥2 ≤ 1 k ‖A‖2‖B‖2.
Proof of Lemma C.2. For convenience, define columns ai := Aei and bi := Bei for i ∈ {1, . . . ,m}. Define importance weighting βi := (‖ai‖/‖A‖F)2, whereby ∑ i βi = 1, and let V
be a random variable with Pr [ V = β−1i aib T i ] = βi,
whereby EV = m∑ i=1 β−1i aib T iβi = m∑ i=1 (Aei)(Bei) T = A [ m∑ i=1 eie T i ] BT = A [I]BT = AB,
E‖V ‖2 = m∑ i=1 β−2i ‖aib T i‖2F βi = m∑ i=1 β−1i ‖ai‖ 2‖bi‖2 = m∑ i=1 ‖A‖2F ‖bi‖2 = ‖A‖2F · ‖B‖2F .
By Lemma C.1, there exist indices (i1, . . . , ik) and matrices V̂j := β−1ij aij b T ij with∥∥∥∥∥∥ABT − 1k ∑ j V̂j ∥∥∥∥∥∥ 2 ≤ ∥∥∥∥∥∥EV − 1k ∑ j V̂j ∥∥∥∥∥∥ 2 = 1 k [ ‖A‖2F ‖B‖2F − ‖AB‖2F ] ≤ 1 k ‖A‖2F ‖B‖2F .
To finish, by the definition of M ,
1
k ∑ j V̂j = 1 k ∑ j β−1ij (Aeij )(Beij ) T = A 1 k ∑ j β−1ij eije T ij BT = A [M ]BT.
A second is to cover the set of matrices W satisfying a norm bound ‖W T‖2,1 ≤ r. The proof here is more succinct and explicit than the one in (Bartlett et al., 2017a, Lemma 3.2). Lemma C.3 (See also (Bartlett et al., 2017a, Lemma 3.2)). Let norm bound r ≥ 0, X ∈ Rn×d, and integer k be given. Define a family of matrices
M := r‖X‖Fk k∑ l=1 sleile T jl ‖Xejl‖ : sl ∈ {±1}, il ∈ {1, . . . , n}, jl ∈ {1, . . . , d}  . Then
|M| ≤ (2nd)k, sup ‖WT‖2,1≤r min Ŵ∈M
‖WXT − ŴXT‖2F ≤ r2‖X‖2F
k .
Proof. Let W ∈ Rm×d be given with ‖W T‖2,1 ≤ r. Define sij := Wij/|Wij |, and note
WXT = ∑ i,j eie T iWeje T jX T = ∑ i,j eiWij(Xej) T = ∑ i,j
|Wij |‖Xej‖2 r‖X‖F︸ ︷︷ ︸ =:qij r‖X‖Fsijei(Xej)T ‖Xej‖︸ ︷︷ ︸ =:Uij .
Note by Cauchy-Schwarz that∑ i,j qij ≤ 1 r‖X‖F ∑ i √∑ j W 2ij‖X‖F = ‖W T‖2,1‖X‖F r‖X‖F ≤ 1,
potentially with strict inequality, thus q is not a probability vector. To remedy this, construct probability vector p from q by adding in, with equal weight, some Uij and its negation, so that the above summation form of WXT goes through equally with p and with q.
Now define iid random variables (V1, . . . , Vk), where
Pr[Vl = Uij ] = pij , EVl = ∑ i,j pijUij = ∑ i,j qijUij = WX T,
‖Uij‖ = ∥∥∥∥∥sijei(Xej)‖Xej‖2 ∥∥∥∥∥
F
· r‖X‖F = |sij | · ‖ei‖2 · ∥∥∥∥∥ Xej‖Xej‖2 ∥∥∥∥∥
2
· r‖X‖F = r‖X‖F,
E‖Vl‖2 = ∑ i,j pij‖Uij‖2 ≤ ∑ ij pijr 2‖X‖2F = r2‖X‖2F .
By Lemma C.1, there exist (V̂1, . . . , V̂k) ∈ Sk with∥∥∥∥∥∥WXT − 1k ∑ l V̂l ∥∥∥∥∥∥ 2 ≤ E ∥∥∥∥∥∥EV1 − 1k ∑ l Vl ∥∥∥∥∥∥ 2 ≤ 1 k E‖V1‖2 ≤ r2‖X‖2F k .
Furthermore, the matrices V̂l have the form
1
k ∑ l V̂l = 1 k ∑ l sleil(Xejl) T ‖Xejl‖ = 1 k ∑ l sleile T jl ‖Xejl‖ XT =: ŴXT, where Ŵ ∈M. Lastly, note |M| has cardinality at most (2nd)k.

heading: D PROOFS FOR SECTION 1.2
text: The bulk of this proof is devoted to establishing the Rademacher bound for computation graphs in Lemma 3.1; thereafter, as mentioned in Section 3, it suffices to plug this bound and the data augmentation bound in Lemma 1.2 into Lemma 1.1, and apply a pile of union bounds.
As mentioned in Section 3, this proof follows the scheme laid out in (Bartlett et al., 2017a), with simplifications due to the removal of “reference matrices” and some norm generality.
Proof of Lemma 3.1. Let cover scale and per-layer scales ( 1, . . . , L) be given; the proof will develop a covering number parameterized by these per-layer scales, and then optimize them to derive the final covering number in terms of . From there, a Dudley integral will give the Rademacher bound. Define b̃i := bi √ n for convenience. As in the statement, recursively define
XT0 := X T, XTi := σi ( [WiΠiDi|〉Fi]XTi−1 ) .
The proof will recursively construct an analogous cover via
X̂T0 := X T, X̂Ti := σi ( [ŴiΠiDi|〉Fi]X̂Ti−1 ) ,
where the choice of Ŵi depends on X̂i−1, and thus the total cover cardinality will product (and not simply sum) across layers. Specifically, the cover Ni for Ŵi is given by Lemma C.3 by plugging in ‖ΠiDiX̂Ti−1‖F ≤ b̃i, and thus it suffices to choose
cover cardinality k := r2i b̃ 2 i
2i , whereby min Ŵi∈Ni ‖WiΠiDiX̂Ti−1 − ŴiΠiDiX̂Ti−1‖ ≤ i.
By this choice (and the cardinality estimate in Lemma C.3, the full cover N satisfies
ln |N | = ∑ i ln |Ni| ≤ ∑ i r2i b̃ 2 i 2i ln(2m2).
To optimize the parameters ( 1, . . . , L), the first step is to show via induction that
‖XTi − X̂Ti ‖F ≤ ∑ j≤i jρj i∏ l=j+1 slρl.
The base case is simply ‖XT0 − X̂T‖ = ‖XT − XT‖ = 0, thus consider layer i > 0. Using the inductive formula for X̂i and the cover guarantee on Ŵi,∥∥∥XTi − X̂Ti ∥∥∥ =∥∥∥σi([WiΠiDi|〉Fi]XTi−1)− σi([ŴiΠiDi|〉Fi]X̂Ti−1)∥∥∥
≤ ρi ∥∥∥[WiΠiDi|〉Fi]hXTi−1 − [ŴiΠiDi|〉Fi]X̂Ti−1∥∥∥
≤ ρi ∥∥∥[WiΠiDi|〉Fi]XTi−1 − [WiΠiDi|〉Fi]X̂Ti−1∥∥∥+ ρi∥∥∥[WiΠiDi|〉Fi]X̂Ti−1 − [ŴiΠiDi|〉Fi]X̂Ti−1∥∥∥
≤ ρi ∥∥[WiΠiDi|〉Fi]∥∥2∥∥∥XTi−1 − X̂Ti−1∥∥∥+ ρi∥∥∥[(Wi − Ŵi)ΠiDiX̂Ti−1|〉(Fi − Fi)X̂Ti−1]∥∥∥ ≤ siρi ∑ j≤i−1 jρj i−1∏ l=j+1 slρl + ρi
∥∥∥(Wi − Ŵi)ΠiDiX̂Ti−1∥∥∥ ≤ ∑ j≤i−1 jρj i∏ l=j+1 slρl + ρi i ≤ ∑ j≤i jρj i∏ l=j+1 slρl.
To balance ( 1, . . . , L), it suffices to minimize a Lagrangian corresponding to the cover size subject to an error constraint, meaning
L(~ , λ) = L∑ i=1 αi 2i + λ  L∑ i=1 iβi −  where αi := r2i b̃2i ln(2m2), βi := ρi L∏ l=i+1 slρl,
whose unique critical point for ~ > 0 implies the choice
i := 1
Z ( 2αi βi )1/3 where Z := 1 ∑ i (2αiβ 2 i ) 1/3,
whereby ‖XTL − X̂TL‖ ≤ automatically, and
ln |N | ≤ Z2 ∑ i r2i b̃ 2 i ln(2m 2) (2αi/βi)2/3
= 1
222/3 2∑ i r 2/3 i b̃ 2/3 i β 2/3 i ln(2m 2)1/3 2∑ i r 2/3 i b̃ 2/3 i ln(2m 2)1/3β 2/3 i
= 24/3 ln(2m2)
2
[∑ i ( rib̃iρi L∏ l=i+1 slρl )2/3]3 =: τ2 2 ,
as desired, with τ introduced for convenience in what is to come.
For the Rademacher complexity estimate, by a standard Dudley entropy integral (Shalev-Shwartz & Ben-David, 2014), setting τ̂ := max{τ, 1/3} for convenience,
nRad(G) ≤ inf ζ
4ζ √ n+12 ∫ √n ζ √ τ̂ d = inf ζ 4ζ √ n+12τ̂ ln( ) ∣∣√n ζ = inf ζ 4ζ √ n+12τ̂(ln √ n−ln ζ),
which is minimized at ζ = 3τ̂ / √ n, whereby
nRad(G) ≤ 12τ̂ + 6τ̂ lnn− 12τ̂ ln(3τ̂ / √ n) = 12τ̂(1− ln(3τ̂)) ≤ 12τ̂ ≤ 12τ + 4.
This now gives the proof of Theorem 1.3.
Proof of Theorem 1.3. With Lemma 1.1, Lemma 1.2, and Lemma 3.1 out of the way, the main work of this proof is to have an infimum over distillation network hyperparameters (~b, ~r, ~s) on the right hand side, which is accomplished by dividing these hyperparameters into countably many shells, and unioning over them.
In more detail, divide (~b, ~r, ~s) into shells as follows. Divide each bi and ri into shells of radius increasing by one, meaning meaning for example the first shell for bi has bi ≤ 1, and the jth shell has bi ∈ (j − 1, j], and similarly for ri; moreover, associate the jth shell with prior weight qj(bi) := (j(j + 1)) −1, whereby ∑ j≥1 qj(bi) = 1. Meanwhile, for si use a finer grid where the first shell has si ≤ 1/L, and the jth shell has si ∈ ((j − 1)/L, j/L), and again the prior weight is qj(si) = (j(j + 1))
−1. Lastly, given a full set of grid parameters (~b, ~r, ~s), associate prior weight q(~b, ~r, ~s) equal to the product of the individual prior weight, whereby the sum of the prior weights over the entire product grid is 1. Enumerate this grid in any way, and define failure probability δ(~b, ~r, ~s) := δ · q(~b, ~r, ~s).
Next consider some fixed grid shell with parameters (~b′, ~r′, ~s′) and letH denote the set of networks for which these parameters form the tightest shell, meaning that for any g ∈ H with parameters (~b, ~r, ~s), then (~b′, ~r′, ~s′) ≤ (~b + 1, ~r + 1, ~s + 1) component-wise. As such, by Lemma 1.1, with probability at least 1− δ(~b′, ~r′, ~s′), each g ∈ H satisfies
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(H) ) + √ k γ Radn(H) )
+ 6
√ ln(q(~b′, ~r′, ~s′)) + ln(1/δ)
2n
( 1 + ∥∥∥∥dµXdνn ∥∥∥∥ ∞ √ n m ) .
To simplify this expression, first note by Lemma 3.1 and the construction of the shells (relying in particular on the finer grid for si to avoid a multiplicative factor L) that
Radm(H) = Õ [ 1√ n (∑ i [ r′ib ′ iρi L∏ l=i+1 s′lρl ]2/3)3/2]
= Õ [ 1√ n (∑ i [ (ri + 1)(bi + 1)ρi L∏ l=i+1 (sl + 1/L)ρl ]2/3)3/2]
= Õ [ 1√ n (∑ i [ ribiρi L∏ l=i+1 slρl ]2/3)3/2] ,
and similarly for Radm(H) (the only difference being √ m replaces √ n). Secondly, to absorb the term ln(q(~b′, ~r′, ~s′)), noting that ln(a) ≤ ln(γ2) + (a− γ2)/(γ2), and also using ρi ≥ 1, then
ln(q(~r′,~b′, ~s′)) = O ln∏ i (ri + 1) 2(bi + 1) 2((si + 1)L) 2  = O L lnL+ ln∏ i r 2/3 i b 2/3 i s 2/3 i  = Õ
L+∑ i ln(r 2/3 i b 2/3 i ) + ln ∏ i s 2/3 i  = Õ L+ ln(γ2) + 1 γ2 ∑ i r2/3i b2/3i +∏ l>i s 2/3 l  
= Õ ( L+ 1
γ2 ∑ i [ ribi L∏ l=i+1 sl ]2/3)
= Õ ( L+ 1
γ2 ∑ i [ ribiρi L∏ l=i+1 slρl ]2/3) .
Together,
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) ) + 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) .
Since h ∈ H was arbitrary, the bound may be wrapped in infg∈H. Similarly, unioning bounding away the failure probability for all shells, since this particular shell was arbitrary, an infimum over shells can be added, which gives the final infimum over (~b, ~r, ~s). The last touch is to apply Lemma 1.2 to bound ‖dµX/dνn‖∞.

heading: E PROOF OF STABLE RANK BOUND, THEOREM 1.4
text: The first step is to establish the sparsification lemma in Lemma 3.2, which in turn sparsifies each matrix product, cannot simply invoke Lemma C.2: pre-processing is necessary to control the elementwise magnitudes of the resulting matrix. Throughout this section, define the stable rank of a matrix W as sr(W ) := ‖W‖2F /‖W‖22 (or 0 when W = 0).
Lemma E.1. Let matrices A ∈ Rd×m and B ∈ Rn×m be given, along with sampling budget k. Then there exists a selection (i1, . . . , ik) of indices and a corresponding diagonal sampling matrix M with at most k nonzero entries satisfying
M := k∑ j=1 Zijeije T ij ‖aij‖ where Zij ≤ ‖A‖F √ m k , and
∥∥ABT −AMBT∥∥2 ≤ 4 k ‖A‖2‖B‖2.
Proof. Let τ > 0 be a parameter to be optimized later, and define a subset of indices S := {i ∈ {1, . . . ,m} : ‖Aei‖ ≥ τ}, with Sc := {1, . . . ,m} \ S. Let Aτ denote the matrix obtained by zeroing out columns not in S, meaning
Aτ := ∑ i∈S (Aei)e T i ,
whereby ‖ABT −AτBT‖F ≤ ‖A−Aτ‖ · ‖B‖ ≤ ‖B‖ √∑ i∈Sc ‖Aei‖2 ≤ τ √ m‖B‖.
Applying Lemma C.2 to AτBT gives
M := ‖Aτ‖2
k k∑ j=1 eije T ij ‖Aτeij‖2 = k∑ j=1 Zijeije T ij ‖Aτeij‖ such that ‖AτBT−AτMBT‖2 ≤ 1 k ‖Aτ‖2‖B‖2,
where Zij is specified by these equalities. To simplify, note ‖Aτ‖ ≤ ‖A‖, and AτM = AM . Combining the two inequalities,
‖ABT −AMBT‖ ≤ ‖ABT −AτBT‖+ ‖AτBT −AτMBT‖ ≤ τ √ m‖B‖+ 1√
k ‖A‖‖B‖.
To finish, setting τ := ‖A‖/ √ mk gives the bound, and ensures that the scaling term Zij satisfies, for any ij ∈ S,
Zij = ‖Aτ‖2 k‖Aτeij‖ ≤ ‖A‖ 2 F kτ = ‖A‖F
√ m
k .
With this tool in hand, the proof of Lemma 3.2 is as follows.
Proof of Lemma 3.2. Let Xj denote the network output after layer j, meaning
XT0 := X T, XTj := σj(WjX T j−1),
whereby ‖XTj ‖F = ‖σj(WjXTj−1)− σj(0)‖F ≤ ‖WjXTj−1‖F ≤ ‖Wj‖2‖XTj−1‖F ≤ ‖X‖F ∏ i≤j ‖Wi‖2.
The proof will inductively choose sampling matrices (M1, . . . ,ML) as in the statement and construct
X̂T0 := X T, X̂Tj := Πjσj(WjMjX̂ T j−1), where Πj denotes projection onto the Frobenius-norm ball of radius ‖X‖F ∏ i≤j ‖Wi‖2 (whereby ΠjXj = Xj), satisfying
∥∥∥Xj − X̂j∥∥∥ F ≤ ‖X‖F  j∏ p=1 ‖Wp‖2  j∑ i=1 √ sr(Wi) ki ,
which gives the desired bound after plugging in j = L.
Proceeding with the inductive construction, the base case is direct since X̂0 = X = X0 and∥∥∥X0 − X̂0∥∥∥ F = 0, thus consider some j > 0. Applying Lemma E.1 to the matrix multiplication WjX̂j−1 with kj samples, there exists a multiset of Sj coordinates and a corresponding sampling matrix Mj , as specified in the statement, satisfying∥∥∥WjX̂Tj−1 −WjMjX̂Tj−1∥∥∥ F ≤ 1√ kj ‖Wj‖F‖X̂j−1‖F ≤ 1√ kj ‖Wj‖F‖X‖F ∏ i<j ‖Wi‖2.
Using the choice X̂Tj := Πjσj(WjMjX̂ T j−1),∥∥∥Xj − X̂j∥∥∥ F = ∥∥∥σj(WjXTj−1)−Πjσj(WjMjX̂Tj−1)∥∥∥ F
≤ ∥∥∥WjXTj−1 −WjMjX̂Tj−1∥∥∥
F = ∥∥∥WjXTj−1 −WjX̂Tj−1 +WjX̂Tj−1 −WjMjX̂Tj−1∥∥∥
F ≤ ∥∥∥WjXTj−1 −WjX̂Tj−1∥∥∥ F + ∥∥∥WjX̂Tj−1 −WjMjX̂Tj−1∥∥∥ F
≤ ∥∥Wj∥∥2∥∥∥Xj−1 − X̂j−1∥∥∥F + 1√kj ‖Wj‖F‖X‖F ∏ i<j ‖Wi‖2 ≤ ∥∥Wj∥∥2 ‖X‖F ∏ i<j ‖Wi‖2 ∑ i<j √ sr(Wi) ki + √ sr(Wj) kj ‖X‖F ∏ i≤j ‖Wi‖2
≤ ‖X‖F ∏ i≤j ‖Wi‖2 ∑ i≤j √ sr(Wi) ki
as desired.
To prove Theorem 1.4 via Lemma 3.2, the first step is a quick tool to cover matrices element-wise.
Lemma E.2. Let A denote matrices with at most k2 nonzero rows and k1 nonzero columns, entries bounded in absolute value by b, and total number of rows and columns each at most m. Then there exists a cover setM⊆ A satisfying
|M| ≤ mk1+k2 ( 2b √ k1k2 )k1k2 , and sup
A∈A min Â∈M
‖A− Â‖F ≤ .
Proof. Consider some fixed set of k2 nonzero rows and k1 nonzero columns, and letM0 denote the covering set obtained by gridding the k1 · k2 entries at scale √k1k2 , whereby
|M0| ≤
( 2b √ k1k2 )k1k2 .
For any A ∈ A with these specific nonzero rows and columns, the Â ∈ M0 obtained by rounding each nonzero entry of A to the nearest grid element gives
‖A− Â‖2 = ∑ i,j (Aij − Âij)2 ≤ ∑ i,j ( √ k1k2 )2 = 2 ∑ i,j 1 k1k2 = 2.
The final coverM is now obtained by unioning copies ofM0 for all ( m k1 )( m k2 ) ≤ mk1+k2 possible submatrices of size k2 × k1.
The proof of Theorem 1.4 now carefully combines the preceding pieces.
Proof of Theorem 1.4. The proof proceeds in three steps, as follows.
1. A covering number is estimate for sparsified networks, as output by Lemma 3.2.
2. A covering number for general networks is computed by balancing the error terms from Lemma 3.2 and its cover computed here.
3. This covering number is plugged into a Dudley integral to obtain the desired Rademacher bound.
Proceeding with this plan, let (X̂T0 , . . . , X̂ T L) be the layer outputs (and network input) exactly as provided by Lemma 3.2. Additionally, define diagonal matrices Dj := ∑ l ! ∈Sj+1 ele T l (with DL = I , where the “!” denotes unique inclusion; these matrices capture the effect of the subsequent sparsification, and can be safely inserted after each Wj without affecting X̂Tj , meaning
X̂Tj = Πjσj(WjMjX̂ T j−1) = Πjσj(DjWjMjX̂ T j−1).
Let per-layer cover precisions ( 1, . . . , L) be given, which will be optimized away later. This proof will inductively construct
X̃T0 := X T, X̃Tj := Πjσj(W̃jX̃ T j−1),
where W̃j is a cover element for DjWjMj , and inductively satisfying ‖X̂Tj − X̃Tj ‖ ≤ ‖X‖Fmj/2 ∑ i≤j i ∏ l≤j l 6=i ‖Wj‖F.
To construct the per-layer cover elements W̃j , first note by the form of Mj (and the scaling Zi provided by Lemma 3.2) that
b := max i,l (DjWjMj)l,i ≤ max i ‖WjMjei‖ ≤ Zi ‖Wjei‖ ‖Wjei‖
≤ ‖Wj‖F √ m
kj−1 .
Consequently, by Lemma E.2, there exists a cover Cj of matrices of the form DjWjMj satisfying |Cj | ≤ mkj+kj−1 ( 2b √ kjkj−1
j
)kjkj−1 ≤ mkj+kj−1 ( 2‖Wj‖F √ kjm
j
)kjkj−1 ,
and the closest cover element W̃jCj to DjWjMj satisfies ‖DjWjMj − W̃j‖F ≤ j .
Proceeding with the induction, the base case has ‖X̂T0 − X̃T0‖ = ‖XT − XT‖ = 0, thus consider j > 0. The first step is to estimate the spectral norm of DjWjMj , which can be coarsely upper bounded via
‖DjWjMj‖22 ≤ ‖DjWjMj‖2F ≤ ∑ i ‖WjMjei‖2 ≤ ∑ i ‖Wj‖2F m kj−1 ≤ ‖Wj‖2Fm.
By the form of X̂j and X̃j ,
‖X̂Tj − X̃Tj ‖ = ‖Πjσj(DjWjMjX̂Tj−1)−Πjσj(W̃jX̃Tj−1)‖
≤ ‖DjWjMjX̂Tj−1 − W̃jX̃Tj−1‖
≤ ‖DjWjMjX̂Tj−1 −DjWjMjX̃Tj−1‖+ ‖DjWjMjX̃Tj−1 − W̃jX̃Tj−1‖
≤ ‖DjWjMj‖2‖X̂Tj−1 − X̃Tj−1‖F + ‖DjWjMj − W̃j‖2‖X̃Tj−1‖F ≤ √ m‖Wj‖F [ ‖X‖Fm(j−1)/2 ∑ i<j i ∏ l<j l 6=i ‖Wj‖F ] + j‖X‖F ∏ i<j ‖Wj‖2
≤ ‖X‖Fmj/2 ∑ i≤j i ∏ l≤j l 6=i ‖Wj‖F,
which establishes the desired bound on the error.
The next step is to optimize kj . Let > 0 be arbitrary, and set −1j := −12L √ m‖X‖F ∏ i 6=j ‖Wi‖F, whereby
‖X̂TL − X̃TL‖F ≤ 2 , |Cj | ≤ mkj+kj−1
( 4mL √ kj‖X‖F ∏ i ‖Wi‖F )kjkj−1 .
The overall network cover N is the product of the covers for all layers, and thus has cardinality satisfying
ln |N | ≤ ∑ j ln |Cj | ≤ 2 ∑ j kj lnm+ ∑ j kjkj−1 ln
( 4mL √ kj‖X‖F ∏ i ‖Wi‖F
j
)
≤ 2 ∑ j kj lnm+ ∑ j 2k2j ln
( 4mL √ kj‖X‖F j ) + ∑ j 2k2j  · ∑ j ln ‖Wj‖F  . To choose (k1, . . . , kL), letting XTL denote the output of the original unsparsified network, note firstly that the full error bound satisfies
‖XTL − X̃TL‖ ≤ ‖XTL − X̂TL‖+ ‖X̂TL − X̃TL‖
≤ ∑ i αi√ ki + 2 , where αi := ‖X‖F ∏ i ‖Wi‖2 √sr(Wi). To choose ki, the approach here is to minimize a Lagrangian corresponding to the cover cardinality, subject to the total cover error being . Simplifying the previous expressions and noting 2kjkj−1 ≤ k2j + k 2 j−1, whereby the dominant term in ln |N | is ∑ j k 2 j , consider Lagrangian
L(k1, . . . , kl, λ) := ∑ i k2i + λ ∑ i αi√ ki − 2  , which has critical points when each ki satisfies
k 5/2 i
αi = λ 4 ,
thus ki := α 2/5 i /Z with Z := 2/(2 ∑ j α 4/5 j )
2. As a sanity check (since it was baked into the Lagrangian), plugging this into the cover error indeed gives
‖XTL − X̃TL‖ ≤ ∑ i αi√ ki + 2 = √ Z ∑ i α 4/5 i + 2 = .
To upper bound the cover cardinality, first note that∑ i k2i = 1 Z2 ∑ i α 4/5 i = 4 4 (∑ i α 4/5 i )5 ,
whereby
ln |N | = Õ [∑ i k2i ] · [∑ i ln ‖Wi‖F ]
= β
4 where β = Õ ( ‖X‖4F [∏ j ‖Wj‖42 ] [∑ i sr(Wi)2/5 ]5 [∑ i ln ‖Wi‖F ]) .
The final step is to apply a Dudley entropy integral (Shalev-Shwartz & Ben-David, 2014), which gives
nRad(F) = inf ζ
( 4ζ √ n+ 12 ∫ √n ζ √ β 2 d ) = inf ζ ( 4ζ √ n+ 12 [ 1 ζ − 1√ n ]√ β ) .
Dropping the negative term gives an expression of the form aζ + b/ζ, which is convex in ζ > 0 and has critical point at ζ2 = b/a, which after plugging back in gives an upper bound 2 √ ab, meaning
nRad(F) ≤ 2 ( 4 √ n · 12 √ β )1/2 = 8 √ 3n1/4β1/4.
Dividing by n and expanding the definition of β gives the final Rademacher complexity bound.

Review summaries:
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. 

The generalization performance of learning algorithms characterizes their ability to generalize their empirical behavior on training examples to unseen test data, which provides an intuitive understanding of how different parameters affect the learning performance and some guides to design learning machines. Different from the traditional error analysis, this paper focuses on bounding the divergence bettween the test error and the training error by the the corresponding distillation error and distillation complexity, e.g., test error  is bounded by training error + distillation error + distillation complexity. The current learning theory analysis may be important to understand the theoretical foundations of distillation strategy in deep networks. However, some theoretical issues should be illustrated to improve its readability, e.g,.