Principle:LaurentMazare Tch rs Sparse Adam Optimization

Knowledge Sources	LaurentMazare_Tch_rs Adam: A Method for Stochastic Optimization
Domains	Optimization, Deep Learning
Last Updated	2026-02-08 00:00 GMT

Overview

Sparse Adam is a variant of the Adam optimizer that efficiently handles sparse gradients by updating only the moment estimates corresponding to non-zero gradient entries.

Description

The standard Adam optimizer maintains two exponential moving averages for every parameter element: the first moment (mean of gradients) and the second moment (mean of squared gradients). When gradients are sparse (i.e., most entries are zero, as commonly occurs with embedding layers), standard Adam still decays all moment estimates at every step, even for entries with zero gradient. This leads to incorrect moment estimates because the decay falsely treats missing gradient information as a zero gradient signal.

Sparse Adam addresses this by applying a key modification:

• Selective moment updates: Only the moment estimates corresponding to non-zero gradient entries are updated. Entries with zero gradients retain their previous moment values without decay.

• Lazy update semantics: For each parameter element, the effective number of updates is tracked. When a gradient entry becomes non-zero after several steps of inactivity, the moment estimates are corrected as if the intermediate zero-gradient steps had not occurred.

• Bias correction: Like standard Adam, Sparse Adam applies bias correction to the first and second moment estimates. The correction factors account for the initialization at zero and the geometric decay:

The result is that Sparse Adam produces identical updates to standard Adam for non-zero gradient entries while avoiding the pathological decay behavior on inactive entries.

Usage

Sparse Adam is particularly important for training models with embedding layers where each mini-batch only updates a small subset of embedding vectors. It is also relevant in any scenario where gradients are structurally sparse, such as in recommender systems, sparse attention mechanisms, or models with conditional computation paths.

Theoretical Basis

Standard Adam Update:

$m_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t$

$v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$

${\hat{m}}_t=\frac{m_t}{1-{\beta }_1^t}$

${\hat{v}}_t=\frac{v_t}{1-{\beta }_2^t}$

${\theta }_{t+1}={\theta }_t-\alpha \frac{{\hat{m}}_t}{\sqrt{{\hat{v}}_t}+ϵ}$

where ${\beta }_1,{\beta }_2$ are decay rates (typically 0.9 and 0.999), $\alpha$ is the learning rate, and $ϵ$ is a small constant for numerical stability.

Sparse Adam Modification:

For each parameter element $i$, define the set of active steps $S_i=\{t:g_{t,i}\ne 0\}$.

The moment updates are applied only when $t\in S_i$:

$m_{t,i}=\{\begin{array}{ll}{\beta }_1{m}_{t-1,i}+(1-{\beta }_1)g_{t,i}& \text{if }{g}_{t,i}\ne 0\\ m_{t-1,i}& \text{if }{g}_{t,i}=0\end{array}$

$v_{t,i}=\{\begin{array}{ll}{\beta }_2{v}_{t-1,i}+(1-{\beta }_2)g_{t,i}^2& \text{if }{g}_{t,i}\ne 0\\ v_{t-1,i}& \text{if }{g}_{t,i}=0\end{array}$

Bias Correction for Sparse Updates:

The bias correction terms adjust for the actual number of active updates rather than the total step count:

${\hat{m}}_{t,i}=\frac{m_{t,i}}{1-{\beta }_1^{|S_i\cap \{1,\dots ,t\}|}}$

${\hat{v}}_{t,i}=\frac{v_{t,i}}{1-{\beta }_2^{|S_i\cap \{1,\dots ,t\}|}}$

This ensures unbiased estimates even when updates are infrequent for certain parameter elements.

Related Pages

• Implementation:LaurentMazare_Tch_rs_Sparse_Adam