Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers

1 INTRODUCTION Thanks to many architectural and algorithmic innovations , the recent decade has witnessed the unprecedented success of deep learning in various high-profile challenges , e.g. , the ImageNet recognition task ( Krizhevsky et al. , 2012 ) , the challenging board game of Go ( Silver et al. , 2017 ) and human-like text generation ( Brown et al. , 2020 ) . Among them , shortcut connections ( He et al. , 2016a ; Srivastava et al. , 2015 ) and normalization layers ( Ioffe & Szegedy , 2015 ; Ba et al. , 2016 ) are two architectural components of modern networks that are critically important for achieving fast training at very high depths , and feature prominently in the ubiquitous ResNet architecture of He et al . ( 2016b ) . Despite the success of ResNets , there is significant evidence to suggest that the primary reason they work so well is that they resemble ensembles of shallower networks during training ( Veit et al. , 2016 ) , which lets them avoid the common pathologies associated with very deep networks ( e.g . Hochreiter et al. , 2001 ; Duvenaud et al. , 2014 ) . In this sense , the question of whether truly deep networks can be efficient and effectively trained on challenging tasks remains an open one . As argued by Oyedotun et al . ( 2020 ) and Ding et al . ( 2021 ) , the multi-branch topology of ResNets also has certain drawbacks . For example , it is memory-inefficient at inference time , as the input to every residual block has to be kept in memory until the final addition . In particular , the shortcut branches in ResNet-50 account for about 40 % of the memory usage by feature maps . Also , the classical interpretation of why deep networks perform well – because of the hierarchical feature representations they produce – does not strictly apply to ResNets , due to their aforementioned tendency to behave like ensembles of shallower networks . Beyond the drawbacks of ResNets , training vanilla deep neural networks ( which we define as networks without shortcut connections or normalization layers ) is an interesting research problem in its own right , and finding a solution could open the path to discovering new model architectures . However , recent progress in this direction has not fully succeeded in matching the generalization performance of ResNets . Schoenholz et al . ( 2017 ) used a mean-field analysis of deep MLPs to choose variances for the initial weights and bias parameters , and showed that the resulting method – called Edge of Chaos ( EOC ) – allowed vanilla networks to be trained at very high depths on small datasets . Building on EOC , and incorporating dynamical isometry theory , Xiao et al . ( 2018 ) was able to train vanilla networks with Tanh units1 at depths of up to 10,000 . While impressive , these EOC-initialized networks trained significantly slower than standard ResNets of the same depth , and also exhibited significantly worse generalization performance . While Oyedotun et al . ( 2020 ) was able to narrow the generalization gap between vanilla networks and ResNets , their experiments were limited to networks with only 30 layers , and their networks required many times more parameters . More recently , Martens et al . ( 2021 ) introduced a method called Deep Kernel Shaping ( DKS ) for initializing and transforming networks ( and their activation functions ) based on an analysis of their initialization-time kernel properties . They showed that their approach enabled vanilla networks to train faster than previous methods , even matching the speed of similarly sized ResNets when combined with stronger optimizers like K-FAC ( Martens & Grosse , 2015 ; Grosse & Martens , 2016 ) or Shampoo ( Anil et al. , 2020 ) . However , their method isn ’ t fully compatible with ReLUs , and in their experiments ( which focused on training speed ) their networks exhibited significantly more overfitting than ResNets . Inspired by both DKS and the line of work using mean-field theory , we propose a new method called Tailored Activation Transformation ( TAT ) . TAT inherits the main advantages of DKS , while working particularly well with the “ Leaky ReLU ” activation function . While being easy to implement and introducing negligible extra computational cost , TAT enables very deep vanilla neural networks to be trained on ImageNet without the use of any additional architectural elements . Using TAT , we demonstrate for the first time that a 50-layer vanilla deep network can nearly match the validation accuracy of its ResNet counterpart when trained on Imagenet . And unlike with the EOC method , validation accuracy we achieve does not decrease with depth ( see Figure 1 ) . Furthermore , TAT can also be applied to ResNets without normalization layers , allowing them to match or even exceed the validation accuracy of standard ResNets of the same width/depth . 2 BACKGROUND . Our main tool of analysis will be kernel functions for neural networks ( Neal , 1996 ; Cho & Saul , 2009 ; Daniely et al. , 2016 ) and the related Q/C maps ( Saxe et al. , 2013 ; Poole et al. , 2016 ; Martens et al. , 2021 ) . In this section , we introduce our notation and some key concepts used throughout . 2.1 KERNEL FUNCTION APPROXIMATION FOR WIDE NETWORKS . For simplicity , we start with the kernel function approximation for feedforward fully-connected networks , and discuss its extensions to convolutional networks and non-feedforward networks later . In particular , we will assume a network that is defined by a sequence of L combined layers ( each of which is an affine transformation followed by the elementwise activation function φ ) as follows : xl+1 = φ ( Wlx l + bl ) ∈ Rdl+1 , ( 1 ) with weights Wl ∈ Rdl+1×dl initialized as Wl iid∼ N ( 0 , 1/dl ) ( or scale-corrected uniform orthogonal matrices ( Martens et al. , 2021 ) ) , and biases bl ∈ Rdl+1 initialized to zero . Due to the randomness of the initial parameters θ , the network can be viewed as random feature model f lθ ( x ) , x ` at each layer l ( with x0 = x ) at initialization time . This induces a random kernel defined as follows : κlf ( x1 , x2 ) = 1 dl f lθ ( x1 ) > f lθ ( x2 ) . ( 2 ) Given these assumptions , as the width of each layer goes to infinity , κlf ( x1 , x2 ) converges in probability ( see Theorem 3 ) to a deterministic kernel κ̃lf ( x1 , x2 ) that has a simple form , and can be computed layer by layer as follows : Σl+1 = Ez∼N ( 0 , Σl ) [ φ ( z ) φ ( z ) > ] , with Σl = [ κ̃lf ( x1 , x1 ) κ̃ l f ( x1 , x2 ) κ̃lf ( x1 , x2 ) κ̃ l f ( x2 , x2 ) ] , ( 3 ) where κ̃0f ( x1 , x2 ) = κ 0 f ( x1 , x2 ) = x > 1 x2/d0 . 1Dynamical isometry is unavailable for ReLU ( Pennington et al. , 2017 ) , even with orthogonal weights . 2.2 LOCAL Q/C MAPS . By equation 3 , any diagonal entry ql+1i of Σ l+1 only depends on the corresponding diagonal entry qli of Σl . Hence , we obtain the following recursion for these diagonal entries , which we call q values : ql+1i = Q ( q l i ) = Ez∼N ( 0 , qli ) [ φ ( z ) 2 ] = Ez∼N ( 0,1 ) [ φ ( √ qliz ) 2 ] , with q0i = ‖xi‖2/d0 ( 4 ) where Q is the local Q map . We note that qli is an approximation of κlf ( xi , xi ) . Analogously , one can write the recursion for the normalized off-diagonal entries , which we call c values , as : cl+1 = C ( cl , ql1 , ql2 ) = E [ z1z2 ] ∼N ( 0 , Σl ) [ φ ( z1 ) φ ( z2 ) ] √ Q ( ql1 ) Q ( ql2 ) , with Σl = [ ql1 √ ql1q l 2c l √ ql1q l 2c l ql2 ] , ( 5 ) where C is the local C map and c0 = x > 1 x2/d0 . We note that cl is an approximation of the cosine similarity between f lθ ( x1 ) and f l θ ( x2 ) . Because C is a three dimensional function , it is difficult to analyze , as the associated q values can vary wildly for distinct inputs . However , by scaling the inputs to have norm √ d0 , and rescaling φ so that Q ( 1 ) = 1 , it follows that qli = 1 for all l. This allows us to treat C as a one dimensional function from [ −1 , 1 ] to [ −1 , 1 ] satisfying C ( 1 ) = 1 . Additionally , it can be shown that C possesses special structure as a positive definite function ( see Appendix A.4 for details ) . Going forward , we will thus assume that q0i = 1 , and that φ is scaled so that Q ( 1 ) = 1 . 2.3 EXTENSIONS TO CONVOLUTIONAL NETWORKS AND MORE COMPLEX TOPOLOGIES . As argued in Martens et al . ( 2021 ) , Q/C maps can also be defined for convolutional networks if one adopts a Delta initialization ( Balduzzi et al. , 2017 ; Xiao et al. , 2018 ) , in which all weights except those in the center of the filter are initialized to zero . Intuitively , this makes convolutional networks behave like a collection of fully-connected networks operating independently over feature map locations . As such , the Q/C map computations for a feed-forward convolutional network are the same as above . Martens et al . ( 2021 ) also gives formulas to compute q and c values for weighted sum operations between the outputs of multiple layers ( without nonlinearities ) , thus allowing more complex network topologies . In particular , the sum operation ’ s output q value is given by q = ∑n i=1 w 2 i qi , and its output c value is given by 1q ∑n i=1 w 2 i qici . In order to maintain the property that all q values are 1 in the network , we will assume that sum operations are normalized in the sense that ∑n i=1 w 2 i = 1 . Following Martens et al . ( 2021 ) , we will extend the definition of Q/C maps to include global Q/C maps , which describe the behavior of entire networks . Global maps , denoted by Qf and Cf for a given network f , can be computed by applying the above rules for each layer in f . For example , the global C map of a three-layer network f is simply Cf ( c ) = C ◦ C ◦ C ( c ) . Like the local C map , global C maps are positive definite functions ( see Appendix A.4 ) . In this work , we restrict our attention to the family of networks comprising of combined layers , and normalized sums between the output of multiple affine layers , for which we can compute global Q/C maps . For all the theorems in this work , we assume this family of networks .

Summary: This paper studied the problem of DNN training and generalization in vanilla architecture (without BN and Skip Connections in ResNets). It follows NTK theory and the approach of applying certain transformations to the activation functions. This work improves an existing work DKS, and solves its incompatibility to ReLU activations ("Leaky" ReLUs). This work introduces the necessary modifications to the Q/C map conditions for using Leaky ReLUs and shows empirical improvement over DKS or an easier method EOC on ImageNet.