Understanding Decoupled and Early Weight Decay

TL;DR

This paper investigates the effects of weight decay in deep learning, especially when applied only at the start or decoupled from the loss, revealing insights into regularization, generalization, and optimizer behavior across vision, NLP, and RL tasks.

Contribution

It provides a comprehensive analysis of decoupled and early weight decay, demonstrating their impact on network norm, regularization, and optimizer dynamics across different domains.

Findings

Applying WD only at the start keeps network norm small and regularizes training.

Decoupled WD prevents gradient mixing in Adam, aiding hyperparameter tuning.

Traditional generalization metrics fail to capture WD effects, but scale-invariant metrics succeed.

Abstract

Weight decay (WD) is a traditional regularization technique in deep learning, but despite its ubiquity, its behavior is still an area of active research. Golatkar et al. have recently shown that WD only matters at the start of the training in computer vision, upending traditional wisdom. Loshchilov et al. show that for adaptive optimizers, manually decaying weights can outperform adding an $l_2$ penalty to the loss. This technique has become increasingly popular and is referred to as decoupled WD. The goal of this paper is to investigate these two recent empirical observations. We demonstrate that by applying WD only at the start, the network norm stays small throughout training. This has a regularizing effect as the effective gradient updates become larger. However, traditional generalizations metrics fail to capture this effect of WD, and we show how a simple scale-invariant metric can. We also show how the growth of network weights is heavily influenced by the dataset and its generalization properties. For decoupled WD, we perform experiments in NLP and RL where adaptive optimizers are the norm. We demonstrate that the primary issue that decoupled WD alleviates is the mixing of gradients from the objective function and the $l_2$ penalty in the buffers of Adam (which stores the estimates of the first-order moment). Adaptivity itself is not problematic and decoupled WD ensures that the gradients from the $l_2$ term cannot "drown out" the true objective, facilitating easier hyperparameter tuning.