Optimizers

Abstract base class for all optimizers.

Concrete optimizers (SGD, Adam, ...) extend this class and implement step(). The base class owns the parameter groups, the per-parameter state map, gradient zeroing and checkpointing through stateDict() / loadStateDict().

Parameters are split into groups, and every group carries its own options (for example a different learning rate per layer). Options a group does not set are taken from the optimizer defaults, and every group is range checked when it is added and again at each step().

A parameter with requiresGrad = false (for example one frozen with Module.freezeParameters()) stays in its group but is skipped by step(). A trainable parameter whose gradient is null (it took no part in the loss, or no backward() has run yet) is skipped as well, as in PyTorch, and gets no optimizer state. A failed step() (a non-finite gradient, an invalid option) changes neither the parameters nor the optimizer state.

AdaDelta optimizer (Zeiler, 2012).

AdaDelta is an extension of Adagrad that avoids its monotonically shrinking learning rate. It keeps exponential moving averages of the squared gradients and of the squared parameter updates, and scales each step by the ratio of their root-mean-squares. The update follows torch.optim.Adadelta:

v = rho * v + (1 - rho) * g^2delta = sqrt(acc + eps) / sqrt(v + eps) * gacc = rho * acc + (1 - rho) * delta^2theta -= lr * delta

When weightDecay is non-zero, weightDecay * theta is added to the gradient (L2 penalty) before the update.

Adagrad (Adaptive Gradient Algorithm) optimizer.

Adagrad adapts the learning rate for each parameter based on the historical sum of squared gradients. Parameters with larger gradients receive smaller effective learning rates, while parameters with smaller gradients receive larger effective learning rates. The update follows torch.optim.Adagrad:

clr = lr / (1 + (t - 1) * lrDecay)sum += g^2theta -= clr * g / (sqrt(sum) + eps)

When weightDecay is non-zero, weightDecay * theta is added to the gradient (L2 penalty) before the update.

Adam (Adaptive Moment Estimation) optimizer (Kingma and Ba, 2015).

Computes adaptive learning rates for each parameter by maintaining running averages of both the gradients and their squared values. The update follows torch.optim.Adam:

m = beta1 * m + (1 - beta1) * gv = beta2 * v + (1 - beta2) * g^2theta -= (lr / (1 - beta1^t)) * m / (sqrt(v / (1 - beta2^t)) + eps)

When weightDecay is non-zero, weightDecay * theta is added to the gradient (classic L2 penalty); use AdamW for decoupled weight decay. With amsgrad, v is replaced by its running maximum in the denominator.

Adamax optimizer, a variant of Adam based on the infinity norm (Kingma and Ba, 2015).

Often used for embeddings and sparse gradients. The update follows torch.optim.Adamax:

m = beta1 * m + (1 - beta1) * gu = max(beta2 * u, |g| + eps)theta -= (lr / (1 - beta1^t)) * m / u

When weightDecay is non-zero, weightDecay * theta is added to the gradient (L2 penalty) before the update.

AdamW (Adam with decoupled weight decay) optimizer (Loshchilov and Hutter, 2019).

AdamW applies weight decay directly to the parameters instead of adding it to the gradient, so the decay is not rescaled by the adaptive denominator. The update follows torch.optim.AdamW:

theta *= 1 - lr * weightDecaym = beta1 * m + (1 - beta1) * gv = beta2 * v + (1 - beta2) * g^2theta -= (lr / (1 - beta1^t)) * m / (sqrt(v / (1 - beta2^t)) + eps)

Averaged Stochastic Gradient Descent (ASGD) optimizer (Polyak and Juditsky, 1992).

Runs SGD with a decaying step size and keeps a running (Polyak-Ruppert) average of the iterates, which often generalizes better than the last iterate. The average starts after t0 steps and is available through ASGD.averagedParameters. The update follows torch.optim.ASGD:

theta = theta * (1 - lambda * eta) - eta * gax += mu * (theta - ax)eta = lr / (1 + lambda * lr * t)^alphamu = 1 / max(1, t - t0)

When weightDecay is non-zero, weightDecay * theta is added to the gradient (L2 penalty) before the update. With maximize, the gradient is negated first.

LAMB (Layer-wise Adaptive Moments optimizer for Batch training) optimizer.

LAMB extends Adam with layer-wise adaptive learning rates, enabling training with very large batch sizes (up to 64K for BERT). It combines the benefits of Adam's per-element adaptivity with LARS-style layer-wise trust ratios. Each parameter tensor is treated as one layer:

r = m_hat / (sqrt(v_hat) + eps) + weightDecay * thetatrust = ||theta|| / ||r||        (1 when either norm is zero)theta -= lr * trust * r

where m_hat and v_hat are the bias-corrected Adam moment estimates.

Reference: "Large Batch Optimization for Deep Learning: Training BERT in 76 Minutes" (You et al., 2019)

LARS (Layer-wise Adaptive Rate Scaling) optimizer.

LARS scales the learning rate per layer based on the ratio of the parameter norm to the gradient norm. This allows training with very large batch sizes (up to 32K) without divergence. Each parameter tensor is treated as one layer:

localLr = eta * ||theta|| / (||g|| + weightDecay * ||theta|| + eps)   (1 when either norm is zero)m = momentum * m + lr * localLr * (g + weightDecay * theta)theta -= m

Reference: "Large Batch Training of Convolutional Networks" (You et al., 2017)

L-BFGS (Limited-memory Broyden–Fletcher–Goldfarb–Shanno) optimizer.

A quasi-Newton method that approximates the inverse Hessian using a limited history of past updates. Unlike first-order methods (Adam, SGD), LBFGS can converge much faster for smooth, deterministic objectives.

Important: LBFGS requires a closure that reevaluates the model and returns the loss. This is because LBFGS may evaluate the function multiple times per step. All parameters are treated as one flat vector, so it uses a single set of hyperparameters: parameter groups with differing options are rejected.

Each step(closure) runs up to maxIter L-BFGS iterations. The step length is chosen by one of two line searches: - lineSearchFn: null (default): the initial step lr is halved (at most 10 times) until the loss decreases. The first step is scaled by min(1, 1 / ||g||_1). maxEval is checked between iterations, so one line search can exceed it by up to nine evaluations. - lineSearchFn: "strong_wolfe": a line search satisfying the strong Wolfe conditions, as in torch.optim.LBFGS. This is the better choice for non-quadratic objectives.

LBFGS only supports float32/float64 parameters on the host; device tensors are rejected with a DeviceError.

Lion (EvoLved Sign Momentum) optimizer.

Introduced by Chen et al. (2023) in "Symbolic Discovery of Optimization Algorithms". Lion is a simpler, more memory-efficient alternative to Adam that keeps a single momentum buffer and applies sign-based updates, so every element moves by the same magnitude lr regardless of its gradient scale.

Update rule (per parameter):

u = sign(beta1 * m + (1 - beta1) * g)theta -= lr * (u + weightDecay * theta)m = beta2 * m + (1 - beta2) * g

Weight decay is decoupled, as in AdamW. Lion typically needs a 3-10x smaller learning rate than Adam (suggested: 1e-4 to 3e-4, versus Adam's 1e-3) and a correspondingly larger weight decay.

Nadam (Nesterov-accelerated Adam) optimizer (Dozat, 2016).

Combines Adam's adaptive learning rates with Nesterov momentum: the first moment estimate is replaced by a look-ahead blend of the next momentum estimate and the current gradient. The momentum coefficient follows the schedule mu_t = beta1 * (1 - 0.5 * 0.96^(t * momentumDecay)) used by torch.optim.NAdam.

By default weightDecay is an L2 penalty added to the gradient. With decoupledWeightDecay: true the parameters are instead multiplied by 1 - lr * weightDecay before the update (NAdamW).

RAdam (Rectified Adam) optimizer.

Fixes Adam's early training variance issue by computing a variance rectification term. When the variance is tractable (high enough SMA), uses the adaptive learning rate; otherwise falls back to SGD with momentum. The update follows PyTorch's torch.optim.RAdam:

rhoInf = 2 / (1 - beta2) - 1rhoT   = rhoInf - 2 t beta2^t / (1 - beta2^t)rhoT > 5:  param -= lr * rect * sqrt(1 - beta2^t) * mHat / (sqrt(v) + eps)otherwise: param -= lr * mHat                      (mHat = m / (1 - beta1^t))

With the default beta2 = 0.999 the first five steps use the SGD-with-momentum branch.

RMSprop (Root Mean Square Propagation) optimizer.

RMSprop adapts the learning rate for each parameter by dividing by a running average of recent gradient magnitudes. This helps with non-stationary objectives and is particularly effective for RNNs.

The update follows PyTorch's torch.optim.RMSprop:

v = alpha * v + (1 - alpha) * g^2avg = centered ? sqrt(v - gAvg^2) + eps : sqrt(v) + eps      (gAvg = EMA of g)buf = momentum * buf + g / avg                               (momentum > 0)param = param - lr * (momentum > 0 ? buf : g / avg)

Rprop (Resilient Backpropagation) optimizer.

Only uses the sign of the gradient, not its magnitude. Step sizes adapt individually per parameter based on sign changes. The update follows PyTorch's torch.optim.Rprop (Rprop- with weight reset on a sign change): for every element, with prev the previous gradient,

g * prev > 0:  step = min(step * etaPlus,  stepMax);  param -= sign(g) * stepg * prev < 0:  step = max(step * etaMinus, stepMin);  param unchanged, prev = 0otherwise:     param -= sign(g) * step

The step sizes start at lr (clamped to [stepMin, stepMax]). Because the learning rate only seeds the step sizes, changing it later (for example with an LR scheduler) does not affect parameters that already have state.

Stochastic Gradient Descent (SGD) optimizer.

Implements vanilla SGD with optional momentum, weight decay, and Nesterov acceleration. The update follows PyTorch's torch.optim.SGD:

d = grad + weightDecay * parambuf = momentum * buf + (1 - dampening) * d      (buf = d on the first step)d = nesterov ? d + momentum * buf : bufparam = param - lr * d

With maximize: true the gradient sign is flipped and the objective is maximized.

SparseAdam optimizer, a variant of Adam designed for sparse gradients.

Only updates the moment estimates for gradient indices that are non-zero, making it suited to parameters with sparse gradient updates such as embedding layers. Gradients are dense tensors here, and an exact zero marks an index that was not touched. The update follows PyTorch's torch.optim.SparseAdam. Weight decay is not supported.

Algorithm: For each parameter element where gradient ≠ 0:

m_t = β₁ * m_{t-1} + (1 - β₁) * g_tv_t = β₂ * v_{t-1} + (1 - β₂) * g_t²θ_t = θ_{t-1} - lr * √(1 - β₂^t) / (1 - β₁^t) * m_t / (√v_t + ε)

For indices where gradient = 0, moment estimates and parameters are unchanged. The step counter t advances on every call to step().