Search papers, labs, and topics across Lattice.
This paper introduces DeltaMomentum, an innovative momentum update rule that incorporates direction-awareness into the optimization process by leveraging a key-value structure derived from the gradients of linear layers. Unlike traditional methods that apply a uniform forgetting rate across all directions, DeltaMomentum adjusts the forgetting rate based on the frequency of direction queries, leading to more efficient training. Experimental results demonstrate that AdamW with DeltaMomentum (DeltaAdamW) achieves significantly lower validation loss in fewer training steps compared to standard AdamW, indicating enhanced optimization performance across various model scales and architectures.
DeltaMomentum accelerates convergence by adapting momentum updates to the frequency of gradient directions, outperforming traditional methods in both speed and efficiency.
Most modern optimizers form their momentum as an exponential moving average (EMA) of past gradients, forgetting every direction at one fixed rate. However, the inputs a deep network sees during training can be highly anisotropic, with a few directions queried frequently while most are seen rarely. Recent methods address this anisotropy by wrapping extra processing around this buffer, leaving the momentum update itself unchanged. We propose DeltaMomentum, which builds direction-awareness into the momentum update rule. The main observation is that the gradient of a linear layer splits into an input that acts as a key and an output-side error that acts as a value. Exploiting the key-value structure, DeltaMomentum updates the momentum buffer by the canonical delta rule, so each direction is forgotten at a rate set by how often it appears. We prove that it is a valid momentum, that it applies the input-side curvature correction without matrix inversion, and that it clears stale directions faster than EMA under both a fixed and a drifting optimum. It is a drop-in replacement for the momentum buffer of any optimizer, its coefficient transfers across widths under $\mu$P, and its extra compute stays between $22.2\%$ and $25.0\%$ of a gated-MLP block's linear cost with no persistent memory. In FineWeb-Edu pretraining, AdamW with DeltaMomentum (DeltaAdamW) reaches AdamW's validation loss in up to $46.39 \pm 4.32\%$ fewer steps at 67M and $22.12 \pm 0.80\%$ at 370M over three seeds, and the gain persists at 1B on a Chinchilla-optimal budget. A Muon baseline tuned under the same protocol sits above DeltaAdamW at both language-model scales, and the gain holds for SGD, ResNet-18, and ViT-Tiny on CIFAR-10. Training-time diagnostics confirm the predicted mechanism, better gradient tracking and healthier input directions.
