Stochastic and adaptive optimizers

The gradient descent note ended with two costs. Every step reads all $n$ examples, and the number of steps grows with the condition number $\kappa$. This note attacks both. The first cost goes away when the gradient is estimated from a small random sample of examples: the steps become cheap and noisy, and the noise changes how the step size must behave. The second is reduced by remembering past gradients (momentum) and by giving each coordinate its own step size (AdaGrad, RMSProp, Adam). Adam trains almost every neural network in the NLP and Generative Models courses, so it is worth knowing exactly what each of its lines does.

Notation. The loss is an average over examples, $Q(w) = \frac1n\sum_{i=1}^n \ell_i(w)$, where $\ell_i$ is the loss on example $i$; for least squares $\ell_i(w) = \frac12(\langle w, x_i\rangle - y_i)^2$ in this note. A minibatch $S$ is a random set of $B$ indices, $g = \frac1B\sum_{i \in S}\nabla\ell_i(w)$ is the stochastic gradient, and one epoch is $n/B$ steps, one pass over the data. $\eta_k$ is the step size at step $k$, $\beta$ the momentum coefficient, and $m$, $v$ running averages of gradients and squared gradients. Products and squares of vectors in the adaptive methods are taken coordinate by coordinate. $H$, $\mu$, $L$ and $\kappa$ are the Hessian, its extreme eigenvalues and their ratio, as before.

A gradient from a sample#

The gradient of an average is the average of the gradients, $\nabla Q(w) = \frac1n\sum_i\nabla\ell_i(w)$. Pick one index $i$ uniformly at random: then $E[\nabla\ell_i(w)] = \nabla Q(w)$, so a single example’s gradient is an unbiased estimate of the full one. Averaging $B$ independently drawn examples keeps the mean and divides the covariance by $B$:

$$E[g] = \nabla Q(w), \qquad \operatorname{Cov}(g) = \frac1B\,\Sigma(w), \qquad \Sigma(w) = \frac1n\sum_{i=1}^n \big(\nabla\ell_i(w) - \nabla Q(w)\big)\big(\nabla\ell_i(w) - \nabla Q(w)\big)^\top .$$

This is the Monte Carlo estimate of an average, and its error falls like $1/\sqrt B$. Stochastic gradient descent (SGD) takes the step $w \leftarrow w - \eta g$ with that estimate. A step costs $O(Bd)$ instead of $O(nd)$. In practice the examples are shuffled once per epoch and taken in consecutive batches, sampling without replacement, which has slightly less noise than the independent draws of the formula.

Worked example. On the diabetes data (442 patients, 10 standardized features, centred target), at $w = 0$ the root mean square error of the minibatch gradient, measured over 5,000 random batches, is $238.6$ for $B = 1$, $75.6$ for $B = 10$ and $23.9$ for $B = 100$; the formula $\sqrt{\operatorname{tr}\Sigma/B}$ gives $238.8$, $75.5$ and $23.9$. Ten times more examples per step buy a $\sqrt{10} = 3.16$ times smaller error.

Why accept noise at all? Because early in training the full gradient is overkill: a sample of a few hundred examples already points in roughly the right direction, and one epoch of SGD takes $n/B$ steps for the price of one full step. Nor is there a reason to optimize the training loss far more precisely than the training data determine the true minimizer; the statistical error of $\hat w$ is already of order $1/\sqrt n$ (Goodfellow, Bengio and Courville, 2016, section 8.1.3).

A constant step stalls in a noise ball#

One dimension. Take $Q(w) = \frac\lambda2 w^2$ and a stochastic gradient $\lambda w + \xi$ with independent noise $\xi$ of mean 0 and variance $s^2$. The update is $w_{k+1} = (1 - \eta\lambda)w_k - \eta\xi_k$. If $0 < \eta\lambda < 2$, the variance $V_k = \operatorname{Var}(w_k)$ obeys $V_{k+1} = (1 - \eta\lambda)^2 V_k + \eta^2 s^2$ and converges to the fixed point

$$V_\infty = \frac{\eta^2 s^2}{1 - (1 - \eta\lambda)^2} = \frac{\eta s^2}{\lambda(2 - \eta\lambda)} \approx \frac{\eta s^2}{2\lambda} \quad (\eta\lambda \ll 1) .$$

The expected excess loss $\frac\lambda2 V_\infty$ is about $\eta s^2/4$: proportional to the step. With $\lambda = s = 1$, $V_\infty = 0.1/1.9 = 0.0526$ for $\eta = 0.1$ and $0.01/1.99 = 0.0050$ for $\eta = 0.01$; 20,000 simulated runs give $0.053$ and $0.005$ (the Code section). A ten times smaller step gives a ten times smaller variance, a ball about $\sqrt{10} \approx 3.2$ times narrower, but approaches it ten times more slowly, since the mean shrinks by $1 - \eta\lambda$ per step.

Least squares. Near the minimum the noise of a single example’s gradient $x_i(\langle w, x_i\rangle - y_i)$ comes mainly from the residual noise: its covariance is about $\sigma^2 H$ when the residuals have variance $\sigma^2$ and are independent of $x$. Applying the one-dimensional result to each eigendirection, where the noise variance is $\sigma^2\lambda_i$, gives the excess loss $\sum_i \eta\sigma^2\lambda_i/4 = \eta\sigma^2\operatorname{tr}(H)/4$ for small steps.

Figure

A constant step stalls in a noise ball

The figure shows the trade on a problem where everything can be computed. SGD with step $0.1$ gets in 50 examples, a tenth of an epoch, to a loss that full gradient descent reaches only after about five epochs, and then stops improving; the smaller step $0.02$ ends about seven times lower ($0.0072$ against $0.050$ over the last epoch; the small-step formula predicts five times) but takes longer to get there. A decreasing step has both advantages. The classical conditions of Robbins and Monro (1951) are

$$\sum_{k} \eta_k = \infty, \qquad \sum_{k} \eta_k^2 < \infty :$$

the steps must add up to an infinite distance, so that any starting point can be left behind, while the accumulated noise, whose variance adds up as $\sum\eta_k^2 s^2$, stays finite. Steps $\eta_k \propto 1/k$ satisfy both. For a $\mu$-strongly convex loss with bounded stochastic gradients, $\eta_k = 2/(\mu(k + 1))$ makes a weighted average of the iterates converge at the rate $O(1/k)$ (Bubeck, 2015, Theorem 6.2), much slower than the linear rate of exact gradient descent: the price of the noise.

In deep learning the decrease is applied by hand as a schedule, for example a tenfold cut after a set number of epochs or a cosine decay; the training note covers them. On the diabetes data, SGD with batches of 32 and a constant step $0.2$ ends 30 epochs with an excess loss of $50.2$ (the mean over five shuffling seeds, against a minimum loss of $1429.8$); cutting the step to $0.02$ after 15 epochs ends at $5.1$. With a constant step of $0.05$ the noise ball is smaller, but the end point, $9.5$, is now limited by the slow progress along the flattest direction of a problem with $\kappa = 470$, and the same cut changes little ($9.3$).

Averaging the iterates also shrinks the ball without touching the step. Polyak averaging reports $\bar w_k = \frac1k\sum_{j \le k} w_j$ instead of the last iterate; the exponential moving average of the weights used to train diffusion models is its modern form (DDPM training in practice).

Momentum: remember the direction#

The zigzag of the previous note had a simple cause: along the steep direction the gradient flips sign at every step, along the flat one it keeps pointing the same way, but with a small magnitude. Averaging recent gradients cancels the first and accumulates the second. Momentum, or the heavy-ball method (Polyak, 1964), keeps a velocity $v$:

$$v_{k+1} = \beta v_k + \nabla Q(w_k), \qquad w_{k+1} = w_k - \eta\, v_{k+1} .$$

Unrolled, $v_{k+1} = \sum_{j \ge 0}\beta^j\nabla Q(w_{k-j})$, an exponentially weighted sum of past gradients. When the gradient is constant, the velocity grows to $\nabla Q/(1 - \beta)$, so with $\beta = 0.9$ the effective step along a consistent direction is ten times $\eta$; components that alternate in sign largely cancel.

On a quadratic the method can be analysed exactly like gradient descent, now as a two-term recursion per eigendirection (Goh, 2017). With the best choice $\eta = 4/(\sqrt L + \sqrt\mu)^2$ and $\beta = \big((\sqrt\kappa - 1)/(\sqrt\kappa + 1)\big)^2$, the error shrinks asymptotically, per step, by

$$\frac{\sqrt\kappa - 1}{\sqrt\kappa + 1} \quad \text{instead of} \quad \frac{\kappa - 1}{\kappa + 1} :$$

the number of steps grows like $\sqrt\kappa$ instead of $\kappa$. For $\kappa = 100$ the factors are $0.818$ and $0.980$. The momentum factor is the spectral radius of the two-term recursion: it sets the rate over many steps, not a contraction at every step. Unlike gradient descent, the heavy ball can first move away from the minimum: from the starting point of the figure below, its error norm grows from $2.88$ to $8.98$ in four steps before it falls. On the quadratic of the figure below, with eigenvalues $4.5$ and $0.045$, a thousandfold reduction of the loss takes gradient descent 173 steps at its best constant step and the optimally tuned heavy ball 39.

Nesterov’s method evaluates the gradient at the point the velocity is about to carry the iterate to, $v_{k+1} = \beta v_k + \nabla Q(w_k - \eta\beta v_k)$, in the form of Sutskever et al. (2013). The look-ahead corrects the velocity before it overshoots. With $\eta = 1/L$ and $\beta = (\sqrt\kappa - 1)/(\sqrt\kappa + 1)$ it needs 17 steps on the same problem, and unlike the heavy ball it keeps the $\sqrt\kappa$ rate on every smooth strongly convex function, not only on quadratics (Bubeck, 2015, Theorem 3.18). No method that only combines gradients can do better in the worst case: for every such method there is a smooth strongly convex quadratic, in a high enough dimension, on which the error after $k$ steps is at least $\big((\sqrt\kappa - 1)/(\sqrt\kappa + 1)\big)^k$ times the initial error (Theorem 3.15 there).

The optimal parameters need $\mu$ and $L$, which are not known. In practice $\beta = 0.9$ and $\eta$ is tuned; with stochastic gradients momentum also averages the noise, but its noise ball is that of the effective step $\eta/(1 - \beta)$, so schedules are still needed. On the diabetes data, momentum $0.9$ with $\eta = 0.02$ (an effective step of $0.2$) ends at $12.6$, and at $5.4$ with the halfway cut.

Per-coordinate steps: AdaGrad and RMSProp#

Momentum treats all coordinates alike. Features on different scales, or rare features that appear in few examples, call for different step sizes per coordinate, which is what standardization achieved in the gradient descent note by changing the features. Adaptive methods achieve it by changing the steps, using the gradients seen so far.

AdaGrad (Duchi, Hazan and Singer, 2011) divides each coordinate’s step by the root of the sum of its squared gradients:

$$G_{k} = G_{k-1} + g_k^2, \qquad w_{k+1} = w_k - \frac{\eta}{\sqrt{G_k} + \varepsilon}\, g_k .$$

A coordinate that has seen large gradients gets small steps; a coordinate of a rare feature, whose gradient is mostly zero, keeps a large step for the few examples where it appears. This suits sparse problems: GloVe word vectors are trained with AdaGrad (GloVe). The weakness is that $G_k$ only grows, so all steps shrink towards zero even when the loss surface stops changing.

RMSProp (Hinton, Srivastava and Swersky, 2012, lecture 6e) replaces the sum by an exponential moving average, $v_k = \rho v_{k-1} + (1 - \rho)g_k^2$ with $\rho = 0.9$, and divides by $\sqrt{v_k}$: the step adapts to the recent size of the gradient instead of its whole history.

Adam#

Adam (Kingma and Ba, 2015, Algorithm 1) keeps two moving averages and steps by their ratio:

$$\begin{aligned} m_k &= \beta_1 m_{k-1} + (1 - \beta_1)\, g_k, & \hat m_k &= m_k / (1 - \beta_1^k), \\ v_k &= \beta_2 v_{k-1} + (1 - \beta_2)\, g_k^2, & \hat v_k &= v_k / (1 - \beta_2^k), \\ w_{k} &= w_{k-1} - \eta\, \hat m_k / \big(\sqrt{\hat v_k} + \varepsilon\big), & & \end{aligned}$$

with $m_0 = v_0 = 0$ and the defaults $\eta = 10^{-3}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$.

Where the bias correction comes from (Kingma and Ba, 2015, section 3). Unrolling the recursion, $m_k = (1 - \beta_1)\sum_{i=1}^k \beta_1^{k-i} g_i$. If the gradients have a constant mean $E[g]$, then $E[m_k] = (1 - \beta_1)\sum_{i=1}^k\beta_1^{k-i}\,E[g] = (1 - \beta_1^k)\,E[g]$, by the geometric sum. The average started at zero underestimates the mean by the factor $1 - \beta_1^k$, which is $0.1$ at the first step, and dividing by it removes the bias; the same holds for $v_k$ with $\beta_2$. Without the correction, the first step would be $\eta\cdot 0.1\,g/\sqrt{0.001\,g^2} = 3.16\,\eta$ per coordinate, and with $\beta_2 = 0.999$ the second-moment estimate would stay too small for hundreds of steps.

Worked example. A first gradient $g_1 = (100, 0.01)$. Then $\hat m_1 = g_1$ and $\hat v_1 = g_1^2$, so the first step is $\eta\,(1, 1)$: each coordinate moves by $\eta$ against the sign of its gradient, although the two gradients differ by a factor of $10^4$. Plain gradient descent would move the first coordinate $10^4$ times farther than the second.

Figure

Momentum and Adam on a stretched bowl

The figure shows what that buys and what it does not. When the valley runs along the coordinate axes, as it does when features have different scales but are uncorrelated, dividing by $\sqrt{\hat v}$ equalizes the curvatures much as standardization did, and Adam needs 13 steps against 173 for gradient descent. When the same valley is rotated, as it is for correlated features, no rescaling of the axes can make it round, and Adam needs 171 steps. Gradient descent and momentum depend only on the geometry of the loss, so they take exactly as many steps in both panels; Adam depends on the coordinate system. Its scaling is a diagonal approximation of the curvature, effective when the curvature is mostly diagonal.

Three caveats from practice and theory:

• The step size still matters. On the diabetes data, Adam with $\eta = 0.01$ ends 30 epochs with an excess loss of $836.6$, and with $\eta = 1$ at $8.4$. The defaults suit networks whose weights are of order 1, not a problem whose coefficients are in the tens.
• Adam need not converge. Reddi, Kale and Kumar (2018) construct convex problems on which Adam does not converge to the optimum, for any fixed $\beta_1$, $\beta_2$ in a range that includes the usual $0.9$ and $0.999$ (Theorems 1–3), because the moving average $v$ forgets a rare large gradient; their AMSGrad uses the running maximum of $v$ instead.
• Different solutions. On problems with many minimizers, adaptive methods can converge to different ones than SGD; Wilson et al. (2017) found several tasks where they generalized worse. Weight decay also interacts with the division by $\sqrt{\hat v}$, which is why AdamW applies it separately (weight decay under Adam).

As a check on the formulas, the update above written in NumPy and torch.optim.Adam produce the same weights after 100 steps on the diabetes problem, to $9 \cdot 10^{-15}$ (the Code section).

What changed#

	Gradient descent	SGD	SGD with momentum	Adam
Gradient	exact, $O(nd)$	from $B$ examples, $O(Bd)$	from $B$ examples	from $B$ examples
Constant step	converges	stalls in a ball, variance $\propto\eta$	stalls, ball of the effective step $\eta/(1 - \beta)$	stalls
Dependence on $\kappa$ (quadratic, exact gradients)	$\kappa$	$\kappa$	$\sqrt\kappa$ with tuned $\eta, \beta$	diagonal part removed, rotated part not
Memory per weight	none	none	one velocity	two moments
Typical use	small smooth problems	convex models, with a schedule	vision networks, with a schedule	Transformers, most deep learning (as AdamW)

Limitations#

• Noise needs a schedule. When the gradient noise stays nonzero at the minimum, no constant step converges; the schedule adds hyperparameters. The exception is interpolation: if one weight vector fits every example exactly, as in least squares with more features than examples, every per-example gradient vanishes there and constant-step SGD can converge to it. Overparametrized models that fit their training set are close to this regime, the setting of the double-descent caveat.
• Momentum’s optimal tuning needs $\mu$ and $L$. The $\sqrt\kappa$ rate holds for the best parameters; with $\beta = 0.9$ and a guessed $\eta$ the gain is smaller, and too much momentum oscillates.
• Adaptive methods see only the diagonal. Correlated features leave a rotated valley that per-coordinate steps cannot straighten.
• No method removes the need to tune $\eta$. Defaults work for some model families and fail by orders of magnitude for others.
• Theory and practice diverge for networks. The guarantees are for convex losses; for neural networks the choices rest on experiments.

What you should remember#

• A minibatch gradient is an unbiased estimate of the full gradient with covariance $\Sigma/B$: error $\propto 1/\sqrt B$, cost $\propto B$.
• With a constant step, SGD settles in a noise ball: excess loss about $\eta s^2/4$ in one dimension, $\eta\sigma^2\operatorname{tr}(H)/4$ for least squares. Decreasing steps with $\sum\eta_k = \infty$, $\sum\eta_k^2 < \infty$ converge; in practice, schedules.
• Momentum, $v \leftarrow \beta v + g$, $w \leftarrow w - \eta v$, accumulates consistent directions and cancels oscillating ones; on quadratics it needs $\sqrt\kappa$ instead of $\kappa$ steps. Nesterov’s look-ahead keeps that rate on all smooth strongly convex functions.
• AdaGrad divides by the root of the accumulated squared gradients (good for sparse features, steps shrink forever); RMSProp by a moving average of them.
• Adam = momentum + RMSProp + bias correction $1/(1 - \beta^k)$; its first step is $\eta\,\mathrm{sign}(g)$. It fixes differences of scale between coordinates, not correlations, and its step size still has to be tuned.

Code#

The noise ball, minibatch noise, momentum and Nesterov against the condition number, Adam on the two bowls, the bias correction, a check against PyTorch and the optimizers on the diabetes data:

import numpy as np, torch
from sklearn.datasets import load_diabetes

rng = np.random.default_rng(0)

# the noise ball in one dimension: Q(w) = w^2 / 2, stochastic gradient w + noise with variance 1
for eta in (0.1, 0.01):
    w = np.zeros(20_000)                                 # 20,000 independent runs from w = 0
    for k in range(int(20 / eta)):
        w = w - eta * (w + rng.normal(size=w.size))
    print(eta, round(w.var(), 4), round(eta / (2 - eta), 4))   # simulated and stationary variance of w
# 0.1 0.053 0.0526 | 0.01 0.005 0.005

# minibatch gradients on diabetes (standardized, centred target) at w = 0: the error shrinks like 1/sqrt(B)
X, y = load_diabetes(return_X_y=True)
X = (X - X.mean(0)) / X.std(0); y = y - y.mean(); n = len(y)
G = X * (-y)[:, None]                                    # per-example gradients of (x^T w - y)^2 / 2 at w = 0
g = G.mean(0)
for B in (1, 10, 100):
    err = [np.linalg.norm(G[rng.integers(0, n, B)].mean(0) - g) for _ in range(5000)]
    print(B, round(float(np.sqrt(np.mean(np.square(err)))), 1), round(float(np.sqrt(np.trace(np.cov(G.T, bias=True)) / B)), 1))
# 1 238.6 238.8 | 10 75.6 75.5 | 100 23.9 23.9

# a quadratic with eigenvalues 0.045 and 4.5 (kappa = 100): steps to shrink the excess loss 1000-fold
mu, L = 0.045, 4.5
c, s = np.cos(np.pi / 6), np.sin(np.pi / 6); R = np.array([[c, -s], [s, c]])
H1 = np.diag([L, mu]); H2 = R @ H1 @ R.T                 # the same bowl, axis-aligned and rotated by 30 degrees
def steps(H, w0, method, eta, beta=0.0, K=2000):
    w, v, m, s2 = w0.copy(), np.zeros(2), np.zeros(2), np.zeros(2)
    q0 = w0 @ H @ w0
    for k in range(1, K + 1):
        g = H @ w
        if method == 'gd': w = w - eta * g
        if method == 'momentum': v = beta * v + g; w = w - eta * v
        if method == 'nesterov':
            g = H @ (w - eta * beta * v); v = beta * v + g; w = w - eta * v
        if method == 'adam':
            m = 0.9 * m + 0.1 * g; s2 = 0.999 * s2 + 0.001 * g * g
            w = w - eta * (m / (1 - 0.9 ** k)) / (np.sqrt(s2 / (1 - 0.999 ** k)) + 1e-8)
        if w @ H @ w < 1e-3 * q0: return k
w0 = np.array([-2.4, 1.6])
rk = np.sqrt(L / mu)
print(round((L - mu) / (L + mu), 3), round((rk - 1) / (rk + 1), 3),
      steps(H1, w0, 'gd', 2 / (L + mu)), steps(H1, w0, 'momentum', 4 / (np.sqrt(L) + np.sqrt(mu)) ** 2, ((rk - 1) / (rk + 1)) ** 2),
      steps(H1, w0, 'nesterov', 1 / L, (rk - 1) / (rk + 1)))
# 0.98 0.818 173 39 17: asymptotic per-step factors of GD and momentum; steps of GD, optimal heavy ball, Nesterov
w, v, norms = w0.copy(), np.zeros(2), []
for _ in range(10):
    v = ((rk - 1) / (rk + 1)) ** 2 * v + H1 @ w; w = w - 4 / (np.sqrt(L) + np.sqrt(mu)) ** 2 * v; norms.append(np.linalg.norm(w))
print(round(float(np.linalg.norm(w0)), 2), round(float(max(norms)), 2), int(np.argmax(norms)) + 1)
# 2.88 8.98 4: the optimal heavy ball first moves away, error norm 2.88 -> 8.98 after 4 steps
for name, H, start in (('aligned', H1, w0), ('rotated', H2, R @ w0)):
    print(name, steps(H, start, 'gd', 2 / (L + mu)), steps(H, start, 'momentum', 0.05, 0.9), steps(H, start, 'adam', 0.2))
# aligned 173 48 13 | rotated 173 48 171

# Adam's first step is eta * sign(g) in every coordinate, whatever the gradient's scale; without bias correction 3.16 times that
g1 = np.array([100.0, 0.01])
m1, v1 = 0.1 * g1, 0.001 * g1 ** 2
print((m1 / 0.1) / np.sqrt(v1 / 0.001), (m1 / np.sqrt(v1)).round(2))   # [1. 1.] [3.16 3.16]

# our Adam equals torch.optim.Adam; then SGD, momentum and Adam on diabetes, 30 epochs of minibatches of 32
Xt, yt = torch.tensor(X), torch.tensor(y)
w_hat = np.linalg.lstsq(X, y, rcond=None)[0]; q_hat = 0.5 * np.mean((X @ w_hat - y) ** 2)
loss = lambda w: 0.5 * ((Xt @ w - yt) ** 2).mean()
wt = torch.zeros(10, dtype=torch.float64, requires_grad=True); opt = torch.optim.Adam([wt], lr=0.1)
wm, m, v = np.zeros(10), np.zeros(10), np.zeros(10)
for k in range(1, 101):
    opt.zero_grad(); loss(wt).backward(); opt.step()
    gm = X.T @ (X @ wm - y) / n; m = 0.9 * m + 0.1 * gm; v = 0.999 * v + 0.001 * gm ** 2
    wm = wm - 0.1 * (m / (1 - 0.9 ** k)) / (np.sqrt(v / (1 - 0.999 ** k)) + 1e-8)
print(f'{np.abs(wt.detach().numpy() - wm).max():.0e}')   # 9e-15
def train(make, cut_at=None, seeds=5):                  # mean excess loss after 30 epochs over 5 shuffling seeds
    out = []
    for seed in range(seeds):
        torch.manual_seed(seed)
        wt = torch.zeros(10, dtype=torch.float64, requires_grad=True); opt = make([wt])
        for epoch in range(30):
            if epoch == cut_at:                          # divide the learning rate by 10 halfway
                for group in opt.param_groups: group['lr'] /= 10
            for idx in torch.randperm(n).split(32):
                opt.zero_grad(); (0.5 * ((Xt[idx] @ wt - yt[idx]) ** 2).mean()).backward(); opt.step()
        out.append(loss(wt).item() - q_hat)
    return round(float(np.mean(out)), 1)
print(round(q_hat, 1), round(float(np.linalg.cond(X.T @ X)), 0))   # 1429.8 470.0: the minimum loss and kappa
for name, make in (('sgd 0.05', lambda p: torch.optim.SGD(p, lr=0.05)), ('sgd 0.2', lambda p: torch.optim.SGD(p, lr=0.2)),
                   ('momentum 0.02', lambda p: torch.optim.SGD(p, lr=0.02, momentum=0.9)),
                   ('adam 0.01', lambda p: torch.optim.Adam(p, lr=0.01)), ('adam 1', lambda p: torch.optim.Adam(p, lr=1.0))):
    print(name, train(make), train(make, cut_at=15))
# sgd 0.05 9.5 9.3 | sgd 0.2 50.2 5.1 | momentum 0.02 12.6 5.4 | adam 0.01 836.6 1083.5 | adam 1 8.4 8.7

Next#

Least squares can now be fitted for any $n$ and $d$, by formula or by descent. Being able to fit it does not make the fit good. With many features, or features that nearly repeat each other, the least-squares note showed coefficients that swing with every draw of the noise and a training $R^2$ that rises while the validation $R^2$ collapses. The next note adds a penalty on the size of the weights: ridge and the lasso, the priors behind them, and how to choose the penalty’s strength by validation.

The optimizers of this note train the networks of the other courses: Adam and AdamW are the defaults of the Generative Models training runs, and the next section of this course uses them for neural networks.

References#

• E. Sokolov et al. Машинное обучение (HSE), lecture notes, lecture 3: estimating the gradient, stochastic gradient descent, momentum, AdaGrad, RMSProp and Adam.
• H. Robbins, S. Monro. A stochastic approximation method. Annals of Mathematical Statistics, 22(3), 1951.
• B. T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5), 1964.
• I. Sutskever, J. Martens, G. Dahl, G. Hinton. On the importance of initialization and momentum in deep learning. ICML, 2013. Section 2 (Nesterov’s method as momentum with a look-ahead gradient).
• G. Goh. Why momentum really works. Distill, 2017. The section on optimal parameters.
• S. Bubeck. Convex optimization: algorithms and complexity. Foundations and Trends in Machine Learning, 8(3–4), 2015. Theorem 3.15 (the lower bound for first-order methods), Theorem 3.18 (Nesterov’s accelerated method), Theorem 6.2 (SGD for strongly convex functions).
• J. Duchi, E. Hazan, Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12, 2011.
• G. Hinton, N. Srivastava, K. Swersky. Neural networks for machine learning, lecture 6: overview of mini-batch gradient descent. Lecture materials, University of Toronto, 2012. Lecture 6e (RMSProp).
• D. P. Kingma, J. Ba. Adam: a method for stochastic optimization. ICLR, 2015. Algorithm 1, Section 3 (initialization bias correction).
• S. J. Reddi, S. Kale, S. Kumar. On the convergence of Adam and beyond. ICLR, 2018. Theorems 1–3, Section 4 (AMSGrad).
• A. C. Wilson, R. Roelofs, M. Stern, N. Srebro, B. Recht. The marginal value of adaptive gradient methods in machine learning. NeurIPS, 2017.
• I. Goodfellow, Y. Bengio, A. Courville. Deep Learning. MIT Press, 2016. Section 8.1.3 (batch and minibatch algorithms), Section 8.3 (SGD, momentum, Nesterov momentum), Section 8.5 (AdaGrad, RMSProp, Adam), Section 8.7.3 (Polyak averaging).