Maximal Initial Learning Rates in Deep ReLU Networks

TY - JOUR

T1 - Maximal Initial Learning Rates in Deep ReLU Networks

AU - Iyer, Gaurav

AU - Hanin, Boris

AU - Rolnick, David

N1 - Publisher Copyright: © 2023 Proceedings of Machine Learning Research. All rights reserved.

PY - 2023

Y1 - 2023

N2 - Training a neural network requires choosing a suitable learning rate, which involves a trade-off between speed and effectiveness of convergence. While there has been considerable theoretical and empirical analysis of how large the learning rate can be, most prior work focuses only on late-stage training. In this work, we introduce the maximal initial learning rate η∗ - the largest learning rate at which a randomly initialized neural network can successfully begin training and achieve (at least) a given threshold accuracy. Using a simple approach to estimate η∗, we observe that in constant-width fully-connected ReLU networks, η∗ behaves differently from the maximum learning rate later in training. Specifically, we find that η∗ is well predicted as a power of (depth ×width), provided that (i) the width of the network is sufficiently large compared to the depth, and (ii) the input layer is trained at a relatively small learning rate. We further analyze the relationship between η∗ and the sharpness λ1 of the network at initialization, indicating they are closely though not inversely related. We formally prove bounds for λ1 in terms of (depth × width) that align with our empirical results.

AB - Training a neural network requires choosing a suitable learning rate, which involves a trade-off between speed and effectiveness of convergence. While there has been considerable theoretical and empirical analysis of how large the learning rate can be, most prior work focuses only on late-stage training. In this work, we introduce the maximal initial learning rate η∗ - the largest learning rate at which a randomly initialized neural network can successfully begin training and achieve (at least) a given threshold accuracy. Using a simple approach to estimate η∗, we observe that in constant-width fully-connected ReLU networks, η∗ behaves differently from the maximum learning rate later in training. Specifically, we find that η∗ is well predicted as a power of (depth ×width), provided that (i) the width of the network is sufficiently large compared to the depth, and (ii) the input layer is trained at a relatively small learning rate. We further analyze the relationship between η∗ and the sharpness λ1 of the network at initialization, indicating they are closely though not inversely related. We formally prove bounds for λ1 in terms of (depth × width) that align with our empirical results.

UR - http://www.scopus.com/inward/record.url?scp=85174401911&partnerID=8YFLogxK

UR - http://www.scopus.com/inward/citedby.url?scp=85174401911&partnerID=8YFLogxK

M3 - Conference article

SN - 2640-3498

VL - 202

SP - 14500

EP - 14530

JO - Proceedings of Machine Learning Research

JF - Proceedings of Machine Learning Research

T2 - 40th International Conference on Machine Learning, ICML 2023

Y2 - 23 July 2023 through 29 July 2023

ER -