Initializing the Layer-wise Learning Rate

Kwang Yong Shin, Suhyun Kim, Soo-Mook Moon

Primary Area: optimization

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Learning Rate, Exploding Gradient, Vanishing Gradient

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: Adjust the layer-wise learning rate opposite to the gradient magnitude at initialization.

Abstract: The standard method to assign learning rates has been to rely on the optimizer and to use a single, global learning rate across all its layers. We propose to assign individual learning rates as well, according to the layer-wise gradient magnitude at initialization. Even if individual layers are initialized to preserve gradient variance, architectural characteristics result in uneven gradient magnitude even when the network has not started training. We interpret this gradient magnitude as a measure of architecture-induced convergence bias, and adjust the layer-wise learning rate opposite to its gradient magnitude at initialization. This relative learning rate is maintained throughout the entire training scheme. Experiments on convolutional and transformer architectures on ImageNet-1k show improved accuracy and training stability.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 8568