Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszcz, Valentin Dalibard, Jascha Sohl-Dickstein, Samuel S. Schoenholz
Introduction
The current standard approach to deep learning relies on a combination of architectural elements including skip connections, normalization layers, and carefully chosen activation functions (such as RELU) to overcome the well-documented optimization difficulties present in traditional deep neural networks (Ioffe and Szegedy, 2015; He et al., 2016a; Szegedy et al., 2017). While this approach has proven very successful, enabling many applications in diverse fields such as vision (e.g. He et al., 2016a; Tan and Le, 2019), language (e.g. Vaswani et al., 2017; Brown et al., 2020), protein folding (Jumper et al., 2021) and reinforcement learning (e.g. Espeholt et al., 2018; Silver et al., 2018), it is not entirely satisfying for at least several reasons.
First, the precise mechanism of action of these elements, as well as their interaction, is still not well understood, despite some recent progress in this area. This lack of understanding makes it difficult to design new network architectures, as architectural choices not only affect the network’s expressivity, but also its trainability, and in ways that are hard to predict. Second, without competitive alternatives to compare to, it’s not clear whether the current standard approach enables deep networks to reach their full potential, or whether it has unseen drawbacks and limitations. For example, while the use of skip connections helps very deep networks to train much faster, this might only be because it makes them behave like an ensemble of shallower networks (Veit et al., 2016). Finally, the extra complexity introduced by these architectural elements, and their non-trivial interactions, makes theoretical analyses much more difficult, potentially holding us back from developing a more fundamental understanding of deep learning. And while existing theoretical analyses can (and often do) drop these elements, they do so at the risk of missing an essential piece of the picture.
In an ideal world, modelling and trainability would be decoupled, so that architectures could be designed with only modelling considerations in mind, and rapid training would be guaranteed as long as they conformed to a well-defined set of rules. One might also hope that the components of such a framework would each have a clear purpose, be theoretically well-understood, and interact with each other in simple and predictable ways.
In the present work we take an important step towards this “ideal world”, while simultaneously providing a competitive alternative to the current standard approach to deep learning. We do so by developing a theoretically well-founded method for constructing deep networks which allows them to be rapidly trained without the use of skip connections, normalization layers, or standard activation functions. Our approach, which we call Deep Kernel Shaping (DKS), requires only minor model class-preserving modifications to the architecture and activation functions, and is fully compatible with existing analysis frameworks such as Neural Tangent Kernel (NTK) theory (Jacot et al., 2018).
As we show in experiments, DKS enables very deep residual networks without normalization layers to be trained using SGD on Imagenet and CIFAR-10 classification tasks at similar speeds to standard ResNetV2 (He et al., 2016b) and Wide-ResNet models (Zagoruyko and Komodakis, 2016). It also achieves the same for networks without skip connections or normalization layers when combined with stronger optimizers like K-FAC (Martens and Grosse, 2015) or Shampoo (Gupta et al., 2018). Moreover, it works well with a large variety of activation functions, including those that traditionally perform very poorly (such as the logistic sigmoid). As a caveat, we observe a small decrease in generalization performance compared to standard ResNets, which we believe can be addressed in future work.
While there have been some recently proposed methods for training very deep networks without skip connections and normalization layers (e.g Schoenholz et al., 2017; Balduzzi et al., 2017; Xiao et al., 2018), to the best of our knowledge, DKS is the first to achieve training speeds competitive with standard ResNet models on a challenging dataset like Imagenet. And while our use of K-FAC plays an important role in these results, our experiments show that K-FAC alone is not enough, even when used in combination with the aforementioned methods.
The starting point for our development of DKS is the work of Poole et al. (2016), who described the approximate initialization-time behavior of fully-connected combined layers (which we define here as an affine layer followed by an element-wise nonlinear layer) using special one-dimensional maps known as “Q and C maps”, and then used the fixed-point behavior of these maps to describe the depth-limiting behavior of a network composed of many such layers in sequence. We also take inspiration from Schoenholz et al. (2017), who applied this analysis framework to design an initialization method which modulated the fixed point behavior of each layer’s C map to slow the loss of “geometric information” with depth, and demonstrated encouraging results training very deep networks without skip networks or normalization layers.
While originally derived within the semi-rigorous framework of “mean field analysis”, it turns out that Q/C maps also describe the approximate behavior of a combined layer’s kernel function in wide networks, and as we will show, can be applied to convolutional layers if one uses a Delta initialization for the filter banks. These maps can be further extended to describe entire networks with arbitrary topologies, where they provide useful information outside of the depth-limiting case considered by Poole et al. (2016). In the case of a fully connected network , its Q map approximates the mapping from to , and its C map approximate the mapping from to .
In deeper networks, the C map can easily become “degenerate”, mapping most of its input domain . The implication of this is that the distance between any pair of output vectors from the network is effectively independent of the distance between the corresponding pair of input vectors. As we argue, both heuristically and using NTK theory, this behavior inevitably leads to very slow training and/or poor generalization under gradient descent.
We thus design DKS to prevent this problem, while also guarding against certain secondary pathologies such as a badly behaved Q map, high approximation error in the Q/C maps themselves, and network behavior that is “too linear” (which limits network expressivity under gradient descent). To do this, we relate the overall “shape” of the network’s C map and its tendency to become degenerate, to its value and derivative at a couple of points, which we in turn relate to the values of derivatives of the C maps for the network’s individual layers. We then control these properties (and a couple of additional ones to address the aforementioned secondary pathologies), by transforming each activation function using a model class preserving scale and shift operation its input and output. This transformation is the same for each nonlinear layer for a given activation function, but depends on the global structure of the network. In theory, we also require that sum operations in the network are “normalized” in a certain way, that a special kind of data preprocessing is used, and that pooling layers are replaced with certain roughly equivalent alternatives, although we find the latter two of these to be non-essential in practice.
In addition to developing DKS, we also use Q/C maps to help explain the effectiveness of standard deep learning techniques such as normalization layers, skip connections, initialization methods, and common activation functions such as RELU and SELU (Klambauer et al., 2017), in terms of their effect on the network’s initialization-time kernel. This is facilitated in part by the connections we establish between Q/C maps and alternative analysis frameworks such as “variance propagation” and “signal propagation” which underlie many of the said techniques.
Outline
This manuscript is organized into five parts.
Part I gives our assumptions and establishes the theoretical concepts used in subsequent parts. We begin in Sections 3 and 4 by defining our notation and stating our initial assumptions on network architecture and initialization. In Section 5 we discuss kernel functions for networks conforming to these assumptions, and how they can be approximated with much simpler functions at initialization time. In Sections 6 and 7 we show how these kernel approximations can be further broken down in terms of a generalized version of the Q/C maps originally proposed in Poole et al. (2016). Derivative computations for Q/C maps are given in Section 8, and Section 9 discusses how to handle sum operations when computing Q/C maps. In Section 10, we show how C maps can be simplified down to one dimensional functions (from three dimensions) using a special type of data preprocessing which is designed to make two of their three inputs constant. And in Section 11 we discuss additional consequences of this, including that C maps become “positive definite functions”.
With the theoretical groundwork established, Part II focuses on identifying desirable Q/C map properties and ways to achieve these. In Section 12 we discuss C map behavior in deep networks and how it can – and usually does – become “degenerate”, leading to slow training and/or poor generalization. We then set out to analyze C maps with the hope of controlling their properties so as to prevent this. To that end, in Section 13 we use the positive definiteness of a C map to show how its deviation from the identity function (which is large in degenerate maps) can be predicted from its derivative at and value at , implying that we can prevent degeneration by enforcing certain conditions on these quantities. In Section 14 we identify another way that a network can fail to be trainable: that its parameters must move very far from their initial values before the network can exhibit any significantly nonlinear behavior. We then show this failure mode can be avoided by enforcing a condition on the C map of each nonlinear layer. In Section 15 we identify the breakdown of our kernel approximations as a third problem that we must avoid, and propose several solutions to this, including a condition to enforce on the network’s Q map.
Having identified three distinct ways that a network can fail to be trainable, and conditions to enforce on the Q and C maps to prevent or mitigate these failures, we proceed with the specification and derivation of DKS in Part III. In Section 16 we list the four conditions on the Q and C maps of the network (or more precisely its “subnetworks”) which we will enforce. Then in Section 17 we show how these conditions can be reduced to ones on the Q/C maps of the individual layers of the network via a special translation mechanism called the “maximal slope function” which encodes structural information about the network (including its depth). In Section 18 we describe our main mechanism of enforcement for these per-layer conditions: scaling and shifting operations applied to the input and output of each nonlinear layer’s activation function (which preserve the model class). Finally, in in Sections 19 and 20 we discuss how to deal with normalization and pooling layers in DKS. With DKS fully derived, we give a step-wise summary of it in Section 21.1, and provide details for the more difficult aspects of its implementation in Section 22. In Section 23 we demonstrate the application of DKS on the various modified ResNet and Wide-ResNet models which we use in our experiments, including ones with skip connections and/or normalization layers removed.
Before proceeding to experiments, in Part IV we delve deeper into the theory underlying DKS, and analyze various related approaches from the perspective of kernel approximations and Q/C maps. In Section 24 we review Neural Tangent Kernel (NTK) theory, and give an elegant expression for the NTK using (extended) C maps. We then show how NTK theory predicts slow training and poor generalization for networks with degenerate C maps, and characterize the form of the NTK for networks constructed using DKS. In Section 25 we review certain previously published methods for understanding the behavior of neural networks at initialization time (such as variance/signal propagation), show how they give rise to what are essentially Q and C maps (but different interpretations for what they actually compute), and advocate for the use of approximate kernel analysis as a more flexible and mathematically rigorous alternative. Exploiting these connections, we then review and analyze some prior methods for initializing and constructing neural networks in Section 26, including standard techniques such as normalization layers and residual networks, as well as methods aimed at replacing them. In each case we argue that the method can interpreted as enforcing some set of conditions on the network’s Q/C map, which is often a strict subset of those enforced by DKS.
Finally, in Part V we discuss experiments and conclude. This begins in Sections 27 and 28, where we describe the setup of our experiments, and discuss their results. Our experiments include comparisons of DKS to standard ResNets, the methods reviewed/analyzed in Section 26, and various “ablated”/modified versions of DKS. We then summarize our conclusions in Section 29, and in Section 30 discuss the limitations of DKS and possible ways to address them in future work.
Table of contents
Part I Theoretical preliminaries
Throughout this work we will assume that the reader is already familiar with convolutional neural networks (Fukushima and Miyake, 1982; LeCun et al., 1998a) for which many overviews and tutorials are available (e.g. Goodfellow et al., 2016, Chapter 9). The purpose of this subsection won’t be to define convolutional network concepts from scratch, but rather to lay out the specific terminology we will use when referring to them.
In this work will consider neural networks consisting of affine layers of the standard fully-connected and convolutional types, and nonlinear layers that compute element-wise activation functions (that are typically nonlinear). We define a combined layer to be an affine layer, immediately followed by a nonlinear layer. (Note that a combined layer is what was traditionally referred to as a “layer” in the neural network literature, before the modern trend of referring to the individual affine and nonlinear parts as their own separate “layers”.)
In this work, the discussion will center around a single neural network which we will refer to simply as the network or sometimes the entire network. We will define a subnetwork as a neural network formed from a subset of the entire network’s layers which preserves all dependency relationships and has a well-defined and singular input and output (unlike the entire network, which can have multiple inputs and outputs in general). Subnetworks can be thought of as performing part of the computation of the network. So for example, if the network consists of a sequence of five layers, then layers 2, 3 and 4 form a subnetwork whose input is the input to layer 2, and whose output is the output of layer 4. But layers 2, 4, and 5 do not form a subnetwork since the dependency of layer 4 on layer 2 is not preserved.
2 Initial architectural assumptions
We observe that a fully-connected layer is equivalent to an convolutional layer with a 1x1 feature map and 1x1 filter size, where the input/output data dimensions are just the input/output channel dimensions. Thus, going forward, we will restrict our analysis to the convolution case, which implicitly handles the fully-connected case via this reduction.
We will also assume, for now, that the network can be entirely built out of three components: combined layers (as define above), non-zero constant scalar multiplications operations applied to individual feature maps, and concatenation operations, which concatenate two feature maps of compatible sizes along their channel dimensions. We will permit a given feature map to act as the input to multiple operations/layers in the network, thus allowing “branching structures” and multiple “output heads”.
The restriction to combined layers isn’t as severe as it might seem, as an isolated affine layer is equivalent to a combined layer with an identity activation function. And while sum operations are not explicitly included among the allowed operations, under certain conditions they can be simulated via a simple construction whose details we will defer to Section 9. This means that our analysis can apply to networks containing actual sum operations, under said conditions. Two or more consecutive nonlinear layers are also not allowed by our assumptions, however one can simply fuse two such layers into a single one by composing their activation functions.
For now we will assume that the network does not contain any pooling layers. We will (partially) relax this assumption later in Section 20.
Parameter distributions
In order to obtain a sufficiently simple characterization of the function computed by a neural network at initialization time, we will make certain assumptions about the distribution of its parameters at initialization.
Our first one will be that the bias vector is initialized to zero. While not strictly necessary to the derivation and viability of DKS, this assumption will simplify our presentation. Our second will be that if the input of one layer depends on the output of another, either directly or indirectly, then the parameters of these layers must be initialized independently from each other. This rules out recurrent neural networks, for example, since parameters are shared across time-steps.
Finally, except where stated otherwise, we will assume the use of a “Delta initialization” (Balduzzi et al., 2017; Xiao et al., 2018), which requires that filter bank tensors are initialized to zero everywhere except for their central location/offset (and have odd-sized filter dimensions to make this possible). As an example, if we have a filter, then only the weights corresponding to entry would be non-zero. Note that for fully-connected layers there is only one location, so that a Delta initialization becomes equivalent to a standard one.
The non-zero weights of a Delta-initialized filter bank form a matrix, where is the input channel dimension and is the output channel dimension. To initialize this matrix we have two options. First, we can use an entry-wise iid Gaussian distribution with mean 0 and variance , which gives rise to the Gaussian Delta initialization. While it might seem restrictive to assume a variance of (instead of for general ), this will simplify our presentation going forward, and other choices can be simulated by rescaling the network’s activation functions (which will be part of DKS).
The second option is to use a scaled-corrected uniform orthogonal (SUO) distribution, which is a special distribution of rescaled orthogonal matrices. When , samples from this distribution can be generated as , where is a matrix with entries sampled iid from . When , we may apply the same procedure but with and reversed, and then transpose the result. The resulting distribution is given by the well-known Haar measure on orthogonal matrices (e.g. Meckes, 2019), and is also sometimes called the uniform distribution. To be consistent with the scaling characteristics of the Gaussian initialization, we further multiply by the scaling factor , which will have an effect only when . We will call Delta initializations that use the SUO distribution Orthogonal Delta initializations.
2 A brief discussion about random orthogonal matrices and the SUO distribution
The scaled-corrected uniform orthogonal distribution, as we have defined it, has the property that it is invariant to pre or post-multiplication of the matrix by a constant square orthonormal matrix (Eaton, 1989, Chapter 7). This implies that left-multiplying an input vector by an unobserved matrix sampled from this distribution erases all information about the vector’s direction. The input vector’s dimension-normalized squared norm (i.e. ) can meanwhile be exactly recovered when , and is equal to the output vector’s dimension-normalized squared norm.
For the computations in the next section to be valid for a given orthogonal weight distribution, we require that the distribution satisfies these properties. However, many randomized procedures used in practice for sampling orthogonal matrices lack the directional invariance property. And even procedures whose distributions do possess it often don’t include the scale correction factor, which is required for the dimension-normalized squared norm to be preserved. Thus, we strongly recommend that anyone implementing DKS use the sampling procedure for orthogonal matrices that we have outlined, unless they are confident that their own procedure gives precisely the same distribution. Note that Saxe et al. (2014) and Xiao et al. (2018) have used distributions over orthogonal matrices to initialize neural networks. It turns out that the formulas they derive also require SUO-distributed weights to be correct, even though they did not state this explicitly.
Finally, note that the entry-wise iid distribution for matrices behaves very similarly to the SUO distribution with respect to multiplication by an input vector, and gives a distribution on the output vector which is identical up to a multiplication by a random scalar (which is distributed according to the chi distribution with degrees of freedom). The output vector’s dimension-normalized squared norm is thus a random multiplicative perturbation of the input vector’s (instead of being equal to it), where the perturbation’s mean and variance are and respectively. From these observations we can see that Gaussian initializations, like SUO ones, give rise to directional invariance, but only approximately preserve the dimension-normalized squared norm (and in a way that gets more precise as grows).
Kernel function approximations for neural networks
The starting point for our analysis of the initialization-time behavior of neural networks will be kernel functions, and the approximations of these that hold at initialization-time when the channel dimensions are large. This type of analysis was originally pioneered by Neal (1996), and developed further in various subsequent works (e.g. Williams, 1997; Rahimi and Recht, 2008; Cho and Saul, 2009; Mairal et al., 2014; Anselmi et al., 2015; Hazan and Jaakkola, 2015; Daniely et al., 2016; Matthews et al., 2018; Lee et al., 2018; Garriga-Alonso et al., 2018; Novak et al., 2018; Arora et al., 2019). In this section we will review these concepts and establish our notation and terminology for the key quantities. We will depart from the index-heavy tensor notation of some previous works (such as Novak et al., 2018) in favor of a more compact one based on matrices.
Before we launch into our full treatment of kernel function approximations for convolutional neural networks, in this subsection we will quickly give a simplified version for the fully-connected case, with the goal of building intuition. Note that the notation defined here is only a special case of the more general notation we will develop in subsequent subsections.
It turns out (e.g. Daniely et al., 2016) that when is a sufficiently wide fully-connected combined layer with iid weights and activation function , is closely approximated with high probability by , where
This can be derived by observing that any two units in ’s nonlinear layer are Gaussian distributed (when conditioned on and ) with mean zero and covariance matrix . And so if we consider enough of these units, their average statistics (given ) converge in probability to the expectation.
Using the notable fact that only depends on (and not the full details of and ), we can then compose these layer-wise kernel approximations to form ones for networks consisting of many such layers.
2 Notation for feature maps and subnetworks
We will represent subnetworks (of which single layers are a special case) by symbols such as “” or “”. Implicit in these representations is a dependence on all the structural details of the subnetwork, including its parameters, its activation functions, and anything else we need in order to construct our various approximations. At the same time, we will use standard functional notation such as when we want to treat as a function from its input to its output.
3 Inner product matrices (IPMs) and Pair-location kernel functions (PKFs)
The entries of an IPM are the (dimension-normalized) inner products between all pairs of column vectors from and , or in other words, the average (across channels) of the entry-wise products between pairs of location vectors from the feature maps and .
’s PKF gives us a “geometric view” of ’s input-output behavior. In particular, because determines the inner-products between all pairs of output vectors (across the different locations and both inputs), it determines the distances between all such vectors via the formula .
4 Initialization-time approximations to the PKF for combined layers
In this subsection we will assume that is a combined layer with element-wise activation function . We will also assume Gaussian-distributed weights, as part of either a Delta or non-Delta initialization scheme. (An extension to SUO-distributed weights given in Subsection 5.9).
We are interested in extracting a simple mathematical approximation of that is valid at initialization time, which we can use in order to construct approximations of the PKF of larger subnetworks. To begin with, we will assume that the convolutional part of uses padding and has a stride of 1, which means that input and output locations will be in one to one correspondence. (This assumption will be relaxed in the next subsection.)
In general, computing for combined layers boils down to direct evaluation of the defining formula, with no simplifications possible. But when ’s initial parameters are distributed as per Section 4, there exists a much simpler function that approximates at initialization time with high probability, which we call the approximate paired-location kernel function (or APKF) of . is obtained from by taking the limit as the output channel dimension go to infinity, and is a good approximation when the actual (finite) output channel dimension is sufficiently large.
As shown by Garriga-Alonso et al. (2018) and Novak et al. (2018), the APKF for convolutional combined layers initialized with a standard Gaussian fan-in initializationThis initialization uses an entry-wise iid Gaussian distribution with mean 0 and variance , where is the filter size times the input channel dimension. (LeCun et al., 1998b) is given by
where is the operator which maps to , with denoting the matrix of patch vectorsA “patch vector” is one formed by concatenating together the subset of columns of corresponding to a particular location visited by the convolutional filter. They have dimension for convolutions. generated from . A key property of is that it only depends on and via the associated IPM .
As discussed in Section 4, we are assuming the use of a Delta initialization scheme in this work. Intuitively, a Delta initialization makes a convolutional layer behave like a set of fully-connected layers that operate independently over locations in the feature map (and share parameters). This results in a simplified form for which is a directly analogous to the kernel approximation for fully-connected combined layers (i.e. Equation 3). It is given byThis formula can be obtained from Equation 4 by observing that a Delta-initialized filter bank behaves like a 1x1 filter, and that is the identity operator in the case of a 1x1 filter (since ).
A minor technical point is that may be singular, in which case will be “degenerate”, and its density function technically undefined. The easiest way this can happen is if . However, one can still meaningfully define a distribution and sample from it using for , which is equivalent to adding to and then letting . With this extended definition of our formulas remain valid.
5 Padding, strides, and dropped locations
If the stride of ’s convolution is not 1, or if it doesn’t use padding and has a filter size larger than , then the locations in the input and output feature maps won’t be in one to one correspondence. Instead, they will be related to each other via a projection function , which maps input locations (given by the entries of ) to their corresponding output locations (given by the entries of ). This results in the following generalized formula for :
When the input and output locations are in one to one correspondence, is just the identity function. Otherwise, essentially “drops” the input locations that are never visited by the center of the filter (i.e. will be independent of the entries of that are “dropped”). We will refer to these as dropped locations, and most of our discussions going forward will assume that the location under consideration has not been dropped at the layer in question. When a location is dropped at some layer, both the exact and approximate PKFs of that layer (and all subsequent layers) will be effectively zero for that location.
6 Deriving APKFs given Gaussian distributed weights
At a high level, the APKF formulas given above for a combined layer can be derived by observing that each pair of outputs from the affine part of are linear combinations of Gaussian random variables (i.e. the filter weights) when conditioned on the two inputs and , and are thus are jointly Gaussian distributed with mean zero. A straightforward computation then shows that the covariance matrix of this distribution is or , where denotes the Kronecker product. Because units in different channels have zero covariance they are independent, and so is equal to an average over output channels of iid random variables, and thus converges in probability to its expectation as the number of output channels goes to infinity. We set equal to this expectation, whose formula then follows from the one for . Probabilistic bounds on the approximation error can then be obtained using concentration inequalities.
7 The APKF Condition and network-level PKF approximations
The main approximation which we will use going forward is that the PKF of each combined layer is equal to its associated APKF at initialization time. Or in other words, that
for each combined layer of the network. We will refer to this as the APKF Condition.
Observe that a combined layer’s APKF depends on and only through the associated IPM . Thus, under the APKF Condition, we can compose APKFs for each combined layer to form an initialization-time approximations of the PKFs for arbitrary subnetworks, which we will call network-level PKF approximations. Extending our notation from the combined layer case, we will denote these approximations by for arbitrary subnetworks . (Note that we rely on the property that subnetworks have a single input and output feature maps for this definition and notation to make sense.)
An additional complication that we must deal with when constructing network-level PKF approximations is the presence of concatenation operations, where is the concatenation of two feature maps and along their channel dimensions. In this cases, we observe that
where and are the number of channels in and respectively.
As we will see in the following sections, network-level PKF approximations are amenable to detailed analysis, and expose several key properties which end up being crucial determinants of network trainability (and which can be controlled through careful interventions).
8 How accurate are these approximations?
As discussed above, the APKF for a combined layer is derived by observing that the entries of are empirical averages of iid variables that converge in probability to their expectations as the output channel dimension goes to infinity. Applying concentration inequalities then leads to statements of the form: “for any and there exists an integer such that if the output channel dimension satisfies then with probability .”. The precise dependency of and on and is of practical interest, as the output channel dimension of real neural network layers is finite, and may not even be particularly large in some cases.
Ultimately, we are interested in bounding the kernel approximation error not just for single combined layers but for entire networks. In general, such bounds will be worse than anything provable for single layers, as approximation error will compound with depth (since the output of one approximation is fed as input into the next). The only work we are aware of that gives such bounds is that of Daniely et al. (2016). In that work, the authors analyze what are essentially networks of fully-connected combined layers arranged in arbitrary topologies, with certain technical conditions imposed on their input data and activation functions. Translating their main result into the language and assumptions of this work yields the following theorem:
where is maximum number of nonlinear layers in any input-output path through the network (i.e. its “depth”), is its number of combined layers, and . Then at initialization time, for all input vectors and to satisfying , we have that
Note that in our notation, both and are matrices, and extracts the -th entry, or in other words, the value of and its approximation. One can estimate the error for diagonal entries simply by setting .
The bound in Theorem 1 predicts an exponential dependence of the minimum required width and depth , and a dependence on the error tolerance . The exponential dependence on means that this bound could never realistically be applied to a moderately deep network running on actual hardware, as the required width would be prohibitive. While it could easily be the case that some choices of give an exponential dependence as the bound predicts, we conjecture that with more carefully designed assumptions on the properties of , a bound with better dependence could be proven. Indeed, Daniely et al. (2016) themselves give a more specialized bound for networks with rescaled RELU activations (which technically violate the hypotheses of Theorem 1 since they are unbounded and not differentiable everywhere), where the required width is only quadratic in .
The main limitation of Theorem 1 is that it applies only to networks of fully-connected combined layers that don’t share weights. We conjecture that a similar result may also hold for networks with convolutional layers and a restricted type of inter-layer weight sharing.
9 The orthogonal initialization case (assuming SUO-distributed weights)
The kernel formulas and theory given so far in this section have all assumed the use of Gaussian Delta initializations. However, our assumptions also permit the use of Orthogonal Delta initializations, which as discussed in Section 4, use the SUO distribution instead of an iid Gaussian one to initialize the non-zero weights of the filter. While some previous works (e.g. Xiao et al., 2018) have used these kinds of kernel approximation formulas in the orthogonal case, and have appealed to the vague notion that random orthogonal matrices “look like” Gaussian-distributed ones in high dimensions, there hasn’t been any mathematically rigorous justification of this practice until the recent work of Martens (2021).
The following theorem, which is adapted from Martens (2021), establishes convergence in probability of the APKF to the associated PKF for a fully-connected combined layer with SUO-distributed weight matrix. Like Theorem 1, it provides an explicit and fairly reasonable convergence rate. An extension of this result to multi-layer networks would likely proceed along similar lines to the argument given in Daniely et al. (2016) for the Gaussian case.
Let be a fully-connected combined layer with an SUO distributed weight matrix , a bias vector equal to , and an activation function satisfying for some (with denoting the supremal value). Denote , and suppose that for we have
The conditions on , , and in the theorem statement will be satisfied as long as is sufficiently large and is not too much larger than . In the case where , the LHS’s of these bounds simplifies to approximately and , respectively. It thus follows that the APKF converges in probability to the PKF as the output dimension goes to .
In the case where , the conditions imply that
which is similar to the width bound from Theorem 1 for .
Note that while this theorem is stated only for fully-connected combined layers, it also applies to convolutional combined layers that use Orthogonal Delta initializations by taking and to be any pair of vectors from the union of the columns and .
Q and C maps for combined layers
Q maps and C maps are mathematical constructs introduced by Saxe et al. (2014) and Poole et al. (2016) that describe the initialization time behavior of deep fully-connected networks. While original derived within the semi-rigorous “signal propagation” framework (which is discussed in Section 25.6), they can also be applied under certain conditions within the more rigorous context of kernel function approximations. In that context, they provide a compact alternative representation of approximate kernel functions that is easier to work with.
As will be discussed later in Part II, the Q/C maps of a network tell us a lot about its trainability. Indeed, they have appeared either implicitly or explicitly, often in simplified forms, in much of the previous work on network design and initialization (as will be made clear in Sections 25 and 26). They are also central to the derivation of DKS, and over the next few sections we will develop the generalized version of them that we will use in this work.
In this section we will formally introduce Q/C maps maps and their associated notation, and give formulas to compute them for combined layers under our stated hypotheses. Note that while the connection between Q/C maps and approximate kernel functions has been previously observed (e.g. Lee et al., 2018), it hasn’t before been carefully worked out, nor has it been generalized to convolutional layers (as we will do here). In the section that follows we will show how Q/C maps can be naturally extended beyond single combined layers to describe the behavior of network-level PKF approximations for arbitrary subnetworks with complex topologies.
Consider a combined layer with as its element-wise activation function, and and as its two inputs. By Equation 5 and basic properties of Gaussian expectations, any given diagonal entry of depends only on the corresponding diagonal entry of , and can be computed as
where is defined as the Q map of . We will call such diagonal entries q values, and note that they are equal to the dimension-normalized squared norms of their associated location vectors under the APKF Condition. Notably, the form of the Q map is the same for each location, and so we may associate them with combined layers in a location-independent way.
2 C maps for combined layers
An off-diagonal entry of has a slightly more complex dependence on in Equation 5, as it depends on both the corresponding entry of , as well as the two associated diagonal entries ( and ) that share a row or column. It is given by
We call such off-diagonal entries m values, and note that they are equal to the dimension-normalized inner product of their two associated location vectors under the APKF Condition.
Following Poole et al. (2016), we focus on “length-normalized” versions of the m values called c values. A c value can be obtained from an m value by dividing it by the square root of the product of its two associated q values. (e.g. in the context of Equation 8.) Under the APKF Condition, c values are equal to the cosine similarity between their two associated location vectors.
c values are computed using C maps, which for a combined layer are given by
where we have used the fact that and are mean-zero Gaussian distributed with covariance matrix \left[\begin{array}[]{cc}q_{1}&\sqrt{q_{1}q_{2}}c_{\operatorname{in}}\\ \sqrt{q_{1}q_{2}}c_{\operatorname{in}}&q_{2}\end{array}\right]=\left[\begin{array}[]{cc}q_{1}&m_{\operatorname{in}}\\ m_{\operatorname{in}}&q_{2}\end{array}\right]. Like the Q map, the C map is the same for each location, and so we may associate a single C map to each combined layer.
3 Q/C maps for more general combined layers?
Note that the existence of Q/C maps, as we have defined them, depends on our stated hypotheses for combined layers. In particular, that they are convolutional (or fully-connected), and use a Delta initialization scheme. While APKFs do exist for certain other layer types and initialization schemes, they may not always give rise to low dimensional maps that fully describe their behavior. For example, if we use a conventional fan-in initialization instead of a Delta initialization for the filter weights, then the resulting APKF (given in Equation 4) implies a more complex dependence of the entries of the output IPM on the input IPM, where output q values will depend on (many) input c values.
Extended Q and C maps
In Poole et al. (2016) and Schoenholz et al. (2017), the neural networks analyzed were assumed to be sequences of fully-connected combined layers, each with the same activation function. Thus, the network’s initialization-time behavior could be approximated using a single per-layer Q/C map composed with itself times, and a dynamical systems analysis of this map could thus be performed. This analysis looked for the map’s stable points and attractors, and characterized its asymptotic behavior as the number of self-compositions (i.e. the network’s depth) went to infinity.
In this work we consider architectures with a more general structure, and with layers that can be convolutional and employ a variety of activation functions. We are also interested in the given architecture’s finite structure, instead of its depth-limiting behavior, as this will allow us to more carefully tailor our manipulations to the given network. To facilitate this, in this section we will extend the notion of Q maps and C maps to arbitrary subnetworks (consisting of potentially many layers) in the natural way.
Going forward, we will refer to Q maps and C maps defined specifically for combined layers as local Q/C maps, and maps defined specifically for larger subnetworks, via the extension procedure defined in the next subsection, as extended Q/C maps. Unqualified, “Q/C maps” will be a general term referring to both.
The definition for extended Q/C maps is the natural generalization of the definition for local Q/C maps, where we replace APKF approximations for combined layers with network-level PKF approximations for subnetworks. In particular, given a subnetwork , an extended Q map maps input q values, corresponding to the diagonal entries of the input IPM , to the associated output q values, corresponding to the diagonal entries of the associated output IPM , as computed by the network-level PKF approximation . The definition for extended C maps is similar.
That these definitions can be made in a location-independent way (as with the definitions of local Q/C maps), follows from the fact that extended Q/C maps can be constructed from local ones via composition and weighted averaging (as will be detailed below), which are both operations that preserve the location-independence property.
2 Computing extended maps
Because Q maps compose with each other, and C maps compose with the combination of both, we can take the per-combined-layer maps and compose them in a way that mirrors the composition of the subnetwork’s combined layers, analogously to how we assembled network-level PKF approximations from APKF approximations of each combined layer. For example, if we have two consecutive combined layers and , and wish to compute the Q and C map for the subnetwork consisting of their composition, this is simply and .
The only complication is that we need to describe how q and c values can be computed when feature maps are concatenated along their channel dimensions, or when they are multiplied by a non-zero scalar constant. To handle the former situation, we recall from Equation 6 that concatenation leads to a weighted averaging of the feature maps’ associated IPMs, with weights given by their respective number of channels. Thus, the q values, which are the diagonal entries of these matrices, average in the same way under concatenation. So given the channel dimensions and , and the q values and , we have that the associated q value of the concatenation is simply
c values are slightly more complicated to deal with, but still relatively straightforward. We note that m values, the unnormalized counterparts of c values, are the off-diagonal entries of the IPMs, and thus exhibit the same kind of averaging as q values. We can thus obtain the c values by first converting them to m values, performing the required weighted average, and then converting back to c values. This gives us the analogous formula
where refers to the -th q value associated with the c value from the -th feature map being concatenated. (Recall that each c value is associated to two q values.)
Note that the property that local C maps send carries over to extended C maps, as this clearly preserved under composition and weighted averages.
To handle multiplication of a feature map by a constant , we note that the IPM of and is equal to times the IPM of and , or in other words: . We thus have that an output q value (or m value) for such an operation is simply times the corresponding input q value (or m value). And an output c value is just equal to the corresponding input c value, since the constant will cancel out when we divide by the geometric mean of the q values.
3 Generalization to subnetworks with isolated affine and nonlinear layers
Because it will simplify Q/C map computations for certain architectures (such as residual networks), we will also generalize Q/C maps to subnetworks that may contain affine or nonlinear layers in isolation (i.e. separated from their parent combined layer). To do this, we will define local Q/C maps for isolated affine and nonlinear layers in a way that is consistent with our previous definitions (with one small proviso), and then use the previous composition argument to extend Q/C maps to larger subnetworks containing such layers.
The isolated affine layer case is trivial, as an affine layer is equivalent to a combined layer with an identity activation function, and so is covered under the previous discussion. It follows that affine layers have local Q and C maps that are the identity function (which can easily be verified by setting in Equations 7 and 19), and can thus be essentially ignored in the extended map computations.
The case of nonlinear layers is more subtle. APKFs, from which local Q and C maps are defined, don’t actually exist for nonlinear layers in isolation. In particular, for arbitrary input vectors it is not the case that one can closely approximate the norm of the output vector given only the norm of the input vector (with high probability). However, when the layer is part of a larger network in which its input vector is always the output of some affine layer (with a suitable parameter distribution), such a prediction can be made, and is given by the APKF for the corresponding combined layer. (To see this, note affine layers have identity Q and C maps and thus the input to the nonlinear layer has the same q and c values as the input to the corresponding combined layer.)
Thus, we can define the local Q and C map for an isolated nonlinear layer to be equal to the local Q and C maps for its associated combined layer, with the proviso that it describes the layer’s kernel behavior only for “typical” input vectors (i.e. those that are produced with high probability by the previous layers’ computation) and not arbitrary input vectors.
Note that these definitions are consistent with our definitions for combined layers, as the composition of the local Q/C map for an affine and nonlinear layer (as we have defined them here) does indeed recover the local Q/C map of the associated combined layer. Also, it should be emphasized that these arguments rely crucially on the fact that nonlinear layers may be “isolated” only from the point of the view of a given subnetwork. From the perspective of the entire network, it is still required that they are always part of a combined layer, or in other words, are always directly preceded by an affine layer.
Q and C map derivative computations
Central to our analysis of Q and C maps are their derivatives, which encode many of the properties that we will care about. In this section we show how to compute them, first for local maps, and then for extended maps of arbitrary subnetworks.
A conceivable approach to computing the derivatives of local maps would be to derive a closed form expression for the required integrals, and then apply standard differentiation techniques. Unfortunately, closed form expressions for these integrals are not generally available for most the activation functions. Instead, following Poole et al. (2016), we will give integral expressions for the derivatives which are similar to the original maps themselves, and which can be efficiently approximated using numerical integration (as discussed in Section 22.2).
Let be a combined layer (or an isolated nonlinear layer) with element-wise activation function .
The derivative for with respect to , which we denote by , can be computed straightforwardly from Equation 7, and is equal to
where is the derivative of . Note that because is continuous, we are still able to compute this expectation, and similar ones to follow, when is undefined on a finite set of inputs (which is permitted under our global assumptions).
The derivative of local C maps with respect to their c value argument has an especially nice form which we make use of later in Section 11.
We begin by defining the following notation:
This function is closely related to the local C map of (given by Equation 19) in the sense that . The derivative of with respect to , which we denote as , is given by
This elegant formula was stated in Poole et al. (2016), although no explicit derivation of it was given. For completeness we provide one in Appendix B.
An immediate consequence of this result is that the -th derivative of with respect to , which we denote by , is equal to
where denotes the -th derivative of . From this it follows that the -th derivative of w.r.t. can be written as
This formula is valid even when , where the -th derivative is defined as the function itself (i.e. ), as is standard convention. When isn’t defined on a measure zero set of points, the formula may still be valid, provided that is continuous.
For example, if is the RELU function, is continuous everywhere and has a derivative everywhere except at , so the formula is valid for . However, is not continuous at , and one can use Equation 27 to show that as , while the formula would wrongly predict a value of .
2 Derivatives of extended maps
Because extended maps can be expressed as compositions and weighted averages of local maps, their derivative computations can be performed straightforwardly using automatic differentiation. The resulting formulae will still depend on the derivatives of local maps, but these can be computed (or numerically approximated) as per the previous subsection.
In such a scheme, composition corresponds to multiplication, and weighted averages correspond to weighted averages (since differentiation is linear). Notably, because q values don’t depend on c values, the derivative of extended C maps with respect to their input c values can be computed as if all the q values in the network are constant (although they still depend on the network’s input in general). So for example, the C map derivative for a composition of many combined layers is just the product of the local C map derivatives for each layer, evaluated at the appropriate values of as per the forward evaluation.
Handling weighted sum operations
An operation commonly performed in neural network models is the (weighted) sum of two or more feature maps. For example, in the ResNet-V2 architecture (which is described in detail in Section 23.1), the input to a “residual block” is added to its output, using what is known as a “residual connection”. Since sum operations are not among those listed as allowed in Section 3.2, it would seem that our assumptions rule out such architectures.
However, for the purposes of our analysis, there is no requirement that a network be formally constructed the same way it would implemented in code or drawn in a diagram; it only matters that it can be constructed in a way that conforms to the assumptions outlined in Section 3.2. With this in mind, we will now describe a way that a certain restricted class of weighted sum operations can be simulated using only directly supported operations. The consequence of this is that our analysis will in fact apply to architectures that contain such sum operations.
Typically, the feature maps that are summed in neural networks are the outputs of a set affine layers that don’t share parameters. (This is true in ResNet-V2 architectures, for example.) In such situations, we can replace the sum with a single affine layer h\left(\left[\begin{array}[]{cccc}Z_{1}^{\top}&Z_{2}^{\top}&\cdots&Z_{n}^{\top}\end{array}\right]^{\top}\right), which is obtained by concatenating the filter banks, the bias vectors, and the input feature maps (i.e. the ’s) together along their respective channel dimensions. (If is followed by a nonlinear layer in the network, then one simply forms a new combined layer consisting of this and .)
While almost good enough, the issue with this construction is that the implied initial distribution of ’s filter bank parameters is not one of the ones described in Section 4, and in particular, the variance/scale is not correct. To account for this, we must renormalize by the new number of channels (after the concatenation), the effect of which is that will instead compute a weighted sum of the form
where is the input channel dimension for .
Fortunately, we can extend this construction to support a more general class of weighted sums (with weights ) by multiplying each by a scalar before concatenating them. Doing so gives
where we have used the fact that the affine ’s are in fact linear (given that the biases are initialized to 0).
If the layer is still in the network after this replacement is performed for some (e.g. because its output is used in more than one place), this creates parameter sharing between and . However, as long as the network with sum operations that we are trying to simulate doesn’t violate our parameter independence assumptions from Section 4, neither will our simulating network.
The existence of this construction thus implies that weighted sums between the outputs of two or more affine layers (and directly followed by an optional nonlinear layer) are supported within our framework, provided that said affine layers don’t share parameters. Note that the weighed sum operation can be performed directly in the model code, and the concatenation-based construction only needs to be referenced in the theoretical analysis.
To deal with sum operations in Q map computations, we observe that the q value of \left[\begin{array}[]{cccc}\alpha_{1}Z_{1}^{\top}&\alpha_{2}Z_{2}^{\top}&\cdots&\alpha_{n}Z_{n}^{\top}\end{array}\right]^{\top} is, according to Equation 20, equal to
where is the q value associated with , and we have used the fact that the q value for is . From this it follows that the output q value from the sum is also , since is just the identity function.
Given uniform q values, a similar derivation based on Equation 21 lets us compute the corresponding c value as
where is the c value associated with . Note that unlike the formula for the q value, this is always a weighted average of the ’s, regardless of the values of the ’s. And in the case where all input q values are equal, it simplifies to .
Uniform q values
In general, C maps are three dimensional functions that depend on an input c value and two associated q values. While simpler objects than a network’s PKF (or even the network-level PKF approximation), they are not yet simple enough for our purposes. In particular, the behavior of C maps depends strongly on the two input q values, and q values can vary significantly between different network inputs and/or feature map locations. Finding a single scheme that controls the behavior of the C map for all conceivable input q pairs is likely impossible in general, and so we look to restrict the possible q values through some sort of active intervention.
The one we propose in this section is a form of input data preprocessing, which ensures that all q values for a given layer are equal (across all possible locations in the feature map and inputs to the network). (Note that this condition does not require that q values be the same across different layers.) We will call this condition uniform q values.
In Poole et al. (2016) it was observed that local Q maps can have stable fixed points, and that if the network consists of a composition of many combined layers of the same type, then its q values will converge with depth to such a point. Thus, a reasonable approximation, especially for deeper layers, is to assume that this convergence has already taken place, and that the q values over the entire network are equal. (Note that this is strictly stronger condition that uniform q values.)
As discussed in Section 7, our setting is different from Poole et al.’s (2016) in that we consider more general architectures, and are interested in the precise behavior of a finite network architecture instead of its depth limiting behavior. Moreover, it may be a poor approximation to assume that q values are close to convergence in the earlier layers of the network, especially if there are no constraints placed on the initial q values (which are determined by the network’s input).
2 Uniform q values via Per-Location Normalization
Our solution to the problem of unpredictable q values is a type of input data pre-processing which we call Per-Location Normalization (PLN). This is related to the data normalization done in Daniely et al. (2016) for fully-connected networks, but generalized to convolutional networks. PLN ensures that each location vector in the network’s input feature map has a dimension-normalized squared norm of 1, or in other words, that the q values for the network’s input layer are all 1. Because subsequent q values are fully determined by previous q values via location-agnostic computations (i.e. Q maps), it thus follows by induction that each layer will have uniform q values under PLN.
PLN can be easily realized through a number of different possible transformations of the network’s input, although care must be taken not to destroy information. The naive approach of normalizing the vector at each location of the input feature map (and multiplying by the square root of the channel dimension) destroys information because the vector goes from having degrees of freedom to degrees of freedom (where is the number of channels). This can be seen most starkly when , in which case all location-wise “vectors” are reduced to scalar values.
Note that this approach to PLN still destroys some information, although it’s only one degree of freedom across , which includes all locations and channels. This can be seen most clearly in the case of only one location vector , in which case the formula becomes
from which we cannot recover the norm of . Thus, it only makes sense to use this form of PLN when there are a large number of locations and/or channels.
For cases where there is only one location (i.e. in a fully-connected network) and the channel dimension is small, one possible alternative is to use a data-independent constant value for the extra channel. Using a value of 1 gives
which doesn’t destroy any information about (since we can invert the last entry to get , and then multiply that by the other entries to recover ). The disadvantage of this approach is that the scale of could differ very significantly from , so that after normalization, the value of the extra channel could either dominate the overall vector, or be minuscule.
Performing PLN may not always be important in practice, as we demonstrate later in our ablation experiments. Moreover, the version we have proposed seems to slightly degrade optimization performance in our benchmarks, possibly because of the extra parameters it adds to the first layer (for the extra channel dimension), or because of the nonlinear warping it applies to the input space. On the other hand, our experiments also demonstrate that very badly scaled input data can sometimes cause DKS to perform poorly, unless PLN is applied as a corrective measure. (See Appendix N.6 for the relevant results.)
There are other ways we can produce normalized vectors without destroying information, such as those discussed in Daniely et al. (2016) for fully-connected networks. Of all of the aspects of DKS, our method of realizing PLN is the least explored, and we wouldn’t be surprised if there was a significantly better way of doing it.
3 Assuming uniform q values going forward
From this point forward we will assume that the uniform q value condition holds. This thus allows us to treat C maps as one dimensional functions, as the q value for each layer will be constant. As we will show in the next section, it also imbues C maps with a set of very useful properties which end up being crucial to our subsequent analysis of them in Section 13.
Additional consequences of uniform q values for C maps
For convenience, when we have uniform q values we will drop the formal dependence of the C map on its input q values, and instead treat these as known constants within the expression (which are both equal to the same value). This allows us to view C maps as essentially one dimensional functions, and we can use notation of the form “” for them going forward.
Under the assumption of uniform q values, several interesting and useful properties of local C maps emerge. Suppose is a combined layer (or an isolated nonlinear layer) with activation function . We begin by setting in Equation 23, which gives
2 Properties of positive definite functions
Positive definite functions have many interesting and useful properties which thus carry over to C maps. These include:
The set of positive definite functions is closed under differentiationClosedness under differentiation can be easily verified by observing that the derivative of with respect to is just , which is also positive definite since when ..
Positive definite functions are non-negative, non-decreasing, and convex on the non-negative part of their domain. (This follows from the fact that their derivatives are also positive definite functions, and thus non-negative for non-negative inputs.)
The set of positive definite functions is closed under composition and weighted averages with non-negative weightsClosedness under composition can be easily verified by substituting one series into the other, expanding, and observing that the coefficients of the resulting series are non-negative combinations of coefficients from the two original series. Similarly, closedness under weighting averaging follows by observing that the coefficients of the series for the weighted average are just weighted averages of the corresponding coefficients from the original two series..
3 Extended C maps are positive definite functions and map c values of 1 to 1
Given that local C maps are positive definite functions and map c values of 1 to 1, it’s easy to show that the same applies to extended C maps. First, the property that c values of 1 map to 1 is clearly preserved under composition and weighted averaging, and thus carries over to extended maps since they are constructed from local maps this way. Second, the property of C maps being positive definite functions also carries over, since positive definite functions are closed under composition and non-negative weighted averages as mentioned above.
Another way that one can show that extended C maps are positive definite is by observing that they describe the exact one-dimensional kernel function of a fully-connected network in the limit of infinite width (where we substitute convolutional layers in the original network with fully-connected layers). Because this kernel depends only on the inner product of its inputs via the function , it is thus invariant to orthogonal transformations of its input, and so by Schoenberg’s Theorem (Schoenberg, 1988) it is a positive definite function of this inner product (i.e. is positive definite). Note that this argument works even for non-smooth activation functions for which Equation 23 may not apply. See Daniely et al. (2016) for more details.
4 A complementary perspective based on “dual activations functions”
Assuming that we normalize each activation function so that its output q value is 1, these dual activation functions can be composed and averaged in order to form what are called “compositional kernels”, which are approximations of the kernel function for the entire network, and are analogous to our network-level PKF approximation in the case where there is only one location (i.e. the fully-connected case).
Given this connection, it may thus be an appealing prospect for us to adopt Daniely et al.’s (2016) framework instead of the one we’ve presented, as it’s very carefully laid out and rigorously developed, and comes packaged with the best known error bounds for initialization-time kernel approximations of neural networks (one of which we adapt in Section 5.8). However, while their framework can deal with local receptive fields, it cannot directly deal with the weight sharing used in convolutional layers, and it would likely require significant work to extend it in that direction.
Indeed, to deal with convolutional layers in a way that allows kernel approximations for individual layers to be naturally composed, one seemingly must define something like our APKFs which keep track of approximations to entire IPMs (which contain inner products between every pair of locations in the feature maps of both inputs). Moreover, without assuming a Delta initialization, a decomposition of the kernel approximations into 1 dimensional functions (such as Q/C maps) becomes impossible, since APKFs for standard initializations involve non-trivial interactions between all the locations in the feature map (as seen in Equation 4).
While we do indeed restrict our attention to Delta initializations in this work, without the PKF/APKF formalism we would not be able to extend our analysis to mean pooling layers, since the kernel approximation for such layers also involves interactions between locations. As we will see later in Section 23, this extension will be necessary later in order to understand how DKS can be applied to standard convolutional neural network architectures.
Part II Desirable Q/C map behavior and how to achieve it
C maps approximate a network’s PKF at initialization time. In this view, c values approximate the cosine similarity between pairs of vectors (corresponding to different locations/inputs), and their evolution via C maps thus describes how these cosine similarities evolve in the network. Given uniform q values, the norms of these vectors are approximately constant (for a given layer), and thus their relative distance is related to their cosine similarity via
A (sub)network’s C map thus provides a complete description of how it warps the geometry of its input space at initialization time.
As we will argue in this section, the preservation of some amount of this geometric information through the network, as indicated by a “well-behaved” C map, is a necessary condition for the network to be trainable. When C maps “degenerate” in certain ways, as we will show they do for standard deep neural networks, it means that the relative distances between the network’s (location-wise) input vectors are hard to infer from the network’s outputs, making gradient-based training difficult.
The local C map of a combined layer with a RELU activation function is given by
This formula is stated in Daniely et al. (2016), and is based on a derivation by Cho and Saul (2009), where it corresponds to a normalized version of the “1st-order arc-cosine kernel function”. Note that while in general C maps depend on the input q value, this formula is valid for any q value, which is a consequence of the fact that RELUs are positively homogeneous (i.e. for all ).
One interesting fact about is that , which can be verified by taking the derivative of Equation 27 and letting . Moreover, because (which is true for general C maps), we have that a deep RELU network consisting of the composition of combined layers will also have the property that .
From this we can see that the entire domain , which makes intuitive sense since the RELU function is non-negative. Other than that, resembles a slightly shifted and rescaled identity function, and so is reasonably well-behaved.
However, if we build a deep network as the composition of many RELU combined layers, compression of the C map’s output becomes much more extreme, with outputs rapidly concentrating around 1 as depth increases. This can be seen below in the plot of the C map for RELU networks of depths 5, 20 and 100 (which we obtain by iterating Equation 27 the required number of times):
Here, the C map for depth 100 has the property that maps the entire interval $$ to [0.996, 1], which represents an extreme amount of compression.
2 Sigmoidal networks (using the erferf\operatorname{erf} activation)
Another example of an activation function whose associated local Q and C maps have closed-form expressions is the classical “error function”, which is given by . This function has a “sigmoidal shape”, which makes it a reasonable stand-in for the more common sigmoidal activation functions like and the logistic sigmoid.
The local Q map for a combined layer with an activation function is given by
These formulas follow from equation 11 of Williams (1997).
Unlike in the RELU case, the local C map depends on the input q value, and so to plot it we must make an assumption about this value. One natural choice is , which gives the following plot:
Visually, this function is almost indistinguishable from the identity function.
To compute the C maps for deeper RELU networks we need to track both the q and c values through each layer, using the previously stated equations. Doing so for depths 50, 150, and 500 gives the following plot:
From this plot we can see that at high depths, the network’s C map has a tendency to compress nearly all input c values to a small region around 0. Moreover, this behavior only becomes more extreme as the depth increases.
3 C map degeneration in more general deep nonlinear networks
The following proposition establishes that the C map degeneration we observed above for deep RELU and tanh networks happens for a larger class of deep networks. Moreover, the point towards which (nearly) all c values get mapped as the depth increases is unique.
Suppose is a deep network consisting of a composition of subnetworks, each with the same C map . Then for all we have
The proof of this proposition is a straightforward generalizationWhile the statement of their claim assumes that each of the subnetworks is a combined layer, the only fact they use about in their proof is that it is positive definite, which holds for more general subnetworks by Section 11.3. of the proof of “Claim 1” from Daniely et al. (2016).
While Proposition 10 describes the convergence of in the limit of infinite depth, it is still informative about at finite depths (which is the setting we actually care about). In particular, it essentially says that for any there is a constant so that for when , nearly all input c values get compressed to a region of radius around . We will call C maps exhibiting this compressive behavior (with a small ) degenerate.
Note that Proposition 10 assumes that is the same for all values of . If, for example, we were to modify the network’s activation functions based on the value of (which DKS will do), the convergence seen in the proposition may not occur.
Also note that the hypothesis in Proposition 10 that each subnetwork has the same C map will rule out many common cases. For example, it is violated for the deep network we looked at before, because the q values are different for each layer (which leads to different local C maps). However, because the q values converge rapidly to a fixed point in such networks, convergence of the c values will still occur. As shown in Appendix J, the “residual blocks” of ResNets (which are repeated many times in sequence) also violate this hypothesis, but their q values do not converge to a fixed point.
In Appendix C (Theorems 40, 41, and 42) we give a much more detailed analysis of the convergence of in terms of the properties of and the location of in $\mathcal{C}^{\prime}(1)\neq 1D$) with precise rates, thus establishing that degeneration can happen very quickly in deep networks. In contrast to the related analyses of Poole et al. (2016), our results apply pre-asymptotically.
4 Types of degeneration and their implications for trainability
As we’ve seen above, degenerate C maps send a large range of input c values to a small (and sometimes point-like) region near some fixed point . This means that the original geometric relationships between the corresponding input vectors are obscured by the action of the network, becoming essentially impossible to recover from its outputs. While it seems intuitively plausible that this would make gradient-based optimization of such networks difficult (as has been argued by Schoenholz et al. (2017)), it’s worth examining the situation in more detail. In this subsection we will give a detailed intuitive argument. A more rigorous argument which confirms these intuitions will appear later in Section 24.
Suppose is some subnetwork of the overall network that we wish to train. There are two basic cases, corresponding to different possible values for .
If sends nearly all input c values to a small neighborhood near , this means that regardless of the original distance of the two associated input vectors, their corresponding output vectors under will be nearly identical (i.e. have a relative distance ). And because this holds for all pairs of vectors, it means that is nearly constant, with only a very weak dependence on its input. This has a different set of consequences for layers in the network before versus layers after.
For layers before there are two cases. If ’s Jacobian is non-negligible for most inputs, then it will have to vary wildly over ’s input space, as this is the only that a function can achieve a nearly constant output while having a non-negligible Jacobian. (Balduzzi et al. (2017) observed a similar phenomenon for early layers in deep RELU networks, likening the gradient function to a “random noise process”.) This will make learning difficult, or at the very least unlikely to generalize, as similar pairs of training cases will produce very different gradients. If on the other hand ’s Jacobian is negligible, this means that gradient magnitudes for layers before will be very small compared to those for other layers. This makes simultaneous optimization of the network’s layers with gradient descent very difficult, and even sophisticated 2nd-order methods may struggle in the more extreme cases. (This is arguably related to the well-known “vanishing gradients” phenomenon identified in Hochreiter et al. (2001).)
Meanwhile, layers after won’t be able to learn anything more than a constant output prediction, as their inputs will be nearly constant. And even if the output produced by has enough variance across the training data to overcome the limits of numerical precision, the part of the network after would need to have a very large Lipschitz constant in order to produce well-separated outputs for different training cases.
If sends nearly all input c values to a small neighborhood around with , then any two input vectors (that aren’t either almost identical or negations of each other) will be mapped by to output vectors that are nearly a constant relative distance apart.
Given this condition, and that is differentiable, it follows that ’s Jacobian must be very large in certain regions of the input space (and possibly everywhere). This means that the gradient magnitudes for layers in the network before will be much larger than those for other layers, making it difficult to optimize them simultaneously. (This is arguably related to the well-known “exploding gradient” phenomenon discussed in Hochreiter et al. (2001).)
A network containing such a subnetwork possibly stands a greater chance of being trainable than in the previous “collapsing case”, since ’s output vectors will still be distinguishable for different inputs vectors. Optimization of the layers after (or towards the end of ) could even conceivably learn a map from these vectors to their associated targets in the training set. However, it is highly unlikely that the resulting model would generalize well, as the similarity between two such output vectors would have no discernible relationship to the similarity of the associated two input vectors.
5 Are well-behaved C maps sufficient?
While a well-behaved (i.e. non-degenerate) C map seems like a necessary condition for trainability (as argued above), it should be noted that without additional hypotheses, no condition on the C map can be a sufficient. For example, because the C map is invariant to the network’s parameterizationThis can be seen by observing that APKFs, from which local Q/C maps and ultimately extended Q/C maps are defined, only depend on the functional behavior of combined layers, and not which variables are formally considered “parameters” from the standpoint of the optimizer. Indeed, one could reparameterize a layer’s weights using any invertible function without changing what it computes at initialization time, or the form of its PKF/APKF., or even whether its parameters are considered trainable at all, it cannot completely predict the performance of a gradient-based optimizer. Even if we assume the standard network parameterization, and a model class which is equivalent to a standard deep nonlinear network, interesting counterexamples to sufficiency still exist, as we will show in Section 14.2.
One way to incorporate hypotheses about the network’s parameterization, and the optimizer used, is to analyze gradient descent training from the perspective of Neural Tangent Kernel (NTK) theory (Jacot et al., 2018). Later in Section 24 we will argue that for a deep fully-connected network, a degenerate C map leads to a form for the NTK which implies very poor generalization and/or slow optimization under gradient descent. Conversely, we will also show that a network constructed using DKS has an NTK which is suggestive of good generalization and fast optimization (although doesn’t necessarily guarantee it). This NTK-based analysis can be viewed as a rigorous version of the intuitive argument given in the previous subsection.
Mathematical analysis of C maps
As we saw in Section 12, C maps can become degenerate in deep networks by mapping nearly all input c values to a point-like region around some , which leads to difficulty when training with standard optimization methods. In this section we will analyze C maps in closer detail, and show how their overall deviation from the identity function (which serves as a measure of degeneration) can be predicted from their slopes at and/or . (We will ultimately design DKS to control these slopes in order to prevent degeneration.) We will also establish connections between the properties of a local C map and its associated activation function, and characterize the slope behavior of degenerate C maps over $$.
Suppose is some subnetwork. The question of how to measure the deviation of from the identity function, which we will use as a measure of “degeneracy”, is an interesting one. Since we ultimately want to forbid extreme behavior of for all input c values, it makes sense to look at the worst-case ones. This suggests the following two options, which compare to the identity function using either its values or its derivatives:
The first of these, while a reasonable choice, could fail to detect small ranges of c values where geometric information is lost due to the slope of getting close to zero. The second option will detect such regions, but unlike the first measure, is insensitive to being shifted by an additive constant, and is only weakly sensitive to it being “angled” away from the identity function. Fortunately, we can avoid having to choose between the two, since as we will see next, they can be simultaneously bounded using a few easily-computed properties of .
2 Bounding deviation from the identity
The following result relates the deviation of from the identity function to the values of , , and . Its proof, which makes strong use of the fact that is a positive definite function, is given in Appendix D.1.
If , then we additionally have that
From this result we see that the first measure of deviation is within a factor 4 of (and also upper bounded by when ), and the second measure is within a factor 3 of when . This suggests that we can control the deviation of from the identity function by simply controlling the distance of and/or from 1.
Note that a simple consequence of this result is that for subnetworks satisfying . It is also true more generally that (as shown in Appendix D.1).
The lower bounds in Theorem 13 do not imply that will look globally nonlinear when is large (in contrast to the theorem’s upper bounds which do imply it will look globally linear when is small). For example, the function is a valid C map and has a derivative at , and yet is very close to linear everywhere except near and when is large.
3 The relationship between identity C maps and linear activations
The local C map for a nonlinear layer with a linear activation function is the identity (which can be easily verified from Equation 19 by taking with ). However, from this observation it’s not immediately obvious that a local C map will converge to the identity function as its associated activation function becomes “more linear”, or vice versa. In this subsection we will show that this is indeed the case for certain carefully chosen measures on both function spaces, and we will give the rate of this convergence.
is arguably the most natural one to use, as it employs a weighting over that precisely reflects the input distribution we would expect given an input q value of 1. This measure is also closely related to Q/C maps in the sense that (with defined as in Equation 22), where is a nonlinear/combined layer with as its activation function.
is linear insofar as it’s close to a multiple of the identity function . With this in mind, and given the fact that , we will measure the level of nonlinearity of according to
where we normalize by to keep invariant to changes in the overall scale of . Note that with this definition, is perfectly linear (i.e. a multiple of ) if and only if .
It turns out that we can relate to properties of , as is established in the following proposition whose proof is given in Appendix D.2:
Suppose is a nonlinear/combined layer with activation function and input q value 1. Then we have
This result shows a strong relationship between the level of linearity of , and the distance between and the identity function (as measured by the infinity norm). Moreover, one converges to 0 (as we vary ) if and only if the other one does.
Proposition 16 can be straightforwardly extended to the case of general input values by modifying the definition of the inner product used for .
4 The relationship between affine C maps and affine activations
An affine function is, by definition, a linear function plus a constant term. Equivalently, it is a function whose derivative is constant (i.e. is a multiple of the constant function ). From this second characterization, we can measure the “non-affineness” of by
Note that with this definition, is perfectly affine if and only if .
It turns out that we can relate to properties of , as is established in the following proposition (whose proof is given in Appendix D.3).
Suppose is a nonlinear/combined layer with as its activation function. Then
From the above expression we can see that becomes more affine as the ratio approaches 1. Moreover, approaches an affine function if and only if itself does, as is a measure of how affine is. (To see this, note that for all since is convex on $C_{f}^{\prime}C^{\prime}_{f}(0)/C^{\prime}_{f}(1)\rightarrow 1C_{f}$ is analytic.)
5 Slope properties of degenerate C maps
In subsection 13.2 we saw how certain conditions on the slope of a C map at and/or ensure that it is well behaved (i.e. not degenerate). In this subsection we will establish the converse: that degenerate C maps necessarily have extreme values for these slopes.
As shown in Section 12, C maps in deep networks can become degenerate in the sense that they map almost their entire input domain (except points very close to ) to a small region around some limiting c value . One way to quantify this behavior is to look at how “flat” the function is up to some c value s.t. , which we can measure using
for . When is degenerate, or in other words very flat, will be very small (for values of not too close to ).
While the interpretation of is clear for , it is less clear for . In the following proposition (proved in Appendix D.4) we show that does indeed work as a measure of flatness for values of less than .
Suppose is a subnetwork, and with . Then
Note that because is convex and non-decreasing on $C_{f}(1)=1F_{f}(c^{\prime})\leqslant F_{f}(c)\leqslant 1c,c^{\prime}|c^{\prime}|\leqslant|c|F_{f}(c)[-c,c]c$.
A more “analytic” way to measure flatness is to look , which intuitively should be small in flat regions of . It turns out that this intuition is basically correct, as we establish in the following proposition (whose proof is given in Appendix D.5).
Suppose is a subnetwork, and with . Then
This bound establishes that will be small whenever is (which is to say, when is degenerate), provided that is not too close to 1.
While Proposition 20 doesn’t address the value of directly, we can still bound it by applying Proposition 20 with for some and then use the fact that for any (which is true because is non-decreasing and convex on $$ by Section 11.2).
In general, we cannot say that much about or when is degenerate. However, the following two propositions (proved in Appendices D.6 and D.7) give us some basic information about these values in certain special cases.
Suppose is a composition of subnetworks each having the C map , and that . Then we have that with , and either or .
Note that the claim made in Proposition 22 does not hold for more general types of networks. For example, if we have a sequence of networks such that , then for all we have , and as .
For RELU combined layers we have (which follows from Equation 27 by taking the derivative and letting ), and thus the bound in Proposition 22 is tight.
Suppose is a subnetwork. For all we have
By taking a small value for , Proposition 25 tells us that for degenerate C maps with (so that ), will be large provided that the flatness measure is small. See Section 12.2 for an example of degenerate C map where is indeed very large.
C map behavior in linear networks and the problem of being “too linear”
In Section 12 we saw that deep nonlinear networks can easily have degenerate C maps, which makes them very hard to train with gradient-based methods. One might wonder if this pathology is reserved to nonlinear networks, or if deep linear networksHere, a linear network is defined as one whose activation functions are a constant multiple of the identity function. also suffer from it.
Given our assumption that the initial biases are zero, it turns out that the answer is no. The local C map for a combined layer with a non-zero linear activation function is equal to the identity. Because identity functions are preserved under composition and weighted averages, it thus follows that the extended C map for any subnetwork is also the identity function, and is therefore well-behaved.
Does this mean that linear networks are easy to train? Well, as discussed in Section 12.5, more hypotheses are required to say anything about trainability. But if one adopts the standard parameterization, very deep linear networks are surprisingly easy to train both in theory and practice using standard techniques (Saxe et al., 2014), provided that they are initialized using orthogonal weights. Linear networks thus represent an interesting example of where our necessary condition for trainability (i.e. having a well-behaved C map) also appears to be sufficient.
Despite how easy they are to train, we obviously can’t use linear networks in practice, as their expressivity is fundamentally limited. But these observations do suggest a possible strategy to address problem of degenerate C maps in nonlinear networks: we can transform the network’s activation functions so that they appear “sufficiently linear” at initialization time. However, as we will see in the next subsections, going overboard on this idea will lead to a special type of untrainability that exists only in networks with very well-behaved C maps.
2 The problem of being “too linear”
As suggested in the previous subsection, one way to achieve a well-behaved C map would be to transform the activation functions in a network to resemble the identity function (or a multiple thereof). We can do this for the RELU activation function (defined by ) by adding a large constant to its input and subtracting the same constant from its output. In other words, we set
which is equivalent to the identity function for all inputs . If is extremely large, say , this means that all practically sized inputs will satisfy this constraint, and thus the function can be treated as the identity for all practical purposes. Moreover, the expectation formulas for local Q and C maps given in Section 6 will produce practically identical results to the identity function case, since the probability mass associated with inputs to will be vanishingly small. (See Section 13.3 for a formal justification of this.) Thus, the C map for networks consisting of the composition of many combined layers with these transformed RELUs will be equal to the identity function up to a vanishingly small approximation error.
Meanwhile, because nonlinear layers are always preceded and followed by linear layers with learnable biases, the network can in principle learn to undo these transformations and thus simulate a standard deep RELU network. The model class is thus technically no different from a standard RELU network, assuming a perfect optimizer. But despite this, these transformed networks will never actually learn nonlinear behavior via standard gradient-based methods in a reasonable amount of time, and so their hypothetical expressive power will fail to be properly utilized. Indeed, unless the optimizer manages to change the parameters by a factor on the order of , the network will behave nearly identically to the corresponding linear network (which computes only affine functions of its input) throughout the entire course of optimization, both in terms of its function values and its gradient/curvature estimates.
The basic problem here is that the transformed network has become “too linear” in the sense that we require a very large change in its parameters to see any significant nonlinear behavior. While such a network may be readily trainable within the class of linear functions (as linear networks are), it will be severely limited compared to a standard RELU network in terms of its effective expressive power under gradient descent optimization. Thus, to achieve trainable networks that are also expressive, one must avoid this failure mode in addition to requiring a well-behaved C map.
3 How to avoid networks that are “too linear”
If a subnetwork has an C map which is very close to the identity function, this will usually mean the local C maps of its nonlinear layers are also very close to the identity function (and perhaps much more so). By Section 13.3, this implies that the activation functions must therefore be very close to linear, so that the network is at risk of being “too linear” (as defined above). We may thus hope to prevent this by insisting that the network’s C map isn’t too close to the identity function. However, this alone won’t be good enough, as shown in the following example.
Consider modifying the transformed RELU activations in the previous example by adding 1 to their output and dividing the result by , so that they compute over their high-probability range of inputs (which is an affine function of ). With this change, gradient-based learning will still be effectively restricted to the class of linear networks (which compute affine functions). Meanwhile, a straightforward calculation via Equations 7 and 19 shows that and for nonlinear/combined layers with activation function , and so differs significantly from the identity.
Fortunately, by leveraging the previous analysis, there is a simple way we can use C map properties to provably avoid networks that are “too linear”. From Section 13.4, the degree of “non-affineness” of , denoted , is given by
As we have by Remark 14, it thus follows that
This shows that we can avoid activation functions that are too affine (which is sufficient to avoid networks that are “too linear”) by requiring that be sufficiently greater than 1 for every nonlinear layer .
Mitigating kernel approximation error
Our analysis of the initialization-time behavior of deep neural networks via Q/C maps relies on the assumption that the APKF approximation, when applied over multiple layers in a nested fashion, is a reasonable one to make. If this isn’t true, then Q/C maps will cease to describe the network’s PKF at initialization time, and our attempts to make the network trainable by controlling their properties will be doomed to failure.
In Section 5.8 we discussed the error bounds from Daniely et al. (2016) in order to help justify our use of nested APKF approximations in deep networks. These bounds make high-probability statements about the error of initialization-time kernel approximations of neural networks, and give a maximum value which shrinks with the square root of the width and grows exponentially with depth. While they represent the best rigorous account of neural network kernel approximations, they are still too pessimistic to be useful in practical settings, either for predicting the approximation error or controlling it.
In this section we will propose a heuristic way of looking at how approximation error originates and evolves across multiple layers which we have found to be quite predictive in practice, and which implies certain error mitigation strategies that we can incorporate into DKS.
Suppose is a subnetwork consisting of a composition of many combined layers. A perturbation to the input q value to , representing the error from approximations made at previous layers, will manifest as a perturbation of ’s output. Up to first order, the size of the latter will be approximated by that of the former, multiplied by ’s derivative.
As discussed in Section 8.2, the derivative of is equal to the product of the derivatives for its constituent local Q maps (i.e. those for each of ’s combined layers), and thus will grow or shrink in an exponential fashion as a function of the depth. Thus, it can easily be the case that deep networks will have very large Q map derivatives, which suggests a very large amplification of error through successive layers network. One way we can avoid this issue is by requiring that derivative of each local Q map, when evaluated at its expected input, is less than or equal to 1.
A closely related perspective, which applies to networks consisting of a sequence of combined layers each with local Q map given by , is that a fixed point of will be attractive if . Attractive fixed points have the property that will converge to as for all values of sufficiently close to , and are thus naturally “robust” to reasonably sized errors in .
See Appendix K for empirical evidence of the relationship between Q map derivatives and kernel approximation error.
2 Minimizing errors using large width and SUO-distributed weights
In addition to controlling the propagation of errors across layers, another way to mitigate error is to increase quality of each of the layer-wise APKF approximations from which the errors first originate. In the case of Gaussian-distributed weights, APKFs use analytic expectations to approximate finite averages (over unit outputs), where each element being averaged is an iid unbiased estimator of the expectation. Increasing the width/channel dimension will thus reduce the variance of the overall estimator, and thus reduce error (as predicted by Theorems 1 and 6).
While not originally conceptualized as such, the use of SUO-distributed weights (as defined in Section 4) provides another way to mitigate the kernel approximation error that is essentially free. When using the SUO distribution, the weights are no longer statistically independent or Gaussian distributed, and so the unit outputs being averaged across are neither iid nor unbiased estimators of the kernel approximation formula. Nonetheless, their average is a consistent (but biased) estimator, whose variance goes to zero as increases (as is established in Theorem 6).
It is well-known that biased estimators can sometimes have lower variance than unbiased ones, and this does seem to be the case here. Recall the discussion from the end of Section 4.2, where it was observed that the distribution of an output vector from a linear layer is identical for the Gaussian and SUO-distributed cases, except that the former introduces a random multiplicative perturbation (with mean 1 and variance ) on the vector’s dimension-normalized squared length (which is an estimator of the associated q value). While this perturbation is required for the implied estimator (after the nonlinearity) to be unbiased, it leads to additional variance. We conjecture that this extra variance is more significant than the bias, and thus SUO-distributed weights yield an overall lower approximation error for APKFs.
Note that the upper bounds given in Theorems 1 and Theorem 6 do not reflect these intuitions, as they suggest an overall lower approximation error for Gaussian-distributed weights. However, because these are only upper bounds without matching lower bounds, and are likely quite loose/pessimistic, one cannot draw any conclusions. Indeed, we conjecture that tighter bounds could be obtained for SUO-distributed weights given a more careful analysis than the one in Martens (2021).
Part III Specification and derivation of Deep Kernel Shaping
Having established a detailed understanding of Q and C maps, and how their properties relate to network trainability, we are now in a position to state and justify the specific conditions which we will attempt to enforce with DKS. Our particular mechanism for doing this, which involves certain transformations of the network’s activation functions, will be described in later sections, and isn’t important for the present discussion.
In the follow series of subsections we describe each of the conditions. Note that because we want the entire network’s capacity to be utilized, and not just the subnetwork corresponding to the most direct input-output path, we will enforce these conditions for all subnetworks of the network.
The uniform q condition ensures that the q values for a given layer are independent of location (in the feature map) and the network’s input. While this alone is a sufficient hypothesis to derive an approach similar to DKS, we will go a step further and standardize to a q value of 1, which will allow us to reuse local Q/C map computations across the entire network. Note that given our use of PLN (which ensures initial q values are 1), this is equivalent to the condition that for all subnetworks .
The choice to standardize to a q value of (as opposed to some other positive constant) is somewhat arbitrary and not particularly important. 2 would have worked equally well, for example. However, the choice of 1 does lead to somewhat simpler expressions for the local Q/C map and their derivatives, and corresponds to an output scale which is in the range of “interesting behavior” for most typical loss functions (such as the commonly used softmax cross-entropy error).
As discussed in Section 15, the size of the error in our kernel approximations can be roughly predicted from the size of the derivatives of the Q maps. Thus, in order for Q/C maps to be an accurate description of the network’s true kernel function, we must keep the size of these derivatives under control. To this end, we will require that for all subnetworks .
We look at the derivative at in particular since this is the input q value we expect in the absence of error, due to the condition . In principle, we also care about the value of for ’s close to 1, which would give us a more complete picture of the approximation error (and perhaps let us to establish rigorous bounds on it). Unfortunately, we don’t yet have a powerful theory for the global properties Q maps like we do for C maps, and so the best we can do is to look at their properties at particular points. That being said, because we know that Q maps are smooth, it’s likely that will be reasonably well approximated by for values of close to 1.
The choice to make equal 1, which corresponds to the error neither growing nor shrinking, is somewhat arbitrary, and other choices for this value are possible. For example, we could minimize instead of setting it to 1, thus suppressing the growth of errors as much as possible. Minimizing did seem to work in our experiments, however we found that for certain activation functions (such as ) it resulted in slower optimization compared to using . (See Appendix N.4 for these experiments.)
We are currently not sure why setting works better than minimizing it for some activation functions. One possible explanation is that allows to transmit information about the overall scale of its input vector as a roughly linear function of . Meanwhile, networks where is minimized will tend to “squash” the range around , making it harder to recover the original input value from the network’s output. One can perhaps draw a rough analogy between this and the preservation of “geometric information” by C maps as discussed in Section 12.
In order to apply the analysis of Section 13 we require that for all subnetworks . While this might seem like an overly stringent requirement, it is worth noting that arbitrary deviation of from the identity function is possible if we don’t place any restrictions on the value of . This is the case even when is enforced, as can be seen in Section 12.1 for deep RELU networks.
As discussed in Section 12, degenerate C maps correspond to networks that are difficult to train with gradient-based methods. Avoiding this degeneration is the central aim of DKS. As argued in Section 13, we can do this for a given by bounding its maximum deviation from identity function (which is the canonical non-degenerate C map).
Given the condition , Theorem 13 says that this deviation is roughly equal to . Thus, we will enforce the condition for all subnetworks , where is a hyper-parameter which we will sometimes refer to as the global slope bound. (Note that is true automatically as consequence of Theorem 13.) This condition can be thought of as imposing a limit on how “non-linear” any given subnetwork is allowed to look.
Even assuming that our kernel approximations are exact, a well-behaved C map is not a sufficient condition for a nonlinear network to be trainable. As discussed in Section 14.2, one way such a network can fail to be trainable is if it’s too far away in parameter-space from a significantly nonlinear function, or in other words is “too linear”. In such cases, a gradient-based optimizer will struggle to utilize the full expressive power of the network.
For neural networks with standard parameterizations, this issue will manifest as nonlinear layers with activation functions that behave too much like affine functions. As argued in Section 14.3, this can be avoided by requiring that be sufficiently larger than 1 for such layers . Thus, it makes sense to balance the condition in Subsection 16.4 with one requiring that is maximized, where the minimum is taken over all nonlinear layers in the network.
6 Choosing the global slope bound ζ𝜁\zeta
Given the above two conditions, the global slope bound corresponds to the maximum value of over all subnetworks (and must be ). Heuristically, the degree to which is greater than 1 tells us how nonlinear the network’s functional mapping is at initialization time. If is too large, then the C map for the network (or one of its subnetworks) will experience the “exploding” type of degeneration discussed in Section 12.4.2, where c values are squashed towards . If it’s too close to 1, then the C map will be very close to the identity, and we run the risk of making the network “too linear” (as per Section 14.2).
In our experiments on 100 layer networks we tried only a few values of before settling on , and in general we found that DKS is reasonably robust to significant variations in (or more precisely, ). For example, and both produced depth 100 networks that trained at competitive speeds, being only somewhat outperformed by networks that used . More extreme choices like and meanwhile produced significantly slower training. See Appendix M.3 for these results. For depths 200 or greater we found that it was sometimes necessary to use a value of less than 1.5 (such as 1.1) to achieve stable training. We speculate that this is because the kernel approximations underlying our Q/C maps tend to break down at very high depths (given our modest layer widths), but that this can be mitigated by making the network “more linear”.
An interesting systematic trend we observed is that larger values (up to a certain limit) tended to produce slightly faster optimization, whereas smaller ones led to slightly improved generalization, possibly because this made the inductive bias of the model (plus optimizer) favor a more linear solution. Relevant experimental data is presented in Appendix M.4.
From global map conditions to local ones
In this section we will describe how the “global” map conditions given in the previous section can be achieved by enforcing an equivalent set of conditions on the local Q/C maps of the network. The particular mechanism we will use to enforce these “local” conditions will be discussed later in Section 18.
Before we can write down the local map conditions we will define a special construction that allows us to relate the slope at 1 of extended C maps to the slope at 1 of local C maps.
Let be an arbitrary subnetwork. As discussed in Section 8.2, we may apply automatic differentiation to compute from the derivatives of the local C maps of ’s constituent layers, where composition corresponds to multiplication and weighted averages (due to concatenations or sum operations) correspond to weighted averages (with the same weights). Since c values of 1 always map to 1 (as argued in Section 11), the expression for the derivative will be a polynomial function of the local C map derivatives at .
If we further assume that there is a constant such that for each nonlinear layer in , then we can express as a polynomial function of , as the local maps for all other layers are the identity function. We will call this function the slope polynomial of and denote it by . Note that since whenever (which we are enforcing), we may thus assume without loss of generality.
Because slope polynomials can be constructed from products and weighted averages of lower degree slope polynomials, and the value of 1 is preserved under multiplication and weighted averages, it follows that for any subnetwork . And since and are positive definition functions of (trivially), and positive definite functions are closed under multiplication and non-negative weighted averaging (as discussed in Section 11.2), it also follows that slope polynomials are positive definition functions, just like C maps. They are thus non-decreasing for , and indeed strictly increasing provided that the subnetwork contains a nonlinear layer. From this it also follows that for all .
As we will be interested in computing the most extreme slope over a network, we will define a related function called the maximal slope function, which is given by , where the maximum is taken over all subnetworks of the entire network. Because subnetworks with no nonlinear layers won’t influence the maximum, and the maximum of a set of strictly increasing functions is strictly increasing, we have that the maximal slope function is strictly increasing provided that the network contains at least one nonlinear layer. And because it is the maximum over a set of continuous functions, the maximal slope functions is also continuous, and therefore invertible, which is a fact we will make use of later.
2 Computing maximal slope functions
The number of distinct subnetworks in a network can be very large, and so computing the maximal slope function naively from the definition can be laborious. Fortunately, we can eliminate most of these subnetworks from consideration immediately.
Observe that if a subnetwork is formed by feeding the output of one subnetwork into the input of another, i.e. , then we have by the chain rule. And because and for all , it thus follows that and for . Therefore, any subnetwork that is part of another subnetwork in this particular sense can be ignored when computing the maximum. Moreover, without assuming any relationship between , , and , if is a factor of , then is also a valid slope polynomial, and therefore for all , thus allowing us to ignore in the maximum.
Note this does not therefore imply that the maximal slope function is always the slope polynomial of the entire networkFor the entire network to even have a slope polynomial requires that it be a valid subnetwork of itself, which is only the case for networks with a singular input and output., as not every subnetwork can be related to the entire network in this way. For example, if we have a very deep nonlinear network with nonlinear layers and a skip connection from the initial input to the final output, so that the final output is times the initial input plus times the output of the nonlinear subnetwork, then the slope polynomial for the nonlinear subnetwork is , while the slope polynomial for the entire network is
which is strictly smaller than for all . (This formula can be derived by following the recipe given in Section 21.3.) For this network, the maximal slope function is in fact .
An even more interesting example is the same network, but with additional nonlinear layer added to the end, after the skip connection. The maximal slope function of this network is , which cannot be reduced to a polynomial as there are settings of for which either input to the is larger.
3 The equivalent local map conditions
Having defined the maximal slope function, we are now in a position to derive the equivalent local map conditions to the global ones given in Section 16.
First, observe that local Q/C maps for affine layers are identity functions, and can essentially be ignored when computing extended Q/C maps. What remains are nonlinear layers and weighted sum operations, and so we will concentrate on these.
If we have that for all nonlinear layers , then the analogous property automatically holds for all subnetworks that don’t contain weighted sums or constant scalar multiplications, as it is clearly preserved under composition and weighted averages (arising due to concatenations). The same reasoning also applies to the condition .
While weighted sum operations are constructed from concatenation operations (which are accounted for in the above argument), the construction also introduces scalar multipliers which can affect the q values. To account for this, we must ensure that the output q value of each weighted sum operation is 1. By Equation 24, this is equivalent to requiring that the squares of the weights sum to 1, assuming that the inputs to the sum have q values of 1. We will call weighted sums satisfying this condition “normalized sums”. Given that a weighted sum is normalized, and that its input q values are 1, it additionally follows (by Equations 24 and 25) that the corresponding output q and c values will be weighted averages of the input q and c values, with weights given by the squares of the weights of the sum itself.
To finally achieve for all subnetworks we must remove any constant scalar multiplication operations from the network, except those that are part of the above normalized sums. With this done, a simple inductive argument then establishes that for all subnetworks .
Given constant q values of 1, and weighted sums that are all normalized, we may compute using the same slope polynomials used to compute , provided that is the same for all nonlinear layers . Thus if we impose the condition for all nonlinear layers , it will follow that for all subnetworks .
Finally, maximizing , while requiring that for all subnetworks , is equivalent to setting for all nonlinear layers (since a single nonlinear layer is a subnetwork), where and is the inverse of the maximal slope function for the network (which exists as long as the network has at least one nonlinear layer).
Summarizing, the equivalent local map conditions are:
for all nonlinear layers , with the additional requirement that all weighted sum operations in the network are normalized (i.e. that the squares of their weights sum to 1).
Activation function transformations
In addition to PLN and the use of Delta initializations, our main mechanism of control over the initialization-time behavior of neural networks will be to apply transformations to their activation functions. In particular, we will apply constant scalar multiplications and shifts to both their inputs and outputs. For most typical activation functions this will give us sufficient control over a combined/nonlinear layer’s local Q/C maps to enforce the “equivalent local map conditions” from the previous section.
Suppose is some element-wise activation function in the network. We propose to make the following replacement:
where , , , and are static scalar constants (that we do not train). Note that these constants are the same for all channels and feature map locations within a given layer, but can differ between layers.
2 Equivalent parameters and preservation of the model class
Provided that each nonlinear layer is both preceded by and followed by an affine layer (which is true for most architectures), this way of transforming the activation functions has the property that it preserves the model class of the original network. By this we mean that for any network with transformed activation functions, there exists an equivalent network with untransformed activations that has precisely the same functional behavior. We will call the filter weights and biases of this second network the equivalent parameters.
3 Our method for transforming activation functions viewed as an initialization scheme
The existence of equivalent parameters, and their relatively straightforward computation, makes it possible to turn our method for transforming activation functions into an initialization scheme for the network’s parameters. One simply computes the constants needed to appropriately transform the activation functions, and then uses them to instead compute the equivalent parameters, starting from a network initialized as per Section 4. If we view this process as a sampling procedure for the network’s parameters, then it corresponds to a distribution with non-trivial correlations between the weights and biases of each affine layer.
Note that while a transformed network and an untransformed network (with equivalent parameters) compute the same function, they correspond to different parameterizations of the same model class, and thus may give rise to different optimization dynamics. Stochastic gradient descent for example, is not invariant to reparameterizations of this type, and so we would expect it to behave differently on either network. The K-FAC optimizer (Martens and Grosse, 2015) on the other hand is approximately invariant reparameterizations involving affine transformations of layer inputs and outputs (Martens and Grosse, 2015; Grosse and Martens, 2016; Luk and Grosse, 2018).
Our experimental results indicate that the transformed networks are easier to optimize with stochastic gradient descent than networks with equivalent parameters. Meanwhile, as predicted by the theory, the optimization performance with K-FAC is roughly the same for both versions. See Appendix N.7 for the relevant results.
4 Achieving local map conditions with activation function transformations
Suppose is a nonlinear layer with activation function which we propose to replace by . The four equivalent local map conditions (from Section 17.3) give rise to a system of four nonlinear equations, with the four scalar constants (, , , and ) as its unknowns. In this subsection we will show how to solve for these constants, assuming that a solution exists.
Thus is fully determined by and , which eliminates a single degree of freedom.
From Equation 7, and basic properties of expectations, we have that
Thus the condition is equivalent to
This fully determines the value of in terms of the other constants, thus eliminating another degree of freedom.
Given the above solutions for and , which we will treat as functions and of and , it remains to solve for the values of and which satisfy the final two conditions and . From the fact that (for our choice of ), these two conditions can be written as:
where the dependence on and is implicit in and .
We are not aware of any closed-form solution for this two dimensional system. However, because it’s only two dimensional, and the expectations required to evaluate it are one dimensional (including those needed to compute and ), we can readily solve it using black-box numerical software, assuming a solution exists. And because the system of equations only depends on the functional form of and no other details about , we only need to solve it once for each distinct activation function in the network. Implementation details are given in Section 22.
5 When will solutions exist?
While we found in our experiments that solutions for and exist for nearly all commonly used nonlinear activation functions, the popular RELU is a notable exception (which we will examine in the next subsection). Thus, it is worth delving deeper into the question of the existence of these solutions.
Noting that will typically be quite close to 1, if we can show that , this will suggest that is achievable by choosing a sufficiently small value of . Intuitively speaking, shrinking allows us to effectively narrow the interval of typical inputs to , meaning that starts to resemble an affine function over this interval (since differentiable functions are, by definition, closely approximated by their 1st-order Taylor approximations within any sufficiently small neighborhood). As discussed in Section 13.4, this means that as , which in turn implies that (as we have by Remark 14 that when ).
The following proposition formalizes this intuition, although is proved (in Appendix F.1) using a different technique.
Let be a nonlinear layer with transformed activation function defined as above, with and chosen as per Section 18.4. If then we have
The hypothesis that is required here since otherwise will tend to the zero function as (whose C map is undefined). Apart from this restriction, there is no obvious requirement on for the condition to hold, and indeed in our preliminary tests we found that we could satisfy this for nearly all reasonable choices of for most activation functions. The role of can thus be thought of selecting the position in ’s graph to “zoom in on”, and gives us the extra flexibility needed to control the value of .
6 The problem with positively homogeneous activation functions
A positively homogeneous activation function of degree is one where for all non-negative scalars . A well-known example for is the RELU activation function, which is given by .
Due to their defining property, positively homogeneous activation functions yield at most three effective degrees of freedom under our parameterized transformation, instead of the typical four. This can be seen by observing that
7 Examples of transformed activation functions
In this subsection we will give some examples of transformed activation functions produced by DKS. Our examples will assume a basic feedforward network of 100 combined layers, and a global slope bound . We will consider the standard tanh and RELU activation functions, as well as Swish (Prajit et al., 2017), SELU (Klambauer et al., 2017), and a commonly used smooth substitute for RELU called “softplus” (which is given by ).
The following table gives the approximate values for the activation function parameters found by DKS:
In the following plots we compare the default and transformed activation functions over the input interval $$ for tanh, softplus, and RELU. Assuming uniform q values of 1, and that the error in our kernel approximations is relatively low, this interval contains all the inputs that our nonlinear units will see at initialization time with overwhelming probability.
We can see from these plots that the transformed activation functions tend to look more like the identity functions than the defaults ones do (over the relevant range of inputs). In fact, they all bare a resemblance to each other (especially tanh, softplus and swish), as can be seen in the following plot:
Addressing normalization layers
Batch Normalization (BN) layers (Ioffe and Szegedy, 2015) are an important component in many neural network architectures, especially convolutional networks. For each unit scalar in their input, BN layers compute a mean and variance of over the training mini-batch, and then output a “normalized” version , where is a small constant. This is this sometimes followed by the application of per-channel learnable bias parameters, which are initialized to zero.
Because they use statistics computed over the mini-batch, BN layers cannot really be described in Q/C map framework we have presented, and are therefore incompatible with DKS. In particular, our formalism assumes that the network’s computation for a single training input depends only on that input, and not on other elements of the mini-batch. To account for such interactions, one would have to introduce hypotheses on the size of the mini-batch and the statistical distribution of its vectors, as the behavior of BN layers are highly dependent on these factors. Moreover, the evolution of q and c values would not happen independently across the mini-batch, which would likely preclude a simple one-dimensional description like Q and C maps.
2 Layer Normalization layers
Note that LN layers were not explicitly defined for convolutional networks in the original paper. Thus, one could also conceivably define them as computing a mean and variance over both locations and channels, instead of individually per location. In this work we will assume our previous definition, and anything we say regarding LN layers from this point will apply only to that definition.
Unlike BN layers, LN layers perform their computations and transformations individually per training case, and do not involve any computations across the mini-batch. Averaging of statistics instead occurs over entries (i.e. channels) of the location vectors, and the same scale and shift is applied to all entries. In general, and will be different for each input to the network, so that the learnable gain and bias cannot ever actually “undo” the normalization for all training cases cases simultaneously. This means that introducing LN layers into a network will fundamentally change its model class.
As we will show next, LN layers can be understood within our Q/C map framework, and are thus compatible with DKS. The formulas for their local Q/C maps are given below.
If we have uniform q values (so that ), then by Equation 26 this implies , so that the above expression for ’s C map simplifies to
When is an affine layer (or a sum over multiple affine layers), or is a combined/nonlinear layer transformed via DKS, we have . In this case, the above expression for ’s C map reduces to the identity function.
and so the application of the LN layer after thus has the effect of ensuring that even when .
Addressing pooling layers
Pooling layers are a type of layer used in certain convolutional network architectures to compress information from a larger feature map into a smaller one (with fewer locations). In this section we will discuss why standard pooling layers aren’t compatible with our Q/C map framework, and describe potential replacements for them which are. We will also give mathematical arguments and empirical evidence suggesting that it may nonetheless be okay to use them with DKS in practice.
Mean-pooling layers function similarly to convolutional layers, except that instead of computing a (learnable) affine function of each “patch” of activation vectors, they simply compute the average of those vectors. Typically these patches don’t overlap, and thus a mean pooling layer reduces the number of locations (while preserving the channels).
In order to simplify the discussion we will restrict our attention to “global” mean-pooling layers, which average over all locations, and are the most common type used in practice. The same basic conclusions will apply to general mean-pooling layers, with somewhat more complicated formulas for the associated kernel functions.
Formally, a global mean-pooling layer computes
From the above equation, the PKF associated with is a -matrix-valued function given by
where is the output channel dimension.
Noting that the input and output channel dimension are equal for mean-pooling layers, we have
and so only depends on the inputs and via their IPM . Thus, can be composed with APKFs to form a network-level PKF approximation.
A possible solution to the issues associated with mean-pooling layers is to replace them with layers that can be more easily handled within our framework, and which ideally don’t shrink the model class. (Expansion of the model class is less objectionable, provided that it doesn’t significantly harm generalization performance.)
Another option is something we call weighted mean-pooling layers, which are defined similarly to regular mean-pooling layers, except that the vector of 1’s in Equation 30 is replaced by a learnable vector of weights , giving
has more favorable properties than the PKF for mean pooling layers given in Equation 31. In particular, since the output q/c value is just the average across locations of the input q/c values, the property of uniform q values of 1 will be preserved, thus enabling our C map analysis to be valid for subsequent layers.
Complicating the story somewhat is the fact that the c values for different locations are averaged together, as our analysis up to this point has assumed them to be separate and independently evolving. This means that geometric information about each individual location is no longer strictly preserved, as the averaging operation makes recovery of the individual c values impossible. It is true however that each location still has a proportional effect on the output, and thus the degeneration discussed in Section 12.4 can still be avoided, as long as the C map of the subnetwork up to is sufficiently well-behaved.
In some of our experiments on convolutional networks we tried using a weighted mean-pooling layer in place of the usual global mean-pooling operation near the end of the network. While this worked well, we found that it didn’t provide any optimization benefit. (See Appendix N.9 for these experiments.) Thus in our main set of experiments with DKS we continued to use standard global mean-pooling layers, despite their apparent incompatibility with our theoretical framework.
2 Max-pooling layers
A max-pooling layer is similar to a mean-pooling layer, but instead of taking the mean of a set of location vectors, it takes the coordinate-wise maximum.
In the previous subsection we saw that the PKF for a mean-pooling layer, despite having a simple form that depended only on the IPM () of its input, had unfavorable properties that made it impossible to properly analyze within our Q/C map framework. The situation with max-pooling layers is arguably even worse, as its PKF has a more general dependence on its input, and thus cannot be composed with APKFs of combined layers to form a network-level PKF approximation. But despite this, we can still make some non-trivial statements about a max-pooling layer’s PKF that will be useful in understanding how DKS may possibly still apply to networks containing such layers.
Consider a patch of locations over which the max operation is applied. If all the location vectors in the patch are nearly equal to each other, then the max operation simply outputs a close approximation of the vector in the center of the patch, and thus has a PKF approximated by the identity (for non-dropped locations). It is therefore reasonable to approximate max-pooling layers as having local Q and C maps equal to the identity in this case. This situation is fairly common when max-pooling layers are used very early in the network, since nearby pixels tend to be similar to each other in natural image data, which, assuming well-behaved C maps, means that the corresponding vectors for subsequent layers will be similar too (as measured by their cosine similarly).
Analogously, if the pixels within a patch fall into two tight clusters, which can happen if the patch overlaps an edge or the corner of an object, then the subsequent vectors will also fall into two tight clusters. If this is the situation, and we assume uniform q values and wide layers, then it can be shown that the output q value of a max pooling layer will be closely approximated by its input q value (so that we can treat the layer as having an identity Q map). This is shown in Appendix G, and relies on the somewhat surprising fact that
for all , along with the mild assumption that the max-pooling layer in question is directly preceded by a convolutional layer. Note that for clusters of 3 or more pixel values this approximation doesn’t work, although the output q value will only deviate from the input q value by a factor that grows slowly with the number of clusters.
As with mean-pooling layers, we could consider replacing max-pooling layers with ones that are handled within our framework. However, unlike the mean operation, the max operation is difficult to elegantly simulate using our standard layer types, and so there are no obvious substitutions that would preserve the model class.
In some architectures, max-pooling layers are used merely to reduce the size of a feature map, with the particular choice of pooling operation (max or mean) being unimportant from a modeling perspective. In such cases it may thus be quite reasonable to replace max-pooling layers with weighed mean-pooling layers.
Summary of our method
In order to apply DKS we must observe certain architectural requirements on the network. These are summarized below:
The network must be constructed from combined layers (defined as an affine layer followed optionally by a nonlinear layer), weighted sums between the output of two or more affine layers (followed optionally by a nonlinear layer), concatenations of two or more feature maps along their channel dimension, mean-pooling layers, and max-pooling layers (although the latter should be used with caution as discussed in Section 20.2).
Batch Normalization layers must not be used. However, Layer Normalization layers are allowed, provided that their associated gain and bias parameters are initialized as per Item 5 of this list. (See Section 19 for additional discussion of normalization layers.)
Nonlinear layers must use element-wise activation functions. Positively homogeneous activation, such as RELU, are allowed but not recommended as they lead to a limited version of DKS (as discussed in Section 18.6).
Multiplication operations, such as those used in attention mechanisms, are also not allowed (although we hypothesize that DKS can be extended to handle these in the future).
The network should not contain any extraneous trainable parameters such as scalar multiplications or shift, unless these parameters have no effect at initialization time (e.g. a shift that is initialized to 0). Constant scalar multiplications are allowed, although these will typically be removed as part of the application of DKS.
Similarly, constant multiplications and shifts are not allowed, except as part of weighted sum operations. (Note that if such constants are normally required for the network to be trainable with standard optimization methods, it’s likely that DKS will render them unnecessary/obsolete.)
2 Execution steps
To apply DKS to a given network one performs the following steps:
Initialize each bias vector to 0, and each weight matrix/filter bank using either a Gaussian Delta initialization, or an Orthogonal Delta initialization (which are both defined in Section 4).
Choose a value larger than 1 for the scalar hyperparameter (such as 1.5 or 1.1). Note that roughly corresponding to the “degree of nonlinearity” of the network. See Section 16.6 for additional discussion of this. As observed in Section 16.6, lower values of tend to be associated with slightly better generalization, at the cost of somewhat slower optimization.
(optional) Apply some version of Per-Location Normalization (PLN) to the input data. Note that this can be done entirely online, as it only requires the current example, and not any aggregate statistics over the entire training set. (See Section 10.2 for more details.)
Remove any constant scalar multiply operations from the network.
Replace any weighted sums between features maps with “normalized sums” of the form , for weights satisfying (which may be chosen freely). Note that if for all and this simplifies to , although other choices are permitted and may indeed be preferable (as demonstrated in Section 23.3).
(optional) Replace any mean-pooling layers with “weighted mean-pooling layers” as defined in Section 20.1.2.
Compute the network’s maximal slope function (or some approximation of this). One can use the recipe given in Section 21.3.
Using the fact that is a 1D monotonically increasing function, compute using binary search (or a similar such method).
For each distinct activation function in the network, do the following:
Given , solve for , , , and as per Section 18.4, using the numerical methods described in Section 22 (or some alternative).
Replace all instances of in the network with .
3 Recipe for computing slope polynomials and maximal slope functions
As per Sections 17.1 and 17.2, the maximal slope function is computed as
where is the network polynomial of (whose computation we will describe below), and the maximum is taken over all subnetworks of the entire network.
Note while the number of distinct subnetworks may be quadratic (or worse) in the depth, when computing the maximum we may ignore any subnetwork that can be composed with another one to form a strictly larger subnetwork, or more generally, any subnetwork whose slope polynomial is a factor of the slope polynomial of another subnetwork.
For a given subnetwork , the computational graph of the network polynomial may be obtained from the computational graph of by recursively applying the following rules (which are essentially just the result of applying automatic differentiation to the graph of and then evaluating the result at to obtain ):
Composition of two subnetworks and maps to .
Concatenation operations (over the channel dimension) between the outputs of subnetworks map to
where is the number of output channels of .
Normalized sums with weights over the outputs of subnetworks mapFor reference, when computing the slope polynomial for a network whose q values may vary between layers (which won’t come up when applying DKS), normalized sums instead map to where is the output q value associated with . to
Layer normalization layers map to the constant 1.
Max-pooling and weighted mean-pooling layers map to the constant 1. (Standard mean-pooling layers can be heuristically mapped to 1, although they technically break our network polynomial formalism.)
The network’s input maps to the constant 1.
Note that this recipe for computing can be generalized to compute for networks in which may be different for each nonlinear layer (i.e. not equal to some common ) by mapping nonlinear layers to instead of .
Provided that the network architecture is compatible with DKS, a quick way to compute slope polynomials is to count the number of nonlinear layers in a given sequence of layers (to get a slope polynomial of for that subnetwork), and then apply the rule for normalized sums where appropriate. See Sections 17.2 and 23.4 for instructive examples of how to the compute maximal slope function for certain architectures.
Some implementation details
Through careful optimization and engineering, and a lot of trial and error, we were able to get the runtime of DKS down to a few seconds for typical large-scale networks. In this section we describe the aspects of this that were the most challenging and non-obvious.
As we saw in Section 18.4, finding the appropriate and constants for a particular activation function amounts to solving a system of two nonlinear equations. Since we don’t have a closed form solution for this, we must resort to numerical methods.
After trying several possibilities, we got the best results using scipy2001. We call this with the arguments jac=False, and leave all other options at their defaults. This invokes an implementation of the modified Powell algorithm (Powell, 1964). Because the implicit regression loss of the system is non-convex in general, the solver sometimes fails to find a solution from its initial guess. Our solution to this is simply to call it repeatedly with different initial guesses until it returns successfully. In our experiments it never took more than 4 calls to find a solution for any of the eleven activation functions we tried. We took our first six initial guesses for from the list
and then generated subsequent ones randomly using numpy.random.uniform(low=-3.0, high=3.0) for .
2 High-quality estimates of the expectations
Application to various modified ResNets
In this section we will discuss the very commonly used ResNet architecture (He et al., 2016a, b) and certain modified versions of it, and then go over the application of DKS to these different versions. In addition to being instructive in the application of DKS, these example will be the primary focus of our later experiments.
In this work we will only consider the “V2” version of the ResNet architecture (He et al., 2016b) as the opposed to the “V1” version (He et al., 2016a), as the former is conceptually simpler and is usually preferred by practitioners. We will also concentrate on the version of ResNet-V2 designed specifically for use in 224x224 Imagenet classification, noting that versions of the architecture for other problems and datasets can differ, especially in terms of their first and last few layers. We will denote by the “depth parameter” of the ResNet architecture, which corresponds to the total number of nonlinear layers plus 1. The standard values for are 50, 101, and 152. The input is assumed to be 224x224 features maps with 3 dimensional pixel features (possibly extended to 4 dimensions if PLN is used as per Section 10.2). This is then fed into a 7x7 convolutional layer with 64 output channels and a stride of 2. Following this is a 3x3 max-pooling layer with a stride of 2. These two early layers are particular to the Imagenet version of ResNet-V2, and have the purpose of reducing the dimension of the feature representation to a smaller size for processing by subsequent layers. Following this is a long sequence of residual blocks that form the large bulk of the network. Each of these is parameterized by an associated output channel dimension, a “bottleneck” channel dimension, and a stride, which can differ from block to block. The particular values for these parameters are determined by . Let be the input channel dimension, the output channel dimension, the bottleneck channel dimension, and the stride associated with a particular residual block. The residual block contains two “branches” from its input that get summed together at the output. The first is called the residual branch, and consists of the following sequence of layers: a Batch Normalization (BN) layer, a RELU nonlinear layer, a 1x1 convolutional layer with stride 1 and output channel dimension , a BN layer, a RELU nonlinear layer, a 3x3 convolutional layer with stride and output channel dimension , a BN layer, a RELU nonlinear layer, and finally a 1x1 convolutional layer with stride 1 and output channel dimension . The second branch is called the shortcut branch, and consists of the identity map if and , or a 2x2 max-pooling layer if and . Otherwise, if (which is the case for transition blocks), it consists of the following sequence of layers: a BN layer, a RELU nonlinear layerThis nonlinear layer in the shortcut branch does not contribute to the total number of nonlinear layers for the purposes of computing . Moreover, it can be identified with the first nonlinear layer of the residual branch (as they compute the same thing), in which case both the shortcut branch and residual branches can be seen as “branching off” from this layer’s output (instead of from the block’s original input). , and a 1x1 convolutional layer with stride 1 and output channel dimension . The shortcut branch is meant to act as an identity function or a reasonable approximation to one, except when its input and output channel dimensions differ (which is only the case for transition blocks). Meanwhile, the residual branch, which always contains 3 nonlinear layers, is what performs the interesting nonlinear computation in the network. After the sequence of residual blocks, there is a BN layer and a RELU nonlinear layer, followed by a “global” mean-pooling layer which reduces the number of locations down to 1. The final layer of the network is a 1x1 convolutional layer operating on this single location (or equivalently a fully-connected layer), whose output channel dimension is the number of classes. Convolutional layers in ResNets typically do not have bias parameters, since these are made pointless by the mean-subtraction done by the BN layers that always immediately follow them. To compensate for this, BN layers will sometimes include trainable gain and/or bias parameters applied after their centering and normalization operations.
2 Modified ResNet architecture
In this subsection we will describe the particular changes we made to the ResNet-V2 architecture in order to conform to the requirements listed in Section 21.1 and thus achieve compatibility with DKS. These changes don’t perfectly preserve the model class, although we tried to make them as innocuous as possible in order to facilitate the fairest comparison to standard ResNets in our experiments. Note that other modification schemes are possible, and the one we present here is not meant to be in any way “canonical” for this or any other architecture. As BN layers are incompatible with DKS we elect to remove them, while adding learnable bias parameters (initialized at zero) back into the convolutional layers. Another option would be to replace BN layers with Layer Normalization layers, as the latter are compatible with DKS. As discussed in Section 18.6, while technically supported, RELU activations force us to use a diminished version of DKS. Thus in our main experiments we often used alternative activation functions instead, including ones with a “RELU-like” shape, such as softplus. Max-pooling layers are provisionally supported by DKS, especially if they occur early in the network and have a relatively small kernel size. (See Section 18.6 for more details about this.) The 3x3 max-pooling layer near the beginning of the network meets these criteria, and so we elect to leave it in.
3 Further modifications made by DKS
Having achieved compatibility by making the above changes, we can now apply DKS to the resulting modified ResNet architecture. In this subsection we will describe the subsequent changes made to the network as part of the execution of DKS itself, whose steps are outlined in Section 21.2. Note that some of these steps are optional, or involve degrees of freedom, and are all designed to preserve (or slightly expand) the model class. First, the mean-pooling layer near the end of the network can optionally be replaced by a weighted mean-pooling layer (as described in Section 20.1.2). While this replacement is necessary for the Q/C map computations to make sense, we found that it didn’t significantly improve optimization performance in our preliminary experiments, and so we didn’t do it in our main ones. One possible explanation for this finding is that because there is only a single nonlinear layer after the mean-pooling layer, the non-uniform q values produced by the latter can have only a limited effect on the network’s overall C map. Next, we must replace the sum operations, which occur in ResNets at the end of each residual block (where the residual and shortcut branches are combined together), with normalized sums. Each normalized sum involves two weights and one constraint (that the squares of the weights sum to 1), and so has one degree of freedom. A natural choice is to set both weights to , which naively seems like the best option for reproducing the behavior of an unmodified ResNet. However, as we will discuss in Section 26.6, for networks that forgo normalization layers and/or use bounded activation functions (as our modified ResNets do), placing more weight on the shortcut branch will result in better behavior that more closely matches that of a standard ResNet. This is confirmed in our experiments in Appendix M.1. The final modification made to the network as part of DKS is to replace all of the activation functions with their transformed versions, as described in Step 9 of Section 21.2.
4 Computing the maximal slope function for modified ResNets
Having described the modifications we made to ResNets to achieve compatibility with DKS, and the further ones made by DKS itself, we are now in a position to compute the maximal slope function following the recipe in Section 21.3. For simplicity, we will assume that the normalized sums at the end of the residual blocks each have a weight of on their residual branches. (A weight of on the shortcut branches is then implied.) The subnetwork before the sequence of residual blocks is just a affine and max-pooling layer and so has a slope polynomial of 1. Consider any non-transition block. The slope polynomial for the residual branch is , as it has 3 nonlinear layers, and is 1 for the shortcut branch. Thus, the overall slope polynomial for the block is . Similarly, the slope polynomial for a transition block (which has a single combined layer in its shortcut branch) is . The subnetwork after the sequence of blocks consists of nonlinear layer, a (possibly weighted) mean-pooling layer, and then an affine layer, and so has a slope polynomial of . Noting that the total number of residual blocks is , and the number transition blocks is 4 for all values of , the overall slope polynomial for the network (which has a single input and output and so is a subnetwork of itself) is
The only subnetworks of that don’t compose with other subnetworks to form larger ones are the residual branches, and so their slope polynomials are the only other ones to consider when computing the maximal slope function . As they are simple compositions of layers with 3 (or 2) nonlinear layers total, their slope polynomials are (or ). Noting that (or ) is a factor of , we may ignore them when computing the maximum, and thus conclude that . It is worthwhile to note the dependency of on the value of the residual branch weight . For we have , and for we have . More generally, will be a degree polynomial in , where the coefficients (which must sum to 1) will more heavily favor high order terms as approaches 1, and low-order terms as approaches 0. Thus, much like , can be thought of as controlling the overall “degree of nonlinearity” of the network (as quantified by ).
5 Equivalent standard convolutional networks
To help demonstrate the power of DKS in our experiments, we will consider a skip-connection-free convolutional network architecture obtained from the above modified ResNet-V2 architecture by simple removal of the shortcut branches. The resulting architecture retains the channel dimensions, strides, etc., of the original ResNet architecture, including it use of “bottleneck” layers in the residual branches, but is otherwise a standard deep convolutional network. Given the straightforward sequential structure of this architecture, with its nonlinear layers, its network polynomial and maximal slope function are simply (which corresponds to the case above).
6 “Wide” ResNet variants for CIFAR-10
For our experiments involving the CIFAR-10 dataset (Krizhevsky and Hinton, 2009) we will make use of Wide Residual Networks (Zagoruyko and Komodakis, 2016), which are a well-known variant of the standard ResNet architecture. The Wide-ResNets we used in our experiments differ from standard ResNets in the following ways:
The initial subnetwork before the sequence of residual blocks consists of just a 3x3 convolutional layer with 16 output channels and a stride of 1. There is no max-pooling layer.
Given per-block parameters and , a residual branch consist of the following sequence: a BN layer, a RELU nonlinear layer, a 3x3 convolutional layer with stride and output channel dimension , a BN layer, and RELU nonlinear layer, and finally a 3x3 convolutional layer with stride 1 and output channel dimension . Note that there are only 2 nonlinear layers instead of the 3 normally present in standard ResNets.
There is a global “width” parameter which acts as multiplier on the output channel dimensions of all the residual blocks. In our experiments this was set to 2.
The scheme for mapping to a configuration for the residual blocks is generalized to work with any value of such that is divisible by 6. Here, represents the number of nonlinear layers plus 3, so that there are total blocks, 3 of which are transition blocks.
As we did for standard ResNets, to achieve compatibility of Wide-ResNets with DKS we will modify the architecture by removing the BN layers, adding back in learnable biases to the convolutional layers, and (possibly) replacing the RELU activation functions with various alternatives. Following a similar derivation to the one in Subsection 23.4, the maximal slope function for these networks is given by
where like before we have assumed weights and for the weighted sum operation at the end of each residual block. We can also define a skip-connection-free version of this architecture by removing the shortcut branches. The maximal slope polynomial associated with such a network is , as there are total nonlinear layers.
Part IV Additional analysis of DKS and related methods
Recent advances in the theoretical understanding of neural network training have shown that highly overparameterized networks behave like linear functions of their parameters over the entire course of training by gradient descent (Jacot et al., 2018; Li and Liang, 2018; Du et al., 2019b, a; Allen-Zhu et al., 2019; Arora et al., 2019). This analysis works by approximating the network function by its own 1st-order Taylor series with respect to its parameters (centered at their initial values), and then showing that the parameters remain close enough to their initial values throughout training that the approximation remains a good one. Under this approximation, which becomes exact as the width of each layer goes to infinity, training a neural network with gradient descent resembles kernel regression, with a kernelNote that the Neural Tangent Kernel is related to but distinct from the kernels we have been analyzing in this work so far. known as the Neural Tangent Kernel (NTK) that is computed from the network’s Jacobian at initialization time. This enables one to accurately predict the functional form of the trained network, and precisely characterize the rate of convergence to this solution by gradient descent. While this type of analysis has been extended to exact and approximate natural gradient descent methods (Zhang et al., 2019b; Cai et al., 2019; Karakida and Osawa, 2020), we will only consider the gradient descent version in this work. Even though real networks trained on challenging datasets like Imagenet are typically not wide enough to satisfy the formal requirements of NTK theory (especially when random dataset transformations are employed), the setting where this 1st-order Taylor approximation works well – known colloquially as the “NTK regime” – may still serve as a rough analogy to more realistic training. It is thus interesting to consider what effect the network’s architecture, activation functions, and initialization has on the NTK, and what this says about training in the NTK regime. In this section we will review the basics of NTK theory, characterize the NTK in terms of the properties of the network’s C map, and show how the C map degeneration which happens naturally in deep networks (as shown in Section 12) leads to a form for the NTK which implies very slow optimization and/or very poor generalization. We will then show how the form of the NTK under DKS, assuming a reasonable choice for the global slope bound , is much nicer, and leaves open the possibility of fast optimization and good generalization. (Although actually proving that it necessarily leads to these things would require assumptions on the dataset, and is beyond the scope of this work.)
For the remainder of this section we will assume that the network in question is a standard feed-forward MLP comprised of fully-connected combined layers, where the last such layer has an identity activation function. We will represent the network as the function for input vector and parameter vector . For notational simplicity we will assume that the network’s output dimension is 1. The parameter vector will be split across layers into segments denoted for . Each corresponds to (as opposed to itself), where is is the weight matrix for layer , and its input dimension. This non-standard parameterization, which is known as the NTK parameterization, is what we apply gradient descent on, and is required for the NTK to have its desired properties. We will not consider bias parameters in this analysis. The training set will consist of input-output pairs satisfying for all (where is the network’s input dimension), and the objective function used to train the network will be the standard mean squared error: . We will denote by (or ) the parameters at iteration of optimization. (or ) will denote their random initial value, which is determined by a Gaussian fan-in initialization applied to the standard parameters (i.e. the original ’s). In addition to our global assumption that the activation functions are infinitely differentiable everywhere except for a finite set of points, we will also assume that they are Lipshitz continuous, which is required in order to apply the results in Jacot et al. (2018).
2 NTK definition
3 Training in the NTK regime: a brief review
There are various NTK-type results that bound the convergence rate of gradient descent in the case of finite width layers (e.g Du et al., 2019b). However, such results are complicated to prove, and seem to be fairly pessimistic in terms of the rate of convergenceWhile these results predict exponential convergence, the associated rate constants are close enough to 1 that convergence requires a prohibitively large number of iterations. they predict, and the width they require. On the other hand, the situation simplifies considerably in the limit of infinite width (for all layers but input and output ones, whose width is fixed), and quite simple and elegant expressions exist for both the convergence rate of gradient descent, and the function computed at the converged solution (Jacot et al., 2018). While the infinite width limit is unrealistic, and totally ignores how kernel approximation error affects the theoretical predictions, we will nonetheless use it in our analysis for the sake of simplicity and clarity.
Before we begin we must define some additional notation. Let
and define . In the infinite width limit, provided that is positive definite (i.e. non-singular), a standard result of NTK theory is that the -th iterate produced by gradient descent (with learning rate ) satisfies
for all valid . If , where denotes the -th largest eigenvalue of , then converges to some . At this solution, the form of is given by taking in the above equation, yielding:
Observing that is the -th row of , we can “stack” both sides of Equation 33 to obtain
where denotes the eigenvector of corresponding to the eigenvalue . Plugging this expression into the objective function and using the fact that the ’s are mutually orthogonal gives the following expression for the training loss:
When , this expression converges to 0 which implies that is indeed a global minimizer of the objective. Moreover, if we employ early stopping, then directions in function space corresponding to eigenvectors with with smaller eigenvalues in will have converged less than the others. As observed by Jacot et al. (2018), this may help explain early stopping’s regularization benefits. A complete picture of the convergence of the objective requires us to know the entire spectrum of and the coefficients . However, assuming that is significantly large, the convergence speed will tend to asymptotically. With the optimal learning rate of this becomes , where is the condition number of .
If we only optimize layer , we may replace by and by in the above formulas (provided that it is positive definite) in order to obtain a description of the resulting convergence. As before, the training error will converge to zero at a speed determined by the eigenvalues of . An analogous statement is also true if we optimize an arbitrary subset of the layers, in which case we replace by and by . Note that because we are assuming infinitely wide layers there is no paradox here; each layer has enough capacity to memorize the training data entirely by itself. The form of the NTK matrix allows us to gain insight into the relative contribution of each layer to the overall solution. A layer whose per-layer NTK matrix is much smallerBy “smaller” we mean that a PSD matrix is smaller than , written , if is positive definite. than the other layers will have much smaller gradients, and the changes to its weights made during training will have a much smaller effect on the overall solution. While training any single layer is sufficient to achieve zero error in the infinite width case, what this arguably means for realistically sized networks is that layers with very small per-layer NTKs will train much slower than other layers.
4 An elegant expression for the limiting NTK using C maps
The inner-product which defines the per-layer NTK is a random variable that depends on the random initial value of the parameters . In the limit as the width of the network goes to infinity, converges in probability to a deterministic function in much the same way that the network’s kernel function does. Indeed, an approximation result directly analogous to Theorem 1 exists for the NTK (Arora et al., 2019, Theorem 3.1). As we are performing our analysis in the infinite width limit, we will take to be this limiting value going forward. Let be the subnetwork that maps the network’s input to the input of the -th combined layer (which is the output of the -th combined layer when ). Jacot et al. (2018) show that
where we note that (since is the identity), and that the quantities inside of are matrices (so that extracts their top corner entry). Let represent the -th combined layer of the network, its output q value (with being the q value for the network’s input), and its activation function. By Equations 19, 22, and 23, we can write the above expression for the NTK as
where is the c value for the network’s input (recalling that by assumption). Denote by the subnetwork that maps the input of to the network’s final output. Since we have it follows that , and so by the chain rule we have . Plugging this into the above equation we arrive at the elegant formula
While this formula has only been proven for deep MLPs (consisting of a composition of a sequence combined layers), we conjecture that it holds for more general architectures.
5 The form of the NTK matrix given a degenerate C map and implications for gradient descent training
In this subsection we will consider the situation where a deep network has a “degenerate” C map that sends nearly all input values to a small region around some value (in the sense of Section 12), and argue that this implies slow optimization and/or poor generalization in the NTK regime. This analysis can be seen as a more rigorous version of the intuitive argument given in Section 12.4, and overlaps with the results of Xiao et al. (2020).
To simplify the discussion, we will assume that each combined layer has the same activation function (except the last one, which is required to be linear), which means that the network’s C map is just the composition of copies of some local C map . We will further assume that is itself “well-behaved” in the sense that is reasonably close to 1, so that the overall C map is degenerate only because is large. Moreover, any sufficiently “deep” subnetwork of is also degenerate, and any sufficiently “shallow” subnetwork is well-behaved. Additionally, we will assume that (without loss of generality), and that there are no two distinct inputs and , from either the training or test set, for which is very close to or (which would imply that either or given our previous assumption that ).
As in Section 12.4 there are two main cases to consider for : the “collapsing case”, where , and the “exploding case”, where with . For the collapsing case we must have since is an attractive fixed point of . And for the exploding case we must have that is a non-attractive fixed point (since can only have one such point by Proposition 10), and so . For the layer index there are three cases to consider: is small so that the layer is “early” in , is small so that the layer is “late” in , and the default case where neither nor are small, so that the layer is the “middle” of the network. The following table gives estimates of the layer-wise NTK matrix for each combination of cases. These estimates are computed in Appendix H.1 using a semi-rigorous style of argument. The results of these computations have been checked numerically for the case of RELU and Erf activation functions (whose C maps have convenient analytic forms). Here, the symbol denotes the matrix of 1’s.
From the above values we can compute an estimate of the overall NTK matrix. This is given the following table, which is computed in Appendix H.2:
There are several implications for gradient descent training that we can infer from the above estimates, all of which are bad. Firstly, in the collapsing case with , and in the exploding case, the magnitude of the per-layer NTK matrices differ substantially over the network. This means that the layers whose per-layer NTKs are not amoung the largest will train very slowly. Given that such layers are only a small fraction of the total, this implies that only a few layers of the network will have the potential to train quickly. While this is technically sufficient to minimize the training loss in the NTK regime, in practice, our networks often won’t be highly overparameterized to the extent required by NTK theory, and so we actually will need to train all of the layers in order to fit the dataset. Insofar as the NTK regime is an analogy to this more realistic setting, this analysis thus predicts slow training. Secondly, in the collapsing case with , we have that the per-layer and overall NTK matrices are approximately rank 1, which implies that they have a very high condition number. This means neither the individual layers, nor the overall network, will train quickly, and so the training loss will take a very long time to be minimized no matter what subset of layers we elect to train. Finally, in all cases we have that the approximate form of the per-layer and overall NTK matrices does not depend on the input training data. Additionally, we have that the vector does not depend on the training data, since by the derivations in Appendices H.1 and H.2, doesn’t depend on or (except to detect when or ). And also won’t depend on the training data, since it will either look like a multiple of the ones vector in the collapsing case, or a completely random vector in the exploding case. It thus follows from Equation 34 that the predictions made by the fully trained network for a test point will not actually depend on the input training data in any significant way, making it impossible for the network to generalize.
The results of this subsection overlap with those of Xiao et al. (2020), who derive approximations to the overall NTK matrix for deep networks (although not for individual layers) using a different style of argument. Their results mostly agree with ours, except that for the case they predict a universal value of (whereas we observe for deep RELU networks), and for they estimate the second (and less significant) term of to be , whereas we predict . Numerical studies we performed on the Q/C maps of deep RELU/ networks seem to confirm our predictions in these cases.
6 The form of the NTK under DKS
The following theorem is proved in Appendix H.3 using Theorem 13 and Equation 35.
Suppose that is the per-layer NTK (for layer ) of a network conforming to the assumptions of Section 24.1 which has been transformed using DKS with global slope bound . Then we have
The bound in this theorem establishes that each per-layer NTK matrix converges to the training data Gram matrix as approaches 1, where X=\left[\begin{array}[]{cccc}x_{1}&x_{2}&\cdots&x_{n}\end{array}\right]. It also allows us to reason about larger values of to a limited extent, although it arguably only becomes interesting when . (We suspect that with a tighter and/or more detailed analysis, interesting statements about the relationship of and the layer-wise NTK could be made for larger values of .) If is low rank, which it will be in the common case that , this means that the will approach a low-rank matrix as approaches 1, which corresponds to slow/impossible training under gradient descent. Intuitively this makes sense, since a value of very close (or equal) to corresponds to a network that looks almost perfectly linear at initialization time (by Theorem 13), and thus could fail to properly train as per the discussion in Section 14.2. Indeed, the foundational works on NTK only predict that the NTK will be positive definite (i.e. full-rank) when the activation functions are non-polynomial (and thus nonlinear) functions, and DKS makes them approach linear functions as . So while a value of very close to 1 is clearly a bad choice, a value somewhat close to (such as 1.5; which we use in most of our experiments) will allow to retain some of the structure of , thereby ensuring that the network’s prediction (given in Equation 34) depends on the training data and thus has the potential to generalize. It will also allow to deviate enough from to be full rank with a potentially small condition number (which would imply fast training). Unfortunately, is difficult to accurately estimate without full knowledge of both and the behavior of the C map over its entire domain, and existing methods to bound (e.g. Du et al., 2019b) seem unlikely to produce useful results in our context. We leave the problem of accurately estimating the value of under DKS to future work.
Variance propagation, signal propagation, and their relationship to approximate kernel analysis
Many previous methods for constructing and initializing neural networks are justified using analysis frameworks which attempt to characterize the initialization-time behavior of neural networks. The two most prominent examples of such frameworks are “variance propagation” and “signal propagation”. In this section we review these frameworks, highlight certain mathematical issues with them, and provide counterexamples to their general claims where possible. We also relate them and their predictions to the kernel approximation framework underlying DKS, and advocate for the latter as a more powerful and mathematically rigorous alternative.
The input weight vector for each unit has norm 1, and yet for all , so that the network will not preserve the per-unit variance of . (Intuitively, this is because and have strong negative correlation.) Note that this example can easily be generalized to arbitrarily wide layers. Another more subtle issue with the analysis in LeCun et al. (1998b) is that the approximation errors arising from the analysis of nonlinear layers can easily accumulate with depth, and may push the activation functions out of their assumed range of inputs.
2 Klambauer et al.’s (2017) modified variance propagation
Klambauer et al. (2017) present a modified version of LeCun et al.’s (1998b) variance propagation analysis, which would seem to address the issues we’ve highlighted. They do this by arguing that as long as the inputs to a fully-connected layer are independent, its outputs (which are fixed linear combinations of its inputs as determined by the weights) will be approximately Gaussian distributed, thanks to the Central Limit Theorem (CLT) and the assumption of wide layers. Using this approximation they then compute the moments of the subsequent nonlinear layer using Gaussian integrals (similar to those that define Q maps), without having to make any strong assumptions on its activation function. Unfortunately, CLT is not actually applicable to arbitrary weighted sums of variables, even when those variables are perfectly iid. For example, if the weights of the sum are , then the output will have the same distribution as the first input unit, which won’t be Gaussian in general. Even if we somehow ruled out such weight vectors as having “low probability”, and focused only on weight vectors for which CLT would apply, there would still be major difficulties to overcome. Firstly, since we need to show that the output units of a fully-connected layer are approximately independent (in order to recursively apply the same analysis to subsequent layers), we would need to show that they are jointly Gaussian distributed with a diagonal covariance matrix. This would require the use of one of the multi-dimensional versions of CLT, all of which require significant additional hypotheses compared to the standard one-dimensional versions. Secondly, the approximate independence provided by CLT would not be sufficient to recursively apply the same argument to subsequent layers, as CLT typically requires exact independence of the variables under summationIn an attempt to preempt this criticism, Klambauer et al. (2017) refer to Bradley (1981), which proves a version of CLT that relaxes the independence assumption. However, this result assumes a very specific type of weak dependence which is unlikely to satisfied in this setting, and also only applies for one-dimensional variables.. Thirdly, because CLT requires that the number of variables under summation is large, it usually won’t be applicable to the first layer of the network (where the input dimension is a fixed property of the training data).
3 The version of variance propagation in Glorot and Bengio (2010)
Glorot and Bengio (2010) present a modified version of LeCun et al.’s (1998b) variance propagation analysis, which has formed the basis of many subsequent analyses over the years. The first change they make is to compute per-unit variances with respect to the joint distribution on network inputs and parameters (as opposed to just the inputs). Their second modification is to directly assume that the network’s activation functions behave like the identity function over typical inputs, thus implying that they preserve the mean and variance of their inputs. Assuming that the weights of a given fully-connected layer are iid with mean zero and variance , and that its input units are mean zero with variance , they show that the per-unit variance of the layer’s output is simply , where is the input dimension. Notably, by computing variances with respect to training cases and parameters, they do not require the input units to a layer to be uncorrelated. (Intuitively, this is because the multiplication by the independent mean-zero random weights causes any two random variable to become decorrelated.) This addresses one of the main problems of the original variance propagation analysis, and allows it to be recursively applied over the entire network without issue. Unfortunately, the modification also introduces a new issue not present in prior analysis: the variances no longer refer to any single network (with a particular parameter setting), but rather to a distribution over networks. This makes the interpretation of these variances unclear, and represents a subtle but serious issue in their analysis. One could possibly argue that, with high probability, a single network sampled from this distribution would have variances similar to those computed over the whole distribution. However, this would require additional hypotheses, since otherwise there are simple counterexamples to the general claim. For example, consider a linear network with fully-connected layers of width 1, where the biases are zero and the weights are sampled iid from . The function computed by this network amounts to just multiplying its input by the product of scalar weights drawn independently from . Variance propagation would predict that such a network will exactly preserve the variance of its input, or in other words, that the product of these weights would be approximately 1. However, the distribution of the product of independent samples from is highly concentrated around zero for even moderate large values of (which can be seen via Monte Carlo simulation), despite the fact that the variance of this product is 1. These counterexamples are not restricted to narrow networks either. If, for example, the weights are drawn iid from a heavy-tailed distribution that is highly concentrated around zero and has variance (where is the width), then even for large there will be an overwhelming probability that all the weights will be close to zero, leading to a network which “squashes” its input. This contradicts the prediction made by variance propagation, which is that such a network would approximately preserve the variance of its input.
4 Extension of variance propagation to RELUs
He et al. (2015) extend the version of variance propagation in Glorot and Bengio (2010) to deal specifically with RELU activation functions, as there is no zero-centered range of inputs for which RELUs resemble the identity function. To do this, they introduce the additional hypotheses that the weights have a symmetric distribution around zero, that the biases are initialized to zero, and that each RELU layer is directly preceded by a fully-connected layer. Given these hypotheses, it follows that the input to each RELU layer is distributed symmetrically around zero, and thus the expected squared valued of a RELU unit will be exactly times its input variance. One can then use this expected squared value in place of the input variance for the variance propagation calculation at the next layer, since multiplication by the mean-zero weights of said layer will restore a mean of zero. While it deals with nonlinear layers in a cleaner fashion (at least for RELU networks), this analysis retains the central issue present in Glorot and Bengio’s (2010) analysis, which is that the variances do not necessarily describe the behavior of a single network. Moreover, while the variance propagation formulas were originally derived for fully-connected networks, He et al. (2015) also applies them to convolutional networks without any additional justification. (This is problematic since the weight sharing violates the iid weights assumption.)
5 Extensions of variance propagation to normalizer-free residual networks
6 Signal propagation (aka mean field analysis)
Closely related to variance propagation is an approach for understanding the initialization-time behavior of neural networks commonly referred to as “signal propagation”Note that some works (e.g. De and Smith, 2020) use the term “signal propagation” to refer to certain versions of what we have been calling “variance propagation”. In such works elements of both types of analysis often appear, and the precise distinction between them becomes a bit blurry. or “mean field analysis” (Poole et al., 2016). In this approach, instead of propagating variances, one propagates per-unit expected squared values , or expected products between corresponding units from two copies of the same network (each fed different inputs). Here, expectations are taken with respect to the distribution on network parameters, and on the two network inputs (which may be correlated). In order to propagate through nonlinear layers, one approximates their input as being Gaussian distributed with mean zero and covariance matrix determined by the expectations from the previous layer. As will be explained in the next subsection, the expectations computed under signal propagation end up being equal to q and m values (or c values, after suitable normalization) as we have defined them in this work. Indeed, Q/C maps were originally derived by Poole et al. (2016) in the context of signal propagation. The mathematical justification of signal propagation given by Poole et al. (2016) in the case of a single fully-connected combined layer is roughly as follows. One starts from the assumption that the entries of the input vector are iid random variables with expected squared values given by . Then, multiplication by an random matrix with mean-zero iid entries of variance produces outputs, each of which has bounded variance , and is a sum of iid terms. For large one applies the Central Limit Theorem (CLT) to get that these sums will be approximately iid Gaussian distributed, with mean zero and variance . It then follows that the entry-wise outputs of the nonlinear layer are approximately iid, with expected squared values given by Gaussian integrals. Expected products are then handled using a straightforward generalisation of this argument. In principle, this single layer argument can be applied recursively to a composition of combined layers, always starting from the hypothesis that the entry-wise inputs to a given layer are iid with some known expected squared value. Unfortunately, this recursive approach runs into the same problems with CLT discussed in the last paragraph of Section 25.2, which cannot be easily repairedTo the best of our knowledge, the only mathematically rigorous CLT-based treatment of the width-limiting behavior of random networks is that of Matthews et al. (2018), which is given in the context of approximate kernel analysis. It’s not immediately obvious if/how Matthews et al.’s (2018) arguments can be used to rigorously justify signal propagation.. As with variance propagation, the interpretation of the expectations computed under signal propagation isn’t clear. In particular, there is no obvious relationship between these expectations, and the properties of a single randomly initialized network. Signal propagation’s two main advantages over variance propagation are that it handles nonlinearities in a much more general and precise way (via Gaussian integrals), and that it also describes the propagation of expected unit products for correlated network inputs. These features make it a much more powerful framework for understanding the understanding the initialization-time behavior of neural networks, and for designing initialization schemes. However, as we will discuss next, approximate kernel analysis has the same advantages while also being mathematically rigorous and more clearly interpretable.
7 Relationship of variance/signal propagation to approximate kernel analysis
When (in the notation of the previous subsection), signal propagation’s defining equations for fully-connected combined layers are precisely equivalent the local Q and C maps computed under approximate kernel analysis. We may thus interpret the quantities propagated by signal propagation as q and m values, and their normalized versions as c values. And for other values of , a similar statement holds for a slightly generalized notion of Q/C maps (as given in Poole et al. (2016)). In some sense, this equivalence acts as a mathematical justification of signal propagation’s equations, although with a different meaning for the quantities being propagated. In particular, the expected squared values computed by signal propagation correspond to q values, and can thus be thought of as approximations of dimension-normalized squared norms of the associated feature map’s vectors. Similarly, the expected products computed by signal propagation correspond to m values, and can thus be viewed as approximations of the dimension-normalized inner-product between two such vectors (or the same vector for two different network inputs). Given this relationship between approximate kernel analysis and signal propagation, we can also relate approximate kernel analysis to Glorot and Bengio’s (2010) version of variance propagation (and its extensions). To so see this, note that insofar as the units in each layer have mean zero (under variance/signal propagation’s assumed distribution), their expected squared values are equal to their variances, in which case variance propagation also computes q values. Moreover, even when the means are not zero, as is the case for RELU networks, one can modify variance propagation to deal directly with expected squared values in a manner similar to He et al. (2015). While approximate kernel analysis provides the same level of description as signal propagation, it has several advantages. The first is that it is based on a rigorous mathematical theory with clearly defined hypotheses and probabilistic error estimates. This allows one to be confident in determining which architectures it can be applied to, and to have a rigorous pathway for extending it to new architectures (which we exploited in our treatment of normalization and pooling layers). The second advantage is that the quantities it computes have a clear relationship to the (high probability) initialization-time behavior of actual randomly initialized networks with definite inputs and weights. The third is that it applies to networks with low dimensional inputs, for which the CLT-based arguments commonly used to justify signal propagation are inapplicable. And while these advantages come at the cost of additional/stronger hypotheses (such as Gaussian or SUO-distributed weights), such hypotheses are likely required in order for the predictions made by the equations to be accurate in general.
8 Extensions of variance/signal propagation to networks with Batch Normalization layers
De and Smith (2020) propose an extension of variance propagation to networks with Batch Normalization (BN) layers, in order to analyze standard ResNets. To do this, they argue that for large mini-batches, BN layers will compute a per-unit empirical variance which closely matches the per-unit variance computed under variance propagation. Thus, after normalization by the square root of this variance, the per-unit output variance of a BN layer will be always be 1, regardless of its per-unit input variance. There appears to be a subtle issue with this argument. As discussed above, variance propagation is a faithful description of a single randomly initialized network (with definite inputs) only insofar as the variances it computes correspond to q values. q values in turn are approximations of dimension-normalized squared norms of entire activation vectors, and have no clear relationship to the properties of individual units within a layer of such a network. So in general, the empirical unit-wise variances computed by a BN layer will not correspond to the variances computed by variance propagation, even approximately. It is conceivable that with additional hypotheses on the batch size and distribution, the network, and the initialization, the empirical distribution of the values of each input unit to a BN layer (taken across the mini-batch, for fixed parameters) might all have roughly the same variance with high probability, in which case the approximation in De and Smith (2020) would be a valid one. However, formalizing this would likely be quite difficult. Yang et al. (2019) propose an extension of signal propagation/mean field analysis to networks where BN layers are inserted between affine and nonlinear layers. To facilitate this, they propagate matrices representing the expected products between different copies of the same unit for each of possible inputs to the network, where is the batch size. For networks without BN layers the propagation equations decomposes nicely in terms of low-dimensional Q/C maps, while for networks with BN layers no such decomposition exists, due to the way different elements of the mini-batch interact in BN layers. Yang et al. (2019) are nonetheless able to analyze the resulting high-dimensional propagation equations using various sophisticated approximations and characterize their asymptotic fixed point behavior. In their approach, the batch size, as well as the distribution used to generate the mini-batch, are encoded via the initial matrix of expectations to the first layer. Thus, their analysis is not dependent on being large, or on any strong distributional assumptions about the mini-batch. However, unlike for networks with element-wise nonlinearities (where approximate kernel analysis gives rise to the same equations as signal propagation), there is no mathematically rigorous derivation of their generalized equations for BN layers. Thus, it remains an open question as to whether these equations are accurate approximations in the sense of Section 5.8. Yang et al. (2019) provide empirical evidence that they are, at least for fairly wide networks with some commonly used activation functions.
Review and analysis of related approaches for constructing and initializing deep neural networks
In this section we will review some existing techniques, both standard and otherwise, for constructing and initializing neural networks in order to make them easier to train. We will further analyze these techniques from the perspective of approximate kernel theory by exploiting the latter’s connections with variance/signal propagation established in Section 25.
The classical fan-in initialization (LeCun et al., 1998b) for fully-connected neural networks samples filter weights iid with mean zero and variance , where is the total input dimension. Here, is precisely the value required for their version of variance propagation to predict constant per-unit variances throughout the entire network, with other values leading to an exponential increase or decrease with depth. However, as LeCun et al.’s (1998b) variance propagation analysis is only a reasonable approximation for activation functions that preserve the mean and variance of their input, their initialization will tend to fail in more realistic settings, especially as the network’s depth increases (He et al., 2015). Glorot and Bengio (2010) use their own version of variance propagation to motivate a similar initialization scheme, where the weight variance is , with being the output dimension. This choice is made as a “compromise” between the following two competing constraints: that the per-unit variances should be uniform across layers, and that the variances of the per-layer gradients should also be similarly uniform. We would argue that is not a good choice in general compared to . For example, if the layer widths alternate between between and for some , running variance propagation across two consecutive combined layers would predict a decrease in the variance by a factor . This will lead to an exponential convergence of the variance towards zero as depth increases. Meanwhile, for the choice , variance propagation (or approximate kernel analysis) predicts no such exponential increase or decrease for any choice of widths. He et al. (2015) propose to use a weight variance of specifically in RELU networks, which compensates for how RELU nonlinear layers decrease the variance by a factor of instead of preserving it. This is based on their expanded version of variance propagation that handles RELU activation functions. Setting aside issues of mathematical rigor and the interpretation of the quantities being propagated, variance propagation and approximate kernel analysis involve similar calculations (as discussed in Section 25.7), and so these three initialization schemes can all be viewed as methods to control the q values of the network. When combined with the normalization of the input vectors (as per Section 10.2), and applied to standard feed-forward fully-connected networks with suitable activation functionsHere, “suitable” means (approximately) mean and variance preserving for the standard fan-in initialization, or RELU for He et al.’s (2015) modified version. Note that the large majority of activation functions do not fall into the former category., the fan-in initialization method and its extensions achieve q values of throughout the network, which is one of the four constraints enforced by DKS. Having q values of 1 ensures that local C maps are the same for each combined layer (assuming they all use the same activation function), and that the final output of the network falls within a reasonable range. If this is not done, q values can grow very large or small with increasing depth, leading to various problems. In particular, very large values can cause bounded monotonic activation functions like to “saturate”, so that local C maps become increasing degenerate with depth. And very small values can cause most activation functions to behave in a way that is “too linear”, which may limit the effective expressivity of the network (as per the discussion in Section 14.2). Notably, the RELU activation function is immune both of these issues due to it being positively homogeneous, which perhaps explains its popularity. However, by not enforcing the other three conditions of DKS, networks using these initializations can still have degenerate network-level C maps, and can experience an exponential accumulation of kernel approximation errors with depth (so that the q values won’t actually be constant in practice). As a concrete example of the former problem, consider the example from Section 12.1 of a standard deep RELU network. This network’s C map doesn’t depend on the input q value at all (as long its uniform), but still develops degenerate behavior at very high depths, leading to a network that is essentially untrainable.Another more subtle issue with these initializations is that q values of 1 will work very badly for certain activation functions. For example, consider the activation function defined by . As increases, an input q value of 1 becomes arbitrarily bad, leading to increasing levels of saturation and consequent C map degeneration. The reason that q values of one work reasonable well in practice is that the most commonly used activation functions in the literature happen to work well it, or have local C map behavior that is insensitive to q values (as is the case for RELUs). Notably, DKS does not suffer from this issue (despite also enforcing q values of 1), as its use of a multiplier on the input of each activation function ensures that 1 will always be optimal (since any other value can be effectively “simulated”).
2 Layer-Sequential Unit-Variance initialization and Within-Layer initialization
Mishkin and Matas (2015) proposed an initialization method called Layer-Sequential Unit-Variance (LSUV), which uses an iterative procedure that starts from a standard random initialization and adjusts the scale of each weight matrix/filter bank to achieve the condition that the variance of the output of each affine layer – taken over the channels, locations and training cases – is approximately equal to 1. These variances are computed by evaluating the network empirically on random mini-batches of training data. By taking in Equation 36 to be the identity function, we have that the average value (across channels) for each location-wise output vector of an affine layer is approximately zero with high probability. The variances computed by LSUV can therefore be interpreted as estimates of the length-normalized squared norms of the location vectors for each layer, except that they are also averaged over locations and network inputs. Thus, we can think of LSUV as enforcing the condition that the “average q value” for the output of each affine layer is equal to 1. This condition is similar to the one that the fan-in initialization (and its variants) are trying to achieve, and thus our discussion and critique of those methods (in Section 26.1) also applies to LSUV. In particular, a network initialized with LSUV can still have degenerate C maps, with all of their consequent problems. Because LSUV uses empirically computed statistics instead of canned formulas, it takes into account the given architecture and network topology, as well as the properties of the input training vectors. By contrast, variants of the fan-in initialization are usually only valid for the particular activation function and network topology they were derived for, despite often being applied more generally. They also implicitly assume that the training input vectors are appropriately normalized, which isn’t always the case in practice. (Note that DKS, while it is also based on formulas instead of empirical evaluations of the network, takes into account the activation functions and network topology, and is packaged with a data preprocessing technique for the input vectors.) The “Within-Layer Initialization” (WLI) of Krähenbühl et al. (2016) can be viewed as a modification of LSUV where one enforces the condition that the mean and variance of the output of each layer and each channel is 0 and 1 respectively (as opposed to LSUV, which considers the average variance across channels). This is done by rescaling the filter bank weights separately for each output channel, and setting the bias appropriately. Because it enforces conditions per-channel instead of averaging across channels, this modification is harder to compare directly to fan-in initializations or DKS. It is perhaps more closely related to Batch Normalization (which is discussed later in this section), as it achieves the same conditions that BN does before the first step of optimization.
3 Self-normalizing neural networks
Self-normalizing neural networks (Klambauer et al., 2017) use Scaled Exponential Linear Unit (SELU) activation functions to achieve a per-unit mean of 0 and variance of 1 asymptotically with depth, as computed under variance propagation. Due to the relationship between variance propagation and Q/C maps (discussed in Section 25.7), this is equivalent to having an attractive fixed point at , and having an attractive fixed point at , where is a SELU nonlinear layer. Assuming the use of PLN, and a standard feed-forward architecture (or normalized sums in more general architectures), these conditions imply that and for all subnetworks , which is two of the four conditions enforced by DKS. As previously discussed, for all subnetworks is a good condition to have, and will prevent extreme q values from developing in deep networks (which can adversely affect C maps). However, it won’t in general guarantee a well-behaved C map in deep networks, even when combined with the condition . For a SELU nonlinear layer we have , which can be computed numerically using Equation 22.2 and the methods described in Section 22.2. Along with the condition , this guarantees a well-behaved C map up to a modest depth. For example, given a standard 100 layer network we have and , so that is reasonably well behaved according to Theorem 13. However, if has 300 layers then we have , which indicates degenerate behavior with (in the sense of Section 12.4).
4 The “Edge of Chaos” (EOC) method
Consider a network defined by a composition of combined layers, each with the same nonlinear activation function . Every combined layer will have the same Q map, which we denote by . As discussed in Section 10.1, will typically have a fixed point which is rapidly converged to under repeated applications, and thus one may approximate the q values as uniform and constant across layers. This allows one to define a local C map which only depends on the input c value, and which is the same for each combined layer. As argued by Schoenholz et al. (2017), will describe the asymptotic dynamics of the c values as they evolve through the layers of the network in the limit as . indicates rapid convergence of the c values to , while indicates rapid convergence to a value . Because of its close proximity to these two undesirable depth-limiting asymptotic behaviors, a network with is said to be “on the edge of chaos”, and will have its c values converge slowly towards 1 at an asymptotic rate which is sub-exponential. In the initialization method proposed by Schoenholz et al. (2017), which we will call the “Edge of Chaos” method (EOC), one initializes the weights using a standard Gaussian fan-in method, with the variance of the weights and biases chosen so that . (Note that also depends on these variances, which is taken into account when computing .) As observed by Schoenholz et al. (2017), there are typically infinitely many combinations for these two variances which achieve (assuming any exist for the given activation function), and so one is chosen arbitrarily. Notably, the condition is based entirely on the properties of , and as such does not depend on the depth of the network. In Xiao et al. (2018), a version of EOC was proposed that used an Orthogonal Delta initialization technique for convolutional layers, with variances chosen so that . In their experiments Xiao et al. (2018) showed that basic convolutional neural networks (without skip-connections or batch normalization) can be successfully trained on CIFAR-10 with the resulting initialization at depths of up to 10,000. This was a remarkable result, as such networks are considered essentially impossible to train at even modestly large depths when initialized using standard methods.
DKS is in many ways a spiritual successor to EOC, and is derived using an extended version of the same basic Q/C map analysis that underlies the latter. Like EOC, DKS also makes use of the Orthogonal Delta initialization technique. But despite these similarities, there are many important differences between the two methods, which we will discuss in sequence below. Firstly, DKS enforces uniform q values via data pre-processing instead of relying on the (presumed) convergent fixed-point behavior of the Q maps to achieve this asymptotically. See Section 10.1 for a more detailed discussion of this point. Secondly, instead of targeting the condition for each combined layer , DKS targets , so that the “degree of nonlinearity” is calibrated to the given architecture. This is motivated by looking at the overall C map behavior of the network, instead of the fixed point convergence behavior of its local C maps. From the perspective of fixed point convergence, DKS achieves exponential rate towards , while EOC achieves sub-exponential rate towards . While this might seem like a point against DKS, one must remember that the precise rate of its exponential convergence will dependIn particular, the worst-case rate of convergence to the fixed point will be given by , which by Equation 39 satisfies . For a layer convolutional network this is , which will be just slightly below when is large (given typical choices for ). on the overall depth , the effect of which is that the c values will be far from converged even by the -th layer. The third difference between DKS and EOC is that DKS manipulate the network by transforming the input and output of the activation functions, as opposed to changing the variances of the weights and biases. If we consider the “equivalent parameters” (as per Section 18.3), DKS is implicitly searching over a space of distributions with more degrees of freedom than the two used by EOC, and one which allows for non-zero correlations between the weights and biases. Fourthly, despite having more degrees of freedom with which to manipulate the network, DKS makes use of all of them in order to enforce a total of four conditions on the Q and C maps of the network (which are listed in Section 17.3). EOC meanwhile only enforces a single condition (), which leaves one of its two degrees of freedom unconstrained. (The effect of this is that EOC has a manifold of possible weight and biases variances from which to choose, and is therefore under-determined.) Fifthly and finally, DKS is applicable to a diverse set of architectures, thanks to our generalized notion of Q/C maps, analysis of pooling layers, and network polynomial construction. As it was originally developed, EOC assumes a strictly feedforward network consisting of a sequence of fully-connected or convolutional combined layers.
It is worth pointing out that the condition enforced by EOC does not necessarily imply that the network will look perfectly linear at initialization time, as is not enforced in EOC as it is in DKS. Nonetheless, depending on the activation function, there may be choices for the weight and bias variances which can make the network look “too linear”, thus leading to very slow training as per the discussion in Section 14.2. As such choices are not explicitly forbidden in EOC, and are indeed compatible with the condition , this may represent a failure mode of the method. Conversely, without , the condition may not sufficient to ensure that the entire network’s C map is well-behaved. For example, given a choice of 1 and 0 for the variances of the weights and biases respectively, we have for unmodified RELU activation functions (by Section 12.1), and yet quickly degenerates as grows, as can be seen from the figures in Section 12.1. Moreover, our experiments in Section 28.2 confirm that standard deep RELU networks (without BN layers or skip connections) are not readily trained at high depths, even with a high-powered optimizer like K-FAC.
5 The Looks Linear method
Balduzzi et al. (2017) use path-weight analysis (which we review in Appendix I) to argue that gradients with respect to the network’s input will decorrelate or “shatter” in deep RELU networks, leading to difficulties when training with gradient descentNote that this prediction agrees with our NTK analysis in the sense that the per-layer NTK matrix of the first layer of a deep RELU network is approximately the identity, although the implications for training are somewhat different; see Appendix I. They also argue that this happens to a much lesser extent in ResNets. Motivated by these observations, and by the fact that this effect doesn’t occur in a purely linear network, Balduzzi et al. (2017) propose the Looks Linear (LL) method for initializing/constructing RELU networks. This method exploits the fact that when is the RELU function in order to construct a network that behaves exactly like a linear one at initialization. In particular, one replaces by for each RELU nonlinear layer (which effectively doubles its output channel dimension), and initializes the weights of the affine layer after each RELU layer according to , where is sampled according to a Delta Orthogonal initialization. For fully-connected combined layers this produces an overall computation , while for convolutional combined layers the computation is similarly linear (although harder to express in standard matrix notation). From the perspective of our analysis, perfectly linear-looking networks have (very) well-behaved C maps, and thus satisfy one of the main necessary conditions for trainability. With such networks there is always the danger that they may be “too linear” in the sense of Section 14.2, but LL-initialized networks avoid this because will become highly nonlinear given a relatively small perturbations of away from . The main two obvious disadvantages to the LL approach are that it only works for RELU networks, and that it doubles the widths of RELU layers (without proportionally increasing the network’s expressivity/capacity). Beyond these things, the main difference between LL and DKS is the precise mechanism used to make the network “look linear”, and the implications that this has for optimization (which is not well understood in either case). In DKS, the degree of nonlinearity of a nonlinear layer, as measured by properties of its C map (such as the slope at ), varies smoothly as a function of the parameters of the transformed activation functions, and so one could argue that it should also vary smoothly as a function of the network’s parameters, resulting in easier optimization. By contrast, the linearity property of the LL method depends on a delicate mirrored symmetry of the weights in each layer, so that relatively small perturbations in these could lead to large changes in the degree of nonlinearity of the network. This sensitivity may make optimization more difficult, and may explain the optimization difficulties we observed in our experiments with the LL approach. (See Section 28.5.5 for more details.)
6 Residual connections
Residual Networks aka ResNets (He et al., 2016a, b), which are described in detail in Section 23.1, have become the dominant neural network architecture for computer vision problems. What makes ResNets so successful isn’t that they are more powerful or expressive than other more traditional deep convolutional architectures like VGG (Simonyan and Zisserman, 2015), but rather that they are easier to train with stochastic gradient descent at very high depths (He et al., 2016a; Szegedy et al., 2017). This easier training is owed to their use of skip connections (aka shortcut connections; which have been a feature of network architectures since the 1990s), Batch Normalization (BN) layers (Ioffe and Szegedy, 2015), RELU nonlinearities, and the surprising interplay between all three of these components (De and Smith, 2020). Moreover, popular new architectures such as Efficient Nets (Tan and Le, 2019) and Transformers (Vaswani et al., 2017) are based on the same high-level residual block structure, and differ only in terms the layers contained in their residual branches. ResNets, and their generalizations, thus represent a solution to the problem of how to achieve fast and stable training of very deep neural networks. And while the nature of this solution is still not totally understood, there has been progress in this direction (of which we will cover only a small subset). Veit et al. (2016) argued that residual networks behave like a ensemble of shallow networks of varying depth throughout training. They gave evidence for this by showing that deep residual networks are highly robust to “lesion” operations which remove or rearrange layers, and that the network’s gradient is dominated by contributions made by paths through the network with fewer nonlinear layers Zhang et al. (2019c) observed that if one removes the BN layers from a ResNet-V2 network and initializes the last convolutional layer of each residual branch to zero (along with a few other smaller tweaks to the architecture and its initialization), the resulting network achieves training speed comparable to a standard ResNet, at least at modest batch sizes. In such networks, the residual blocks act as identity functions at initialization tune, only becoming nonlinear as training progresses. Subsequent work showed that one could achieve similar results in BN-free networks simply by using learnable weights on the residual branches that are initialized to zero (De and Smith, 2020; Bachlechner et al., 2020). More recently, it was found by Shao et al. (2020) that the branch sum can use static (non-learnable) weights, where the relative size of the weight on the residual branch is set to a small value (that can vary between blocks). To help explain these findings for BN-free networks, De and Smith (2020) applied a version of variance propagation to argue that the per-unit output variance of a residual block will be roughly 1 plus its per-unit input variance, so that the -th residual block has a variance proportional to . Then, because the output variance of each residual branch is constant (due to the use of BN), it follows that the relative contribution to the block’s output made by the residual branch shrinks as . This, they argue, leads to a network which behaves more like an linear function than it otherwise would. In Appendix J we make this argument more rigorous by computing q values in a (nearly) standard ResNet, and showing that their growth over layers leads to a better behaved C map. We also show that an identical C map can be obtained in a network without normalization layers via careful selection of weights on the residual and shortcut branches.
7 Normalization layers
Normalization layers (the two most common types of which are defined in Section 19) have become a standard component in neural networks, since the introduction of Batch Normalization (BN) by Ioffe and Szegedy (2015). In addition to the important and specific role they play in ResNet-style architectures (as discussed in Section 26.6 and Appendix J), these layers have been observed to make deep neural networks easier to train on their own. In this subsection we will discuss possible explanations for this, with a particular focus on ones arising from Q/C map analysis. We will also give some arguments for why normalization layers alone are insufficient to enable fast training of deep networks.
As discussed in Section 19.2.1, a Layer Normalization (LN) layer has the property that for all , provided that its learnable gain and bias are set to their initial values. Additionally, when applied after a combined/nonlinear layer , the C map of the composition has the property that , regardless of C map behavior of . Thus, when used after each nonlinear layer in a network initialized as per Section 4, LN layers achieve uniform q values of 1 throughout the network, and also for all subnetworks , which are two of the four conditions enforced by DKS. (Note that if LN layers are instead inserted before each nonlinear layer, they will not achieve the latter condition.) The way LN layers achieve these conditions differs from DKS in at least two ways. First, they perform direct calculation of the relevant quantities instead of using q and c values as approximations. This allows them to work with arbitrary initializations of the parameters (including badly scaled ones), poorly scaled input data, and without any explicit knowledge of the network’s structure or activation functions. Second, LN layers continue to enforce a version of these conditions throughout training, or at least as long as their learnable gain and bias remain close to their initial values of 1 and 0.As discussed previously (e.g. Subsection 26.1), uniform q values of 1 is a useful property to have, but is far from sufficient to ensure trainability. The condition for all subnetworks is meanwhile only one half of the two conditions required by Theorem 13 to ensure well-behaved C maps, and arguably the less important of the two. To make this discussion more concrete, we will consider how putting LN layers after each nonlinear layer will effect the C map of a standard deep RELU network. For a RELU nonlinear layer we have by Equation 27 that and (which follows from Equation 27 by taking the derivative and letting ). Taking the derivative in Equation 29 we have , and so by the chain rule . Thus we see that while the use of an LN layer after a RELU layer gives us , it comes at the price of increasing the C map slope from 1 to . The following plot shows the extended C map for a RELU network with 20 combined layers, with and without LN layers used after each nonlinear layer:
From this plot we can see that the C map for the network with LN layers has a much larger output range. However, for the vast majority of its input domain, the output is restricted to a small region around 0, and it is still highly degenerate in the sense of Section 12.4 (and thus suggestive of poor training). Beyond their affect on the initial behavior of the network, LN layers may also have an independent and possibly beneficial effect on optimization, as they change the relationship of the loss and the parameters. Ba et al. (2016) argue that LN layers lead to a Fisher information matrix with more favorable properties for optimization. Another intuition is that an LN layer decouples the scale and direction of the weights of its immediately preceding affine layer, which may encourage faster optimization with gradient descent. Networks with LN layers may also be also be smoother when considered as functions of either their parameters or their inputs, since the change in the output of an LN layer is always bounded. Despite these intuitions, as far as we know there has yet to be strong theoretical or empirical evidence in favor of a specific optimization benefit to LN layers beyond their affect on the network’s initial behavior.
As discussed in Section 19.1, Batch Normalization (BN) layers cannot be analyzed within the Q/C map framework we have presented. Despite this, we can still make some observations regarding their effect on network behavior in the context of our previous discussions. As shown in Section 12.4, one common way that a deep network can become difficult/impossible to train is when all input vectors map to approximately the same output vector (as measured by cosine similarity) at deeper layers of the network. This happens naturally in deep RELU networks, where BN is typically applied. Placing BN layers throughout the network may mitigate this particular pathology by ensuring that the empirical distribution of each unit over the mini-batch has a large variance compared to its mean. However, this won’t obviously do anything to help with the opposite problem discussed in Section 12.4, where output vectors appear “random”, and in particular fail to reflect the geometric relationships between the original input vectors. These intuitions are confirmed by Yang et al.’s (2019) signal propagation analysis of RELU networks with BN layers (which we discuss in Section 25.8). In particular, Yang et al. (2019) predict that the distances between the output vectors (generated from different inputs) will converge to a constant as depth increases, and that this leads to an exponential increase in the norm of the gradient. Like with LN layers, placing BN layers after each affine layer makes the network insensitive to the scale of its weight parameters, which can thus correct for badly scaled initial weights. One can perhaps also view BN layers as ensuring that the “average q value” across the mini-batch is 1, although this is an imperfect analogy since BN layers operate on a per-channel basis instead of averaging over channels. Similar to LN layers, BN layers may also have an effect on optimization which is independent from their effect on the network’s initial behavior. Evidence for this includes the fact that various methods which modify networks and their initializations in an attempt to eliminate the need for BN (such as those we’ve previously discussed) fail to achieve the same optimization performance under SGD, except perhaps at small mini-batches sizes (where classical optimization considerations like curvature matter a lot less, as argued in Zhang et al. (2019a)). In support of the optimization-effect hypothesis, Santurkar et al. (2018) argue that BN layers make a network’s output a smoother function of its parameters, and that this helps improves the performance of gradient descent. Li and Arora (2019) argue that gradient descent applied to networks with BN layers behaves similarly to gradient descent applied to a normalizer-free network with a decaying learning rate, thus allowing gradient descent with a constant learning rate to converge in the stochastic setting (where it otherwise might not). Finally, Grosse (2021) argues that placing a BN layer after an affine layer will make the network invariant to scaling and shifting of ’s input, and that this leads to a curvature matrix for ’s parameters which is better conditioned.
Part V Experiments and conclusions
In this section we will describe and justify the setup we will use in our experiments, which will depart somewhat from common practice.
The benchmark training problem we use in all of our experiments is image classification, on either the Imagenet (Deng et al., 2009) and CIFAR-10 (Krizhevsky and Hinton, 2009) datasets. The training objective is the average loss over the training set, with the loss given by the cross-entropy error between network’s output (interpreted as “logits” of a softmax) and the dataset labels. We also measure top-1 classification accuracy, and report this instead of the loss in our plots due to its higher interpretability. For Imagenet, we use an image preprocessing and random augmentation pipeline similar to the one from Szegedy et al. (2015) to obtain images of size . The training set is obtained from the standard Imagenet training set, minus the last 10000 cases (which is used as a new validation set), and the test set is obtained from the usual Imagenet validation set. Training accuracy is reported using the examples actually used during training, which are subject to random augmentation. Test accuracy is meanwhile reported using examples from the test set without random augmentation. For CIFAR-10, we apply the standard preprocessing consisting of mean subtraction and normalization of each color channel. The training and test sets are their standard versions. For both datasets we apply Per-Location Normalization (as described in Section 10.2) as a final stage of processing before feeding the inputs to the network. This is done for all approaches and experiments unless stated otherwise, in the interest of fairness.
2 Focusing on optimization speed
While we will report test set accuracy in many of our experiments, our primary focus will be on optimization speed, as measured using training accuracy. Moreover, the decisions we make while designing out experiments will be in the interest of obtaining the cleanest and fairest comparison for optimization speed, and we will tune various components (like the learning rate schedule, regularization, etc) with this in mind. In this subsection we will explain our rationale for this decision. The current standard approach to deep learning is to train normalized residual architectures with RELU nonlinearities, such as ResNets or Transformers, with basic optimizers like SGD or Adam. Alternative approaches (such as normalization-free networks using Fix-up (Zhang et al., 2019c), standard deep convolutional networks initialized with EOC, or DKS) can underperform the standard approach in one of two basic ways. First, they can yield networks whose training plateaus earlier, resulting in underfitting, or whose training is just much slower overall. And second, they can yield a worse inductive bias for typical training problems (like Imagenet classification), resulting in increased overfitting. In our initial experiments we found that while alternative deep learning approaches are typically affected by both of these problems, slower training is by far the more significant one, particularly for networks without skip connections. Moreover, the resulting underfitting problem (given a finite optimization step budget) led to a commensurate degradation in test set performance. (These findings echo those of He et al. (2016a) and Szegedy et al. (2017), who observed that the main benefit of adding skip connections was faster optimization.) Thus, by focusing on training speed, we are isolating what is the more serious problem currently affecting alternative methods, and the one which arguably should be addressed before attempting to close the generalization gap. We believe that the increased generalization gap we observed on Imagenet for alternative approaches such as ours is small enough that it can be overcome through the use of additional regularization strategies, dataset augmentation, scheduling of the optimizer hyperparameters, architectural tweaks, etc. We will leave this to future work. This position is echoed by Zhang et al. (2019c), and was arguably validated in the recent work of Brock et al. (2021) on “Normalizer-Free Networks”, which used a combination of these techniques to close the generalization gap for one such alternative approach. It is also our view that overfitting may become less of a concern as the machine learning community moves beyond supervised benchmark problems like Imagenet classification, and towards giant/streaming datasets and unsupervised methods.
3 Network architectures and regularization
In our experiments we train standard and modified ResNets and Wide-ResNet models for Imagenet and CIFAR-10 image classification. (A detailed description of all the relevant architectures is given in Section 23.) We will place particular emphasis on “ablated” versions of ResNets, where Batch Normalization (BN) layers and/or the skip connections are removed, leaving everything else unchanged. The motivation for doing this is that we want to facilitate the fairest possible comparison to the standard deep learning approach. In particular, since we are focused mostly on optimization speed in our comparisons, we want to use models that are provably no more powerful than standard ResNets in terms of the class of functions they can express, so that the fundamental data fitting problem doesn’t become any “easier”. In the interest of making our experiments fair we also didn’t include the standard L2 regularization that is often used when training ResNets. This is because the effect of L2 regularization on the effective capacity of a model is highly dependent on the model’s parameterization, and this will vary significantly across the different approaches we consider. For example, due to the way BN layers are invariant to scalar multiplication of their inputs, one can rescale the weights of any affine layer that precedes a BN layer without changing the overall output of the network. Thus, networks with BN layers can effectively “cheat” the L2 regularization penalty in a way that networks without BN layers cannot. In our experiments we found that the removal of L2 regularization did have a small but still significant effect on the test set performance of standard ResNets, which is reflected in our reported results. Note that our purpose in experimenting with these modified ResNets is not to show they are a good replacement for standard ResNets in practice. Rather, our purpose is to determine the extent to which we can replace the ingredients of the standard deep learning approach with various alternatives (that preserve the model class), while retaining its fast training capabilities. If we were primarily interested in maximizing test set performance in our evaluations, then we would be free to design a network architecture best suited to DKS, and to include whatever regularization scheme we found to be most effective. And while this does seem like an interesting direction to explore, as has been done in the context of other alternative approaches to deep learning (Brock et al., 2021), it is beyond the scope of the present work.
4 Automatic learning rate schedules with Fire PBT
Achieving a near optimal rate of convergence for standard ResNet training with SGD seems to require a carefully designed learning rate schedule (and not just a fixed value), especially for more difficult datasets like Imagenet. Through extensive and costly trial and error, the community has produced learning rate schedules which seem to work well on certain standard problems, such as Imagenet classification with ResNets. These typically involve a quick “warm-up” of the learning rate from a moderate starting value to a larger one, followed by a decay or step-wise descent towards zero. In our experiments we consider a large variety of approaches for training deep networks, most of which depart from the standard one along directions such as architectural choices, optimizers, initialization, etc. There is no reason to think that a learning rate schedule tuned for standard ResNet training with SGD should perform well for all such approaches, and this was borne out in our initial experiments. (By contrast, it seemed like the momentum hyperparameter was much less important.) Thus, in order to conduct fair experiments, which are minimally confounded by hyperparameter tuning, we need a way of determining a near-optimal learning rate schedule for each approach. And this should ideally be done in an automatic way in order to reduce the role of experimenter bias. Recently, Dalibard and Jaderberg (2021) proposed an alternative version of Population Based Training (PBT, Jaderberg et al., 2017) called FIRE PBT, which is designed specifically for the dynamic adjustment of optimizer hyperparameters. Like many other methods for automatically tuning the learning rate, standard PBT falls into the trap of being too greedy, and tends to lower the learning rate too quickly for the sake of short-term improvements in the loss (Wu et al., 2018). FIRE PBT is designed to tackle this issue, using a strategy which we will now briefly explain. Both PBT and FIRE PBT work by having many workers independently train neural networks, each with their own values for the hyperparameters. Both methods also associate a fitness to each of their workers which guides an evolutionary procedure. In PBT, this fitness is simply the current value of the objective function (which can be defined on the training or test sets). In FIRE PBT these fitnesses are altered in order to promote population members which may have a worse objective but are promising in other ways. In particular, a separate class of workers, called evaluators, periodically copy the model parameters of other workers, change the hyperparameters (e.g. decay the learning rate), and measure the rate at which the objective function improves while training with the new hyperparameters. The higher the rate of improvement as measured by the evaluator, the higher the fitness FIRE PBT will associate to the original worker whose model parameters were copied. This approach encourages workers to use “non-greedy” hyperparameters (such as high learning rates), if it is shown that doing so leads to better performance after training with different hyperparameters (such as lower learning rates) in the long run. In their experiments, Dalibard and Jaderberg (2021) showed that FIRE PBT worked very well at automatically generating learning rate schedules on the fly for standard ResNet training with SGD, matching or exceeding the performance of the previously mentioned community-tuned schedules. In our initial experiments we found that this capability carried over nicely to non-standard deep learning approaches as well, and so we decided to use it in all of our subsequent experiments. We now discuss the technical settings related to our use of FIRE PBT. We follow the presentation of Dalibard and Jaderberg (2021). Each experiment uses 36 workers. We divide them into three sub-populations , each of size 8, and the evaluator set which includes the remaining 12 workers. We train for a maximum of 200,000 steps when training on ImageNet and for a maximum of 25,000 when training on CIFAR-10. Hyperparameters We optimise the learning rate hyperparameter. When using SGD or Adam, the learning rate is initially sampled log-uniformly in the range . When using K-FAC, we instead use the range . Objective function We evaluate the current model by reporting the current negated training loss. Ready A member of the population is deemed ready to exploit and explore every 500 steps when training on ImageNet and 50 steps when training on CIFAR-10. Exploit We use a truncation selector: If a population member has a fitness in the bottom 25% of the population, it copies the neural network weights and hyperparameters of a random member in the top 25% of the population. Explore We multiply the learning rate by a value sampled at uniform random between the following two value: [0.8, 1.25]. We further set the FIRE PBT hyperparameter of min_steps_before_eval to 5000 and 10000 for and respectively when training on ImageNet, and to 250 and 500 when training on CIFAR-10. The training curves plotted in our results section use the values recorded by the sequence of workers that led to the best eventual objective value (negative loss). In Appendix L we plot the learning rate schedule found by FIRE PBT for some of our main experiments. We note that apart from some small fluctuations, these schedules are fairly simple and natural looking, and typically involve an initial rapid increase in the learning rate, followed by a gradual decay. Thus, we don’t believe that the qualitative nature of our results is highly dependent on our use of FIRE PBT versus a simpler approach for learning rate tuning.
5 Optimizers
In our experiments we used SGD (with momentum), K-FAC (Martens and Grosse, 2015), Adam (Kingma and Ba, 2014), and Shampoo (Gupta et al., 2018; Anil et al., 2020) as optimizers, with the majority just using SGD and K-FAC. Our motivation for considering stronger optimizers is that alternative deep learning approaches such as ours seem to benefit substantially from using them. For all optimizers we used a momentum parameter of 0.9, and adjusted the learning rate dynamically throughout training using FIRE PBT. For Imagenet experiments we used a batch size of 512 with all optimizers, and for CIFAR-10 we used a batch size of 1024. For Adam we used a value of for the “” parameter, which performed slightly better than the default value of . For K-FAC, we used a 0.99 exponential decay of the curvature matrix, and computed its inverse every 50 iterations. We initialized K-FAC’s damping parameter to , and exponentially decayed it at the rate 0.98 every 50 iterations to a minimum value of . Finally, we enforced a maximum norm of on all updates, with the norm computed using K-FAC’s approximate curvature matrix (as in Ba et al. (2017)). For Shampoo we used an epsilon parameter of and an exponential decay factor if 0.99 for the second moments. In order to achieve optimization performance which was competitive with K-FAC, we used an “exponent multiplier” of 3 (which increases the exponent of all factors of the preconditioner by a factor of 3, with 1 being the default value), and enabled “grafting” (which uses Adagrad (Duchi et al., 2011) to compute the magnitude of the update for each parameter tensor, and the usual Shampoo formula to compute its direction).
6 Hardware and implementation details
All of our experiments were implemented in TensorFlow (Abadi et al., 2015). Each of the 36 workers used by FIRE PBT ran on an 16 chip 32 core Cloud TPU v3 Pod (Google,, 2018). For multi-core TPU Pods, each core ran a “replica” of the entire gradient computation on its assigned subset of the training mini-batch, with gradients and other key optimization quantities being averaged across the cores to simulate a single core computation. As long as training cases are independent of each other in the forward pass, this simulation is exact. However, this independence is slightly violated for networks with BN layers, and the resulting simulation is thus imperfect. Handling BN in this way in the multi-core setting has nonetheless become standard practice, and is even thought to be beneficial as it increases the “noise” originating from BN layers, which is thought to have a regularizing effect.
Experimental results
In this section we will present our main experimental results as a series of plots of training/test top-1 accuracy vs iteration number, with some discussion. Most of our experiments will use the standard RELU, tanh, and softplus activation functions, the latter of which is a smooth analogue of the RELU function defined by . \localtableofcontents
We use a value of in all experiments for DKS’s global slope bound parameter , unless stated otherwise. Except for DKS and the Looks Linear method, whenever using RELU activation functions we multiply the network’s initial weights by . This has become standard practice in the literature following He et al. (2015), and can be interpreted as making the local Q map of combined RELU layers equal to the identity. We don’t do this for DKS or the Looks Linear method since those methods achieve identity local Q maps through other means. Unless otherwise indicated, all result will be given for a skip connection-free BN-free modified ResNet-101 architecture trained on Imagenet. “Standard ResNet” will refer to a standard unmodified ResNet with RELU activation functions, initialized with the standard Gaussian Fan-in initialization (with a multiplier). For networks trained on CIFAR-10 we will use a modified Wide-Resnet with 250 layers and a width multiplier of 2.
2 DKS with skip-free nets vs standard baselines
In this subsection we present our main results, in which we compare DKS networks without skip connections or BN layers, to both standard ResNets, and various “ablated” ResNets that are missing skip connections or BN layers (or both). From this first plot we can see that, with K-FAC, DKS enables skip-free BN-free networks to train as fast as a standard ResNet on Imagenet, which is the first time this has been demonstrated to the best of our knowledge. Meanwhile, the ablated ResNets exhibit significantly slower optimization or underfitting. We also see that DKS underperforms for RELU compared to other activation functions, perhaps for the reasons discussed in Section 18.6. The story is somewhat different for SGD training. With SGD and no skip connections, DKS networks fail to match the training speed of standard ResNets, although they still outperform the ablated ResNets. Interestingly, RELUs give the same performance with DKS as the other activation functions do in this setting. For test set performance with K-FAC training we observe increased overfitting with DKS compared to standard ResNets, resulting in an overall lower test accuracy. Notably however, test accuracy is still higher than for the ablated ResNets. Once again the story is somewhat similar for SGD training, although with a larger performance gap vs standard ResNets due to the additional effect of underfitting from using SGD (without skip connections) instead of K-FAC. Note that the test error numbers for standard ResNet training with SGD are a few percentage points the commonly reported values. This is for a number of reasons, including the fact that we don’t include L2 regularization (as discussed in Section 27.3), that we configured FIRE PBT to maximize training speed and not test set performance, and that we use PLN to process the data. (Because these things affect DKS as well, we believe the comparison to still be fair.) The remaining results in this subsection are analogous to the previous ones, but use CIFAR-10 with modified/ablated Wide-ResNets models. The observations from these results are similar, although we note that the performance gap between the DKS networks and ablated Wide-ResNets is considerably larger, likely due to the higher depth (250 vs 100) used in these experiments.
3 DKS with and without skip connections
In this subsection we compare the performance, with and without skip connections, of BN-free networks constructed with DKS. We use weights of and for the residual and shortcut branches respectively (so that the all sums in the network are normalized as per Section 21.2). The value was selected from several candidate options in order to maximize training speed, as shown in Appendix M.1. When using K-FAC we see that the training speed remains the same whether or not we use skip connections, except in the case of RELU activation functions. For RELUs, skip connections seem to help significantly, closing the performance gap with the other activation functions. With SGD the story is different, and skip connections allow us to match the training speed of standard ResNets with DKS, at least when using softplus or RELU activation functions. For K-FAC, the improvement to test set accuracy from using skip connections with DKS appears to be minimal, with the notable exception of RELU networks (where the improvement is likely due to improved fitting/optimization, as opposed to improved generalization). By contrast, in the context of SGD training we see a significant improvement to the test set accuracy from using skip connections with DKS. Although again, this is likely due to improved fitting enabled by the use of skip connections with SGD, rather than improved generalization.
4 DKS with different activation functions
In this subsection we compare performance of DKS with twelve different activation functions. In addition to certain well-known mathematical functions, we also include SELU (Klambauer et al., 2017), Softsign (Bergstra et al., 2009), Swish (Ramachandran et al., 2017; Elfwing et al., 2018), Elu (Clevert et al., 2016), and BentId (defined by ). For K-FAC we see fairly similar training speeds for each of the twelve activation functions, with RELU being the notable outlier. For SGD, there is a larger deviation in performance observed for the different options, and RELU is notably no longer an outlier. Results for test set accuracy were qualitatively very similar, and so we won’t report them here.
5 Comparisons to other approaches
In this subsection we compare DKS to various other approaches for initializing and constructing neural networks. We will focus primarily on skip connection-free BN-free networks, except when comparing to Fix-up (which requires the use of skip connections). We will omit test set accuracy in these comparisons, as we found that it gave qualitatively similar results to training accuracy. (This is likely because nearly all competing methods yield significant underfitting for skip-free BN-free networks, which overwhelms any possible advantage they might have in terms of generalization.)
The Gaussian fan-in initialization (aka “variance scaling initialization” or “Lecun initialization”), which is discussed in Section 26.1, is the default initialization method used in many modern neural network frameworks, and is the first method we compare to. From these results we can see that DKS significantly outperforms this canonical approach, whose poor performance in this setting is not surprising given the analysis of Section 12. Note that it is common in practice to use a truncated Gaussian distribution or uniform distribution to sample the weights in a fan-in initialization, instead of the usual Gaussian distribution. When used with an appropriate rescaling term, these distributions produce weights with the same variance as the standard Gaussian distribution, although they won’t necessarily give rise to the same approximate kernel functions. We ran additional experiments using these distributions, and found that they gave similar results to those presented above.
Glorot initialization (aka “Xavier initialization”) is a commonly used modification of the Gaussian fan-in initialization which we discuss in Section 26.1. As with the Gaussian fan-in method, it is also often used with a truncated Gaussian or uniform distribution, the latter of which we will present results for. (We also performed experiments using truncated and non-truncated Gaussian distributions for the weights, which yielded similar findings.) From these results we can see that the Glorot approach is significantly outperformed by DKS, and completely fails to produce a trainable network for both the RELU and softplus activation functions.
The LSUV and WLI approaches, which are discussed in Section 26.2, represent the first generation of methods which attempt to capture the benefits of Batch Normalization through initialization. They are fairly similar in their implementation, which is why we consider them together here. From these plots we can see that these methods outperform simple initializations schemes like fan-in and Glorot, but are still significantly outperformed by DKS.
Self-normalizing neural networks (which we discuss in Section 26.3) use SELU activation functions, together with a standard Gaussian fan-in initialization, to achieve certain conditions under variance propagation which are essentially equivalent to two of the four conditions enforced by DKS. From these results we see that DKS applied to a softplus network matches or exceeds the optimization performance of a self-normalizing network. DKS also improves the performance of a SELU network optimized with K-FAC, although slightly degrades it for SGD.
The Looks Linear method, which is discussed in Section 26.5, is an approach for constructing and initializing RELU networks which makes them behave like perfectly linear functions at initialization time, without the use of skip connections. The method is somewhat difficult to fairly compare to other ones, as it involves doubling the channel dimension of each layer, while using a form of weight sharing which makes the resulting network less expressive than a standard one of the same dimensions. Our imperfect solution to this problem is to use the original dimensions when constructing networks with DKS, which will disadvantage DKS in the comparison. We had some trouble optimizing the networks constructed with the Looks Linear method. K-FAC would quickly diverge for all the hyperparameter settings we tried, perhaps because it broke the delicate symmetry of the initial weights too quickly, leading to extreme nonlinear behavior. We had more luck with Adam and SGD, although we found that it was necessary to threshold the maximum update magnitude at 1 to achieve stable optimization (which is an approach known as “clipping” (Pascanu et al., 2013)). Because we couldn’t get K-FAC to work well with the Looks Linear method, we used it with Adam instead in our first comparison. We note that with Adam, DKS performs similarly to the Looks Linear method, but when used with K-FAC, DKS significantly outperforms it. For SGD both methods seem to perform similarly, and notably better than both the fan-in/Glorot initializations, and also the LSUV/WLI methods. We also conducted experiments with CIFAR-10, which yielded similar results. These are given below without commentary.
The Edge of Chaos (EOC) method (described in detail in Section 26.4) is the closest approach to ours in the existing literature, and the one which directly inspired it. The version in Xiao et al. (2018), which we will use here, involves two ingredients: choosing variances for the weight and bias distributions so that for each local C map , and using the Delta Orthogonal initialization for the weights (which is rescaled to achieve the target variance). A clean comparison to EOC is somewhat difficult, as it is not fully specified. In particular, for most activation functions there are infinitely many combinations of the two variances which achieve . And for the RELU activation function, the condition holds for any weight variance (given zero bias variance), so that the method reduces to an Orthogonal Delta initialization. Xiao et al. (2018) focused their experiments on tanh networks, and following their advice we will take the variance of the weights and biases to be and (respectively) for tanh nets, where is the input channel dimension for the given layer. We will also consider RELU networks, with a weight variance of and a bias variance of 0. From these results we see that DKS significantly outperforms EOC in terms of optimization speed (for both Imagenet and CIFAR-10), which in turn outperforms the simple Fan-in initialization method.
Fix-up, which we briefly discuss in Section 26.6, is a recent method for constructing and initializing networks with residual connections which is designed to eliminate the need for normalization layers. It involves initializing the weights of the final convolutional layer in each residual block to zero (so that the residual blocks behave like identity functions at initialization), using a special formula for the variance of the weights distribution, as well as introducing learnable scalar multiplication and bias operations throughout the network. We were not able to get K-FAC to work well with Fix-up. This might have been due to a bad interaction with K-FAC and the extra parameters introduced by Fix-up (as K-FAC is designed specifically for the standard neural network parameters). Another possible explanation is that, like networks created with the Looks Linear method, the larger steps taken by K-FAC cause Fix-up networks to transition too quickly to extreme nonlinear behavior (after being essentially linear at initialization). As Fix-up requires the use a skip connections, for the sake of fairness we compared it to BN-free networks constructed with DKS that also used skip connections. And because we couldn’t get K-FAC to work well with Fix-up, we instead used Adam with Fix-up in our first comparison. (While this may seem unfair, we note that for networks with skip connections, K-FAC and SGD perform similarly, as shown in Subsection 28.3.) From these results we see that Fix-up performs similarly to DKS for RELU activation functions, but falls behind for tanh and softplus.
6 Meta-parameter studies
The influence of various training “meta-parameters” on the optimization and generalization performance of DKS networks is considered in Appendix M. These meta-parameters include the weight on the residual branch when using skip connections, DKS’s parameter, and the choice of optimizer. Our conclusions from these studies are summarized as follows:
When using DKS with skip connections, a weight of on the residual branch works the best overall among several other sensible options, although this is likely to be contingent on details of the architecture (such as depth).
In terms of optimization performance, typically works better than values that are much larger, or much closer to 1, although the difference isn’t very big. In terms of generalization performance, somewhat smaller values (such as 1.1) may work slightly better.
For networks without skip connections, K-FAC is the best optimizer in terms of speed, followed closely by Shampoo. Following that are Adam and then SGD, which both perform significantly worse than Shampoo in this setting. For networks with skip connections, the gap between K-FAC and SGD narrows substantially.
7 Ablations and modifications of DKS
Various ablations and modifications of DKS are considered in Appendix N. The overall conclusion of these studies is that each component of DKS, except perhaps for PLN (assuming reasonably well scaled input data), is required to achieve the highest optimization speed. When considering test error the conclusions are similar but somewhat muted, with the single exception that using weighted mean-pooling layers with K-FAC seems to improve test set performance while degrading training set performance.
Conclusions
In this work we developed Deep Kernel Shaping (DKS), a method for making neural networks easier to train via model class preserving transformations. We showed how our method controls the shape of the network’s initialization-time kernel, by way of our generalized Q/C map analysis, in order to prevent certain common pathologies associated with slow optimization and poor generalization. In our experiments we showed that DKS allows deep networks without skip connections or normalization layers to be trained at similar speeds to ResNets on Imagenet, assuming the use of K-FAC or Shampoo. To the best of our knowledge this is a unprecedented result. We also applied our generalized Q/C map analysis to explain the effectiveness of previously proposed methods for training deep networks, such as skip connections, normalization layers, and popular initialization schemes. By demystifying trainability in deep networks, and disentangling it from model design, we hope that DKS will enable deep networks to reach new heights of performance, flexibility, and ease of use. There is even the potential that DKS may unlock a new class of neural models untrainable with standard tools like normalization layers and skip connections, possibly when used in combination with strong optimizers like K-FAC or Shampoo. Finally, because of their sensitivity to the strength of the optimizer, deep skip connection-free networks constructed with DKS have the potential to serve as new benchmark problem for neural network optimizers. This should be a welcome development to the area, as ResNets are often still used for benchmarking optimizers, despite the fact that it is known to be impossible to significantly outperform well-tuned SGD when training them at small/medium batch sizes (Zhang et al., 2019a).
Limitations and future directions
We end by discussing some limitations of DKS, along with possible ways to overcome them in future work.
DKS currently doesn’t support layers with multiplicative units, such as the self-attention layers in Transformers (Vaswani et al., 2017). This is because we don’t have a kernel approximation for such layers that would yield one-dimensional Q/C maps (or something similar). A possible way around this would be to generalize C maps to higher dimensional inputs, and develop new theory along the lines of Section 13 to control their shape. Another possibility would be to find some weight initialization for self-attention layers which would give rise to one dimensional maps.
While DKS supports pooling layers in practice (based on our experiments), our theoretical treatment of these layers in Section 20 is rudimentary and incomplete. Moreover, mean-pooling layers aren’t really supported at all within our framework, since they make it impossible to achieve uniform q values. And while weighted mean pooling layers can serve as a reasonable replacement (as discussed in Section 20.1.2), their Q/C map interpretation, and the quality of their kernel function approximations, are both somewhat dubious. One possible way to improve this situation would be to develop better replacements for pooling layers that are compatible with DKS. Another would be to extend our analysis to handle non-uniform q values, possibly as part of a generalized higher-dimensional version of Q/C map analysis. The fact that pooling layers seem to work reasonably well with DKS in practice hints that this should be possible.
To match the training speed of standard ResNets on skip connection-free networks using DKS we were required to train with K-FAC or Shampoo. From a practical perspective this is somewhat unsatisfying, as those methods are significantly more complex than SGD, and introduce additional computational overheads (although these can be largely mitigated through various strategies such as those proposed in Martens and Grosse (2015), Ba et al. (2017), and Anil et al. (2020)). An interesting direction for future work would be to try to achieve rapid training of skip connection-free networks with a much simpler optimizer like SGD or Adam (e.g. by modifying DKS somehow), or to explain the importance of stronger optimizers for training such networks. Recent work arguing for the optimality of approximate natural gradient methods like K-FAC in in the NTK setting (Zhang et al., 2019b; Karakida and Osawa, 2020) may be a good starting point for the latter direction.
In our experiments with DKS we consistently observed increased overfitting compared to standard ResNets, resulting in top-1 test set accuracy on ImageNet that was lower by a few percent. This echoes similar observations made in related works such as Zhang et al. (2019c), and could by caused by a number of things, including the loss of noise from BN layers, or a subtle change in the inductive bias of the model. Addressing this remains an important direction for future work.
As discussed in Section 18.6, while RELU activation functions can be used with DKS, one can only enforce three of the four Q/C map conditions. And while RELU networks with DKS perform well in most settings in our experiments despite this limitation, they perform poorly with skip connection-free networks trained using K-FAC (relative to other activation functions). Fortunately, since we have consistently strong performance for other activation functions, including RELU-like ones such as softplus, this arguably isn’t a serious issue. Indeed, the main reason for using RELUs over other activation functions is that they are an important ingredient in the standard recipe for achieving fast and stable training of very deep networks, for which DKS is an alternative.
Recurrent neural networks (RNNs) are currently not supported by DKS. This is due to their sharing of parameters across time steps, which invalidates the kernel approximations that underlie our analysis. However, it is conceivable that a more advanced theory could be used to extend DKS to RNNs, and preliminary experiments we conducted with DKS on RNNs gave positive results, suggesting that it might already work well in practice.
Q/C map analysis is formally justified using kernel function approximations for neural networks at initialization time. The accuracy of these approximations is predicted by bounds such as those reviewed in Section 5.8. Currently, the best known bounds seem to be overly pessimistic, and in order to guarantee reasonable approximation error, require that the width (or channel dimension) grow exponentially with the network’s depth. We conjecture that much stronger bounds exist, although they might require the introduction of additional hypotheses, such as that the networks are constructed using DKS (or something similar).
We would like to thank Alex Botev, Alex Graves, Georg Ostrovski, Guodong Zhang, Ilja Kuzborskij, Koray Kavukcuoglu, Neil Rabinowitz, Soham De, Yann Dauphin, and Yee Whye Teh for their guidance, helpful discussions, and feedback on early drafts. We would also like to thank the entire team at DeepMind for supporting this project.
Part VI Appendix
In addition to approximating the contents of the PKF , which can be thought of as entry-averages (across channels) of element-wise products between pairs of vectors in ’s output feature map, we may sometimes be interested in approximating the entry-averages of such vectors themselves. For a given -dimensional vector with associated q value in ’s output, we have the intuitive approximation
Meanwhile, under the APKF Condition we have that
Comparing diagonal entries of both sides implies
B Mathematical details for Section 8.1.2
In this section we will derive the formula
and the derivative of is taken with respect to . Taking the derivative inside of the expectation we have
C A detailed analysis of C map convergence in deep networks
Unless specified otherwise, will denote , i.e. composed with itself times.
Proof Let . Let us consider . Then
Then . Moreover, , and thus
Proof We use induction. The inequality holds for by assumption. Assume now that for some and that . Then , where the last inequality follows from the fact that is an increasing function.
Assume that an increasing, strictly convex map satisfies and for some . Then for every in .
Proof The function is increasing, so , and it is strictly convex, so for all and . Using , and we obtain .
C.2 Definitions and preliminaries
We will call a C map nontrivial if and only if it is neither a constant function nor the identity function. As we are assuming uniform q values, we have by Section 11 that C maps are positive definite functions, from which various properties immediately follow. The following proposition lists the ones we will use in this section:
A C map satisfies the following properties:
and ,
for is an increasing, convex function,
if is not a constant function it is strictly increasing for ,
if is nontrivial, it is strictly increasing and strictly convex for .
If is a nontrivial -map, then for all there holds .
which contradicts the strict convexity of on $$ (Proposition 31, iv)).
Suppose is a nontrivial C map. Then we have
for every . Moreover, if this inequality becomes an equality for any , then is an odd function.
Proof Because is positive definite, it can be written as
for . We can then decompose this as
where is an even function, and is an odd function. Thus, for
as and . Therefore
This inequality can become an equality for any nonzero if and only if . But if for any such , then , which ends the proof.
For there holds for .
There holds for .
Proof By convexity (Proposition 31), . Combining this and Equation 38 we get
Let be a C map and . Then
Proof The first inequality follows from the fact that is an increasing function on $\mathcal{C}(0)\geq 0-c\in(0,1)\mathcal{C}^{\prime}(-c)\geq 1\mathcal{C}(-c)<-c\mathcal{C}^{\prime}\mathcal{C}^{\prime}(x)>1x\in[-c,1]\mathcal{C}(1)-\mathcal{C}(-c)>1-(-c)\mathcal{C}(1)=1-\mathcal{C}(-c)>c-c>\mathcal{C}(-c)\geq\mathcal{C}(-c)-\mathcal{C}(0)\mathcal{C}(0)\geq 0\frac{\mathcal{C}(-c)-\mathcal{C}(0)}{-c}<1-c>0\mathcal{C}(-c)<1\mathcal{C}^{\prime}\mathcal{C}^{\prime}(x)<1x\in[0,-c]\frac{\mathcal{C}(-c)-\mathcal{C}(0)}{-c}<1$ follows from the Mean Value Theorem.
Let be a nontrivial C map. Then either has no fixed points on , or is an odd function.
Proof By Corollary 35, for . Thus, if , then for all . In the case by Lemma 33 there holds for all . But, by strict concavity, for all . Thus if has a fixed point , then and thus , which (by Lemma 33) implies that is an odd function.
Suppose is a nontrivial C map. Then one of the three following alternatives holds:
The map has precisely one fixed point. This happens if and only if , and the fixed point is .
The map has precisely two fixed points. This happens if and only if and is not an odd function, and the fixed points are and .
The function has precisely fixed points. This happens if and only if is an odd function and the fixed points are , and .
Proof We will treat each of i), ii) and iii) separately.
By Proposition 31 iv), the function is strictly convex on $\mathcal{C}(1)=1\mathcal{C}^{\prime}(1)\leq 1\mathcal{C}(x)>xx\in[0,1)g(x):=\mathcal{C}(x)-xg(1)=0gg^{\prime}(x)=\mathcal{C}^{\prime}(x)-1<0x\in[0,1)\mathcal{C}[-1,0)\mathcal{C}(0)>0\mathcal{C}[-1,0)$.
Assume, that . Consider . There holds , and . Thus is an increasing function in some neighbourhood of (by assumption is an analytic function, hence all of its derivatives must be continuous). Then for sufficiently small. On the other hand, , and thus by continuity of there exists such that . The function is strictly convex on . By Corollary 37, these are the only fixed points of , as we assumed that is not an odd function.
We treat each of the implications in ”if an only if” separately
By definition , so if is an odd function, and . By Proposition 31 iv), the function is strictly convex on $(-1,0)\mathcal{C}(x)=x(-1,0)\cup(0,1)$.
Is a direct consequence of Lemma 33, as .
The point from Corollary 38 ii) has to be an attractor. (In other words, .)
Proof As is positive definite, for positive ’s. There holds thus, by the Mean Value Theorem there exists a point in the interval , such that . By convexity, is an increasing function on , so as , there holds . The following theorems give bounds on the convergence rate of c values under repeated applications of a C map satisfying i) or ii) of Corollary 38.
Let a nontrivial C map satisfy . Then the unique (see Corollary 38) fixed point of is a linear global attractor of the whole set $$, and the following set inequalities holds
Proof By Corollary 32 there holds for . For values of we use Lemma 33. This yields . Thus we can apply Lemma 28 with and combined with Lemma 29 on the whole set $$ and the Theorem follows.
Let be a non-trivial C map satisfying and is not an odd function (i.e. satisfying ii) of Corollary 38), and let be the unique fixed point of in the interval . Then for all
for there holds ,
for there holds ,
for , if there holds , where .
Moreover, we have that , , and are all bounded strictly between 0 and 1.
Proof We are going to use an appropriate linear estimate in each of the three dynamical regimes, and then invoke Lemma 28
Note that by Lemma 30. There holds for , as and is a convex function in $\mathcal{C}\mathcal{C}_{1}(c)\equiv\frac{\mathcal{C}(c_{0})-\mathcal{C}(c^{*})}{c_{0}-c^{*}}(c-c^{*})+c^{*}\mathcal{C}_{2}\equiv\mathcal{C}\mathcal{C}_{1}(c^{*})=\mathcal{C}_{2}(c^{*})=c^{*}\mathcal{C}_{1}(c_{0})=\mathcal{C}_{2}(c_{0})
Note that by Corollary 39. There holds for , as and is a convex function in $\mathcal{C}^{\prime}(c)<\mathcal{C}^{\prime}(c^{*})c\in[0,c^{*}]\mathcal{C}(c^{*})-\mathcal{C}(c)\leq\mathcal{C}^{\prime}(c)(c^{*}-c)\mathcal{C}(c^{*})=c^{*}c^{*}-\mathcal{C}(c)\leq\mathcal{C}^{\prime}(c)(c^{*}-c)\mathcal{C}(c)\geq\mathcal{C}^{\prime}(c)(c-c^{*})+c^{*}\mathcal{C}_{1}\equiv\mathcal{C}\mathcal{C}_{2}(c)\equiv\mathcal{C}^{\prime}(c^{*})(c-c^{*})+c^{*}\mathcal{C}_{2}\mathcal{C}^{\prime}(c_{0})>0\mathcal{C}^{n}_{1}(c_{0})\geq\mathcal{C}^{n}_{2}(c_{0})n=1,2,\ldots\alpha=\mathcal{C}^{\prime}(c^{*})\beta=\left(1-\mathcal{C}^{\prime}(c^{*})\right)c^{*}\frac{\beta}{1-\alpha}=c^{*}\mathcal{C}^{n}_{2}(c_{0})=\left(\mathcal{C}^{\prime}(c^{*})\right)^{n}(c_{0}-c^{*})+c^{*}$, which finishes the proof in this case.
Note, that by Lemma 36. By applying Corollary 34 with we obtain for , and therefore for . Similarly to the previous two regimes, we apply Lemma 29 to and , and we get for all such that . Note, that is an increasing function, as . We apply Lemma 28 to function with and , which yields for every . Thus we have whenever , which finishes the proof in this case.
Let be an odd nontrivial C map. Let . Then is a linear global attractor of the whole set and for any there holds
Proof We only need to prove the inequality for positive -s, due to symmetry and the fact that (the positive semi-axis is invariant under the map ). By Proposition 31, iv) the map is strictly convex for , thus for and the theorem follows from Lemma 28 combined with Lemma 29.
D Mathematical details for Section 13
Given our running assumption of uniform q values, C maps are positive definite functions (as established in Section 11). This means that we can write
for some coefficients , so that the derivative of can similarly be written as
, and
for all .
Observing that (which follows from ) we also have that
Similarly, for the second measure of deviation we have
Finally, if , we can relate the two key quantities and as follows:
The theorem then follows directly from the above inequalities.
D.2 Proof of Proposition 16
It is well known that has a basis known as the Hermite polynomials (e.g. Wikipedia contributors, 2021), which is orthonormal (i.e. for and ), and has many other properties useful properties. Two of which we will make use of here is that is the identity function (i.e. ), and that is constant and equal to 1 (i.e. ). Because form an orthonormal basis of , we can represent in terms of this basis as
Using the bilinearity of inner products, and the fact that , we have
and thus as claimed. Plugging this into Theorem 13 it further follows that
D.3 Proof of Proposition 18
Analogously to , can be written as
Using Equation 23, and the identities from the previous subsection, we have
where is a combined layer with as its activation function. Thus,
D.4 Proof Proposition 19
Because is positive definite we can write it as for some with . Using this we can rewrite as follows:
D.5 Proof of Proposition 20
As in Appendix D.4 we have for some with , and . Let , so that . We observe that for , and thus is a decreasing function for . Let . Since , we have . Using this fact, and that for all , it thus follows that
D.6 Proof of Proposition 22
Let be the C map for each of the subnetworks. Because is positive definite we can write it as for some with . Since , we have that is an attractive fixed point of , and thus . (noting that is true because is positive definite and thus convex on $\mathcal{C}(1)=1C^{\prime}_{f}(1)=\mathcal{C}^{\prime}(1)^{D}\mathcal{C}^{\prime}(1)\leqslant 1-1\mathcal{C}c_{1}\equiv\mathcal{C}(-1)>-1c_{1}c_{1}=1$ we have
and thus . Because for all it thus follows that for odd , and so is an even function and therefore . It remains to consider the case . Let be a subnetwork of consisting of compositions of the subnetworks that define . Under this definition we have . Since we have that as for any so that . Thus by Proposition 20 (and Remark 21 to handle the case ) it follows that (since ). By the chain run we therefore have that as .
D.7 Proof of Proposition 25
Since , we have by definition that
Because is positive definite and thus convex on $c\in$ that
Taking and using using and the above inequality, we thus have
E Mathematical details for Section 20.1.2
where is the number of channels. This is the claimed formula for . It remains to establish the conditions under which will concentrate around its expectation. This is more difficult than in the combined layer case, as the different output channels of are not independent given , meaning that the variance of will not necessarily shrink as grows. Instead, our strategy going forward will be to compute
Thus, by taking it follows that
F Mathematical details for Section 18.5
Taking in the numerator clearly gives . To handle the denominator, we will make use of the “delta method” from statistics (which is derived using a Taylor series argument), which says that if is a continuous with respect to , , and is distributed as , then converges in distribution to as . It thus follows that
and so we conclude that .
G Mathematical details for Section 20.2
In this section we will argue that a max pooling layer approximately preserves its input q values, under the condition the vectors within a given patch of its input feature map always fall into two tight clusters. To simplify the argument, we will assume that these clusters have zero variance (or in other words that there are only two distinct input vectors in each patch), which means that our conclusions will only hold approximately. In addition, we will assume uniform input q values to (which is guaranteed by DKS), and that is directly preceded by a convolutional layer initialized with a Gaussian Delta or Gaussian fan-in scheme. To begin, we observe that for all , and thus we may assume without loss of generality that the input q values are 1. We may also assume without loss of generality that the number of locations in each patch is two, since extra vectors that are duplicates of the first two will not affect the maximum. Consider a single patch in ’s input. For each channel , we denote by x^{(i)}=\left[\begin{array}[]{cc}x_{1}^{(i)}&x_{2}^{(i)}\end{array}\right] the two inputs in said patch for that channel. When conditioned on the input to , we have that the are iid Gaussian vectors with mean zero and variance matrix \Sigma=\left[\begin{array}[]{cc}1&c\\ c&1\end{array}\right] for some . This follows from the fact the ’s are linear combinations of iid mean zero Gaussian random variables (the weights), and that the ’s output q values are 1 by assumption. The dimension-normalized squared norm of ’s output vector associated with this patch (which is what the corresponding output q value approximates) is given by
where , and where is the number of channels (which is the same both for ’s input and output). So, to argue that ’s output q value is 1, it suffices to show that the mean of is indeed 1, and that it concentrates around this mean as grows. To compute the mean we will make use of the following result from Nadarajah and Kotz (2008):
Suppose and are Gaussian random variables wth means and , variances and , and correlation coefficient . Then we have
where and are the pdf and cdf of the standard normal distribution, and where
In our case we have , , and . Substituting these into the above expression yields
H Mathematical details for Section 24
In this section we will estimate the per-layer NTK matrices for each of the different cases described in Section 24.5. For the sake of simplicity we will argue in a semi-rigorous manner, employing fuzzy notions like “very large”, “small”, “(not) approximately equal”, “not too close”, “reasonably smooth”, etc. Note that in infinite depth limit these fuzzy notions all become precise, and in particular, “approximately equal” becomes “equal”, “reasonably smooth” becomes “smooth”, etc. Note that all of our conclusions from this analysis have been verified by our numerical studies. Let for some pair of inputs and taken from the training set. Note that by assumption we have . Following the notation of Section 24.4, for the -th combined layer of the network we will denote by the subnetwork that maps ’s input to the input of , and by the subnetwork that maps the input of to ’s final output. Suppose that is early in the network, so that is small and is large. This means that is a shallow subnetwork of , and so is well-behaved, while is a deep subnetwork, and so is degenerate. Because is well-behaved it is reasonably smooth (e.g. by Theorem 13), and so (since ). Then since is degenerate, this implies by Proposition 20 that . Meanwhile, we trivially have that (since ). Thus, by Equation 35, and also . Since and are general distinct inputs from the training set we thus have that . And because is large this will be very small when , very large when , and equal to the identity matrix when . Now suppose is a layer later in the network, so that is large and is small. This means that is a deep subnetwork of , and so is degenerate, while is a shallow subnetwork, and so is well-behaved. Because is degenerate we have (as by assumption). Since is well-behaved it is reasonably smooth, and thus . By Equation 35 we therefore have that , and as in the previous case. Since and are distinct general inputs from the training set we thus have that , where denotes the matrix of 1’s and we have used the fact that is a fixed point of to get that . If , then the estimate for simplifies to . Finally, suppose is a layer in the middle of the network, so that both and are large. This means that both and are deep subnetworks of , and so and are degenerate. Because is degenerate we have (as by assumption). Thus, by Equation 35 we have that , and as in the previous case. There are three scenarios to consider. If and , then since by assumption and for all by Proposition 20, it follows that . So in this case we have , which will be very large since and is large. If and , then since as is large and is non-negative and non-decreasing (by Section 11.2) we have , and thus . So in this scenario we have since and is large. If and , then we have , but this doesn’t help us estimate , since the value of may be highly sensitive to the distance of from 1 (because is degenerate and can thus have extreme behavior near ). Since we have that , and that is non-negative and non-decreasing on $C_{g_{i}}(c_{0})K_{i}c^{*}=1\mathcal{C}^{\prime}(1)=1K_{i}\approx I+\alpha_{i}(E-I)0=\alpha_{1}\leqslant\alpha_{2}\leqslant\cdots\leqslant\alpha_{D}=1$ are constants, and we conjecture that this holds in general.
H.2 Estimating the overall NTK matrix
Having computed an estimate of the per-layer NTK matrix in each case for and , it remains to estimate the full NTK matrix for each . When and we have for early layers and middle layers, and for later layers. If is the number of later layers we thus have
When and we have for early and middle layers, and for later layers. Moreover, since is an attractive fixed point of we have . If is the number of later layers we thus have
Finally, when and , we have for early layers, for some for middle layers (which is only a conjectured formula), and for later layers. Or in general, we have for some for all layers. This gives
where (so that ). From our empirical studies of the C map of deep RELU networks we observe that , and for some other example C maps we observe (the latter of which is consistent with the estimate given in Xiao et al. (2020)).
H.3 Proof of Theorem 27
Let and for be defined as in Section 24.4 for the network in question, and let and be two input data vectors. Under DKS we have , , and . By Theorem 13 this implies that
where . As under DKS, it then follows from Equation 35 that
where we have used the general facts that on the second line, and for on the third line.
I Path-weight analysis and its relationship to approximate kernel analysis
Path-weight analysis is a method for analyzing deep fully-connected RELU networks developed in a series works (Balduzzi et al., 2015; Balduzzi, 2016; Balduzzi et al., 2017). It is capable of approximating some of the same quantities computed by Q/C maps. Unlike those methods, it is not based on the “propagation” of anything through the layers of the network, but instead exploits the special structure of RELU networks to decompose their computation in terms of a collection of simple “paths” that are easier to analyze. In this section we will give a quick derivation of path-weight analysis, identify a possible issue with it, and then discuss how its predictions relate to those made by Q/C maps. For simplicity, we will assume that the sub-network has only 1 output unit, and biases of zero. We will start by defining a path as a sequence of (scalar-valued) units chosen from each layer of the network (which includes the input layer). (So for example, we might choose the 4-th unit from the input layer, the 2-nd unit from the next layer, etc). An active path is one where every RELU unit in is active in the sense of having a non-negative input value (i.e. so that the RELU function is in its “linear region”). We then define a path-weight for a path as the product of the weights that connect the units of in ’s graph representation. Given these definitions, its not hard to see that , where is the entry of corresponding to the (single) input unit in , and the sum is taken over all active paths in . Observe that the total number of paths is simply the product of the input dimensions for each fully-connected layer. Thus, if the weights are chosen according to a standard Gaussian fan-in initialization (so that they are iid with mean zero and variance , where is the layer’s input dimension) it follows that
where is the number of weights along . Using the fact that the expected number of active paths is the same starting from any input unit (due to symmetry), we thus have that
for standard feed-forward RELU networks with R nonlinear layers. They then argue that this decorrelation (or “shattering”) of the input gradients will make the network difficult or impossible to successfully train. As we saw in Section 24, Q/C map analysis can be used to compute the NTK matrix (and per-layer NTK matrices for ), which is a matrix of estimates of the inner products between parameter gradients for different training inputs (assuming a squared error loss). For deep normalized RELU networks such as we are in the collapsing case with and (by Section 12.1), and so assuming that is large we have by Section 24.5 that is approximately the identity for early layers, and approximately equal to a matrix of ones for later layers. This predicts that substantial optimization will only occur in early layers, and that this is unlikely to yield any significant generalization. Note that this prediction is somewhat different than the one made by Balduzzi et al. (2017) insofar as the training loss will in fact be minimized given sufficiently wide layers. The per-layer NTK matrix for the first linear layer, denoted , is closely related to the matrix of input gradients estimated above under path-weight analysis. Reassuringly, the predictions agree in the sense that they are both (close to) an identity matrix.
J Analyzing (nearly) standard ResNets using Q/C maps
As discussed in Section 26.6, ResNets represent a very popular solution to the same problem that DKS is aimed at solving: how to construct a very deep network that can be trained with a gradient-based optimizer. It thus worth understanding how the effectiveness of ResNets can be explained within the framework of Q/C maps. Since our analysis framework doesn’t handle BN layers, we can’t apply it directly to standard ResNets. As a compromise, we will instead analyze a ResNet which is modified to use Layer Normalization (LN) layers in their place, as these are handled within our framework, and perform a somewhat similar function. (See Section 19.2 for a discussion of normalization layers.) Note that Transformer models (Vaswani et al., 2017), which also employ a residual structure, already use LN layers in place of BN layers, although BN layers remain the more popular option for convolutional residual networks. In order to apply our Q/C map analysis to convolutional layers, we will assume the use of a Delta initialization (as opposed to a Gaussian fan-in initialization). And to simulate a weight variance of 2 as used in He et al.’s (2015) initialization scheme we will use “normalized RELU” activations, which are obtained from standard RELU activation functions by multiplication of their input (or output) by . The Q map for a normalized RELU layer is the identity function, which can be straightforwardly derived from Equation 7. So, since each (normalized) RELU nonlinear layer is immediately preceded by an LN layer, we have that its output q value will always be 1 (as the output q value of a LN layer is always 1). It thus follows that the output q value of a residual branch will always be 1 (since affine layers preserve q values). Meanwhile, for non-transition blocks (i.e. those with equal input and output channel dimensions), shortcut branches compute the identity function, so that their Q map is the identity. And for transition blocks, shortcut branches consist of an LN layer, a RELU layer, and an affine layer, so that their output q value is always 1. By Equation 24, the output q value of residual block is the sum of the q values for the two branches, which will therefore be 1 plus the block’s input q value for non-transition blocks (mirroring De and Smith’s (2020) variance propagation analysis), and for transition blocks. From these we observations it follows that q values will grow as a sequence with each successive block, until a transition block is encountered, at which point the q value is reset to 2. For the standard values 50, 101, and 152 of the ResNet-V2 “depth” parameter , there is a sequence of residual blocks, 4 of which are transition blocks (which includes the first block). Thus, the sequence of input q values for the blocks are , for some integers (where the input q values to the transition blocks are , , , and ). For we have , for we have , and for we have . Let denote the local C map for a RELU nonlinear layer, which is given in Equation 27, and notably doesn’t depend on the input q value (which will always be the case for positively homogeneous activation functions). Further, let denote the C map for non-transition blocks with input q value , and denote the C map for transition blocks (which doesn’t depend on the input q value). From the above analysis we have
As discussed in Section 19.2, LN layers always output q values of 1. And because their inputs always comes directly from an affine layer in this network (or a sum over these), we have that their C maps are the identity function. Thus, by Equation 25 it follows that
Given these identities (and Equation 27), we can compute and plot for each possible :
While there is some compression of the range of c values as depth increases in ResNets, it is much milder than what we see for standard deep RELU networks (e.g. in the plots of Section 12.1). For example, $[0.94,1]D=152[0.996,1]1/\sqrt{2}$ on both branches. Or, we could add an LN layer to the shortcut branch of every block. In either case, we have q values of 1 throughout the network, which leads to the following redefinition:
Intuitively, this definition places more weight onto the nonlinear contribution than we had previously. Plotting in this scenario for different ’s yields the following:
While less extreme than standard (non-residual) RELU networks, we still see significantly more compression than before, thus reinforcing the importance of growing q values for the trainability of ResNets. It’s worth noting one can achieve the same effect in a network with constant q values equal to 1 (such as the ones constructed with DKS) by careful choice of weights on the sums at the end of each residual block. For example, one can recover the original form of (seen in Equation 42) by using a weight of for the residual branch of the corresponding block (and a weight of on the shortcut branch). Doing this for all blocks exactly recovers the C map of a standard ResNet, as neither nor depend on the q values.
K Empirical evidence for the relationship between Q map derivatives and kernel approximation error
In this section we will provide empirical evidence for the relationship between Q map derivatives and kernel approximation error that we posited in Section 15. To do this, we will examine the effect of changes to the local Q map conditions used in DKS on the accuracy of the predictions made by Q maps for an example network. In particular, we will consider the skip-free BN-free modified ResNet used in our main experiments (from Section 28.2) with the softplus activation function, and a depth parameter of 50. The local Q maps of this network’s combined layers are equal to the same function , and we will consider the effect of using values of , , and for in DKS’s local map conditions. As we still have with this change, and the network’s input q values are still 1 (due to our use of PLN), we thus have constant q values of 1 for all layers. We can empirically estimate the accuracy of this prediction by measuring how much deviates from 1 for location vectors from the network’s feature maps (computed at initialization time). The following plots show these “empirical q values”, averaged across locations and 192 training examples, versus the layer index for which they are computed. Vertical lines indicate the standard deviation.
From these first two plots we can see that average empirical q value remains close to 1 for both the and cases. Although in the latter case we see higher variance, especially for deeper layers. In the next plot, we see that the empirical q values rapidly diverge from 1 when , thus confirming the intuitions given in Section 15.
L Example learning rate schedules from FIRE PBT
In this section we present the learning rate schedules that were found by FIRE PBT for our main Imagenet experiments from Section 28.2.
M Meta-parameter studies
In this section we will experimentally study the effect of various training “meta-parameters” on the optimization and generalization performance of networks constructed with DKS. These will include the weight on the residual branch when using skip connections, DKS’s global slope bound parameter , and the choice of optimizer. Except when otherwise indicated, all experiments will use the same default settings (as stated in Section 28.1) as our main set of experiments.
In this subsection we compare different values for the weights of the residual branches BN-free networks with skip connections constructed using DKS. To satisfy the condition that the branch sums at the end of each residual block are normalized, we set their weights to and for the residual and shortcut branches (respectively). In addition to running experiments using the same weights for all blocks, we also tried setting the weights individually according to the recipe at the end of Appendix J, so as to recover the C map of an (almost) standard ResNet. (Note that this requires a generalized version of the maximal slope functions given for our modified ResNets in Section 23.4, but is otherwise a straightforward change.) From these results we see that the value seems to work best when using SGD. When using K-FAC, the difference in optimization speed between the three largest options is much smaller, and so we will use as the default value for all optimizers. Note that while is (arguably) the best amoung the values we tried for this network, there is no reason to think that this value will be the best choice for other residual architectures (or the same architecture for a different depth parameter).
M.2 DKS with different optimizers
In this subsection we compare the optimization performance of different optimizers on skip-free BN-free networks constructed with DKS. These results show that K-FAC and Shampoo have a large advantage over Adam and SGD in this setting. Moreover, Adam has a small advantage over SGD, and K-FAC has a small advantage over Shampoo. The picture looks different for networks with skip connections, and K-FAC and SGD yield fairly similar optimization speeds for two out of the three activation functions we tried.
M.3 Sweeping ζ𝜁\zeta values
In this subsection we study the influence of the global slope bound on the optimization speed of skip-free BN-free tanh networks constructed with DKS. In particular, we compare the default choice of to various “extreme” values, which are either very large (corresponding to highly nonlinear network behavior), or are very close to 1 (corresponding to very linear behavior). From these results we can see that gives the fastest optimization performance for K-FAC and the fastest short-term optimization performance for SGD. , which corresponds to a very linear network, gives very slow optimization performance, perhaps for the reasons discussed in Section 14.2. Somewhat surprisingly, the choice yields quite respectable (although still suboptimal) performance.
M.4 The influence of ζ𝜁\zeta on generalization
In this subsection we study the influence of the global slope bound on the generalization performance for networks trained with K-FAC. From these results we can see that while the value tends to give faster optimization, slightly lower values are associated with improved generalization.
N Experiments with ablations and modifications of DKS
In this section we consider various ablations and modifications of DKS. The overall conclusion of these studies is that each component of DKS, except perhaps for PLN (assuming reasonably well scaled input data), is required to achieve the highest optimization speed. When considering test error the conclusions are similar but somewhat muted, with the single exception that using weighted mean-pooling layers with K-FAC seems to improve test set performance while degrading training set performance. For plots that contain solid and dotted lines of the same color, solid lines will correspond to the default unmodified version of DKS, while dotted lines will correspond to the ablated/modified version. Except when otherwise indicated, all experiments will use the same default settings (as stated in Section 28.1) as our main set of experiments. We will omit results for test error, except in those cases where it gives qualitatively different results from the training error.
In this subsection we consider replacing the Orthogonal Delta initialization used in DKS with various alternative weight initialization schemes. Note that while the use of Delta weight initializations is required by the Q/C map theory that underlies DKS, in practice one can still try DKS with any other weight initialization scheme. From these results we can see that while the Orthogonal Delta initialization gives the best results for both K-FAC and SGD, the results for the Orthogonal (non-Delta) and Gaussian Delta initializations are very close. One might be tempted to conclude from these findings that the weight initialization is relatively unimportant for skip-free BN-free networks in general. However, as we can see from the following plot, there is a much larger gap in performance between the different options when we don’t use DKS’s activation function transformations:
In this subsection we consider modifying DKS to only enforce the condition that for all subnetworks (which is equivalent to doing the same for all nonlinear layers ). Note that this conditions is roughly analogous to what normalization layers, the “He initialization method” for RELUs, and the LSUV/WLI initializations are trying to achieve. To achieve this condition we use only the output scale parameter (denoted in Section 18) in the transformed activation functions. From these results we see that the condition is clearly not enough by itself to achieve fast optimization in BN-free networks, with or without skip connections. (And as we will see in Section N.5, enforcing this condition by itself may actually do more harm than good.)
In this subsection we consider removing the condition from the set of four conditions that we enforce in DKS. The remaining three conditions are achieved by setting three of the four activation function parameters (defined in Section 18), with the input shift parameter being left out (which is equivalent to taking ). From these results we see that this condition appears to be important in most training scenarios, but not all of them.
In this subsection we consider the effect of minimizing for each subnetwork in DKS instead of enforcing the condition . This is accomplished by minimizing for each nonlinear layer in the network. This modification is motivated by the observation that minimizing should, according to the reasoning of Section 16.2, minimize the total kernel approximation error. From these results we can see that minimizing works overall worse than simply setting . The reasons for this remain unclear.
N.5 Removing activation function transformations completely
In this subsection we consider completely removing the activation function transformations from DKS. What remains is the Delta Orthogonal initialization for the weights (and a zero initialization of the biases), normalized sums between residual and shortcut branches, and the use of PLN. From these results we see that activation function transformations are important in all the scenarios we tested, except when training softplus networks with skip connections using K-FAC. Moreover, they seem to be especially important when training with SGD. Comparing the results here for tanh networks to those given in Subsection N.2, we can see that training speed becomes worse if we enforce the condition by itself (versus enforcing no conditions at all). While potentially counterintuitive, this isn’t actually surprising. Indeed, there is no reason to think that the local C map for an untransformed tanh layer will be more favorable given an input q value of 1 compared to some other value. (In this case, the “other value” is the fixed point of a tanh layer’s local Q map.) Note that for the fully transformed tanh layers generated by DKS this consideration is moot, since any fixed input q value is essentially equivalent to all other choices due to the flexibility afforded by the activation function’s input scale parameter (denoted in Section 18). We also have results for skip-free BN-free networks trained on CIFAR-10, which tell a similar story. These are given below:
N.6 Removing Per-Location Normalization (PLN)
In this subsection we consider using DKS without the Per-Location Normalization (PLN) data pre-processing step described in Section 10.2. From these results we can see that, somewhat surprisingly, the use of PLN actually harms optimization performance, especially for SGD. We speculate about possible reasons for this in Section 10.2. Whatever the reasons, we know that they are contingent on the default properties of the input training data. We can demonstrate this by rerunning the same experiment without PLN for input data that is scaled by a factor of 100 (after the usual Imagenet pre-processing and augmentation). From these results we can see that if the original input data is badly scaled, using PLN will have a positive effect on optimization performance.
N.7 Using equivalent parameters instead of activation transformations
In this subsection we consider the effect of achieving the four conditions of DKS via “equivalent parameters” (as defined in Section 18.3), instead of explicit transformations to the activation functions. This change brings DKS much closer to being a pure initialization approach, but introduces a reparameterization which can have implications for optimization. From these results we can see that optimization performance with K-FAC is mostly unaffected by this change, while with SGD it becomes much worse. This difference in behavior between the two optimizers isn’t surprising, since as discussed in Section 18.3, K-FAC is essentially invariant to the kind of reparameterization being performed here, while SGD is not.
N.8 Replacing max-pooling layers with convolutions
As discussed in Section 20.2, max-pooling layers are not fully compatible with the Q/C map theory underlying DKS, and so we must be cautious when applying DKS to networks containing them. The networks used in our main experiments contain a max-pooling layer (near the beginning), and in this subsection we will justify this decision by considering the effect of replacing that layer with a standard convolutional layer of same kernel size, stride, etc. From these results we see roughly similar optimization performance with and without this replacement, with K-FAC becoming slightly slower, and SGD becoming slightly faster.
N.9 Replacing mean pooling layers with weighted mean-pooling layers
In this subsection we consider the effect using a weighted mean-pooling layer, as defined in Section 20.1.2, in place of the standard mean-pooling layer normally present near the end of our modified ResNet architecture. As discussed in Section 20.1.2, standard mean-pooling layers are not compatible with the theory underlying DKS, while weighted mean-pooling layers are, at least to some extent. From these results we see that optimization performance gets slightly worse with K-FAC, and slightly better with SGD, and that generalization improves slightly as well for both optimizers. Thus, even though we didn’t use weighted mean-pooling layers in our main set of experiments, they are probably worth trying when using DKS.