GradInit: Learning to Initialize Neural Networks for Stable and
  Efficient Training
  Innovations in neural architectures have fostered significant breakthroughs
in language modeling and computer vision. Unfortunately, novel architectures
often result in challenging hyper-parameter choices and training instability if
the network parameters are not properly initialized. A number of
architecture-specific initialization schemes have been proposed, but these
schemes are not always portable to new architectures. This paper presents
GradInit, an automated and architecture agnostic method for initializing neural
networks. GradInit is based on a simple heuristic; the norm of each network
layer is adjusted so that a single step of SGD or Adam with prescribed
hyperparameters results in the smallest possible loss value. This adjustment is
done by introducing a scalar multiplier variable in front of each parameter
block, and then optimizing these variables using a simple numerical scheme.
GradInit accelerates the convergence and test performance of many convolutional
architectures, both with or without skip connections, and even without
normalization layers. It also improves the stability of the original
Transformer architecture for machine translation, enabling training it without
learning rate warmup using either Adam or SGD under a wide range of learning
rates and momentum coefficients. Code is available at
https://github.com/zhuchen03/gradinit.
