Natural Neural Networks

Introduction

Deep networks have proven extremely successful across a broad range of applications. While their deep and complex structure affords them a rich modeling capacity, it also creates complex dependencies between the parameters which can make learning difficult via first order stochastic gradient descent (SGD). As long as SGD remains the workhorse of deep learning, our ability to extract high-level representations from data may be hindered by difficult optimization, as evidenced by the boost in performance offered by batch normalization (BN) on the Inception architecture .

Though its adoption remains limited, the natural gradient appears ideally suited to these difficult optimization issues. By following the direction of steepest descent on the probabilistic manifold, the natural gradient can make constant progress over the course of optimization, as measured by the Kullback-Leibler (KL) divergence between consecutive iterates. Utilizing the proper distance measure ensures that the natural gradient is invariant to the parametrization of the model. Unfortunately, its application has been limited due to its high computational cost. Natural gradient descent (NGD) typically requires an estimate of the Fisher Information Matrix (FIM) which is square in the number of parameters, and worse, it requires computing its inverse. Truncated Newton methods can avoid explicitly forming the FIM in memory , but they require an expensive iterative procedure to compute the inverse. Such computations can be wasteful and do not take into account the smooth change of the Fisher during optimization or the highly structured nature of deep models.

Inspired by recent work on model reparametrizations , our approach starts with a simple question: can we devise a neural network architecture whose Fisher is constrained to be identity? This is an important question, as SGD and NGD would be equivalent in the resulting model. The main contribution of this paper is in providing a simple, theoretically justified network reparametrization which approximates via first-order gradient descent, a block-diagonal natural gradient update over layers. Our method is computationally efficient due to the local nature of the reparametrization, based on whitening, and the amortized nature of the algorithm. Our second contribution is in unifying many heuristics commonly used for training neural networks, under the roof of the natural gradient, while highlighting an important connection between model reparametrizations and Mirror Descent . Finally, we showcase the efficiency and the scalability of our method across a broad-range of experiments, scaling our method from standard deep auto-encoders to large convolutional models on ImageNet, trained across multiple GPUs. This is to our knowledge the first-time a (non-diagonal) natural gradient algorithm is scaled to problems of this magnitude.

The Natural Gradient

This section provides the necessary background and derives a particular form of the FIM whose structure will be key to our efficient approximation. While we tailor the development of our method to the classification setting, our approach generalizes to regression and density estimation.

Namely, each iterate $\theta^{(t+1)}$ is the solution to an auxiliary optimization problem, where $\alpha$ controls the distance between consecutive iterates, using an $L_{2}$ distance. In contrast, the natural gradient relies on the KL-divergence between iterates, a more appropriate distance measure for probability distributions. Its metric is determined by the Fisher Information matrix,

i.e. the covariance of the gradients of the model log-probabilities wrt. its parameters. The natural gradient direction is then obtained as $\nabla_{N}=F_{\theta}^{-1}\nabla$. See for a recent overview of the topic.

2 Fisher Information Matrix for MLPs

We start by deriving the precise form of the Fisher for a canonical multi-layer perceptron (MLP) composed of $L$ layers. We consider the following deep network for binary classification, though our approach generalizes to an arbitrary number of output classes.

Let us define $\delta_{i}$ to be the backpropagated gradient through the $i$-th non-linearity. We ignore the off block-diagonal components of the Fisher matrix and focus on the block $F_{W_{i}}$, corresponding to interactions between parameters of layer $i$. This block takes the form:

where $\mathbf{vec}(X)$ is the vectorization function yielding a column vector from the rows of matrix $X$.

Assuming that $\delta_{i}$ and activations $h_{i-1}$ are independent random variables, we can write:

Projected Natural Gradient Descent

The building block of WNN is the following neural layer,

2 Updating the Whitening Coefficients

As the whitened model parameters $\Omega$ evolve during training, so do the statistics $\mu_{i}$ and $\Sigma_{i}$. For our model to remain well conditioned, the whitening coefficients must be updated at regular intervals, while taking care not to interfere with the convergence properties of gradient descent. This can be achieved by coupling updates to $\Phi$ with corresponding updates to $\Omega$ such that the overall function implemented by the MLP remains unchanged, e.g. by preserving the product $V_{i}U_{i-1}$ before and after each update to the whitening coefficients (with an analoguous constraint on the biases).

Unfortunately, while estimating the mean $\mu_{i}$ and $diag(\Sigma_{i})$ could be performed online over a mini-batch of samples as in the recent Batch Normalization scheme , estimating the full covariance matrix will undoubtedly require a larger number of samples. While statistics could be accumulated online via an exponential moving average as in RMSprop or K-FAC , the cost of the eigen-decomposition required for computing the whitening matrix $U_{i}$ remains cubic in the layer size.

In the simplest instantiation of our method, we exploit the smoothness of gradient descent by simply amortizing the cost of these operations over $T$ consecutive updates. SGD updates in the whitened model will be closely aligned to NGD immediately following the reparametrization. The quality of this approximation will degrade over time, until the subsequent reparametrization. The resulting algorithm is shown in the pseudo-code of Algorithm 1. We can improve upon this basic amortization scheme by including a diagonal scaling of $U_{i}$ based on the standard deviation of layer $i$ activations, after each gradient update, thus mimicking the effect of a diagonal natural gradient method. For this update to be valid, this enhanced version of the method, denoted PRONG+, scales the rows of $V_{i}$ accordingly so as to preserve the feed-forward computation of the network. This can be implemented by combining PRONG with batch normalization.

3 Duality and Mirror Descent

There is an inherent duality between the parameters $\Omega$ of our whitened neural layer and the parameters $\theta$ of a canonical model. Indeed, there exist linear projections $P_{\Phi}(\theta)$ and $P_{\Phi}^{-1}(\Omega)$, which map from canonical parameters $\theta$ to whitened parameters $\Omega$, and vice-versa. $P_{\phi}(\theta)$ corresponds to line 10 of Algorithm 1, while $P^{-1}_{\Phi}(\Omega)$ corresponds to line 7. This duality between $\theta$ and $\Omega$ reveals a close connection between PRONG and Mirror Descent .

4 Related Work

This work extends the recent contributions of in formalizing many commonly used heuristics for training MLPs: the importance of zero-mean activations and gradients , as well as the importance of normalized variances in the forward and backward passes . More recently, Vatanen et al. extended their previous work by introducing a multiplicative constant $\gamma_{i}$ to the centered non-linearity. In contrast, we introduce a full whitening matrix $U_{i}$ and focus on whitening the feedforward network activations, instead of normalizing a geometric mean over units and gradient variances.

The recently introduced batch normalization (BN) scheme quite closely resembles a diagonal version of PRONG, the main difference being that BN normalizes the variance of activations before the non-linearity, as opposed to normalizing the latent activations by looking at the full covariance. Furthermore, BN implements normalization by modifying the feed-forward computations thus requiring the method to backpropagate through the normalization operator. A diagonal version of PRONG also bares an interesting resemblance to RMSprop , in that both normalization terms involve the square root of the FIM. An important distinction however is that PRONG applies this update in the whitened parameter space, thus preserving the natural gradient interpretation.

K-FAC is also closely related to PRONG and was developed concurrently to our method. In one of its implementations, it targets the same block diagonal as PRONG while also exploiting the low rank structure of these blocks for efficiency, reminiscent of TONGA. Their method however operates online via low-rank updates to each block, similar to the preconditioning used in the Kaldi speech recognition toolkit . This is in contrast to our approach based on amortization. They also consider the covariance of the backpropagated gradients $\delta_{i}$, while PRONG only looks at the covariance of activations $h_{i}$. K-FAC further proposes a tri-diagonal variant which decorrelates gradients across neighboring layers, though resulting in a more complex algorithm.

A similar algorithm to PRONG was later found in , where it appeared simply as a thought experiment, but with no amortization or recourse for efficiently computing $F$.

Experiments

We begin with a set of diagnostic experiments which highlight the effectiveness of our method at improving conditioning. We also illustrate the impact of the hyper-parameters $T$ and $\epsilon$, controlling the frequency of the reparametrization and the size of the trust region. Section 4.2 evaluates PRONG on unsupervised learning problems, where models are both deep and fully connected. Section 4.3 then moves onto large convolutional models for image classification.

Conditioning. To provide a better understanding of the approximation made by PRONG, we train a small 3-layer MLP with tanh non-linearities, on a downsampled version of MNIST (10x10) . The model size was chosen in order for the full Fisher to be tractable. Fig. 2(a-b) shows the FIM of the middle hidden layers before and after whitening the model activations (we took the absolute value of the entries to improve visibility). Fig. 2c depicts the evolution of the condition number of the FIM during training, measured as a percentage of its initial value (before the first whitening reparametrization in the case of PRONG). We present such curves for SGD, RMSprop and PRONG. The results clearly show that the reparametrization performed by PRONG improves conditioning (reduction of more than 95%). These observations confirm our initial assumption, namely that we can improve conditioning of the block diagonal Fisher by whitening activations alone.

Sensitivity of Hyper-Parameters. Figures 3a- 3b highlight the effect of the eigenvalue regularization term $\epsilon$ and the reparametrization interval $T$. The experiments were performed on the best performing auto-encoder of Section 4.2 on the MNIST dataset. Figures 3a- 3b plot the reconstruction error on the training set for various values of $\epsilon$ and $T$. As $\epsilon$ determines a maximum multiplier on the learning rate, learning becomes extremely sensitive when this learning rate is highUnstable combinations of learning rates and $\epsilon$ are omitted for clarity.. For smaller step sizes however, lowering $\epsilon$ can yield significant speedups often converging faster than simply using a larger learning rate. This confirms the importance of the manifold curvature for optimization (lower $\epsilon$ allows for different directions to be scaled drastically different according to their corresponding curvature). Fig 3b compares the impact of $T$ for models having a proper whitened initialization (solid lines), to models being initialized with a standard “fan-in” initialization (dashed lines) . These results are quite surprising in showing the effectiveness of the whitening reparametrization as a simple initialization scheme. That being said, performance can degrade due to ill conditioning when $T$ becomes excessively large ($T=10^{5}$).

2 Unsupervised Learning

Following Martens , we compare PRONG on the task of minimizing reconstruction error of an 8-layer auto-encoder on the MNIST dataset. The encoder is composed of 4 densely connected sigmoidal layers, with a number of hidden units per layer in $\{1\text{k},500,250,30\}$, and a symmetric (untied) decoder. Hyper-parameters were selected by grid search, based on training error, with the following grid specifications: training batch size in $\{32,64,128,256\}$, learning rates in $\{10^{-1},10^{-2},10^{-3}\}$ and momentum term in $\{0,0.9\}$. For RMSprop, we further tuned the moving average coefficient in $\{0.99,0.999\}$ and the regularization term controlling the maximum scaling factor in $\{0.1,0.01\}$. For PRONG, we fixed the natural reparametrization to $T=10^{3}$, using $N_{s}=100$ samples (i.e. they were not optimized for wallclock time). Reconstruction error with respect to updates and wallclock time are shown in Fig. 3 (c,d).

We can see that PRONG significantly outperforms the baseline methods, by up to an order of magnitude in number of updates. With respect to wallclock, our method significantly outperforms the baselines in terms of time taken to reach a certain error threshold, despite the fact that the runtime per epoch for PRONG was 3.2x that of SGD, compared to batch normalization (2.3x SGD) and RMSprop (9x SGD). Note that these timing numbers reflect performance under the optimal choice of hyper-parameters, which in the case of batch normalization yielded a batch size of $256$, compared to $128$ for all other methods. Further breaking down the performance, $34\%$ of the runtime of PRONG was spent performing the whitening reparametrization, compared to $4\%$ for estimating the per layer means and covariances. This confirms that amortization is paramount to the success of our method.We note that our implementation of the whitening operations is not optimized, as it does not take advantage of GPU acceleration, as opposed to the neural network computations. Therefore, runtime of our method is expected to improve as we move the eigen-decompositions to GPU.

3 Supervised Learning

The next set of experiments addresses the problem of training deep supervised convolutional networks for object recognition. Following , we perform whitening across feature maps only: that is we treat pixels in a given feature map as independent samples. This allows us to implement the whitened neural layer as a sequence of two convolutions, where the first is by a 1x1 whitening filter. PRONG is compared to SGD, RMSprop and batch normalization, with each algorithm being accelerated via momentum. Results are presented on both CIFAR-10 and the ImageNet Challenge (ILSVRC12) datasets . In both cases, learning rates were decreased using a “waterfall” annealing schedule, which divided the learning rate by $10$ when the validation error failed to improve after a set number of evaluations. On CIFAR-10, validation error was estimated every $10^{3}$ updates and the learning rate decreased by a factor of $10$ if the validation error failed to improve by $1\%$ over $4$ consecutive evaluations. For ImageNet, we employed a more aggressive schedule which required that the validation error improves by $1\%$ after each epoch.

The model used for our CIFAR experiments consists of 8 convolutional layers, having $3\times 3$ receptive fields. $2\times 2$ spatial max-pooling was applied between stacks of two convolutional layers, with the exception of the last convolutional layer which computes the class scores and is followed by global max-pooling and soft-max non-linearity. This particular choice of architecture was inspired by the VGG model and held fixed across all experiments. The number of filters per layer is as follows: $64,64,128,128,256,256,512,10$. The model was trained on $24\times 24$ random crops with random horizontal reflections. Model selection was performed on a held-out validation set of $5$k examples. Results are shown in Fig. 4.

With respect to training error, PRONG and batch normalization seem to offer similar speedups compared to SGD with momentum. Our hypothesis is that the benefits of PRONG are more pronounced for densely connected networks, where the number of units per layer is typically larger than the number of maps used in convolutional networks. Interestingly, PRONG generalized better, achieving $7.32\%$ test error vs. $8.22\%$ for batch normalization. This could reflect the findings of , which showed how NGD can leverage unlabeled data for better generalization: the “unlabeled” data here comes from the extra perturbations in the training set when estimating the whitening matrices.

3.2 ImageNet Challenge Dataset

Our final set of experiments aims to show the scalability of our method: we thus apply our natural gradient algorithm to the large-scale ILSVRC12 dataset (1.3M images labelled into 1000 categories) using the Inception architecture . In order to scale to problems of this size, we parallelized our training loop so as to split the processing of a single minibatch (of size $256$) across multiple GPUs. Note that PRONG can scale well in this setting, as the estimation of the mean and covariance parameters of each layer is also embarassingly parallel. Eight GPUs were used for computing gradients and estimating model statistics, though the eigen decomposition required for whitening was itself not parallelized in the current implementation.

For all optimization algorithms, we considered initial learning rates in $\{10^{-1},10^{-2},10^{-3}\}$ and used a value of $0.9$ as the momentum coefficient. For PRONG we tested reparametrization periods $T\in\{10,10^{2},10^{3},10^{4}\}$, while typically using $N_{s}=0.1T$. Eigenvalues were regularized by adding a small constant $\epsilon\in\{1,10^{-1},10^{-2},10^{-3}\}$ before scaling the eigenvectors The grid was not searched exhaustively as the cost would have been prohibitive. As our main focus is optimization, regularization consisted of a simple $L_{2}$ weight decay parameter of $10^{-4}$, with no Dropout .. Given the difficulty of the task, we employed the enhanced PRONG+ version of the algorithm, as simple periodic whitening of the model proved to be unstable.This instability may have been compounded by momentum, which was initially not reset after each model reparametrization when using standard PRONG.

Figure 4 (c-d) shows that batch normalisation and PRONG+ converge to approximately the same top-1 validation error ($28.6\%$ vs $28.9\%$ respectively) for similar cpu-time. In comparison, SGD achieved a validation error of $32.1\%$. PRONG+ however exhibits much faster convergence initially: after $10^{5}$ updates it obtains around $36\%$ error compared to $46\%$ for BN alone. We stress that the ImageNet results are somewhat preliminary. While our top-1 error is higher than reported in ($25.2\%$), we used a much less extensive data augmentation pipeline. We are only beginning to explore what natural gradient methods may achieve on these large scale optimization problems and are encouraged by these initial findings.

Discussion

We began this paper by asking whether convergence speed could be improved by simple model reparametrizations, driven by the structure of the Fisher matrix. From a theoretical and experimental perspective, we have shown that Whitened Neural Networks can achieve this via a simple, scalable and efficient whitening reparametrization. They are however one of several possible instantiations of the concept of Natural Neural Networks. In a previous incarnation of the idea, we exploited a similar reparametrization to include whitening of backpropagated gradientsThe weight matrix can be parametrized as $W_{i}=R_{i}^{T}V_{i}U_{i-1}$, with $R_{i}$ the whitening matrix for $\delta_{i}$.. We favor the simpler approach presented in this paper, as we generally found the alternative less stable with deep networks. Ensuring zero-mean gradients also required the use of skip-connections, with tedious book-keeping to offset the reparametrization of centered non-linearities .

Maintaining whitened activations may also offer additional benefits from the point of view of model compression and generalization. By virtue of whitening, the projection $U_{i}h_{i}$ forms an ordered representation, having least and most significant bits. The sharp roll-off in the eigenspectrum of $\Sigma_{i}$ may explain why deep networks are ammenable to compression . Similarly, one could envision spectral versions of Dropout where the dropout probability is a function of the eigenvalues. Alternative ways of orthogonalizing the representation at each layer should also be explored, via alternate decompositions of $\Sigma_{i}$, or perhaps by exploiting the connection between linear auto-encoders and PCA. We also plan on pursuing the connection with Mirror Descent and further bridging the gap between deep learning and methods from online convex optimization.

We are extremely grateful to Shakir Mohamed for invaluable discussions and feedback in the preparation of this manuscript. We also thank Philip Thomas, Volodymyr Mnih, Raia Hadsell, Sergey Ioffe and Shane Legg for feedback on the paper.

References