Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers

Introduction

Thanks to many architectural and algorithmic innovations, the recent decade has witnessed the unprecedented success of deep learning in various high-profile challenges, e.g., the ImageNet recognition task (Krizhevsky et al., 2012), the challenging board game of Go (Silver et al., 2017) and human-like text generation (Brown et al., 2020). Among them, shortcut connections (He et al., 2016a; Srivastava et al., 2015) and normalization layers (Ioffe & Szegedy, 2015; Ba et al., 2016) are two architectural components of modern networks that are critically important for achieving fast training at very high depths, and feature prominently in the ubiquitous ResNet architecture of He et al. (2016b).

Despite the success of ResNets, there is significant evidence to suggest that the primary reason they work so well is that they resemble ensembles of shallower networks during training (Veit et al., 2016), which lets them avoid the common pathologies associated with very deep networks (e.g. Hochreiter et al., 2001; Duvenaud et al., 2014). Moreover, ResNets without normalization layers could lose expressivity as the depth goes to infinity (Hayou et 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-inefficient at inference time, as the input to every residual block has to be kept in memory until the final addition. In particular, the shortcut branches in ResNet-50 account for about 40% of the memory usage by feature maps. Also, the classical interpretation of why deep networks perform well – because of the hierarchical feature representations they produce – does not strictly apply to ResNets, due to their aforementioned tendency to behave like ensembles of shallower networks. Beyond the drawbacks of ResNets, training vanilla deep neural networks (which we define as networks without shortcut connections or normalization layers) is an interesting research problem in its own right, and finding a solution could open the path to discovering new model architectures. However, recent progress in this direction has not fully succeeded in matching the generalization performance of ResNets.

Schoenholz et al. (2017) used a mean-field analysis of deep MLPs to choose variances for the initial weights and bias parameters, and showed that the resulting method – called Edge of Chaos (EOC) – allowed vanilla networks to be trained at very high depths on small datasets. Building on EOC, and incorporating dynamical isometry theory, Xiao et al. (2018) was able to train vanilla networks with Tanh unitsDynamical isometry is unavailable for ReLU (Pennington et al., 2017), even with orthogonal weights. 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 a 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 ReLUs, and in their experiments (which focused on training speed) their networks exhibited significantly more overfitting than ResNets.

Inspired by both DKS and the line of work using mean-field theory, we propose a new method called Tailored Activation Transformation (TAT). TAT inherits the main advantages of DKS, while working particularly well with the “Leaky ReLU” activation function. 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 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.

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.

For simplicity, we start with the kernel function approximation for feedforward fully-connected networks, and discuss its extensions to convolutional networks and non-feedforward networks later. In particular, we will assume a network that is defined by a sequence of $L$ combined layers (each of which is an affine transformation followed by the elementwise activation function $\phi$) as follows:

2 Local Q/C maps

By equation 3, any diagonal entry $q_{i}^{l+1}$ of $\Sigma^{l+1}$ only depends on the corresponding diagonal entry $q_{i}^{l}$ of $\Sigma^{l}$. Hence, we obtain the following recursion for these diagonal entries, which we call q values:

where $\mathcal{Q}$ is the local Q map. We note that $q_{i}^{l}$ is an approximation of $\kappa_{f}^{l}({x}_{i},{x}_{i})$. Analogously, one can write the recursion for the normalized off-diagonal entries, which we call c values, as:

where $\mathcal{C}$ is the local C map and $c^{0}={x}_{1}^{\top}{x}_{2}/d_{0}$. We note that $c^{l}$ is an approximation of the cosine similarity between $f_{{\theta}}^{l}({x}_{1})$ and $f_{{\theta}}^{l}({x}_{2})$. Because $\mathcal{C}$ is a three dimensional function, it is difficult to analyze, as the associated q values can vary wildly for distinct inputs. However, by scaling the inputs to have norm $\sqrt{d_{0}}$, and rescaling $\phi$ so that $\mathcal{Q}(1)=1$, it follows that $q_{i}^{l}=1$ for all $l$. This allows us to treat $\mathcal{C}$ as a one dimensional function from $$to$$ satisfying $\mathcal{C}(1)=1$. Additionally, it can be shown that $\mathcal{C}$ possesses special structure as a positive definite function (see Appendix A.4 for details). Going forward, we will thus assume that $q_{i}^{0}=1$, and that $\phi$ is scaled so that $\mathcal{Q}(1)=1$.

3 Extensions to convolutional networks and more complex topologies

As argued in Martens et al. (2021), Q/C maps can also be defined for convolutional networks if one adopts a Delta initialization (Balduzzi et al., 2017; Xiao et al., 2018), in which all weights except those in the center of the filter are initialized to zero. Intuitively, this makes convolutional networks behave like a collection of fully-connected networks operating independently over feature map locations. As such, the Q/C map computations for a feed-forward convolutional network are the same as above. Martens et al. (2021) also gives formulas to compute q and c values for weighted sum operations between the outputs of multiple layers (without nonlinearities), thus allowing more complex network topologies. In particular, the sum operation’s output q value is given by $q=\sum_{i=1}^{n}w_{i}^{2}q_{i}$, and its output c value is given by $\frac{1}{q}\sum_{i=1}^{n}w_{i}^{2}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 $\sum_{i=1}^{n}w_{i}^{2}=1$.

Following Martens et al. (2021), we will extend the definition of Q/C maps to include global Q/C maps, which describe the behavior of entire networks. Global maps, denoted by $\mathcal{Q}_{f}$ and $\mathcal{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 $\mathcal{C}_{f}(c)=\mathcal{C}\circ\mathcal{C}\circ\mathcal{C}(c)$. Like the local C map, global C maps are positive definite functions (see Appendix A.4). In this work, we restrict our attention to the family of networks comprising of combined layers, and normalized sums between the output of multiple affine layers, for which we can compute global Q/C maps. And all of our formal results will implicitly assume this family of networks.

4 Q/C maps for rescaled ResNets

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 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 at the end of each block with a normalized sum. These networks, which we will call rescaled ResNets, are defined by the following recursion:

where $\mathcal{R}$ is the residual branch, and $w$ is the shortcut weight (which must be in $$). Applying the previously discussed rules for computing Q/C maps, we get$q_{i}^{l}=1$for all$l$ and

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 vector, so that $\mathcal{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 $\mathcal{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 first 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.

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 $\mathcal{C}^{\prime}(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 affine transformation to the output of activation functions to achieve $\mathcal{Q}(1)=1$ and $\mathcal{C}(0)=0$, while Lu et al. (2020) applied them to achieve $\mathcal{Q}(1)=1$ and $\mathcal{C}^{\prime}(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 $\mathcal{C}_{f}(0)=0$ and $\mathcal{C}_{f}^{\prime}(1)=\zeta$ (for some modest constant $\zeta>1$) directly on the network’s global C map $\mathcal{C}_{f}$. They show that these conditions, along with the positive definiteness of C maps, cause $\mathcal{C}_{f}$ to be close to the identity and thus well-behaved. In addition to these C map conditions, DKS enforces that $\mathcal{Q}(1)=1$ and $\mathcal{Q}^{\prime}(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 $\hat{\phi}(x)=\gamma(\phi(\alpha x+\beta)+\delta)$. with non-trainable parameters $\alpha$, $\beta$, $\gamma$ and $\delta$. The hyperparameter $\zeta$ 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 if $\zeta=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.

Tailored Activation Transformation (TAT)

The reason why DKS is not fully compatible with ReLUs is that they are positive homogeneous, i.e. $\phi(\alpha x)=\alpha\phi(x)$ for $\alpha\geq 0$. This makes the $\gamma$ 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 $\mathcal{Q}^{\prime}(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 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, 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 of this choice. The resulting method, which we name Tailored Activation Transformation (TAT) achieves competitive generalization performance with ResNets in our experiments.

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

Note that the condition $\mathcal{Q}(q)=q$ is actually stronger than DKS’s Q map conditions ($\mathcal{Q}(1)=1$ and $\mathcal{Q}^{\prime}(1)=1$), and has the potential to reduce kernel approximation errors in finite width networks compared to DKS, as it provides a better guarantee on the stability of $\mathcal{Q}_{f}$ w.r.t. random perturbations of the q values at each layer. Additionally, because the form of $\mathcal{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).

In summary, we have the following three conditions:

For the hyperparameter $0\leq\eta\leq 1$, we note that a value very close to 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 $\eta=0.9$ or $0.95$, which seems to work well in most settings. Once we decide on $\eta$, we can solve the value $\alpha$ using binary search by exploiting the closed-form form of $\mathcal{C}$ in equation 9 to efficiently compute $\mathcal{C}_{f}(0)$. For instance, if $f$ is a $100$ layer vanilla network, one can compute $\mathcal{C}_{f}(0)$ as follows:

which is a function of $\alpha$. This approach can be generalized to more advanced architectures, such as rescaled ResNets, as discussed in Appendix B.

2 Tailored Activation Transformation for Smooth Activation Functions

Unlike LReLU, most activation functions don’t have closed-form formulas for their local C maps. As a result, the computation of $\mathcal{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 $\mathcal{C}_{f}$ is to the identity, while maintaining the condition $\mathcal{C}_{f}^{\prime}(1)=1$, is to modulate its second derivative $\mathcal{C}_{f}^{\operatorname*{{\prime\prime}}}(1)$. The validity of this approach is established by the following theorem: {thm}[] Suppose $f$ is a network with a smooth activation function. If $\mathcal{C}_{f}^{\prime}(1)=1$, then we have

Given $\mathcal{C}(1)=1$ and $\mathcal{C}^{\prime}(1)=1$, a straightforward computation shows that $\mathcal{C}_{f}^{\operatorname*{{\prime\prime}}}(1)=L\mathcal{C}^{\operatorname*{{\prime\prime}}}(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:

To achieve these we adopt the same activation transformation as DKS: $\hat{\phi}(x)=\gamma(\phi(\alpha x+\beta)+\delta)$ for non-trainable scalars $\alpha$, $\beta$, $\delta$, and $\gamma$. We emphasize that these conditions cannot be used with LReLU, as LReLU networks have $\mathcal{C}^{\operatorname*{{\prime\prime}}}(1)=\infty$. By equation 4 and basic properties of expectations, we have

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 modified Powell’s method (Powell, 1964), to obtain a solution. See https://github.com/deepmind/dks for a complete implementation.

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 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 $\eta$ 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 $\sigma_{w}$. We initialize biases iid from $\mathcal{N}(0,\sigma_{b}^{2})$. We use $(\sigma_{w},\sigma_{b})=(1,0)$ in all experiments (unless explicitly stated otherwise), with the single exception that we use $(\sigma_{w},\sigma_{b})=(\sqrt{2},0)$ in standard ResNets, as per standard practice (He et al., 2015). For all other details see Appendix D.

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

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 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.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 $w=0.8$.

2 The difficulty of removing shortcut connections

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

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. 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 first 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 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.8$ and TReLU, we can exactly match or even surpass the performance of standard ResNets.

4 Comparisons with existing approaches

Comparison with EOC. Our first comparison is between TAT and EOC on vanilla deep networks. For EOC with ReLUs we set $(\sigma_{w},\sigma_{b})=(\sqrt{2},0)$ to achieve $\mathcal{Q}(1)=1$ as in He et al. (2015), since ReLU networks always satisfy $\mathcal{C}^{\prime}(1)=1$ whenever $\sigma_{b}=0$. For Tanh activations, a comprehensive comparison with EOC is more difficult, as there are infinitely many choices of $(\sigma_{w},\sigma_{b})$ that achieve $\mathcal{C}^{\prime}(1)=1$. Here we use $(\sigma_{w},\sigma_{b})=(1.302,0.02)$We also ran experiments with $(\sigma_{w},\sigma_{b})=(1.0,0.0)$, and the scheme described in Pennington et al. (2017) and Xiao et al. (2018) for dynamical isometry. The results were worse than those reported in the table., as suggested 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 a significant 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 LReLUsFor DKS, we set the negative slope as a parameter and adopt the transformation $\hat{\phi}(x)=\gamma(\phi_{\alpha}(x+\beta)+\delta)$., and smooth functions like the SoftPlus and Tanh. For smooth activations, we perform a grid search over $\{0.2,0.3,0.5\}$ for $\tau$ in TAT, and $\{1.5,10.0,100.0\}$ for $\zeta$ 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. $\mathcal{Q}(q)=q$ for all $q$ vs $\mathcal{Q}(1)=1$ and $\mathcal{Q}^{\prime}(1)=1$, and invariance of $\mathcal{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).

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 fix it during training. In our comparisons with PReLU we consider two different initializations: (which recovers the standard ReLU), and $0.25$, which was used in He et al. (2015). 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 first initialize the negative slope parameter with TAT, and then optimize it during training), however we did not see any benefit from doing this in our experiments.

5 The role of the optimizer

One interesting phenomenon we observed in our experiments, which echoes the findings of Martens et al. (2021), is that a strong optimizer such as K-FAC significantly 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 efficiency 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.

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

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 a 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 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 insufficient.

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://github.com/deepmind/dks. Secondly, we have given all important details of our experiments in Appendix D.

References

Appendix A Background

with probability at least $1-\delta$. 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 $\delta$.

Although Theorems A.1 and A.1 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.2 Degenerate 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. (2016): {prop} Suppose $f$ is a deep network consisting of a composition of $L$ combined layers. Then for all $c\in(-1,1)$ we have

While the above result doesn’t characterize the rate of convergence $\mathcal{C}_{f}(c)$ to a constant function, Poole et al. (2016) show that if $\mathcal{C}^{\prime}(1)\neq 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.3 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 $\phi$:

For a 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:

A.4 Some useful properties of C maps

In this section we will assume that $q_{1}=q_{2}=1$.

As discussed in Section 2.3, global C maps are computed by recursively taking compositions and weighted averages (with non-negative weights), starting from $\mathcal{C}$. Because all of the above properties are preserved under these operations, it follows that global C maps inherit them from $\mathcal{C}$.

Appendix B Additional details and pseudocode for activation function transformations

In the main text of this paper we have used the condition $\mathcal{C}^{\prime}_{f}(1)=\zeta$ in DKS, $\mathcal{C}_{f}(0)=\eta$ in TAT for Leaky ReLUs, and $\mathcal{C}^{\prime\prime}_{f}(1)=\tau$ in TAT for smooth activation functions. However, the condition used by Martens et al. (2021) in DKS was actually $\mu_{f}^{1}(\mathcal{C}^{\prime}(1))=\zeta$, where $\mu_{f}^{1}$ is the so-called “maximal slope function”:

where “$g\subseteq f$” denotes that $g$ is a subnetworkA subnetwork of $f$ is defined as a (non-strict) connected subset of the layers in $f$ that constitute a neural network with a singular input and output layer. So for example, layers 3, 4 and 5 of a 10 layer MLP form a subnetwork, while layers 3, 4, and 6 do not. of $f$. (That $\mathcal{C}^{\prime}_{g}(1)$ is fully determined by $\mathcal{C}^{\prime}(1)$ follows from the fact that $\mathcal{C}_{g}$ can be written in terms of compositions, weighted average operations, and applications of $\mathcal{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 $\mathcal{C}^{\prime}_{g}(1)$ as a polynomial function of $\mathcal{C}^{\prime}(1)$.)

The motivation given by Martens et al. (2021) for looking at $\mathcal{C}^{\prime}_{g}(1)$ over all subnetworks $g\subseteq f$ (instead of just $\mathcal{C}^{\prime}_{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 $\mu_{f}^{1}(\mathcal{C}^{\prime}(1))=C^{\prime}_{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 $\mathcal{C}_{f}(0)=\eta$ and $\mathcal{C}^{\prime\prime}_{f}(1)=\tau$ 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 $\mathcal{C}_{f}(0)=\eta$ (used in TAT for Leaky ReLUs) by the condition $\mu_{f}^{0}(0)=\eta$, where we define the maximal c value function $\mu_{f}^{0}$ of $f$ by

where $\alpha$ is the negative slope parameter (which determines $\mathcal{C}$ in LReLU networks [via $\phi_{\alpha}$] and thus each $\mathcal{C}_{g}$).

We will similarly replace the condition $\mathcal{C}^{\prime\prime}_{f}(1)=\tau$ (used in TAT for smooth activations) by the condition $\mu_{f}^{2}(\mathcal{C}^{\prime\prime}(1))=\tau$, where we define the maximal curvature function $\mu_{f}^{2}$ of $f$ by

where each $\mathcal{C}^{\prime\prime}_{g}(1)$ is determined by $\mathcal{C}^{\prime\prime}(1)$. That each $\mathcal{C}^{\prime\prime}_{g}(1)$ is a well-defined function of $\mathcal{C}^{\prime\prime}(1)$ follows from the fact that C maps always map the value 1 to 1, the aforementioned relationship between $\mathcal{C}_{g}$ and $\mathcal{C}$, and the fact that we have $\mathcal{C}^{\prime}(1)=1$ under TAT (so that $\mathcal{C}_{h}^{\prime}(1)=1$ for all subnetworks $h$). These facts allow us to write $\mathcal{C}^{\prime\prime}_{g}(1)$ as a constant multiple of $\mathcal{C}^{\prime\prime}(1)$ using the linearity of 2nd derivatives and the 2nd-order chain rule (which is given by $(a\circ b)^{\prime\prime}(x)=a^{\prime\prime}(b(x))b^{\prime}(x)^{2}+a^{\prime}(b(x))b^{\prime\prime}(x)$).

Given these new conditions for TAT, it remains to compute their left hand sides so that we may ultimately solve for the required quantities ($\alpha$ or $\mathcal{C}^{\prime\prime}(1)$). In Section 2.3 we discussed how a (sub)network $f$’s C map $\mathcal{C}_{f}$ can be computed in terms of the local C map $\mathcal{C}$ by a series of composition and non-negative weighted sum operations. We can define a generalized version of this construction $U_{f,r}$ which replaces $\mathcal{C}$ with an arbitrary non-decreasing function $r$, so that $U_{f,\mathcal{C}}(c)=\mathcal{C}_{f}(c)$. A recipe for computing $U_{f,r}$ is given in Appendix B.4.

Given $U_{f,r}$, we define the subnetwork maximizing function $M$ by

With this definition, it is not hard to see that if $r_{0}(x)=\mathcal{C}(x)$, $r_{1}(x)=\mathcal{C}^{\prime}(1)x$, and $r_{2}(x)=\mathcal{C}^{\prime\prime}(1)+x$, then $\mu_{f}^{0}(\alpha)=M_{f,r_{0}}(0)$ (where the dependence on $\alpha$ is implicit through the dependence of $\mathcal{C}$ on $\phi_{\alpha}$), $\mu_{f}^{1}(\mathcal{C}^{\prime}(1))=M_{f,r_{1}}(1)$, and $\mu_{f}^{2}(\mathcal{C}^{\prime\prime}(1))=M_{f,r_{2}}(0)$. Thus, it suffices 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 simplified 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\circ h,r}=U_{g,r}\circ U_{h,r}$, it thus follows that $U_{g\circ h,r}(x)\geq 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 sufficiently deep residual branch of a residual block in a rescaled ResNet won’t compose with any subnetwork to form a larger one.

Having shown how to efficiently compute $M_{f,r}$, and thus both of $\mu_{f}^{0}$ and $\mu_{f}^{2}$, it remains to show how we can invert them to find solutions for $\alpha$ and $\mathcal{C}^{\prime\prime}(1)$ (respectively). Fortunately, this turns out to be easy, as both functions are strictly monotonic in their arguments ($\alpha$ and $\mathcal{C}^{\prime\prime}(1)$), provided that $f$ contains at least one nonlinear layer. Thus, we may apply a simple 1-dimensional root-finding approach, such as binary search.

To see that $\mu_{f}^{0}(\alpha)$ is a strictly decreasing function of $\alpha$ (or in other words, a strictly increasing function of $-\alpha$), we observe that it is a maximum over terms of the form $U_{g,\mathcal{C}}(0)$, which are all either strictly decreasing non-negative functions of $\alpha$, or are identically zero. These properties of $U_{g,\mathcal{C}}(0)$ follow from the fact that it involves only applications of $\mathcal{C}$, along with compositions and non-negative weighted averages, and that $\mathcal{C}(c)$ is a strictly decreasing function of $\alpha$ for all $c\in$ (in Leaky ReLU networks). A similar argument can be used to show that $\mu_{f}^{2}(\mathcal{C}^{\prime\prime}(1))$ is a strictly increasing function of $\mathcal{C}^{\prime\prime}(1)$ (and is in fact equal to a non-negative multiple of $\mathcal{C}^{\prime\prime}(1)$).

As defined, $U_{f,r}$ is computed from $f$ by taking the computational graph for $C_{f}$ and replacing the local C map $\mathcal{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\circ h$ of two subnetworks $g$ and $h$ maps to $U_{g,r}\circ U_{h,r}$.

Affine layers map to the identity function.

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

where $x_{1},x_{2},...,x_{n}$ are the respective inputs to the $U_{g_{i},r}$’s.

In the special case of computing $U_{f,r_{2}}(0)$, one gets the following simplified list of rules:

Composition $g\circ h$ of two subnetworks $g$ and $h$ maps to $U_{g,r_{2}}(0)+U_{h,r_{2}}(0)$

Nonlinear layers map to $\mathcal{C}^{\prime\prime}(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

Note that this second procedure will always produce a non-negative multiple of $\mathcal{C}^{\prime\prime}(1)$, provided that $f$ 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 $\mu_{f}^{2}$ for a rescaled ResNet $f$ with shortcut weight $w$ and residual branch $\mathcal{R}$ (as defined 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 $\mu_{f}^{2}(\mathcal{C}^{\prime\prime}(1))=\max\{U_{f,r_{2}}(0),U_{\mathcal{R},r_{2}}(0)\}$.

Because each residual branch has a simple feedforward structure with three nonlinear layers, it follows that $U_{\mathcal{R},r_{2}}(0)=3\mathcal{C}^{\prime\prime}(1)$. And because each shortcut branch $\mathcal{S}$ has no nonlinear layers, it follows that $U_{\mathcal{S},r_{2}}(0)=0$. Applying the rule for weighted averages to the output of each block $\mathcal{B}$ we thus have that $U_{\mathcal{B},r_{2}}(0)=w^{2}U_{\mathcal{S},r_{2}}(0)+(1-w^{2})U_{\mathcal{R},r_{2}}(0)=3(1-w^{2})\mathcal{C}^{\prime\prime}(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)\cdot 3(1-w^{2})\mathcal{C}^{\prime\prime}(1)=L(1-w^{2})\mathcal{C}^{\prime\prime}(1)$. We therefore conclude that $\mu_{f}^{2}(\mathcal{C}^{\prime\prime}(1))=\max\{3,L(1-w^{2})\}\mathcal{C}^{\prime\prime}(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

B.6 Pseudocode

Appendix C Technical Results and Proofs

In this proof we will use the notation $\mathcal{Q}_{\phi}$ and $\mathcal{C}_{\phi}$ to denote the local Q and C maps for networks that use a given activation function $\phi$.

First, we note that LReLU is basically the weighted sum of identity and ReLU. In particular, we have the following equation:

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

and for the last part of equation 21 we have

Plugging equation 22 and equation 23 back into equation 21, we get

Rearranging this gives the claimed formula. ∎

By equation 7, the C map for a residual block $\mathcal{B}$ of the hypothesized rescaled ResNet is given by

Taking $w=\sqrt{\frac{2\alpha}{1+\alpha^{2}}}$, one obtains the following

[] Suppose $f$ is vanilla network consisting of $L$ combined layers with the TReLU activation function (so that $\mathcal{C}_{f}(0)=\eta\in(0,1)$). Then $\mathcal{C}_{f}$ converges to a limiting map on $(-1,1)$ as $L$ goes to infinity. In particular,

where $T$ is such that $\psi(0,T)=\eta$, and where $\psi$ is the solution of the following ordinary differential equation (ODE) with the first argument being the initial condition (i.e. $\psi(c,0)=c$), and the second argument being time:

First, we notice that the local C map for TReLU networks can be written as a difference equation:

Importantly, $\mathcal{C}$ is a monotonically increasing function of $c$, whose derivative goes to zero only as $\alpha\in$ goes to $1$. Thus, to achieve $\mathcal{C}_{f}(0)=\eta$ in the limit of large $L$, we require that $\tfrac{(1-\alpha)^{2}}{\pi(1+\alpha^{2})}$ goes to 0. This implies that the above difference equation converges to the ODE in equation 28.

Because the function $\sqrt{1-x^{2}}-x\cos^{-1}(x)$ is continuously differentiable in $$, and its derivative$-\cos^{-1}(x)$is bounded, one can immediately show that it is globally Lipschitz, and the ODE has a unique solution$\psi(c_{0},t)$ according to Theorem 3.2 of Khalil (2008).

Now, we are only left to find the time $T$ such that $\mathcal{C}_{f}^{\infty}(0)=\psi(0,T)=\eta$. To that end, we notice that

because $g(1)=0$ and $g^{\prime}(x)=-\cos^{-1}(x)<0$ on $(-1,1)$. This implies that the $\psi(0,t)$ is a monotonically increasing continuous function of $t$. Since $\psi(0,0)=0$, to establish the existence of $T$ it suffices to show that $\psi(0,\infty)\geq 1$.

Because $\mathcal{C}_{f}$ is a positive definite function (by Section A.4) we have that it can be written as $\mathcal{C}_{f}(c)=\sum_{n}^{\infty}b_{n}c^{n}$ for $b_{n}\geq 0$. Given $\mathcal{C}_{f}(1)=\mathcal{C}_{f}^{\prime}(1)=1$, we have

Hence, $1-\mathcal{C}^{\prime}_{f}(0)=1-b_{1}\leq 2b_{0}=2\mathcal{C}_{f}(0)$. Now we are ready to bound the deviation of $\mathcal{C}_{f}(c)$ from identity:

From our assumption that $\alpha\geq 0$ it follows that $0\leq\mathcal{C}^{\prime}(c)\leq 1$ for all $c\in$. 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\leq\mathcal{C}_{f}(c)\leq 1$ for all $c\in$. An immediate consequence of this is that $\mathcal{C}_{f}(c)$ is non-decreasing, and that

Next, we bound the deviation of $\mathcal{C}_{f}^{\prime}(c)$ from $1$:

From the previous fact that $0\leq\mathcal{C}^{\prime}_{f}(c)\leq 1$ for all $c\in$ we also have that $\max_{c\in}\left|\mathcal{C}_{f}^{\prime}(c)-1\right|\leq 1$. This completes the proof. ∎

$\mathcal{C}_{f}$ 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 $\mathcal{C}^{\prime}_{f}(0)\geq\mathcal{C}^{\prime}_{f}(1)-C^{\operatorname*{{\prime\prime}}}_{f}(1)=1-C^{\operatorname*{{\prime\prime}}}_{f}(1)$. By equation 33, we have

[] Suppose $f$ is some function computed by a neural network with the ReLU activation. Then for any negative slope parameter $\alpha\neq\pm 1$, we can compute $f$ using an LReLU neural network of the same structure and double the width of the original network.

The basic intuition behind this proof is that a ReLU unit can always be “simulated” by two LReLU units as long as $\alpha\neq\pm 1$, due to the following formula:

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:

where $x\in^{d}$ is the input, and $w\in^{md}$, $b\in^{m}$ and $a\in^{m}$ are weights, biases of the input layer and weights of output layer, respectively. For LReLU with negative slope $\alpha$, one can construct the following network

If we choose $w_{r}^{\prime}=w_{r}=-w_{r+m}^{\prime}$, $b_{r}^{\prime}=b_{r}=-b_{r+m}^{\prime}$, $a_{r}^{\prime}=\frac{1}{1-\alpha^{2}}a_{r}$ and $a_{r+m}^{\prime}=\frac{\alpha}{1-\alpha^{2}}a_{r}$, we have

This immediately suggests that $f(w^{\prime},b^{\prime},a^{\prime},{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. ∎

Appendix D Experiment details

For input preprocessing on ImageNet we perform a random crop of size $224\times 224$ to each image, and apply a random horizontal flip. 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 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 budgetWe later found that cosine learning rate annealing (Loshchilov & Hutter, 2016) is slightly better for most settings, but this did not change our conclusions., unless specified 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.9$. For K-FAC, we used a 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 to $0.9$.

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\times k$ matrix (with $k$ being the input dimension), samples from this distribution can be generated by computing $\left(XX^{\top}\right)^{-1/2}X$, where $X$ is an $m\times k$ matrix with entries sampled independently from $\mathcal{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\{\sqrt{m/k},1\}$, which will have an effect only when $k\leq m$. For convolutional networks, we initialize only the weights in the center of each filter 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.

Appendix E Additional Experimental Results

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 a fixed $\eta=0.9$ for TReLU and we compute the $l$-th “empirical c value” $\hat{c^{l}}=\frac{{x_{1}^{l}}^{\top}x_{2}^{l}}{\|x_{1}^{l}\|\|x_{2}^{l}\|}$ for each layer index $l$, where $x_{1}^{0}$ and $x_{2}^{0}$ are random vectors chosen so that $\|x_{1}^{0}\|^{2}=\|x_{2}^{0}\|=d_{0}$ and ${x_{1}^{0}}^{\top}x_{2}^{0}=0$ (so that $\hat{c^{0}}=0$). As shown in Figure 4(a) and 4(c), the approximation error is relatively small even for networks with width $30$.

We also included the results for networks using DKS (with $\zeta=10$) and the SoftPlus activation function. Figure 4(b) and 4(d) reports empirical c values as a function of layer index $l$, with $x_{1}^{0}$ and $x_{2}^{0}$ chosen so that $\hat{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 significantly from the theoretical predictions. (The DKS conditions implies $\mathcal{C}(c)\leq c$ for any $c\in$, 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. (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.

E.2 Results on CIFAR-10

In addition to our main results on the ImageNet dataset, we 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 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 specific, 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.0001$. For data preprocessing we include basic data augmentations such as 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 of EOC network drops dramatically while the accuracy of the TAT network remains roughly unchanged.

E.3 Reducing the overhead of K-FAC

In our main experiments the per-step wall-clock time of K-FAC was roughly $2.5\times$ 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 $1.25\times$ 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 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 fitting/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 final 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.5 Closing the Remaining Gap using Wider Networks

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 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 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 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% to 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 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 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.7 Comparison of different initializations

In all of our experiments we use the Orthogonal Delta initialization introduced by Balduzzi et al. (2017) and Xiao et al. (2018). 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 beneficial. In this subsection we examine this choice more closely by comparing it to a traditional Gaussian fan-in initialization (with $\sigma_{w}^{2}=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.