Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers
Authors: Guodong Zhang
Aleksandar Botev
James Martens
Venue: ICLR 2022 Poster (2022)
Keywords: Neural Network Training
Kernel Approximation for Neural Networks
Neural Network Initialization
Generalization

Abstract: Training very deep neural networks is still an extremely challenging task. The common solution is to use shortcut connections and normalization layers, which are both crucial ingredients in the popular ResNet architecture. However, there is strong evidence to suggest that ResNets behave more like ensembles of shallower networks than truly deep ones. Recently, it was shown that deep vanilla networks (i.e.~networks without normalization layers or shortcut connections) can be trained as fast as ResNets by applying certain transformations to their activation functions. However, this method (called Deep Kernel Shaping) isn't fully compatible with ReLUs, and produces networks that overfit significantly more than ResNets on ImageNet. In this work, we rectify this situation by developing a new type of transformation that is fully compatible with a variant of ReLUs -- Leaky ReLUs. We show in experiments that our method, which introduces negligible extra computational cost, achieves validation accuracies with deep vanilla networks that are competitive with ResNets (of the same width/depth), and significantly higher than those obtained with the Edge of Chaos (EOC) method. And unlike with EOC, the validation accuracies we obtain do not get worse with depth.

Published as a conference paper at ICLR 2022
DEEP LEARNING WITHOUT SHORTCUTS: SHAPING THE KERNEL WITH TAILORED RECTIFIERS
Guodong Zhang Aleksandar Botev?, James Martens Uniiversity of Toronto, Veector Institute, 'DeepMind gdzhangécs toronto. edu, {botev, bamesmartensj@googll.
ABSTRACT
Training very deep neural networks is still an extremely challenging task. The common solution is to use shortcut connections and normalization layers, which are both crucial ingredients in the popular ResNet architecture. However, there is strong evidence to suggest that ResNets behave more like ensembles of shallower networks than truly deep ones, Recently, was shown that deep vanilla networks (i.e. networks without normalization layers or shortcut connections) can be trained as fast as ResNets by applying certain transformations to their activation functions. However, this method (called Deep Kernel Shaping) isn't 't fully compatible with ReLUs, and produces networks that overfit significantly more than ResNets on ImageNet. In this work, rectify this situation by developing transformation that is fully compatible with a variant of ReLUs - Leaky ReLUs. We show in experiments that our method, which introduces negligible extra com- putational cost, achieves validation accuracies with deep vanilla networks that are competitive with ResNets (of the same width/depth), and significantly higher than those obtained with the Edge of Chaos (EOC) method. And unlike with EOC, the validation accuracies we obtain do not get worse with depth.
type of
INTRODUCTION
0.75
Thanks many architectural and algorithmic innova- tions, the recent decade has witnessed the unprecedented success of deep learning in various high-profile chal- lenges, e.g., the ImageNet recognition task (Krizhevsky al., 2012), the challenging board game of Go (Sil- ver 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 (loffe & Szegedy, 2015; Ba et al., 2016) are two architectural components of modern networks that critically important for achieving fast training at very high depths, and feature prominently in the ubiquitous ResNet architecture of He et al. (2016b).
0.70
Validation
0.65
50 layers with TAT
0.60
layers with EDC 101 layers with TAT
layers with
0.55
200 layers with TAT
layers with
0.50
20000
40300
60000
Inúcc
1000co
Iteretions
Figure 1: Top-1 ImageNet validation accuracy of vanilla deep networks initialized using either EOC (with ReLU) or TAT (with LReLU) and trained with K-FAC.
Despite the success of ResNets, there 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). Moreover, ResNets without normalization layers could lose expressivity as the depth goes to infinity (Hayou al., 2021). 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-inefficien 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 define as networks without shorteut connections

---

Published as a conference paper at ICLR 2022
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 mcan-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) aallowed 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 units* 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. Qi et al. (2020) proposed to enforce the convolution kernels to be near isometric, but the gaps with ResNets are still significant on ImageNet. 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 method called Deep Kernel Shaping (DKS) for initializing and transforming networks 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) or Shampoo (Anil et al., 2020). However, their method isn't fully compatible with ReLUU,, 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 mcan-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. TAT enables very deep vanilla neural networks to be trained on ImageNet without the use of any additional architectural elements, while only introducing negligible extra computational cost. 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 even exceed the validation accuracy of standard ResNets of the same width/depth. A frame-ff open source implementation of DKS and TAT is available at https: //g thub. com/ deepmi nd/ dks.
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 sequence of L combined layers (each of which is an affine transformation followed by the elementwise activation function Φ) as follows:
1+1 = @ (Wiz² + bi) E Rd++1
(1)
with weights Wi E Rdd-1xd initialized as Wi ld matrices (Martens al., 2021)), and biases by Rd initialized to zero. Due to the randomness
~ (0, 1/d:) (or cale-ccoreettd uniform orthogonal
of the initial parameters 0, the network can be viewed as random feature model (x(r) 2 x at each
layer/ / (with x0 x) initialization time. This induces a random kernel defined follows:
x2) = (x1) fxz)//d.
(2)
Given these assumptions, as the width of each layer goes to infinity, (11 x2) converges in proba- bility (see Theorem 3) to a deterministic kernel nty (1),))))))) that can be computed layer by layer:
Ti (11, R (21,
E = EZ-N((. [)(=)((=)T] with Σ
(3)
(01,
(122
where &9 (11, .222 = x (21. 21 /2/
Dynamiccal isometry is unavailable for ReLU (Pennington et al., 2017), even with orthogonal weights
2

---

Published as a conference paper at ICLR 2022
2.2 LOCAL Q/C MAPS
By equation 3, any diagonal entry q, of E(+1 only depends on the corresponding diagonal entry g of Σ'. Hence, we obtain the following recursion for these diagonal entries, which we call values:
% 2(gh) (( [o(=)² EL-,((..))
σ 222)2
with of /|x:11//do
(4)
where @ is the local Q map. We note that of is an approximation of k'y (Ty, v;). Analogously, one can write the recursion for the normalized off-diagonal entries, which we call values, as:
1 s1 (NN(,,)))) (+=2)) with L'
= = (2)
(5)
V ((())
Voigge
where C is the local C map and cº I /2/do- We note that el is an approximation of the cosine similarity between fz()) ) and f22)) Because is a three dimensional function, it is difficult to analyze, as the associated values can vary wildly for distinct inputs. However, by scaling the inputs to have norm /dg, and rescaling ø so that 2(1) = 1, it follows that oj = for all 1. This allows us treat C as a one dimensional function from |-1, 1] to [-1, satisfying C(1) 1. Additionally, 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 qf 1, and that e is scaled so that @(1)
2.3 EXTENSIONS TO CONVOLUTIONAL NETWORKS AND MORE COMPLEX TOPOLOGIES
As argued in Martens et al. (2021), Q/C maps 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 fced-forward convolutional network are the same as above. Martens et al. (2021) also gives formulas to compute q and с 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 s output value is given by = >i=1 wfq.. and its output value is given by }
q E wfg. In order to maintain the property that all values are in the network, we will assume that sum operations are normalized in the sense that Σx= w
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 2, and Cj for a given network f. can be computed by applying the above rules for each layer in f. For example, the global C map a three-layer network f is simply (c) = o C o 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 implicitly assume this family of networks.
compute global Q/C maps. And all of our formal results will
2.4 Q/C MAPS FOR RESCALED RESNETS
ResNets consist of a sequence of residual blocks, each of which computes the sum of residual branch (which consists of a small multi-layer convolutional network) and a shortcut branch (which copies the block's input). In order to analyze ResNets we will consider the modified version used in Shao et al. (2020) and Martens et al. (2021) which removes the normalization layers found in the residual branches, and replaces the sum the end of each block with normalized sum. These networks, which we will call rescaled ResNets, are defined by the following recursion:
wx' v 2 RIx)),
(6)
where R is the residual branch, and w is the shortcut weight (which must be in [-1, 1D. Applying the previously discussed rules for computing Q/C maps, we get of for all I and
I+1 = wdd + (1 11--"a Cx
3 EXISTING SOLUTIONS AND THEIR LIMITATIONS
Global Q/C maps can be intuitively understood as a way of characterizing signal propagation through the network, f at initialization time. The q value approximates the squared magnitude of the activation
3

---

Published as a conference paper at ICLR 2022
vector, that Q describe the contraction or expansion of this magnitude through the action of f. On the other hand, the с value approximates the cosine similarity of the function values for different inputs, that C describes how well f preserves this cosine similarity from its input to its output.
Standard initializations methods (LeCun et al., 1998; Glorot & Bengio, 2010; He et al., 2015) are motivated through an analysis of how the variance of the activations evolves throughout the network. This can be viewed as a primitive form of Q map analysis, and from that perspective, these methods are trying ensure that q values remain stable throughout the network by controlling the local Q map. This is necessary for trainability, since very large or tiny q values can cause numerical issues, saturated activation functions (which have implications for C maps), and problems with scale-sensitive losses. C map is also necessary
However, as was first observed by Schoenholz et al. (2017), well-behaue for trainability. When the global map is close to constant function (i.e. degenerate) on (-1,1) 1), which easily happens in deep networks (as discussed in Appendix A.2), this means that the network's output will appear either constant or random looking, and won't convey any useful information about the input. Xiao et al. (2020) and Martens et al. (2021) give more formal arguments for why this leads to slow optimization and/or poor generalization under gradient descent.
Several previous works (Schoenholz et al., 2017; Yang & Schoenholz, 2017: Hayou et al., 2019) attempt to achieve well rell-behaved global C map by choosing the variance of the initial weights and biases in each layer such that (1)) = 1 procedure which is referred to Edge of Chaos (EOC). However, this approach only slows down the convergence (with depth) of the с values from exponen- tial to sublinear (Hayou et al., 2019), and does not solve the fundamental issue of degenerate global C maps for very deep networks. In particular, the global C map of a deep network with ReI.U and EOC initialization rapidly concen- trates around 1 as depth increases (see Figure 2). While EOC allows very deep vanilla networks to be trained, the training speed and generalization performance is typically much worse than for comparable ResNets. Klambauer et al. (2017) applied an affine transformation to the output of activation functions to achieve Q(1) 1 and C(O) = 0, while Lu et al. (2020) applied them to achieve (11) = 1 and (1( = 1, although the effect of both approaches is similar to EOC.
0.2
output value
ReLu with depth 10
RoLU with depth 50
94
TRoLU with depth
TReLU with depth 10
TRoLU with depth 50
TReLU with depth 1000
"5.00 50.00 -0.75
-0.50
-0.25
0.00 0.25 0.50 0.75 1.00
input value
Figure 2: Global maps for ReLU networks (EOC) and TReLU networks (0) = 0.5). The global C map of TReL.U network con- verges to well-behavior function as depth increases (proved in Proposition 3).
To address these problems, Martens et al. (2021) introduced DKS, which enforces the conditions (0) = 0 and C2 (1) = ( (for some modest constant < > 1) directly on the network's global C map Cy. They show that these conditions, along with the positive definiteness of C maps, cause Cy to be close to the identity and thus well-behaved, In addition to these C map conditions, DKS enforces that Q(1) = 1 and (1(1) = which lead to constant values of in the network, and lower kernel approximation error (respectively). DKS enforces these Q/C map conditions by applying model class- presssrreerrinn transformation ((x) = ((ó(zz + 3) + 5). with non- -ra-trriieble parameters a, 3, ? and 8. The hyperparameter Ç is chosen to be sufficiently greater than 1 (e.g. 1.5) in order to prevent the transformed activation functions from looking "nearly linear" (as they would be exactly linear <= 1). which Martens et al. (2021) argue makes it hard for the network to achieve nonlinear behavior during training. Using DKS, they were able to match the training speed of ResNets on ImageNet with vanilla networks using K-FAC. However, DKS is not fully compatible with ReLUs, and the networks in their experiments fell substantially short of ResNets terms of generalization performance.
4 TAILORED ACTIVATION TRANSFORMATION (TAT)
The reason why DKS is not fully compatible with ReLUs is that they are positive homogeneous, iec ø(az)= = aσ(() for a 2 O. This makes the parameter of the transformed activation function redundant, thus reducing the degrees of freedom with which to enforce DKS's four Q/C map conditions. Martens et al. (2021) attempt to circumvent this issue by dropping the condition Q'
1, which leads to vanilla deep networks that are trainable, but slower to optimize compared to using DKS with other activation functions. This is a significant drawback for DKS, as the best generalizing deep models often use ReLU-family activations. We therefore set out to investigate other possible remedies either in the form of different activation functions, new Q/C map conditions, both. To this end, we adopt a ReLU Rammlly activation function with an extra degree of freedom (known as "I "Leaky ReL.U*), and modify the Q/C map conditions in order preserve certain desirable properties
4

---

Published as a conference paper at ICLR 2022
Table 1: Comparison of different methods applied to a network f.
EOC (smooth)
EOC (LReLU)
DKS
TAT (smooth)
TAT (LReLU)
((1) =1
2(1==
((q) = q
pe exists
((g) =
(((1= =1
g(1)=
C₁(1) = 1
(1(1) = 1
Cy(0) =0
c/(1==
(1()) 5 C](1==T Cy(0) 1)
of this choice. The resulting method, which we name Tailored Activation Transformation (TAT)
achieves competitive generalization performance with ResNets in our experiments.
4.1 TAILORED ACTIVATION TRANSFORMATION FOR LEAKY RELUS
One way of addressing the issue of DKS's partial incompatibility with ReLUS is to consider a slightly different activation functiin namely the Leaky ReLU (LRcLU) (Maas et al., 2013):
)a = m 0} + a min 0}.
(8)
where a is the negative slope parameter. While using LReLUs with a + 0 in place of ReLUs changes the model class, it doesn't 't limit the model's expressive capabilities compared ReLU, as assuming o +1, one can simulate a ReLU network with a LReLU network of the same depth by doubling the number of neurons (see Proposition 4). Rather than using a fixed value for ax we will use it as
an extra parameter to satisfy our desired Q/C map conditions. Define da(x) =V
V I+ax 1aa= By
Lemma 1, the local Q and C maps for this choice of activation function are:
(1-x)² (()) = q and C(c) = c =(I+02)
(VI-c² ccos ccccccsss (c)) -
Note that the condition Q(q) = is actually stronger than DKS's Q map conditions (Q(1) and @'(1) 1), and has the potential to reduce kernel approximation errors in finite width networks compared to DKS, as it provides better guarantee the stability of O, w.r.t. random perturbations of the q values at each layer. Additionally, because the form of does not depend on either of the layer's input q values, it won't be affected by such perturbations at all. (Notably, if one uses the negative slope parameter to transform LReLUs with DKS, these properties will not be achieved.) In support of these intuitions is the fact that better bounds the kernel approximation error exist for ReI.U networks than for general smooth ones (as discussed in Appendix A.
Another consequence of using 6a (z) for our activation function is that we have C' (1(1)= = as in EOC. If combined with the condition C(O) O (which is used to achieve (0) 0 in DKS) this would imply by Theorem that C is the identity function, which by equation 9 is only true when a = 1, thus resulting in linear network. In order to avoid this situation, and the closely related one where 6. appears "nearly linear", we instead choose the value of a so that Cy 17, for a hyperparameter 0 < 17 S 1. As shown in the following theorem, controls how close C is to the identity, thus allowing us to achieve well-behaver global C map without making éa nearly linear:
Theorem 1. For network with d. (ar) as its activation function (with 2 0), we have
max (C((c) - cl S min (4C, 1 C,()), cE[-1,1) cE[-1,1]
max [C)(()) S min (4Cr 13 (10)
Another motivation for using (x) as an activation function is given by the following proposition:
Proposition 1. The global C map of feedforward network with (x)x) as its activation function
equal to that of a rescaled ResNet of the same depth (see Section 2.4) with normalized ReLU
activation <(x) = /2 max(x, 0), shortcut weight v /Ita. and residual branch R consisting of combined layer (or just a normalized ReLU activation) followed by an affine layer.
This result implies that at initialization, vanilla network using &a behaves similarly to ResNet, property that is quite desirable given the success that ResNets have already demonstrated.
In summary, we have the following three conditions:
Q(q) q. C,(1) 11, Cy(0) =1,
(11)
which we achieve by picking the negative slope parameter so that Cy(0) = 7. We define the Tailored Rectifier (TReLU to be 3h with a chosen in this way. Note that the first two conditions
5

---

Published as a conference paper at ICLR 2022
also true when applying the EOC method to LReLUs, and its only the third which sets TReLU apart. While this might seem like a minor difference, it actually matters a lot to the behavior of the global C map. This can be seen in Figure 2 where the value quickly converges towards 1 with depth under FOC, resulting in a degenerate global C map. By contrast, the global C map of TReLU for a fixed " converges rapidly to nice function, suggesting very deep vanilla network with TReLU has the same well-behaved global C map as shallow network. We prove this in Proposition 3 by showing the local C map in equation converges to an ODE as we increase the depth. For direct comparison of all Q/C map conditions, we refer the readers to Table 1.
For the hyperparameter S 77 S 1, we note that a value very close will produce network that is "nearly linear", while a value very close to I will give rise to degenerate C map. In practice we use 7) = 0.9 or 0.95, which seems to work well in most settings. Once we decide on 77, we can solve the value a using binary search by exploiting the closed-form form of C in equation 9 to efficiently compute (O). For instance, f is 100 layer vanilla network, one can compute (0) as follows:
100 times
Cy(0)= =
C(O),
(12)
which is a function of a. This approach can be generalized more advanced architectures, such as rescaled ResNets, discussed in Appendix B.
4.2 TAILORED ACTIVATION TRANSFORMATION FOR SMOOTII ACTIVATION FUNCTIONS
Unlike LReLU, most activation functions don't have closed-form formulas for their local C maps. As result, the computation of (O) involves the numerical approximation of many two-dimensional integrals high precision (as in equation 5), which can be quite expensive. One alternative way to control how close Cy is to the identity, while maintaining the condition second derivative (1). The validity of this approach is established by the following theorem:
(1) = 1, is to modulate its
Theorem 2. Suppose f is a network with smooth activation function. Cf (1) 1. then we have
max Cy(c) c s (1), cE|-1.1]
max cEI-1.11
(g((c) 11 (1)
(13)
Given C(1) = 1 and (1(1) = 1, straightforward computation shows that (1) (1) if f is an L-layer vanilla network. (See Appendix B for discussion of how to do this computation for more general architectures.) From this we obtain the following four local Q/C map conditions:
@(1) 1, C((1==1 1,
=/L, C'(1)= 11
(14)
To achieve these we adopt the same activation transformation as DKS: 6(z) Vó(ax + 3) + 8) for non-trainable scalars a, B, 8, and 7. We emphasize that these conditions cannot be used with LReLU, as I.Rel.U networks have C" (1) co. By equation 4 and basic properties of expectations, we have
1 = 2(1 1.,.N(
[o(=)]] = = [[ő(az B) 8)21
(15)
so that 7 12 [[ó(az + 9) + -3/2 To obtain the values for a, 8 and 8, we can treat the remaining conditions as a threc-dimensional nonlinear system, which can be written as follows:
E [)<)) 2(11) 1,
(16)
E.-A(o.)) (=')) C* (1) T/L, Be-A((.)) [0(;)]] = CL( =1
We do not have closed-form solution of this system. However, each expectation is one dimensional integral, and so can be quickly evaluated to high precision using Gaussian quadrature. One can then black-box nonlinear equation solvers, such modified Powell's method (Powell, 1964), obtain a solution. See https: / / g thub. com/ deepmi nd/dks. dks for a complete implementation.
5 EXPERIMENTS
Our main experimental evaluation of TAT and competing approaches is on training deep convolutional networks for ImageNet classification (Deng et al., 2009). The goal of these experiments is not to achieve tfate-of-the-a but rather to compare TAT as fairly as possible with existing methods, and standard ResNets in particular. To this end, we use ResNet V2 (He et al., 201 6b)) as the main reference

---

Published as a conference paper at ICLR 2022
architecture, from which we obtain rescaled ResNets (by removing normalization layers and weighing the branches as per equation 6), and vanilla networks (by further removing shortcuts). For networks add dropout to the penultimate layer for regularization, was done
without batch normalization, in Brock et al. (2021b). We train the models with 90 epochs and a batch size of 1024, unless stated otherwise. For TReLU, for all methods is the Orthogonal Delta initialization, with an extra multiplier given by Tw- We initialize biases iid from (0, of)). We use (o Tb) (1, 0) in all experiments (unless explicitly stated otherwise), with the single exception that we use (ow.; oT)) = (v2, 0) standard ResNets, as per standard practice (He et al., 2015). For all other details see Appendix D.
obtain y by grid search in (0.9, 0.95). The weight initialization used
5.1 TOWARDS REMOVING BATCH NORMALIZATION
Two crucial components for the successful training of very deep neural networks are shortcut connections and batch normalization (BN) layers. As argued in De & Smith (2020) and Shao et al. (2020), BN implicitly biases the residual blocks toward the identity function, which makes the network better behaved at initialization time, and thus easier to train. This suggests that one can compensate for the removal of BN layers, at least in terms of their effect on the behaviour of the network at initialization time, by down-scaling the residual branch of each residual block. Arguably, almost all recent work on training deep networks without normalization layers (Zhang al., 2018 8; Shao et al., 2020; Bachlechner et al., 2020; Brock et al., 202 has adopted this idea by introducing multipliers on the residual branches (which may or may not be optimized during training).
In Table 2, we show that one can close most of the gap with standard ResNets by simply adopting the modification in equation 6 without using BN lay- ers. By further replacing ReI.U with TReLU, we can exactly match the per- formance of standard ResNets. With K-FAC as the optimizer, the rescaled ResNet with shortcut weight w 0.9 is only 0.5 shy of the validation accuracy (76.4) of the standard ResNet. Further replacing ReLU with TReLU, we match the performance of standard ResNet with shortcut weight ws 0.8.
Table 2:
Top-1 validation accuracy of rescaled ResNet50 with varying shorteut weights. We set 17 = 0.9 for TReLU.
Standard ResNet
Rescaled ResNet (w)
Optimizer
Activation
0.0
0.5
0.8
0.9
ReLUU TReLUU
72.6
74.5
75.6
75.9
K-FAC
76.4
74.6
75.5
76.4
75.9
ReLU
63.7
72.4
73.9
75.0
SGD
76.3
TReLU
71.0
72.6
76.0
74.8
5.2 THE DIFFICULTY OF REMOVING SHORTCUT CONNECTIONS
While the aforementioned works have shown that it possible achieve competitive results without nor- malization layers, they all rely on the use of shortcut connections to make the network look more linear initialization. A natural question to ask is whether
Table 3: ImageNet 1 validation accuracies of shortcut-free networks on ImageNet.
Depth
Optimizers
vanilla
BN
IN
K-FAC
72.6
72.8
72.7
50
SGD
63.7
72.6
58.1
normalization layers could compensate for the re- moval of shortcut connections. We address this ques- tion by training shortcut-free networks with either BN or Layer Normalization (LN) layers. As shown in Table 3, these changes do not seem to make significant difference, especially with strong optimizers like K-FAC. These findings are in agreement with the analyses of Yang et al. (2019) and Martens et al. (2021), who respectively showed that deep shortcut- free networks with BN layers still suffer from exploding gradients, and deep shortcut- free networks with LN layers still have degenerate C maps.
K-FAC
71.8
67.6
72.0
101
SGD
41.6
43.4
28.6
5.3 TRAINING DEEP NEURAL NETWORKS WITHOUT SHORTCUTS
The main motivation for developing TAT is to help deep vanilla networks achieve generalization performance similar to standard ResNets. our investigations we include rescaled ResNets with a shorteut weight of either 0 (i.e. vanilla networks) or 0.8. In Table 4 we can see that with strong optimizer like K-FAC, we can reduce the gap on the 50 layer network to only 1.8% accuracy when training for 90 epochs, and further down to 0.6% when training for 180 epochs. For 101 layers, the gaps are 3.6% and 1.7% respectively, which we show can be further reduced with wider networks (see Table 9). To our knowledge, this is the first time that deep vanilla network has been trained to such a high validation accuracy on ImageNet. In addition, our networks have fewer parameters and run faster than standard ResNets, and use less memory at inference time due the removal of
7

---

Published as a conference paper at ICLR 2022
Table 4: ImageNet top-1 validation accuracy. For rescaled ResNets (w 0.0 or 0.8), we do not include any normalization layer. For standard ResNets, batch normalization is included. By default, ReLU activation used for standard ResNet while we use TReLU for rescaled networks.
90 epochs
180 epochs
Depth
Optimizer
ResNet
0.0
0.8
ResNet
0.0
0.8
K-FAC
76.4
74.6
76.4
76.6
76.0
77.0
50
SGD
76.3
71.0
76.0
76.6
72.3
76.8
K-FAC
77.8
74.2
77.8
77.6
75.9
78.4
101
SGD
77.9
70.0
77.3
77.6
73.8
77.4
shortcut connections and BN layers. The gaps when using SGD as the optimizer are noticeably larger, which we further explore in Section 5.5. Lastly, using rescaled ResNets with a shorteut weight of 0.8 and TReLU, we can exactly match or even surpass the performance of standard ResNets.
5.4 COMPARISONS WITH EXISTING APPROACHES
Comparison with EOC. Our first comparison
Table 5: ImageNet top-1 validation accuracy compari- son between EOC and TAT on deep vanilla networks.
is between TAT and EOC on vanilla deep net- works. For EOC with ReLUs we set (ow, ob)
Depth
Optimizer
Method
(L.)eelUU
Tanh
(V2,0) 0) achieve ((1) = 1 as in He et al. (2015), since ReLU networks always satisfy ('(1)= whenever a O. For Tanh acti- vations, comprehensive comparison with EOC is more difficult, as there are infinitely many choices of (6) that achieve C'(1)
EOC
72.6
70.6
K-FAC
TAT
74.6
73.1
50
BOC
63.7
55.7
SGDD
TAT
71.0
69.5
EOC
71.8
69.2
1. Here we use (Tws σ6) = (1.302, 0.02)". as sug- gested in Hayou et al. (2019). In Table 5, we can see that in all the settings, networks constructed with TAT outperform EOC-initialized networks by significant margin, especially when using SGD. Another observation is that the accuracy of EOC-initialized networks drops as depth increases.
K-FAC
TAT
74.2
72.8
101
BOC
41.6
54.0
SGD
TAT
70.0
69.0
Comparison with DKS. The closest approach to TAT in the existing literature is DKS, whose similarity and drawbacks are discussed in Section 4. We compare TAT to DKS on both LReLUs", and smooth functions like the SoftPlus and Tanh. For smooth activations, we perform grid search over f0.2, 0.3, 0.5} for T in TAT, and {1.5, 10.0, 100.0} for Ç in DKS, and report only the best performing one. From the results shown in Table 7, we observe that TAT, together with LReLU (i.e. TReLU), performs the best in nearly all settings we tested, and that its advantage becomes larger when we remove dropout. One possible reason for the superior performance of TReLU networks the stronger Q/C map conditions that they satisfy compared to other activations (i.e. @(q) g for all q vs Q(1) = and (1( = 1, and invariance of to the input q value), and the extra resilience to kernel approximation eтror that these stronger conditions imply. In practice, we found that TReLU indeed has smaller kernel approximation error (compared to DKS with smooth activation functions, see Appendix E. and works equally well with Gaussian initialization (see Appendix E.7).
Comparison with PReLU. The Parametric ReLU (PReL.U) introduced in He al. (2015) differs from LReLU by making the negative slope a trainable parameter. Note that this is distinct from what we are doing with TReLU, since there we compute the negative slope pa- rameter ahead of time and fix it during training. In our comparisons with PReLU we consider two different initializations: O (which recovers the standard ReLU),, and 0.25, which was used in He et al. (2015). We report the results on deep vanilla networks Table (see Appendix E.6 for results rescaled ResNets). For all settings, emphasizing the importance of the initial negative slope value. In principle, these two methods can be combined together (i.e. we could first initialize the negative slope parameter with TAT, and then optimize during training), however we did not see any benefit from doing this in our experiments.
Table 6: Comparison with PRcLU.
Depth
Optimizer
TReLU
PReLUb.
2ReLU
K-FAC
74.6
72.5
73.6
50
SGD
71.0
66.7
67.9
K-FAC
74.2
71.9
72.8
101
SGD
70.0
$4.3
66.3
method outperforms PReLU by large margin,
We also ran experiments with (aw, σ2) = (1.0, 0.0), and the scheme described in Pennington al. (2017) and Xiao et al. (2018) for dynamical isometry. The results were worse than those reported the table.
For DKS, we set the negative slope a parameter and adopt the transformation 6(x) 7 (a + (8) 6).
8

---

Published as a conference paper at ICLR 2022
Table 7: Comparisons between TAT and DKS. The numbers on the right hand of / are results without dropout. The methods with are introduced in this paper.
Shortcut Weight
TAT
DKS
Depth Optimizer
LReLU
SoftPlus
Tanh"
LReLU
SoftPlus
Tanh
0.0
74.6/74.2
74.4/74.2
73.1/72.9
74.3/74.3
74.3/73.7
72.9/72.9
K-FAC
u 0.8
76.4/75.9
76.4/75.0
74.8/74.4
76.2/76.2
76.3/75.1
74.7/74.5
50
0.0
71.1/71.1
70.2/70.0
69.5/69.5
70.4/70.4
71.8/71.4
69.2/69.2
SGD
W 0.8
76.0/75.8
74.3/73.8
72.4/72.2
73.4/73.0
75.2/74.1
72.8/72.8
W 0.0
74.2/74.2
74.1/73.4
72.8/72.5
73.5/73.5
73.9/73.1
72.5/72.4
K-FAC
0.8
77.8/77.0
76.6/75.7
75.8/75.1
76,8/76.7
76.8/75.6
75.9/75.7
101
W 0.0
70.0/70.0
70.3/68.8
69.0/67.8
68.3/68.3 74.9/74.6
68.3/68.3
69.8/69.8
SGD
0.8
77.3/76.0
75.3/75.3
73.8/73.5
76.3/75.1
74.6/74.6
5.5
THE ROLE OF THE OPTIMIZER
Table 8: Batch size scaling.
One interesting phenomenon we observed in our experiments, which echoes the find- ings of Martens et al. (2021), is that a strong optimizer such as K-FAC significantly out- performs SGD on vanilla deep networks in terms of training speed. One plausible expla- nation is that K-FAC works better than SGD in the large-batch setting, and our default batch size of 1024 is already beyond SGD's "critical batch size", at which scaling efficiency begins to drop. Indeed, it shown by Zhang et al. (2019) that optimization algorithms that employ preconditioning, such as Adam and K-FAC, result in much larger critical batch sizes.
Batch size
Optimizer
128
256
512
1024
2048
4096
K-FAC
74.5
74.4
74.5
74.6
74.2
72.0
SGD
72.7
72.6
72.7
71.0
69.3
62.0
LARS
72.4
72.3
72.6
71.8
71.3
70.2
0.80
To investigate this further, we tried batch sizes between 128 and 4096 for training 50-layer vanilla TReLU net- works. As shown in Table 8, K-FAC performs equally well for all different batch sizes except 4096 (where we see increased overfitting), while the performance of SGD starts to drop when we increase the batch size past 512. Surprisingly. we observe similar trend for the LARS optimizer (You et al., 2019), which was designed for large- batch training. Even at the smallest batch size we tested (128), K-FAC still outperforms SGD by a gap of 1.8% within our standard epoch budget. We conjecture the rea- son behind this to be that vanilla networks without normalization and shortcuts give rise to loss landscapes with worse curvature properties compared to ResNets, and that this slows down simpler optimizers like SGD. To investigate further, we also ran SGD (with batch size of 512) and K-FAC for up to 360 epochs with a "one-cycle" cosine learning rate schedule (Loshchilov & Hutter, 2016) that decreases the learning rate to to 0 by the final epoch. As shown in Figure 3, SGD does indeed eventually catch up with K-FAC (using cosine scheme), requiring just over double the number of epochs achieve the same validation accuracy. While one may argue that K-FAC introduces additional computational overhead at each step, thus making a head-to-head comparison versus SGD unfair, we note that this overhead can amortized by not updating K-FAC's preconditioner matrix at every step. In our experiments we found that this strategy allowed K-FAC to achieve similar per-step runtime to SGD, while retaining its optimization advantage on vanilla networks. (See Appendix E.3.)
Valedation arc
0.70
EFAC with epoctt
5/2 with 90 eprchi
RTAC with 180 epcchs
SCD with 100 epochs
K7AC With 260 specto
5GD with 350 opocns
0:35
Cases processed
Figure 3: Training speed comparison be- tween K-FAC and SGD 50 layer vanilla TReLU network.
6 CONCLUSIONS
In this work we considered the problem of training and generalization in vanilla deep neural networks (i.e. those without shortcut connections and normalization layers). To address this we developed a novel method that modifies the activation functions in way tailored to the specific architecture, and which enables us to achieve generalization performance on par with standard ResNets of the same width/depth. Unlike the most closely related approach (DKS), our method is fully compatible with ReLU-family activation functions, and in fact achieves its best performance with them. By obviating the need for shorteut connections, we believe our method could enable further research into deep models and their representations. In addition, our method may enable new architectures to be trained for which existing techniques, such as shortcuts and normalization layers, are insufficient.

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
Here we discuss our efforts to facilitate the reproducibility of this paper. Firstly, we have made an open Python implementation of DKS and TAT, supporting multiple tensor programming frameworks, available at https: /gi sitthub. com/ deepmind/ d/k.. Secondly, we have given all important details of our experiments in Appendix D.
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, arXiv preprint arXiv:2002.09018, 2020.
Jimmy Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXhv:1607.06450 2016.
arXiv preprint
Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using kronecker- factored approximations. In International Conference on Learning Representations, 2017.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. arXiv preprint arXiv:2003.04887, 2020.
David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? In International Conference on Machine Learning, pp. 342 350. PMLR, 2017.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://gite tttpbb.tttuutttgg com/ homm/google//// jax.
Andrew Brock, Soham De, and Samuel L Smith. Characterizing signal propagation to close the per- formance gap in unnormalized resnets. In International Conference Learning Representations, 2021a.
Andrew Brock, Soham De, Samuel L. Smith, and Karen Simonyan. Brigh-performacce large-scale image recognition without normalization, arXiv preprint arXiv:2102.06171, 2021b.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. arXiv preprint 2rXiv:2005. 2020.
Sam Buchanan, Dar Gilboa, and John Wright. Deep networks and the multiple manifold problem. In International Conference on Learning Representations, 2020.
Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. Advances in Neural Informa- tion Processing Systems, 22:342-350, 2009.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and dual view expressivity. Advances In Neural Information Processing Systems, 11-2253-2261 2016.
Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function in deep networks. Advances in Neural Information Processing Systems, 33, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Xiaohan Ding, Xiangyu Zhang, Ningning Ma. Jungong Han. Guiguang Ding, and Jian Sun. Repvgg: Making styletyle convnets great again. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 13733-137423 2021.
10

---

Published as a conference paper at ICLR 2022
David Duvenaud, Oren Rippel, Ryan Adams, and Zoubin Ghahramani. Avoiding pathologies in very deep networks. In Artificial Intelligence and Statistics, pp. 202-210. PMLR, 2014.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the impact of the activation function on deep neural networks training. International conference on machine learning, pp. 2672-2680.
PMLR, 2019.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, pp. 1324 1332. PMLR, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pp. 630-645. Springer, 201
Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin. Haiku: Sonnet for JAX, 2020.
URL http::ttt/..... / / thub. com/ deepm nd 3am-hai ku.
Matteo Hessel, David Budden, Fabio Viola, Mihaela Rosca, Eren Sezener, and Tom Hennigan. Optax: composable gradient transformation and optimisation, in jax!, 2020. URL http://ttt/// thub. ccom/deppe nd/ opt.ax.
Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448 456.. PMLR, 2015.
Hassan K. Khalil. Nonlinear systems third edition. 2008
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self Self-normalizing neural networks. Advances in neural information processing systems, 30, 2017.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Advances in neural information processing systems, 2012.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade. Springer, 1998.
Iyya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts, arXiv preprint arXiv:1608.03000 2016.
Yao Lu, Stephen Gould, and Thalaiyasingam Ajanthan. Bidirectional self self-normalizing neural networks. arXiv preprint arXiv:2006. 12169, 2020.
Andrew I. Maas, Awni Y Hanmun, and Andrew Y Ng. Rectifier nonlinearities improve neural network acoustic models. International Conference on Machine Learning, 2013.
James Martens. On the validity of kernel approximations for orthogonally- initialized neural networks. arXiv preprint arXiv: 00480587 2021.
11

---

Published as a conference paper at ICLR 2022
James Martens and Roger Grosse. Optimizing neural networks with Aroneceer-factored approximate curvature. In International conference on machine learning, pp. 2408-2417 PMLR, 2015.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Rapid training deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprini arXiv: 21 10.01765, 2021.
Radford M Neal. Bayesian learning for neural networks. Lecture notes in statistics, 118, 1996.
Oyebade K Oyedotun, Djamila Aouada, Bjöm Ottersten, al. Going deeper with neural networks without skip connections. In 2020 IEEE International Conference on Image Processing (ICIP), pp. 1756-1760. IEEE, 2020.
Jeffrey Pennington, Samuel S Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 7788-44995 2017.
Ben Poole. Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity deep neural networks through transient chaos. Advances in neural information processing systems, 29:3360-3368, 2016.
Michael JD Powell. An efficient method for finding the minimum of function of several variables without calculating derivatives. The Computer Journal, 7(2):155-162, 1964.
Haozhi Qi, Chong You, Xiaolong Wang, Yi Ma, and Jitendra Malik. Deep isometric learning for visual recognition. In International Conference on Machine Learning, pp. 7835. PMLR, 2020.
Andrew M Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 2013.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.
Jie Shao, Kai Hu, Changhu Wang, Xiangyang Xue, and Bhiksha Raj. Is normalization indispensable for training deep neural network? Advances in Neural Information Processing Systems, 33, 2020.
David Silver, Julian Schritwiessr, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lueas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676)354-3397 2017.
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint 1rX5vv1505.00387, 2015.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-2826, 2016.
Andreas Veit, Michael J Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow networks. Advances in neural information processing systems, 29:550-558. 2016.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Penning- ton. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, 2018.
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generaliza- tion in deep neural networks. In International Conference on Machine Learning, pp. 10462-10472. PMLR, 2020.
Greg Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. In Advances in Neural Information Processing Systems, volume 30, 2017.
12

---

Published as a conference paper at ICLR 2022
Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. A mean field theory of batch normalization. ArXiv, abs/1902.08129, 2019.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: /904,0002, 2019.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 2016. British Machine Vision Association, 2016.
Guodong Zhang. Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 2019.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In International Conference on Learning Representations, 2018
13

---

Published as a conference paper at ICLR 2022
A BACKGROUND
A. KERNEL FUNCTION APPROXIMATION ERROR BOUNDS
In Section 2.1, we claimed that the kernel defined in equation would converge to deterministic kernel, as the width of each layer goes to infinity. To be specific, one has the following result bounding the kernel approximation error.
Theorem 3 (Adapted from Theorem 2 of Daniely et al. (2016)). Consider fully--onnected network of depth with weights initialized independently using a standard Gaussian fan-in initialization. Further suppose that the activation function & is C-bounded (i.e. ||0l| I C, l ll S C and || I C for some constant C) and satisfies Eg-NN(. (0.1))(())))) 1, and that the width of each layer is greater than or equal to (4C*)² log(ST./6) /e2 Then at initialization time, for inputs Dn and 22 satisfying
||>|||" || dim(x1), we have that
| (21, (21, Se
with probability at least
1-
The bound in Theorem 3 predicts an exponential dependence on the depth L of the minimum required width of each layer. However, for a network with ReLU activations, this dependence only quadratic in L, as is established in the following theorem:
Theorem (Adapted from Theorem of Daniely et al. (2016)). Consider a fully-cconeccte network of depth L with ReLU activations and weights initialized independently using He initialization (He et al,, 2015), and suppose that the width of each layer is greater than or equal to I.2 log T/ / Then at initialization time, for inpuls 2ı and 22 satisfying | [2=2|l2 dim(x1), and S T
we have that
(2)()1)22) (+(2), Se
with probability at least 1 1-
According Lemma D.1 of Buchanan al. (2020), the requirement of the width for ReLU networks could further be reduced to linear in the depth T., but with a worse dependency on 6.
Although Theorems 3 and are only applicable to Gaussian initializations, a similar bound has been given by Martens (2021) for scaled uniform orthogonal initializations in the case that L = 1. Moreover, Martens (2021) conjectures that their result could extended to general values of L.
A.2 DEGENERATE MAPS FOR VERY DEEP NETWORKS
Daniely et al. (2016), Poole et al. (2016), and Martens et al. (2021) have shown that without very careful interventions, C maps inevitably become "degenerate" in deep networks, tending rapidly towards constant functions on (-1, 1) as depth increases. The following proposition is restatement of Claim from Daniely et al. (2016):
Proposition 2. Suppose is a deep network consisting of composition of L combined layers. Then for all e E (-1,1) we have
lim L-+00
for some c* E {[0,10 1].
While the above result doesn't characterize the rate of convergence (c) to constant function, Poole et al. (2016) show that if (1()) * 1, happens exponentially fast as a function of L in the asymptotic limit of large L. Martens et al. (2021) gives a similar result which holds uniformly for all L, and for networks with more general repeated structures.
A.3 C. MAP DERIVATIVE
Poole et al. (2016) gave the following nice formula for the derivative of map of a combined layer with activation function Øi
41
(d (vgizi) 8' (VTa (czn + vI vI 22)))))) (17)
C' C'(cqq1.9)) 92)
E B
~((.1)
((g))
14

---

Published as a conference paper at ICLR 2022
For rigorous proof of this result we refer the reader to Martens et al. (2021).
One can iterate this formula to obtain a similar equation for higher-order derivatives:
(1/2)) Ex N(0.1) [o(0] (V1121) 6( (V02 (es) + v1-c22)) C (c. 92) VQ(ai ((()
(18)
A.4 SOME USEFUL PROPERTIES OF MAPS
In this section we will assume that Q Q2 1.
Observe that (1) = Bz-M(0.1) (x)) and that C maps 10 |-1, (which follows from its interpretation as computing cosine similarities for infinitely wide networks). Moreover, C is positive definite function, which means that it can be written as _ c t for b,, 2 0 (Daniely al., 2016; Martens al., 2021). Note that for smooth activation functions, positive definiteness can be easily verified by Taylor-expanding ((c) about 0 and using
c (0) (e..)((.1) (000) (=) 2 O.
(19)
As discussed Section 2.3, global C maps are computed by recursively taking compositions and weighted averages (with non-negative weights), starting from C. Because all of the above properties are preserved under these operations, it follows that global maps inherit them from C.
B ADDITIONAL DETAILS AND PSEUDOCODE FOR ACTIVATION FUNCTION TRANSFORMATIONS
B. TAKING ALL SUBNETWORKS INTO ACCOUNT
In the main text of this paper we have used the condition C) Ç DKS, C₁(0) = 17 in TAT for Leaky ReLUs, and (1) in TAT for smooth activation functions. However, the condition used by Martens al. (2021) in DKS was actually (1) = C. where u² is the so-called "maximal slope function":
(C((
max C2 (1), 3:95
where "g C f" denotes that g is subnetwork of f. (That C2 (1) is fully determined by C1( follows from the fact that Cg can be written in terms of compositions, weighted average operations, and applications of C, and that C maps always preserve the value 1. Using the chain rule, and the linearity of derivatives, these facts allow one to write C, 1) as a polynomial function of (1(1).)
The motivation given by Martens et al. (2021) for looking at (1) over all subnetworks S f (instead of just (1)) is that we want all layers of f. in all of its subnetworks, to be readily trainable. For example, a very deep and untrainable MLP could be made have a reasonable global C map simply by adding skip connection from its input to its output, but this won't do anything to address the untrainability of the layers being "skipped around" (which form a subnetwork).
In the main text we ignored this complication in the interest of shorter presentation, and because we happened to have (C'( = (1) for the simple network architectures focused on in this work. To remedy this, in the current section we will discuss how to modify the conditions Cy( and (1) used in TAT that they take into account all subnetworks. This will be done using natural generalization of the maximal slope function from DKS. We will then address the computational challenges that result from doing this.
A subnetwork of f is defined as (non-strict) connected subset of the layers in f that constitute neural network with singular input and output layer. So for example, layers 3, 4 and 5 of 10 layer MI.P form subnetwork, while layers 3, 4, and do not.
15

---

Published as a conference paper at ICLR 2022
To begin, we will replace the condition ()) 1y (used in TAT for Leaky ReLUs) by the condition ( (0) 7), where we define the maximal с value function p4g of by
(a)(a)
max Cg)
9:
where a is the negative slope parameter (which determines C in LReLU networks [via 6xJ and thus each Cg).
We will similarly replace the condition (1) - (used in TAT for smooth activations) by the condition (1)) T, where define the maximal curvature function nz of f by
(C)
max (1)
g:
where each (1)) (1) is determined by (1). That each (1) is well-defined function of (1) follows from the fact that maps always map the value to 1, the aforementioned relationship between C. and C, and the fact that we have (1(1) = 1 under TAT (so that C₁ (1) 1 for all subnetworks h). These facts allow us to write (1) as a constant multiple of (1(1) using the linearity of 2nd derivatives and the 2nd-order chain rule (which is given by (a (x) = a" (=((()))((x)) + a' (b(x))b((z))
B.2 COMPUTING 14) AND n3 IN GENERAL
Given these new conditions for TAT, it remains to compute their left hand sides so that we may ulti- mately solve for the required quantities (a C" (1)). In Section 2.3 we discussed how a (sub)network f's C map C can be computed in terms of the local C map C by series of composition and non- negative weighted sum operations. We can define a generalized version of this construction Uy.r which replaces C with an arbitrary non-decreasing function so that Uf,((c) = Cy(c). A recipe for computing U is given in Appendix B.4.
Given f.r, we define the subnetwork maximizing function M by Mf,r(x) = max Ug.((z)) -:
With this definition, it is not hard to see that (z)) C(x), rı (z) C'()xx, and (2()) (1)++z then (a) = Mr (0) (where the dependence on a is implicit through the dependence of C on Φ.). n2 (1) M (1), and n²) (C" computing (and inverting) My for general networks f and non-decreasing functions r.
Mf,,, (0). Thus, it suffices to derive scheme for
Naively, computing My, could involve a very large maximization and be quite computationally expensive. But analogously to the maximal slope function computation described in Martens et al. (2021), the computation of M, can simplified substantially, so that we rarely have to maximize over more than a few possible subnetworks. In particular, since Ug. (x) is a non-decreasing function of v for all y (which follows from the fact that non-decreasing), and Ugoh.r = Ug.r Un. thus follows that U (z) maximization, we can ignore any subnetwork in ff which composes with another subnetwork (not necessarily in f) to form a strictly larger subnetwork isomorphic to one in f. This will typically be the vast majority of them. Note that this does not therefore imply that M = Ujr since not all subnetworks compose in this way. For example, sufficiently deep residual branch of residual block in rescaled ResNet won't compose with any subnetwork to form a larger one.
(x), (2) for all x. This means that for the purposes of the
B.3 SOLVING FOR AND (I)
Having shown how to efficiently compute Mrr, and thus both of ne and n²z. it remains to show how we can invert them to find solutions for a and (1) (respectively). Fortunately, this turns out to be easy, as both functions are strictly monotonic in their arguments (a and C"(()), provided that if contains at least one nonlincar layer. Thus, we may apply a simple I-dimensional root-fim mont-findig approach, such as binary search.
To see that (a) is a strictly decreasing function of a (or in other words, a strictly increasing function of aa), we observe that it is a maximum over terms of the form (0), which are all either strictly decreasing non-negative functions of a, or are identically zero. These properties of U Ugc((0)
16

---

Published as a conference paper at ICLR 2022
follow from the fact that it involves only applications of C, along with compositions and non -egativee weighted averages, and that C(c) is a strictly decreasing function of a for all c E [-1,1] (in Leaky ReLU networks). A similar argument be used to show that μ (C is strictly increasing function of (1) (and in fact equal a non-negative multiple of (1)))).
B.4 RECIPE FOR COMPUTING U,r
As defined, Up .. is computed from f by taking the computational graph for Cj and replacing the local map C with r wherever the former appears. So in partieular, one can obtain computational graph for frr (x) from f's computational graph by recursively applying the following rules:
1. Composition g o h of two subnetworks and h maps to Ugr Un.--
2. Affine layers map to the identity function.
3. Nonlinear layers map to r.
4. Normalized sums with weights W1, W2, ..., Wn over the outputs of subnetworks 91 .2. map to the function
... In.
W w/U (2 + w (
where 11, .2, Fn are the respective inputs to the Ug.r's.
5. I's input layer maps to æ.
In the special case of computing (0), one gets the following simplified list of rules:
1. Composition g h of two subnetworks g and h maps Ug.rz (0) + Unrr( (0)
2. Affine layers map to O.
3. Nonlinear layers map to C" (1).
4. Normalized sums with weights w1, W2 Wn. over the outputs of subnetworks 91 92: ... In. map to the function
wf u (0) + UZU 92.12 12 (0) + + (0).
5. s's input layer maps to
Note that this second procedure will always produce a non-negative multiple of (1),, provided that contains at least one nonlinear layer.
B.5 RESCALED RESNET EXAMPLE
In this subsection we will demonstrate how to apply the above rules to compute the maximal curvature function u² for a rescaled ResNet f with shortcut weight w and residual branch R (as defined in equation 6). We note that this computation also handles the case of a vanilla network by simply taking = O.
First, we observe that all subnetworks in compose to form larger ones in f. except for itself, and for the residual branches of its residual blocks. We thus have that p2 (1)) max (U,,,, (0), ())).
Because each residual branch has simple feedforward structure with three nonlinear layers, follows that R, (0) = (1)( And because each shortcut branch S has no nonlinear layers, it follows that Us.r, (0) = 0. Applying the rule for weighted averages to the output of each block B we thus have that UBr,,00 (0) w² Us, (0) + (1 w² -UR (0) 3(1 w²) Given a network with L nonlinear layers, we have 1./3 blocks, and since the blocks compose in feedforward manner
thus follows that U, (0) = (L/3) 3(1 - (1) = L(1 ( We therefore conclude
that (²²(c"()) (1)) = (3, L w²) ")) (I)
The rescaled ResNets used in our experiments have slightly more complex structure (based on the ResNet-50 and ResNet- 101 architectures). with a nonlinear layer appearing after the sequence of residual blocks, and with four of their blocks being "transition blocks", whose shorteut branches
17

---

Published as a conference paper at ICLR 2022
contain nonlinear layer. In these networks, the total number of residual blocks is given by (L 1223333 Following a similar argument to the one above we have that
(0) L - 3 - (1) 44 3 (1)) (1)
=[(L [1 w²) Au 11) = [[L (0l1-66(111 6) 5))
and thus
(²)(² = -) 51) =
[L 6) - w²) (1)
B.6 PSEUDOCODE
Algorithm TAT for LReLU.
Require:
The target value for (a)
1: Use the steps from Subsection B.2 to construct a procedure for computing the maximal e value function μ (a) for general 2 0. Note that the local map C, which (a)(a) depends, can be computed efficiently for (tranfformed LReLUS using equation 9.
2: Perform a binary search to find the negative slope a such that 1 = 17.
3: Using the found output the transformed activation function given by (x)) = / de (z).
Algorithm 2 TAT for smooth activations.
Require: The target value t of (1)
Require: The original activation function ((x)
1: Use the steps from Subsection B.2 construct procedure for computing the maximal curvature function 1 (C" (1)) for general (1(1)= 2 O.
2: Perform a binary search to find ("(1) such that (1))
3: Using a numerical solver, solve the three-dimensional nonlinear system in equation 16 (but with the value of (1) found above instead of L) to obtain values for o, 3, % and 8.
4: Using the solution from the last step, output the transformed activation function given by
((x) (éb(r 3) 8).
C TECHNICAL RESULTS AND PROOFS
Lemma 1. For networks using the activation function (a(x) v 1 (z(x, the local and C
maps are given by
(1 ax) (VI- c² ccos (c))
2(g) =q and
(20)
a(1 a2)
Proof. In this proof we will use the notation 2s and Cs to denote the local Q and maps for networks that use given activation function ø
First, we note that LReLU basically the weighted sum of identity and ReLU. In particular, we have the following equation:
a)))(x)) = max(2, 0} + min(2, 0}. (x) ar + (1-o(1=))
Second, we have that 2s. (g) E2-,((,11) |q= 2 0] + a a²E [q>01[> 0]] 1-a2
(from which 2g (g) =q immediately follows).
18

---

Published as a conference paper at ICLR 2022
It then follows from equation 5, and the fact that local C maps invariant to multiplication of the activation function by a constant, that
Co (c) Coa (c)
2
E Ba 20.
0a (z1) da (ez V 222)
+ a²
2
[a (1- (1- (c) (1)
(21)
1 + a²
2
2a(1 22
[c22] + v1 22220
1 + a²
From Daniely et al. (2016) we have that
Co. (c) vI c² + (π cos co)-)(cc) (c))e
(22)
and for the last part of equation 21 have
(2..1) (ez) + v e 22) (21) E ((0.)) [e2 =120
(23)
Plugging equation 22 and equation 23 back into equation 21, we get
C₁. (c) I+a* a a²c+ (1 a1 -I- + (a - cos (c))c
2
cas-i(c))e
+ a(1 a)c
2
π
(24)
(1-a) (V1 - c² + c(T cos cos (c))) + 2Tac
a1+
Rearranging this gives the claimed formula.
The global C map of a feedforward network with 0a(x) as its activation activationfuunction equal to that of a rescaled ResNet of the same depth (see Section 2.4) with normalized ReLU activation <(x) = 22 max(x, (r. 0), shortcut weight / Ifa and residual branch R. consisting of combined layer (or just a normalized ReLU activation) followed by affine layer.
Proposition 1.
Proof. By equation 7, the C map for residual block B of the hypothesized rescaled ResNet is given by
Cs(c) w + (1 - C²)
(25)
The global map of this network is given by L compositions of this function, while the global C map of the hypothesized feedforward network is given by I. compositions of Cz,, (c). So to prove the claim it suffices to show that Cc(c) Cc. (c).
Taking w V
2a T.a2. one obtains the following
2a (1 a1- VI c² c(TT cos (c))
CB(C) = I+ +
(26)
1+
π
which is exactly the same Co, (c) as given in Lemma 1. This concludes the proof.
Proposition 3. Suppose f is vanilla network consisting of combined layers with the TReLU
activation function (so that Cy(0) = E (0, 1)). Then Cj converges to a limiting map on (-1, 1) as goes infinity. In particular.
lim =y(c)=v(e,TT) (c,,T),
(27)
where T. is such that (0, T) v, and where vp is the solution of the following ordinary differential equation (ODE) with the first argument being the initial condition (i.e. vb(c, 0) c), and the second argument being time:
dx(()) = /1- x0))2 - xxt) cos (x())).
(28)
dt
19

---

Published as a conference paper at ICLR 2022
Proof. First, we notice that the local C map for TReLU networks can be written difference equation:
a
C(c) = T1+ (V1-c² ccos c (())
(29)
Importantly, C is monotonically increasing function of c, whose derivative goes to zero only as E [0, 1] goes to 1. Thus, to achieve (0) = rj in the limit of large L, we require that (1-a)² to O. This implies that the above difference equation converges to the ODE in equation 28.
m[1+a2] goes
Because the function v1 - x2 ax cos - (x: is continuously differentiable in [-1, 1], and its derivative cos (x)) bounded, one can immediately show that it is globally Lipschitz, and the ODE has unique solution @C t) according to Theorem 3.2 of Khalil (2008).
Now, we are only left to find the time T such that C7 (0, T) = 17. To that end, we notice that
g(x) = v - x² - x cos (σ) > 0, for x E (-1,1)
(30)
because g(1) 0 and gxx) = cos (x) < 0 on (-1, 1). This implies that the (500 t) is monotonically increasing continuous function of t. Since (0, 0) = 0, to establish the existence of T suffices to show that (0, oo) 2
To this end we first observe that
g(x) 2 2/2 (1 (1-x) 2))
(31)
which follows by defining h(x) = g(z) - 2/2
3 (1 z/)3/2 and observing that h(1) = 0 and h (z())
= -oss (cos-((x)= + / 1/2 < on (-1, 1). Given this, it is sufficient to show that the solution Ø for the ODE 2v2 (1-110 (1-x) x/3/2 satisfies (0, oo) = 1. The solution & turns out to have a closed-form of (0, () = 1 - 3
(2(+3) ) and thus (0, x)) 2 (0, xc) 1. This completes the proof.
Theorem 1. For a network with d. (z) as its activation function (with 2 0), we have
max (Cy(c) cl S min (4C, (0), 1 + C)))),, max (c)(c) - s min (4C,(0),1) (10) cEI-1,1]
cE1-1,11
Proof. Because Cy is a positive definite function (by Section A.4) we have that it can be written as Cy(c) = - cn for bn 0. Given Cy(1) C1)
have
8
20
2
Σ
bn Σ
nbn =1
=
bp Σ
1)bn
2ho
b 2 1.
(32)
n
Hence, 1 (0) b2 S 2hg (0). Now we are ready bound the deviation of Cc from identity:
max rE[-1.1]
-c( -cl
max
by
Onc b1)
cE dEi-1i 1,1]
80
max cE cE|-1,1]
by
l/l + (1 - (1-billd
(33)
n=
= bo +
bn + 1 - b1 =2 - b1)
1=2
2(1 (0) S 44,((0)
Using equation 20 we have that
=1 a)=
c1
a2) cos 1c)
From our assumption that a 2 0 it follows that 0 < C'(c) < 1 for all c E [-1, 1]. Since the property
of having derivative bounded between 0 and 1 is closed under functional composition and positive
20

---

Published as a conference paper at ICLR 2022
weighted averages, it thus follows that 0 I Cf(c) for c [-1, 1]. An immediate consequence of this is that (c(c) is non-decreasing, and that
max (c((c) - - +1. cE|-1.1]
(34)
Next, we bound the deviation of (c) from 1:
max cE[-1.1] cE[-1.11
(c)
max
nbn - - (1-bby b1)
max cE|-1,1]
nb,/c" m- + (1-b))
(35)
Σ nbn + b = =2(1 - bi)
n=2
= 2(1 -2((0) S (C)(
From the previous fact that 0 & ()) < for all e E [-1.1] we also have that
max cj(c)) - < 1. This completes the proof.
Theorem 2. Suppose f is network with smooth activation function. (f 1. then we have ()(c) - 1| (1)
max Cy(c) c s (1), cEI-1.11 cE[-1,1]
(13)
Proof. C is a positive definite function by Section A.4. So by the fact that positive definite functions are non-negative, non-decreasing, and convex on the non-negative part of their domain, we obtain that c (0) 2
10----4 C1) By equation 33, we have
max (Cy(c) cl s 2(1 ()) - (1)
(36)
cE(-1,1]
Further by equation 35, we also have
(c(()) - 2(1 -2 < (1) cE[-1,1]
max
(37)
This completes the proof.
Proposition 4. Suppose f is some function computed by neural network with the ReLU activation. compute using LReLU neural network
Then for any negative slope parameter + +1, of the same structure and double the width of the original network.
Proof. The basic intuition behind this proof that a ReLU unit can always be "simulated" by two LReLU units long + =1, due the following formula:
(x)
a (0a(x) + ada(=
We will begin by proving the claim in the case of a network with one hidden layer. In particular, we assume the ReLU network has m hidden units:
m
fl,
Σ a, b,),
(38)
where x E Rd is the input, and w E mid b E 1 and E R" are weights, biases of the input layer and weights of output layer, respectively. For LReLU with negative slope a, one can construct the following network
2m Σ r=l
f(w,, fod,,''c d,,d,z)=
Ch x b,).
(39)
21

---

Published as a conference paper at ICLR 2022
If we choose w. wp WW by b, - -b/ -bpmm of
I-az Or and d/4m I-a2 Or, we have
(us, T I + b,,) + a Pa a (w'+m T I + by km)
a²
1 - a² ar I + b,) - 1 a² a (w, + b,) = a, x by). (40)
This immediately suggests that f(w', d''''''''' ,x) = f(w,f((ww.k b,
Since deeper networks, and one with more complex topologies, can be constructed by composing and summing shallower ones, the general claim follows.
D EXPERIMENT DETAILS
For input preprocessing ImageNet we perform a random crop of size 224 224 to each image, and apply a random horizontal flip. In all experiments, we applied L2 regularization only to the weights (and not the biases or batch normalization parameters). We selected the L2 constant by grid search from {0.00005, 0.00002, 0.0}. For networks without batch normalization layers we applied dropout to the penultimate layer, with the dropout rate chosen by grid search from {0.2, 0.03. In addition, we used label smoothing (Szegedy et al., 2016) with a value of 0.1.
For each optimizer we used a standard learning rate warm-up scheme which linearly increases the learning rate from O to the "initial learning rate" in the first 5 epochs, and then decays the learning rate by a factor of 10 at 4/9 and 7/9 of the total epoch budget', unless specified other- wise. The initial learning rate was chosen by grid search from (1.0, 0.1, 0.03, 0.01) for SGD, (0.003, 0.001, 0.0003, 0.0001, 1 0. 0.000033 for K-FAC, and (10.0, 3.0, 1.0, 0.3, 0.1) for LARS. For all optimizers we set the momentum constant 0.9. For K-FAC, we used fixed damping value of 0.001, and a norm constraint value of 0.001 (see Ba et al. (2017) for a description of this parameter). We also updated the Fisher matrix approximation every iteration, and computed the Fisher inverse every 50 iterations, unless stated otherwise. For LARS, we set the "trust" coefficient to 0.001. For networks with batch normalization layers, we set the decay value for the statistics 0.9.
For initialization of the weights we used the scale-corrected uniform orthogonal (SUO) distribu- tion (Martens al., 2021) for all methods/models, unless stated otherwise. For X k matrix (with & being the input dimension), samples from this distribution can be generated by computing (XXT) 1/2 X, where X is an m x & matrix with entries sampled independently from (0, 1). When m > k. we may apply the same procedure but with k and m reversed, and then transpose the result. The resulting matrix is further multiplied by the scaling factor max /m/k, 1}. which will have an effect only when k S m. For convolutional networks, we initialize only the weights in the center of each filter non-zero values, which is technique known as Delta initialization (Balduzzi et al., 2017; Xiao et al., 2018), or Orthogonal Delta initialization when used with orthogonal weights (as we do in this work).
We implemented all methods/models with JAX (Bradbury et al., 2018) and Haiku (Hennigan et al., 2020). We used the implementation of SGD and LARS from Optax (Hessel et al., 2020). We used the JAX implementation of K-FAC available at https: / wtt/ demm//eeemmddd hwd.d//faaaa /kac_jaaaaaaaa//aaa/aa//aaa//aaa
E ADDITIONAL EXPERIMENTAL RESULTS
E. EMPIRICAL C VALUES FOR FINITE-WIDTH NETWORKS
The computation of cosine similarities performed by C maps is only an approximation for finite width networks, and it is natural to ask how large the approximation error is. To answer this question, we compare the theoretical predictions with the empirical simulations on fully-connect networks of different depths and widths. In particular, we use fixed = 0.9 for TReLU and we compute the 1-th "empirical value" 2 s for each layer index I, where æ² and 2 are random vectors
'ee later found that cosine learning rate annealing (Loshchilov & Hutter, 2016) is slightly better for most settings, but this did change our conclusions.
22

---

Published as a conference paper at ICLR 2022
valuu computed Hayer
value computed i-llaye
0.8
0.8
0.6
0.6
width 30
width - 30
witth 100
width 100
0.0
40
60
100
0.0
20
80
100
depth
depth
(a) TReLU, Gaussian
(b) DKS SoftPlus, Gaussian
c vallue computed on ielyyr
computed H-pyer
0.0
0.8
0.6
0.6
0.4
0.4
width
width 100
withh 100
width 300
width 300
0.0
100
depth
depth
(c) TReLU, Orthogonal
(d) DKS Softplus, Orthogonal
Figure Empirical values for TAT and DKS, which are averaged over 100 pairs of inputs and different randomly-inian networks. We include the results for both Gaussian fan-in and Orthogonal initialization. Vertical lines indicate the standard deviation. TReLU has smaller kernel approximation error and is robust to Gaussian initialization. For TReLU, also plot the evolution of the values (black dashed line) predicted compute analytically for TReLU).
by the C map (which
| do and 5 =2 = 0 (so that e 0). As shown in Figure 4a and 4c, the approximation error is relatively small even for networks with width 30.
chosen that |
We also included the results for networks using DKS (with Ç 10) and the SoftPlus activation function. Figure 4b and 4d reports empirical с values as a function of layer index I, with 2f and
xa chosen that e 0.8. With Gaussian initialization, the standard deviations than TReLU, U, and the average values for widths 30 and 100 deviate significantly from the theoretical predictions. (The DKS conditions implies C(c) S for any E [0. 1], which suggests the value should decrease monotonically.) By comparison, the error seems to be much smaller for orthogonal initialization, which is consistent with the better performance orthogonal initialization reported by Martens et al. (2021). (By contrast, we show in Appendix E.7 that Gaussian initialization performs on par with orthogonal initialization for TReLU.) In addition, we note that the standard deviations increase along with the depth for both Gaussian and orthogonal initializations.
much larger
E.2 RESULTS ON CIFAR-10
0.95
In addition to our main results on the ImageNet dataset, also compared TAT to EOC on CIFAR- 10 (Krizhevsky et al., 2009) using vanilla networks derived from a Wide ResNet reference architecture (Zagoruyko & Komodakis, 2016). In particular, we start with Wide ResNet with a widening factor 2. and remove all the batch normalization layers and shortcut connections. We trained these networks with the K-FAC optimizer for 200 epochs using standard piecewise constant learning rate schedule. To be spe- eific, decay the learning rate factor of 10 75 and 150 epochs. For K- FAC, we set the damping value 0.01 and norm constraint value 0.0001.
0.85
0.70
layers
0.50
100 layers with EOC
304 layers with TAT
layers
0.50
2500
5000
12500 15000 20000
Iterations
Figure 5: CIFAR-10 validation accuracy of ResNets with ReLU activation function initialized using either EOC or TAT (ours).
For data preprocessing include basic data
23

---

Published as a conference paper at ICLR 2022
augmentations such random crop and horizontal flip during training. As shown in Figure 5, TAT outperforms EOC significantly. As we increase the depth from 100 to 304, the accuracy EOC network drops dramatically while the accuracy of the TAT network remains roughly unchanged.
E.3
REDUCING THE OVERHEAD OF K-FAC
0.750
0.750
0.725
0.725
0.700
0.700
aC
0.675
0.675
0.630
0.650
0.625
0.625
0.600
Desth-50, KFFAC pariod-10. perlod-200
D6sch=50, K-FAC w/ inv perlod=50
0.600
Depthe550 K-FAC w/ hwr_perodd-10 inc_perind=200
Deathu50, IC K-FAC wi mvv_peridd-50
0.575
50p SGD
0.575
Degth-5D.
Desch=50, SGD w/ 180 apochs
Depthw50. SGD w/ 180 epochs
0.530
50000
100000
150000
200030
0.550
10000
20000
30000
40000
50000
lterations
Wall-clock time (seconds)
Figure Top-1 validation accuracy on ImageNet as function of number ilerations (left) wall-clock time (right) with K-FAC optimizer. One can reduce the computational overhead significantly by updating curvature matrix approximation and its inverse less frequently.
In our main experiments the per-step wall-clock time of K-FAC was roughly 2.5x that of SGD. However, this gap can be decreased significantly by reducing the frequency of the updates of K-FAC's approximate curvature matrix and its inverse. For example, if we update the curvature approximation every 10 steps, and the inverses every 200 steps, the average per-step wall-clock time of K-FAC reduces by half to a mere x.25x that of SGD. Importantly, as can be seen on Figure 6, this does not appear to significantly affect optimization performance.
E.4 DISENTANGLING TRAINING AND GENERALIZATION
In main experiments only reported validation accuracy on ImageNet, making it hard to tell whether the superior performance of TAT vs EOC is due to improved fitting/otimizzation speed, or improved gen- eralization. Here, compare training accuracies of EOC-initialized networks (with ReLU) and networks with TReLU, in exactly the same experimental setting as Figure 1. We train each network on ImageNet us- ing K-FAC for 90 epochs. For each setting, we plot the training accuracy for the hyperparameter combi- nation that gave the highest final validation accuracy. As shown in Figure 7, the EOC-initialized networks achieve competitive (if not any better) training accu- racy, suggesting that the of TReLU improves the generalization performance and not optimization performance.
0.90
0.85
Training acc
0.80
0.75
layers
50 layers with EOC layers
101 layers with EDC
0.65
200 layers
200 layers with EDC
0.60
0
20000
40000
60000
80000
100090
Iterations
Figure ImageNet training accuracy of deep vanilla networks with either SoC-initialized ReL.U networks or T'RelU networks.
E.5 CLOSING THE REMAINING GAP USING WIDER NETWORKS
In all of our main experiments we used networks derived from standard ResNets (by removing nor- malization layers and/or shortcut connections). By construction, these have the same layer widths as stan- dard ResNets. A natural question to ask is whether using wider networks would change our results. For example, it's possible that vanilla networks with TAT would benefit more than ResNets from increased width, since higher width would make the kernel approximations more accurate, and could also help compensate for the minor loss of expressive power due the removal of shortcut connections.
Table 9: The effect of increasing width on Im-
ageNet validation accuracy. use vanilla net- works for EOC and TAT (ours).
Depth
Width
EOC
TAT
ResNets
1x
72.0
76.0
76.7
50
2x
73.5
77.3
77.9
1x
62.4
76.5
77.9
101
2x
66.5
77.6
78.6
24

---

Published as a conference paper at ICLR 2022
With layers double the width of standard ResNets, it becomes too expensive store and invert Kronecker factors used in K-FAC. Therefore, we only train these wider networks with SGD. In order to mitigate the slower convergence of SGD for vanilla networks (see Section 5.5), we train them for 360 epochs at a batch size of 512. Note that due to increased overfitting we observed in ResNets after 360 epochs (resulting in lower validation accuracy) we only trained them for 90 epochs. As shown in Table 9, doubling the width does indeed narrow the remaining validation accuracy gap between ResNets and vanilla TAT networks. In particular, the gap goes from 0.7% 0.6% for depth 50 networks, and from 1.4% to 1 % for depth 101 networks.
E.6 COMPARISON WITH PRELU ON RESCALED RESNETS
In Table of the main text we compare PReLU and TReLU deep vanilla networks. Here we extend this comparison to rescaled ResNets with shortcut weight of 0.8. For PReLU, we again include two different initializations: one with negative slope (effectively ReLU), and another with 0.25 negative slope (which was used in He et al. (2015)). We report the full results in Table 10. For all settings, TAT outperforms PReLU by large margin, suggesting that better-initialized negative slope crucial for both rescaled ResNets and deep vanilla networks.
Table 10:
Comparison with PReLU with rescaled
ResNets (w = 0.8).
Depth
Optimizer
TReLU
PReLUU..
PReI.Uo.25
K-FAC
76.4
75.7
73.6
50
SGD
76.0
71.5
71.5
K-FAC
77.8
76.4
76.8
101
SGD
77.3
73.1
73.4
E.7 COMPARISON OF DIFFERENT INITIALIZATIONS
In all of our experiments we use the Orthogonal Delta initialization intro- duced by Balduzzi et al. (2017) and Xiao et al. (2018). This because it's technically required in order to apply the extended Q/C map analy- sis of Martens et al. (2021) (which underlies DKS and TAT) to convo- lutional networks, and because it is generally thought to be beneficial. In this subsection examine this choice more closely by comparing it traditional Gaussian fan-in ini- tialization (with o² = 2 for ReLUs). We consider standard ResNets and deep vanilla networks using either EOC (with ReLUs) or TAT with (with LReLU), Surprisingly, it turns that the Orthogonal Delta initialization does not have any clear advantage over the Gaussian fan-in approach, at least in terms of validation accuracy after 90 epochs.
Table 11: Comparison of Orthogonal Delta and Gaussian fan-in initialization.
Depth
Optimizer
Init
ResNet
EOC
TAT
Orth Delta
76.4
72.6
74.6
K-FAC
Gaussian
76.5
72.5
74.8
50
Orth Delta
76.3
63.7
71.0
SGD
Gaussian
76.6
63.1
68.7
Orth Delta
77.8
71.8
74.2
K-FAC
Gaussian
77.8
72.3
74.1
101
Orth Delta
77.9
41.6
70.0
SGD
Gaussian 78.0
41.1
68.7
25


Reviews:
[{"review_id": "qA0SQ9lvJJM", "reviewer": "Reviewer_8mde", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "Strong points:\n\n1. A study of an interesting problem with rigorous theoretical analysis.\n\n2. Both theoretical extension and empirical improvement.\n\n3. Solves the ReLU incompatibility issue of existing DKS work with new Q/C map conditions and LReLu, which also looks nicer than the the DKS's model class-preserving transformation.\n\nWeak points:\n\nI only have one major concern. The claims that DHS is not compatible are a bit too strong, at least the experiments (LReLu+DKS) did not show that way. At the end of Section 3, it says \"However, DKS is not fully compatible with ReLUs,\", which needs an exact explanation with the DKS condition.  DHS is compatible with LReLUs, and works pretty well, as in Table 7. So overall,  I feel some of the claims frequently mentioned throughout the paper are a little bit exaggerating. DHK, in their paper, claims quite differently, \"small decrease in generalization performance\" for removing BN, and \"similar results\" for removing skip connections. \n \nThe other claim is ** Using TAT, we demonstrate for the first time that a 50-layer vanilla deep network can nearly match the validation\naccuracy of its ResNet counterpart when trained on Imagenet.**. What \"for the first time, ..., nearly...\" means? I can tell, from Table 7, that  DKS is also very close, even with SoftPlus if you insist LReLU+DKS is your contribution. \n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "I would suggest the authors be careful about their claims. It appears to me that some extension from DKS motivated by the similarity between the model class-preserving transformation and Leaky ReLU is nice, but is not world-changing. \n\nI appreciate the detailed derivations. \n\nMy overall rating is 6.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gHhZoy3N82C", "reviewer": "Reviewer_HqiL", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "## Pros\n\n1. This paper is well-organized and easy to follow.  The main contribution of this work is to adapt DKS to models with ReLU-family activations. The authors have explored the reason why DKS is not fully compatible with ReLU activations, then providing reasonable solutions by changing some Q/C Maps conditions. The proposed methods can be easily applied to activation layers and only involve a few extra computations. \n\n2. Effectiveness of  the proposed TReLU has been demonstrated by a series of ablation studies. With K-FAC optimizer, the TReLU ResNet can maintain performance without any residual connections and normalization in some cases. \n\n## Cons\n\n1. The problem of removing residual connections on NN is interesting, and this paper has provided a good solution. However, existing methods seem not to be so bad as claimed in the paper. As discussed in the paper, the main drawback of DKS is incompatibility with ReLU-family activations. But as shown in Table.7, DKS+LReLU+K-FAC+ResNet50 seems to slightly outperform TAT counterparts, which means the mentioned incompatibility of DKS can be alleviated by using other methods like K-FAC optimizer and thus making TAT not so significant. Anyway, I thought DKS is a very recent work compared to this paper, and the proposed TAT can also yield comparable results.\nTherefore,  I still lean to accept this work.", "summary_of_the_paper": "This paper mainly discusses the training of neural networks without residual connections. To close the gap between residual-free and regular models, an activation transformation technique named \"Tailored Activation Transformation (TAT)\" is introduced.  Compared to the state-of-art method DKS,  the proposed TAT can yield better results for models using ReLU-family activation. Overall the motivation is well-discussed and sufficient ablation studies are provided. ", "summary_of_the_review": "The paper provides sufficient discussion and experiments for improving DKS. The proposed approaches yield good results on models without residuals and normalizations.  However, when compared to existing DKS methods, the improvements seem to be somewhat marginal, thus restricting the significance of this paper.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OYupl6TudEm", "reviewer": "Reviewer_vixN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "Strength: (1) The paper is well-organized --- although the technical backgrounds are complicated, the authors try to introduce them in reasonably short sections. (2) The paper succeeds in increasing the performance of deep vanilla networks with ReLU-family activation, matching the performance of deep residual networks. \n\nWeakness: (1) The current draft is unfriendly to practitioners. Maybe the authors would like to include an algorithm box at the beginning to explain the proposed method (without understanding how to derive it). (2) Although the proposed method supports ReLU-family activation (TReLU), the performance gain is very marginal compared to smooth activations with DKS (at most 0.3%, Table 7). Thus, it raises the concern whether ReLU-family activation is needed. Is there any tradeoff using ReLU-family activation verse using smooth activation? ", "summary_of_the_paper": "The paper extends the Deep Kernel Shaping (DKS) method to Tailored Activation Transformation (TAT). As a result, deep vanilla neural networks with ReLU-family activations match the performance of deep residual networks.", "summary_of_the_review": "Although the authors have well-organized the paper, the current draft is hard to follow due to its technical complicity. Moreover, given the marginal improvement compared to DKS, a more extensive discussion on smooth and ReLU activations is desirable.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3SOwqQqO1qR", "reviewer": "Reviewer_TjGY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "This paper makes solid theoretical contributions on analyzing the initial conditioning of feedforward networks without shortcuts. Most of the analysis are derived directly from the concurrent work (DKS). It provides insights for proper initialization and configuring the activations for feedforward networks to achieve rapid training speed, improves our understanding of neural networks and has the potential to impact future design choices of them. Although the main point of the paper is to remove shortcuts, the analysis can also be applied to rescale the shortcuts of ResNets and achieve even better results than standard ResNets (with BN) when BN is not used (Table 4). \n\nI feel the only weaknesses of the paper are the clarity and the significance of the results. Limited by the pages, many important details are either abbreviated or require reading the 172-page DKS paper. The results of shortcut-free networks are still slightly worse than ResNets, and the merits of removing shortcuts do not seem to offset the drawbacks. A minor concern is the current analysis cannot be applied to Batch Normalizations and attention mechanisms, but I believe the theoretical contribution is enough.\n\nIn addition, I have following questions or suggestions:\n\n1. For each experimental result, how many different random seeds are used, and what are the standard errors of these results?\n\n2. Is training using K-FAC really faster than SGD? Most tables show better results achieved with K-FAC than SGD under the same number of epochs, and it has been shown in Figure 5 that K-FAC can be made 2x as fast without losing accuracy, but the wall-clock time comparison with SGD is not shown. How is the wall-clock time like for the 90-epoch K-FAC and 180-epoch SGD? \n\n3. The paper did not clearly specify how the weights are initialized. It is mentioned in Section D.6 in the appendix, but it was not clearly specified in Section 2.1. It would be better to include some form of pseudo code to specify the weight initialization and setting the parameters of activations. It is also not clearly described in Section 4.2 that the smooth activations are transformed.\n\n4. The method itself introduces extra hyperparameters that require additional hyperparameter search, such as $\\eta$. This increases the cost of hyperparamter search. Could there be some ablation studies on the impact of $\\eta$?\n\n5. How is the result like for even deeper networks, e.g., ResNet-152?\n", "summary_of_the_paper": "Based on the analysis in the concurrent work (DKS), this work develops a new set of conditions for the Q/C maps to set the non-trainable parameters for the activations (Leaky ReLU and transformed smooth activations), which enables training deep feedforward networks at comparable speed as ResNets on ImageNet, and achieves significantly better test accuracy than those obtained with EOC. ", "summary_of_the_review": "This paper makes solid theoretical contributions on analyzing the initial conditioning of feedforward networks without shortcuts. The analysis can even be used to rescale the shortcuts of BN-free ResNets and achieve even better results than standard ResNets (with BN). However, the method also introduces an extra hyperparameter $\\eta$, and results of shortcut-free networks are still slightly worse than ResNets. The clarity of the paper also needs improvements.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper seeks to find an answer to some quite interesting research question: can deep vanilla networks without skip connections or normalization layers be trained as fast and accurately as ResNets? In this regard, the authors extend Deep Kernel Shaping and show that a vanilla network with leaky RELU-family activations can match the performance of a deep residual network.

Four reviewers unanimously suggested acceptance of the paper. There were concerns about the clarity or marginal performance improvement. However, they all including myself agree: achieving the competitive performance with the vanilla deep model itself can be seen as a big contribution and the clarity has been improved to some extent through revision.
Decision: Accept (Poster)