Deep Learning Optimization Algorithms

An optimizer turns gradients into parameter updates. The gradient tells us the local direction of increasing loss, while the optimizer determines the step using the learning rate and, in adaptive methods, statistics from earlier gradients.

Batch, stochastic, and mini-batch descent

Full-batch descent uses every training example per update. With deterministic computations, it removes noise from sampling a batch, but a large learning rate can still make updates unstable. Stochastic descent uses one example; mini-batch descent uses a subset to balance efficient matrix operations with frequent updates. For independently sampled examples, shuffling each epoch is a useful default; preserve required sequence structure for time-dependent data. Do not silently discard the final incomplete batch unless that is intended.

For a per-example loss, a batch \(B_t\) uses \(g_t=|B_t|^{-1}\sum_{i\in B_t}\nabla_\theta\mathcal L_i(\theta_{t-1})\). One pass through all training examples is an epoch; it usually contains several updates. With five examples and batch size two, the batches have sizes 2, 2, and 1, so the last gradient is averaged over one example, not two.

This one-epoch example fits \(\hat y=wx\) to \(y=2x\) with mean squared error and no intercept. Examples occupy columns. Each batch recomputes predictions using the latest weight, so sequential mini-batch updates are not equivalent to one full-batch update.

import numpy as np

X = np.array([[1., 2., 3., 4., 5.]])
Y = 2*X
w = np.zeros((1, 1))
order = np.random.default_rng(0).permutation(X.shape[1])
batch_size, lr = 2, 0.01
for step, start in enumerate(range(0, X.shape[1], batch_size), 1):
    indices = order[start:start+batch_size]
    xb, yb = X[:, indices], Y[:, indices]
    error = w.T @ xb-yb
    dw = 2*(xb @ error.T)/xb.shape[1]
    w = w-lr*dw
    print("step", step, "size", xb.shape[1], "w", round(w.item(), 6))
# step 1 size 2 w 0.68
# step 2 size 2 w 0.9044
# step 3 size 1 w 0.992048

The shuffled order is fixed for reproducibility. This illustrates batching and updates, not convergence or performance on new data. A partial batch still triggers an update; when reporting an epoch-wide mean loss, weight batch means by their numbers of examples. Changing batch size also changes updates per epoch, so record whether schedules and budgets use epochs, examples, or optimizer steps.

How to find a learning rate and shape it over training is covered in Learning Rate Schedules and Warmup in Deep Learning.

Exponentially weighted averages

Let \(\theta\) be the parameter vector, \(g_t\) the gradient evaluated at the parameters before update \(t\), and \(\alpha\) the learning rate. The index \(t\) starts at 1. An exponentially weighted average uses \(v_t=\beta v_{t-1}+(1-\beta)x_t\), where \(0\leq\beta<1\) controls how much history remains. Starting from \(v_0=0\), the observation weights sum to \(1-\beta^t\); dividing by that sum gives \(\hat v_t=v_t/(1-\beta^t)\). This removes the zero-initialization effect. If the observations have a common mean, the corrected estimate has that mean in expectation; it need not equal the current observation or track a changing mean without delay.

Momentum

Momentum averages gradients with \(v_t=\beta v_{t-1}+(1-\beta)g_t\), then updates \(\theta_t=\theta_{t-1}-\alpha v_t\). Consistent gradient components reinforce one another while alternating components can cancel. Depending on the learning rate and momentum setting, this can reduce oscillation and improve progress. We use the normalized-average convention here; implementations that use \(v_t=\beta v_{t-1}+g_t\) require a corresponding learning-rate rescaling for comparison. With both buffers initialized to zero and a fixed \(\beta\), the summed-gradient convention needs \(\alpha_{\rm sum}=(1-\beta)\alpha_{\rm average}\) to produce the same updates. For \(g_1=[2,20]\), \(v_0=0\), \(\beta=0.9\), and \(\alpha=0.1\), we get \(v_1=[0.2,2]\) and subtract \([0.02,0.2]\) from the parameters. Momentum can temporarily move against the current negative gradient: if a scalar buffer is 2 and the new gradient is −1, the next buffer is \(0.9(2)+0.1(-1)=1.7\). Subtracting it increases the current loss to first order. Averaging history does not guarantee descent at every step.

RMSprop

RMSprop averages squared gradients: \(s_t=\beta s_{t-1}+(1-\beta)g_t^2\), then updates \(\theta_t=\theta_{t-1}-\alpha g_t/(\sqrt{s_t}+\epsilon)\). Squares, roots, and division act coordinate by coordinate; \(\epsilon>0\) stabilizes division near zero. A large denominator reduces the effective learning rate, but the actual update also depends on the numerator. With the same \(g_1=[2,20]\), \(s_0=0\), and \(\beta=0.9\), the squared-gradient average is \([0.4,40]\). With \(\alpha=0.1\) and negligible \(\epsilon\), RMSprop subtracts about \([0.3162,0.3162]\). The normalization offsets the tenfold difference between the gradient components in this first step.

Adam

Adam maintains a weighted average of gradients, \(v_t\), and of squared gradients, \(s_t\), with separate coefficients \(\beta_1\) and \(\beta_2\). With stationary gradients these estimate first and second raw moments; during training they are moving summaries of a changing gradient stream, and \(s_t\) is not the gradient variance. Starting both arrays at zero, Adam corrects their weights and updates \(\theta_t=\theta_{t-1}-\alpha\hat v_t/(\sqrt{\hat s_t}+\epsilon)\). The following runnable example performs one update. Both state arrays have the same shape as the parameters.

import numpy as np

param = np.array([1.0, 1.0])
grad = np.array([2.0, 20.0])
v = np.zeros_like(param)
s = np.zeros_like(param)
beta1, beta2 = 0.9, 0.999
learning_rate, epsilon, t = 0.1, 1e-8, 1

v = beta1 * v + (1 - beta1) * grad
s = beta2 * s + (1 - beta2) * grad**2
v_hat = v / (1 - beta1**t)
s_hat = s / (1 - beta2**t)
update = learning_rate * v_hat / (np.sqrt(s_hat) + epsilon)
param = param - update
print("corrected gradient average:", v_hat)
print("corrected squared-gradient average:", s_hat)
print("update:", np.round(update, 6))
print("parameters:", np.round(param, 6))
# corrected gradient average: [ 2. 20.]
# corrected squared-gradient average: [  4. 400.]
# update: [0.1 0.1]
# parameters: [0.9 0.9]

AdamW: decoupled weight decay

For an L2 penalty \(\lambda\lVert\theta\rVert_2^2/2\), the penalty gradient is \(\lambda\theta\). In plain SGD without momentum, adding that gradient is equivalent to multiplying the old parameters by \(1-\alpha\lambda\) as part of the update. In coupled Adam, the modified gradient enters both moment estimates. Its effect cannot be read from the adaptive denominator alone or represented as a separate uniform shrinkage.

AdamW keeps the loss gradient in the moment estimates and applies decay directly to the old parameters:

\[\theta_t=\theta_{t-1}-\alpha\frac{\hat v_t}{\sqrt{\hat s_t}+\epsilon}-\alpha\lambda\theta_{t-1}\]

import numpy as np

def adam_step(param, grad, v, s, t, lr, decay=0.0, decoupled=False,
              beta1=0.9, beta2=0.999, epsilon=1e-8):
    g = grad if decoupled else grad + decay * param
    new_v = beta1 * v + (1 - beta1) * g
    new_s = beta2 * s + (1 - beta2) * g**2
    v_hat = new_v / (1 - beta1**t)
    s_hat = new_s / (1 - beta2**t)
    update = lr * v_hat / (np.sqrt(s_hat) + epsilon)
    base = (1 - lr * decay) * param if decoupled else param
    return base - update, new_v, new_s

param = np.array([1.0, 1.0])
grad = np.array([2.0, 20.0])
zeros = np.zeros_like(param)
for label, decay, decoupled in (("Adam", 0.0, False),
                               ("Adam + L2", 0.1, False),
                               ("AdamW", 0.1, True)):
    result, _, _ = adam_step(param, grad, zeros, zeros, 1, 0.1,
                             decay=decay, decoupled=decoupled)
    print(label, np.round(result, 6))
# Adam [0.9 0.9]
# Adam + L2 [0.9 0.9]
# AdamW [0.89 0.89]

In this first-step example, coupled L2 changes both moments, yet normalization leaves the printed update almost unchanged because both gradient components stay positive and \(\epsilon\) is small. AdamW additionally subtracts 0.01 from each old parameter. The function computes new arrays without changing the supplied state, so all three runs start from identical values. Applying decay to the already updated parameter would also shrink the adaptive step and would not match the displayed equation.

Decoupling removes decay from the moment estimates, but its per-step strength remains \(\alpha\lambda\). Learning-rate schedules and the number of steps therefore still affect cumulative decay. Excluding biases and normalization parameters is a common configuration to try, not a proof that regularizing them can never help. Compare configurations under the intended validation protocol.

To use adam_step in training, initialize the two moment arrays once, compute the loss gradient at the current parameters, assign all three returned arrays, and increment t once per optimizer update. Carry the moments and counter across batches and epochs. Resetting them each step repeatedly applies the first-step rule and discards Adam’s history. Adam stores two moment values per parameter, which also matters for memory use.

When resuming a run, restore optimizer moments and step counters along with model parameters and the schedule state. Loading only weights starts a different optimization trajectory.

Choosing an optimizer

Setting	Candidates and checks
Convolutional vision model	Compare momentum SGD and AdamW with separately tuned learning rates and schedules.
Transformer	AdamW is a common starting point; use a recipe matched to the model and batch size.
Small dense model, initial experiment	Adam at \(\alpha=10^{-3}\) is one trial value; search smaller and larger values.
Sparse inputs or embeddings	Check the gradient representation. PyTorch sparse gradient tensors require a supported optimizer such as SparseAdam; sparse input data alone do not imply sparse gradients.
Fine-tuning a pretrained model	Start from the model’s documented recipe; full-model and adapter-only training can require different learning rates.

There is no fixed ordering of these optimizers across tasks. Validation quality, time to reach it, and the tuning budget all matter. A comparison that tunes one optimizer carefully and uses defaults for another does not isolate the effect of the update rule.

Learning-rate schedules

Optimizer comparisons depend on the learning-rate schedule as well as the update rule. Warmup gradually raises the rate during the initial updates; decay reduces it later. A cosine decay phase gives a smooth reduction, though schedules with restarts raise it again. These are tunable choices, not guarantees of stable training. Use the linked schedule guide for formulas and match the comparison budget to the question: equal updates and equal wall-clock time need not rank methods the same way.

How to allocate a tuning budget across these choices is the subject of A Hyperparameter Tuning Strategy That Fits Your Budget.

How to compare optimizers fairly

• Use the same initialization, data order, batch size, and evaluation schedule.
• Tune learning rates separately for each optimizer.
• Compare validation quality and time or updates to reach it.
• Keep loss reduction and regularization conventions consistent.
• Monitor effective update size, not only raw gradients.

Use the comparison to select an optimizer for the architecture, data, regularization, and budget at hand. Repeat promising configurations with more than one seed before treating a small validation gap as a reliable advantage.

How these models are pretrained at scale is covered in LLM Pretraining: Objectives, Data, and Scaling Laws.

References

• Polyak (1964). Some Methods of Speeding Up the Convergence of Iteration Methods. USSR Computational Mathematics and Mathematical Physics.
• Tieleman and Hinton (2012). Lecture 6.5 – RMSprop. COURSERA: Neural Networks for Machine Learning.
• Sutskever, Martens, Dahl, and Hinton (2013). On the Importance of Initialization and Momentum in Deep Learning. ICML.
• Kingma and Ba (2015). Adam: A Method for Stochastic Optimization. ICLR.
• Loshchilov and Hutter (2019). Decoupled Weight Decay Regularization. ICLR.