Title: THE INTRIGUING ROLE OF MODULE CRITICALITY IN THE GENERALIZATION OF DEEP NETWORKS

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.

1. INTRODUCTION:

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. Given a ball with radius r (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 r. Hence, non-critical modules are the ones with a 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 to increase when the ball comes too close to the initial weight. The x axis is simply chosen to be parallel to θ F -θ 0 and the y axis is a compact representation of all other dimensions generated by adding Gaussian noise to the points on the convex combination of θ 0 and θ F and evaluating the loss. The sign on the y axis is decided based on the sign of the inner product of the noise to θ 0 .

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 ] be the convex combination between initial weights θ 0 i 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).

θ α i = (1 -α)θ 0 i + αθ F i , α ∈ [0, 1

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 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 L S (f ) be the loss of function f on the training set S with m samples. We are mainly interested in the classification task where

L S (f ) = 1 m (x,y)∈S 1[f (x)[y] ≤ max j =y f (x)[j]]. For any γ > 0, we also define margin loss L S,γ (f ) = 1 m (x,y)∈S 1[f (x)[y] ≤ γ + max j =y f (x)[j]]. Let L D (f )

be the loss of function f on population data distribution D defined similar to L S (f ). We will denote the function parameterized by Θ by f Θ .

2. TOWARDS UNDERSTANDING MODULE CRITICALITY:



2.1. SETTING:

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 θ 0 i , θ 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.

2.2. ROBUSTNESS TO REWINDING:

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.

2.3. WHAT MEASURES FAIL TO DISTINGUISH CRITICAL LAYERS:

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.

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 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 -α)θ 0 i + αθ F i , α ∈ [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:

EQUATION

We also define the network criticality as the sum of the module criticality over modules of the network:

EQUATION

Here, L S 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 i th module is replaced by θ α i and all other modules are fixed at their values in the end of training, Θ F -i .

θ α i = (1 -α)θ 0 i + αθ F i ,

where θ 0 i is the value of the weight matrix at initialization and θ F i is the trained value. Intuitively, network criticality measure is sum of module criticalities. This is 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 deriving a generalization bound using the PAC-Bayesian framework (McAllester, 1999) . Given a prior Figure 4 : Performance degradation as we move on convex combination path from final to 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 to 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 a single (non-critical) layer and the yellow (dashdot) curve corresponds to moving on the convex path for a critical layer in ResNet-18 architecture. 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 Q i 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 u i ∼ N (0, σ i I). We use Θ α to refer to the case where all the parameters θ i are replaced with θ αi i and matrix U includes all the noise u i . 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 S m ∼ D the following generalization bound holds:

E U [L D (f Θ α +U )] ≤ E U [L S (f Θ α +U )] + 1 4 d i=1 k i log 1 + α 2 i θ F i -θ 0 i 2 Fr kiσ 2 i + log m δ + O(1) m -1 ,

where k i is the number of parameters in module i. For example, for a convolution module with kernel size q i × q i and number of output channels c i ,

k i = q 2 i c i-1 c i .

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

L D (f Θ F ) ≤ E U [L D (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 L D (f Θ F ) ≤ E U [L D (f Θ α +U )]

where u i ∼ N (0, σ i I) , then with probability 1 -δ over the choice of the training set S m ∼ D, the following generalization bound holds

L D (f Θ ) ≤ + 1 4 µ (f Θ ) + log m δ + O(1) m -1 ,

where µ (f Θ ) is calculated as follows:

µ (f Θ ) = min 0≤α,σ≤1 i α 2 i θ F i -θ 0 i 2 Fr σ 2 i : E U [L S (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.

4. EXPERIMENTS:

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. 

L D (f Θ ) -L S (f Θ ) Product of Frobenius Norms (PFN) Π i θ F i Fr Product of Spectral Norms (PSN) Π i θ F i 2 Distance to Initialization (DtI) i θ 0 i -θ F i 2 Fr

Number of Parameters (NoP) Total number of parameters in the network Sum of Spectral Norms (SoSP)

Total number of parameters ×( i θ 0 i -θ F i 2 ) PAC Bayes (at error threshold 0.1)

i θ 0 i -θ F i 2 Fr /σ 2 i Network Criticality Measure (at error threshold 0.1) i α 2 i θ 0 i -θ F i 2 Fr /σ 2 i

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.

5. CONCLUSION:

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 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 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 S m ∼ 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 [L S (f Θ )] E Θ∼Q [L D (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 [L D (f Θ )] ≤ KL -1 E Θ∼Q [L S (f Θ )] KL(Q||P ) + log m δ m -1 ≤ E Θ∼Q [L S (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 Q i 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 u i ∼ N (0, σ i I). We use Θ α to refer to the case where all θ i are replaced with θ αi i and matrix U includes all u i . Then the training loss term can be decomposed as

E Θ∼Q [L S (f Θ )] = E ui∼N (0,σiI) [L S (f Θ α +U )] ≤ L S (f Θ F ) + E ui∼N (0,σiI) [L S (f Θ α +U )] -L S (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:

EQUATION

)

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 Σ -1 2 Σ 1 + (µ 2 -µ 1 ) Σ -1 2 (µ 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 (θ 0 i , σ 2 P,i I) and let the posterior Q be a Gaussian distribution such that for each module i, the distribution is

N ((1 -α)θ 0 i + αθ F i ), σ 2 Q,i I).

We can then write the KL-divergence KL(Q||P ) as

EQUATION

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 π 2 j 2 , then all bounds hold with probability 1 -j∈N 6 π 2 j 2 = 1 -δ. We can apply the same idea to every module such that the bound holds with probability 1 -δ

d i=1 6 π 2 j 2 i .

If we choose σ 2 Q,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 σ 2 P,i to have the form exp 4m-ji ki + 1 , for some integer j i , we can always find choose j i such that

EQUATION

Therefore, the KL-divergence can be bounded as

KL(Q||P ) ≤ 1 2 i k i k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr -k i + k i log σ 2 P,i σ 2 Q,i = 1 2 i k i log σ 2 P,i σ 2 Q,i ≤ 1 2 i k i log   exp(1/k i ) k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr k i σ 2 Q,i   ≤ 1 2 i k i log   exp(1/k i ) k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr k i σ 2 Q,i   ≤ 1 2 i 1 + k i log 1 + α 2 i θ F i -θ 0 i 2 Fr k i σ 2 Q,i

.

Note that in order to achieve the inequality in Equation 5, j i should be chosen as

j i = 4m k i + 1 + log k i k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr ≤ 5m + log k i k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr .

Given that each such bound should hold with probability 1 -δ

d i=1 6 π 2 j 2 i

, the log term in the bound can be written as

log m δ + i log(π 2 j 2 i /6) ≤ log m δ + 2 i log 7m + 2 log k i k i σ 2 Q,i + α 2 i θ F i -θ 0 i 2 Fr .

Putting everything together proves the theorem statement. 2018)). Let f Θ : X → R C 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

P U ∼Q max x∈X |f Θ+U (x) -f Θ (x)| ≤ γ 4 ≥ 1 2 , we have L D (f Θ ) ≤ L S,γ (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 → R C 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

EQUATION

where k i is the number of parameters in module i. For example, for a convolution module with kernel size q i × q i and number of output channels c i ,

k i = q 2 i c i-1 c i .

Proof. First, we express the value of the O(1) term in Theorem B.2 , which is equal to 2 as follows:

EQUATION

We note that for any Θ, Θ , if

max x∈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.

L D (f Θ F ) ≤ L D, γ 4 (f Θ α ) (8) ≤ L S, 3γ 4 (f Θ α ) + 1 2 d i=1 k i log 1 + α 2 i θ F i -θ 0 i 2 Fr kiσ 2 i + log m δ + 2 m -1 ≤ L S,γ (f Θ α ) + 1 2 d i=1 k i log 1 + α 2 i θ F i -θ 0 i 2 Fr kiσ 2 i + 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 → R C be a neural network with ReLU activations and depth d. Then for any Θ, x ∈ X , and any perturbation U s.t. u i 2 ≤ θ i 2 , the change in the output of the network can be bounded as follows

EQUATION

)

We next use the following lemma by Pitas et al. (2017) that bounds the magnitude of the Gaussian perturbation u i for each convolutional module based on the standard deviation of the perturbation.

Lemma B.4 (Pitas et al. (2017) ). Let u i be a Gaussian perturbation for each module i of a convolutional model. Let N be the image size, q i , c i be the kernel size and the number of output channels at module i respectively. We have that

P u i 2 ≥ σ i q i (2 √ c i ) + t ≤ 2N 2 e -t 2 2q 2 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.

u i 2 ≤ σ i q i (2 √ c i + 2 log(4dN 2 )) ≤ 2σ i q i ( √ c i + log(4dN 2 )) ≤ 4σ i q i c i log(4dN 2 ).

Combining this with perturbation bound in Equation 9, we have that

f Θ α +U -f Θ α 2 ≤ eB d i=1 u i 2 d j =i θ α j 2 ≤ 4eB d i=1 σ i q i c i log(4dN 2 ) d j =i θ α i 2 ≤ γ 8 ,

where the last inequality can be achieved with

EQUATION

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 + α 2 i θ F i -θ 0 i 2 Fr k i σ 2 i ≤ log 1 + α i θ F i -θ 0 i Fr 2 q 2 i c 2 i σ 2 i ≤ log 1 + [32edBα i θ F i -θ 0 i Fr d i=1 θ α i 2 q i c i log(4dN 2 )] 2 q 2 i c 2 i γ 2 = log 1 + [32edBα i θ F i -θ 0 i Fr d i=1 θ α i 2 log(4dN 2 )] 2 c i γ 2 .

Then, 2 can also be simplified as follows.

2 = 1 + i log 7m + 2 log k i k i σ 2 i + α 2 i θ F i -θ 0 i 2 Fr ≤ 1 + i log   7m + 2 log    k i k i γ 2 / 32edB d i=1 θ α i 2 q i c i log(4dN 2 ) 2 + α 2 i θ F i -θ 0 i 2 Fr       ≤ 1 + i log   7m + 2 log    k i γ 2 / 32edB d i=1 θ α i 2 log(4dN 2 ) 2 + α 2 i θ F i -θ 0 i 2 Fr       ≤ 1 + i log 7m + 2 log(k i ) -4 log α i θ F i -θ 0 i Fr + γ/32edB d i=1 θ α i 2 log(4dN 2 ) .

C DETAILS ON EXPERIMENTAL SET-UP:

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). 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.

D ADDITIONAL EXPERIMENTS:

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. 

:

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 layered model such as VGG, module definition is equivalent to definition of a layer.

A 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.

Figures and Tables:

Figure 1: Module Criticality: Loss values in the valleys that connect the initial weights θ 0 to the final weights θ F of a non-critical (left) and a critical (right) module in the ResNet18 architecture.Given a ball with radius r (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 r. Hence, non-critical modules are the ones with a 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 to increase when the ball comes too close to the initial weight. The x axis is simply chosen to be parallel to θ F -θ 0 and the y axis is a compact representation of all other dimensions generated by adding Gaussian noise to the points on the convex combination of θ 0 and θ F and evaluating the loss. The sign on the y axis is decided based on the sign of the inner product of the noise to θ 0 .

Figure 2: Analysis of rewinding modules to initialization for the ResNet-18 architecture. Each row represents a module in ResNet18-v1 and each column represents a particular training epoch to 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.

Figure 3: Spectrum of a non-critical and a critical module during different epochs of training.

Figure 5: Network criticality as a function of error threshold for networks trained on CIFAR10.

A DETERMINISTIC GENERALIZATION BOUND FOR CONVOLUTIONAL NETWORKS 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. (

Figure 7: Criticality pattern of Resnet18 when trained with Fixup initialization and with the Adam optimizer

Quantities of Interest

Table content: <table><tr><td>Generalization Error (GE)</td></tr></table>

Measuring complexity of different architectures trained on CIFAR10.

Table content: <table><tr><td>Network</td><td>GE</td><td>PFN</td><td>PSN</td><td>DtI</td><td colspan="3">NoP SoSP PAC Bayes Net. Criticality</td></tr><tr><td>ResNet18</td><td>4.61%</td><td>1e22</td><td>4e14</td><td colspan="2">3430 1.1e7 1.3e9</td><td>6.9e5</td><td>2.2e5</td></tr><tr><td>ResNet34</td><td>6.3%</td><td>2e37</td><td>3e24</td><td colspan="2">4768 2.1e7 3.7e9</td><td>9.1e5</td><td>1.7e5</td></tr><tr><td>ResNet50</td><td>6.6%</td><td>4e56</td><td colspan="3">4e20 10018 2.3e7 3.3e9</td><td>1.6e6</td><td>1.8e5</td></tr><tr><td>ResNet101</td><td>6.4%</td><td colspan="4">8e110 3e32 18730 4.2e7 9.8e9</td><td>2.8e6</td><td>6.3e5</td></tr><tr><td>DenseNet121</td><td>7.8%</td><td colspan="4">2e129 7e42 21359 6.8e6 2.0e9</td><td>1.2e6</td><td>4.1e5</td></tr><tr><td>VGG11</td><td>8.51%</td><td>1e11</td><td>1e6</td><td colspan="2">2106 2.8e7 1.3e9</td><td>1.0e6</td><td>2.8e5</td></tr><tr><td>VGG16</td><td>7.47%</td><td>5e15</td><td>2e8</td><td colspan="2">2341 3.4e7 2.1e9</td><td>1.2e6</td><td>2.70e5</td></tr><tr><td>FCN (I)</td><td colspan="2">29.83% 3e20</td><td>2e7</td><td colspan="2">75221 2.0e7 4.6e8</td><td>9.0e6</td><td>5.7e6</td></tr><tr><td>FCN (II)</td><td colspan="2">26.45% 3e21</td><td>1e7</td><td colspan="2">81258 5.0e7 2.0e9</td><td>9.5e6</td><td>6.2e6</td></tr><tr><td>Kendall's τ</td><td>-</td><td colspan="2">-0.22 -0.33</td><td>0.38</td><td>0.16 -0.53</td><td>0.42</td><td>0.55</td></tr></table>

Measuring complexity of different architectures trained on CIFAR100. , 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.

Table content: <table><tr><td>Network</td><td>GE</td><td>PFN</td><td>PSN</td><td>DtI</td><td>NoP</td><td colspan="3">SoSP PAC Bayes Net. Criticality</td></tr><tr><td>ResNet18</td><td>30.6%</td><td>2e22</td><td>9e14</td><td>4855</td><td colspan="2">1.1e7 1.4e9</td><td>3.4e6</td><td>1.6e6</td></tr><tr><td>ResNet34</td><td>29.3%</td><td>1e37</td><td>1e22</td><td>6017</td><td colspan="2">2.1e7 3.6e9</td><td>6.7e6</td><td>1.4e6</td></tr><tr><td>ResNet50</td><td>31.1%</td><td>7e57</td><td>9e23</td><td colspan="3">12715 2.3e7 4.1e9</td><td>5.8e6</td><td>1.9e6</td></tr><tr><td>ResNet101</td><td colspan="3">25.5% 2e112 5e36</td><td colspan="3">21233 4.2e7 1.1e10</td><td>6.4e6</td><td>1.3e6</td></tr><tr><td colspan="4">DenseNet121 35.3% 1e131 1e49</td><td/><td colspan="2">6.9e6 2.4e9</td><td>4.5e6</td><td>2.6e6</td></tr><tr><td>VGG11</td><td>43.5%</td><td>6e13</td><td>1e8</td><td>5059</td><td colspan="2">2.8e7 1.7e9</td><td>3.6e6</td><td>2.4e5</td></tr><tr><td>VGG16</td><td colspan="2">32.69% 8e19</td><td>4e12</td><td>7010</td><td colspan="2">3.4e7 3.6e9</td><td>8.1e6</td><td>4.6e6</td></tr><tr><td>FCN (I)</td><td colspan="2">57.59% 8e27</td><td colspan="4">1e10 246636 2.0e7 8.3e8</td><td>2.4e7</td><td>1.2e7</td></tr><tr><td>FCN (II)</td><td colspan="2">53.17% 1e29</td><td colspan="4">1e10 329296 5.0e7 3.7e9</td><td>3.3e7</td><td>2.3e7</td></tr><tr><td>Kendall's τ</td><td>-</td><td colspan="2">-0.27 -0.47</td><td>0.33</td><td>0</td><td>-0.42</td><td>0.22</td><td>0.55</td></tr><tr><td colspan="2">architectures over others</td><td/><td/><td/><td/><td/><td/></tr></table>

VGGs: Architectures bySimonyan & 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 byHuang 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: 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).

Table content: <table><tr><td>-ResNet18: ResNet18 architechtures studied in the paper have 1 convolutional layer in</td></tr><tr><td>Stage 0 (64 ouput channels), Stage 1 has 2 residual blocks (64 output channels), Stage</td></tr><tr><td>2 has one downsample block and one residual block (128 output channels), Stage 3</td></tr><tr><td>has one downsample block and one residual block (256 output channels) and Stage 4</td></tr><tr><td>again has one downsample block and a residual block (512 output channels).</td></tr><tr><td>-ResNet34: ResNet34 architectures in this paper have 5 stages. Stage 0 has 1 convo-</td></tr><tr><td>lutional layer with 64 output channels followed by a ReLU activation. Stage 1 has 3</td></tr><tr><td>residual blocks (64 output channels), Stage 2 has 1 downsample block and 3 residual</td></tr><tr><td>blocks (128 output channels), Stage 3 has 1 downsample block and 5 residual blocks</td></tr><tr><td>(256 output channels) and, Stage 4 has 1 downsample block and 2 residual blocks</td></tr><tr><td>(512 output channels).</td></tr><tr><td>-ResNet50:</td></tr></table>

Measuring complexity of different architectures trained on CIFAR10. The ResNet18 architectures have different channel widths. The network ResNet18 (1x width) has 16,16,32,64,128 channels in the five stages.

Table content: <table><tr><td>ResNet18 (x)</td><td>GE</td><td colspan="2">PFN PSN</td><td>DtI</td><td colspan="4">NoP SoSP PAC Bayes Net. Criticality</td></tr><tr><td>1x width</td><td colspan="6">6.27% 4e17 1e12 1409 6.9e5 6.0e7</td><td>3.0e5</td><td>1.0e5</td></tr><tr><td>2x width</td><td colspan="6">5.13% 8e19 4e13 2253 2.7e6 2.9e8</td><td>4.2e5</td><td>1.3e5</td></tr><tr><td>4x width</td><td colspan="6">4.61% 1e22 4e14 3430 1.1e7 1.3e9</td><td>6.9e5</td><td>2.2e5</td></tr><tr><td>8x width</td><td colspan="6">2.88% 1e24 2e15 5365 4.4e7 5.6e9</td><td>2.7e6</td><td>6.7e5</td></tr><tr><td>Kendall's τ</td><td>-</td><td>-1</td><td>-1</td><td>-1</td><td>-1</td><td>-1</td><td>-1</td><td>-1</td></tr></table>

Measuring complexity of different architectures trained on CIFAR100. The ResNet18 architectures have different channel widths. The network ResNet18 (1x width) has 16,16,32,64,128 channels in the five stages.

Table content: <table><tr><td>ResNet18 (x)</td><td>GE</td><td>PFN PSN</td><td>DtI</td><td colspan="2">NoP SoSP PAC Bayes Net. Criticality</td></tr><tr><td>1x width</td><td colspan="4">30.4% 6e18 3e12 2650 7.1e5 7.3e7</td><td>3.2e6</td><td>1.5e6</td></tr><tr><td>2x width</td><td colspan="4">31.8% 1e21 4e13 4248 2.8e6 3.4e8</td><td>2.7e6</td><td>1.3e6</td></tr><tr><td>4x width</td><td colspan="4">30.6% 2e22 9e14 4855 1.1e7 1.4e9</td><td>3.4e6</td><td>1.6e6</td></tr><tr><td>8x width</td><td colspan="4">28.4% 3e25 1e18 9269 4.4e7 8.0e9</td><td>6.9e6</td><td>2.8e6</td></tr><tr><td>Kendall's τ</td><td>-</td><td colspan="3">-0.33 -0.33 -0.33 -0.33 -0.33</td><td>-0.66</td><td>-0.66</td></tr></table>

Peer review (iclr2020):
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. The 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. The 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. Overall, a good paper.