[Paper Review] WNGrad: Learn the Learning Rate in Gradient Descent

Xiaoxia Wu, Rachel Ward+1 more|arXiv (Cornell University)|Mar 7, 2018

Stochastic Gradient Optimization Techniques21 references55 citations

TL;DR

WNGrad introduces a dynamic learning-rate update rule that adapts based on gradient observations, achieving robustness to Lipschitz constants and near-optimal convergence in batch and stochastic settings.

ABSTRACT

Adjusting the learning rate schedule in stochastic gradient methods is an important unresolved problem which requires tuning in practice. If certain parameters of the loss function such as smoothness or strong convexity constants are known, theoretical learning rate schedules can be applied. However, in practice, such parameters are not known, and the loss function of interest is not convex in any case. The recently proposed batch normalization reparametrization is widely adopted in most neural network architectures today because, among other advantages, it is robust to the choice of Lipschitz constant of the gradient in loss function, allowing one to set a large learning rate without worry. Inspired by batch normalization, we propose a general nonlinear update rule for the learning rate in batch and stochastic gradient descent so that the learning rate can be initialized at a high value, and is subsequently decreased according to gradient observations along the way. The proposed method is shown to achieve robustness to the relationship between the learning rate and the Lipschitz constant, and near-optimal convergence rates in both the batch and stochastic settings ($O(1/T)$ for smooth loss in the batch setting, and $O(1/\sqrt{T})$ for convex loss in the stochastic setting). We also show through numerical evidence that such robustness of the proposed method extends to highly nonconvex and possibly non-smooth loss function in deep learning problems.Our analysis establishes some first theoretical understanding into the observed robustness for batch normalization and weight normalization.

Motivation & Objective

• Motivate and address the difficulty of selecting learning-rate schedules in stochastic gradient methods when key loss-function constants are unknown.
• Propose a reparametrization-inspired learning-rate update that starts large and adapts based on observed gradients.
• Establish convergence guarantees for WNGrad in both batch (nonconvex smooth) and stochastic (convex, not necessarily smooth) settings.
• Demonstrate robustness and practical performance through numerical experiments on standard datasets (MNIST and CIFAR-10).

Proposed method

• Introduce WNGrad with updates x_{k+1} = x_k - (1/b_k) ∇f(x_k) and b_{k+1} = b_k + (1/b_k) ∥∇f(x_k)∥^2.
• Show that b_k grows to a level that stabilizes when b_k ≥ L, leading to robustness to the unknown Lipschitz constant L.
• Prove global convergence for smooth f: min_k ∥∇f(x_k)∥^2 ≤ ε with T bounds that depend on f(x_1), f*, and L.
• Prove stochastic convergence: f( x̄_k ) - f* ≤ G^2(D^2+2)/(γ√k) + (b_1 ∥x_1 - x*∥^2)/(2k) under convexity and variance assumptions.
• Relate WNGrad to AdaGrad-like behavior without square-root computations, highlighting efficiency and scale-invariance.

Experimental results

Research questions

• RQ1 Can a single, dynamically updated learning-rate parameter b_k yield convergence guarantees in both batch and stochastic gradient settings without knowing the Lipschitz constant?
• RQ2 Does the proposed b_k update provide robustness to the choice of learning-rate scale and improve generalization in nonconvex neural network problems?
• RQ3 What are the theoretical convergence rates for WNGrad in batch (smooth) and stochastic (convex, not necessarily smooth) settings?
• RQ4 How does WNGrad perform empirically on standard datasets (MNIST, CIFAR-10) compared to SGD and adaptive methods?
• RQ5 How do momentum variants (WN-Adam, WNGrad-Momentum) affect robustness to learning-rate scale?

Key findings

• WNGrad converges to a point with ∥∇f(x_T)∥^2 ≤ ε in O((f(x_1)−f*+L)^2/ε) iterations for smooth f in the batch setting.
• In stochastic settings, b_k grows like O(√k/G), yielding the optimal O(1/√T) rate for convex losses.
• The scheme is scale-invariant: resizing f by a constant scales does not change the iterates under WNGrad.
• WNGrad with momentum or Adam variants maintains robustness to Lipschitz-constant scale, outperforming standard SGD/Adam in some settings in experiments.
• Numerical experiments on MNIST and CIFAR-10 show WNGrad exhibits robustness to gradient Lipschitz constants and competitive generalization behavior compared to SGD.