Осторожное затухание весов

Аннотация

Мы представляем Cautious Weight Decay (CWD) — однострочную, независимую от оптимизатора модификацию, которая применяет затухание весов только к тем координатам параметров, чьи знаки совпадают с обновлением оптимизатора. В отличие от стандартного раздельного затухания, которое неявно оптимизирует регуляризованную или ограниченную целевую функцию, CWD сохраняет исходную функцию потерь и допускает биуровневую интерпретацию: оно вызывает скользящий режим при достижении стационарного многообразия, позволяя искать локально Парето-оптимальные стационарные точки неизмененной целевой функции. На практике CWD представляет собой простую замену для оптимизаторов, таких как AdamW, Lion и Muon, не требующую новых гиперпараметров или дополнительной настройки. Для предварительного обучения языковых моделей и классификации на ImageNet CWD последовательно улучшает итоговую функцию потерь и точность на масштабах от миллионов до миллиардов параметров.

English

We introduce Cautious Weight Decay (CWD), a one-line, optimizer-agnostic modification that applies weight decay only to parameter coordinates whose signs align with the optimizer update. Unlike standard decoupled decay, which implicitly optimizes a regularized or constrained objective, CWD preserves the original loss and admits a bilevel interpretation: it induces sliding-mode behavior upon reaching the stationary manifold, allowing it to search for locally Pareto-optimal stationary points of the unmodified objective. In practice, CWD is a drop-in change for optimizers such as AdamW, Lion, and Muon, requiring no new hyperparameters or additional tuning. For language model pre-training and ImageNet classification, CWD consistently improves final loss and accuracy at million- to billion-parameter scales.