The intriguing role of module criticality in the generalization of deep networks
Authors: Niladri Chatterji
Behnam Neyshabur
Hanie Sedghi
Venue:  (2020)
Keywords: Module Criticality Phenomenon
Complexity Measure
Deep Learning

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.

Published as a conference paper at ICLR 2020
THE INTRIGUING ROLE OF MODULE CRITICALITY IN THE GENERALIZATION OF DEEP NETWORKS
Niladri S. Chatterji
Behnam Neyshabur
Hanie Sedghi
University of California, Berkeley chatterji @ berkeley.edu
Google
Google
neyshabur google.com
hsedghi @ google.com
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, re- sults in large drop in the network's performance. Our analysis reveals interesting properties of the loss landscape which leads us to propose complexity measure, called module criticality, based on the shape of the valleys that connect the ini- tial and final values of the module parameters. We formulate how generalization relates the module criticality, and show that this measure is able explain the superior generalization performance of some architectures over others, whereas, earlier measures fail to do so.
INTRODUCTION
Neural networks have had tremendous practical impact in various domains such revolutionizing many tasks in computer vision, speech and natural language processing. However, many aspects of their design and analysis have remained mysterious 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 unified view that provides a complete and satisfactory answer. In order to attain a unified view on superiority one architecture over another terms of generalization performance, need to come up with measure that effectively captures this.
Analyzing the generalization behavior 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 either based on covering number and Rademacher complexity (Neyshabur et al., 2015; Bartlett et al., 2017; Neyshabur 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 have an exponential dependency on the depth and focus on fully connected networks. More re- cently, Long & Sedghi (2019) provided generalization bounds for convolutional neural networks (CNNs) and fully connected networks used in practice and their bounds have linear dependency the depth.
al., 2019; Long & Sedghi, 2019). Earlier results
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 perfor- mance of candidate architectures for 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 ResNets.

---

Published as a conference paper at ICLR 2020
non-critical - 0/1 loss
critical - 0/1 loss
0.4
0.96
0.4
0.96
0.80
0.80
0.2
0.2
00
0
0.64
0
0 F
0.64
0.0
0.48
0.0
0.48
-0.2
0.32
-0.2
0.32
0.16
-0.4
-0.4
0.00
0.00
0.2
0.4
0.6
0.8
1.0
0.2
0,4
0.6
0.8
1.0
Figure 1: Module Criticality: Loss values in the valleys that connect the initial weights 6 to the final weights R of non-critical (left) and a critical (right) module the ResNet18 architecture. Given a ball with radius (length of the red line), module criticality can be defined as how far one can push the ball in the valley towards initialization (length of the white dashed line) divided by the radius Hence, non-critical modules are the ones with wide valley connecting the initial weight vector to the final one whereas in critical modules, the valley either becomes too sharp or the loss values start increase when the ball comes too close to the initial weight. The axis is simply chosen to be parallel to oF d and the axis is a compact representation of all other dimensions generated by adding Gaussian noise to the points the convex combination of 6 and of and evaluating the loss. The sign on the axis is decided based on the sign the inner product of the noise 8
Every DNN architecture is a computation graph where each node is module'. We are interested in understanding how different modules in the network interact with each other and influence gen- eralization 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 characteristies 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 performance is negligible. They show that various conventional DNN architectures exhibit this phenomenon.
Let informally define what mean by the measure "module criticality" (see Figure I). For each module, we move on a line from its final trained value to its initialization value (convex combination path) while keeping all other modules fixed their trained value. Then measure the performance drop. Let 0% (1 a a)00 + a0..' a E [0, 1] be the convex combination between initial weights 80 and the final weights 8." at module i, where a; is the minimum value between and 1 when performance (train error) of the network drops by at most threshold value If a; 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, also wish to incorporate the robustness 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 formal definition).
In this paper, we seek study this phenomenon 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
A module is node in the computation graph that has incoming edges from other modules and outgoing edges other nodes and performs linear transformation module definition equivalent to definition of layer.
inputs. For layered model such as VGG,
A convex combination of two points is linear combination of them where the enefficients are non-negative and to 1, Every combination of two points lies on the line segment between the two.
2

---

Published as a conference paper at ICLR 2020
Training Epochs
0
1
0.7
0.6
Train Error
3
0.5
20
0.4
100
0.3
0.2
190
0.1
0.0
alkz.cooo
com2
alled convo
2o0s bokk cone
wxx..com
Paanos 2222 ghs 22
Anuos EMIQ ghs 2nuo3
linear tha bay
22222 2012
"orooo +62s
dep bokk 2ooe
blkk
bik
251
EMI'
IMI p
Figure 2: Analysis of rewinding modules to initialization for the ResNet- 18 architecture. Each row represents module in Res and each column represents a particular training epoch which this module is rewound to. The difference from analysis of Zhang et al. (2019a) is that we rewind each module, whereas Zhang et al. (2019a) rewind the entire ResNet blocks.
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 PAC- PAC-Bayesian anal- ysis. In Section we show that the overall network criticality measure and not the number of critical modules that controls generalization. If the network criticality measure is smaller for architecture, it has better generalization performance. In Section 4, we demonstrate through various experiments that proposed measure is able to distinguish between different network architec- tures 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]] respectively. For n E N. we use [n] denote the set f1....,n) n.. Let Ls(()) be the loss of function f on the training set S with m samples. We are mainly interested in classification task where Ls( = m (G [x)(x] maxj S [x]]. For any 7 O, also define margin loss (f, - E ge) [x) < + maxj [x]][i]] Let LD())
be the loss of function on population data distribution D defined similar to (f)). We will denote the function parameterized by θ by fe.
TOWARDS UNDERSTANDING MODULE CRITICALITY
2.1 SETTING
DNN architecture is directed acyclic computation graph which may or may not be layered. In order to have unifying definition between different architectures, we use the notion of a "module". module is a node the computation graph that has incoming edges from other modules and outgoing edges to other nodes, and performs linear transformation model such as a VGG, a module equivalent to layer. On the other hand, in a ResNet some modules are parallel to each other. For example, downsample module and the concatenation of two convolutional modules in ResNet18- 1-v1 architecture. Note that, similar to conventional definitions the non- lineariyyy (such ReLU) is not part of the module.
its inputs. For a layered
Let θ = (......, 0z) correspond to all parameters a DNN with d modules, where 8, refers to the weight matrix (or operator matrix in case of convolution) at module i and 06. 0f refer the value
3

---

Published as a conference paper at ICLR 2020
Sjage2.block2 wttp:/bppccbbbby.ogog.....
Stage2 wwwpbbpww.......c.........
180000
200000
Rndinit
Rodinit
140000
150000
120000
100000
100000
80000
30
60000
110
40000
190
50000
190
0
0
(a) A non-critical module
(b) critical module
Figure 3: Spectrum of non-critical and a critical module during different epochs of training.
of weight matrix at initialization and the end of training, respectively. For sequential architectures, is equal to the depth of the network but that is not necessarily true for general architectures such ResNet.
2.2 ROBUSTNESS TO REWINDING
Consider the following perturbation to a trained network at some training epoch as considered by Zhang 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 each module (each convolutional mod-
rewind the whole ResNet block once, whereas, ule) separately. This rewind analysis is shown in Figure 2 for ResNet18-vl. Each column represents module in ResNet18-v/ and each row represents a particular training epoch 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 "Stagg2. conv" 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 differ- ent performance criteria (train loss, train error and test error) and observe that they exhibit a similar trend.
A stable phenomena: The plots in Figure capture 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 the appendix illustrate the pattern when we changed the initializa- tion to Fixup (Zhang et al., 90),,,, and when we replace SGD with Adam (Kingma & Ba, 2014) respectively.
2.3 WHAT MEASURES FAIL TO DISTINGUISH CRITICAL LAYERS
Spectrum of weight matrices: 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 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.
4

---

Published as a conference paper at ICLR 2020
Change in the activation patterns: We investigated the change in the activation patterns of 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 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 critical module, the similarity between the activation patterns between the two networks degrades gradually rather than abruptly. See Figure 9 in Appendix E.
3 GENERALIZATION BOUNDS BASED ON MODULE CRITICALITY
Our goal is to understand criticality and how it affects the generalization performance of a DNN. Inspired by the rewind initialization experiments of Zhang 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 8 with 0 (1 a + att" a E [0. 1], and keep all other layers fixed. Then we look at the effect of this perturbation on the performance of the network.
Figure depicts how the train error, test error and train loss change as we decrease the value of a in O when i refers to critical module (yellow dashdot curve), 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 in- terested in also accounting for the width of the valley 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 σ, around each point in this path. PAC-Bayesian analysis, looks for 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, are interested in moving from the final value 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 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 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
test data.
Definition 3.1 (Module and Network Criticality). Given an 0 and network fe, we define the module criticality for module as follows:
a
oF
Hr.e(fe) = min
:uw..oo Se
(I)
o²
also define the network criticality as the sum of the module criticality over modules of the net- work:
Mc(Se) Σ (fe)
Here, Ls denotes the empirical zero-one loss over the training set, fop,0F, is the DNN's function value where weight matrix corresponding to 4th module is replaced by 0e and all other modules are fixed at their values in the end of training, OF,, 0r = (1 a) + ad, where 00 is the value of the weight matrix at initialization and 6." is the trained value.
Intuitively, network criticality measure is sum of module criticalities. This also theoretically derived using the analysis below.
3.1 A PAC-BAYESIAN GENERALIZATION BOUND
We attempt to understand the relationship between module criticality, and generalization by deriv- ing generalization bound using the PAC-Bayesian framework (McAllester, 1999). Given prior
5

---

Published as a conference paper at ICLR 2020
Network Non-Critical
0.75
0.8
0.00
Train Error
Test Error
0.5
Train Loss
Critical
0.50
0.01
0.25
0.2
30P
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0
move convex combination path from final initial value of modules. We find that along this path the training error (as well as test error and train loss) increases monotonically from the final weights initial weights. The blue (solid) curve is when we replace all the parameters in the network by the convex combination between their initial and final value simultaneously. the red (dashed) curve corresponds to moving on the convex path for single (non-critical) layer and the yellow (dashdot) curve corresponds to moving on the convex path for critical layer in ResNet- 18 architecture.
Figure 4: Performance degradation
distribution over the parameters that is picked in advance before observing training set, 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 possible to the initialization value before harming the performance. For such a; we can then define the posterior Q; for module i to be Gaussian distribution centered at 0f with covariance matrix of that is, as if we have additive noise u, ~ (000 a/). We use Oa to refer to the case where all parameters 0, replaced with 8. and matrix U includes all the noise ut Then the following theorem holds.
Theorem 3.2. For any data distribution D, number of samples m N for any 0 < 8 < 1, for any < 0r S and any 0 S ag S with probability over the choice of the training set Sm D the following generalization bound holds:
Ee= k log (1+ afllof -oeeo + log (3) Ö(1)
k2o
Eo(LD(fe=tu)l S Eu
m
where k; is the number of parameters in module i. For example, for a convolution module with kernel size q 9: and number of output channels c, k; c4-1-
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 not exactly a generalization bound on original network but rather on a perturbed network. We conjecture that for almost any realistic distribution D, any random O any OF achieved by known gradient based optimization algorithms, any 0 I a S 1 and any o 2 O, the test error does not improve by taking convex combination of parameters and their initial values followed by Gaussian perturbation. Therefore, we have that (D(err) S EEy (fer+t) ]]. 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 E N. For any e > 0, for any 0 < 8, if Lp(fe=) I BuLLn(fer )] where u, (0,,o;1), then with probability 1 - 8 over the choice of the training set Sm ~ D, the following generalization bound holds
AP((Se)
log (T) O(1)
Co(se) S (
m -
where v¹ (fe) is calculated as follows:
0 af 10.F5
('(fe) = -min
Ey Ec
050,051
of

---

Published as a conference paper at ICLR 2020
Note that the above bound uses slightly different notion of network criticality compared to Def- inition 3.1, as the bound requires finding a and a values simultaneously for all modules, whereas, Definition 3.1 allows to decouple the search over and σ.
Deterministic generalization bound for convolutional networks Although PAC-Bayest -ayyssan bounds are data-dependent and hence numerically superior, they provide less insight about the un- derlying 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, de- terministic 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, build on the results of Theorem 3.2 present norm-based deterministic bound using module criticality.
In this section, 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 correctly rank the generalization performance of different architectures better than measures proposed earlier.
EXPERIMENTS
We perform several experiments to compare our network criticality measure to earlier complexity measures in the literature. Our experiments are performed on the CIFARIO and CIFARI00 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 σ; and a; per Definition 3.1. In our experiments, each module is 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 a linear transformation.
Table Quantities of Interest
Generalization Error (GE)
LD(fe) Ls(fe)
Product of Frobenius Norms (PFN)
11, 00 Fr
Product of Spectral Norms (PSN)
11, 1
Distance to Initialization (Dtl)
>. |00 0
Number of Parameters (NoP)
Total number of parameters in the network
Sum Spectral Norms (SoSP)
Total number of parameters x (> 10 8 ll))
PAC Bayes (at error threshold 0.1)
>, 0 TTT /////
Network Criticality Measure
E. of 12/
(at error threshold 0.1)
First, as a sanity check we use our complexity measure (lower is better) to compare between a ResNet18 trained on true labels and ResNet18 trained on data where 20% of the labels are ran- domly 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, Table 2 we compare the generalization performance of several conventional DNN architec- tures trained on the CIFAR10 dataset. There is particular ranking of the networks based on their generalization error and it is desirable for complexity measure to capture this ranking. Therefore, 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 ' T # of pairs where the rankings agree # of pairs where the rankings disagree # pairs
7

---

Published as a conference paper at ICLR 2020
Network Criticality Measure [107]
True Labels
Netwark Criticality Measure [109]
ResNet18
20
20% Corrupted Labels
ResNet34
100
VGG16
15
FCN (I)
60
40
200
0.025
0.050
0.075
0,100
0.125
0.150
0.025
0.125
0.150
Error Threshold
Error Threshold
(a) Comparing ResNet18 trained true la- bels vs. corrupted labels
(b) Comparing ResNet18, ResNet34, VGG16 and FCN (I).
Figure 5: Network criticality function of error threshold for networks trained on CIFARIO.
This coefficient lies between -1 1 and 1, where denotes high correlation between the two set of rankings. Table 2 shows that the Kendall's 5 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 CIFARI00 dataset (see Table 3). We note that network criticality measure correctly predicts the ranking of gemeralization performance for ResNet101, ResNet34, ResNet18, ResNet50, VGG16 and the 3-layer fully con- nected networks (FCN). We find that the generalization error ResNet101 the lowest which correctly captured only by complexity measure and not by any other measure. Further, the Kendall's 's correlation coefficient between the ranking based the generalization error and our net- work 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 VGG1 and DenseNet12 21 relative to the other DNN architectures and the relative ranking between FCN (I) and FCN (II).
Table 2: Measuring complexity of different architectures trained on CIFARIO.
Network GE
PFN
PSN
Dtl
NoP
SoSP
PAC Bayes
Net. Criticality
ResNet18
4.61%
1e22
4e14
3430
1.1e7
1.3e9
6.9e5
2.2e5
ResNet34
6.3%
2e37
3e24
4768
2.1e7
3.7e9
9.1e5
1.7e5
ResNet50
6.6%
4e56
4e20
10018
2.3e7
3.3e9
1.6e6
1.8e5
ResNet101
6.4%
8e110
3e32
18730
4.2e7
9.8e9
2,8e6
6.3e5
DenseNet121
7.8%
2e129
7e42
21359
6.8e6
2.0e9
1,2e6
4.1e5
VGGI1
8.51%
lell
le6
2106
2.8c7
1.3c9
1,0e6
2.8e5
VGG16
7.47%
5c15
2c8
2341
3.4e7
2.1c9
1.2c6
2.70e5
FCN (I)
29.83%
3e20
2e7
75221
2.0e7
4.6e8
9.0e6
5.7e6
Kendall's 26.45%
FCN (II)
3e21
1e7
81258
5.0e7
2.0e9
9.5e6
6.2e6
-0.22
-0.33
0.38
0.16
-0.53
0.42
0.55
We also perform additional set of experiments in Appendix D where calculate these complex- ity measures on four ResNet18 networks with changing channel widths.
5 CONCLUSION
In this paper, we studied the module criticality phenomenon and proposed a complexity measure based module criticality that is able correctly predict the superior performance of some DNN
8

---

Published as a conference paper at ICLR 2020
Table Mcasuring complexity of different architectures trained on CIFARI00.
Network
GE
PFN
PSN
Dtl
NoP
SoSP
PAC Bayes
Net. Criticality
ResNet18
30.6%
2e22
9e14
4855
1.1e7
1.4e9
3,4e6
1.6e6
ResNet34
29.3%
1e37
1e22
6017
2.1e7
3.6e9
6.7e6
1.4e6
ResNer50
31.1%
7e57
9e23
12715
2.3e7
4.1e9
5.8e6
1.9e6
ResNet101
25.5%
2e112
5e36
21233
4.2e7
1.1e10
6.4e6
1.3e6
DenseNet121
35.3%
1c131
1c49
23702
6.9c6
2.4c9
4.5c6
2.6e6
VGGI1
43.5%
6c13
1e8
5059
2.8c7
1.7c9
3.6e6
2.4c5
VGG16
32.69%
8e19
4e12
7010
3.4e7
3.6e9
8.1e6
4.6e6
FCN (I)
57.59%
8e27
1e10
246636
2.0e7
8.3e8
2.4e7
1.2e7
FCN (II)
53.17%
le29
le10
329296
5.0e7
3.7e9
3.3e7
2.3e7
Kendall's
-0.27
-0.47
0.33
0
-0.42
0.22
0.55
architectures others, for specific task. We believe module criticality be used as road-map for designing new task-specific architectures. Proposing new regularizers that improve generaliza- tion performance by bounding criticality spreading among various modules of the network 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 select architectures that generalize well on the unseen test set.
ACKNOWLEDGEMENTS
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 Google AI, Brain team.
REFERENCES
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. In Proceedings of the International Conference Ma- chine Learning, 2018.
Peter Bartlett.
The sample complexity of pattern classification with neural networks: the size of the weights is more important than the size of the network. IEEE Transactions on Information Theory, 44(2):525-532 1998.
Peter Bartlett, Dylan Foster, and Matus Telgarsky, Spectrally-normmaizze normalizzed margin bounds for neural networks. In Proceedings of the Advances in Neural Information Processing Systems, 2017.
Eric Baum and David Haussler. What size net gives valid generalization? In Proceedings of the Advances in Neural Information Processing Systems. 1989.
Gintare Karolina Dziugaite and Daniel Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of Uncertainity in Artificial Intelligence, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE International Conference Computer Vision, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings the IEEE Conference Computer Vision and Pattern Recognition, 2016.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks, In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2017.

---

Published as a conference paper at ICLR 2020
Maurice Kendall. A new measure of rank correlation. Biometrika, 30(1/2):81-93, 1938.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6988, 2014.
Simon Komblith, Mohammad Norouzi, Honglak Lec, and Geoffrey Hinton. Similarity of neural network representations revisited. In Proceedings of the International Conference Machine Learning, 2019.
Solomon
Kullback and Richard Leibler. On information and sufficiency. The Annals of Mathemati- cal Statistics, 22( (21))99--86 1951.
John Langford and Rich Caruana. (not) bounding the true error. In Proceedings of the Advances in Neural Information Processing Systems, 2002.
Philip Long and Hanie Sedghi. Generalization bounds for deep convolutional neural networks. arXiv preprint ICG5:I900..26000 2019.
David McAllester. PAC-bayesian model averaging, In Proceedings of the Conference on Computa- tional Learning Theory, 1999.
Vaishnavh Nagarajan and Zico Kolter. Deterministic MAC-Bayesian generalization bounds for deep networks via generalizing noise-resilience. In Proceedings of the International Conference on Learning Representations, 2019a.
Vaishnavh Nagarajan and Zico Kolter. Generalization in deep networks: The role of distance from initialization. arXiv preprint urXiv/190.01162., 2019b.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Norm-based capacity control in neural networks. In Proceedings of the Conference Learning Theory, 2015.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. Exploring gen- cralization deep learning. In Proceedings of the Advances in Neural Information Processing Systems, 2017.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-Bayesian approach to spectrally-normalizzed margin bounds for neural networks. In Proceedings of the International Conference on Learning Representations, 2018
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards understanding the role of peeemerramttinniionnn in generalization of neural networks. In Proceed- ings of the International Conference on Learning Representations, 2019.
Konstantinos Pitas, Mike Davies, and Pierre Vandergheynst.. PAC-Bayesian margin bounds for con- volutional neural networks. arXiv preprint arXiv: 2017.
Hanie Sedghi, Vineet Gupta, and Philip Long. The singular values of convolutional layers. In Proceedings of the International Conference on Learning Representations, 2019.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In Proceedings of the International Conference Learning Representations, 2015.
Colin Wei and Tengyu Ma. Data-dependent sample complexity of deep neural networks via lipschitz augmentation. arXiv preprint arXiv: 1905.03684, 2019.
Chiyuan Zhang, Samy Bengio, and Yoram Singer. Are layers created equal? arXiv preprint arXiv:1902.01996, 2019a.
Hongyi Zhang, Yann Dauphin, and Tengyu Ma. Residual learning without normalization via bet- ter initialization. In Proceedings of the International Conference on Learning Representations, 2019b.
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan Adams, and Peter Orbanz. Non-vacuous gen- cralization bounds at the imagenet scale: a PAC- BayyBsa compression approach. In Proceedings of the International Conference Learning Representations, 2019.
10

---

Published as a conference paper at ICLR 2020
A PROOF OF THEREORM 3.2
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 prior distribution P on parameters that should be chosen in advance and before observing the training set. Through- out this section given two scalar p.q [0, 1] let KL(pllq) 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 E N, train- ing set Sm ~ D, and prior distribution P on parameters θ posterior distribution Q, for any 0 < 6, with probability 1 ó over the draw of training data we have that
KL Ee Bo- (e-o(<o([ee) < KI(QIIP) + log
m-1 where KL is the Bailbback-eeiblers (KL) divergence (Kullback & Leibler, 1951).
Following Dziugaite & Roy (2017),
use the inequality KL (glc)
sup {p E [0, 1] (L(())p) s c) S a vc/2 to achieve a simple bound on the test error:
[L(Q|IP) + log m
EaLolLo(fe)] s -1 [
m -1
KL(Q|IP) + log 5
EEagl[s((fe) +
2
The intuition from Figure suggests that moving the parameters of each module as close as possible to the initialization value before harming performance. For such ag Qi for module to be a Gaussian distribution centered at 0 with covariance matrix of I, that is, as if we have additive noise u, NN(0, o,I). We use O to refer to the case where all 0; replaced with p. and matrix U includes all u, Then the training loss term can be decomposed as
then define the posterior
He=])Cs(fe)] = =Bu.-
{Ls(Jor) + Bu, [on Ls
where the second term on the right hand side of the inequality captures the flatness of the point Oa 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 Ō(1) term in Theorem 3.2, which is equal to e as follows:
k;
=
Σ log 7m + 2og
K + a 10p 80 11
(3)
Now we proceed with the proof.
The KL-divergence between two k-dimensional Gaussian distributions is given by the formula:
(L(N(HI, Ex) Ep)) 2 tr (- + (#2 =1) Σn (#2 41) A + In det Σg det Σ
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 the distribution is (V and let the posterior Q be Gaussian distribution such that for each module i, the distribution is (11 a)op a00"),00
0Z (). We then write the KL--divegeence KL(Q[IP) as
k + af ||0f - 0/122
1 Σ (QIPP) 2
h2 + Ki log
(4)
of
22.1

---

Published as a conference paper at ICLR 2020
Since prior should be decided before observing the training P.P directly. However, one can optimize for op,, over a pre-defined set values and use a union bound argument to get the generalization bound for the best rr in that set. We covering approach suggested by Langford Caruana (2002). For b,, > 0, if one chooses the variance of prior to be exp(-cj cj for j E N such that for each j the bound holds with probabilityl *
not allowed optimize for
6 then all bounds hold with probability EjEN =ł = 1 - 8. We can apply the same idea to every
module such that the bound holds with probability 1 II 11==1791 ン
4m + 1). Otherwise, the bound holds since the right hand side is greater than one. Given (from Equation 3) 2 0, if we choose oZ to have the form exp
If we choose o S 1 then we have opi S exp
Am--i h 1), for some integer ji, can always find choose such that hx
koq2i + aj 00 6.11 S koo I exp(L/k) (Kogg 110 6P
?
(5)
Therefore, the KL-divergence can be bounded as
K + a? 11
0
o
KL(QIIP) S
k; + k; log
k:o + a 110
0
Or. 02.5
2 Σ
his log
exp(11/:) (ki0g, + ot 100f - -0%1
log
exp(1/ki) (K102 as 110f- 00(1(²)
I Σ k log
ki0g,i
(1+ of 10, 11
$ Σ + Kú log
1+
kiog,
Note that in order to achieve the inequality in Equation
5, ja should be chosen as
J: = Am + 1 log
k;
)
k.
S 5m + log
R:
K of 1100 09
koo . + a² 11
09 l1
Given that each such bound should hold with probability 1 81111 ,5 the log term in the bound can be written as
log 8 + E [ogg(13j7/0) < log
m + + > log 7m + 2 log x
k;
kj0z. + at 110f 80|
Putting everything together proves the theorem statement.
B A DETERMINISTIC GENERALIZATIN BOUND FOR CONVOLUTIONAL NETWORKS
We start by stating generalization bound given by Neyshabur et al. (2018) with slight improve- ment in the constants.
Lemma B.1 (Neyshabur et al. (2018).. Let fe + RC be any predictor function with parameters O and P be prior distribution on parameters O Then for any 7: m, 5 O, with probability 1
12

---

Published as a conference paper at ICLR 2020
over the training set S of size m, for any parameter @ and any perturbation distribution Q over
parameters such that Pu-Q mmaxzex lfe+u(z) - ff(()) $ 31 Σ 2' we have
+1 2KL(O + UI) + 1 log m
2(m 1)
The lemma above gives data-independent deterministic bound which depends on the maximum change of the output function over the domain after 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 be an N N image whose norm is bounded by B, fe X + R be the predictor function with parameters which is a DNN of depth d made of convolutional blocks. Then for any margin 7, sample size m, O, with probability 1 over the training set S, any parameter O and any oz > O such that maxzEx (fe(x) Sea (x) S E. we have
1+ [32edBa; || 0011, 11 |0 /ooooodd))))
Ei ki log
+ log ( O(1)
(D(Se) S (S,,(
(6)
where kę is the number of parameters in module i. For example, for convolution module with kernel size qi x qi and number of output channels o, K = 9/
Proof. First, we express the value of the Õ(1) term in Theorem B.2, which is equal to 62 as follows:
=-1 Σ log 7m + 2og
h;
kn [1 |00; 1l2 log log(4dN2) + a 110 80
We note that for any θ O', if m (X)) fel < 7/2 then [(fe) < L. (f6). The reason that the output for each class can change by most 7/2 and therefore the label only change for the data points that are within y 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.
Lp(fer) < LD, (fe=)
(8)
[+ afllo? e0 kio? -
Dd Dim kg log
+ log ma
+ E2
< Sg (Jen) +
af - 1+
1/ Ddd kg log
+ log oh +
< (f.r)feu)) +
m 1
where eg is given above in Equation
Therefore, it suffices to find the value of a; under which the assumption norm of perturbation in function space holds and then simplify the following upper bound given the desired value of o,.
In order to find the desired value of bation lemma by Neyshabur et al. (2018) 1) to bound the change the output a network based on the magnitude of the perturbation:
use the following two lemmas. First we adopt the pertur-
Lemma B.3 (Neyshabur et al. (2018)). Let norm of input x bounded by B. For any B 0, let fe X - HC be a neural network with ReLU activations and depth d. Then for any θ, X, and any perturbation U s.I. ||uz||2 & ||0:||2. the change in the output of the network can be bounded as follows
|
Ilfe+u - fell2 < eB II
10
Σ
10:112
13

---

Published as a conference paper at ICLR 2020
We next use the following lemma by Pitas al. (2017) that bounds the magnitude of the Gaussian perturbation u, for each convolutional module based on the standard deviation the perturbation.
Lemma B.4 (Pitas et al. (2017)). Let u, be Gaussian perturbation for each module of a convo- lutional model. Let N be the image size, qi Gs be the kernel size and the number of output channels at module respectively. We have that
P [[luella 2 (1(C))) +1] S 2N 207
The lemma above suggests that by taking union bounds over all modules, we 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.
||uil|2 c.q((2/e + 22 S ( Vlog(AdN2)) 40,% V (2gg4)22).
Combining this with perturbation bound in Equation 9, we have that
I fer S eB
Julll [ S 4eB Σ σ. V log(4dN2) jyt
II |10 S
Y
8'
where the last inequality can be achieved with
0,
(10)
32edB II |105 d / log (4dN3)
Therefore, this value for σ, 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 62, by showing the simplification from inserting the value for o, from Equation 10 as follows.
log 1+ af 11f - 10
0i 10.0 0.11
s log
1
k,o?
ofo
[32ddBa; 110g- - -8% II
|0 / log (og(4))))]
$ log
1+
ef2222
[32edBa; 1|0,- - -09 II
| Vlog(4dN3)3
log
1
c27
Then, 62 can also be simplified as follows.
K; kio? + at 110f 12
62 1+ Σ log
7m + 2lo
< 11 Σ log
7m + 2 log
2
1 33edBB [I. |10 a: log log(4dNN2) (4dN2))
+ a a 0g
$1+ + Σ log
7m + 2lo
/ (32edB II |109/12 vlog(4dN2)) + + a - 0 2
a: 10f - 0.01 + w/32edB 10:01 / Iog looglea Laddee L
< 11 Σ log
7m + 2 log(k;) loo
14

---

Published as a conference paper at ICLR 2020
C DETAILS ON EXPERIMENTAL SET-UP
For all our experiments, we use the CIFARIO and CIFARION 100 datasets. To train our networks we used Stochastic Gradient Descent (SGD) with momentum 0.9 minimize multi-class cross-entropy loss. On CIFARIO each model is trained until the cross-entropy loss on the training dataset falls below 0.19. While on CIFARIO0 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- learning rate scheduling with a starting learning rate of 0.1 and with decrease by a multiplicative factor of 0.2 every 60 epochs. FCN was trained with an initial learning rate of 0. with 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.
constant
We mainly study three types of neural network architectures:
Fully Connected Networks (FCNs): The FCNs consist of 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 followed by a batch normalization layer and ReLU activation. The final output layer (that follows the ReL.U 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 convolu- tional layers, followed by multiple fully connected layers and 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 final classifier layer (with output dimension 10 or 100). We study the DenseNet with 121 layers.
ResNets: Architectures used are ResNets V1 (He al., 2016). All convolutional layers (except downsample convolutional layers) have kernel size 3 x with stride 1. Downsample convolutions have stride 2. All the ResNets have five stages (0-4) where each stage has multiple residua//donnsample blocks. These stages are followed by maxpool layer and final linear layer. Here are further details about the ResNets used the paper:
ResNet18: ResNet18 8 architechtures studied the paper have convolutional layer in Stage 0 (64 ouput channels), Stage has 2 residual blocks (64 output channels), Stage 2 has one downsample block and one residual block (128 output channels), Stage has one downsample block and one residual block (256 output channels) and Stage 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 convo- lutional layer with 64 output channels followed by a ReLU activation. Stage 1 has residual blocks (64 output channels), Stage 2 has downsample block and 3 residual blocks (128 output channels), Stage has downsample block and 5 residual blocks (256 output channels) and, Stage 4 has downsample block and 2 residual blocks (512 output channels).
ResNet50: ResNet50 architectures in this paper again have stages. Stage 0 has 1 con- volutional layer with 64 output channels followed by a ReLU activation. Stage has downsample block and 2 residual blocks (256 output channels), Stage 2 has I down- sample block and residual blocks (512 output channels), Stage 3 has downsample block and 5 residual blocks (1024 output channels) and, Stage 4 has 1 downsample block and 2 residual blocks (2048 output channels).
ResNet101: 01: Res architectures in this paper again have 5 stages. Stage 0 has convolutional layer with 64 output channels followed by ReLU activation. Stage has downsample block and 2 residual blocks (256 output channels), Stage 2 has downsample block and 3 residual blocks (512 output channels), Stage 3 has down- sample block and 22 residual blocks (1024 output channels) and, Stage 4 has down- sample block and 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 (hor- izontally) and cropped. Global mean and standard deviation are computed on all training images
15

---

Published as a conference paper at ICLR 2020
and applied normalize the inputs. While training ResNet18 the CIFARIO 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.
D ADDITIONAL EXPERIMENTS
In these set of experiments compare the generalization performance of four ResNet18 archi- tectures where we vary the number of output channels in each stage. In the ResNet18 (Ix width) network the number of output channels 16, 16, 32, 64, 128 in the five stages respectively. The other ResNet18s 8s have their output channels scaled by factors of 2,4 and 8 in each stage. Table summarizes our results for networks trained the CIFARIO dataset and Table 5 summarizes results for networks trained on the CIFARIO 100 dataset.
We find that separating these networks based complexity measure is 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.
Table 4: Measuring complexity of different architectures trained on CIFARIO. The ResNet18 ar- chitectures have different channel widths. The network ResNet18 (Ix width) has 16,16,32,64,128 channels in the five stages.
ResNet18 (x)
GE
PFN
PSN
Du
NoP
SoSP
PAC Bayes
Net. Criticality
width
6.27%
4e17
1e12
1409
6.9e5
6.0e7
3.0e5
1.0e5
2x width
5.13%
8e19
4e13
2253
2.7e6
2.9e8
4.2e5
1.3e5
4x width
4.61%
1e22
4e14
3430
1.1c7
1.3c9
6.9e5
2.2c5
8x width
2.88%
lc24
2c15
5365
4.4e7
5.6c9
2.7e6
6.7c5
Kendall's s
-1
-1
-1
-1
-1
-1
Table 5: Measuring complexity of different architectures trained CIFAR100. The ResNet18 architectures have different channel widths. The network ResNet18 (Ix width) has 16,16,32,64,,12 channels in the five stages.
ResNet18 (x)
GE
PFN
PSN
DtI
NoP
SoSP
PAC Bayes
Net. Criticality
Ix width
30.4%
6e18
3e12
2650
7.1e5
7.3e7
3.2e6
1.5e6
2x width
31.8%
le21
4e13
4248
2.8e6
3.4e8
2.7e6
1.3e6
4x width
30.6%
2e22
9e14
4855
1.1c7
1.4c9
3.4e6
1.6c6
8x width
28.4%
3e25
1e18
9269
4.4e7
8.0e9
6.9e6
2.8e6
Kendall's
-0.33
-0.33
-0.33
-0.33
-0.33
-0,66
-0.66
16

---

Published as a conference paper at ICLR 2020
E
FIGURES
07
Training Epochs
Train Erro
3-
0.5
0.4
20
100
0.2
190
200
sty0
stgl.bik)
stg bik2
stg2. bik1
stg2.blk2 b2k2
stg2.blk1
stg3 atl3.bll2
stg4 bik1
stg4 bikr
final linear
0.0
(a) Rewind analysis of Zhang et al. (2019a)
Training Epochs
0
2
1
0.6
Train Error
20
100
0.3
190
0.0
TADos
mooon 17
conve
cooea 2n00 1Gis
cone
coon 1a000 20w
Soooo 2010
coge 2010 2615
conn
sooos 2 gōy
coooe 2010 g6p
sooee tho pōys
Touoss 1210 p5is
crne puys sotoe ZNI
cone mn pays
1x10 16w
To0 16p
ahe tōps
THe 261
g0ys
twu ghis
Exiq pass
20ps
We rewind each module, whereas, Zhang et (2019a) rewind each block
Training Epochs
Training Epochs
0
-307 T0ll
3-
20-
100
100
200
201 pe0.caai
bit Mit..coov
o0l bir Lot.22
sojt biz conv2
seg bil Loccv
downse molee
comnz
tnd pou cooone
ao cone mo 2h22 2 conve (mml soun
the Cosooo
tie 2o00222
cooee tn tn 20p
correree toor tis
Toio cun gfp
Counge ain pdge
somul yay
Call phus phi
ths
dhe
Tes
naie
slge
(c) The effect rewinding on train loss
(d) The effect of rewinding on test error
Figure 6: Analysis of rewinding modules to initialization for the ResNet- architecture. Each row represents a layer in ResNet18-v1 and each column represents a particular training epoch that the module is rewound
17

---

Published as a conference paper at ICLR 2020
Training Epochs
0
0.
Train Error
05
3
100
180
190
Stg0 conv1
stg bllbbl11. conov
stg1 bik1 conv2
stg1 bik. convi
stg1 bllbblk22 conv2
stg1 bll.blkk convi
stg1. bllk conv2
stgl bik4 conv2
stg1 blk4 conv2
stg2 blkb conv2
stg2 bik1 conv2
stg2 wwl.bl2..ccc convl
stg2 blk- conv2
stg2 b2l.bk33 conv2
stg2 bik blk3.conv2 conc2
stg2 bik4 conv2
stg? bik4 conv2
stg3 blkk1 conv2
stg3. 22..bik1. conv2
stg3 bl2.blk222 convl
stg3 bik conv2
stg blk.blkk.... convl
stg3 bikk conv2
stg3 blk4 conv2
stg3. bl..blk4 conv2
(a) Fixup initialization
Training Epochs
Train Error
2
-06
3-
20
100
190
02
200
stgo convz
stg bllbblk1 convi
stg blk.blkk1kkc conv2
stgl. bikk conv1
stg blkk conv2
stg2 bik. conv2
stg2 blk1. conv2
stg2 bikk www.b/kkkkkww.... w.w.dammppleee
stg2 blkk. conv2
St22 bll.blk2. con.2c
stg 1.kk convi
stg b.l.blkk conv2
www.dommppppoo ample
stg b.l.k22.. convi
stg blk conv2
stg4 bik1. conv2
stg4. bik1 conv2
stg4 blk1 downsample
stg4 blkk conv1
St94 blkk conv2
final linear
(b) Adam optimizer
Figure 7: Criticality pattern of Resnet18 when trained with Fixup initialization and with the Adam optimizer
Training Epochs
Operator Norm
3
20
100
190
3.9 conve
stg1 blk2 conv1
stg blk1 conv2
stg1 blk. conv1
stg1. blkk conv2
stg2 blk1 conv2
stg2 bik1 conv2
dgg2 blkl..1kk.. downsample
2 blkz blk2.ccovz 2onvi
conv2
egg blke1 conv2
wile biki con22
Coumsammee blk2 convi
canve
2os tikt 2anve
conva cowneaaae olk2 allca yelt
sx22 blk2
2
eaa bikz
a a pōs alt
3222
p6g
Figure Operator norm of difference from initialization
18

---

Published as a conference paper at ICLR 2020
S000 b00% todas
Mest Coom bede
MoDoiu booe tetas
Pia" p9pos" padess
topouu rood potms
pia' booee cadays
pus poeee peess
Appoi" pooo podes
sooo 2000 podias
"pou"" pooys pallas
puo 1000 todes
*prru copoe's patwes
pie 1ete pedm
putour 2000 20ta
"pon"" ssooe's Labves
"ptout capae's wttees
pio" CwwN' cddes
pan' (rwo yodlas
2eoe zodes
catiss
(a) Rewind an ambient module
(b) Rewind critical module
Figure 9: Similarity in activation patterns when an ambient critical module is rewound. Darker green denotes higher similarity
critical cross-
critical 0/1 loss
56
0.4
0.96
0.4
48
0.2
0.80
0.2
40
0.64
32
0.0
0.48
0.0
24
-0.2
0.32
-0.2
16
0.16
8
-0.4
-0.4
0.00
0
0.25
0.50
0.75
1.00
0.25
0.50
0.75
1.00
critical cross-entropy
non-critical 0/1 loss
56
0.4
0.96
0.4
48
0.2
0.80
0.2
40
0.64
32
0.0
0.48
0.0
24
-.2
0.32
-0.2
16
0.16
8
-0.4
-0.4
0.00
0.25
0.50
0.75
1.00
0.25
0.50
0.75
1.00
a
Figure 10: 0/1 loss and cross-entropy loss for critical and non-critical modules, for given different values of and a in Definition 3.1.
19


Reviews:
[{"review_id": "SJlrtEmAKS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces a new way to reason about neural network generalization using a module criticality measure. The measure is tangible and intuitive. It leads to some formal bounds on the generalization of deep networks, and is able to better rank trained image classification architectures than previous measures.\n\nI am leaning to accept, as I expect this to be a significant theoretical contribution with several potential practical applications. With a few additional details, this could be a very strong submission:\n\n(1)\tChoice of module decomposition. Having each module be a single convolutional or fully-connected layer makes intuitive sense, but is there some theoretical motivation for this choice? If the only requirement for a module is that it includes some linear transformation, in the extreme, a module could consist of a single weight, or the entire network. Would those choices change the generalization bounds or relative criticality across different architectures?\n(2)\tScope of experimental results. The ranking results would be much more compelling if they included a broader range of architectures, including more recent models with more branching, e.g., DenseNet. Is there some reason ResNet101 has higher generalization error than 18 and 34? Net. Criticality for ResNets is inversely correlated with the number of layers; is there an explanation for this? Is this true for other very deep models?\n(3)\tPractical use. To compute the criticality measure, we must train the model; but, if we train the model, we can compute generalization directly. So, what is the practical application of the measure? Is there some way it could be used to save computation? Could it help in the case of a small validation dataset, which we do not want to look at many times during model selection?\n\nMinor typos:\n-\tSection 2.2: “An stable phenomena”\n-\tSection 2.3: “…an the…”\n-\tIn appendix: “ResNet101: ResNet34 architectures…”\n\n----------------------------\n\nAfter rebuttal:\n\nThe authors have addressed my concerns, and I've increased my rating. There are still a few points I'd like to see addressed in the final version:\n\n1. The fact that the approach cannot yet be applied to batch normalization is a big practical drawback. Some discussion of approaches you tried, why they didn't work, and possible future directions for overcoming this would be appreciated.\n\n2. Clarify in the paper that the \"PAC Bayes\" approach used for comparison (Table 1) is a your method, i.e., an ablated version of criticality. As is, someone reading the paper quickly may think all you've done is add an alpha parameter to an existing \"PAC Bayes\" approach, which does fairly well on its own.\n\n3. Visualizing the experimental tables as scatterplots could make them easier for a reader to interpret.\n\n\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SJeMgKD5tS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper builds upon the \"module criticality\" phenomenon and proposes a quantitative approach to measure this at the module and the network level.  A module's criticality is low if when it is switched to its initialization value, the error does not change drastically. \n\nThe paper uses a convex combination of the initial weights and the final weights of a layer/module to define an optimization path to traverse. The authors quantitatively define the module criticality such that it depends on how much closer the weights can get to the initial weights on this path while still being robust to random permutations. The network critically is defined as the sum of the module criticality measure of all the layers. \n\nEmpirical results on CIFAR10 show that the network's criticality is reflective of the generalization performance. For example, increasing resnet depth leads to improved generalization and low criticality. Though intuitively, it is not clear why moving closer to the initial values and thus lower average criticality indicated better generalization. It will be useful ot have a discussion on this issue. Results on other datasets will also be useful.\n\nOverall, the network criticality measure appears a useful tool to predict the generalization performance compared to other measures such as distance from initialization, weight spectrum, and others. ", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper."}, "extra_text": {}}, {"review_id": "ByxzqndmsB", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe paper studies this per module compared to previous works where the entire architecture is rewounded. The paper also studies this for ResNet models which are more widely/practically used than just the fully connected layers alone. This helps better understand the model as a whole. \n\nThe authors do a good robust experimental study for different network initialization, various CNN models like ResNet18, 34, 101, VGG16 and also FCN. The results demonstrate the module criticality to be a good metric for generalization of models.\n\nOverall, a good paper.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper analyses the importance of different DNN modules for generalization performance, explaining why certain architectures may be much better performing than others. All reviewers agree that this is an interesting paper with a novel and important contribution. 
Decision: Accept (Spotlight)