class

SparseAdam

extendsOptimizer

SparseAdam(params: Iterable[Parameter] | Iterable[dict[str, object]], lr: float = 0.001, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-08, maximize: bool = False)

source edit

Lazy Adam for gradients that touch a few rows of a parameter.

The variant of Adam meant for embedding tables: each step updates the moments and the values of only the rows the gradient touches — a row being a slice along the first dimension. A row absent from a batch keeps its moments and its value; it does not keep moving on its old momentum the way it would under lucid.optim.Adam.

For each touched row, with the parameter's step count $t$ (one count per parameter, advanced by every step in which it has a gradient):

\begin{aligned} m &\leftarrow m + (1 - \beta_1)(g - m) \\ v &\leftarrow v + (1 - \beta_2)(g^2 - v) \\ \theta &\leftarrow \theta - \frac{\alpha \sqrt{1 - \beta_2^t}}{1 - \beta_1^t} \cdot \frac{m}{\sqrt{v} + \epsilon} \end{aligned}

Parameters

paramsiterable of Parameter or iterable of dict

Parameters to optimise, or a list of parameter-group dicts.

lrfloat= 0.001

Learning rate

\alpha

(default: 1e-3).

betastuple of float= (0.9, 0.999)

Coefficients

(\beta_1, \beta_2)

for the first- and second-moment estimates (default: (0.9, 0.999)).

epsfloat= 1e-08

Term

\epsilon

for numerical stability (default: 1e-8).

maximizebool= False

If True, maximise the objective with respect to params by negating the gradient on the touched rows before updating the moments and parameter values (default: False).

Attributes

param_groupslist of dict

Parameter groups with keys "params", "lr", "betas", "eps", and "maximize".

defaultsdict

Default hyperparameter values.

statedict

The live state of each parameter that has stepped, by flat parameter index: "step" (int), "exp_avg" and "exp_avg_sq" (tensors shaped like the parameter). It is what state_dict saves and load_state_dict restores.

Notes

Lucid's gradients are dense — nn.Embedding(sparse=True) is not supported — so a row counts as touched when its gradient holds a nonzero value. The reference framework takes only a sparse gradient and updates every row it lists, so the two differ for a row that was looked up but whose gradient is exactly zero: the reference decays its moments, Lucid leaves it alone. Within a touched row every element is updated, zero or not, as the reference's embedding gradient lists whole rows.

A parameter whose gradient is None is skipped and its step count does not advance; one whose gradient is all zero advances its count and changes nothing else, as in the reference.

The moments follow their parameter through a conversion (module.to(), .double(), .half()) between steps.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> import lucid.optim as optim
>>> embedding = nn.Embedding(4, 2)
>>> optimizer = optim.SparseAdam(embedding.parameters(), lr=0.1)
>>> before = embedding.weight.detach().clone()
>>> embedding(lucid.tensor([0, 2])).sum().backward()
>>> optimizer.step()
>>> (embedding.weight != before).any(dim=1).tolist()
[True, False, True, False]
>>> optimizer.state[0]["step"]
1

Used by 2

Constructors

dunder

init

→None

__init__(params: Iterable[Parameter] | Iterable[dict[str, object]], lr: float = 0.001, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-08, maximize: bool = False)

source edit

Initialise the SparseAdam. See the class docstring for parameter semantics.

Instance methods

step

→Tensor or None

step(closure: _OptimizerClosure = None)

source edit

Perform a single SparseAdam optimisation step.

Each parameter with a gradient advances its step count and updates the rows its gradient touches (see the class notes); parameters whose gradient is None are skipped.

Parameters

closurecallable= None

A closure that re-evaluates the model and returns the loss. If provided, it is called before the parameter updates.

Returns

Tensor or None

The loss returned by closure, or None if no closure was provided.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> import lucid.optim as optim
>>> embedding = nn.Embedding(100, 8)
>>> optimizer = optim.SparseAdam(embedding.parameters(), lr=1e-3)
>>> optimizer.zero_grad()
>>> loss = embedding(lucid.tensor([[3, 17, 3]])).sum()
>>> loss.backward()
>>> optimizer.step()

>>> import lucid >>> import lucid.nn as nn >>> import lucid.optim as optim >>> embedding = nn.Embedding(4, 2) >>> optimizer = optim.SparseAdam(embedding.parameters(), lr=0.1) >>> before = embedding.weight.detach().clone() >>> embedding(lucid.tensor([0, 2])).sum().backward() >>> optimizer.step() >>> (embedding.weight != before).any(dim=1).tolist() [True, False, True, False] >>> optimizer.state[0]["step"] 1

>>> import lucid >>> import lucid.nn as nn >>> import lucid.optim as optim >>> embedding = nn.Embedding(100, 8) >>> optimizer = optim.SparseAdam(embedding.parameters(), lr=1e-3) >>> optimizer.zero_grad() >>> loss = embedding(lucid.tensor([[3, 17, 3]])).sum() >>> loss.backward() >>> optimizer.step()