Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective

• URL: http://arxiv.org/abs/2410.05192v2
• Date: Tue, 29 Oct 2024 06:26:00 GMT
• Title: Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective
• Authors: Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, Tengyu Ma,
• Abstract summary: Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can continue indefinitely without a pre-specified compute budget. We show that pretraining loss exhibits a river valley landscape, which resembles a deep valley with a river at its bottom. Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints' decay phases and keeps only one main branch.
• Score: 66.80315289020487
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Training language models currently requires pre-determining a fixed compute budget because the typical cosine learning rate schedule depends on the total number of steps. In contrast, the Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can in principle continue indefinitely without a pre-specified compute budget. Then, given any compute budget, one can branch out from the main branch at a proper at any time with a rapidly decaying learning rate to produce a strong model. Empirically, WSD generates a non-traditional loss curve: the loss remains elevated during the stable phase but sharply declines during the decay phase. Towards explaining this phenomenon, we conjecture that pretraining loss exhibits a river valley landscape, which resembles a deep valley with a river at its bottom. Under this assumption, we show that during the stable phase, the iterate undergoes large oscillations due to the high learning rate, yet it progresses swiftly along the river. During the decay phase, the rapidly dropping learning rate minimizes the iterate's oscillations, moving it closer to the river and revealing true optimization progress. Therefore, the sustained high learning rate phase and fast decaying phase are responsible for progress in the river and the mountain directions respectively, and are both critical. Our analysis predicts phenomenons consistent with empirical observations and shows that this landscape can emerge from pretraining on a simple bi-gram dataset. Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints' decay phases and keeps only one main branch, where we resume from a decayed checkpoint. WSD-S empirically outperforms WSD and Cyclic-Cosine in obtaining multiple language model checkpoints across various compute budgets in a single run for parameters scaling from 0.1B to 1.2B.

Related papers

• Inference and Interference: The Role of Clipping, Pruning and Loss Landscapes in Differentially Private Stochastic Gradient Descent [13.27004430044574]
  Differentially private gradient descent (DP-SGD) is known to have poorer training and test performance on large neural networks. We compare the behavior of the two processes separately in early and late epochs. We find that while DP-SGD makes slower progress in early stages, it is the behavior in the later stages that determines the end result.
  arXiv  Detail & Related papers  (2023-11-12T13:31:35Z)
• Critical Parameters for Scalable Distributed Learning with Large Batches and Asynchronous Updates [67.19481956584465]
  It has been experimentally observed that the efficiency of distributed training with saturation (SGD) depends decisively on the batch size and -- in implementations -- on the staleness. We show that our results are tight and illustrate key findings in numerical experiments.
  arXiv  Detail & Related papers  (2021-03-03T12:08:23Z)
• Acceleration via Fractal Learning Rate Schedules [37.878672787331105]
  We show that the learning rate schedule remains notoriously difficult to understand and expensive to tune. We reinterpret an iterative algorithm from the numerical analysis literature as what we call the Chebyshev learning rate schedule for accelerating vanilla gradient descent. We provide some experiments and discussion to challenge current understandings of the "edge of stability" in deep learning.
  arXiv  Detail & Related papers  (2021-03-01T22:52:13Z)
• Implicit bias of deep linear networks in the large learning rate phase [15.846533303963229]
  We characterize the implicit bias effect of deep linear networks for binary classification using the logistic loss in a large learning rate regime. We claim that depending on the separation conditions of data, the gradient descent iterates will converge to a flatter minimum in the catapult phase.
  arXiv  Detail & Related papers  (2020-11-25T06:50:30Z)
• Direction Matters: On the Implicit Bias of Stochastic Gradient Descent with Moderate Learning Rate [105.62979485062756]
  This paper attempts to characterize the particular regularization effect of SGD in the moderate learning rate regime. We show that SGD converges along the large eigenvalue directions of the data matrix, while GD goes after the small eigenvalue directions.
  arXiv  Detail & Related papers  (2020-11-04T21:07:52Z)
• Consistency Guided Scene Flow Estimation [159.24395181068218]
  CGSF is a self-supervised framework for the joint reconstruction of 3D scene structure and motion from stereo video. We show that the proposed model can reliably predict disparity and scene flow in challenging imagery. It achieves better generalization than the state-of-the-art, and adapts quickly and robustly to unseen domains.
  arXiv  Detail & Related papers  (2020-06-19T17:28:07Z)
• Bridging the Gap Between Training and Inference for Spatio-Temporal Forecasting [16.06369357595426]
  We propose a novel curriculum learning based strategy named Temporal Progressive Growing Sampling to bridge the gap between training and inference for S-temporal sequence forecasting. Experimental results demonstrate that our proposed method better models long term dependencies and outperforms baseline approaches on two competitive datasets.
  arXiv  Detail & Related papers  (2020-05-19T10:14:43Z)
• On Learning Rates and Schr\"odinger Operators [105.32118775014015]
  We present a general theoretical analysis of the effect of the learning rate. We find that the learning rate tends to zero for a broad non- neural class functions.
  arXiv  Detail & Related papers  (2020-04-15T09:52:37Z)
• The Break-Even Point on Optimization Trajectories of Deep Neural Networks [64.7563588124004]
  We argue for the existence of the "break-even" point on this trajectory. We show that using a large learning rate in the initial phase of training reduces the variance of the gradient. We also show that using a low learning rate results in bad conditioning of the loss surface even for a neural network with batch normalization layers.
  arXiv  Detail & Related papers  (2020-02-21T22:55:51Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.