GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

Innovations in neural architectures have fostered significant breakthroughs\nin language modeling and computer vision. Unfortunately, novel architectures\noften result in challenging hyper-parameter choices and training instability if\nthe network parameters are not properly initialized. A number of\narchitecture-specific initialization schemes have been proposed, but these\nschemes are not always portable to new architectures. This paper presents\nGradInit, an automated and architecture agnostic method for initializing neural\nnetworks. GradInit is based on a simple heuristic; the norm of each network\nlayer is adjusted so that a single step of SGD or Adam with prescribed\nhyperparameters results in the smallest possible loss value. This adjustment is\ndone by introducing a scalar multiplier variable in front of each parameter\nblock, and then optimizing these variables using a simple numerical scheme.\nGradInit accelerates the convergence and test performance of many convolutional\narchitectures, both with or without skip connections, and even without\nnormalization layers. It also improves the stability of the original\nTransformer architecture for machine translation, enabling training it without\nlearning rate warmup using either Adam or SGD under a wide range of learning\nrates and momentum coefficients. Code is available at\nhttps://github.com/zhuchen03/gradinit.\n