D EEP L EARNING WITHOUT S HORTCUTS : S HAPING THE K ERNEL WITH T AILORED R ECTIFIERS
A BSTRACT
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 overﬁt signiﬁcantly 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 signiﬁcantly 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
I NTRODUCTION
[Figure]
trained with K-FAC.
Thanks to many architectural and algorithmic innovations, the recent decade has witnessed the unprecedented success of deep learning in various high-proﬁle challenges, e.g., the ImageNet recognition task (Krizhevsky et al., 2012), the challenging board game of Go (Silver et al., 2017) and human-like text generation (Brown et al., 2020).
Among them, shortcut connections (He et al., 2016a; Srivastava et al., 2015) and normalization layers (Ioffe & Szegedy, 2015; Ba et al., 2016) are two architectural components of modern networks that are critically important for achieving fast training at very high depths, and feature prominently in the ubiquitous
ResNet architecture of He et al.
(2016b).
Despite the success of ResNets, there is signiﬁcant 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 inﬁnity (Hayou et al., 2021).
In this sense, the question of whether truly deep networks can be efﬁcient 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-inefﬁcient at inference time, as the input to every residual block has to be kept in memory until the ﬁnal addition.
In particular, the shortcut branches in ResNet-50 account for about 40% of the memory usage by feature maps.
Also, the classical interpretation of why deep networks perform well – because of the hierarchical feature representations they produce – does not strictly apply to ResNets, due to their aforementioned tendency to behave like ensembles of shallower networks.
Beyond the drawbacks of ResNets, training vanilla deep neural networks (which we deﬁne as networks without shortcut connections or normalization layers) is an interesting research problem in its own right, and ﬁnding a solution could open the path to discovering new model architectures.
However, recent progress in this direction has not fully succeeded in matching the generalization performance of ResNets.
Schoenholz et al.
(2017) used a mean-ﬁeld analysis of deep MLPs to choose variances for the initial weights and bias parameters, and showed that the resulting method – called Edge of Chaos (EOC) – allowed vanilla networks to be trained at very high depths on small datasets.
Building on EOC, and incorporating dynamical isometry theory, Xiao et al.
(2018) was able to train vanilla networks with Tanh units 1 at depths of up to 10,000.
While impressive, these EOC-initialized networks trained signiﬁcantly slower than standard ResNets of the same depth, and also exhibited signiﬁcantly worse generalization performance.
Qi et al.
(2020) proposed to enforce the convolution kernels to be near isometric, but the gaps with ResNets are still signiﬁcant 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 a method called Deep Kernel
Shaping (DKS) for initializing and transforming networks based on an analysis of their initializationtime 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 ReLUs, and in their experiments (which focused on training speed) their networks exhibited signiﬁcantly more overﬁtting than ResNets.
Inspired by both DKS and the line of work using mean-ﬁeld 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 ﬁrst time that a 50-layer vanilla deep network can nearly match the validation accuracy of its ResNet counterpart when trained on ImageNet.
And unlike with the EOC method, validation accuracy we achieve does not decrease with depth (see Figure 1).
Furthermore, TAT can also be applied to ResNets without normalization layers, allowing them to match or even exceed the validation accuracy of standard ResNets of the same width/depth.
A multi-framework open source implementation of DKS and TAT is available at https://github.com/deepmind/dks .
B ACKGROUND
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. K ERNEL F UNCTION A PPROXIMATION FOR W IDE N ETWORKS
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 deﬁned by a sequence of L combined layers (each of which is an afﬁne transformation followed by the elementwise activation function φ ) as follows:
[Equation]
with weights W l ∈ R d l +1 ×
d l initialized as W l iid ∼ N (0 , 1 /d l ) (or scale-corrected uniform orthogonal matrices (Martens et al., 2021)), and biases b l ∈ R
d l +1 initialized to zero.
Due to the randomness of the initial parameters θ , the network can be viewed as random feature model f lθ ( x ) (cid:44) x (cid:96) at each layer l (with x 0 = x ) at initialization time.
This induces a random kernel deﬁned as follows:
[Equation]
Given these assumptions, as the width of each layer goes to inﬁnity, κ lf ( x 1 , x 2 ) converges in probability (see Theorem 3) to a deterministic kernel ˜ κ l f ( x 1 , x 2 ) that can be computed layer by layer:
[Equation]
where ˜ κ 0 f ( x 1 , x 2 )
= κ 0 f ( x 1 , x 2 ) =
x (cid:62) 1 x 2 /d 0 .
2. L OCAL Q/C MAPS
By equation 3, any diagonal entry q l +1 i of Σ l +1 only depends on the corresponding diagonal entry q li of Σ l .
Hence, we obtain the following recursion for these diagonal entries, which we call q values :
[Equation]
where Q is the local Q map .
We note that q li is an approximation of κ lf ( x i , x i ) .
Analogously, one can write the recursion for the normalized off-diagonal entries, which we call c values , as:
[Equation]
where C is the local C map and c
= x (cid:62) 1 x 2 /d 0 .
We note that c l is an approximation of the cosine similarity between f lθ ( x 1 ) and f lθ ( x 2 ) .
Because C is a three dimensional function, it is difﬁcult to analyze, as the associated q values can vary wildly for distinct inputs.
However, by scaling the inputs to have norm √ d 0 , and rescaling φ so that Q (1) = 1 , it follows that q li = 1 for all l .
This allows us to treat C as a one dimensional function from [ − 1 , 1] to [ − 1 , 1] satisfying C (1) = 1 .
Additionally, it can be shown that C possesses special structure as a positive deﬁnite function (see Appendix A.4 for details).
Going forward, we will thus assume that q 0 i = 1 , and that φ is scaled so that Q (1) = 1 .
2. E XTENSIONS TO CONVOLUTIONAL NETWORKS AND MORE COMPLEX TOPOLOGIES
As argued in Martens et al.
(2021), Q/C maps can also be deﬁned 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 ﬁlter are initialized to zero.
Intuitively, this makes convolutional networks behave like a collection of fully-connected networks operating independently over feature map locations.
As such, the Q/C map computations for a feed-forward convolutional network are the same as above.
Martens et al.
(2021) also gives formulas to compute q and c values for weighted sum operations between the outputs of multiple layers (without nonlinearities), thus allowing more complex network topologies.
In particular, the sum operation’s output q value is given by q = (cid:80) ni
=
w
i q i , and its output c value is given by 1 q (cid:80) ni
=
w
i q i c i .
In order to maintain the property that all q values are 1 in the network, we will assume that sum operations are normalized in the sense that (cid:80) ni
=
w
i = 1 .
Following Martens et al.
(2021), we will extend the deﬁnition of Q/C maps to include global Q/C maps , which describe the behavior of entire networks.
Global maps, denoted by Q f and C f for a given network f , can be computed by applying the above rules for each layer in f .
For example, the global C map of a three-layer network f is simply C f ( c ) =
C ◦ C ◦ C ( c ) .
Like the local C map, global
C maps are positive deﬁnite 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 afﬁne layers, for which we can compute global Q/C maps.
And all of our formal results will implicitly assume this family of networks.
2. Q/C MAPS FOR RESCALED R ES N ETS
ResNets consist of a sequence of residual blocks, each of which computes the sum of a 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 modiﬁed 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 at the end of each block with a normalized sum.
These networks, which we will call rescaled ResNets , are deﬁned by the following recursion:
[Equation]
where R is the residual branch, and w is the shortcut weight (which must be in [ − 1 , 1] ).
Applying the previously discussed rules for computing Q/C maps, we get q l i = 1 for all l and
[Equation]
E XISTING S OLUTIONS AND T HEIR L IMITATIONS
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 vector, so that Q f describe the contraction or expansion of this magnitude through the action of f .
On the other hand, the c value approximates the cosine similarity of the function values for different inputs, so that C f 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 to 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.
However, as was ﬁrst observed by Schoenholz et al.
(2017), a well-behaved C map is also necessary for trainability.
When the global C map is close to a constant function (i.e. degenerate) on ( − 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.
[Figure]
The global C map of a TReLU network con- verges to a
Several previous works (Schoenholz et al., 2017; Yang &
Schoenholz, 2017; Hayou et al., 2019) attempt to achieve a well-behaved global C map by choosing the variance of the initial weights and biases in each layer such that
C (cid:48) (1) = 1 – a procedure which is referred to as Edge of
Chaos (EOC).
However, this approach only slows down the convergence (with depth) of the c values from exponential 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 ReLU and EOC initialization rapidly concentrates 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 afﬁne transformation to the output of activation functions to achieve Q (1) = 1 and C (0)
= 0 , while Lu et al.
(2020) applied them to achieve Q (1)
= 1 and C (cid:48) (1)
= 1 , although the effect of both approaches is similar to EOC.
To address these problems, Martens et al.
(2021) introduced DKS, which enforces the conditions
C f (0) = 0 and C (cid:48)
f
(1) =
ζ (for some modest constant ζ > 1 ) directly on the network’s global C map
C f .
They show that these conditions, along with the positive deﬁniteness of C maps, cause C f to be close to the identity and thus well-behaved.
In addition to these C map conditions, DKS enforces that Q (1) = 1 and Q (cid:48) (1) = 1 , which lead to constant q values of 1 in the network, and lower kernel approximation error (respectively).
DKS enforces these Q/C map conditions by applying a model class-preserving transformation
[Equation]
+ δ ) .
with non-trainable parameters α , β , γ and δ .
The hyperparameter ζ is chosen to be sufﬁciently 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 if ζ = 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 in terms of generalization performance.
T AILORED A CTIVATION T RANSFORMATION (TAT)
The reason why DKS is not fully compatible with ReLUs is that they are positive homogeneous, i.e. φ ( αx ) =
αφ ( x ) for α ≥ 0 .
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 (cid:48) (1) = 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 signiﬁcant 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, or both.
To this end, we adopt a ReLU-family activation function with an extra degree of freedom (known as “Leaky ReLU”), and modify the Q/C map conditions in order to preserve certain desirable properties
[Table]
of this choice.
The resulting method, which we name Tailored Activation Transformation (TAT) achieves competitive generalization performance with ResNets in our experiments.
4. T AILORED A CTIVATION T RANSFORMATION FOR L EAKY R E LU S
One way of addressing the issue of DKS’s partial incompatibility with ReLUs is to consider a slightly different activation function – namely the Leaky ReLU (LReLU) (Maas et al., 2013):
[Equation]
(8) where α is the negative slope parameter .
While using LReLUs with α (cid:54) = 0 in place of ReLUs changes the model class, it doesn’t limit the model’s expressive capabilities compared to ReLU, as assuming α (cid:54) = ± 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 ﬁxed value for α , we will use it as an extra parameter to satisfy our desired Q/C map conditions.
Deﬁne ˜
φ α ( x ) =
(cid:113) 2
+ α 2 φ α ( x ) .
By
Lemma 1, the local Q and C maps for this choice of activation function are:
[Equation]
Note that the condition Q ( q ) =
q is actually stronger than DKS’s Q map conditions ( Q (1) = 1 and
Q (cid:48) (1) = 1 ), and has the potential to reduce kernel approximation errors in ﬁnite width networks compared to DKS, as it provides a better guarantee on the stability of Q f w.r.t.
random perturbations of the q values at each layer.
Additionally, because the form of C 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 on the kernel approximation error exist for
ReLU networks than for general smooth ones (as discussed in Appendix A.1).
Another consequence of using ˜ φ α ( x ) for our activation function is that we have C (cid:48) (1) = 1 as in EOC.
If combined with the condition C (0) =
(which is used to achieve C f (0)
= 0 in DKS) this would imply by Theorem 1 that C is the identity function, which by equation 9 is only true when α
= 1 , thus resulting in a linear network.
In order to avoid this situation, and the closely related one where ˜ φ α appears “nearly linear”, we instead choose the value of α so that C f (0) =
η , for a hyperparameter 0 ≤ η ≤ 1 .
As shown in the following theorem, η controls how close C f is to the identity, thus allowing us to achieve a well-behaved global C map without making ˜ φ α nearly linear:
Theorem 1.
For a network f with ˜
φ α ( x ) as its activation function (with α ≥ 0 ), we have
[Equation]
Another motivation for using ˜ φ α ( x ) as an activation function is given by the following proposition:
Proposition 1.
The global C map of a feedforward network with ˜ φ α ( x ) as its activation function is 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 (cid:113)
α 1+ α 2 , and residual branch R consisting of a combined layer (or just a normalized ReLU activation) followed by an afﬁne layer.
This result implies that at initialization, a vanilla network using ˜ φ α behaves similarly to a ResNet, a property that is quite desirable given the success that ResNets have already demonstrated.
In summary, we have the following three conditions:
[Equation]
which we achieve by picking the negative slope parameter α so that C f (0)
=
η .
We deﬁne the
Tailored Rectiﬁer ( TReLU ) to be ˜ φ α with α chosen in this way.
Note that the ﬁrst two conditions are 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 c value quickly converges towards 1 with depth under
EOC, resulting in a degenerate global C map.
By contrast, the global C map of TReLU for a ﬁxed η converges rapidly to a nice function, suggesting a very deep vanilla network with TReLU has the same well-behaved global C map as a shallow network.
We prove this in Proposition 3 by showing the local C map in equation 9 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 0 ≤ η ≤ 1 , we note that a value very close to 0 will produce a network that is “nearly linear”, while a value very close to 1 will give rise to a degenerate C map.
In practice we use η = 0 .
9 or 0 .
95 , which seems to work well in most settings.
Once we decide on η , we can solve the value α using binary search by exploiting the closed-form form of C in equation 9 to efﬁciently compute C f (0) .
For instance, if f is a 100 layer vanilla network, one can compute C f (0) as follows:
[Equation] [Equation]
which is a function of α .
This approach can be generalized to more advanced architectures, such as rescaled ResNets, as discussed in Appendix B.
4. T AILORED A CTIVATION T RANSFORMATION FOR S MOOTH A CTIVATION F UNCTIONS
Unlike LReLU, most activation functions don’t have closed-form formulas for their local C maps.
As a result, the computation of C f (0) involves the numerical approximation of many two-dimensional integrals to high precision (as in equation 5), which can be quite expensive.
One alternative way to control how close C f is to the identity, while maintaining the condition C (cid:48)
f (1) = 1 , is to modulate its second derivative C (cid:48)(cid:48)
f
The validity of this approach is established by the following theorem:
Theorem 2.
Suppose f is a network with a smooth activation function.
If C (cid:48) f (1) = 1 , then we have
[Equation]
Given C (1) = 1 and C (cid:48) (1) = 1 , a straightforward computation shows that C (cid:48)(cid:48) f (1) =
L C (cid:48)(cid:48) (1) if f is an
L -layer vanilla network.
(See Appendix B for a discussion of how to do this computation for more general architectures.)
From this we obtain the following four local Q/C map conditions:
[Equation]
To achieve these we adopt the same activation transformation as DKS: ˆ φ ( x ) =
γ ( φ ( αx + β ) + δ ) for non-trainable scalars α , β , δ , and γ .
We emphasize that these conditions cannot be used with LReLU, as LReLU networks have C (cid:48)(cid:48) (1)
= ∞ .
By equation 4 and basic properties of expectations, we have
[Equation]
so that γ = E z ∼N (0 , 1) (cid:2) ( φ ( αz + β ) + δ ) 2 (cid:3) − 1 / 2 .
To obtain the values for α , β and δ , we can treat the remaining conditions as a three-dimensional nonlinear system, which can be written as follows:
[Equation]
We do not have a closed-form solution of this system.
However, each expectation is a one dimensional integral, and so can be quickly evaluated to high precision using Gaussian quadrature.
One can then use black-box nonlinear equation solvers, such as modiﬁed Powell’s method (Powell, 1964), to obtain a solution.
See https://github.com/deepmind/dks for a complete implementation.
E XPERIMENTS
Our main experimental evaluation of TAT and competing approaches is on training deep convolutional networks for ImageNet classiﬁcation (Deng et al., 2009).
The goal of these experiments is not to achieve state-of-the-art, 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., 2016b) as the main reference 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 without batch normalization, we add dropout to the penultimate layer for regularization, as was done in Brock et al.
(2021b).
We train the models with 90 epochs and a batch size of 1024 , unless stated otherwise.
For TReLU , we obtain η by grid search in { 0 . 9 , 0 . 95 } .
The weight initialization used for all methods is the Orthogonal Delta initialization, with an extra multiplier given by σ w .
We initialize biases iid from N (0 , σ 2 b ) .
We use ( σ w , σ b ) =
(1 , 0) in all experiments (unless explicitly stated otherwise), with the single exception that we use ( σ w , σ b ) = ( √ 2 , 0) in standard ResNets, as per standard practice (He et al., 2015).
For all other details see Appendix D.
5. T OWARDS 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 et al., 2018;
Shao et al., 2020; Bachlechner et al., 2020; Brock et al., 2021a;b) has adopted this idea by introducing multipliers on the residual branches (which may or may not be optimized during training).
[Table]
In Table
[Table]
we show that one can close most of the gap with standard ResNets by simply adopting the modiﬁcation in equation 6 without using BN layers.
By further replacing ReLU with
TReLU , we can exactly match the performance of standard ResNets.
With
K-FAC as the optimizer, the rescaled
ResNet with shortcut weight w = 0 .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 w = 0 . 8 .
5. T HE DIFFICULTY OF REMOVING SHORTCUT CONNECTIONS
[Table]
While the aforementioned works have shown that it is possible to achieve competitive results without normalization layers, they all rely on the use of shortcut connections to make the network look more linear at initialization.
A natural question to ask is whether normalization layers could compensate for the removal of shortcut connections.
We address this question 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 a signiﬁcant difference, especially with strong optimizers like K-FAC.
These ﬁndings 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.
5. T RAINING D EEP N EURAL N ETWORKS WITHOUT S HORTCUTS
The main motivation for developing TAT is to help deep vanilla networks achieve generalization performance similar to standard ResNets.
In our investigations we include rescaled ResNets with a shortcut weight of either 0 (i.e. vanilla networks) or 0.8.
In Table 4 we can see that with a 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 ﬁrst time that a 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 to the removal of
[Table]
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 shortcut weight of 0 .and
TReLU
[Table]
we can exactly match or even surpass the performance of standard ResNets.
5. C OMPARISONS WITH EXISTING APPROACHES
[Table]
Comparison with EOC.
Our ﬁrst comparison is between TAT and EOC on vanilla deep networks.
For EOC with ReLUs we set ( σ w , σ b ) =
( √ 2 , 0) to achieve Q (1) = 1 as in He et al.
(2015), since ReLU networks always satisfy
C (cid:48) (1) =
whenever σ b = 0 .
For Tanh activations, a comprehensive comparison with EOC is more difﬁcult, as there are inﬁnitely many choices of ( σ w , σ b ) that achieve C (cid:48) (1) = 1 .
Here we use ( σ w , σ b ) =
(1 . 302 , 0 . 02) 2 , as suggested in Hayou et al.
In Table 5, we can see that in all the settings, networks constructed with TAT outperform EOC-initialized networks by a signiﬁcant margin, especially when using SGD.
Another observation is that the accuracy of EOC-initialized networks drops as depth increases.
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 3 , and smooth functions like the SoftPlus and Tanh.
For smooth activations, we perform a grid search over { 0 . 2 , 0 . 3 , 0 . 5 } for τ 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 is the stronger Q/C map conditions that they satisfy compared to other activations (i.e. Q ( q )
=
q for all q vs Q (1) = 1 and Q (cid:48) (1)
= 1 , and invariance of C to the input q value), and the extra resilience to kernel approximation error 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.1) and works equally well with Gaussian initialization (see Appendix E.7).
[Table]
Comparison with PReLU.
The Parametric
ReLU (PReLU) introduced in He et 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 parameter ahead of time and ﬁx it during training.
In our comparisons with PReLU we consider two different initializations: 0 (which recovers the standard ReLU), and 0 .
25 , which was used in
He et al.
We report the results on deep vanilla networks in Table 6 (see Appendix E.6 for results on rescaled ResNets).
For all settings, our method outperforms PReLU by a large margin, emphasizing the importance of the initial negative slope value.
In principle, these two methods can be combined together (i.e. we could ﬁrst initialize the negative slope parameter with TAT, and then optimize it during training), however we did not see any beneﬁt from doing this in our experiments.
[Table]
5. T HE ROLE OF THE OPTIMIZER
[Table]
One interesting phenomenon we observed in our experiments, which echoes the ﬁndings of Martens et al.
(2021), is that a strong optimizer such as K-FAC signiﬁcantly outperforms SGD on vanilla deep networks in terms of training speed.
One plausible explanation 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 efﬁciency begins to drop.
Indeed, it was 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.
[Figure]
To investigate this further, we tried batch sizes between 128 and 4096 for training 50-layer vanilla TReLU networks.
As shown in Table 8, K-FAC performs equally well for all different batch sizes except 4096 (where we see increased overﬁtting), while the performance of SGD starts to drop when we increase the batch size past 512.
Surprisingly, we observe a similar trend for the LARS optimizer (You et al., 2019), which was designed for largebatch 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 reason 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 a 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 ﬁnal 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 to 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 a similar per-step runtime to SGD, while retaining its optimization advantage on vanilla networks.
(See Appendix E.3.)
C ONCLUSIONS
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 modiﬁes the activation functions in a way tailored to the speciﬁc 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 shortcut 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 insufﬁcient.
R EPRODUCIBILITY S TATEMENT
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://github.com/deepmind/dks .
Secondly, we have given all important details of our experiments in Appendix D
A B ACKGROUND A. K ERNEL F UNCTION A PPROXIMATION E RROR B OUNDS
In Section 2.1, we claimed that the kernel deﬁned in equation 2 would converge to a deterministic kernel, as the width of each layer goes to inﬁnity.
To be speciﬁc, one has the following result bounding the kernel approximation error.
Theorem 3 (Adapted from Theorem 2 of Daniely et al.
Consider a fully-connected network of depth L with weights initialized independently using a standard Gaussian fan-in initialization.
Further suppose that the activation function φ is C -bounded (i.e. (cid:107) φ (cid:107) ∞
≤ C , (cid:107) φ (cid:48) (cid:107) ∞ ≤ C and (cid:107) φ (cid:48)(cid:48) (cid:107) ∞
≤ C for some constant C ) and satisﬁes E z ∼N (0 , 1) [ φ ( z ) 2 ] = 1 , and that the width of each layer is greater than or equal to (4 C 4 )
L log(8 L/δ ) /(cid:15) 2 .
Then at initialization time, for inputs x 1 and x 2 satisfying (cid:107) x 1 (cid:107)
=
(cid:107) x 2 (cid:107) 2 = dim( x 1 ) , we have that
[Equation]
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 is only quadratic in L , as is established in the following theorem:
Theorem 4 (Adapted from Theorem 3 of Daniely et al.
Consider a fully-connected network of depth L with ReLU activations and weights initialized independently using a He initialization (He et al., 2015), and suppose that the width of each layer is greater than or equal to L 2 log( L/δ ) /(cid:15) 2 .
Then at initialization time, for inputs x 1 and x 2 satisfying (cid:107) x 1 (cid:107)
=
(cid:107) x 2 (cid:107) 2 = dim( x 1 ) , and (cid:15) (cid:46) 1 L , we have that
[Equation]
with probability at least 1 − δ
[Equation]
According to Lemma D.1 of Buchanan et al.
(2020), the requirement of the width for ReLU networks could further be reduced to linear in the depth L , but with a worse dependency on δ .
Although Theorems 3 and 4 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 be extended to general values of L .
A. D EGENERATE C 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 a restatement of Claim 1 from Daniely et al.
Proposition 2.
Suppose f is a deep network consisting of a composition of L combined layers.
Then for all c ∈ ( − 1 , 1) we have
[Equation]
for some c ∗ ∈ [0 , 1] .
While the above result doesn’t characterize the rate of convergence C f ( c ) to a constant function,
Poole et al.
(2016) show that if C (cid:48) (1) (cid:54)
= 1 , it 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. C MAP DERIVATIVE
Poole et al.
(2016) gave the following nice formula for the derivative of C map of a combined layer with activation function φ :
[Equation]
For a rigorous proof of this result we refer the reader to Martens et al.
One can iterate this formula to obtain a similar equation for higher-order derivatives:
[Equation]
A. S OME USEFUL PROPERTIES OF C MAPS
In this section we will assume that q 1 = q 2 = 1 .
Observe that C (1) =
E z ∼N (0 , 1) (cid:104) φ ( z ) 2 (cid:105) =
1 and that C maps [ − 1 , 1] to [ − 1 , 1] (which follows from its interpretation as computing cosine similarities for inﬁnitely wide networks).
Moreover, C is a positive deﬁnite function, which means that it can be written as (cid:80) ∞ n =
b n c n for b n ≥
(Daniely et al., 2016; Martens et al., 2021).
Note that for smooth activation functions, positive deﬁniteness can be easily veriﬁed by Taylor-expanding C ( c ) about c = 0 and using
[Equation]
As discussed in 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 C maps inherit them from C .
B A DDITIONAL DETAILS AND PSEUDOCODE FOR ACTIVATION FUNCTION TRANSFORMATIONS B. T AKING ALL SUBNETWORKS INTO ACCOUNT
In the main text of this paper we have used the condition C (cid:48)
f (1) =
ζ in DKS, C f (0)
=
η in TAT for
Leaky ReLUs, and C (cid:48)(cid:48) f (1) =
τ in TAT for smooth activation functions.
However, the condition used by Martens et al.
(2021) in DKS was actually µ 1 f ( C (cid:48) (1))
= ζ , where µ 1 f is the so-called “maximal slope function”:
[Equation]
where “ g ⊆ f ” denotes that g is a subnetwork 4 of f .
(That C (cid:48) g (1) is fully determined by C (cid:48) (1) follows from the fact that C g 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 (cid:48) g (1) as a polynomial function of C (cid:48) (1) .)
The motivation given by Martens et al.
(2021) for looking at C (cid:48) g
(1) over all subnetworks g ⊆
f
(instead of just C (cid:48) f (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 to have a reasonable global C map simply by adding a 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 a shorter presentation, and because we happened to have µ 1 f ( C (cid:48) (1))
=
C (cid:48)
f (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 C f (0) =
η and C (cid:48)(cid:48) f
(1) =
τ used in TAT so that they take into account all subnetworks.
This will be done using a natural generalization of the maximal slope function from DKS.
We will then address the computational challenges that result from doing this.
To begin, we will replace the condition C f (0) =
η (used in TAT for Leaky ReLUs) by the condition µ 0 f (0) =
η , where we deﬁne the maximal c value function µ 0 f of f by
[Equation]
where α is the negative slope parameter (which determines C in LReLU networks [via φ α ] and thus each C g ).
We will similarly replace the condition C (cid:48)(cid:48)
f (1) =
τ (used in TAT for smooth activations) by the condition µ 2 f ( C (cid:48)(cid:48) (1))
=
τ , where we deﬁne the maximal curvature function µ 2 f of f by
[Equation]
where each C (cid:48)(cid:48) g (1) is determined by C (cid:48)(cid:48) (1) .
That each C (cid:48)(cid:48) g (1) is a well-deﬁned function of C (cid:48)(cid:48) (1) follows from the fact that C maps always map the value 1 to 1, the aforementioned relationship between C g and
C , and the fact that we have C (cid:48) (1)
= 1 under TAT (so that C (cid:48) h (1) =
1 for all subnetworks h ).
These facts allow us to write C (cid:48)(cid:48) g (1) as a constant multiple of C (cid:48)(cid:48) (1) using the linearity of 2nd derivatives and the 2nd-order chain rule (which is given by ( a ◦ b ) (cid:48)(cid:48)
( x
[Equation]
=
a
[Equation]
( x ) ).
B. C OMPUTING µ 0 f AND µ 2 f IN GENERAL
Given these new conditions for TAT, it remains to compute their left hand sides so that we may ultimately solve for the required quantities ( α or C (cid:48)(cid:48) (1) ).
In Section 2.3 we discussed how a (sub)network f ’s C map C f can be computed in terms of the local C map C by a series of composition and nonnegative weighted sum operations.
We can deﬁne a generalized version of this construction U f,r which replaces C with an arbitrary non-decreasing function r , so that U f, C ( c ) =
C f ( c ) .
A recipe for computing U f,r is given in Appendix B.4.
Given U f,r , we deﬁne the subnetwork maximizing function M by
[Equation]
With this deﬁnition, it is not hard to see that if r 0 ( x ) =
C ( x ) , r 1 ( x ) = C (cid:48) (1) x , and r 2 ( x ) =
C (cid:48)(cid:48) (1)+ x , then µ 0 f ( α ) =
M f,r 0 (0) (where the dependence on α is implicit through the dependence of C on φ α ), µ 1 f ( C (cid:48) (1))
=
M f,r 1 (1) , and µ 2 f ( C (cid:48)(cid:48) (1)) =
M f,r 2 (0) .
Thus, it sufﬁces to derive a scheme for computing (and inverting)
M f,r for general networks f and non-decreasing functions r .
Naively, computing M f,r 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 f,r can simpliﬁed substantially, so that we rarely have to maximize over more than a few possible subnetworks.
In particular, since U g,r ( x ) is a non-decreasing function of x for all g (which follows from the fact that r is non-decreasing), and U g ◦ h,r
= U g,r ◦
U h,r , it thus follows that
U
g ◦ h,r ( x ) ≥ U g,r ( x ) , U h,r ( x ) for all x .
This means that for the purposes of the maximization, we can ignore any subnetwork in f 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 f,r
= U f,r , since not all subnetworks compose in this way.
For example, a sufﬁciently deep residual branch of a residual block in a rescaled ResNet won’t compose with any subnetwork to form a larger one.
B. S OLVING FOR α AND C (cid:48)(cid:48) (1)
Having shown how to efﬁciently compute M f,r , and thus both of µ 0 f and µ 2 f , it remains to show how we can invert them to ﬁnd solutions for α and C (cid:48)(cid:48) (1) (respectively).
Fortunately, this turns out to be easy, as both functions are strictly monotonic in their arguments ( α and C (cid:48)(cid:48) (1) ), provided that f contains at least one nonlinear layer.
Thus, we may apply a simple 1-dimensional root-ﬁnding approach, such as binary search.
To see that µ 0 f ( α ) is a strictly decreasing function of α (or in other words, a strictly increasing function of − α ), we observe that it is a maximum over terms of the form U g, C (0) , which are all either strictly decreasing non-negative functions of α , or are identically zero.
These properties of U g, C (0) follow from the fact that it involves only applications of C , along with compositions and non-negative weighted averages, and that C ( c ) is a strictly decreasing function of α for all c ∈ [ − 1 , 1] (in Leaky
ReLU networks).
A similar argument can be used to show that µ 2 f ( C (cid:48)(cid:48) (1)) is a strictly increasing function of C (cid:48)(cid:48)
(1) (and is in fact equal to a non-negative multiple of C (cid:48)(cid:48) (1) ).
B. R ECIPE FOR COMPUTING U f,r
As deﬁned, U f,r is computed from f by taking the computational graph for C f and replacing the local C map C with r wherever the former appears.
So in particular, one can obtain a computational graph for U f,r ( x ) from f ’s computational graph by recursively applying the following rules:
Composition g ◦ h of two subnetworks g and h maps to U g,r ◦ U h,r . 2.
Afﬁne layers map to the identity function.
Nonlinear layers map to r . 4.
Normalized sums with weights w 1 , w 2 , ..., w n over the outputs of subnetworks g 1 , g 2 ,
[Equation]
., g n
[Equation]
map to the function
[Equation]
where x 1 , x 2 , ..., x n are the respective inputs to the U g i ,r ’s.
f ’s input layer maps to x .
In the special case of computing U f,r 2 (0) , one gets the following simpliﬁed list of rules:
Composition g ◦ h of two subnetworks g and h maps to U g,r 2 (0)
+ U h,r 2 (0) 2.
Afﬁne layers map to 0 .
Nonlinear layers map to C (cid:48)(cid:48) (1) .
Normalized sums with weights w 1 , w 2 , ..., w n over the outputs of subnetworks g 1 , g 2 , ..., g n , map to the function
[Equation]
f ’s input layer maps to x .
Note that this second procedure will always produce a non-negative multiple of C (cid:48)(cid:48) (1) , provided that f contains at least one nonlinear layer.
B. R ESCALED R ES N ET EXAMPLE
In this subsection we will demonstrate how to apply the above rules to compute the maximal curvature function µ 2 f for a rescaled ResNet f with shortcut weight w and residual branch R (as deﬁned in equation 6).
We note that this computation also handles the case of a vanilla network by simply taking w = 0 .
First, we observe that all subnetworks in f compose to form larger ones in f , except for f itself, and for the residual branches of its residual blocks.
We thus have that µ 2 f ( C (cid:48)(cid:48) (1))
=
max { U f,r 2 (0) , U R ,r 2 (0) } .
Because each residual branch has a simple feedforward structure with three nonlinear layers, it follows that U R ,r 2 (0) =
3 C (cid:48)(cid:48) (1) .
And because each shortcut branch S has no nonlinear layers, it follows that U S ,r 2 (0)
= 0 .
Applying the rule for weighted averages to the output of each block B we thus have that U B ,r 2 (0) =
w 2 U S ,r 2 (0) + (1 − w 2 ) U R ,r 2 (0) =
3(1 − w 2 ) C
(cid:48)(cid:48) (1) .
Given a network with L nonlinear layers, we have L/ 3 blocks, and since the blocks compose in a feedforward manner it thus follows that U f,r 2 (0) =
( L/ 3) · 3(1 − w 2 ) C
(cid:48)(cid:48) (1) =
L (1 − w 2 ) C
(cid:48)(cid:48) (1) .
We therefore conclude that µ 2 f ( C (cid:48)(cid:48) (1))
= max { 3 , L (1 − w 2 ) }C (cid:48)(cid:48) (1) .
The rescaled ResNets used in our experiments have a 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 a four of their blocks being “transition blocks”, whose shortcut branches contain a nonlinear layer.
In these networks, the total number of residual blocks is given by ( L − 2) / 3 .
Following a similar argument to the one above we have that
[Equation]
and thus
[Equation]
B. P SEUDOCODE
Algorithm 1 TAT for LReLU.
Require: The target value η for µ 0 f ( α ) 1: Use the steps from Subsection B.2 to construct a procedure for computing the maximal c value function µ 0 f ( α ) for general α ≥ 0 .
Note that the local C map C , on which µ 0 f ( α ) depends, can be computed efﬁciently for (transformed) LReLUs using equation 9.
2: Perform a binary search to ﬁnd the negative slope α such that µ 0 f ( α ) =
η .
3: Using the found α , output the transformed activation function given by ˜ φ α ( x ) = (cid:113) 2 1+ α 2 φ α ( x ) .
Algorithm
Require: The target value τ
of C (cid:48)(cid:48)
f (1) Require:
The original activation function
φ (
x ) 1: Use the steps from Subsection B.2 to construct a procedure for computing the maximal curvature function µ 2 f ( C (cid:48)(cid:48) (1)) for general C (cid:48)(cid:48) (1) ≥ 0 .
2: Perform a binary search to ﬁnd C (cid:48)(cid:48) (1) such that µ 2 f ( C (cid:48)(cid:48) (1))
=
τ . 3: Using a numerical solver, solve the three-dimensional nonlinear system in equation 16 (but with the value of C (cid:48)(cid:48) (1) found above instead of τ/L ) to obtain values for α , β , γ , and δ .
Using the solution from the last step, output the transformed activation function given by ˆ
φ
[Equation]
C T ECHNICAL R ESULTS AND P ROOFS
Lemma 1.
For networks using the activation function ˜
φ α ( x ) =
(cid:113) 2
+ α 2 φ α ( x ) , the local Q and C maps are given by
[Equation]
Proof.
In this proof we will use the notation Q φ and C φ to denote the local Q and C maps for networks that use a given activation function φ .
First, we note that LReLU is basically the weighted sum of identity and ReLU.
In particular, we have the following equation:
[Equation]
Second, we have that Q φ
[Equation]
qz
I [ z <
(cid:3)
=
1+ α 2 2 q
It then follows from equation 5, and the fact that local C maps are invariant to multiplication of the activation function by a constant, that
[Equation]
From Daniely et al.
we have that
[Equation]
and for the last part of equation 21 we have
[Equation]
Plugging equation 22 and equation 23 back into equation
[Equation]
we get
[Equation]
Rearranging this gives the claimed formula.
Proposition 1.
The global C map of a feedforward network with ˜ φ α ( x ) as its activation function is 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 (cid:113)
α 1+ α 2 , and residual branch R consisting of a combined layer (or just a normalized ReLU activation) followed by an afﬁne layer.
Proof.
By equation 7, the C map for a residual block B of the hypothesized rescaled ResNet is given by
[Equation]
The global C 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 L compositions of C ˜
φ α ( c ) .
So to prove the claim it sufﬁces to show that C B ( c )
=
C ˜
φ α ( c ) .
Taking w = (cid:113)
α 1+ α 2 , one obtains the following
[Equation]
which is exactly the same as C ˜
φ α ( c ) as given in Lemma 1.
This concludes the proof.
Proposition 3.
Suppose f is vanilla network consisting of L combined layers with the TReLU activation function (so that C f (0)
=
η ∈ (0 , 1) ).
Then C f converges to a limiting map on ( − 1 , 1) as
L goes to inﬁnity.
In particular,
[Equation] [Equation]
where T is such that ψ (0 , T )
=
η , and where ψ is the solution of the following ordinary differential equation (ODE) with the ﬁrst argument being the initial condition (i.e. ψ ( c, 0)
=
c ), and the second argument being time:
[Equation] [Equation]
Proof.
First, we notice that the local C map for TReLU networks can be written as a difference equation:
[Equation]
Importantly, C is a monotonically increasing function of c , whose derivative goes to zero only as α ∈ [0 , 1] goes to 1 .
Thus, to achieve C f (0)
=
η in the limit of large L , we require that (1 − α ) 2 π (1+ α 2 ) goes to 0.
This implies that the above difference equation converges to the ODE in equation 28.
Because the function √ 1 − x 2 − x cos − 1 ( x ) is continuously differentiable in [ − 1 , 1] , and its derivative − cos − 1 ( x ) is bounded, one can immediately show that it is globally Lipschitz, and the ODE has a unique solution ψ ( c 0 , t ) according to Theorem 3.2 of Khalil (2008).
Now, we are only left to ﬁnd the time T such that C ∞ f (0) =
ψ (0 , T ) =
η .
To that end, we notice that
[Equation]
30) because g (1) =
and g (cid:48)
( x ) = − cos − 1 ( x ) < 0 on ( − 1 , 1) .
This implies that the ψ (0 , t ) is a monotonically increasing continuous function of t .
Since ψ (0 , 0) = 0 , to establish the existence of T it sufﬁces to show that ψ (0 , ∞ ) ≥ 1 .
To this end we ﬁrst observe that
[Equation]
(31) which follows by deﬁning h ( x ) =
g ( x ) − 2 √ 23 (1 − x ) 3 / 2 and observing that h (1) = 0 and h (cid:48) ( x ) = − cos − 1 ( x )+ √ 2(1 − x ) 1 / 2 < 0 on ( − 1 , 1) .
Given this, it is sufﬁcient to show that the solution ˜ ψ for the ODE ˙ x = 2 √ 23 (1 − x ) 3 / 2 satisﬁes ˜ ψ (0 , ∞ )
= 1 .
The solution ˜ ψ turns out to have a closed-form of ˜ ψ (0 , t ) = 1 − ( 3 √ 2 t +3 ) 2 , and thus ψ (0 , ∞ ) ≥ ˜ ψ (0 , ∞ )
= 1 .
This completes the proof.
Theorem 1.
For a network f with ˜
φ α ( x ) as its activation function (with α ≥ 0 ), we have
[Equation]
Proof.
Because C f is a positive deﬁnite function (by Section A.4) we have that it can be written as
C f ( c ) = (cid:80) ∞
n b n c n for b n ≥ 0 .
Given C f (1) =
C (cid:48)
f (1) = 1 , we have ∞ (cid:88) n = b n = ∞ (cid:88) n = nb n ==
⇒
b 0 = ∞ (cid:88) n = ( n − 1) b n = ⇒
2 b 0 + b 1 ≥ 1 .
Hence, 1 − C (cid:48)
f (0)
= 1 − b 1 ≤ 2 b 0 = 2 C f (0) .
Now we are ready to bound the deviation of C f ( c ) from identity:
[Equation]
Using equation 20 we have that
[Equation]
From our assumption that α ≥ 0 it follows that 0 ≤ C (cid:48) ( c ) ≤ 1 for all c ∈ [ − 1 , 1] .
Since the property of having a derivative bounded between 0 and 1 is closed under functional composition and positive weighted averages, it thus follows that 0 ≤ C f ( c ) ≤ 1 for all c ∈ [ − 1 , 1] .
An immediate consequence of this is that C f ( c ) is non-decreasing, and that
[Equation]
Next, we bound the deviation of C (cid:48)
f ( c ) from 1 :
[Equation]
From the previous fact that 0 ≤ C (cid:48) f ( c ) ≤ 1 for all c ∈ [ − 1 , 1] we also have that max c ∈ [ − 1 , 1] (cid:12)(cid:12)(cid:12) C (cid:48)
f ( c ) − 1 (cid:12)(cid:12)(cid:12) ≤ 1 .
This completes the proof.
Theorem 2.
Suppose f is a network with a smooth activation function.
If C (cid:48) f (1) = 1 , then we have
[Equation]
Proof.
C f is a positive deﬁnite function by Section A.4.
So by the fact that positive deﬁnite functions are non-negative, non-decreasing, and convex on the non-negative part of their domain, we obtain that C (cid:48)
[Equation]
(0) ≥ C
[Equation] [Equation]
(1) = 1 − C (cid:48)(cid:48)
f (1) .
By equation 33, we have
[Equation]
Further by equation 35, we also have
[Equation]
This completes the proof.
Proposition 4.
Suppose f is some function computed by a neural network with the ReLU activation.
Then for any negative slope parameter α (cid:54) = ± 1 , we can compute f using an LReLU neural network of the same structure and double the width of the original network.
Proof.
The basic intuition behind this proof is that a ReLU unit can always be “simulated” by two
LReLU units as long as α (cid:54) = ± 1 , due to the following formula:
[Equation]
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:
[Equation]
where x ∈ R
d is the input, and w ∈ R md , b ∈ R m and a ∈ R m are weights, biases of the input layer and weights of output layer, respectively.
For LReLU with negative slope α , one can construct the following network
[Equation]
[Equation]
If we choose w (cid:48) r = w r = − w (cid:48) r + m , b (cid:48)
r = b r = − b (cid:48) r + m , a (cid:48) r = 1 1 − α 2 a r and a (cid:48) r + m = α 1 − α 2 a r , we have a (cid:48)
r φ α ( w (cid:48) r (cid:62) x
+ b (cid:48) r ) + a (cid:48)
r + m φ α ( w (cid:48) r + m
[Equation]
(cid:48)
[Equation]
This immediately suggests that f ( w (cid:48) , b (cid:48) , a (cid:48) , x ) =
f ( w, b, a, x ) .
Since deeper networks, and one with more complex topologies, can be constructed by composing and summing shallower ones, the general claim follows.
D E XPERIMENT DETAILS
For input preprocessing on ImageNet we perform a random crop of size 224 × 224 to each image, and apply a random horizontal ﬂip.
In all experiments, we applied L 2 regularization only to the weights (and not the biases or batch normalization parameters).
We selected the L 2 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 . 0 } .
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 0 to the “initial learning rate” in the ﬁrst 5 epochs, and then decays the learning rate by a factor of 10 at 4 / 9 and 7 / 9 of the total epoch budget 5 , unless speciﬁed otherwise.
The initial learning rate was chosen by grid search from { 1 . 0 , 0 . 3 , 0 . 1 , 0 . 03 , 0 . 01 } for SGD, { 0 . 003 , 0 . 001 , 0 . 0003 , 0 . 0001 , 0 . 00003 } for K-FAC, and { 10 . 0 , 3 . 0 , 1 . 0 , 0 . 3 , 0 . 1 } for LARS.
For all optimizers we set the momentum constant to 0 .
For K-FAC, we used a ﬁxed 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” coefﬁcient to 0 .
For networks with batch normalization layers, we set the decay value for the statistics to 0 .
For initialization of the weights we used the scale-corrected uniform orthogonal (SUO) distribution (Martens et al., 2021) for all methods/models, unless stated otherwise.
For a m × k matrix (with k being the input dimension), samples from this distribution can be generated by computing (cid:0) XX (cid:62) (cid:1) − 1 / 2 X , where X is an m × k matrix with entries sampled independently from N (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 { (cid:112) m/k, 1 } , which will have an effect only when k ≤ m .
For convolutional networks, we initialize only the weights in the center of each ﬁlter to non-zero values, which is a 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://github.com/deepmind/kfac_jax .
E A DDITIONAL E XPERIMENTAL R ESULTS E. E MPIRICAL C VALUES FOR FINITE - WIDTH NETWORKS
The computation of cosine similarities performed by C maps is only an approximation for ﬁnite 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 a ﬁxed η = 0 .
9 for TReLU
and we compute the l -th “empirical c value” ˆ
[Equation]
for each layer index l , where x 01 and x 02 are random vectors
[Figure]
chosen so that (cid:107) x 01 (cid:107)
=
(cid:107) x 02 (cid:107) =
d 0 and x 01 (cid:62) x 02 =
(so that ˆ c 0 = 0 ).
As shown in Figure 4a and 4c, the approximation error is relatively small even for networks with width 30 .
We also included the results for networks using DKS (with ζ = 10 ) and the SoftPlus activation function.
Figure 4b and 4d reports empirical c values as a function of layer index l , with x 01 and x 02 chosen so that ˆ c 0 = 0 . 8 .
With Gaussian initialization, the standard deviations are much larger than TReLU , and the average values for widths 30 and 100 deviate signiﬁcantly from the theoretical predictions.
(The DKS conditions implies C ( c ) ≤ c for any c ∈ [0 , 1] , which suggests the c value should decrease monotonically.)
By comparison, the error seems to be much smaller for orthogonal initialization, which is consistent with the better performance of orthogonal initialization reported by
Martens et al.
(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.
E. R ESULTS ON CIFAR-
[Figure]
In addition to our main results on the ImageNet dataset, we also compared TAT to EOC on CIFAR10 (Krizhevsky et al., 2009) using vanilla networks derived from a Wide ResNet reference architecture (Zagoruyko & Komodakis, 2016).
In particular, we start with a Wide ResNet with a widening factor of 2 , and remove all the batch normalization layers and shortcut connections.
We trained these networks with the K-FAC optimizer for 200 epochs using a standard piecewise constant learning rate schedule.
To be speciﬁc, we decay the learning rate by a factor of 10 at 75 and 150 epochs.
For K-FAC, we set the damping value to 0 .
01 and norm constraint value to 0 .
For data preprocessing we include basic data augmentations such as random crop and horizontal ﬂip during training.
As shown in Figure 5, TAT outperforms EOC signiﬁcantly.
As we increase the depth from 100 to 304 , the accuracy of EOC network drops dramatically while the accuracy of the TAT network remains roughly unchanged
[Figure]
E. R EDUCING THE OVERHEAD OF K-FAC
[Figure]
In our main experiments the per-step wall-clock time of K-FAC was roughly 2 .
5 × that of SGD.
However, this gap can be decreased signiﬁcantly 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
[Figure]
and the inverses every 200 steps, the average per-step wall-clock time of K-FAC reduces by half to a mere 1 .
25 × that of SGD.
Importantly, as can be seen on Figure 6, this does not appear to signiﬁcantly affect optimization performance.
E. D ISENTANGLING T RAINING AND G ENERALIZATION
[Figure]
In our main experiments we only reported validation accuracy on ImageNet, making it hard to tell whether the superior performance of TAT vs EOC is due to improved ﬁtting/optimization speed, or improved generalization.
Here, we 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 using K-FAC for 90 epochs.
For each setting, we plot the training accuracy for the hyperparameter combination that gave the highest ﬁnal validation accuracy.
As shown in Figure 7, the EOC-initialized networks achieve competitive (if not any better) training accuracy, suggesting that the use of TReLU improves the generalization performance and not optimization performance.
E. C LOSING THE R EMAINING G AP USING W IDER N ETWORKS
[Table] [Table]
[Table] [Table]
In all of our main experiments we used networks derived from standard ResNets (by removing normalization layers and/or shortcut connections).
By construction, these have the same layer widths as standard 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 beneﬁt 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 to the removal of shortcut connections.
With layers double the width of standard ResNets, it becomes too expensive to 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 overﬁtting 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% to 0.6% for depth 50 networks, and from 1.4% to 1% for depth 101 networks.
E. C OMPARISON WITH PR E LU ON R ESCALED R ES N ETS
[Table]
In Table 6 of the main text we compare PReLU and TReLU on deep vanilla networks.
Here we extend this comparison to rescaled ResNets with a shortcut weight of w = 0 . 8 .
For PReLU, we again include two different initializations: one with 0 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 a large margin, suggesting that a better-initialized negative slope is crucial for both rescaled ResNets and deep vanilla networks.
E. C OMPARISON OF DIFFERENT INITIALIZATIONS
[Table]
In all of our experiments we use the
Orthogonal Delta initialization introduced by Balduzzi et al.
(2017) and
Xiao et al.
This is because it’s technically required in order to apply the extended Q/C map analysis of Martens et al.
(2021) (which underlies DKS and TAT) to convolutional networks, and because it is generally thought to be beneﬁcial.
In this subsection we examine this choice more closely by comparing it to a traditional Gaussian fan-in initialization (with σ 2 w = 2 for ReLUs).
We consider standard ResNets and deep vanilla networks using either EOC (with ReLUs) or TAT with (with LReLU).
Surprisingly, it turns out 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.