The intriguing role of module criticality in the generalization of deep networks

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 . Every DNN architecture is a computation graph where each node is a module1 . We are interested in understanding how different modules in the network interact with each other and influence generalization performance as a whole . To do so , we delve deeper into the phenomenon of “ module criticality ” which was reported by Zhang et al . ( 2019a ) . They observed that modules of the network present different robustness characteristics to parameter perturbation . Specifically , they look into the following perturbation : Rewind one module back to its initialization value while keeping all other modules fixed ( at the final trained value ) . They note that the impact of this perturbation on network performance varies between modules and depends on which module was rewound . Some modules are “ critical ” meaning that rewinding their value to the initialization harms the network performance , while for others the impact of this perturbation on performance is negligible . They show that various conventional DNN architectures exhibit this phenomenon . Let us now informally define what we mean by the measure “ module criticality ” ( see Figure 1 ) . For each module , we move on a line from its final trained value to its initialization value ( convex combination2 path ) while keeping all other modules fixed at their trained value . Then we measure the performance drop . Let θαi = ( 1 − α ) θ0i + αθFi , α ∈ [ 0 , 1 ] be the convex combination between initial weights θ0i and the final weights θ F i at module i , where αi is the minimum value between 0 and 1 when performance ( train error ) of the network drops by at most a threshold value . If αi is small we can move a long way back to initialization without hurting performance and the “ module criticality ” of this module would be low . Further , we also wish to incorporate the robustness to noise ( that is , the valley width ) for the module along this path . If the module is robust to noise along this path ( that is , the valley is wide ) then the module criticality would again be low ( see Definition 3.1 for a formal definition ) . In this paper , we seek to study this phenomenon in depth and shed some light on it by showing that conventional complexity measures can not capture criticality ( see Section 2 ) . Next , we theoretically formulate this phenomenon and analyze its role in generalization . Through this analysis , we provide 1A module is a node in the computation graph that has incoming edges from other modules and outgoing edges to other nodes and performs a linear transformation on its inputs . For layered model such as VGG , module definition is equivalent to definition of a layer . 2A convex combination of two points is a linear combination of them where the coefficients are non-negative and sum to 1 . Every convex combination of two points lies on the line segment between the two . a new generalization measure that captures the dissimilarity of different modules and depicts how it influences the generalization of the corresponding DNN . Intuitively , the closer we can get to initialization for each module , the better the generalization . We analyze the relation between generalization and module criticality through a PAC-Bayesian analysis . In Section 3 we show that it is the overall network criticality measure and not the number of critical modules that controls generalization . If the network criticality measure is smaller for an architecture , it has better generalization performance . In Section 4 , we demonstrate through various experiments that our proposed measure is able to distinguish between different network architectures in terms of their generalization performance . Moreover , the network criticality measure is able to correctly rank the generalization performance of different architectures better than the measures proposed earlier . Notation : We use upper case letters for matrices . The operator norm and Frobenius norm of M are denoted by ‖M‖2 , ‖M‖Fr respectively . For n ∈ N , we use [ n ] to denote the set { 1 , . . . , n } . Let LS ( f ) be the loss of function f on the training set S with m samples . We are mainly interested in the classification task where LS ( f ) = 1m ∑ ( x , y ) ∈S 1 [ f ( x ) [ y ] ≤ maxj 6=y f ( x ) [ j ] ] . For any γ > 0 , we also define margin loss LS , γ ( f ) = 1m ∑ ( x , y ) ∈S 1 [ f ( x ) [ y ] ≤ γ + maxj 6=y f ( x ) [ j ] ] . Let LD ( f ) be the loss of function f on population data distribution D defined similar to LS ( f ) . We will denote the function parameterized by Θ by fΘ . 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 θ0i , θ F i refer to the value of weight matrix at initialization and the end of training , respectively . For sequential architectures , d is equal to the depth of the network but that is not necessarily true for a general architectures such as a ResNet . 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 .

Summary: 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.