Skip to content

PigeonColumba compressaTheme

Lesson 4 of 10

The Adam optimizer

Plain gradient descent uses one learning rate for every parameter. Adam keeps two running averages per parameter, the gradient (m, momentum) and the squared gradient (v), and steps by m / sqrt(v): large where the gradient is consistent, small where it is noisy or huge.

For step t = 1, 2, ...:

m = b1 * m + (1 - b1) * g
v = b2 * v + (1 - b2) * g * g
m_hat = m / (1 - b1**t), v_hat = v / (1 - b2**t) (bias correction: both start at 0)
p -= lr * m_hat / (sqrt(v_hat) + eps)

Implement it as a class that updates arrays in place, like lab/nn.py:Adam (which also adds decoupled weight decay).

Your code

import numpy as np


class Adam:
    def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
        """params: list of numpy arrays to update in place."""
        raise NotImplementedError

    def step(self, grads):
        """grads: list of arrays, same shapes as params."""
        raise NotImplementedError

Checking your code and saving a draft need JavaScript, which is turned off in this browser. Turn it on and reload the page to check your code. You can still edit it here, and the hints and the reference solution are below.

Pigeon's hints (2)

Hint 1 of 2

Store one m and one v array per parameter, initialized with np.zeros_like.

Hint 2 of 2

Use in-place operations (p -= ...) so the caller's arrays change.

The reference solution

import numpy as np


class Adam:
    def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
        self.p, self.lr, self.b1, self.b2, self.eps = params, lr, b1, b2, eps
        self.m = [np.zeros_like(x) for x in params]
        self.v = [np.zeros_like(x) for x in params]
        self.t = 0

    def step(self, grads):
        self.t += 1
        for p, g, m, v in zip(self.p, grads, self.m, self.v):
            m *= self.b1; m += (1 - self.b1) * g
            v *= self.b2; v += (1 - self.b2) * g * g
            mh = m / (1 - self.b1 ** self.t)
            vh = v / (1 - self.b2 ** self.t)
            p -= self.lr * mh / (np.sqrt(vh) + self.eps)

Where this is used

Knowledge distillation (Hinton)Runs hereTrain a small model to match a big model's softened output probabilities instead of (or as well as) the one-hot labels.