Source-linked AI summary
Natural Neural Networks
Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, Koray Kavukcuoglu
TL;DR
Deep-network optimization is hindered by parameter dependencies, while natural-gradient methods are costly because they require Fisher-matrix estimation and inversion. The paper introduces whitened neural networks and PRONG, which preserve feed-forward computation while amortizing whitening-based reparametrizations. The approach improves Fisher conditioning and scales from auto-encoders to multi-GPU ImageNet models, with reported gains across unsupervised and supervised settings.
Problem
Deep networks can be difficult to optimize with first-order SGD, while natural-gradient methods face high computational cost from Fisher-matrix estimation and inversion.
Method
Whitened neural networks reparametrize layer weights using estimated activation statistics, and PRONG amortizes these whitening operations over multiple updates.
Results
The method improves Fisher conditioning by more than 95%, outperforms baselines by up to an order of magnitude in updates, and scales to multi-GPU ImageNet models.
Takeaways & Limitations
Natural Neural Networks provide a simple, scalable whitening-based route toward efficient natural-gradient training across unsupervised and supervised tasks.
Takeaways & Limitations
Full covariance estimation requires more samples, and computing whitening matrices requires an eigen-decomposition whose cost is cubic in layer size.
Abstract
from arXiv · showhide
We introduce Natural Neural Networks, a novel family of algorithms that speed up convergence by adapting their internal representation during training to improve conditioning of the Fisher matrix. In particular, we show a specific example that employs a simple and efficient reparametrization of the neural network weights by implicitly whitening the representation obtained at each layer, while preserving the feed-forward computation of the network. Such networks can be trained efficiently via the proposed Projected Natural Gradient Descent algorithm (PRONG), which amortizes the cost of these reparametrizations over many parameter updates and is closely related to the Mirror Descent online learning algorithm. We highlight the benefits of our method on both unsupervised and supervised learning tasks, and showcase its scalability by training on the large-scale ImageNet Challenge dataset.
1 Introduction
Deep networks’ parameter dependencies can make first-order SGD difficult to optimize, motivating natural-gradient approaches despite their high computational cost. The paper introduces a whitened reparametrization that approximates layerwise natural-gradient updates efficiently and scales to ImageNet.
- Deep network dependencies can make learning difficult with first-order SGD, while batch normalization demonstrates optimization-related performance gains.
- Natural gradient uses the probabilistic-manifold geometry and KL divergence, but typically requires estimating and inverting a parameter-sized Fisher Information Matrix.
- The proposed reparametrization seeks a neural architecture with identity-constrained Fisher, approximating block-diagonal natural-gradient updates through local whitening and amortized computation.
- The method scales from deep auto-encoders to large multi-GPU convolutional models on ImageNet, representing the authors’ first scaling of a non-diagonal natural-gradient algorithm to this magnitude.
2 The Natural Gradient
The paper formulates natural-gradient optimization through the Fisher geometry of neural networks and derives its layerwise structure for multilayer perceptrons. It focuses on blockwise Fisher terms and a whitening-related independence assumption that can improve conditioning.
- The method is developed for classification but is stated to generalize to regression and density estimation.
- The optimization problem fits model parameters to an empirical distribution under log-loss.
- SGD follows the expected per-example loss gradient and updates parameters using a learning rate.
- Natural gradient replaces SGD’s L2 iterate distance with KL divergence, using the Fisher Information Matrix as its metric.
- For an MLP, the Fisher metric is derived from layer weights, biases, nonlinear activations, and backpropagated gradients.
- The analysis focuses on within-layer Fisher blocks, vectorizes matrix parameters, and ignores off-block-diagonal interactions.
- Assuming backpropagated gradients and preceding activations are independent, the paper hypothesizes that enforcing whitened activation covariance can improve Fisher conditioning.
3 Projected Natural Gradient Descent
Whitened Neural Networks reparametrize layers by centering and whitening hidden representations, while PRONG updates these coefficients periodically and preserves the network’s feed-forward function. This improves Fisher conditioning and connects projected updates to Mirror Descent.
- 3.1 A Whitened Neural Layer: Whitened Neural Networks use per-layer centering and ZCA-whitening matrices to approximately whiten internal hidden representations.The centering parameters track layer means, while whitening matrices are obtained from covariance eigendecompositions.
- 3.2 Updating the Whitening Coefficients: The whitening coefficients are estimated from model statistics rather than learned by loss minimization, making block-diagonal Fisher terms identity by construction.The coefficients are treated as constants by the optimizer and improve conditioning with respect to the whitened parameters.
- 3.2 Updating the Whitening Coefficients: Coupled updates to whitening and model parameters preserve the implemented network function while the internal representation changes.The procedure preserves products such as ViUi−1 and applies an analogous constraint to biases.
- 3.2 Updating the Whitening Coefficients: Full covariance estimation requires more samples than mean or diagonal-variance estimation, and computing each whitening eigendecomposition remains cubic in layer size.These costs motivate amortizing reparametrizations over consecutive updates.
- 3.2 Updating the Whitening Coefficients: PRONG updates whitening coefficients every T steps, amortizing covariance estimation and eigendecomposition costs across intervening SGD updates.At each reparametrization, it estimates means and covariances, updates the whitening coefficients, projects parameters, and then performs SGD updates in the whitened coordinates.
- 3.3 Duality and Mirror Descent: PRONG has a parameter duality with canonical models and is closely connected to Mirror Descent through projections between canonical and whitened parameter spaces.The whitened formulation permits projection costs to be amortized over several updates while gradients are computed directly in the dual parameter space.
4 Experiments
The experiments evaluate PRONG across conditioning diagnostics, unsupervised auto-encoding, and supervised convolutional classification, including large-scale ImageNet training. PRONG improves conditioning and convergence, while PRONG+ reaches comparable ImageNet validation error with faster initial convergence.
- 4.1 Introspective Experiments: More than 95% reduction in the Fisher condition number confirms that whitening activations improves conditioning in the diagnostic MLP experiment.The condition number was measured relative to its initial value before PRONG’s first whitening reparametrization.
- 4.1 Introspective Experiments: Lowering ϵ can significantly speed convergence at smaller step sizes, whereas excessively large reparametrization intervals can degrade performance through ill conditioning.The experiments varied the trust-region control ϵ and reparametrization interval T.
- 4.2 Unsupervised Learning: PRONG outperforms baseline methods by up to an order of magnitude in the number of updates needed for the 8-layer MNIST auto-encoder.Despite a 3.2x SGD runtime per epoch, PRONG reaches error thresholds faster in wall-clock time; whitening consumed 34% of its runtime.
- 4.2 Unsupervised Learning: Whitening overhead was not GPU-optimized, so moving eigendecompositions to the GPU is expected to reduce PRONG’s runtime.The reported implementation spent 34% of runtime on whitening reparametrization.
- 4.3.1 CIFAR-10: On CIFAR-10, PRONG and batch normalization provide similar training-error speedups over SGD with momentum, while PRONG achieves 7.32% test error versus 8.22% for batch normalization.The convolutional architecture was held fixed across the CIFAR experiments.
- 4.3.2 ImageNet Challenge Dataset: The ImageNet results use PRONG+ because simple periodic whitening was unstable, and the authors describe these findings as preliminary.The comparison used an aggressive validation-based learning-rate schedule and a less extensive data-augmentation pipeline than prior work.
- 4.3.2 ImageNet Challenge Dataset: On ImageNet, PRONG+ and batch normalization reach approximately the same top-1 validation error, 28.9% and 28.6% respectively, for similar CPU time.After 10^5 updates, PRONG+ obtains around 36% error compared with 46% for batch normalization alone.
5 Discussion
The discussion presents Whitened Neural Networks as one scalable instantiation of Natural Neural Networks that improves convergence through whitening reparametrization. It also identifies possible links to compression, generalization, spectral dropout, and online convex optimization as directions for further study.
- 5 Discussion: Whitened Neural Networks improve convergence through a simple, scalable, and efficient whitening reparametrization, but represent only one possible Natural Neural Network design.The discussion contrasts the presented approach with a less stable alternative that also whitened backpropagated gradients.
- 5 Discussion: Maintaining whitened activations may benefit model compression and generalization by producing ordered representations associated with the eigenspectrum of each layer’s covariance.The paper proposes spectral dropout and alternative orthogonalization schemes as directions for exploration.