[Paper Review] Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks

Tim Salimans, Diederik P. Kingma|arXiv (Cornell University)|Feb 25, 2016

Generative Adversarial Networks and Image Synthesis27 references928 citations

TL;DR

Weight normalization reparameterizes weights as w = (g/||v||) v to decouple length and direction, improving gradient conditioning and speeding up stochastic gradient descent across architectures, including RNNs and RL, with lower overhead than batch normalization.

ABSTRACT

We present weight normalization: a reparameterization of the weight vectors in a neural network that decouples the length of those weight vectors from their direction. By reparameterizing the weights in this way we improve the conditioning of the optimization problem and we speed up convergence of stochastic gradient descent. Our reparameterization is inspired by batch normalization but does not introduce any dependencies between the examples in a minibatch. This means that our method can also be applied successfully to recurrent models such as LSTMs and to noise-sensitive applications such as deep reinforcement learning or generative models, for which batch normalization is less well suited. Although our method is much simpler, it still provides much of the speed-up of full batch normalization. In addition, the computational overhead of our method is lower, permitting more optimization steps to be taken in the same amount of time. We demonstrate the usefulness of our method on applications in supervised image recognition, generative modelling, and deep reinforcement learning.

Motivation & Objective

• Motivate faster and more stable optimization for deep neural networks by improving gradient conditioning.

Proposed method

• Reparameterize each weight vector w as w = (g/||v||) v, introducing parameters v and g.
• Compute gradients w.r.t. v and g using the derived equations, enabling backpropagation through the reparameterization.
• Show that this decouples weight norm from direction, yielding a gradient covariance closer to identity and faster convergence.
• Relate weight normalization to batch normalization, noting it is deterministic and minibatch-independent.
• Provide a practical initialization scheme to stabilize training and discuss extensions like mean-only batch normalization.
• Demonstrate applicability across supervised, generative, and reinforcement learning tasks.

Experimental results

Research questions

• RQ1 Does weight normalization speed up training and improve convergence compared to standard parameterization?
• RQ2 How does weight normalization compare to batch normalization in terms of speed, overhead, and robustness to learning rate?
• RQ3 Can weight normalization be effectively applied to RNNs/LSTMs and reinforcement learning settings where batch normalization is problematic?
• RQ4 What initialization strategies are necessary to realize stable training with weight normalization?

Key findings

• Weight normalization speeds up convergence and provides robust optimization with a wider range of learning rates than standard parameterization.
• Mean-only batch normalization combined with weight normalization yields best CIFAR-10 test error in their setup (7.31%).
• Weight normalization achieves similar acceleration to batch normalization but with lower computational overhead and no minibatch dependencies.
• Applied to CNNs (CIFAR-10), CVAEs (MNIST, CIFAR-10), DRAW (MNIST), and DQN (Atari), weight normalization consistently speeds training and improves or matches performance.
• In CIFAR-10 experiments, mean-only batch normalization with weight normalization outperforms standard BN, fully normalized BN, and normal parameterizations in test error across their architecture variant."],
• table_headers:[