3. Calculus and Optimization

Derivatives and local approximation

A derivative describes how a function changes under a small change in its input. For a scalar function,

\[f'(x)=\lim_{h\to0}\frac{f(x+h)-f(x)}h, \qquad f(x+h)=f(x)+f'(x)h+o(|h|).\]

The notation \(o(|h|)\) means an error whose ratio to \(|h|\) tends to zero. The derivative supplies the best first-order, linear approximation near \(x\).

The notation \(O(\|h\|^p)\) instead means an error bounded by a constant times \(\|h\|^p\) for sufficiently small \(h\). For sequences, \(O(1/T)\) describes a bound up to a fixed multiplicative constant as \(T\) grows.

Vectors are columns throughout. A partial derivative varies one coordinate while holding the others fixed. For a differentiable \(f:\mathbb R^d\to\mathbb R\), the gradient collects these derivatives:

\[\nabla f(x)= \begin{bmatrix}\partial f/\partial x_1\\\vdots\\\partial f/\partial x_d\end{bmatrix}, \qquad f(x+h)=f(x)+\nabla f(x)^\top h+o(\|h\|).\]

Unless stated otherwise, \(\|\cdot\|\) is the Euclidean norm. Matrix norms are induced Euclidean operator norms unless specified otherwise. The inner product \(\nabla f(x)^\top h\) is the predicted change in the function. When the gradient is nonzero, the unit direction \(-\nabla f(x)/\|\nabla f(x)\|\) gives the greatest first-order decrease.

For a differentiable vector-valued map \(F:\mathbb R^d\to\mathbb R^m\), the derivative is represented by the Jacobian:

\[(J_F)_{ij}=\frac{\partial F_i}{\partial x_j}, \qquad J_F\in\mathbb R^{m\times d}, \qquad dF=J_F\,dx.\]

The local approximation is \(F(x+h)=F(x)+J_F(x)h+o(\|h\|)\). Differentiability requires it to work for every sufficiently small perturbation. Existence of coordinate partial derivatives alone is insufficient; continuous partial derivatives in a neighborhood are sufficient.

The Jacobian uses numerator layout: output coordinates index rows and input coordinates index columns. For a scalar output, the derivative \(Df\) is a row, while the gradient is its transpose:

\[Df(x)=\nabla f(x)^\top, \qquad df=\nabla f(x)^\top dx.\]

For a scalar loss depending on a matrix \(W\), its gradient has the same shape as \(W\). The defining identity is

\[d\ell=\langle\nabla_W\ell,dW\rangle_F =\sum_{i,j}\frac{\partial\ell}{\partial W_{ij}}\,dW_{ij}.\]

The same entrywise pairing applies to higher-order tensors. The general derivative-array convention is developed in Linear Algebra.

The chain rule and backpropagation

If \(G:\mathbb R^d\to\mathbb R^m\) and \(F:\mathbb R^m\to\mathbb R^p\), the chain rule is

\[J_{F\circ G}(x)=J_F(G(x))J_G(x).\]

For a scalar loss \(\ell\) depending on \(z=G(x)\), this becomes

\[\nabla_x\ell=J_G(x)^\top\nabla_z\ell.\]

Thus a gradient at the output of an operation can be propagated to its input. Backpropagation repeatedly applies this rule through a network. If a value is used along several paths, the contributions from those paths add.

Example: one layer with a tanh activation

Let \(x\in\mathbb R^d\) be an input, \(W\in\mathbb R^{m\times d}\) the weight matrix, \(b\in\mathbb R^m\) a bias, and \(y\in\mathbb R^m\) a fixed target. Define

\[z=Wx+b, \qquad a=\tanh z, \qquad \ell=\frac12\|a-y\|^2.\]

Here \(z\) is the pre-activation, the vector before applying the nonlinearity; \(a\) is the output. The function \(\tanh\) acts on each coordinate separately.

First differentiate the loss with respect to the output:

\[\nabla_a\ell=a-y.\]

Next propagate through \(a_i=\tanh z_i\). Since \(da_i/dz_i=1-a_i^2\), define

\[\boxed{ \delta:=\nabla_z\ell, \qquad \delta_i=\frac{\partial\ell}{\partial z_i} =(a_i-y_i)(1-a_i^2). }\]

The vector \(\delta\) measures the sensitivity of the loss to each pre-activation. It combines the output residual \(a-y\) with the local slope of \(\tanh\). In vector notation, \(\delta=(a-y)\odot(1-a\odot a)\), where \(\odot\) means entrywise multiplication.

Finally, propagate through \(z=Wx+b\):

\[\nabla_W\ell=\delta x^\top, \qquad \nabla_b\ell=\delta, \qquad \nabla_x\ell=W^\top\delta.\]

For instance, \(\partial z_i/\partial W_{ij}=x_j\), so \(\partial\ell/\partial W_{ij}=\delta_i x_j\). The matrix gradient therefore has shape \(m\times d\), matching \(W\). When training this layer, \(W,b\) are parameters and \(x\) is held fixed; the input gradient is useful when propagating further into an earlier layer.

Values move forward through the three operations, and gradients of \(\ell\) move back along the same edges. The single vector \(\delta\) yields the weight, bias, and input gradients.

Automatic differentiation

Automatic differentiation (AD) evaluates derivatives by composing the derivative rules of a program’s elementary operations. Addition, multiplication, matrix products, and nonlinear functions each supply a local rule; the chain rule connects them. The resulting derivative concerns the real-valued computation represented by the program, evaluated using its computed intermediate values.

Finite differences infer a derivative from nearby function evaluations and introduce a perturbation size and truncation error. Symbolic differentiation transforms expressions into derivative expressions. AD propagates derivative information through a computation while retaining its intermediate structure. It can therefore differentiate a long composition without expanding it into a single formula, and has no finite-difference truncation error. These distinctions and the two accumulation modes are developed in Baydin et al.’s Automatic Differentiation in Machine Learning: a Survey.

For the truncated logistic map \(l_4\), symbolic differentiation of the closed form produces a long expression, a finite difference only approximates the derivative, and AD transforms the program itself and is exact up to rounding. Figure 2 of Baydin et al.

Andrej Karpathy’s The spelled-out intro to neural networks and backpropagation: building micrograd (2 h 25 min) develops these ideas by implementing a scalar reverse-mode engine and training a small network. The accompanying lecture notebook and micrograd code expose the graph, local derivative rules, and gradient accumulation directly.

Computational graphs and local derivative rules

A computational graph records which operations depend on which values. Nodes may represent scalars, vectors, or tensors. A node is evaluated only after its inputs are available; this gives a topological order. An executed finite loop can be represented by successive copies of its operations. A shared intermediate appears once but may supply several later operations.

For a node \(z=\phi(a,b)\), the local differential is

\[dz=J_{\phi,a}\,da+J_{\phi,b}\,db.\]

For scalar multiplication, \(z=ab\), this reads \(dz=b\,da+a\,db\). Every differentiation mode uses this same rule. What changes is the information propagated through it and the direction of traversal. The ordinary values computed by the program are called primal values, distinguishing them from derivative information.

Tangents travel with the evaluation, from inputs to output; adjoints travel against it and accumulate with \(\mathrel{+}=\). Both directions use the same local Jacobians \(J_{\phi,a}\) and \(J_{\phi,b}\).

Forward mode: sensitivity to a chosen input change

Suppose a program computes \(F(x)\). In addition to its value, we may want the first-order effect of a particular change to its inputs. For \(F:\mathbb R^d\to\mathbb R^m\), specify that change by a direction \(v\in\mathbb R^d\) and consider the path \(x(\tau)=x+\tau v\). The scalar \(\tau\) measures how far we move along this path. For instance, \(v=(1,-1)^\top\) means increasing the first input and decreasing the second by the same amount.

Every intermediate value \(z\) then changes with \(\tau\). Its tangent is its rate of change at the original input:

\[\dot z=\left.\frac{d}{d\tau}z(x+\tau v)\right|_{\tau=0}.\]

Equivalently, \(z(x+\tau v)=z(x)+\tau\dot z+o(|\tau|)\). Thus \(z\) records the current value and \(\dot z\) records its sensitivity to the chosen input change. All dots refer to this same \(\tau\); \(\dot z\) has the same shape as \(z\). It is \(\tau\dot z\) that approximates the actual change in \(z\).

Forward-mode AD computes the pair \((z,\dot z)\) at each operation. Initialize \(\dot x=v\), set the tangents of fixed constants to zero, and use the local chain rule in the ordinary evaluation order:

\[\dot z=J_{\phi,a}\dot a+J_{\phi,b}\dot b.\]

For a scalar product \(z=ab\), this becomes \(\dot z=b\dot a+a\dot b\): changes in either factor contribute. The intermediate tangents let each operation pass this sensitivity to the next. At the output,

\[\dot F=J_F(x)v\in\mathbb R^m.\]

This Jacobian–vector product (JVP) supplies the approximation \(F(x+\tau v)\approx F(x)+\tau J_F(x)v\). AD obtains its coefficient from local derivative rules, without evaluating the program at a nearby input or choosing a finite-difference step size. The direction \(v\) specifies the sensitivity question; it need not be a unit vector or an optimizer update.

Here \(f(x_1,x_2)=x_1x_2+\sin x_1\), with intermediate values \(w_1,\dots,w_5\). Tangents flow upward from the seeds \(\dot w_1,\dot w_2\), which choose the input direction; \(\dot w_1=1\), \(\dot w_2=0\) gives \(\partial f/\partial x_1\).

Choosing \(v=e_j\) isolates the \(j\)th input and returns the \(j\)th Jacobian column. For scalar \(f\), one direction gives the scalar \(\dot f=\nabla f(x)^\top v\); separate coordinate seeds recover the full gradient. Forward mode is useful when only a few input directions matter, including sensitivities to a scalar parameter or matrix-free derivative products; Choosing between forward and reverse mode below gives concrete cases.

An ordinary forward pass evaluates values such as activations and a loss. Forward-mode differentiation additionally propagates tangents. A usual neural-network training step evaluates values forward and uses reverse mode to obtain all parameter gradients of its scalar loss; it does not require a separate forward tangent for every parameter.

Reverse mode: sensitivity of a chosen output

After computing the primal values, choose an output seed \(u\in\mathbb R^m\) and the scalar quantity \(u^\top F(x)\), with \(u\) fixed. Reverse mode asks how each input affects this chosen output quantity. An intermediate’s adjoint \(\bar z\) is the gradient of that scalar objective with respect to \(z\), represented using the same shape as \(z\). Initialize the output adjoint to \(u\), all other adjoints to zero, and traverse the graph in reverse order. The node \(z=\phi(a,b)\) contributes

\[\bar a\mathrel{+}=J_{\phi,a}^\top\bar z, \qquad \bar b\mathrel{+}=J_{\phi,b}^\top\bar z.\]

The final input adjoint is

\[\bar x=J_F(x)^\top u\in\mathbb R^d.\]

The usual name vector–Jacobian product (VJP) refers to \(u^\top J_F\); with column vectors, the returned adjoint represents its transpose \(J_F^\top u\). For a scalar loss, seeding \(\bar f=1\) returns the entire column gradient \(\nabla f(x)\).

The addition in \(\mathrel{+}=\) is essential. When a value feeds several operations, each outgoing path contributes to its sensitivity. Overwriting one contribution with another loses part of the chain rule. Backpropagation through a neural network is this reverse accumulation applied to its loss computation.

For the same function, reverse mode starts from the seed \(\bar f=\bar w_5=1\) at the output. Because \(x_1\) feeds both the sine and the product, its adjoint is the sum of two contributions, \(\bar w_1^a\) and \(\bar w_1^b\).

A shared intermediate, differentiated both ways

Consider the scalar program

\[s=x_1x_2,\qquad r=\sin s,\qquad f(x)=s+r.\]

Its graph makes the two uses of \(s\) explicit:

flowchart LR
    X1["x₁"] --> S["s = x₁x₂"]
    X2["x₂"] --> S
    S --> R["r = sin s"]
    S --> F["f = s + r"]
    R --> F

At \(x=(2,3)^\top\), choose \(v=(1,-1)^\top\). The question is how \(f\) changes along \((x_1,x_2)=(2+\tau,3-\tau)\). Seed \(\dot x_1=1\) and \(\dot x_2=-1\). Forward mode evaluates

\[\begin{aligned} s&=6,&\dot s&=3(1)+2(-1)=1,\\ r&=\sin6,&\dot r&=\cos6\,\dot s=\cos6,\\ f&=6+\sin6,&\dot f&=1+\cos6. \end{aligned}\]

The current output is \(f(x)\approx5.72058\) and its tangent is \(\dot f\approx1.96017\). In particular,

\[f(2+\tau,3-\tau) =f(2,3)+(1+\cos6)\tau+o(|\tau|).\]

A small positive \(\tau\) therefore increases the output by approximately \(1.96017\tau\). Each intermediate tangent was needed to compute this final slope through the product, sine, and addition.

The horizontal axis moves along the chosen input direction. Forward AD computes the slope of the orange tangent line at \(\tau=0\). Multiplying that slope by a small displacement predicts the output change; the blue curve shows the actual output along the path.

For the reverse pass, start with \(\bar f=1\). The final addition contributes \(1\) to both \(\bar s\) and \(\bar r\). The sine operation then adds \(\cos6\,\bar r\) to \(\bar s\), giving \(\bar s=1+\cos6\). Finally, the product contributes

\[\nabla f(x)=\bar x =(1+\cos6)\begin{bmatrix}3\\2\end{bmatrix}.\]

The two answers agree through \(\dot f=\bar x^\top v\): forward mode answers one directional-sensitivity question, while reverse mode gives all input sensitivities of this scalar output in one reverse traversal.

Backpropagation in PyTorch

For the tanh layer above, automatic differentiation recovers the same \(\delta\), weight gradient, and bias gradient:

import torch

W = torch.tensor([[0.4, -0.2], [0.1, 0.3]],
                 dtype=torch.float64, requires_grad=True)
b = torch.zeros(2, dtype=torch.float64, requires_grad=True)
x = torch.tensor([1., 2.], dtype=torch.float64)
y = torch.tensor([0.2, -0.1], dtype=torch.float64)

z = W @ x + b
a = torch.tanh(z)
loss = 0.5*((a - y)**2).sum()
grad_W, grad_b, delta = torch.autograd.grad(loss, (W, b, z))

assert torch.allclose(delta, (a - y)*(1 - a*a))
assert torch.allclose(grad_W, torch.outer(delta, x))
assert torch.allclose(grad_b, delta)
print("gradient with respect to z:", delta)

torch.tensor creates a tensor from data. The argument dtype=torch.float64 requests double precision, and requires_grad=True tells PyTorch’s automatic differentiation engine, autograd, to record the operations applied to that tensor. torch.zeros, torch.tanh, and torch.outer behave like their NumPy counterparts, and torch.allclose compares tensors within a tolerance. torch.autograd.grad runs reverse mode from loss and returns derivatives with respect to the requested inputs, here including the intermediate z. A usual training loop instead calls loss.backward(), which adds each parameter’s gradient to its .grad field. At a nondifferentiable point, such as the kink of ReLU, a framework uses a specified convention; the returned value does not establish the existence of an ordinary derivative there. Appendix I lists every library function used in the chapter’s examples, with links to their documentation.

PyTorch records these nodes for the code above; they can be read from loss.grad_fn and its next_functions. No node is recorded for x or y, which do not require gradients. loss.backward() adds the results into W.grad and b.grad; torch.autograd.grad returns them instead.

Choosing between forward and reverse mode

Both modes follow the program’s own operations, so their cost is tied to the cost of evaluating the function once. Write \(\operatorname{cost}(F)\) for the number of elementary operations in one evaluation of \(F\). Each local derivative rule costs at most a small constant multiple of the operation it differentiates. Consequently:

• one forward-mode pass computes \(F(x)\) together with one Jacobian–vector product \(J_F(x)v\in\mathbb R^m\);
• one reverse-mode pass computes \(F(x)\) and then one vector–Jacobian product \(J_F(x)^\top u\in\mathbb R^d\);
• each costs \(c\cdot\operatorname{cost}(F)\) operations for a modest constant \(c\). Baydin et al. report \(c<6\), typically between 2 and 3.

Neither pass forms the Jacobian. They differ in which question one pass answers. A forward pass answers “how do all outputs change when the input moves along the one direction \(v\)?” A reverse pass answers “how does the one output combination \(u^\top F\) depend on every input?” Building the full \(m\times d\) Jacobian therefore takes \(d\) forward passes, one per input coordinate, each giving a column; or \(m\) reverse passes, one per output coordinate, each giving a row.

One forward pass returns one column of the \(m\times d\) Jacobian, and one reverse pass returns one row. For a scalar loss the Jacobian is a single row, so one reverse pass returns the whole gradient.

Example. A network whose loss depends on \(d=10^6\) parameters is a map \(\mathbb R^{10^6}\to\mathbb R\). One reverse pass returns the entire gradient at the cost of a few loss evaluations. Forward mode would need \(10^6\) passes, one per parameter. This is why neural networks are trained with reverse mode, that is, backpropagation.

Memory. Reverse mode runs backward through the computation, so it needs the intermediate values of the forward evaluation; storing all of them can take memory proportional to the length of the computation. Checkpointing stores only selected intermediate values and recomputes the others during the backward pass, trading extra work for less memory; see Appendix C. Forward mode carries each tangent alongside its value and can discard intermediates as soon as they have been used.

When forward mode is used. Forward mode is the natural choice when only a few input directions matter:

• Sensitivity to a few parameters. The change in an entire simulated trajectory, a rendered image, or a whole vector of predictions when one physical constant or one hyperparameter changes is obtained from one forward pass per parameter. Franceschi et al. (2017) compute hyperparameter gradients this way when there are few hyperparameters and a long training run.
• Matrix-free solvers. Newton–Krylov and related methods for systems \(F(x)=0\) need only products \(J_F(x)v\), never the full Jacobian.
• Hessian–vector products. Applying forward mode to a reverse-mode gradient computation gives \(H(x)v\) for a few times the cost of a gradient, without forming the \(d\times d\) Hessian (Pearlmutter, 1994). Such products drive Newton-type methods and are used to measure the curvature of neural-network losses. Appendix C contains a PyTorch version.
• Long computations with few inputs. Because it need not store intermediate values, forward mode can be used when the memory required by reverse mode is prohibitive and the number of input directions is small.

Appendix C contains explicit NumPy traversals of both modes, the tensor-valued rules, PyTorch’s jvp and vjp interfaces, and a cost table.

Curvature and Taylor’s theorem

The Hessian of a scalar function is the Jacobian of its gradient:

\[H(x)=\nabla^2f(x), \qquad H_{ij}(x)=\frac{\partial}{\partial x_j} \left(\frac{\partial f}{\partial x_i}\right).\]

If the second partial derivatives are continuous in a neighborhood, mixed partials agree and \(H(x)\) is symmetric. The second directional derivative along \(v\) is \(v^\top H(x)v\). Positive eigenvalues describe upward curvature, negative eigenvalues downward curvature, and eigenvalues near zero weak curvature in the corresponding directions. A symmetric matrix \(H\) is positive semidefinite, written \(H\succeq0\), when \(v^\top Hv\ge0\) for every \(v\). It is positive definite, written \(H\succ0\), when this inequality is strict for every nonzero \(v\). The notation \(A\preceq B\) means \(B-A\succeq0\).

The three Hessians share their eigenvectors; darker regions have smaller values of \(\tfrac12h^\top Hh\). Along an eigenvector with eigenvalue \(\lambda\) the value is \(\tfrac12\lambda\|h\|^2\), so positive eigenvalues curve upward, a negative eigenvalue curves downward, and a zero eigenvalue leaves a flat direction.

The quadratic approximation

A function is \(C^k\) when its derivatives through order \(k\) exist and are continuous. For a \(C^2\) scalar function near \(x\),

\[f(x+h)=f(x)+\nabla f(x)^\top h +\frac12h^\top H(x)h+o(\|h\|^2).\]

The gradient supplies the linear change; the Hessian supplies its first curvature correction. For a quadratic function this expansion is exact, with no remainder.

The second-order model of \(f(x)=\sin x+x/2\) at \(x=0.6\) follows the curve farther than the tangent does. On logarithmic axes (right), halving the step divides the first-order error by about \(4\) and the second-order error by about \(8\).

The fundamental theorem of calculus applied along \(x+th\) also gives

\[f(x+h)-f(x)=\int_0^1\nabla f(x+th)^\top h\,dt.\]

This identity connects derivative bounds to approximation errors and underlies the descent lemma below. The segment must stay inside the domain. Appendix A develops arbitrary-order Taylor formulas, the integral remainder and its proof, convergence of Taylor series, and tensor-valued expansions.

Optimization problems and notation

The rest of the chapter is about minimizing a function. It addresses three questions:

1. What is a minimizer, and how can one be recognized? This section defines minimizers; Optimality conditions relates them to gradients and Hessians.
2. What makes a function easy or hard to minimize? Convexity, Lipschitz conditions, and curvature introduces the properties that the guarantees rely on.
3. How fast do specific methods make progress? The sections from Gradient descent and step sizes onward analyze gradient descent, Newton’s method, stochastic gradients, and adaptive optimizers under those properties.

In the layer example, \(x\) denoted the input. From here onward, \(x\) denotes the parameters being optimized; a neural network can be represented by concatenating its trainable tensors. An optimization problem has the form

\[\min_{x\in C} f(x),\]

where \(C\) is the feasible set; this set is unrelated to the smoothness classes \(C^k\). An unconstrained problem has \(C=\mathbb R^d\). A global minimizer \(x_\star\) satisfies \(f(x_\star)\le f(x)\) for every \(x\in C\); a local minimizer satisfies this only for feasible points in some neighborhood. A strict minimizer gives a strict inequality for other points in that neighborhood.

The infimum \(f_{\inf}=\inf_{x\in C}f(x)\) is the greatest lower bound on the objective values; it need not be attained. For instance, \(e^x\) on \(\mathbb R\) has infimum zero and no minimizer. Even when an algorithm’s objective values converge to the infimum, its iterates may diverge.

A function can have strict and non-strict local minimizers besides its global one (left). On \(\mathbb R\), \(e^x\) approaches its infimum \(0\) without attaining it (right).

A continuous objective attains its minimum on a nonempty compact feasible set. Existence on unbounded domains requires additional assumptions; see Appendix B.

When a minimum is attained, write \(f_\star=f(x_\star)\), \(R_0=\|x_0-x_\star\|\), and \(\Delta_0=f(x_0)-f_\star\). For nonconvex stationarity bounds, only a finite lower bound \(f_{\inf}\) is needed; there we use \(\Delta_0=f(x_0)-f_{\inf}\).

Optimization notation

The convergence statements and optimizer equations use the following conventions. An update with index \(t=0\) starts at \(x_0\) and produces \(x_1\). Each symbol is also defined where it is first used, so the table is for reference rather than memorization.

Bounds on gradients, gradient noise, and function values are separate assumptions.

Convexity, Lipschitz conditions, and curvature

Three properties determine how much a gradient method can guarantee:

• Convexity makes local information global. A tangent never lies above a convex function, so a point with zero gradient is a global minimizer.
• Smoothness with constant \(L\) bounds how fast the gradient can change. It determines how large a step can be trusted: roughly \(1/L\).
• Strong convexity with constant \(\mu\) makes the function curve upward at least as fast as a parabola. The size of the gradient then indicates how far the current point is from the minimizer.

The ratio \(\kappa=L/\mu\) of the two curvature bounds will govern the speed of gradient descent.

Convex and strongly convex functions

A set \(C\) is convex if \((1-t)x+ty\in C\) whenever \(x,y\in C\) and \(0\le t\le1\). A function on a convex domain is convex if

\[f((1-t)x+ty)\le(1-t)f(x)+tf(y).\]

Its graph lies below every chord. Equivalently, its epigraph \(\{(x,a):f(x)\le a\}\) is convex. The function is strictly convex if the chord inequality is strict whenever \(x\ne y\) and \(0<t<1\). For differentiable \(f\), ordinary convexity is equivalent to

\[f(y)\ge f(x)+\nabla f(x)^\top(y-x).\]

Thus tangent affine functions are global lower bounds. For \(C^2\) functions on an open convex domain, convexity is equivalent to \(\nabla^2f(x)\succeq0\) everywhere.

A function is \(\mu\)-strongly convex, with \(\mu>0\), if \(x\mapsto f(x)-\tfrac\mu2\|x\|^2\) is convex. In the differentiable case,

\[f(y)\ge f(x)+\nabla f(x)^\top(y-x) +\frac\mu2\|y-x\|^2.\]

For \(C^2\) functions this is equivalent to \(\nabla^2f(x)\succeq\mu I\). Strong convexity implies strict convexity and hence uniqueness of a minimizer when one exists. Strict convexity alone supplies no uniform positive curvature: \(x^4\) is strictly convex but has zero second derivative at zero and is not strongly convex on \(\mathbb R\).

Let \(f\) be differentiable and \(\mu\)-strongly convex on \(\mathbb R^d\), with minimizer \(x_\star\) and \(f_\star=f(x_\star)\). For every \(x\),

\[\frac\mu2\|x-x_\star\|^2 \le f(x)-f_\star \le\frac{\|\nabla f(x)\|^2}{2\mu}.\]

These inequalities say that three measures of being far from optimal control one another: the distance \(\|x-x_\star\|\), the objective gap \(f(x)-f_\star\), and the gradient size \(\|\nabla f(x)\|\). Both follow from the strong-convexity inequality above, used in two different ways.

Left inequality: a small gap forces a small distance. Use the strong-convexity inequality with the minimizer as base point, replacing \(x\) by \(x_\star\) and \(y\) by \(x\). Since \(\nabla f(x_\star)=0\), the linear term vanishes:

\[f(x)\ge f_\star+\frac\mu2\|x-x_\star\|^2.\]

The function lies above a parabola of curvature \(\mu\) centered at the minimizer.

Right inequality: a small gradient forces a small gap. Now fix \(x\) and read the strong-convexity inequality as a lower bound that holds for every \(y\):

\[f(y)\ge q_x(y):=f(x)+\nabla f(x)^\top(y-x)+\frac\mu2\|y-x\|^2.\]

Take the minimum of both sides over \(y\). The left side’s minimum is \(f_\star\). The quadratic \(q_x\) is minimized where its gradient \(\nabla f(x)+\mu(y-x)\) vanishes, at \(y=x-\nabla f(x)/\mu\), and its minimum value is \(f(x)-\|\nabla f(x)\|^2/(2\mu)\). Hence \(f_\star\ge f(x)-\|\nabla f(x)\|^2/(2\mu)\), which rearranges to the right inequality.

A convex function lies below its chords and above its tangents, and its epigraph is convex (left). For a strongly convex function with \(\mu=0.5\) (right), the quadratic lower bound \(q_x\) has minimum \(f(x)-\|\nabla f(x)\|^2/(2\mu)\), which lies below \(f_\star\): this is the right inequality.

Combining the two inequalities gives \(\|x-x_\star\|\le\|\nabla f(x)\|/\mu\). The right inequality has two uses later. It justifies stopping when the gradient is small, because for a strongly convex objective a small gradient certifies a small objective gap. It is also the step that turns the descent lemma into geometric convergence of gradient descent. Without strong convexity, a small gradient need not mean a small gap. The convex function \(f(x)=\sqrt{1+x^2}/100\) has \(|f'(x)|<0.01\) everywhere, yet \(f(1000)-f_\star\approx9.99\).

The geometric definitions and their equivalences are developed in Boyd and Vandenberghe’s Convex Optimization, Chapters 2–3. The arguments here use the Euclidean norm; changing the norm changes the constants and the associated geometry.

Three different Lipschitz assumptions

A Lipschitz bound limits how fast a quantity can change with the input. The three conditions below apply such a bound to the function value, the gradient, and the Hessian, respectively.

For \(f(y)=\sin 1.2y+0.3y\) and the orange point \(x\), the smallest valid constants are \(G=1.5\), \(L=1.44\), and \(\rho=1.728\). Each assumption confines the graph to its shaded region: a cone around the value \(f(x)\), a quadratic band around the tangent, or a cubic band around the second-order Taylor model.

In optimization, smooth usually means a Lipschitz gradient, with a stated constant \(L\). Merely being \(C^\infty\) does not imply any global Lipschitz bound. The function \(x^4\) is \(C^\infty\) but has unbounded curvature on \(\mathbb R\).

On an open convex domain, a differentiable function with \(\|\nabla f(x)\|\le G\) is \(G\)-Lipschitz, by integration along segments; the converse follows from directional derivatives. For a \(C^2\) function, \(L\)-smoothness is equivalent to

\[\|\nabla^2f(x)\|\le L, \quad\text{or equivalently}\quad -LI\preceq\nabla^2f(x)\preceq LI.\]

It does not require convexity. Combining smoothness and strong convexity gives

\[\mu I\preceq\nabla^2f(x)\preceq LI, \qquad \kappa=\frac L\mu\ge1,\]

where \(\kappa\) is a curvature condition number.

Examples. The function \(|x|\) is globally 1-Lipschitz and nonsmooth. The function \(x^2/2\) is globally 1-smooth and 1-strongly convex, but is not globally Lipschitz as a function. The function \(\sin x\) has a 1-Lipschitz gradient but is not convex. A quadratic has a constant Hessian, so its Hessian Lipschitz constant is zero.

Always attach a domain to a Lipschitz claim. The function \(x^4\) is smooth on each bounded interval with a constant depending on that interval. A strongly convex function on all of \(\mathbb R^d\) cannot also be globally Lipschitz as a function. Bounded-gradient assumptions in constrained strongly convex results therefore concern the relevant feasible region.

Optimality conditions

For an unconstrained differentiable function, an interior local minimum must satisfy \(\nabla f(x_\star)=0\). For \(C^2\) functions, a necessary second-order condition is \(\nabla^2f(x_\star)\succeq0\). Conversely,

\[\nabla f(x_\star)=0, \qquad \nabla^2f(x_\star)\succ0\]

imply a strict local minimum. A positive semidefinite Hessian alone is inconclusive: \(x^4\), \(-x^4\), and \(x^3\) all have zero first and second derivatives at zero but different local behavior.

A stationary point has zero gradient. It can be a minimum, maximum, or saddle. A negative Hessian eigenvalue at a stationary point gives a direction of decrease; an indefinite Hessian gives both increasing and decreasing directions. For a differentiable convex function, every stationary point is globally optimal.

On the graph of \(z=x^2-y^2\) the gradient vanishes at the origin, and the Hessian \(\operatorname{diag}(2,-2)\) is indefinite: the surface curves upward along \(x\) and downward along \(y\).

Over a convex feasible set \(C\), the first-order condition for a differentiable convex objective is

\[\nabla f(x_\star)^\top(x-x_\star)\ge0 \qquad\text{for every }x\in C.\]

The gradient need not vanish at a constrained optimum. For example, minimizing \(f(x)=x\) over \(x\ge0\) gives \(x_\star=0\) with \(f'(x_\star)=1\).

Gradient descent and step sizes

An iterative method chooses a direction \(p_t\) and a step size \(\eta_t>0\):

\[x_{t+1}=x_t+\eta_t p_t.\]

A descent direction satisfies \(\nabla f(x_t)^\top p_t<0\). Differentiability ensures decrease for sufficiently small steps along such a direction. It does not say that a unit step decreases the function, or that a sequence of decreasing values reaches the global minimum.

Gradient descent uses \(p_t=-\nabla f(x_t)\). Its update also minimizes a local linear model with a quadratic movement penalty:

\[x_{t+1}=\arg\min_z\left\{ \nabla f(x_t)^\top(z-x_t)+\frac{1}{2\eta_t}\|z-x_t\|^2 \right\}.\]

To see why, complete the square:

\[\nabla f(x_t)^\top(z-x_t)+\frac{1}{2\eta_t}\|z-x_t\|^2 =\frac{1}{2\eta_t}\bigl\|z-\bigl(x_t-\eta_t\nabla f(x_t)\bigr)\bigr\|^2-\frac{\eta_t}{2}\|\nabla f(x_t)\|^2.\]

The last term does not depend on \(z\), so the minimizer is \(z=x_t-\eta_t\nabla f(x_t)\), the gradient step. Equivalently, adding the constant \(f(x_t)\), gradient descent replaces \(f\) near \(x_t\) by the quadratic model

\[\widehat f_t(z)=f(x_t)+\nabla f(x_t)^\top(z-x_t)+\frac{1}{2\eta_t}\|z-x_t\|^2\]

and moves to the model’s minimizer. The model has the correct value and slope at \(x_t\), and it curves upward equally in every direction, with curvature \(1/\eta_t\). The step size is therefore the reciprocal of an assumed curvature: a small step assumes that \(f\) bends sharply and moves cautiously. If the assumed curvature is at least the true curvature bound \(L\), the model lies above \(f\), as the descent lemma below shows, so moving to its minimizer cannot increase \(f\). Newton's method uses the same idea with the Hessian in place of the uniform curvature \(1/\eta_t\).

Each dashed model matches the value and slope of \(f(x)=\log\cosh x+0.1x^2\) at \(x_t\) and has curvature \(1/\eta\); the gradient step moves to its minimizer (open circle). Here \(L=1.2\). With \(\eta\le1/L\) the model lies above \(f\) and \(f\) decreases; with \(\eta=2.5/L\) the step lands above \(f(x_t)\).

The descent lemma

Smoothness makes the upper-model picture precise. If the gradient cannot change faster than \(L\), then \(f\) cannot bend away from its tangent faster than a parabola of curvature \(L\), in either direction.

Theorem. If \(f\) has an \(L\)-Lipschitz gradient on a convex domain, then for any segment in that domain,

\[\left|f(y)-f(x)-\nabla f(x)^\top(y-x)\right| \le\frac L2\|y-x\|^2.\]

Proof. Put \(s=y-x\) and subtract the linear term from the integral formula:

\[f(y)-f(x)-\nabla f(x)^\top s =\int_0^1(\nabla f(x+ts)-\nabla f(x))^\top s\,dt.\]

The integrand has absolute value at most \(Lt\|s\|^2\). Integrating gives the bound. \(\square\)

For a gradient step \(y=x-\eta\nabla f(x)\), the upper bound becomes

\[\boxed{ f(x-\eta\nabla f(x)) \le f(x)-\eta\left(1-\frac{L\eta}{2}\right)\|\nabla f(x)\|^2. }\]

Consequently, \(0<\eta<2/L\) gives strict descent whenever the gradient is nonzero. The common choice \(0<\eta\le1/L\) gives the simpler bound

\[f(x_{t+1})\le f(x_t)-\frac\eta2\|\nabla f(x_t)\|^2.\]

In words, each such step lowers the objective by at least a fixed multiple of the squared gradient norm: large gradients guarantee large progress, and progress slows only where the gradient is small. Each gradient-descent convergence proof in the next section starts from this inequality.

For one gradient step on \(f(x)=\log\cosh x\) from \(x=1.5\), with \(L=1\), the guarantee is largest at \(\eta=1/L\) and reaches zero at \(\eta=2/L\). The actual decrease is larger because \(f\) curves less than \(L\) away from the origin.

At \(\eta=2/L\) the guarantee can fail to give progress; above it divergence is possible. These are sufficient bounds using a global \(L\), not necessary limits for every particular point or direction.

Quadratics expose the stability threshold

For \(f(x)=\tfrac12x^\top Qx-b^\top x\) with \(Q\succ0\), let \(x_\star=Q^{-1}b\) and \(e_t=x_t-x_\star\). Then

\[e_{t+1}=(I-\eta Q)e_t.\]

In a Hessian eigenvector with eigenvalue \(\lambda_i\), the error is multiplied by \(1-\eta\lambda_i\). Convergence from every initial point occurs exactly when \(0<\eta<2/\lambda_{\max}(Q)\). Small-curvature directions may decay slowly while large-curvature directions oscillate.

In the Hessian’s eigenvector coordinates (right), gradient descent on a quadratic splits into independent coordinates: the steep one oscillates because \(1-\eta\lambda_{\max}<0\), while the flat one creeps toward the optimum. The left panel shows the same run in the original coordinates. From Gabriel Goh’s Why Momentum Really Works (Distill, 2017).

import numpy as np

Q = np.diag([1., 20.])
x0 = np.array([1., 1.])
L = np.linalg.eigvalsh(Q)[-1]

for eta in [0.5/L, 1/L, 1.9/L, 2/L, 2.1/L]:
    x = x0.copy()
    for _ in range(80):
        x -= eta * (Q @ x)
    value = 0.5 * x @ Q @ x
    print(f"eta={eta:.3f}, f={value:.6g}, x={x}")

np.linalg.eigvalsh returns the eigenvalues of a symmetric matrix in ascending order, so its last entry is \(L=\lambda_{\max}(Q)=20\). At \(2/L\), the second coordinate changes sign without shrinking. At \(2.1/L\), its magnitude grows. The objective is a sum over directions, so a step that is safe for the low-curvature coordinate can still destabilize the other coordinate.

Backtracking when the smoothness constant is unknown

A line search chooses a step using function evaluations. For a descent direction \(p\), Armijo backtracking starts at a trial step and repeatedly multiplies it by \(\tau\in(0,1)\) until

\[f(x+\eta p)\le f(x)+c\eta\nabla f(x)^\top p, \qquad 0<c<1.\]

Differentiability guarantees acceptance at a sufficiently small step. At a zero gradient, gradient descent stops. Appendix D gives the accepted-step bound, convergence argument, and an executable example.

With \(c=0.3\), \(\tau=\tfrac12\), and \(\eta_0=1\), backtracking rejects two trial steps and accepts the third, the first for which \(f(x+\eta p)\) lies on or below the dashed line. The shaded interval marks every step that would pass.

Convergence of gradient descent

The three theorems below use the same algorithm and the same descent inequality; they differ only in their assumptions, and each stronger assumption buys a stronger conclusion:

There are several different claims an analysis might establish:

• Objective convergence: \(f(x_t)\to f_\star\).
• Iterate convergence: \(x_t\to x_\star\).
• First-order stationarity: \(\|\nabla f(x_t)\|\to0\).
• Finite-time complexity: a bound on the number of iterations needed to make a specified error at most \(\varepsilon\).

They are not interchangeable. A bound on the smallest gradient norm seen so far is also different from a bound on the last iterate. The deterministic results below use exact gradients; their constants and conclusions change with noise.

Smooth functions without convexity

Theorem. Let \(f:\mathbb R^d\to\mathbb R\) be \(L\)-smooth and bounded below by \(f_{\inf}\). For gradient descent with fixed \(0<\eta\le1/L\),

\[\sum_{t=0}^{T-1}\|\nabla f(x_t)\|^2\le\frac{2\Delta_0}{\eta}, \qquad \boxed{ \min_{0\le t<T}\|\nabla f(x_t)\|^2 \le\frac{2\Delta_0}{\eta T}. }\]

Proof. Sum the descent inequality over \(t\); the objective differences telescope, and \(f(x_T)\ge f_{\inf}\). The smallest of \(T\) nonnegative terms is at most their average. \(\square\)

For \(\eta=1/L\), finding an iterate with \(\|\nabla f(x_t)\|\le\varepsilon\) needs at most \(O(L\Delta_0/\varepsilon^2)\) iterations. The squared gradient has an \(O(1/T)\) bound; the gradient norm has an \(O(1/\sqrt T)\) bound. The finite sum over all iterations also implies \(\|\nabla f(x_t)\|\to0\).

This does not establish global optimality, avoidance of saddles, or convergence of the iterates. Starting exactly at a nonoptimal stationary point leaves gradient descent there. Bounded iterates have stationary accumulation points by continuity of the gradient, but further assumptions are needed to obtain a single limiting point.

Five runs of gradient descent with \(\eta=0.02\) on \(f(x,y)=(x^2-1)^2+0.3x+1.2y^2\) start at the open circles; dots mark every sixth update. Every path reaches a stationary point, but the start decides which one: three of the five end at the local minimizer.

Varying steps require additional care: even positive steps can shrink so quickly that gradient descent stops short of a stationary point. Appendix D gives the bound and a simple counterexample.

Smooth convex functions

Theorem. Let \(f\) be convex and \(L\)-smooth on \(\mathbb R^d\), with a minimizer \(x_\star\). For fixed \(0<\eta\le1/L\) and \(T\ge1\),

\[\boxed{f(x_T)-f_\star\le\frac{R_0^2}{2\eta T}.}\]

In particular, \(\eta=1/L\) gives \(LR_0^2/(2T)\).

Proof. Convexity bounds \(f(x_t)-f_\star\le\nabla f(x_t)^\top(x_t-x_\star)\). Combining it with smoothness at \(x_{t+1}\) and using \(x_t-x_{t+1}=\eta\nabla f(x_t)\) yields

\[f(x_{t+1})-f_\star \le\frac{\|x_t-x_\star\|^2-\|x_{t+1}-x_\star\|^2}{2\eta}.\]

Summing bounds the sum of the first \(T\) objective gaps by \(R_0^2/(2\eta)\). The objective values decrease, so the last gap is no larger than their average. \(\square\)

On the convex, not strongly convex function \(\log\cosh(1.2x_1)+\tfrac12\log\cosh(3x_2)\), gradient descent with \(\eta=1/L\) never moves away from \(x_\star\). Each dashed circle is centered at \(x_\star\) and passes through one iterate, and every later iterate lies inside it. The proof’s telescoping sum adds up these decreases in squared distance.

Convexity supplies a link from a local gradient to a global objective gap. Without it, the same distance argument does not work. For further first-order complexity results, see Bubeck’s Convex Optimization: Algorithms and Complexity, Chapter 3.

Smooth strongly convex functions

Let \(f\) be \(L\)-smooth and \(\mu\)-strongly convex on \(\mathbb R^d\), with \(0<\mu\le L\) and minimizer \(x_\star\). Strong convexity implies \(\|\nabla f(x)\|^2\ge2\mu(f(x)-f_\star)\), the right inequality of Convex and strongly convex functions: a large gap forces a large gradient, and hence a large guaranteed decrease. With \(\eta=1/L\), the descent lemma therefore gives

\[f(x_{t+1})-f_\star \le\left(1-\frac\mu L\right)(f(x_t)-f_\star),\]

and hence

\[\boxed{ f(x_T)-f_\star \le\left(1-\frac\mu L\right)^T\Delta_0 \le e^{-T/\kappa}\Delta_0. }\]

This is linear, or geometric, convergence: each iteration contracts the error by a fixed factor. It does not mean that error decreases linearly as a function of time. To obtain objective gap at most \(\varepsilon<\Delta_0\) requires \(O(\kappa\log(\Delta_0/\varepsilon))\) iterations. The distance also converges because \(\|x_T-x_\star\|^2\le2(f(x_T)-f_\star)/\mu\).

Here gradient descent with \(\eta=1/L\) runs on a quadratic whose Hessian has eigenvalues \(0.01\), \(0.1\), and \(1\), so \(\kappa=100\). The colored bars split the objective gap by eigendirection, and the bars below show how long each component takes to become negligible. The stiffest component vanishes after one step; the flattest shrinks by the factor \(1-\mu/L=0.99\) per step and sets the rate. From Goh (2017), with the step-size slider set to \(1\); Goh writes \(w\) for the parameters.

Newton’s method

Gradient descent models \(f\) with the same curvature \(1/\eta\) in every direction. Newton’s method instead uses the curvature that the Hessian assigns to each direction, through the local quadratic approximation

\[m_t(s)=f(x_t)+\nabla f(x_t)^\top s+\tfrac12s^\top H(x_t)s.\]

When \(H(x_t)\succ0\), its minimizer solves

\[H(x_t)p_t=-\nabla f(x_t), \qquad x_{t+1}=x_t+p_t.\]

Solve the linear system instead of explicitly computing \(H^{-1}\). A positive definite quadratic objective is solved in one exact Newton step. For a general objective, the accuracy of the quadratic model determines how far the step can be trusted.

Started from the same \(x_0\) with small steps, Newton’s method (red) accounts for curvature and heads more directly to the minimizer than gradient descent (green).

If the Hessian is positive definite and \(\rho\)-Lipschitz near a minimizer, with smallest eigenvalue at least \(\mu>0\), then sufficiently close full Newton steps satisfy

\[\|x_{t+1}-x_\star\|\le\frac{\rho}{2\mu}\|x_t-x_\star\|^2.\]

This is local quadratic convergence of the distance, distinct from an \(O(1/T^2)\) objective bound. Farther away, a line search or trust region controls how much of the quadratic model to trust. An indefinite Hessian may give a direction that increases the objective. Appendix F gives the precise neighborhood assumptions, proof, and curvature approximations.

Constraints and nonsmooth objectives

The smooth unconstrained setting has useful extensions. Projected gradient descent takes a gradient step and projects the parameters back onto a nonempty closed convex feasible set \(C\):

\[x_{t+1}=\Pi_C(x_t-\eta_t\nabla f(x_t)), \qquad \Pi_C(z)=\arg\min_{u\in C}\|u-z\|^2.\]

For a convex function with kinks, a subgradient supplies a supporting affine lower bound, \(f(y)\ge f(x)+g^\top(y-x)\). Subgradient methods generally need different step schedules and averaging to obtain their guarantees. A proximal step instead keeps a nonsmooth penalty intact while approximating the differentiable part of the objective. Appendix E develops these methods, their convergence bounds, soft-thresholding, and constrained optimality conditions.

At the kink \(x_0\) of this convex piecewise-linear function, the one-sided slopes are \(0.3\) and \(1\). Each red line through \((x_0,f(x_0))\) stays below the graph, so its slope, \(0.5\) or \(0.8\), is a subgradient; so is every slope in \([0.3,1]\).

Stochastic gradients and convergence in expectation

An objective may be an expectation or a large finite sum:

\[f(x)=\mathbb E_\xi[\ell(x;\xi)], \qquad\text{or}\qquad f(x)=\frac1n\sum_{i=1}^n\ell_i(x).\]

Computing a full gradient can be expensive. Stochastic gradient descent (SGD) uses a random estimate \(g_t\):

\[x_{t+1}=x_t-\eta_tg_t.\]

Let \(\mathcal F_t\) represent all information available before drawing the gradient estimate at iteration \(t\). A standard assumption is conditional unbiasedness with bounded conditional variance:

\[\mathbb E[g_t\mid\mathcal F_t]=\nabla f(x_t), \qquad \mathbb E[\|g_t-\nabla f(x_t)\|^2\mid\mathcal F_t]\le\sigma^2.\]

For an expected loss, identifying \(\nabla f\) with the expected sample gradient requires interchanging differentiation and expectation; sufficient conditions appear in Appendix B. For a finite average, sampling an index uniformly and independently of the past gives an unbiased sample gradient. Averaging \(B\) conditionally independent estimates reduces the variance bound to \(\sigma^2/B\). Sampling and dependence assumptions matter; random reshuffling is not identical to fresh independent sampling at each step.

With the same constant step \(\eta=0.12\) on \(\tfrac12(x_1^2+6x_2^2)\), SGD follows the general route of gradient descent but keeps fluctuating around \(x_\star\). Its gradients add independent zero-mean noise with standard deviation \(2.2\) in each coordinate.

A basic stochastic convergence bound

For an \(L\)-smooth objective bounded below, a constant step \(0<\eta\le1/L\) gives

\[\frac1T\sum_{t=0}^{T-1}\mathbb E\|\nabla f(x_t)\|^2 \le\frac{2\Delta_0}{\eta T}+L\eta\sigma^2.\]

The first term falls with the number of updates; the second reflects gradient noise. Reducing the step lowers the noise contribution but slows initial progress. For a planned horizon \(T\), choosing \(\eta=1/(L\sqrt T)\) gives an \(O(1/\sqrt T)\) bound on the average expected squared gradient norm. It is neither a last-iterate guarantee nor a global optimality guarantee.

A constant step can leave persistent fluctuations around the optimum. Near a minimizer the true gradient is small, but the noise in its estimate is not, so every update still moves the iterate by a random amount proportional to \(\eta\). On \(f(x)=x^2/2\) with gradient noise of variance \(\sigma^2\), the exact limiting mean square is \(\eta\sigma^2/(2-\eta)\) for \(0<\eta<2\). Appendix D derives both results, the convex bounds, and the classical diminishing-step conditions.

For SGD on a quadratic, with momentum off and a step size near \(1.9\), each bar splits the expected objective gap into the part plain gradient descent would also have (darker, decaying geometrically) and the variance added by gradient noise (lighter, growing from zero). Once the first part has decayed, the gap levels off; small dots follow a single run. From Goh (2017), with the momentum slider set to \(0\).

Stateful optimizers and adaptive scaling

The optimizers in this section modify gradient descent in two ways. Momentum methods average gradients over time, to damp oscillation and build speed along consistent directions. Adaptive methods rescale each coordinate by the recent size of its gradients, so that coordinates with very different gradient scales take comparable steps. Adam and AdamW do both.

An optimizer maps gradients and stored state to a parameter update. Automatic differentiation supplies derivatives of the loss computation; the optimizer decides how those derivatives affect the parameters. Its state can contain past gradients, squared gradients, curvature estimates, and a step counter. Two runs at the same parameter vector can therefore take different steps if their optimizer states differ.

Here \(g_t\) is the current batch gradient at \(x_t\), unless a look-ahead point is specified. Parameter tensors retain their shapes; vector formulas apply entrywise across their entries.

Squares, square roots, and divisions in adaptive updates are coordinatewise. Their denominator constant is \(\epsilon_{\mathrm{opt}}>0\), and fixed moving-average coefficients satisfy \(\beta,\beta_1,\beta_2\in[0,1)\).

Momentum as a gradient filter

A common gradient-buffer convention is

\[u_{t+1}=\beta u_t+g_t, \qquad x_{t+1}=x_t-\eta_tu_{t+1}, \qquad u_0=0,\quad 0\le\beta<1.\]

Unrolling the recurrence gives \(u_{t+1}=\sum_{j=0}^t\beta^{t-j}g_j\). Directions that persist across updates reinforce each other, while rapidly alternating directions partially cancel. The effective memory length is of order \(1/(1-\beta)\) updates.

Some libraries instead use \(u_{t+1}=\beta u_t+(1-\beta)g_t\), an exponential moving average of the gradients. With zero initialization, the two buffers differ by the factor \(1-\beta\), so a learning rate tuned for one convention does not carry over to the other.

With a constant learning rate \(\eta\) and \(x_{-1}=x_0\), the first convention is the heavy-ball method: each update adds a fraction \(\beta\) of the previous displacement,

\[x_{t+1}=x_t-\eta g_t+\beta(x_t-x_{t-1}).\]

With \(\beta=0.99\) and step size \(0.02\), heavy-ball momentum first overshoots across a long, curved valley; the accumulated velocity then carries the iterates along its floor. From Goh’s Why Momentum Really Works, whose step size \(\alpha\) is the \(\eta\) used here.

With a changing learning rate, a gradient buffer and a fixed displacement coefficient define different recurrences. Appendix G gives both forms and the Nesterov look-ahead variant, which evaluates the gradient at an extrapolated point. Acceleration guarantees require a specified extrapolation rule, described next.

Nesterov acceleration

Acceleration evaluates gradients at an extrapolated point. For a convex \(L\)-smooth \(f\) on \(\mathbb R^d\) with a minimizer, set \(y_0=x_0\), \(a_0=1\), and

\[\begin{aligned} x_{t+1}&=y_t-\frac1L\nabla f(y_t),\\ a_{t+1}&=\frac{1+\sqrt{1+4a_t^2}}2,\\ y_{t+1}&=x_{t+1}+\frac{a_t-1}{a_{t+1}}(x_{t+1}-x_t). \end{aligned}\]

This is the smooth case of FISTA. It gives \(f(x_T)-f_\star\le2LR_0^2/(T+1)^2\), improving the general convex gradient-descent rate. Its objective values need not decrease at every step. The guarantee belongs to this specified rule and its assumptions; adding an arbitrary momentum buffer does not establish it. Appendix E gives the composite version and its reference.

Both methods start from the same \(x_{t-1}\) and \(x_t\), with extrapolation coefficient \(0.8\) and step \(0.5\). Heavy ball uses the gradient at \(x_t\); the accelerated method first extrapolates to \(y_t\) and uses the gradient there, which pulls back the overshoot.

AdaGrad and RMSProp

AdaGrad accumulates squared gradients and uses a separate scale for each coordinate:

\[v_{t+1}=v_t+g_t\odot g_t, \qquad x_{t+1}=x_t-\eta_t \frac{g_t}{\sqrt{v_{t+1}}+\epsilon_{\mathrm{opt}}}, \qquad v_0=0.\]

Coordinates with a large accumulated squared gradient receive smaller effective steps. Coordinates that are rarely updated can retain relatively large steps. Since the accumulator never decreases, its memory does not adapt to a new scale as quickly as a moving average. The original analysis connects this geometry to online convex optimization: Duchi, Hazan, and Singer (2011).

RMSProp replaces the cumulative sum by an exponential moving average:

\[v_{t+1}=\beta_2v_t+(1-\beta_2)(g_t\odot g_t), \qquad x_{t+1}=x_t-\eta_t \frac{g_t}{\sqrt{v_{t+1}}+\epsilon_{\mathrm{opt}}}, \qquad v_0=0.\]

This tracks recent squared-gradient scale. The displayed form is uncentered RMSProp without a separate momentum buffer. Centered variants estimate a variance by subtracting a squared mean; the quantity \(v_t\) itself is an uncentered second moment, not a variance. The PyTorch RMSProp specification states these implementation choices explicitly.

When the gradients of one coordinate shrink tenfold after 200 updates (\(\eta=0.01\)), AdaGrad’s accumulated sum keeps its steps small, while RMSProp with \(\beta_2=0.97\) forgets the old scale and its steps recover. The large first RMSProp steps come from the zero initial average.

Adam and bias correction

Adam combines an exponential average of gradients with an exponential average of their squares:

\[\begin{aligned} m_{t+1}&=\beta_1m_t+(1-\beta_1)g_t,\\ v_{t+1}&=\beta_2v_t+(1-\beta_2)(g_t\odot g_t), \qquad m_0=v_0=0. \end{aligned}\]

The two averages have different roles. The first retains a signed direction; the second measures coordinatewise magnitude. It does not estimate the Hessian and it does not, without subtracting the square of the mean, estimate gradient variance.

Zero initialization gives both averages too little total weight early in the run. With fixed coefficients, the accumulated weights after update \(t\) are \(1-\beta_1^{t+1}\) and \(1-\beta_2^{t+1}\).

The bias corrections divide by these weights:

\[\widehat m_{t+1}=\frac{m_{t+1}}{1-\beta_1^{t+1}}, \qquad \widehat v_{t+1}=\frac{v_{t+1}}{1-\beta_2^{t+1}}.\]

The Adam step is

\[\boxed{x_{t+1}=x_t-\eta_t \frac{\widehat m_{t+1}} {\sqrt{\widehat v_{t+1}}+\epsilon_{\mathrm{opt}}}.}\]

For the first update (\(t=0\)), the denominators are \(1-\beta_1\) and \(1-\beta_2\). Correction removes the initialization deficit in the weights; it does not remove the lag of a moving average or make the final ratio unbiased. The original method is due to Kingma and Ba (2015).

For a constant gradient \(g=1\) with \(\beta_1=0.9\) and \(\beta_2=0.999\), the uncorrected \(v\) lags far behind \(m\), and the ratio \(m/\sqrt v\) would make early steps up to about \(6.6\) times too large. After correction the ratio is \(1\) from the first update.

For a single fresh update, \(\widehat m_1=g_0\) and \(\widehat v_1=g_0^2\). When a coordinate’s gradient magnitude is large relative to \(\epsilon_{\mathrm{opt}}\), its first update is approximately \(-\eta_0\operatorname{sign}(g_{0,i})\). This illustrates how different Adam’s scaling is from ordinary gradient descent.

Although each denominator is positive, the direction \(\widehat m_{t+1}\) can differ from the current gradient. Even with exact gradients, a particular Adam update need not be a descent step. Classical SGD convergence bounds also do not transfer automatically through a gradient-dependent preconditioner. Reddi, Kale, and Kumar give convex counterexamples to unconditional Adam convergence and develop AMSGrad, which modifies the second-moment rule. Guarantees depend on the variant, schedule, and assumptions.

L2 regularization and AdamW

For an explicit penalty \(r(x)=\tfrac\lambda2\|x\|^2\) with \(\lambda\ge0\), plain SGD—or exact gradient descent when \(g_t=\nabla f(x_t)\)—gives

\[x_{t+1}=x_t-\eta_t(g_t+\lambda x_t) =(1-\eta_t\lambda)x_t-\eta_tg_t.\]

In this unpreconditioned, momentum-free case, adding an L2 penalty is the same algebraic update as multiplicative weight decay. With a diagonal preconditioner \(P_t\), the regularized-gradient update instead contains \(-\eta_t\lambda P_tx_t\), so coordinates decay at different rates. For Adam, adding \(\lambda x_t\) to the gradient also changes both moment histories.

AdamW computes moments from the loss gradient and applies decoupled weight decay:

\[\boxed{ x_{t+1}=(1-\eta_t\lambda)x_t -\eta_t\frac{\widehat m_{t+1}} {\sqrt{\widehat v_{t+1}}+\epsilon_{\mathrm{opt}}}. }\]

Here \(\lambda\) is the decay coefficient in this convention; some descriptions parameterize shrinkage differently. Decoupling means that the adaptive denominator does not scale the shrinkage term. It does not mean that total shrinkage is independent of the learning-rate schedule: without gradient updates, it is \(\prod_{t<T}(1-\eta_t\lambda)\). See Loshchilov and Hutter (2019).

Two parameters start at \(1\), and their loss gradients are pure zero-mean noise of scale \(1\) and \(10^{-3}\); \(\eta=10^{-3}\), \(\lambda=1\), and \((\beta_1,\beta_2)=(0.9,0.999)\). With the penalty inside the gradient, the adaptive denominator rescales the decay, and the parameter with tiny gradients is pulled to zero much faster. AdamW shrinks both at the rate \(e^{-\eta\lambda t}\).

Parameter groups can have different learning rates and decay coefficients. Excluding biases or normalization parameters from decay is a modeling choice, not part of the mathematical definition. Because Adam’s state remembers past gradients, a parameter whose current gradient is zero can still move by more than the decay; only decay moves it when its stored first moment is also zero.

Optimizer paths on a quadratic

Consider

\[f(x)=\frac12(x_1^2+20x_2^2),\qquad \nabla f(x)=\begin{bmatrix}x_1\\20x_2\end{bmatrix}, \qquad x_\star=0.\]

The Hessian eigenvalues are \(1\) and \(20\), so \(L=20\), \(\mu=1\), and \(\kappa=20\). Each contour line joins points with the same objective value; darker bands have smaller values. The narrow direction has larger curvature. Every run below starts at \(x_0=(3,1)^\top\) and takes 40 updates with exact gradients and zero initial optimizer buffers.

The dashed box in each panel is magnified about three times in the strip below it. Dots darken with time; the open circle is the start and the filled circle the iterate after 40 updates. Gradient descent and heavy-ball momentum use \(\eta=0.09\), momentum with the unnormalized buffer and \(\beta=0.65\). Nesterov uses the displayed smooth FISTA recurrence with \(1/L=0.05\). Adam uses \(\eta=0.18\), \((\beta_1,\beta_2)=(0.9,0.99)\), \(\epsilon_{\mathrm{opt}}=10^{-8}\), and no weight decay.

Gradient descent multiplies the two error coordinates by \(0.91\) and \(-0.8\), respectively: slow movement down the valley accompanies alternating movement across it. Momentum combines current gradients with past directions and can overshoot the minimizer. Here the accelerated method removes the second coordinate in its first step because \(1-20/L=0\); later extrapolation acts along the first coordinate. Adam’s coordinatewise normalization changes its direction and its scale relative to the raw gradient.

On a logarithmic axis, the objective gaps of the same runs show oscillation near the minimum that is hard to see in the contours. These finite runs illustrate particular settings; they do not rank the methods across problems. Learning rates differ, and the horizontal axis counts updates rather than elapsed time.

Appendix H contains the code for these trajectories. Full Newton would solve this positive definite quadratic in one exact step; on a large problem, constructing and solving with the Hessian has a different cost from a gradient evaluation.

Schedules, batches, and clipping

A learning-rate schedule specifies \(\eta_t\) as training progresses. Common choices include constant steps, step decay, warmup followed by cosine decay, and inverse-square-root decay. A microbatch is processed in one forward/backward computation; an optimizer update changes parameters and stored state; an epoch is one pass through the dataset. A schedule’s clock can count updates, examples, or tokens, so its unit must be specified.

The schedules are drawn as multiples of their peak rate over \(T\) updates. Step decay divides the rate by \(10\) every \(0.4T\); both warmups rise linearly over \(0.06T\) before cosine or inverse-square-root decay.

Gradient accumulation combines microbatch gradients before a single parameter update. If microbatch \(j\) has summed loss \(S_j\) over \(n_j\) scored items, the effective-batch mean gradient is

\[g_t=\frac{\sum_j\nabla S_j(x_t)}{\sum_j n_j}.\]

All microbatches use the same parameters. An unweighted average of microbatch means is correct only when their counts agree. The optimizer’s moments, step counter, and weight decay advance once per accumulated update. Changing the batch size also changes the number of updates per epoch and the duration of momentum memory in examples.

Global norm clipping replaces \(g_t\) by \(g_t\min\{1,c/\|g_t\|\}\) for threshold \(c>0\), with zero mapped to zero. It preserves direction and caps the gradient magnitude. Applying it once after accumulation differs from clipping each microbatch first. For AdamW it bounds the gradient entering the moments, rather than directly bounding the final parameter displacement.

Appendix G contains schedule formulas, an accumulation-and-clipping example, distributed averaging, and additional optimizer families. A basic forward/backward/update loop is developed in Numerical Computing with NumPy and PyTorch.

Convergence rates and their scope

Bounds such as \(O(1/T)\) and \(O(1/\sqrt T)\) are called sublinear rate bounds; a bound \(e_T\le Mq^T\) with a constant \(M>0\) and \(q\in(0,1)\) is geometric. These are upper bounds, so a particular problem may converge faster. For a positive error sequence tending to zero, local superlinear convergence means \(e_{t+1}/e_t\to0\); local quadratic convergence means \(e_{t+1}\le M'e_t^2\) eventually, for some constant \(M'\). Every rate must specify its error quantity.

On a logarithmic axis, a geometric rate is a straight line that eventually falls below every sublinear rate, and a quadratic rate doubles the number of correct digits at each step. The curves match the rates in the table below.

In the table, \(R_0\) is the initial distance to a minimizer, \(\Delta_0\) is an initial objective gap or gap to a finite lower bound, and \(\kappa=L/\mu\). Constants refer to the domains required by the preceding theorems.

Step sizes. A constant safe step can converge exactly for deterministic smooth objectives. Nonsmooth or noisy updates often require averaging or diminishing steps because their error bounds also contain accumulated squared-step terms. Diminishing steps alone do not guarantee convergence: the total step length must permit continued progress.

Stopping criteria. A stopping criterion should match the theorem. An ordinary gradient norm measures unconstrained first-order stationarity; constrained problems require a measure that accounts for feasible directions (Appendix E). A small change in objective or a small step can also result from an excessively small learning rate. Objective accuracy, gradient accuracy, and parameter accuracy require different certificates.

Cost per iteration. Iteration complexity counts updates or oracle calls. It does not by itself measure elapsed work: a stochastic gradient, full gradient, proximal solve, and Hessian factorization can have very different costs. Likewise, a proof requiring global smoothness cannot be applied merely because a Hessian is bounded at the starting point. Local assumptions suffice only when the analysis also keeps all required iterates and trial segments inside the region where those assumptions hold.

Appendices

import numpy as np

X = np.array([[1., 2.], [0., -1.]])
E = np.array([[0.1, 0.], [0.2, -0.1]])
linear = X @ E + E @ X
quadratic = E @ E
assert np.allclose((X + E) @ (X + E), X @ X + linear + quadratic)

import numpy as np

x = np.array([2.0, 3.0])
v = np.array([1.0, -1.0])
s = x[0] * x[1]
r = np.sin(s)
value = s + r

# Sensitivity along x(tau) = [2 + tau, 3 - tau].
dot_s = x[1] * v[0] + x[0] * v[1]
dot_r = np.cos(s) * dot_s
dot_f = dot_s + dot_r

# Reverse: accumulate both contributions at the shared node s.
bar_s = 0.0
bar_r = 1.0
bar_s += 1.0
bar_s += np.cos(s) * bar_r
gradient = bar_s * np.array([x[1], x[0]])

assert np.isclose(dot_f, gradient @ v)
print(round(value, 8), round(dot_f, 8), gradient.round(8))
# 5.7205845 1.96017029 [5.88051086 3.92034057]

tau = 0.01
x_new = x + tau * v
s_new = x_new[0] * x_new[1]
actual_change = s_new + np.sin(s_new) - value
predicted_change = tau * dot_f
print("predicted change:", predicted_change)
print("actual change:", actual_change)
# Approximately 0.019602 versus 0.019419.

import torch
from torch.func import grad, jvp, vjp

def F(x):
    s = x[0] * x[1]
    return torch.stack((s, torch.sin(s)))

def f(x):
    return F(x).sum()

x = torch.tensor([2.0, 3.0], dtype=torch.float64)
v = torch.tensor([1.0, -1.0], dtype=torch.float64)
u = torch.ones(2, dtype=torch.float64)
_, Jv = jvp(F, (x,), (v,))
_, pullback = vjp(F, x)
(JTu,) = pullback(u)
gradient = grad(f)(x)

assert torch.allclose(gradient, JTu)
assert torch.allclose(u @ Jv, JTu @ v)
print(Jv, JTu)
# tensor([1.0000, 0.9602], dtype=torch.float64)
# tensor([5.8805, 3.9203], dtype=torch.float64)

import torch
from torch.func import grad, jvp

def f(x):
    s = x[0]*x[1]
    return s + torch.sin(s)

x = torch.tensor([2., 3.], dtype=torch.float64)
v = torch.tensor([1., -1.], dtype=torch.float64)
_, Hv = jvp(grad(f), (x,), (v,))

# Independent analytic product for f(x) = s + sin(s), s = x[0]*x[1].
s = x[0] * x[1]
q = torch.stack((x[1], x[0]))
expected = -torch.sin(s) * q * (q @ v) + (1 + torch.cos(s)) * v.flip(0)
assert torch.allclose(Hv, expected, atol=1e-12, rtol=1e-12)
print(Hv)
# tensor([-1.1219, 2.5190], dtype=torch.float64)

import numpy as np

q = np.array([1., 10.])

def f(x):
    # Evaluate log(cosh(x)).
    return np.sum(np.logaddexp(x, -x) - np.log(2.)) + 0.05*np.sum(q*x*x)

def grad(x):
    return np.tanh(x) + 0.1*q*x

x = np.array([4., -3.])
values, accepted_steps = [f(x)], []
for _ in range(60):
    g = grad(x)
    if np.linalg.norm(g) < 1e-8:
        break
    eta = 10.0
    for _ in range(80):
        candidate = x - eta*g
        if f(candidate) <= f(x) - 1e-4*eta*(g @ g):
            break
        eta *= 0.5
    else:
        raise RuntimeError("No acceptable step found")
    x = candidate
    values.append(f(x))
    accepted_steps.append(eta)

print("x:", x, "gradient norm:", np.linalg.norm(grad(x)))
print("first accepted steps:", accepted_steps[:8])
assert np.all(np.diff(values) <= 1e-12)

import numpy as np

def grad(x):
    return np.sign(x) * np.sqrt(abs(x))

x_fixed = x_adaptive = 1.0
nu, L_nu = 0.5, np.sqrt(2.0)
for _ in range(40):
    x_fixed -= 0.4*grad(x_fixed)
    g = grad(x_adaptive)
    if g != 0:
        eta = L_nu**(-1/nu) * abs(g)**((1-nu)/nu)
        x_adaptive -= eta*g

print("fixed step:", x_fixed, "Hölder step:", x_adaptive)

import numpy as np

rng = np.random.default_rng(12)
runs, T, sigma, eta = 4000, 2000, 0.5, 0.1
x_fixed = np.full(runs, 3.0)
x_decay = x_fixed.copy()
for t in range(T):
    noise = sigma*rng.normal(size=runs)
    x_fixed -= eta*(x_fixed + noise)
    x_decay -= (x_decay + noise)/(t + 2)

print("fixed-step mean square:", np.mean(x_fixed**2))
print("fixed-step theoretical limit:", eta*sigma**2/(2-eta))
print("diminishing-step mean square:", np.mean(x_decay**2))

import numpy as np

def subgradient_run(T, schedule):
    x, weighted_sum, total_weight = 0.95, 0.0, 0.0
    tail = []
    for t in range(T):
        eta = schedule(t)
        weighted_sum += eta*x       # Average the pre-update iterates.
        total_weight += eta
        x -= eta*np.sign(x)         # Choose subgradient 0 at x=0.
        if t >= T - 6:
            tail.append(x)
    return x, weighted_sum/total_weight, np.array(tail)

for label, schedule in [
    ("constant", lambda t: 0.2),
    ("diminishing", lambda t: 0.5/(t + 1)**0.75),
]:
    last, average, tail = subgradient_run(4000, schedule)
    print(label, "last:", last, "average:", average, "tail:", tail)

import numpy as np

q = np.array([1., 3., 8.])
b = np.array([2., -3., 4.])
eta = 1/q.max()
x = np.zeros(3)
for _ in range(200):
    x = np.clip(x - eta*(q*x - b), 0., 1.)

gradient = q*x - b
mapping = (x - np.clip(x - eta*gradient, 0., 1.))/eta
print("solution:", x)
print("gradient:", gradient, "gradient mapping:", mapping)
assert np.allclose(x, np.clip(b/q, 0., 1.))

import numpy as np

rng = np.random.default_rng(4)
A = rng.normal(size=(80, 30))/np.sqrt(80)
A *= np.geomspace(1., 12., 30)
truth = np.zeros(30)
truth[:4] = [2., -1., 0.5, 1.5]
b = A @ truth + 0.02*rng.normal(size=80)
lam = 0.05
L = np.linalg.norm(A, 2)**2

def objective(x):
    return 0.5*np.sum((A @ x - b)**2) + lam*np.abs(x).sum()

def prox_step(z):
    u = z - (A.T @ (A @ z - b))/L
    return np.sign(u)*np.maximum(np.abs(u) - lam/L, 0.)

def run(accelerated, T=400):
    x = np.zeros(A.shape[1])
    y, a = x.copy(), 1.0
    values = [objective(x)]
    for _ in range(T):
        old_x = x.copy()
        x = prox_step(y)
        if accelerated:
            next_a = (1 + np.sqrt(1 + 4*a*a))/2
            y = x + ((a - 1)/next_a)*(x - old_x)
            a = next_a
        else:
            y = x.copy()
        values.append(objective(x))
    return x, np.array(values)

x_pg, values_pg = run(False)
x_fast, values_fast = run(True)
print("proximal gradient:", values_pg[-1], "FISTA:", values_fast[-1])
print("nonzero coordinates:", np.count_nonzero(np.abs(x_fast) > 1e-8))
assert np.all(np.diff(values_pg) <= 1e-10)

import numpy as np

Q = np.diag([1., 4., 9.])
b = np.array([1., 2., 3.])
A = np.ones((1, 3))
c = np.array([1.])
K = np.block([[Q, A.T], [A, np.zeros((1, 1))]])
solution = np.linalg.solve(K, np.concatenate([b, c]))
x_star, multiplier = solution[:3], solution[3:]
print(x_star, multiplier)
assert np.allclose(A @ x_star, c)
assert np.allclose(Q @ x_star - b + A.T @ multiplier, 0)

import numpy as np

def grad(x):
    return x**3 + x - 1

def hess(x):
    return 3*x*x + 1

# A separate bisection computation supplies a reference minimizer.
lo, hi = 0.0, 1.0
for _ in range(60):
    mid = (lo + hi)/2
    if grad(mid) < 0:
        lo = mid
    else:
        hi = mid
x_star = (lo + hi)/2

x = 0.9
for _ in range(6):
    error = abs(x - x_star)
    next_x = x - grad(x)/hess(x)
    next_error = abs(next_x - x_star)
    if error > 1e-7:   # Print the ratios before the error becomes extremely small.
        print("error:", error, "next_error/error^2:", next_error/error**2)
    x = next_x
assert abs(grad(x)) < 1e-12

import numpy as np

def adamw_step(x, g, m, v, t, eta=0.05, beta1=0.9, beta2=0.999,
               eps=1e-8, weight_decay=0.01):
    m_next = beta1*m + (1 - beta1)*g
    v_next = beta2*v + (1 - beta2)*g*g
    m_hat = m_next/(1 - beta1**(t + 1))
    v_hat = v_next/(1 - beta2**(t + 1))
    x_next = (1 - eta*weight_decay)*x - eta*m_hat/(np.sqrt(v_hat) + eps)
    return x_next, m_next, v_next

x = np.array([2., -1.])
m, v = np.zeros_like(x), np.zeros_like(x)
q = np.array([1., 100.])
initial_value = 0.5*np.sum(q*x*x)
for t in range(250):
    g = q*x
    x, m, v = adamw_step(x, g, m, v, t)
print("initial objective:", initial_value, "final:", 0.5*np.sum(q*x*x))

# With fresh moments, coupled L2 and decoupled decay already differ.
x0 = np.array([1., 2.])
zeros = np.zeros_like(x0)
coupled, _, _ = adamw_step(x0, 0.1*x0, zeros, zeros, 0,
                           eta=0.01, weight_decay=0.)
decoupled, _, _ = adamw_step(x0, zeros, zeros, zeros, 0,
                             eta=0.01, weight_decay=0.1)
print("Adam with L2:", coupled, "AdamW:", decoupled)
assert np.allclose(decoupled, 0.999*x0)
assert not np.allclose(coupled, decoupled)

import torch

A = torch.tensor([[1., 2.], [3., -1.], [2., 0.], [-1., 1.], [0., 4.]],
                 dtype=torch.float64)
b = torch.tensor([1., 0., 2., -1., 3.], dtype=torch.float64)
x = torch.tensor([0.3, -0.2], dtype=torch.float64, requires_grad=True)
cuts = [slice(0, 2), slice(2, 5)]       # Unequal microbatch counts.

full_loss = 0.5*((A @ x - b)**2).mean()
full_gradient = torch.autograd.grad(full_loss, x)[0]
optimizer = torch.optim.AdamW([x], lr=0.01, weight_decay=0.1)
optimizer.zero_grad(set_to_none=True)
for cut in cuts:
    summed_loss = 0.5*((A[cut] @ x - b[cut])**2).sum()
    (summed_loss/len(b)).backward()    # Global count, not this slice's count.
assert torch.allclose(x.grad, full_gradient)
print("full and accumulated gradients:", full_gradient, x.grad)

raw_norm = torch.nn.utils.clip_grad_norm_([x], max_norm=1.0)
optimizer.step()                      # One moment and parameter update.
print("raw gradient norm:", raw_norm.item(), "new parameters:", x.detach())

import numpy as np
import matplotlib.pyplot as plt

Q = np.diag([1., 20.])
x0 = np.array([3., 1.])
T = 40
histories = {}
for name, eta in [("GD", .09), ("Momentum", .09),
                  ("Nesterov", .05), ("Adam", .18)]:
    x = x0.copy()
    u = np.zeros_like(x)
    m, v = u.copy(), u.copy()
    y, a = x.copy(), 1.0
    path = [x.copy()]
    for t in range(T):
        g = Q @ x
        if name == "GD":
            x = x - eta*g
        elif name == "Momentum":
            u = .65*u + g
            x = x - eta*u
        elif name == "Nesterov":
            old_x = x.copy()
            x = y - eta*(Q @ y)
            next_a = (1 + np.sqrt(1 + 4*a*a))/2
            y = x + (a - 1)/next_a*(x - old_x)
            a = next_a
        else:
            m = .9*m + .1*g
            v = .99*v + .01*g*g
            m_hat = m/(1 - .9**(t + 1))
            v_hat = v/(1 - .99**(t + 1))
            x = x - eta*m_hat/(np.sqrt(v_hat) + 1e-8)
        path.append(x.copy())
    histories[name] = np.array(path)

xx, yy = np.meshgrid(np.linspace(-1.25, 3.35, 300),
                     np.linspace(-1.4, 1.4, 250))
levels = [.05, .2, .5, 1, 2, 4, 8, 14, 22]
fig, axes = plt.subplots(2, 2, figsize=(9, 6), sharex=True, sharey=True)
for ax, (name, path) in zip(axes.flat, histories.items()):
    ax.contour(xx, yy, (xx**2 + 20*yy**2)/2,
               levels=levels, colors="0.8")
    ax.plot(path[:, 0], path[:, 1], "o-", markersize=3)
    ax.plot(0, 0, "k*")
    ax.set(title=name, xlabel="x1", ylabel="x2")
    ax.set_aspect("equal")
fig.tight_layout()

fig, ax = plt.subplots(figsize=(8, 4))
for name, path in histories.items():
    values = .5*np.einsum("ti,ij,tj->t", path, Q, path)
    ax.semilogy(np.arange(T + 1), values, label=name)
    print(name, "final objective:", values[-1])
ax.set(xlabel="Optimizer updates", ylabel="Objective gap")
ax.legend()
fig.tight_layout()
plt.show()