Module sparse_optimizer 

Sparse optimizer step (CSR/COO + lazy-state sparse SGD / Adam)

This module implements sparse optimizer updates that touch only the coordinates with non-zero gradients. It is the CPU reference implementation of the kind of update used to train very large embedding tables, where on every step only a handful of the millions of rows actually receive a gradient. Touching only those coordinates (both for the parameter update and for the optimizer moment state) is the entire point of “sparse / lazy” optimizers.

Sparse gradient representations

• CooGradient – a coordinate list for a 1-D parameter / embedding row view. It stores the touched indices, their gradient values and the logical dimension dim of the dense vector it represents. It is kept in a canonical form: indices are strictly ascending (hence sorted and unique).
• CsrGradient – compressed sparse row for a 2-D embedding table. It stores row_offsets, col_indices and values together with the dense shape. Within every row the column indices are strictly ascending.

The two representations convert losslessly to one another: a CsrGradient of shape (r, c) corresponds to a CooGradient over the flattened dim = r * c view using linear indices row * c + col. See CsrGradient::to_coo / CsrGradient::from_coo.

Sparse SGD

SparseSgd performs params[i] -= lr * grad[i] only for the non-zero coordinates i, optionally with (coupled) weight decay and momentum applied lazily on the touched coordinates. Untouched coordinates are left bit-identical.

Lazy Adam (the hard part)

SparseAdam maintains, per coordinate, the first/second moment estimates m[i], v[i] and last_step[i] – the global optimizer step at which the coordinate was last updated. Two precisely-defined semantics are supported, selected through LazyAdamMode:

LazyAdamMode::Lazy (pure TensorFlow LazyAdam)

When coordinate i is touched at global step t with gradient g, the moments are updated with the current gradient only, using a single decay factor regardless of how long the coordinate was dormant:

  m[i] = beta1 * m[i] + (1 - beta1) * g
  v[i] = beta2 * v[i] + (1 - beta2) * g^2

Bias correction uses the global step t:

  m_hat = m[i] / (1 - beta1^t)
  v_hat = v[i] / (1 - beta2^t)
  params[i] -= lr * m_hat / (sqrt(v_hat) + eps)

This is cheap and matches TensorFlow’s LazyAdam, but the moments do not account for the EMA decay that conceptually elapsed while the coordinate was dormant.

LazyAdamMode::DormancyDecay (lazy with dormancy catch-up)

Here we model exactly what dense Adam would have done to the moments had the gradient been zero during the dormant steps. If the coordinate was last updated at step s = last_step[i] and is now touched at step t, dense Adam would multiply the old moments by beta^(t - s) (one factor of beta per elapsed step, each carrying a zero gradient) before incorporating the new gradient. Hence:

  gap  = t - s                      (>= 1)
  m[i] = beta1^gap * m[i] + (1 - beta1) * g
  v[i] = beta2^gap * v[i] + (1 - beta2) * g^2

followed by the same global-step bias correction and update as above. For a coordinate touched on every step gap == 1, so this reduces exactly to the standard Adam recursion; for an intermittently-touched coordinate the moments m[i], v[i] at touch time are identical to those that dense Adam (fed explicit zero gradients on the dormant steps) would hold. The parameter trajectory still differs from dense Adam, because lazy Adam — by design — performs no parameter update on the dormant steps.

In both modes bias correction uses the global optimizer step count t (number of step calls so far), never a per-coordinate visit count.

Structs

CooGradient

Sparse gradient for a 1-D dense parameter / embedding-row view, stored as a coordinate list (COO).

CsrGradient

Sparse gradient for a 2-D embedding table, stored in compressed sparse row (CSR) form.

SparseAdam

Sparse / lazy Adam optimizer over a 1-D dense parameter vector.

SparseAdamConfig

Configuration for SparseAdam (and its 2-D table variant SparseAdamTable).

SparseAdamTable

Sparse / lazy Adam optimizer over a 2-D embedding table.

SparseSgd

Sparse SGD optimizer over a 1-D dense parameter vector.

SparseSgdConfig

Configuration for SparseSgd (and its 2-D table variant SparseSgdTable).

SparseSgdTable

Sparse SGD optimizer over a 2-D embedding table (CSR gradients).

Enums

LazyAdamMode

Semantics used by SparseAdam for catching up the moment estimates of a coordinate that has been dormant. See the module-level documentation for the exact update equations of each variant.