All you need is a good init

Introduction

Deep nets have demonstrated impressive results on a number of computer vision and natural language processing problems. At present, state-of-the-art results in image classification (Simonyan & Zisserman (2015); Szegedy et al. (2015)) and speech recognition (Sercu et al. (2015)), etc., have been achieved with very deep ($\geq 16$ layer) CNNs. Thin deep nets are of particular interest, since they are accurate and at the same inference-time efficient (Romero et al. (2015)).

One of the main obstacles preventing the wide adoption of very deep nets is the absence of a general, repeatable and efficient procedure for their end-to-end training. For example, VGGNet (Simonyan & Zisserman (2015)) was optimized by a four stage procedure that started by training a network with moderate depth, adding progressively more layers. Romero et al. (2015) stated that deep and thin networks are very hard to train by backpropagation if deeper than five layers, especially with uniform initialization.

On the other hand, He et al. (2015) showed that it is possible to train the VGGNet in a single optimization run if the network weights are initialized with a specific ReLU-aware initialization. The He et al. (2015) procedure generalizes to the ReLU non-linearity the idea of filter-size dependent initialization, introduced for the linear case by (Glorot & Bengio (2010)). Batch normalization (Ioffe & Szegedy (2015)), a technique that inserts layers into the the deep net that transform the output for the batch to be zero mean unit variance, has successfully facilitated training of the twenty-two layer GoogLeNet (Szegedy et al. (2015)). However, batch normalization adds a 30% computational overhead to each iteration.

The main contribution of the paper is a proposal of a simple initialization procedure that, in connection with standard stochastic gradient descent (SGD), leads to state-of-the-art thin and very deep neural netsThe code allowing to reproduce the experiments is available at https://github.com/ducha-aiki/LSUVinit. The result highlights the importance of initialization in very deep nets. We review the history of CNN initialization in Section 2, which is followed by a detailed description of the novel initialization method in Section 3. The method is experimentally validated in Section 4.

Initialization in neural networks

After the success of CNNs in IVSRC 2012 (Krizhevsky et al. (2012)), initialization with Gaussian noise with mean equal to zero and standard deviation set to 0.01 and adding bias equal to one for some layers become very popular. But, as mentioned before, it is not possible to train very deep network from scratch with it (Simonyan & Zisserman (2015)). The problem is caused by the activation (and/or) gradient magnitude in final layers (He et al. (2015)). If each layer, not properly initialized, scales input by $k$, the final scale would be $k^{L}$, where $L$ is a number of layers. Values of $k>1$ lead to extremely large values of output layers, $k<1$ leads to a diminishing signal and gradient.

Glorot & Bengio (2010) proposed a formula for estimating the standard deviation on the basis of the number of input and output channels of the layers under assumption of no non-linearity between layers. Despite invalidity of the assumption, Glorot initialization works well in many applications. He et al. (2015) extended this formula to the ReLU (Glorot et al. (2011)) non-linearity and showed its superior performance for ReLU-based nets. Figure 1 shows why scaling is important. Large weights lead to divergence via updates larger than the initial values, small initial weights do not allow the network to learn since the updates are of the order of 0.0001$\%$ per iteration. The optimal scaling for ReLU-net is around 1.4, which is in line with the theoretically derived $\sqrt{2}$ by He et al. (2015). Sussillo & Abbott (2014) proposed the so called Random walk initialization, RWI, which keeps constant the log of the norms of the backpropagated errors. In our experiments, we have not been able to obtain good results with our implementation of RWI, that is why this method is not evaluated in experimental section.

Hinton et al. (2014) and Romero et al. (2015) take another approach to initialization and formulate training as mimicking teacher network predictions (so called knowledge distillation) and internal representations (so called Hints initialization) rather than minimizing the softmax loss.

Srivastava et al. (2015) proposed a LSTM-inspired gating scheme to control information and gradient flow through the network. They trained a 1000-layers MLP network on MNIST. Basically, this kind of networks implicitly learns the depth needed for the given task.

Independently, Saxe et al. (2014) showed that orthonormal matrix initialization works much better for linear networks than Gaussian noise, which is only approximate orthogonal. It also work for networks with non-linearities.

The approach of layer-wise pre-training (Bengio et al. (2007)) which is still useful for multi-layer-perceptron, is not popular for training discriminative convolution networks.

Layer-sequential unit-variance initialization

To the best of our knowledge, there have been no attempts to generalize Glorot & Bengio (2010) formulas to non-linearities other than ReLU, such as tanh, maxout, etc. Also, the formula does not cover max-pooling, local normalization layers Krizhevsky et al. (2012) and other types of layers which influences activations variance. Instead of theoretical derivation for all possible layer types, or doing extensive parameters search as in Figure 1, we propose a data-driven weights initialization.

We thus extend the orthonormal initialization Saxe et al. (2014) to an iterative procedure, described in Algorithm 1. Saxe et al. (2014) could be implemented in two steps. First, fill the weights with Gaussian noise with unit variance. Second, decompose them to orthonormal basis with QR or SVD-decomposition and replace weights with one of the components.

The LSUV process then estimates output variance of each convolution and inner product layer and scales the weight to make variance equal to one. The influence of selected mini-batch size on estimated variance is negligible in wide margins, see Appendix.

The proposed scheme can be viewed as an orthonormal initialization combined with batch normalization performed only on the first mini-batch. The similarity to batch normalization is the unit variance normalization procedure, while initial ortho-normalization of weights matrices efficiently de-correlates layer activations, which is not done in Ioffe & Szegedy (2015). Experiments show that such normalization is sufficient and computationally highly efficient in comparison with full batch normalization.

The LSUV algorithm is summarized in Algorithm 1. The single parameter $\textit{Tol}_{\textit{var}}$ influences convergence of the initialization procedure, not the properties of the trained network. Its value does not noticeably influence the performance in a broad range of 0.01 to 0.1. Because of data variations, it is often not possible to normalize variance with the desired precision. To eliminate the possibility of an infinite loop, we restricted number of trials to $T_{\textit{max}}$. However, in experiments described in paper, the $T_{\textit{max}}$ was never reached. The desired variance was achieved in 1-5 iterations.

We tested a variant LSUV initialization which was normalizing input activations of the each layer instead of output ones. Normalizing the input or output is identical for standard feed-forward nets, but normalizing input is much more complicated for networks with maxout (Goodfellow et al. (2013)) or for networks like GoogLeNet (Szegedy et al. (2015)) which use the output of multiple layers as input. Input normalization brought no improvement of results when tested against the LSUV Algorithm 1,

LSUV was also tested with pre-initialization of weights with Gaussian noise instead of orthonormal matrices. The Gaussian initialization led to small, but consistent, decrease in performance.

Experimental validation

Here we show that very deep and thin nets could be trained in a single stage. Network architectures are exactly as proposed by Romero et al. (2015). The architectures are presented in Table 1.

First, as a ”sanity check”, we performed an experiment on the MNIST dataset (Lecun et al. (1998)). It consists of 60,000 28x28 grayscale images of handwritten digits 0 to 9. We selected the FitNet-MNIST architecture (see Table 1) of Romero et al. (2015) and trained it with the proposed initialization strategy, without data augmentation. Recognition results are shown in Table 2, right block. LSUV outperforms orthonormal initialization and both LSUV and orthonormal outperform Hints initialization Romero et al. (2015). The error rates of the Deeply-Supervised Nets (DSN, Lee et al. (2015)) and maxout networks Goodfellow et al. (2013), the current state-of-art, are provided for reference.

Since the widely cited DSN error rate of 0.39%, the state-of-the-art (until recently) was obtained after replacing the softmax classifier with SVM, we do the same and also observe improved results (line FitNet-LSUV-SVM in Table 2).

2 CIFAR-10/100

We validated the proposed initialization LSUV strategy on the CIFAR-10/100 (Krizhevsky (2009)) dataset. It contains 60,000 32x32 RGB images, which are divided into 10 and 100 classes, respectively.

The FitNets are trained with the stochastic gradient descent with momentum set to 0.9, the initial learning rate set to 0.01 and reduced by a factor of 10 after the 100th, 150th and 200th epoch, finishing at 230th epoch. Srivastava et al. (2015) and Romero et al. (2015) trained their networks for 500 epochs. Of course, training time is a trade-off dependent on the desired accuracy; one could train a slightly less accurate network much faster.

Like in the MNIST experiment, LSUV and orthonormal initialized nets outperformed Hints-trained Fitnets, leading to the new state-of-art when using commonly used augmentation – mirroring and random shifts. The gain on the fine-grained CIFAR-100 is much larger than on CIFAR-10. Also, note that FitNets with LSUV initialization outperform even much larger networks like Large-All-CNN Springenberg et al. (2014) and Fractional Max-pooling Graham (2014a) trained with affine and color dataset augmentation on CIFAR-100. The results of LSUV are virtually identical to the orthonormal initialization.

Analysis of empirical results

For the FitNet-1 architecture, we have not experienced any difficulties training the network with any of the activation functions (ReLU, maxout, tanh), optimizers (SGD, RMSProp) or initialization (Xavier, MSRA, Ortho, LSUV), unlike the uniform initialization used in Romero et al. (2015). The most probable cause is that CNNs tolerate a wide variety of mediocre initialization, only the learning time increases. The differences in the final accuracy between the different initialization methods for the FitNet-1 architecture is rather small and are therefore not presented here.

The FitNet-4 architecture is much more difficult to optimize and thus we focus on it in the experiments presented in this section.

We have explored the initializations with different activation functions in very deep networks. More specifically, ReLU, hyperbolic tangent, sigmoid, maxout and the VLReLU – very leaky ReLU (Graham (2014c)) – a variant of leaky ReLU ( Maas et al. (2013), with a large value of the negative slope 0.333, instead of the originally proposed 0.01) which is popular in Kaggle competitions Dieleman (2015), Graham (2014b)).

Testing was performed on CIFAR-10 and results are in Table 3 and Figure 2. Performance of orthonormal-based methods is superior to the scaled Gaussian-noise approaches for all tested types of activation functions, except tanh. Proposed LSUV strategy outperforms orthonormal initialization by smaller margin, but still consistently (see Table 3). All the methods failed to train sigmoid-based very deep network. Figure 2 shows that LSUV method not only leads to better generalization error, but also converges faster for all tested activation functions, except tanh.

We have also tested how the different initializations work ”out-of-the-box” with the Residual net training He et al. (2015); a residual net won the ILSVRC-2015 challenge. The original paper proposed different implementations of residual learning. We adopted the simplest one, showed in Table 1, FitResNet-4. The output of each even convolutional layer is summed with the output of the previous non-linearity layer and then fed into the next non-linearity. Results are shown in Table 4. LSUV is the only initialization algorithm which leads nets to convergence with all tested non-linearities without any additional tuning, except, again, sigmoid. It is worth nothing that the residual training improves results for ReLU and maxout, but does not help tanh-based network.

2 Comparison to batch normalization (BN)

LSUV procedure could be viewed as batch normalization of layer output done only before the start of training. Therefore, it is natural to compare LSUV against a batch-normalized network, initialized with the standard method.

It is not clear from the paper Ioffe & Szegedy (2015) where to put the batch-normalization layer – before input of each layer as stated in Section 3.1, or before non-linearity, as stated in section 3.2, so we have conducted an experiment with FitNet4 on CIFAR-10 to clarify this.

Results are shown in Table 5. Exact numbers vary from run to run, but in the most cases, batch normalization put after non-linearity performs better.

In the next experiment we compare BN-FitNet4, initialized with Xavier and LSUV-initialized FitNet4. Batch-normalization reduces training time in terms of needed number of iterations, but each iteration becomes slower because of extra computations. The accuracy versus wall-clock-time graphs are shown in Figure 3. LSUV-initialized network is as good as batch-normalized one.

However, we are not claiming that batch normalization can always be replaced by proper initialization, especially in large datasets like ImageNet.

3 Imagenet training

We trained CaffeNet (Jia et al. (2014)) and GoogLeNet (Szegedy et al. (2015)) on the ImageNet-1000 dataset( Russakovsky et al. (2015)) with the original initialization and LSUV. CaffeNet is a variant of AlexNet with the nearly identical performance, where the order of pooling and normalization layers is switched to reduce the memory footprint.

LSUV initialization reduces the starting flat-loss time from 0.5 epochs to 0.05 for CaffeNet, and starts to converge faster, but it is overtaken by a standard CaffeNet at the 30-th epoch (see Figure 4) and its final precision is 1.3% lower. We have no explanation for this empirical phenomenon.

On the contrary, the LSUV-initialized GoogLeNet learns faster than hen then original one and shows better test accuracy all the time – see Figure 5. The final accuracy is 0.680 vs. 0.672 respectively.

4 Timings

A significant part of LSUV initialization is SVD-decomposition of the weight matrices, e.g. for the fc6 layer of CaffeNet, an SVD of a 9216x4096 matrix is required. The computational overhead on top of generating almost instantly the scaled random Gaussian samples is shown in Table 6. In the slowest case – CaffeNet – LSUV initialization takes 3.5 minutes, which is negligible in comparison the training time.

Conclusions

LSUV, layer sequential uniform variance, a simple strategy for weight initialization for deep net learning, is proposed. We have showed that the LSUV initialization, described fully in six lines of pseudocode, is as good as complex learning schemes which need, for instance, auxiliary nets.

The LSUV initialization allows learning of very deep nets via standard SGD, is fast, and leads to (near) state-of-the-art results on MNIST, CIFAR, ImageNet datasets, outperforming the sophisticated systems designed specifically for very deep nets such as FitNets( Romero et al. (2015)) and Highway( Srivastava et al. (2015)). The proposed initialization works well with different activation functions.

Our experiments confirm the finding of Romero et al. (2015) that very thin, thus fast and low in parameters, but deep networks obtain comparable or even better performance than wider, but shallower ones.

The authors were supported by The Czech Science Foundation Project GACR P103/12/G084 and CTU student grant SGS15/155/OHK3/2T/13.

References

Appendix A Technical details

We have selected tanh activation as one, where LSUV initialization shows the worst performance and tested the influence of mini-batch size to training process. Note, that training mini-batch is the same for all initializations, the only difference is mini-batch used for variance estimation. One can see from Table 7 that there is no difference between small or large mini-batch, except extreme cases, where only two sample are used.

A.2 LSUV weight standard deviations in different networks

Tables 8 and 9 show the standard deviations of the filter weights, found by the LSUV procedure and by other initialization schemes.

A.3 Gradients

To check how the activation variance normalization influences the variance of the gradient, we measure the average variance of the gradient at all layers after 10 mini-batches. The variance is close to $10^{-9}$ for all convolutional layers. It is much more stable than for the reference methods, except MSRA; see Table 10.