Backpropagation & Optimizers

Background

Rumelhart, Hinton, and Williams (1986) popularized backpropagation for training multi-layer networks
Based on the chain rule of calculus applied to computational graphs
Adam optimizer (Kingma & Ba, 2014) is now the default optimizer for most deep learning tasks

The Concept

Backpropagation computes gradients of the loss with respect to every weight in the network by applying the chain rule backward from the output. Optimizers (SGD, Adam, AdaGrad) use these gradients to update weights, iteratively reducing the loss.

How it Links

Chain rule is the mathematical backbone - connects to automatic differentiation in PyTorch/JAX
Learning rate scheduling (cosine annealing, warmup) builds on the basic optimizer framework
Vanishing/exploding gradients connect to initialization strategies, batch norm, and residual connections

Concrete Example

Training a network to classify photos: forward pass computes a wrong prediction (cat predicted as dog). Backprop asks: by how much did each weight in each layer contribute to that error? Each weight gets a gradient proportional to its blame. Adam updates each weight scaled by its recent gradient history.

Prerequisites

Multi-Layer Perceptron, Math you need: linear algebra & calculus, Bias, variance, over/under-fitting

Learning outcomes

Derive reverse-mode automatic differentiation as the chain rule applied right-to-left on a DAG, with cost equal to one forward pass per output.
State the SGD, momentum, Nesterov, Adagrad, RMSprop, Adam, and AdamW updates and explain when each is preferred.
Prove (sketch) Adam's bias correction and explain why Adam + L2 differs from AdamW.
Pick a learning-rate schedule (cosine annealing, warmup) and explain why warmup helps with Adam in transformers.
Apply gradient clipping and mixed-precision training, and recognize when L-BFGS or K-FAC is the right choice.

TL;DR

Backpropagation is reverse-mode automatic differentiation specialized to feed-forward DAGs [1]. It computes every parameter gradient in time proportional to one forward pass. The optimizer family layered on top is a tower of corrections: SGD is the baseline; momentum (Polyak 1964) [9] averages gradients to dampen noise; Nesterov (1983) [8] evaluates the gradient at the look-ahead point; Adagrad [6] rescales each parameter by its accumulated squared gradient; RMSprop [10] replaces the sum with an exponential moving average; Adam [2] combines momentum with RMSprop and adds bias correction; AdamW [3] decouples weight decay from the adaptive update. Cosine annealing [4] and warmup [5] are the standard schedules; gradient clipping and mixed precision [11] are the standard production fixes.

Intuition

Training a neural network is gradient descent on a high-dimensional, non-convex loss surface. Two pieces have to work together:

A way to compute the gradient. Backprop computes $\nabla_\theta \mathcal{L}$ for all parameters at once in a single sweep over the network, reusing intermediate results.
A rule for taking the next step. SGD steps directly along the negative gradient. Everything fancier (momentum, Adam, AdamW) is a different way to pick the step size and direction given a stream of stochastic gradient estimates.

The optimizers in this section all share the same skeleton: maintain some running statistics of past gradients (mean, second moment, both, or neither), use those statistics to scale the next step, and trade more state for faster convergence and more stable training.

The key insight behind reverse-mode autodiff

For a scalar loss $\mathcal{L}$ that depends on parameters through a DAG, reverse-mode autodiff (backprop) computes all $\partial \mathcal{L} / \partial \theta_i$ in a single backward pass over the graph. The cost is at most a small constant (typically 2 to 4) times the forward pass, regardless of the number of parameters. This is what makes training networks with billions of parameters feasible.

Forward-mode autodiff has the opposite trade-off: cheap per-input partial $\partial f / \partial x_j$ , but you need one pass per input. For ML where there is one scalar loss and millions of parameters, reverse mode wins by orders of magnitude.

Math

Reverse-mode autodiff: chain rule on a DAG

Setup

Model the computation as a directed acyclic graph $G = (V, E)$ where each node $v$ stores a value $a_v$ computed from its parents $\mathrm{pa}(v)$ via some elementary operation:

a_v = f_v\!\bigl( a_{u_1}, a_{u_2}, \dots \bigr), \qquad u_i \in \mathrm{pa}(v).

Designate one node as the loss $\mathcal{L} = a_{v^\star}$ and a subset of nodes as parameters $\theta$ . We want $\partial \mathcal{L} / \partial a_\theta$ for every parameter.

Forward pass

Evaluate nodes in topological order. Cache every $a_v$ . Time and memory are both $O(|V|)$ .

Backward pass

Define the adjoint $\bar{a}_v = \partial \mathcal{L} / \partial a_v$ . The adjoint of the loss node is $\bar{a}_{v^\star} = 1$ . For every other node $v$ , by the chain rule applied to the children $\mathrm{ch}(v)$ :

\bar{a}_v = \sum_{u \in \mathrm{ch}(v)} \bar{a}_u \cdot \frac{\partial f_u}{\partial a_v}\bigl( a_{\mathrm{pa}(u)} \bigr).

Process nodes in reverse topological order. For each node, fetch the adjoints of its children (already computed) and multiply by the local Jacobian. The local Jacobian is computed at the cached forward values, which is why the forward pass had to run first.

Cost

Every edge in $E$ is touched once in the forward pass and once in the backward pass. So time is $\Theta(|E|)$ for both passes, and the backward pass is at most a small constant times the forward pass (typically 2-4× because of cache pressure and slightly more arithmetic per local Jacobian).

Memory

The naive implementation stores every intermediate activation ( $O(|V|)$ memory). For very deep networks this is the dominant cost; checkpointing recomputes activations in segments to trade compute for memory, taking the cost from $O(L)$ memory and $O(L)$ time to $O(\sqrt{L})$ memory and $O(L)$ time. [14]

Why it works for arbitrary DAGs

Backprop is just reverse-mode autodiff specialized to feed-forward networks. For arbitrary DAGs (residual connections, attention, weight tying), the same algorithm applies: process in reverse topological order and accumulate the adjoint contributions from every child. Modern frameworks (PyTorch’s autograd, JAX’s grad, TensorFlow’s GradientTape) all implement this.

From SGD to Nesterov: the first-order family

Notation: parameters $\theta \in \mathbb{R}^p$ , stochastic gradient at step $t$ is $g_t = \nabla_\theta \mathcal{L}_t(\theta_t)$ , learning rate $\eta$ , momentum coefficient $\rho \in [0, 1)$ .

SGD

\theta_{t+1} = \theta_t - \eta\, g_t.

Pros: cheap (no extra state), well-understood theory. Cons: noisy on stochastic gradients, slow on ill-conditioned problems (the step size must be at most $1/L$ where $L$ is the largest curvature, but the slowest direction has curvature $\mu$ , so convergence rate is $O((1 - \mu/L)^t)$ ).

Polyak’s heavy-ball momentum

v_{t+1} = \rho\, v_t + g_t, \qquad \theta_{t+1} = \theta_t - \eta\, v_{t+1}.

The velocity $v$ is an exponential moving average of past gradients with effective horizon $1/(1-\rho)$ . This dampens oscillations across narrow valleys (the perpendicular components average to zero) while accumulating progress along the valley floor.

Convergence rate on quadratic objectives: $O((\sqrt{\kappa} - 1)^2 / (\sqrt{\kappa} + 1)^2)$ where $\kappa = L/\mu$ is the condition number, roughly $\sqrt{\kappa}$ times faster than vanilla GD on ill-conditioned quadratics [9]. Typical $\rho$ for deep learning: 0.9 (effective horizon 10 steps).

Nesterov’s accelerated gradient

Evaluate the gradient at the look-ahead point $\theta_t - \eta \rho v_t$ rather than at $\theta_t$ :

g_t = \nabla \mathcal{L}\bigl(\theta_t - \eta \rho v_t\bigr), \qquad v_{t+1} = \rho v_t + g_t, \qquad \theta_{t+1} = \theta_t - \eta v_{t+1}.

The look-ahead lets the algorithm “see” the curvature ahead and correct earlier. On smooth convex problems Nesterov’s method achieves the optimal $O(1/t^2)$ rate (vs $O(1/t)$ for plain GD) [8]. In practice on deep nets the win over plain momentum is small but consistent and free [7].

Adaptive optimizers: Adagrad, RMSprop, Adam, AdamW

The first-order family scales every parameter with the same learning rate $\eta$ . Adaptive methods give each parameter its own learning rate, scaled by past gradient magnitude.

Adagrad

G_t = G_{t-1} + g_t \odot g_t, \qquad \theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{G_t + \epsilon}} \odot g_t.

(Element-wise.) Each parameter’s effective learning rate decays like $1/\sqrt{t}$ along directions that have seen large gradients. Strong on sparse gradients (NLP with bag-of-words features, recommender embeddings) where rare features get to keep a large step. Convergence guarantees on online convex problems [6]. Weakness: $G_t$ grows unboundedly and the effective LR eventually dies.

RMSprop

Replace the sum with an exponential moving average:

v_t = \beta_2\, v_{t-1} + (1 - \beta_2)\, g_t \odot g_t, \qquad \theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{v_t + \epsilon}} \odot g_t.

This fixes the dying-LR problem. Hinton introduced it in lecture 6.5 of his 2012 Coursera class [10]; it never had a paper but is widely used.

Adam

Combine RMSprop (second-moment EMA) with momentum (first-moment EMA):

m_t = \beta_1 m_{t-1} + (1 - \beta_1)\, g_t,

v_t = \beta_2 v_{t-1} + (1 - \beta_2)\, g_t \odot g_t,

with bias correction (see below):

\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t},

and update

\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon}\, \hat{m}_t.

Defaults from Kingma & Ba [2]: $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\epsilon = 10^{-8}$ , $\eta = 10^{-3}$ . Adam is the default optimizer for most deep-learning workloads outside CV.

Why bias correction

At $t = 1$ , $m_1 = (1 - \beta_1) g_1 \approx 0.1\, g_1$ , biased toward zero. Computing $\hat{m}_t = m_t / (1 - \beta_1^t)$ undoes the bias. Without it, the early steps would be a factor of $1 - \beta_1^t \approx 0.1$ smaller than intended, making the network train too slowly for the first few hundred steps; over-time it self-corrects but the warm-up period is significant. The same logic applies to $\hat{v}_t$ .

AdamW: decoupled weight decay

The original Adam paper recommended adding L2 regularization through the gradient: $g_t \mathrel{+}= \lambda \theta_t$ before the moment updates. Loshchilov & Hutter showed this is wrong: the L2 term gets divided by $\sqrt{\hat{v}_t}$ along with everything else, so parameters with large second moments end up with effectively less weight decay than parameters with small second moments. The fix is to decouple weight decay from the gradient:

\theta_{t+1} = \theta_t - \eta\, \Bigl( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \;+\; \lambda\, \theta_t \Bigr).

The decay term is a uniform pull toward the origin, applied directly to $\theta$ without going through the adaptive scaling. AdamW [3] is the default for transformer training (GPT, BERT, ViT).

Schedules and tricks

A lot of “what makes training work” lives in the learning-rate schedule and a small set of stability tricks.

Cosine annealing with warm restarts

The cosine schedule (Loshchilov & Hutter 2016) [4] decays the learning rate from $\eta_{\max}$ to $\eta_{\min}$ as

\eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max} - \eta_{\min})\,\bigl(1 + \cos(\pi t / T)\bigr).

The variant SGDR resets $t$ at the end of each cycle (warm restart); each restart escapes flat regions and explores new local minima. Now the de facto schedule for image classification.

Linear warmup

Start at $\eta = 0$ and linearly ramp to $\eta_{\max}$ over the first ~5% of training. Standard for transformers and large-batch CV (Goyal et al. 2017 trained ImageNet in 1 hour on 256 GPUs by carefully scheduling warmup with batch size; see “linear scaling rule”) [5]. Empirically the largest stability benefit is for Adam in transformers, where the second-moment estimate $\hat{v}_t$ is noisy at the start and a too-large learning rate quickly diverges.

Gradient clipping

When a single batch produces a huge gradient (a step into a non-smooth region of the loss), clip its norm:

g \;\leftarrow\; g \cdot \min\!\Bigl( 1,\; \tau / \lVert g \rVert \Bigr).

Standard for RNNs (where exploding gradients are routine) and for transformers under mixed precision. Typical $\tau$ : 1.0 or 5.0.

Mixed precision

Store activations and gradients in fp16 (or bf16) but keep a master copy of weights and the optimizer state in fp32. Use a dynamic loss-scaling factor to keep small gradients out of the fp16 subnormal range. The result: ~2× memory and ~1.5-2× throughput at no measurable accuracy loss [11]. Implemented in torch.cuda.amp and tf.keras.mixed_precision.

Second-order methods

Briefly: when batch size is large enough that one batch is essentially the full gradient, second-order methods can win. Two examples:

L-BFGS [12]: maintains an implicit low-rank approximation of the inverse Hessian using the last $m$ pairs of $(\Delta \theta, \Delta g)$ . Memory $O(mp)$ . Used for batch optimization in scientific computing and sometimes for fine-tuning small models.
K-FAC [13]: approximates the Fisher information matrix as a Kronecker product of layer-local matrices, then preconditions the gradient. Sometimes useful on small networks but adds significant complexity. Mostly displaced in practice by Adam + careful schedules.

For most modern training, the first-order family (Adam/AdamW with cosine + warmup) is the right default; second-order methods are a niche tool.

Optimizer comparison

Optimizer	State per parameter	Adaptive LR?	Typical hyperparameters	When to use	Watch out for
SGD	0	no	lr 0.01-0.1	CV from scratch with momentum	Slow on ill-conditioned problems.
SGD + momentum	1 (velocity)	no	lr 0.01-0.1, momentum 0.9	ResNet-style CNNs, ImageNet baselines	Tune lr per architecture.
Nesterov	1	no	lr 0.01-0.1, momentum 0.9	Slight upgrade over momentum	Same as momentum.
Adagrad	1 (G_t)	yes (per-param)	lr 0.01-0.1	Sparse gradients (NLP bag-of-words)	LR dies as G_t grows; rarely used today.
RMSprop	1 (v_t EMA)	yes	lr 1e-3, beta2 0.99	RNNs and pre-Adam workloads	Mostly displaced by Adam.
Adam	2 (m_t, v_t)	yes	lr 1e-3, beta1 0.9, beta2 0.999	Default for transformers and most non-CV	Adam + L2 is not weight decay; use AdamW.
AdamW	2	yes	lr 1e-3 to 5e-4, weight_decay 0.01-0.1	Default for transformers, ViT, LLM finetuning	Decay weight is decoupled, not part of the gradient.

Implementations

NumPy implementations of SGD, momentum, and Adam side-by-side, then PyTorch’s torch.optim and JAX/Optax. Each version trains the same tiny MLP on XOR.

python

import numpy as np

# Each optimizer is a (state -> next state) pure function.

def sgd_step(params, grads, state, lr=0.1):
  return [p - lr*g for p, g in zip(params, grads)], state

def momentum_step(params, grads, state, lr=0.1, rho=0.9):
  if state is None: state = [np.zeros_like(p) for p in params]
  new_v = [rho*v + g for v, g in zip(state, grads)]
  new_p = [p - lr*v for p, v in zip(params, new_v)]
  return new_p, new_v

def adam_step(params, grads, state, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
  if state is None:
      state = ([np.zeros_like(p) for p in params],
               [np.zeros_like(p) for p in params],
               0)
  m, v, t = state
  t += 1
  new_m = [b1*mi + (1-b1)*g  for mi, g in zip(m, grads)]
  new_v = [b2*vi + (1-b2)*g*g for vi, g in zip(v, grads)]
  bc1 = 1 - b1**t
  bc2 = 1 - b2**t
  new_p = [p - lr * (mi/bc1) / (np.sqrt(vi/bc2) + eps)
           for p, mi, vi in zip(params, new_m, new_v)]
  return new_p, (new_m, new_v, t)

About 25 lines of pure NumPy. Each optimizer is a state transformer; this is essentially what Optax does in JAX.

Run a head-to-head numpy comparison of SGD, momentum, and Adam on a tiny MLP for XOR. Watch how Adam reaches a low loss in far fewer steps:

Idle (packages: numpy)

import numpy as np

def init(seed=0):
  rng = np.random.default_rng(seed)
  W1 = rng.normal(0, np.sqrt(2/2), size=(8, 2))
  b1 = np.zeros(8)
  W2 = rng.normal(0, np.sqrt(1/8), size=(1, 8))
  b2 = np.zeros(1)
  return [W1, b1, W2, b2]

def relu(z):  return np.maximum(0., z)
def sigm(z):  return 1./(1. + np.exp(-z))

def forward(X, p):
  Z1 = X @ p[0].T + p[1]
  H1 = relu(Z1)
  Z2 = H1 @ p[2].T + p[3]
  Y  = sigm(Z2)
  return Y, (X, Z1, H1, Z2, Y)

def backward(y, p, c):
  X, Z1, H1, Z2, Yhat = c
  n = X.shape[0]
  dZ2 = (Yhat - y) / n
  dW2 = dZ2.T @ H1
  db2 = dZ2.sum(0)
  dH1 = dZ2 @ p[2]
  dZ1 = dH1 * (Z1 > 0)
  dW1 = dZ1.T @ X
  db1 = dZ1.sum(0)
  return [dW1, db1, dW2, db2]

def loss(y, yhat):
  return float(-np.mean(y*np.log(yhat+1e-9) + (1-y)*np.log(1-yhat+1e-9)))

X = np.array([[0.,0],[0,1],[1,0],[1,1]])
y = np.array([[0.],[1],[1],[0]])

def train(opt_step, init_state, T=2000):
  p = init()
  s = init_state(p)
  losses = []
  for t in range(T):
      Yhat, c = forward(X, p)
      if t % 200 == 0:
          losses.append((t, loss(y, Yhat)))
      g = backward(y, p, c)
      p, s = opt_step(p, g, s)
  return losses

def sgd(p, g, s):
  return [pi - 0.5*gi for pi, gi in zip(p, g)], s

def mom(p, g, s):
  v = [0.9*vi + gi for vi, gi in zip(s, g)]
  return [pi - 0.5*vi for pi, vi in zip(p, v)], v

def adam(p, g, s):
  m, v, t = s
  t += 1
  m = [0.9*mi + 0.1*gi for mi, gi in zip(m, g)]
  v = [0.999*vi + 0.001*(gi*gi) for vi, gi in zip(v, g)]
  bc1, bc2 = 1 - 0.9**t, 1 - 0.999**t
  new_p = [pi - 1e-2 * (mi/bc1) / (np.sqrt(vi/bc2) + 1e-8)
           for pi, mi, vi in zip(p, m, v)]
  return new_p, (m, v, t)

print("step |   SGD     mom     Adam")
print("-----+--------------------------")
ls_sgd  = train(sgd,  lambda p: None)
ls_mom  = train(mom,  lambda p: [np.zeros_like(pi) for pi in p])
ls_adam = train(adam, lambda p: ([np.zeros_like(pi) for pi in p],
                                [np.zeros_like(pi) for pi in p], 0))
for (t, ls), (_, lm), (_, la) in zip(ls_sgd, ls_mom, ls_adam):
  print(f"{t:4d} | {ls:.4f}  {lm:.4f}  {la:.4f}")

(no output yet)

Interactive demo

A two-dimensional toy: gradient descent on $f(x) = x^2$ . Step the ball forward, change the learning rate, and watch the algorithm converge or oscillate. The same dynamics generalize to high-dimensional loss surfaces.

Gradient descent on f(x) = x²

Learning rate0.10x = 2.5000 f(x) = 6.2500 ∇f(x) = 5.0000

Real-world uses

Domain	Optimizer	Schedule	Why
Modern image classification (ResNet, EfficientNet)	SGD + momentum 0.9	Cosine annealing, warmup ~5%	Tradition + good generalization on CV; many papers report SGD beats Adam at final test accuracy.
Transformers (BERT, GPT)	AdamW	Linear warmup + linear or cosine decay	Adaptive LR is essential for the wide range of gradient scales across attention vs FFN layers; AdamW because plain Adam mishandles weight decay.
Vision Transformers and hybrids	AdamW	Warmup + cosine	Same reasoning as transformers; CNN tricks (SGD + momentum) do not transfer to ViTs.
LLM fine-tuning (full-parameter)	AdamW	Constant or short cosine	Fine-tuning is short enough that fancy schedules do not matter much; stability does.
LLM fine-tuning (LoRA / PEFT)	AdamW (only on adapter weights)	Constant	Most parameters are frozen; the optimizer state is small enough to ignore.
Reinforcement learning (policy gradient)	Adam	Constant	Noisy gradients benefit from adaptive scaling; weight decay is rarely needed when episodes provide implicit regularization.

Pitfalls

Adam + L2 is NOT what you want; use AdamW

Adding $\lambda \theta$ to the gradient before the moment update means the decay term gets divided by $\sqrt{\hat{v}_t}$ along with the rest. Parameters with large second moments end up with effectively less decay than parameters with small ones, the opposite of what L2 regularization is supposed to do. AdamW decouples the decay so it acts uniformly on $\theta$ . For BERT/GPT/ViT-class models the gap between Adam+L2 and AdamW is several percent on downstream tasks [3].

Learning-rate cliff

Too small a learning rate stalls (you watch the loss plateau, then declare the model “doesn’t work”); too large diverges (loss to infinity in a few steps). The sweet spot can span an order of magnitude or two. The cheapest way to find it: an LR-finder sweep (Smith 2017): exponentially ramp the LR from very small to very large over a single epoch, plot the loss, pick the largest stable LR.

Warmup is not optional for Adam in transformers

Without warmup, the first few hundred Adam steps have a noisy second-moment estimate $\hat{v}_t$ , and the resulting per-parameter step sizes are wildly miscalibrated. Combined with the typical large LR of $10^{-3}$ for transformers, this makes early training diverge frequently. A linear warmup from 0 to peak over the first ~5% of steps fixes this [5]. CNNs trained with SGD+momentum do not need warmup nearly as urgently because momentum is much less sensitive to early-step noise.

Momentum interacts with batch size

The “linear scaling rule” (Goyal et al. 2017) [5]: when you multiply batch size by $k$ , multiply the learning rate by $k$ and re-tune the warmup. Otherwise momentum overshoots the new gradient regime. This is the trick that lets large-batch training keep pace with small-batch baselines.

Gradient clipping after, not before, the LR schedule

The conventional order in PyTorch is: loss.backward() -> clip_grad_norm_ -> optimizer.step() -> scheduler.step(). Swapping clip_grad_norm_ with optimizer.step() clips the wrong quantity. This is a common bug in tutorial code.

Pop quiz

Check yourself

Score: 0 / 5(0 answered)

1. What is the asymptotic cost of one backward pass relative to one forward pass?
2. Why do we apply the bias correction m_hat = m / (1 - beta1^t) in Adam?
3. Why is AdamW preferred over Adam + L2 regularization?
4. Why does linear warmup help Adam in transformer training?
5. Which technique is most effective at preventing gradient explosion in RNNs?

Interview Q&A

Derive backprop on a 2-layer MLP from first principles.onsite

Forward pass. For input $x \in \mathbb{R}^{d_0}$ , weights $W_1 \in \mathbb{R}^{d_1 \times d_0}$ , $W_2 \in \mathbb{R}^{d_2 \times d_1}$ , and biases $b_1, b_2$ :

z_1 = W_1 x + b_1, \qquad h_1 = \sigma(z_1), \qquad z_2 = W_2 h_1 + b_2,

with loss $\mathcal{L}(z_2, y)$ .

Adjoints. Let $\delta_2 = \partial \mathcal{L} / \partial z_2$ . Then by the chain rule

\frac{\partial \mathcal{L}}{\partial W_2} = \delta_2\, h_1^\top, \qquad \frac{\partial \mathcal{L}}{\partial b_2} = \delta_2, \qquad \frac{\partial \mathcal{L}}{\partial h_1} = W_2^\top \delta_2.

For the hidden layer, $\delta_1 = \partial \mathcal{L} / \partial z_1 = (W_2^\top \delta_2) \odot \sigma'(z_1)$ , and

\frac{\partial \mathcal{L}}{\partial W_1} = \delta_1\, x^\top, \qquad \frac{\partial \mathcal{L}}{\partial b_1} = \delta_1.

Why this is reverse-mode autodiff. Each adjoint $\delta_\ell$ is computed from the next-layer adjoint $\delta_{\ell+1}$ and a local Jacobian (the activation derivative and the weight matrix transpose). Processed in reverse topological order, the cost is one local Jacobian-vector product per node, equal to the forward pass cost up to a constant. [1]

Derive Adam's update rule with bias correction.onsite

Maintain two EMAs:

m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t, \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t \odot g_t,

with $m_0 = v_0 = 0$ .

Bias derivation. Unroll: $m_t = (1 - \beta_1) \sum_{s=1}^{t} \beta_1^{t-s} g_s$ . If $g_s$ is approximately constant at value $\bar{g}$ , then $\mathbb{E}[m_t] = (1 - \beta_1^t) \bar{g}$ , biased toward zero by the factor $1 - \beta_1^t$ . Define $\hat{m}_t = m_t / (1 - \beta_1^t)$ to undo this; analogously $\hat{v}_t = v_t / (1 - \beta_2^t)$ .

Update.

\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon}\, \hat{m}_t.

The numerator $\hat{m}_t$ is a momentum-style EMA of past gradients; the denominator $\sqrt{\hat{v}_t}$ is a per-parameter normalization (RMSprop-style). Defaults: $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\eta = 10^{-3}$ , $\epsilon = 10^{-8}$ . [2]

Why does AdamW outperform Adam + L2 on transformer training?research

In Adam + L2, the regularizer is added to the gradient: $g_t \mathrel{+}= \lambda \theta_t$ . Then the moment updates and the divide by $\sqrt{\hat{v}_t}$ are applied. The effect of the decay on parameter $i$ is

- \eta\, \frac{\lambda \theta_{t,i}}{\sqrt{\hat{v}_{t,i}} + \epsilon}.

Parameters with large $\hat{v}_{t,i}$ (typically large outgoing weights, or attention queries/keys) see effectively smaller decay; parameters with small $\hat{v}_{t,i}$ see effectively larger decay. The intended uniform pull toward zero is corrupted into a non-uniform pull weighted by gradient magnitudes.

AdamW decouples by writing

\theta_{t+1,i} = \theta_{t,i} - \eta\, \frac{\hat{m}_{t,i}}{\sqrt{\hat{v}_{t,i}} + \epsilon} - \eta\, \lambda\, \theta_{t,i}.

The decay term is the same $- \eta \lambda \theta$ for every parameter. Empirically this matters: Loshchilov & Hutter and many subsequent papers show 0.5-2% absolute accuracy gains on ImageNet ViTs and language model perplexity. AdamW is the default for every transformer optimizer in HuggingFace and PyTorch out of the box [3].

Sketch SGD's convergence proof on a smooth convex objective.research

Assume $\mathcal{L}: \mathbb{R}^p \to \mathbb{R}$ is convex, $L$ -smooth ( $\lVert \nabla \mathcal{L}(x) - \nabla \mathcal{L}(y) \rVert \le L \lVert x - y \rVert$ ), and the stochastic gradient $g_t$ satisfies $\mathbb{E}[g_t] = \nabla \mathcal{L}(\theta_t)$ and $\mathbb{E}\lVert g_t \rVert^2 \le G^2$ .

For SGD with step size $\eta_t$ :

\mathbb{E}\lVert \theta_{t+1} - \theta^\star \rVert^2 = \mathbb{E}\lVert \theta_t - \theta^\star \rVert^2 - 2 \eta_t\, \mathbb{E}\langle g_t, \theta_t - \theta^\star \rangle + \eta_t^2 \mathbb{E}\lVert g_t \rVert^2.

By convexity, $\langle \nabla \mathcal{L}(\theta_t), \theta_t - \theta^\star \rangle \ge \mathcal{L}(\theta_t) - \mathcal{L}(\theta^\star)$ . Take expectations and rearrange:

\mathbb{E}[\mathcal{L}(\theta_t) - \mathcal{L}(\theta^\star)] \le \frac{\mathbb{E}\lVert \theta_t - \theta^\star \rVert^2 - \mathbb{E}\lVert \theta_{t+1} - \theta^\star \rVert^2}{2 \eta_t} + \frac{\eta_t G^2}{2}.

Sum from $t = 1$ to $T$ with $\eta_t = c / \sqrt{t}$ and apply Jensen on the average iterate:

\mathbb{E}[\mathcal{L}(\bar{\theta}_T) - \mathcal{L}(\theta^\star)] = O\!\bigl( G \lVert \theta_1 - \theta^\star \rVert / \sqrt{T} \bigr).

So SGD on smooth convex problems converges at rate $1/\sqrt{T}$ in function value. With strong convexity the rate improves to $1/T$ ; deep learning has neither, but the bound’s intuition (smaller step + decreasing $\eta_t$ implies smaller variance in the limit) carries over and motivates LR decay schedules [14].

Build it yourself

A short project that takes about an evening:

Implement the from-scratch numpy MLP from the previous page, plus the SGD / momentum / Adam optimizers above. Train on MNIST.
Reproduce the loss curves: SGD at $\eta = 0.5$ vs momentum at the same $\eta$ vs Adam at $\eta = 10^{-3}$ . Plot all three on one axis.
Add cosine annealing and a 5% warmup. Compare to constant LR.
Add gradient clipping at $\tau = 1.0$ and try to find an LR that diverges without it but converges with it.
Stretch: implement L-BFGS with a memory of $m = 5$ and run it on a small linear regression. Notice it converges in a handful of iterations, but the per-iteration cost is much higher than Adam.

References

Rumelhart, D. E., Hinton, G. E., and Williams, R. J. (1986). Learning representations by back-propagating errors. Nature 323, 533-536. doi:10.1038/323533a0
Kingma, D. P. and Ba, J. (2014). Adam: A Method for Stochastic Optimization. ICLR 2015 (arXiv preprint 2014). arXiv:1412.6980
Loshchilov, I. and Hutter, F. (2017). Decoupled Weight Decay Regularization. ICLR 2019 (arXiv preprint 2017). arXiv:1711.05101
Loshchilov, I. and Hutter, F. (2016). SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR 2017 (arXiv preprint 2016). arXiv:1608.03983
Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. (2017). Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. arXiv preprint. arXiv:1706.02677
Duchi, J., Hazan, E., and Singer, Y. (2011). Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research 12, 2121-2159. https://jmlr.org/papers/v12/duchi11a.html
Sutskever, I., Martens, J., Dahl, G., and Hinton, G. (2013). On the importance of initialization and momentum in deep learning. Proc. 30th International Conference on Machine Learning (ICML), 1139-1147. https://proceedings.mlr.press/v28/sutskever13.html
Nesterov, Y. (1983). A method for solving the convex programming problem with convergence rate O(1/k^2). Doklady Akademii Nauk SSSR 269(3), 543-547. https://www.mathnet.ru/eng/dan46009
Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics 4(5), 1-17. doi:10.1016/0041-5553(64)90137-5
Tieleman, T. and Hinton, G. (2012). Lecture 6.5 - rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning. https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., and Wu, H. (2017). Mixed Precision Training. ICLR 2018 (arXiv preprint 2017). arXiv:1710.03740
Liu, D. C. and Nocedal, J. (1989). On the limited memory BFGS method for large scale optimization. Mathematical Programming 45(1), 503-528. doi:10.1007/BF01589116
Martens, J. and Grosse, R. (2015). Optimizing Neural Networks with Kronecker-factored Approximate Curvature. ICML 2015. arXiv:1503.05671
Goodfellow, I., Bengio, Y., and Courville, A. (2016). Deep Learning. MIT Press. https://www.deeplearningbook.org

Interactive Demo

Mode:

Click "Step Forward" to begin

Learning rate: 0.01Example weight update: 0.005000

Mode: ForwardActive layer: none

Tip: Click Forward to watch activations propagate, then Backward to see gradients flow in reverse

Multi-Layer Perceptron Regularization in Deep Nets