THE INTRIGUING ROLE OF MODULE CRITICALITY IN THE GENERALIZATION OF DEEP NETWORKS
Year: 2020

Abstract:
We study the phenomenon that some modules of deep neural networks (DNNs) are more critical than others. Meaning that rewinding their parameter values back to initialization, while keeping other modules fixed at the trained parameters, results in a large drop in the network’s performance. Our analysis reveals interesting properties of the loss landscape which leads us to propose a complexity measure, called module criticality, based on the shape of the valleys that connect the initial and final values of the module parameters. We formulate how generalization relates to the module criticality, and show that this measure is able to explain the superior generalization performance of some architectures over others, whereas, earlier measures fail to do so.

text: We study the phenomenon that some modules of deep neural networks (DNNs) are more critical than others. Meaning that rewinding their parameter values back to initialization, while keeping other modules fixed at the trained parameters, results in a large drop in the network’s performance. Our analysis reveals interesting properties of the loss landscape which leads us to propose a complexity measure, called module criticality, based on the shape of the valleys that connect the initial and final values of the module parameters. We formulate how generalization relates to the module criticality, and show that this measure is able to explain the superior generalization performance of some architectures over others, whereas, earlier measures fail to do so.

heading: 1 INTRODUCTION
text: Neural networks have had tremendous practical impact in various domains such as revolutionizing many tasks in computer vision, speech and natural language processing. However, many aspects of their design and analysis have remained mysterious to this date. One of the most important questions is “what makes an architecture work better than others given a specific task?” Extensive research in this area has led to many potential explanations on why some types of architectures have better performance; however, we lack a unified view that provides a complete and satisfactory answer. In order to attain a unified view on superiority of one architecture over another in terms of generalization performance, we need to come up with a measure that effectively captures this.
Analyzing the generalization behavior of neural networks has been an active area of research since Baum & Haussler (1989). Many generalization bounds and complexity measures have been proposed so far. Bartlett (1998) emphasized the importance of the norm of the weights in predicting the generalization error. Since then various analysis have been proposed. These results are either based on covering number and Rademacher complexity (Neyshabur et al., 2015; Bartlett et al., 2017; Neyshabur et al., 2019; Long & Sedghi, 2019; Wei & Ma, 2019), or they use approaches similar to PAC-Bayes (McAllester, 1999; Dziugaite & Roy, 2017; Neyshabur et al., 2017; 2018; Arora et al., 2018; Nagarajan & Kolter, 2019a; Zhou et al., 2019). Recently authors have emphasized on the role of distance to initialization rather than norm of the weights in generalization (Dziugaite & Roy, 2017; Nagarajan & Kolter, 2019b; Neyshabur et al., 2019; Long & Sedghi, 2019). Earlier results have an exponential dependency on the depth and focus on fully connected networks. More recently, Long & Sedghi (2019) provided generalization bounds for convolutional neural networks (CNNs) and fully connected networks used in practice and their bounds have linear dependency on the depth.
Despite the success of earlier works in capturing the dependency of generalization performance of a model on different parameters, they fail at the following task: Rank the generalization performance of candidate architectures for a specific task such that the ranking aligns well with the ground truth. Moreover, majority of these bounds are proposed for fully connected modules and it is not straightforward to evaluate them for different architectures such as ResNets.
Every DNN architecture is a computation graph where each node is a module1. We are interested in understanding how different modules in the network interact with each other and influence generalization performance as a whole. To do so, we delve deeper into the phenomenon of “module criticality” which was reported by Zhang et al. (2019a). They observed that modules of the network present different robustness characteristics to parameter perturbation. Specifically, they look into the following perturbation: Rewind one module back to its initialization value while keeping all other modules fixed (at the final trained value). They note that the impact of this perturbation on network performance varies between modules and depends on which module was rewound. Some modules are “critical” meaning that rewinding their value to the initialization harms the network performance, while for others the impact of this perturbation on performance is negligible. They show that various conventional DNN architectures exhibit this phenomenon.
Let us now informally define what we mean by the measure “module criticality” (see Figure 1). For each module, we move on a line from its final trained value to its initialization value (convex combination2 path) while keeping all other modules fixed at their trained value. Then we measure the performance drop. Let θαi = (1 − α)θ0i + αθFi , α ∈ [0, 1] be the convex combination between initial weights θ0i and the final weights θ F i at module i, where αi is the minimum value between 0 and 1 when performance (train error) of the network drops by at most a threshold value . If αi is small we can move a long way back to initialization without hurting performance and the “module criticality” of this module would be low. Further, we also wish to incorporate the robustness to noise (that is, the valley width) for the module along this path. If the module is robust to noise along this path (that is, the valley is wide) then the module criticality would again be low (see Definition 3.1 for a formal definition).
In this paper, we seek to study this phenomenon in depth and shed some light on it by showing that conventional complexity measures cannot capture criticality (see Section 2). Next, we theoretically formulate this phenomenon and analyze its role in generalization. Through this analysis, we provide
1A module is a node in the computation graph that has incoming edges from other modules and outgoing edges to other nodes and performs a linear transformation on its inputs. For layered model such as VGG, module definition is equivalent to definition of a layer.
2A convex combination of two points is a linear combination of them where the coefficients are non-negative and sum to 1. Every convex combination of two points lies on the line segment between the two.
a new generalization measure that captures the dissimilarity of different modules and depicts how it influences the generalization of the corresponding DNN. Intuitively, the closer we can get to initialization for each module, the better the generalization.
We analyze the relation between generalization and module criticality through a PAC-Bayesian analysis. In Section 3 we show that it is the overall network criticality measure and not the number of critical modules that controls generalization. If the network criticality measure is smaller for an architecture, it has better generalization performance. In Section 4, we demonstrate through various experiments that our proposed measure is able to distinguish between different network architectures in terms of their generalization performance. Moreover, the network criticality measure is able to correctly rank the generalization performance of different architectures better than the measures proposed earlier.
Notation: We use upper case letters for matrices. The operator norm and Frobenius norm of M are denoted by ‖M‖2, ‖M‖Fr respectively. For n ∈ N, we use [n] to denote the set {1, . . . , n}. Let LS(f) be the loss of function f on the training set S with m samples. We are mainly interested in the classification task where LS(f) = 1m ∑ (x,y)∈S 1[f(x)[y] ≤ maxj 6=y f(x)[j]]. For any γ > 0,
we also define margin loss LS,γ(f) = 1m ∑
(x,y)∈S 1[f(x)[y] ≤ γ + maxj 6=y f(x)[j]]. Let LD(f) be the loss of function f on population data distribution D defined similar to LS(f). We will denote the function parameterized by Θ by fΘ.

heading: 2 TOWARDS UNDERSTANDING MODULE CRITICALITY

heading: 2.1 SETTING
text: A DNN architecture is a directed acyclic computation graph which may or may not be layered. In order to have a unifying definition between different architectures, we use the notion of a “module”. A module is a node in the computation graph that has incoming edges from other modules and outgoing edges to other nodes, and performs a linear transformation on its inputs. For a layered model such as a VGG, a module is equivalent to a layer. On the other hand, in a ResNet some modules are parallel to each other. For example, a downsample module and the concatenation of two convolutional modules in a ResNet18-v1 architecture. Note that, similar to conventional definitions the non-linearity (such as a ReLU) is not part of the module.
Let Θ = (θ1, . . . , θd) correspond to all parameters of a DNN with d modules, where θi refers to the weight matrix (or operator matrix in case of convolution) at module i and θ0i , θ F i refer to the value
of weight matrix at initialization and the end of training, respectively. For sequential architectures, d is equal to the depth of the network but that is not necessarily true for a general architectures such as a ResNet.

heading: 2.2 ROBUSTNESS TO REWINDING
text: Consider the following perturbation to a trained network at some training epoch as considered by Zhang et al. (2019a). For each module in the network, rewind its value back to its value at this training epoch while keeping the values of all other modules fixed (at their final trained value). Next, measure the change in performance of the model before and after this manipulation. We repeat a similar analysis that differs from that of Zhang et al. (2019a) in one detail. Zhang et al. (2019a) rewind the whole ResNet block at once, whereas, we rewind each module (each convolutional module) separately. This rewind analysis is shown in Figure 2 for ResNet18-v1. Each column represents a module in ResNet18-v1 and each row represents a particular training epoch to which this module is rewound to. Similar to earlier analysis, we observe that for many modules of the network, this manipulation does not influence the network performance drastically, while, for some others the impact is more pronounced. For example, in Figure 2 we look at the effect of rewinding on train error. The “Stage2.block1.conv2” module is critical, whereas, most other modules, once rewound, do not affect the performance. In Figure 6 in Appendix E we plot the effect of rewinding on different performance criteria (train loss, train error and test error) and observe that they exhibit a similar trend.
A stable phenomena: The plots in Figure 2 capture a network trained with SGD with weights initialized using the standard Kaiming initialization (He et al., 2015). To ensure that the observed phenomenon is not an artifact of the training method and the initialization scheme, we repeated the experiments with different initialization and optimization methods. We saw a similar pattern. For example, Figures 7a, 7b in the appendix illustrate the pattern when we changed the initialization to Fixup (Zhang et al., 2019b), and when we replace SGD with Adam (Kingma & Ba, 2014) respectively.

heading: 2.3 WHAT MEASURES FAIL TO DISTINGUISH CRITICAL LAYERS
text: Spectrum of weight matrices: We explore the change in the spectrum of different weight matrices on rewinding and note that the spectrum for a critical and non-critical module look similar. This is shown in Figure 3. We calculate the spectrum of convolutional layers using the algorithm by Sedghi et al. (2019).
Distance to initialization: Next, we analyze the operator norm of difference from initialization for each module. Figure 8 in the appendix depicts this and reveals no difference between critical and non-critical modules. A similar plot was explored by Zhang et al. (2019a), where they find that the Frobenius norm and the infinity norm also fail to capture criticality.
Change in the activation patterns: We investigated the change in the activation patterns of a network when we rewind a module. To do this, we study the similarity between two networks: 1. The original trained network and 2. The network with a rewound module. We use CKA (Kornblith et al., 2019) as the measure of similarity. For a non-critical module, the original and rewound networks are similar and in case of a critical module, the similarity between the activation patterns between the two networks degrades gradually rather than abruptly. See Figure 9 in Appendix E.

heading: 3 GENERALIZATION BOUNDS BASED ON MODULE CRITICALITY
text: Our goal is to understand criticality and how it affects the generalization performance of a DNN. Inspired by the rewind to initialization experiments of Zhang et al. (2019a), we take one step further and consider changing the value of each module, to the convex combination of its initial and final value. That is, for each module i, we replace θi with θαi = (1−α)θ0i +αθFi , α ∈ [0, 1], and keep all other layers fixed. Then we look at the effect of this perturbation on the performance of the network.
Figure 4 depicts how the train error, test error and train loss change as we decrease the value of α in θαi when i refers to a critical module (yellow dashdot curve), a non-critical module (red dashed curve) and all modules (blue solid curve). We find that along this convex combination path all these performance measures degrade monotonically (increase in error and loss), as we move from the final weights to the initial weights.
The above experiment shows the effect of moving along a convex combination between module’s initial and trained value. To capture the relation between criticality and generalization, we are interested in also accounting for the width of the valley as we move from the final value to the initial value. In particular, we are interested in analyzing what happens if we are moving inside a ball of some radius σi around each point in this path. PAC-Bayesian analysis, looks for a ball around final value of parameters such that the loss does not change if we move in this ball. Bringing this idea together with the one mentioned above, we are interested in moving from the final value to the initial value in a valley of some radius, and want to find out how far we can move on this path. Intuitively, being able to move closer to initialization values indicate that the effective function class is smaller and hence the network should generalize better. For example, in the extreme case where none of the weights change from their initialization value the function class would be a single function (the initial function) and the generalization error would be very low (∼0%) as both the train and test error would be very high (but equal). In this paper, we consider the case of trained models where train error is very low and hence low generalization error means good performance on test data. Definition 3.1 (Module and Network Criticality). Given an > 0 and network fΘ, we define the module criticality for module i as follows:
µi, (fΘ) = min 0≤αi,σi≤1
{ α2i ∥∥θFi − θ0i ∥∥2Fr
σ2i : Eu∼N (0,σ2i )[LS(fθαi +u,ΘF−i)] ≤
} , (1)
We also define the network criticality as the sum of the module criticality over modules of the network:
µ (fΘ) = d∑ i=1 µi, (fΘ). (2)
Here, LS denotes the empirical zero-one loss over the training set, fθαi ,ΘF−i is the DNN’s function value where weight matrix corresponding to ith module is replaced by θαi and all other modules are fixed at their values in the end of training, ΘF−i. θ α i = (1−α)θ0i +αθFi , where θ0i is the value of the weight matrix at initialization and θFi is the trained value.
Intuitively, network criticality measure is sum of module criticalities. This is also theoretically derived using the analysis below.

heading: 3.1 A PAC-BAYESIAN GENERALIZATION BOUND
text: We attempt to understand the relationship between module criticality, and generalization by deriving a generalization bound using the PAC-Bayesian framework (McAllester, 1999). Given a prior
distribution over the parameters that is picked in advance before observing a training set, a posterior distribution over the parameters that could depend on the training set and a learning algorithm, the PAC-Bayesian framework bounds the generalization error in terms of the Kullback-Leibler (KL) divergence (Kullback & Leibler, 1951) between the posterior and the prior distribution. We use PAC-Bayesian bounds as they hold for any architecture.
The intuition from Figure 4 suggests moving the parameters of each module as close as possible to the initialization value before harming the performance. For such αi, we can then define the posterior Qi for module i to be a Gaussian distribution centered at θαi with covariance matrix σ 2 i I , that is, as if we have additive noise ui ∼ N (0, σiI). We use Θα to refer to the case where all the parameters θi are replaced with θαii and matrix U includes all the noise ui. Then the following theorem holds. Theorem 3.2. For any data distribution D, number of samples m ∈ N , for any 0 < δ < 1, for any 0 < σi ≤ 1 and any 0 ≤ αi ≤ 1, with probability 1− δ over the choice of the training set Sm ∼ D the following generalization bound holds:
EU [LD(fΘα+U )] ≤ EU [LS(fΘα+U )] +
√√√√√ 14 ∑di=1 ki log ( 1 + α2i‖θFi −θ0i‖2Fr kiσ2i ) + log ( m δ ) + Õ(1)
m− 1 ,
where ki is the number of parameters in module i. For example, for a convolution module with kernel size qi × qi and number of output channels ci, ki = q2i ci−1ci.
The exact bound including the constants and the proof of the theorem above is given in Appendix A. Theorem 3.2 already gives us some insight into generalization of the original network. However, it is not exactly a generalization bound on the original network but rather on a perturbed network. We conjecture that for almost any realistic distribution D, any random Θ0, any ΘF achieved by known gradient based optimization algorithms, any 0 ≤ α ≤ 1 and any σ ≥ 0, the test error does not improve by taking a convex combination of parameters and their initial values followed by Gaussian perturbation. Therefore, we have that LD(fΘF ) ≤ EU [LD(fΘα+U )]. The following corollary restates Theorem 3.2 by using this assumption and optimizing over α and σ in the bound. Corollary 3.3. For any data distribution D, number of samples m ∈ N. For any > 0, for any 0 < δ, if LD(fΘF ) ≤ EU [LD(fΘα+U )] where ui ∼ N (0, σiI) , then with probability 1 − δ over the choice of the training set Sm ∼ D, the following generalization bound holds
LD(fΘ) ≤ +
√ 1 4µ ′ (fΘ) + log ( m δ ) + Õ(1)
m− 1 ,
where µ′ (fΘ) is calculated as follows:
µ′ (fΘ) = min 0≤α,σ≤1 {∑ i α2i ∥∥θFi − θ0i ∥∥2Fr σ2i : EU [LS(fΘα+U )] ≤ } .
Note that the above bound uses a slightly different notion of network criticality compared to Definition 3.1, as the bound requires finding α and σ values simultaneously for all modules, whereas, Definition 3.1 allows us to decouple the search over α and σ.
Deterministic generalization bound for convolutional networks Although PAC-Bayesian bounds are data-dependent and hence numerically superior, they provide less insight about the underlying reason that results in generalization. For example, the flatness of the solution after adding Gaussian perturbation can be computed numerically. But computing this value does not reveal what properties of the network enforce the loss surface around a point to be flat. On the other hand, deterministic norm-based generalization bounds are numerically much looser yet they provide better insights into the dependence of generalization on different network parameters. In Appendix B, we build on the results of Theorem 3.2 to present a norm-based deterministic bound using module criticality.
In this section, we intuitively justified network criticality measure and related the generalization of a DNN to the network criticality measure in Corollary 3.3. In the next section, we empirically show that the network criticality measure is able to correctly rank the generalization performance of different architectures better than measures proposed earlier.

heading: 4 EXPERIMENTS
text: We perform several experiments to compare our network criticality measure to earlier complexity measures in the literature. Our experiments are performed on the CIFAR10 and CIFAR100 datasets. For all experiments, implementation and architecture details are presented in Appendix C.
Table 1 summaries the quantities that are calculated in this section. The quantity SoSP was proposed by Long & Sedghi (2019). For the last two measures, we calculate σi and αi as per Definition 3.1. In our experiments, each module is a single convolutional or linear layer. This represents a natural choice where each module is a linear transformation (with respect to the parameters). This choice also leads to the lowest number of modules such that each module is a linear transformation.
First, as a sanity check we use our complexity measure (lower is better) to compare between a ResNet18 trained on true labels and a ResNet18 trained on data where 20% of the labels are randomly corrupted. As seen in Figure 5a, our measure is able to correctly capture that the network trained with true label generalizes better than the one trained on corrupted labels (4.62% error vs. 35% error).
Next, in Table 2 we compare the generalization performance of several conventional DNN architectures trained on the CIFAR10 dataset. There is a particular ranking of the networks based on their generalization error and it is desirable for a complexity measure to capture this ranking. Therefore, we compare the rankings proposed by network criticality measure and complexity measures from the literature with the empirical rankings obtained in the experiment. To do this, we calculate the Kendall’s τ correlation coefficient (Kendall, 1938) which is defined as follows:
Kendall’s τ = # of pairs where the rankings agree−# of pairs where the rankings disagree
# pairs .
This coefficient lies between −1 and 1, where 1 denotes a high correlation between the two set of rankings. Table 2 shows that the Kendall’s τ coefficient between our network criticality measure and the generalization error is higher than all other complexity measures that we compared to.
We find that our measure correctly ranks the generalization performance of the networks – ResNet18, ResNet101, VGG16 and FCN (I). It fails to correctly identify the correct rank of ResNet34, ResNet50, DenseNet121, VGG11 and FCN (II).
We also repeat the above experiment for the same networks trained on the CIFAR100 dataset (see Table 3). We note that network criticality measure correctly predicts the ranking of generalization performance for ResNet101, ResNet34, ResNet18, ResNet50, VGG16 and the 3-layer fully connected networks (FCN). We find that the generalization error of ResNet101 is the lowest which is correctly captured only by our complexity measure and not by any other measure. Further, the Kendall’s τ correlation coefficient between the ranking based of the generalization error and our network criticality measure is 0.55 which is again higher than this coefficient for any other complexity measure. Our measure only fails to capture the ranking of VGG11 and DenseNet121 relative to the other DNN architectures and the relative ranking between FCN (I) and FCN (II).
We also perform an additional set of experiments in Appendix D where we calculate these complexity measures on four ResNet18 networks with changing channel widths.

heading: 5 CONCLUSION
text: In this paper, we studied the module criticality phenomenon and proposed a complexity measure based on module criticality that is able to correctly predict the superior performance of some DNN
architectures over others, for a specific task. We believe module criticality can be used as a road-map for designing new task-specific architectures. Proposing new regularizers that improve generalization performance by bounding criticality or spreading it among various modules of the network is an exciting direction for future work. Our measure could also be potentially used in architecture search where we could calculate this score over the training set to select architectures that generalize well on the unseen test set.

heading: ACKNOWLEDGEMENTS
text: We would like to thank Samy Bengio and Chiyuan Zhang for valuable conversations, and Yann Dauphin for his help with the implementation of Fixup initialization. We would also like to thank Chiyuan Zhang for sharing the code for the paper “Are all layers created equal?” Part of this work was performed while the author NC was an intern at Google AI, Brain team.

heading: A PROOF OF THEREORM 3.2
text: We start by stating the PAC-Bayes theorem which bounds the generalization error of any posterior distribution Q on parameters Θ that can be reached using the training set given a prior distribution P on parameters that should be chosen in advance and before observing the training set. Throughout this section given two scalar p, q ∈ [0, 1] let KL(p||q) denote the KL divergence between two Bernoulli distributions with success probabilities p and q respectively. Theorem A.1 (McAllester (1999)). For any data distribution D, number of samples m ∈ N, training set Sm ∼ D, and prior distribution P on parameters Θ, posterior distribution Q, for any 0 < δ, with probability 1− δ over the draw of training data we have that
KL ( EΘ∼Q[LS(fΘ)] ∣∣∣∣∣∣∣∣EΘ∼Q[LD(fΘ)]) ≤ KL(Q||P ) + log mδm− 1 where KL is the Kullback-Leibler (KL) divergence (Kullback & Leibler, 1951).
Following Dziugaite & Roy (2017), we use the inequality KL−1(q|c) = sup {p ∈ [0, 1] : KL(q||p) ≤ c} ≤ q + √ c/2 to achieve a simple bound on the test error:
EΘ∼Q[LD(fΘ)] ≤ KL−1 ( EΘ∼Q[LS(fΘ)] ∣∣∣∣KL(Q||P ) + log mδm− 1 )
≤ EΘ∼Q[LS(fΘ)] +
√ KL(Q||P ) + log mδ
2(m− 1) .
The intuition from Figure 4 suggests that moving the parameters of each module as close as possible to the initialization value before harming performance. For such αi, we can then define the posterior Qi for module i to be a Gaussian distribution centered at θαi with covariance matrix σ 2 i I , that is, as if we have additive noise ui ∼ N (0, σiI). We use Θα to refer to the case where all θi are replaced with θαii and matrix U includes all ui. Then the training loss term can be decomposed as
EΘ∼Q[LS(fΘ)] = Eui∼N (0,σiI)[LS(fΘα+U )] ≤ LS(fΘF ) + ∣∣∣∣Eui∼N (0,σiI)[LS(fΘα+U )]− LS(fΘF )∣∣∣∣,
where the second term on the right hand side of the inequality captures the flatness of the point Θα by adding Gaussian noise and measuring the change in the loss. Therefore, searching over the posterior corresponds to finding a flat solution in the valley that connects the initial and final points. Next, we use this intuition to prove a generalization bound based on module criticality.
First, we express the value of the Õ(1) term in Theorem 3.2 , which is equal to as follows:
= ∑ i log
( 7m+ 2 log ( ki
kiσ2i + α 2 i ∥∥θFi − θ0i ∥∥2Fr )) . (3)
Now we proceed with the proof.
The KL-divergence between two k-dimensional Gaussian distributions is given by the formula:
KL(N (µ1,Σ1)||N (µP ,ΣP )) = 1
2
[ tr ( Σ−12 Σ1 ) + (µ2 − µ1)>Σ−12 (µ2 − µ1)− k + ln(
det Σ2 det Σ1 )
] .
The above equation can be further simplified for Gaussian distributions with diagonal covariance matrices. Let the prior P be a Gaussian distribution such that for each module i, the distribution is N (θ0i , σ2P,iI) and let the posterior Q be a Gaussian distribution such that for each module i, the distribution is N ((1− α)θ0i + αθFi ), σ2Q,iI). We can then write the KL-divergence KL(Q||P ) as
KL(Q||P ) = 1 2 ∑ i
[ kiσ 2 Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr σ2P,i − ki + ki log ( σ2P,i σ2Q,i )] . (4)
Since prior should be decided before observing the training set, we are not allowed to optimize for σP,i directly. However, one can optimize for σP,i over a pre-defined set of values and use a union bound argument to get the generalization bound for the best σP,i in that set. We use a covering approach suggested by Langford & Caruana (2002). For b, > 0, if one chooses the variance of prior to be exp(− j + b) for j ∈ N such that for each j the bound holds with probability1 − 6π2j2 , then all bounds hold with probability 1− ∑ j∈N 6 π2j2 = 1− δ. We can apply the same idea to every
module such that the bound holds with probability 1− δ ∏d i=1
6 π2j2i . If we choose σ2Q,i ≤ 1 then we have σP,i ≤ exp ( 4m ki + 1 ) . Otherwise, the bound holds since the right hand side is greater than one. Given (from Equation 3) ≥ 0, if we choose σ2P,i to have the form exp ( 4m−ji ki + 1 ) , for some integer ji, we can always find choose ji such that
kiσ 2 Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr ≤ kiσ2P,i ≤ exp(1/ki)(kiσ2Q,i + α2i ∥∥θFi − θ0i ∥∥2Fr) . (5) Therefore, the KL-divergence can be bounded as
KL(Q||P ) ≤ 1 2 ∑ i
[ ki kiσ 2 Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr kiσ2Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr − ki + ki log ( σ2P,i σ2Q,i )]
= 1
2 ∑ i
[ ki log ( σ2P,i σ2Q,i )]
≤ 1 2 ∑ i ki log
exp(1/ki) ( kiσ 2 Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr) kiσ2Q,i  ≤ 1
2 ∑ i ki log
exp(1/ki) ( kiσ 2 Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr) kiσ2Q,i  ≤ 1
2 ∑ i 1 + ki log
( 1 + α2i ∥∥θFi − θ0i ∥∥2Fr kiσ2Q,i ) .
Note that in order to achieve the inequality in Equation 5, ji should be chosen as
ji =
⌊ 4m
ki + 1 + log
( ki
kiσ2Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr )⌋ ≤ 5m+ log ( ki kiσ2Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr ) .
Given that each such bound should hold with probability 1− δ ∏d i=1
6 π2j2i , the log term in the bound can be written as
log m δ + ∑ i log(π2j2i /6) ≤ log m δ + 2 ∑ i log
( 7m+ 2 log ( ki
kiσ2Q,i + α 2 i ∥∥θFi − θ0i ∥∥2Fr )) .
Putting everything together proves the theorem statement.

heading: B A DETERMINISTIC GENERALIZATION BOUND FOR CONVOLUTIONAL NETWORKS
text: We start by stating a generalization bound given by Neyshabur et al. (2018) with a slight improvement in the constants.
Lemma B.1 (Neyshabur et al. (2018)). Let fΘ : X → RC be any predictor function with parameters Θ and P be a prior distribution on parameters Θ. Then for any γ,m, δ > 0, with probability 1− δ
over the training set S of size m, for any parameter Θ and any perturbation distribution Q over parameters such that PU∼Q [ maxx∈X |fΘ+U (x)− fΘ(x)| ≤ γ4 ] ≥ 12 , we have
LD(fΘ) ≤ LS,γ(fΘ) +
√ 2KL(Θ + U ||P ) + 1 + log mδ
2(m− 1) .
The lemma above gives a data-independent deterministic bound which depends on the maximum change of the output function over the domain after a perturbation. We combine Lemma B.1 with Theorem 3.2 and prove a bound on the perturbation which leads to the following theorem. Theorem B.2. Let input x be an N × N image whose norm is bounded by B, fΘ : X → RC be the predictor function with parameters Θ which is a DNN of depth d made of convolutional blocks. Then for any margin γ, sample size m, δ > 0, with probability 1 − δ over the training set S, any parameter Θ and any αi > 0 such that maxx∈X |fΘ(x)− fΘα(x)| ≤ γ8 , we have
LD(fΘ) ≤ LS,γ(fΘ)+
√√√√√∑di=1 ki log ( 1 + [32edBαi‖θFi −θ0i‖Fr ∏ i6=j‖θαi ‖2 √ log(4dN2)]2 ciγ2 ) + log ( m δ ) + Õ(1)
m− 1 ,
(6) where ki is the number of parameters in module i. For example, for a convolution module with kernel size qi × qi and number of output channels ci, ki = q2i ci−1ci.
Proof. First, we express the value of the Õ(1) term in Theorem B.2 , which is equal to 2 as follows:
2 = 1 + ∑ i log 7m+ 2 log  ki kiγ2/ ( 16e ∏ j 6=i ‖θαi ‖2 log(4dN2) )2 + α2i ∥∥θFi − θ0i ∥∥2Fr   .
(7)
We note that for any Θ,Θ′, if maxx∈X ‖fΘ(X)− fΘ‖∞ ≤ γ/2 then L(fΘ) ≤ Lγ(f ′Θ). The reason is that the output for each class can change by at most γ/2 and therefore the label can only change for the data points that are within γ of the margin.
We start using the assumptions on the perturbation bound. Combining the results from Theorem B.1 and Theorem 3.2, we can get the following bound.
LD(fΘF ) ≤ LD, γ4 (fΘα) (8)
≤ LS, 3γ4 (fΘα) +
√√√√√ 12 ∑di=1 ki log ( 1 + α2i‖θFi −θ0i‖2Fr kiσ2i ) + log mδ + 2
m− 1
≤ LS,γ(fΘα) +
√√√√√ 12 ∑di=1 ki log ( 1 + α2i‖θFi −θ0i‖2Fr kiσ2i ) + log mδ + 2
m− 1 where 2 is given above in Equation 7.
Therefore, it suffices to find the value of σi under which the assumption on norm of perturbation in function space holds and then simplify the following upper bound given the desired value of σi.
In order to find the desired value of σi we use the following two lemmas. First we adopt the perturbation lemma by Neyshabur et al. (2018) to bound the change in the output a network based on the magnitude of the perturbation:
Lemma B.3 (Neyshabur et al. (2018)). Let norm of input x be bounded by B. For any B > 0, let fΘ : X → RC be a neural network with ReLU activations and depth d. Then for any Θ, x ∈ X , and any perturbation U s.t. ‖ui‖2 ≤ ‖θi‖2, the change in the output of the network can be bounded as follows
‖fΘ+U − fΘ‖2 ≤ eB d∏ i=1 ‖θi‖2 d∑ j=1 ‖ui‖2 ‖θi‖2 . (9)
We next use the following lemma by Pitas et al. (2017) that bounds the magnitude of the Gaussian perturbation ui for each convolutional module based on the standard deviation of the perturbation.
Lemma B.4 (Pitas et al. (2017)). Let ui be a Gaussian perturbation for each module i of a convolutional model. Let N be the image size, qi , ci be the kernel size and the number of output channels at module i respectively. We have that
P [ ‖ui‖2 ≥ σi ( qi(2 √ ci) + t )] ≤ 2N2e − t2 2q2 i .
The lemma above suggests that by taking union bounds over all modules, we can ensure that with probability 1/2 we have that for any module i, the following upper bound on the spectral norm of the perturbation holds.
‖ui‖2 ≤ σiqi(2 √ ci + √ 2 log(4dN2)) ≤ 2σiqi( √ ci + √ log(4dN2)) ≤ 4σiqi √ ci log(4dN2).
Combining this with perturbation bound in Equation 9, we have that
‖fΘα+U − fΘα‖2 ≤ eB d∑ i=1 ‖ui‖2 d∏ j 6=i ∥∥θαj ∥∥2 ≤ 4eB d∑ i=1 σiqi √ ci log(4dN2) d∏ j 6=i ‖θαi ‖2 ≤ γ 8 ,
where the last inequality can be achieved with
σi = γ 32edB ∏d i=1 ‖θαi ‖2 qi √ ci log(4dN2) . (10)
Therefore, this value for σi ensures the assumption on norm of perturbation in function space in Theorem B.2 holds and hence completes the proof.
Moreover, we show how we get the value of 2, by showing the simplification from inserting the value for σi from Equation 10 as follows.
log ( 1 + α2i ∥∥θFi − θ0i ∥∥2Fr kiσ2i ) ≤ log ( 1 + [ αi ∥∥θFi − θ0i ∥∥Fr]2 q2i c 2 iσ 2 i )
≤ log ( 1 + [32edBαi ∥∥θFi − θ0i ∥∥Fr∏di=1 ‖θαi ‖2 qi√ci log(4dN2)]2
q2i c 2 i γ 2
)
= log ( 1 + [32edBαi ∥∥θFi − θ0i ∥∥Fr∏di=1 ‖θαi ‖2√log(4dN2)]2
ciγ2
) .
Then, 2 can also be simplified as follows.
2 = 1 + ∑ i log
( 7m+ 2 log ( ki
kiσ2i + α 2 i ∥∥θFi − θ0i ∥∥2Fr ))
≤ 1 + ∑ i log 7m+ 2 log  ki kiγ2/ ( 32edB ∏d i=1 ‖θαi ‖2 qi √ ci log(4dN2) )2 + α2i ∥∥θFi − θ0i ∥∥2Fr   ≤ 1 + ∑ i log 7m+ 2 log  ki γ2/ ( 32edB ∏d i=1 ‖θαi ‖2 √ log(4dN2) )2 + α2i ∥∥θFi − θ0i ∥∥2Fr   ≤ 1 + ∑ i log ( 7m+ 2 log(ki)− 4 log ( αi ∥∥θFi − θ0i ∥∥Fr + γ/32edB d∏ i=1 ‖θαi ‖2 √ log(4dN2) )) .

heading: C DETAILS ON EXPERIMENTAL SET-UP
text: For all our experiments, we use the CIFAR10 and CIFAR100 datasets. To train our networks we used Stochastic Gradient Descent (SGD) with momentum 0.9 to minimize multi-class cross-entropy loss. On CIFAR10 each model is trained until the cross-entropy loss on the training dataset falls below 0.19. While on CIFAR100 each model is trained until the cross-entropy loss on the training dataset falls below 0.25. The ResNets, DenseNets and VGGs were trained using a stage-wise constant learning rate scheduling with a starting learning rate of 0.1 and with a decrease by a multiplicative factor of 0.2 every 60 epochs. FCN was trained with an initial learning rate of 0.1 with a decrease by a multiplicative factor of 0.2 every 200 epochs. Batch size of 128 was used for all models and weight decay with factor 5e-4 was used to train all networks.
We mainly study three types of neural network architectures:
• Fully Connected Networks (FCNs): The FCNs consist of 2 fully connected layers. FCN (I) contains 5000 and 1000 hidden units respectively while FCN (II) contains 10000 and 2000 hidden units respectively. Each of these hidden layers is followed by a batch normalization layer and a ReLU activation. The final output layer (that follows the ReLU activation in the second layer) has an output dimension of 10 or 100 (number of classes).
• VGGs: Architectures by Simonyan & Zisserman (2015) that consists of multiple convolutional layers, followed by multiple fully connected layers and a final classifier layer (with output dimension 10 or 100). We study the VGG with 11 and 16 layers.
• DenseNets: Architectures by Huang et al. (2017) that consists of multiple convolutional layers, followed by a final classifier layer (with output dimension 10 or 100). We study the DenseNet with 121 layers.
• ResNets: Architectures used are ResNets V1 (He et al., 2016). All convolutional layers (except downsample convolutional layers) have kernel size 3×3 with stride 1. Downsample convolutions have stride 2. All the ResNets have five stages (0-4) where each stage has multiple residual/downsample blocks. These stages are followed by a maxpool layer and a final linear layer. Here are further details about the ResNets used in the paper:
– ResNet18: ResNet18 architechtures studied in the paper have 1 convolutional layer in Stage 0 (64 ouput channels), Stage 1 has 2 residual blocks (64 output channels), Stage 2 has one downsample block and one residual block (128 output channels), Stage 3 has one downsample block and one residual block (256 output channels) and Stage 4 again has one downsample block and a residual block (512 output channels).
– ResNet34: ResNet34 architectures in this paper have 5 stages. Stage 0 has 1 convolutional layer with 64 output channels followed by a ReLU activation. Stage 1 has 3 residual blocks (64 output channels), Stage 2 has 1 downsample block and 3 residual blocks (128 output channels), Stage 3 has 1 downsample block and 5 residual blocks (256 output channels) and, Stage 4 has 1 downsample block and 2 residual blocks (512 output channels).
– ResNet50: ResNet50 architectures in this paper again have 5 stages. Stage 0 has 1 convolutional layer with 64 output channels followed by a ReLU activation. Stage 1 has 1 downsample block and 2 residual blocks (256 output channels), Stage 2 has 1 downsample block and 3 residual blocks (512 output channels), Stage 3 has 1 downsample block and 5 residual blocks (1024 output channels) and, Stage 4 has 1 downsample block and 2 residual blocks (2048 output channels).
– ResNet101: ResNet101 architectures in this paper again have 5 stages. Stage 0 has 1 convolutional layer with 64 output channels followed by a ReLU activation. Stage 1 has 1 downsample block and 2 residual blocks (256 output channels), Stage 2 has 1 downsample block and 3 residual blocks (512 output channels), Stage 3 has 1 downsample block and 22 residual blocks (1024 output channels) and, Stage 4 has 1 downsample block and 2 residual blocks (2048 output channels).
The ResNets, DenseNets and VGGs in the paper are trained without batch normalization.
During training, images are padded with 4 pixels of zeros on all sides, then randomly flipped (horizontally) and cropped. Global mean and standard deviation are computed on all training images
and applied to normalize the inputs. While training a ResNet18 on the CIFAR10 dataset with 20% of the labels randomly corrupted, we do not augment the training set with images that are randomly flipped and cropped. We also do not use weight decay during training these networks.

heading: D ADDITIONAL EXPERIMENTS
text: In these set of experiments we compare the generalization performance of four ResNet18 architectures where we vary the number of output channels in each stage. In the ResNet18 (1x width) network the number of output channels are 16, 16, 32, 64, 128 in the five stages respectively. The other ResNet18s have their output channels scaled by factors of 2,4 and 8 in each stage. Table 4 summarizes our results for networks trained on the CIFAR10 dataset and Table 5 summarizes are results for networks trained on the CIFAR100 dataset.
We find that separating these networks based on a complexity measure is a much more challenging as these four networks differ by just the channel widths at the different stages. We find that on this task all complexity measures that we studied (including ours) does poorly. It is an interesting question for future research to see if our complexity measure can be refined to separate these networks as well.
E FIGURES

Review summaries:
The paper introduces concept of "module criticality" to understand the role played by several modules in a model and how this affects the generalization of the models. This is quite an important problem to study as this helps to develop better understanding of the current architectures and potentially reduce their size without suffering the accuracy drop. This is a great theoretical contribution.