Random Walk Initialization for Training Very Deep Feedforward Networks

Introduction

Since the early 90s, it has been appreciated that deep neural networks suffer from a vanishing gradient problem (Hochreiter, 1991), (Bengio et al., 1993), (Bengio et al., 1994), (Hochreiter et al., 2001). The term vanishing gradient refers to the fact that in a feedforward network (FFN) the back-propagated error signal typically decreases (or increases) exponentially as a function of the distance from the final layer. This problem is also observed in recurrent networks (RNNs), where the errors are back-propagated in time and the error signal decreases (or increases) exponentially as a function of the distance back in time from the current error. Because of the vanishing gradient, adding many extra layers in FFNs or time points in RNNs does not usually improve performance.

Although it can be applied to both feedforward and recurrent networks, the analysis of the vanishing gradient problem is based on a recurrent architecture (e.g. (Hochreiter, 1991)). In a recurrent network, back-propagation through time involves applying similar matrices repeatedly to compute the error gradient. The outcome of this process depends on whether the magnitudes of the leading eigenvalues of these matrices tend to be greater than or less than oneExcluding highly non-normal matrices.. Eigenvalue magnitudes greater than one produce exponential growth, and less than one produces exponential decay. Only if the magnitude of the leading eigenvalues are tightly constrained can there be a useful “non-vanishing” gradient. Although this fine-tuning can be achieved by appropriate initialization, it will almost surely be lost as the optimization process goes forward.

Interestingly, the analysis is very different for an FFN with randomly initialized matrices at each layer. When the error gradient is computed in a FFN, a different matrix is applied at every level of back-propagation. This small difference can result in a wildly different behavior for the magnitude of the gradient norm for FFNs compared to RNNs. Here we show that correctly initialized FFNs suffer from the vanishing gradient problem in a far less drastic way than previously thought, namely that the magnitude of the gradient scales only as the square root of the depth of the network.

Different approaches to training deep networks (both feedforward and recurrent) have been studied and applied, such as pre-training (Hinton & Salakhutdinov, 2006), better random initial scaling (Glorot & Bengio, 2010),(Sutskever et al., 2013), better optimization methods (Martens, 2010), specific architectures (Krizhevsky et al., 2012), orthogonal initialization (Saxe et al., 2013), etc. Further, the topic of why deep networks are difficult to train is also an area of active research (Glorot & Bengio, 2010), (Pascanu et al., 2012), (Saxe et al., 2013), (Pascanu et al., 2014).

Here, we address the vanishing gradient problem using mathematical analysis and computational experiments that study the training error optimized deep-networks. We analyze the norm of vectors that result from successive applications of random matrices, and we show that the analytical results hold empirically for the back-propagation equations of nonlinear FFNs with hundreds of layers. We present and test a basic heuristic for initializing these networks, a procedure we call Random Walk Initialization because of the random walk of the log of the norms (log-norms) of the back-propagated errors.

Analysis and Proposed Initialization

We focus on feedforward networks of the form

with $\mathbf{h}_{d}$ the vector of hidden activations, $\mathbf{W}_{d}$ the linear transformation, and $\mathbf{b}_{d}$ the biases, all at depth $d$, with $d=0,1,2,\ldots,D$. The function $f$ is an element-wise nonlinearity that we will normalize through the derivative condition $f^{\prime}(0)=1$, and $g$ is a scale factor on the matrices. We assume that the network has $D$ layers and that each layer has width $N$ (i.e. $\mathbf{h}_{d}$ is a length $N$ vector). Further we assume the elements of $\mathbf{W}_{d}$ are initially drawn i.i.d. from a Gaussian distribution with zero mean and variance $1/N$. Otherwise the elements are set to 0. The elements of $\mathbf{b}_{d}$ are initialized to zero. We define $\mathbf{h}_{0}$ to be the inputs and $\mathbf{h}_{D}$ to be the outputs.

We assume that a task is defined for the network by a standard objective function, $E$. Defining $\boldsymbol{\delta}_{d}\equiv\left.\frac{\boldsymbol{\partial}\mathbf{E}}{\boldsymbol{\partial}\mathbf{a}}\right|_{d}$, the corresponding back-propagation equation is

The evolution of the squared magnitude of the gradient vector, $|\boldsymbol{\delta}_{d}|^{2}$, during back-propagation can be written as

where we have defined, for reasons that will become apparent,

The entire evolution of the gradient magnitude across all $D$ layers of the network is then described by

Because the matrices $\mathbf{W}$ are initially random, we can think of the $z$ variables defined in equation (6) as random variables. Then, $Z$, given by equation (7), is proportional to a product of random variables and so, according to the central limit theorem for products of random variables, $Z$ will be approximately log-normal distributed for large $D$. This implies that the distribution for $Z$ is long-tailed. For applications to neural network optimization, we want a procedure that will regularize $Z$ in most cases, resulting in most optimizations making progress, but are willing to tolerate the occasional pathological case, resulting in a failed optimization. This means that we are not interested in catering to the tails of the $Z$ distribution. To avoid issues associated with these tails, we choose instead to consider the logarithm of equation (7),

The sum in this equation means that we can think of $\ln(Z)$ as being the result of a random walk, with step $d$ in the walk given by the random variable $\ln(z_{d})$. The goal of Random Walk Initialization is to chose $g$ to make this walk unbiased. Equivalently, we choose $g$ to make $\ln(Z)$ as close to zero as possible.

2 Calculation of the Optimal g𝑔g Values

Here $z$ is a random variable determined by

Writing $z=\eta/N$, expanding the logarithm in a Taylor series about $z=1$ and using the mean and variance of the distribution $\chi^{2}_{N}$, we find

From equation (10), this implies, to the same degree of approximation, that the optimal $g$ is

The slope of the variance of the random walk of $\ln(Z)$ is given to this level of approximation by

Note that this is inversely proportional to $N$. These expressions are only computed to lowest order in a $1/N$ expansion, but numerical studies indicate that they are reasonably accurate (more accurate than expressions that include order $1/N^{2}$ terms) over the entire range of $N$ values.

We can compute $\langle\ln(z)\rangle$ to leading order in $1/N$ using a Taylor series as above, but expanding around $z=1/2$ in this case, to obtain

However, unlike in the linear case, this expression is not a good approximation over the entire $N$ range. Instead, we computed $\langle\ln(z)\rangle$ and $\langle(\ln(z))^{2}\rangle$ numerically and fit simple analytic expressions to the results to obtain

3 Computational Verification

The random walks that generate $Z$ values according to equation (8) are shown in the top panel of Figure 1 for a linear network (with random vectors back-propagated). In this case, the optimal $g$ value, given by equation (13) was used, producing an unbiased random walk (middle panel of Figure 1). The linear increase in the variance of the random walk across layers is well predicted by variance computed in equation (14).

Results of Training Deep Networks with Random Walk Initialization

2 Experimental Methods

To assess the quality of the training error for deep nonlinear FFNs set up with Random Walk Initialization, we ran experiments on both the MNIST and TIMIT datasets with a standard FFN defined by equations (1-2). In particular we studied the classification problem for both MNIST and TIMIT, using cross-entropy error for multiclass classification, and we studied reconstruction of MNIST digits using auto-encoders, using mean squared error. For the TIMIT study, the input features were 15 frames (+/- 7 frames of context, with $\Delta$ and $\Delta\Delta$). In these studies, we focused exclusively on training error, as the effect of depth on generalization is a different problem (though obviously important) from how one can train deep FFNs in the first place.

The general experimental procedure was to limit the number of parameters, e.g. 4e6 parameters, and distribute them between matrices and biases of each layer in a network. The classification experiments used constant width layers, and for these experiments the actual number of parameters was the first value above the parameter limit, $p_{lim}$, such that a constant integer value of $N$ was possible. Thus as a network got deeper, its layers also became more narrow. For example, for the MNIST dataset, at $p_{lim}=4e6$ , for $D=4$, $N=1228$ and for $D=512$, $N=88$. For the MNIST auto-encoder experiments, $p_{lim}=16e6$, and the code layer was 30 linear units. The size of each layer surrounding this middle encoding layer was chosen by picking a constant increase in layer size such that the total number of parameters was first number above $p_{lim}$ that led to an integral layer width for all layers. For example, at $D=4$, the layer sizes were $[9816\;30\;9816\;784]$, while for $D=128$ the layer sizes were $[576\;567\;558\;...\;48\;39\;30\;38\;48\;...\;558\;567\;576\;784]$. In these auto-encoder studies we used the $\tanh$ nonlinearity, and varied the $g$ parameter per experiment, but not per layer.

Our experiments compared one depth to another so we varied the learning rates quite a bit to ensure fairness for both shallow and deep networks. In particular, we varied the minimal and maximal learning rates per experiment. In essence, we had an exponential learning rate schedule as a function of depth, with the minimal and maximal values of that exponential set as hyper-parameters. More precisely, we denote the maximum depth in an experiment as $D_{max}$ (e.g. if we compared networks with depths $[4\;8\;16\;32\;64\;128]$ in a single experiment, then $D_{max}=128$). Let $\lambda_{in}$ and $\lambda_{out}$ be the learning rate hyper-parameters for the input and output layers, respectively. The exponential learning rate schedule with decay $\tau$ and scale $\alpha$, as a function of depth, took the form

Then for a given network with depth $D$, potentially smaller than $D_{max}$, the learning rates were set for the actual experiment as

A key aspect of this learning rate scheme is that shallower networks are not overly penalized with tiny learning rates in the early layers. This is because the decay starts with layer $D$ getting learning rate $\lambda_{out}$ and goes backwards to the first layer, which gets a learning rate $\lambda_{D_{max}-D}$. This means that for networks more shallow than $D_{max}$, $\lambda_{1}$ could be much larger than $\lambda_{in}$; only if $D=D_{max}$ did $\lambda_{1}=\lambda_{in}$. Some experiments had $\lambda_{in}<\lambda_{out}$, some had $\lambda_{in}>\lambda_{out}$, and we also tested the standard $\lambda_{in}=\lambda_{out}$ (no learning rate schedule as a function of depth for all experiments). For the very deep networks, varying the learning rates as a function of depth was very important, although we do not study it in depth here. Finally, the learning rates in all layers were uniformly decayed by a multiplicative factor of 0.995 at the end of each training epoch.

Beyond setting learning rate schedules, there were no bells and whistles. We trained the networks using standard stochastic gradient descent (SGD) with a minibatch size of 100 for 500 epochs of the full training dataset. We also used gradient clipping, in cases when the gradient became very large, although this was very uncommon. The combination of hyper-parameters: the varied learning rates, depths, and $g$ values resulted in roughly 300-1000 optimizations for each panel displayed in Figure 3 and Figure 4.

3 Performance Results

Discussion

This study revealed a number of points about training very deep networks. First, one should be careful with biases. Throughout our experiments, we initialized the biases to zero, though we always allowed them to be modified. For the most part, use of biases did not hurt the results. However, care must be taken with the learning rates because the optimization may use the biases to quickly match the target mean across examples. If this happens, the careful initialization may be destroyed and forward progress in the optimization will cease. Second, learning rates in very deep networks are very important. As can be seen from Figure 4D (e.g. $D=32$), the exact learning rate scheduling made a huge difference in performance. Third, we suspect that for extremely deep networks (e.g. 1000 layers as in Figure 4C), curvature of the error landscape may be extremely problematic. This means that the network is so sensitive to changes in the first layer that effective optimization of the 1000 layer network with a first-order optimization method is impossible. Indeed, we set $g=1.05$ in Figure 4C precisely to deal with this issue.

Our experimental results show that even though depth did not clearly improve the training error, the initialization scheme was nevertheless effective at allowing training of these very deep networks to go forward. Namely, almost all models with correctly chosen $g$ that were not broken, due to a mismatch of learning rate hyper-parameters to architecture, reached zero or near-zero training classification error or extremely low reconstruction error, regardless of depth. Further research is necessary to determine whether or not more difficult or different tasks can make use of very deep feedforward networks in a way that is useful in applied settings. Regardless, these results show that initializing very deep feedforward networks with Random Walk Initialization, $g$ set according to Figure 2 or as described in the section Calculation of the Optimal $g$ Values, as opposed to $g=1$, is an easily implemented, sensible default initialization.

We thank Quoc Le and Ilya Sutskever for useful discussions.

References