Universal Dynamics of Warmup Stable Decay: understanding WSD beyond Transformers

TL;DR

This paper investigates the universal training dynamics of the Warmup Stable Decay learning rate scheduler across different model architectures, revealing shared geometric properties of loss landscapes in large language models and CNNs.

Contribution

It demonstrates that WSD's effectiveness and training behavior are consistent across transformer-based language models and CNNs, suggesting shared geometric landscape features.

Findings

WSD exhibits similar path features in both language models and CNNs.

Training signals and sharpness dynamics are qualitatively similar across architectures.

Shared geometric characteristics of loss landscapes are observed in diverse nonconvex problems.

Abstract

The Warmup Stable Decay (WSD) learning rate scheduler has recently become popular, largely due to its good performance and flexibility when training large language models. It remains an open question whether the remarkable performance of WSD - using a decaying learning rate for only a fraction of training compared to cosine decay - is a phenomenon specific to transformer-based language models that can potentially offer new theoretical insights into their training dynamics. Inspired by the usage of learning rate schedulers as a new lens into understanding landscape geometry (e.g., river valley, connected minima, progressive sharpening), in this work we compare the WSD path of the Adam optimizer on a Pythia-like language model to that of a small CNN trained to classify CIFAR10 images. We observe most training signals, optimizer path features, and sharpness dynamics to be qualitatively similar in such architectures. This consistency points to shared geometric characteristics of the loss landscapes of old and new nonconvex problems, and hints to future research questions around the geometry of high dimensional optimization problems.