Adaptive moment estimation optimizer

Until now, we have used the stochastic gradient descent (SGD) optimizer to update parameters. Its logic is very simple: multiply the gradient by a fixed learning rate, and that is it.

But it has a flaw: one size fits all. All the parameters share one learning rate, and it never changes from start to finish. In practice, different parameters, and different stages of training, need different step sizes:

Some parameters always have large gradients and others always have small ones. With the same learning rate, the large ones tend to overshoot and the small ones barely move;
In the early stage of training we can usually take bold steps, while near convergence we need to fine-tune carefully.

In this chapter, we introduce a smarter optimizer: the adaptive moment estimation (Adam) optimizer. It adjusts the effective step size of each parameter individually. Put simply: when the gradient has stayed small for a long time, Adam increases the effective step size; when the gradient has stayed large for a long time, Adam decreases it. Adam is the most commonly used optimizer for training neural network models today.

from abc import ABC, abstractmethod

import numpy as np

np.random.seed(42)

Tensor¶

class Tensor:

    def __init__(self, data):
        self.data = np.array(data)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = set()

    def backward(self):
        topo = []
        visited = set()
        stack = [(self, False)]

        while stack:
            node, expanded = stack.pop()
            if node in visited:
                continue

            if expanded:
                visited.add(node)
                topo.append(node)
            else:
                stack.append((node, True))
                for p in node.parents:
                    if p not in visited:
                        stack.append((p, False))

        self.grad = np.ones_like(self.data)
        for t in reversed(topo):
            if t.gradient_fn is not None:
                t.gradient_fn()

        for t in topo:
            t.gradient_fn = lambda: None
            t.parents = set()

    def __add__(self, other):
        p = Tensor(self.data + other.data)

        def gradient_fn():
            self.grad += p.grad
            other.grad += p.grad

        p.gradient_fn = gradient_fn
        p.parents = [self, other]
        return p

    def __str__(self):
        return f'Tensor({self.data})'

Dataset¶

Abstract class¶

class Dataset(ABC):

    def __init__(self, batch_size=1):
        self.batch_size = batch_size
        self.load()
        self.train()

    @abstractmethod
    def load(self):
        pass

    def train(self):
        self.data = self.train_data

    def eval(self):
        self.data = self.test_data

    def all(self):
        x, y = self.data
        return Tensor(x), Tensor(y)

    def __len__(self):
        x, *_ = self.data
        return len(x) // self.batch_size

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y = self.data
        return Tensor(x[s]), Tensor(y[s])

Character-encoding dataset¶

class CharDataset(Dataset):

    def __init__(self, filename, batch_size=1, context_size=64, stride=None, split=0.9):
        self.filename = filename
        self.context_size = context_size
        self.stride = stride if stride is not None else context_size // 2
        self.split = split
        super().__init__(batch_size)

    def load(self):
        with open(self.filename, encoding="utf-8") as f:
            text = f.read()

        self.vocab = sorted(set(text))
        self.vocab_size = len(self.vocab)
        self.stoi = {ch: i for i, ch in enumerate(self.vocab)}
        self.itos = {i: ch for i, ch in enumerate(self.vocab)}
        self.tokens = self.encode(text)

        split = int(len(self.tokens) * self.split)
        self.train_data = self._pack(self.tokens[:split])
        self.test_data = self._pack(self.tokens[split:])

    def _pack(self, tokens):
        x, y = [], []
        for i in range(0, len(tokens) - self.context_size - 1, self.stride):
            x.append(tokens[i: i + self.context_size])
            y.append(tokens[i + 1: i + self.context_size + 1])
        return x, y

    def encode(self, symbols):
        return [self.stoi[s] for s in symbols]

    def decode(self, tokens):
        return "".join(self.itos[t] for t in tokens)

Model¶

Abstract layer¶

class Layer(ABC):

    def __call__(self, *args):
        return self.forward(*args)

    @abstractmethod
    def forward(self, *args):
        pass

    @property
    def parameters(self):
        return []

Linear layer¶

class Linear(Layer):

    def __init__(self, in_size, out_size):
        self.weight = Tensor(np.random.randn(out_size, in_size) * np.sqrt(2 / in_size))
        self.bias = Tensor(np.zeros(out_size))

    def forward(self, x: Tensor):
        p = Tensor(x.data @ self.weight.data.T + self.bias.data)

        def gradient_fn():
            self.weight.grad += p.grad.T @ x.data
            self.bias.grad += np.sum(p.grad, axis=0)
            x.grad += p.grad @ self.weight.data

        p.gradient_fn = gradient_fn
        p.parents = [x]
        return p

    @property
    def parameters(self):
        return [self.weight, self.bias]

Composite layer¶

class Composite(Layer, ABC):

    def __init__(self, layers):
        super().__init__()
        self.layers = list(layers)

    @property
    def parameters(self):
        return [p for l in self.layers for p in l.parameters]

Sequential layer¶

class Sequential(Composite):

    def forward(self, x: Tensor):
        for l in self.layers:
            x = l(x)
        return x

Embedding layer¶

class Embedding(Layer):

    def __init__(self, vocab_size, embedding_size, std=0.02):
        super().__init__()
        self.weight = Tensor(np.random.randn(vocab_size, embedding_size) * std)

    def forward(self, x: Tensor):
        p = Tensor(self.weight.data[x.data])

        def gradient_fn():
            np.add.at(self.weight.grad, x.data, p.grad)

        p.gradient_fn = gradient_fn
        p.parents = [self.weight]
        return p

    @property
    def parameters(self):
        return [self.weight]

Activation functions¶

Hyperbolic tangent (Tanh)¶

class Tanh(Layer):

    def forward(self, x: Tensor):
        a = Tensor(np.tanh(x.data))

        def gradient_fn():
            x.grad += a.grad * (1 - a.data ** 2)

        a.gradient_fn = gradient_fn
        a.parents = [x]
        return a

Loss functions¶

Abstract class¶

class Loss(ABC):

    def __call__(self, p: Tensor, y: Tensor):
        return self.loss(p, y)

    @abstractmethod
    def loss(self, p: Tensor, y: Tensor):
        pass

Cross-entropy loss function¶

class CELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        if y.data.ndim == softmax.ndim:
            labels = y.data
        else:
            labels = np.zeros_like(softmax)
            rows = np.arange(y.data.size)
            labels[rows, y.data] = 1

        log = np.log(np.clip(softmax, 1e-10, 1.0))
        ce = Tensor(-np.mean(np.sum(labels * log, axis=-1)))

        def gradient_fn():
            p.grad += (softmax - labels) / labels.shape[0]

        ce.gradient_fn = gradient_fn
        ce.parents = [p]
        return ce

Optimizers¶

Before introducing the new optimizer, we also first extract an abstract class for optimizers.

Abstract class¶

The abstract class defines two methods:

step (step method): an abstract method that updates all the parameters from the gradients;
zero_grad (zero-gradient method): zeroes the gradients of all the parameters before each iteration.

class Optimizer(ABC):

    def __init__(self, parameters, lr):
        self.parameters = list(parameters)
        self.lr = lr

    @abstractmethod
    def step(self):
        pass

    def zero_grad(self):
        for p in self.parameters:
            p.grad = np.zeros_like(p.data)

Stochastic gradient descent optimizer¶

We also change the existing stochastic gradient descent optimizer to inherit from the abstract class. It only needs to implement the step method.

class SGDOptimizer(Optimizer):

    def step(self):
        for p in self.parameters:
            p.data -= p.grad * self.lr

Adaptive moment estimation optimizer¶

The moment in the name of the adaptive moment estimation optimizer refers to two statistics of the gradient. Adam maintains two moments for each parameter individually:

$m$ (the first moment): the exponentially weighted moving average (EWMA) of the gradient. It acts like inertia on the gradient, making the update direction smoother so that it is not thrown off by a single random fluctuation;
$v$ (the second moment): the exponentially weighted moving average of the squared gradient, which reflects the magnitude of the gradient. For a parameter whose gradient is always large, $v$ is large and the actual step size is automatically reduced; conversely, for a parameter whose gradient is always small, $v$ is small and the actual step size is automatically enlarged.

Exponentially weighted moving average¶

The purpose of a weighted moving average is to make the changes of the gradient (or the squared gradient) smoother. The formulas are:

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

(1)

v_t = \beta_2 v_{t-1} + (1 - \beta_2){g_t}^2

(2)

Where:

$m_t$ , $v_t$ : the exponentially weighted averages at the current step $t$ ;
$m_{t-1}$ , $v_{t-1}$ : the exponentially weighted averages at the previous step (step $t-1$ );
$g_t$ , $g_t^2$ : the current gradient and squared gradient;
$\beta_1$ , $\beta_2$ : the decay rates of the gradient (hyperparameters, usually 0.9 and 0.999), which decide the weight given to past information.

Step size scaling¶

When updating, $m$ decides the direction, and dividing by $\sqrt{v}$ scales the step size:

p \;{-}{=}\; lr \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}

(3)

Where:

$p$ : the parameter;
$lr$ : the learning rate;
$\hat{m}$ : the exponentially weighted average of the gradient;
$\hat{v}$ : the exponentially weighted average of the squared gradient;
$\epsilon$ : a tiny number that prevents division by 0.

Both $\hat{m}$ and $\hat{v}$ here need bias correction.

Bias correction¶

In the early stage of training, $m_0$ and $v_0$ are usually initialized to 0. As a result, in the first few steps the values of $m_t$ and $v_t$ are far too small (close to 0). To remove this error, we usually apply bias correction. Taking $m$ as an example, the formula is:

\hat{m_t} = \frac{m_t}{1 - {\beta_1}^t}

(4)

Where:

$\hat{m_t}$ : the corrected exponentially weighted average;
$m_t$ : the exponentially weighted average;
$\beta_1$ : the decay rate of the gradient;
$t$ : the number of steps the optimizer has updated so far.

As $t$ grows, ${\beta_1}^t$ approaches 0 and the denominator approaches 1, so the correction naturally fades away.

⚠️ Note: The step count $t$ here is the number of times the optimizer has updated the parameters. It is not the same thing as the “time step” of reading one character in the RNN of the previous chapter.

Gradient clipping¶

The adaptive moment estimation optimizer also has a gradient clipping method, clip_grad_norm.

The gradient computation of an RNN passes through many time steps and occasionally becomes suddenly huge (exploding gradients), throwing the parameters far off. The purpose of gradient clipping is this: if the total norm of the gradients of all the parameters taken together exceeds a threshold, they are all scaled down proportionally, keeping the step size within a safe range.

class AdamOptimizer(Optimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8):
        super().__init__(parameters, lr)
        self.b1, self.b2 = betas
        self.eps = eps
        # First moment
        self.m = [None] * len(parameters)
        # Second moment
        self.v = [None] * len(parameters)
        # Number of update steps
        self.t = 0

    def step(self):
        self.t += 1
        for i, p in enumerate(self.parameters):
            if p is not None:
                if self.m[i] is None:
                    self.m[i] = np.zeros_like(p.data)
                    self.v[i] = np.zeros_like(p.data)

                # Compute the first moment
                self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * p.grad
                # Compute the second moment
                self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (p.grad ** 2)
                # Bias correction
                m = self.m[i] / (1 - self.b1 ** self.t)
                v = self.v[i] / (1 - self.b2 ** self.t)
                # Weight decay (AdamW)
                self._apply_weight_decay(p)
                # Update the parameters
                p.data -= self.lr * m / (np.sqrt(v) + self.eps)

    # Weight decay (hook method)
    def _apply_weight_decay(self, p):
        pass

    # Gradient clipping method
    def clip_grad_norm(self, max_norm=1.0):
        sq = 0.0
        for p in self.parameters:
            sq += float(np.sum(p.grad ** 2))

        if np.sqrt(sq) > max_norm > 0:
            scale = max_norm / (np.sqrt(sq) + 1e-6)
            for p in self.parameters:
                p.grad *= scale

Adaptive moment estimation with weight decay optimizer¶

The adaptive moment estimation with weight decay (AdamW) optimizer is a small improvement on Adam, and is standard equipment for training large models such as GPT. On top of Adam, it adds weight decay: at every update, the weights themselves are also pulled a tiny bit toward 0.

The purpose is to prevent overfitting: if no weight is allowed to grow too large, the model is less likely to “memorize” the details of the training data.

Note that only matrices (ndim >= 2) are decayed; one-dimensional parameters such as biases are not. This is also a rule of thumb from practice.

class AdamWOptimizer(AdamOptimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
        super().__init__(parameters, lr, betas, eps)
        self.weight_decay = weight_decay

    # Weight decay
    def _apply_weight_decay(self, p):
        if p.data.ndim >= 2:
            p.data -= p.data * self.weight_decay * self.lr

Recurrent neural network¶

Recurrent cell¶

class Cell(Composite):

    def __init__(self, in_size, out_size):
        self.input = Linear(in_size, out_size)
        self.hidden = Linear(out_size, out_size)
        self.tanh = Tanh()

        super().__init__([self.input,
                          self.hidden,
                          self.tanh])

    def forward(self, x: Tensor, h: Tensor):
        p = self.input(x)
        h = self.hidden(h)
        return self.tanh(p + h)

RNN model¶

class RNN(Composite):

    def __init__(self, vocab_size, embedding_size):
        self.embedding_size = embedding_size

        self.embedding = Embedding(vocab_size, embedding_size)
        self.cell = Cell(embedding_size, embedding_size * 2)
        self.output = Linear(embedding_size * 2, vocab_size)

        super().__init__([self.embedding,
                          self.cell,
                          self.output])

    def step(self, token: Tensor, h: Tensor):
        p = self.embedding(token)
        h = self.cell(p, h)
        return self.output(h), h

    def forward(self, x: Tensor, h: Tensor = None):
        batch_size, sequence_len = x.data.shape
        if h is None:
            h = Tensor(np.zeros((batch_size, self.embedding_size * 2)))

        ps = []
        for s in range(sequence_len):
            p, h = self.step(Tensor(x.data[:, s]), h)
            ps.append(p)

        return ps, h

Trainer¶

Abstract class¶

class Trainer(ABC):

    def __init__(self, model, loss_fn, optimizer):
        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    @abstractmethod
    def train(self, dataset, epochs):
        pass

    @abstractmethod
    def test(self, dataset):
        pass

RNN trainer¶

In the trainer, before calling the step method step, we first call the gradient clipping method clip_grad_norm to prevent exploding gradients.

class RNNTrainer(Trainer):

    def train(self, dataset, epochs):
        dataset.train()

        for epoch in range(epochs):
            for i in range(len(dataset)):
                feature, label = dataset[i]

                self.optimizer.zero_grad()
                prediction, _ = self.model(feature)
                loss = Tensor(0.0)
                for t in range(len(prediction)):
                    loss += self.loss_fn(prediction[t], Tensor(label.data[:, t]))
                loss.backward()
                # Gradient clipping
                self.optimizer.clip_grad_norm()
                self.optimizer.step()

    def test(self, dataset):
        dataset.eval()

        feature, label = dataset.all()
        prediction, _ = self.model(feature)
        loss = Tensor(0.0)
        for t in range(len(prediction)):
            loss += self.loss_fn(prediction[t], Tensor(label.data[:, t]))
        return prediction, loss

    def generate(self, dataset, prompt, steps=256):
        tokens = dataset.encode(prompt)

        h = Tensor(np.zeros((1, self.model.embedding_size * 2)))
        for token in tokens:
            p, h = self.model.step(Tensor([token]), h)

        for _ in range(steps):
            exp = np.exp(p.data[0] - np.max(p.data[0]))
            probs = exp / np.sum(exp)
            token = np.random.choice(len(probs), p=probs)
            tokens.append(token)
            p, h = self.model.step(Tensor([token]), h)

        return dataset.decode(tokens)

Hyperparameters¶

Learning rate¶

LEARNING_RATE = 0.001

Batch size¶

BATCH_SIZE = 4

Number of epochs¶

EPOCHS = 2

Context size¶

CONTEXT_SIZE = 32

Embedding size¶

EMBEDDING_SIZE = 16

Modeling¶

We change the optimizer to the adaptive moment estimation with weight decay (AdamW) optimizer.

DATA_FILE = "../../../../data/tinyshakespeare.txt"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = RNN(dataset.vocab_size, EMBEDDING_SIZE)
loss_fn = CELoss()
# Adaptive moment estimation with weight decay (AdamW) optimizer
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = RNNTrainer(model, loss_fn, optimizer)

Training¶

trainer.train(dataset, EPOCHS)

Evaluation¶

prediction, loss = trainer.test(dataset)
print(f'prediction:\t{len(prediction)} steps, each {prediction[0].data.shape}\ntotal loss:\t{loss}')

prediction:	32 steps, each (6970, 65)
total loss:	Tensor(65.93744561104684)

Inference¶

print(trainer.generate(dataset, prompt="ROMEO:"))

ROMEO:
I not beting you at sheenes hew so ene I's have boo?

GRUMPIO:
Westeesilv bees
tha my sapin, it and ofte, couther me it,
Ane sceat, bot, mance?

Prevees live hem?

GOMYe leald het emen, Beran: I but graed cruith, quep a you geed, I
Anthing will sunf a Sor

In this chapter, we upgraded the optimizer from the “one size fits all” SGD to the “tailor-made” AdamW:

Adam: adjusts the effective step size of each parameter individually (momentum for the direction, scale for the step size);
Gradient clipping: prevents exploding gradients;
Weight decay: prevents overfitting.

💡 Tip: In general, Adam and AdamW converge faster than SGD early in training, but their final generalization may be worse than that of a carefully tuned SGD. For Transformer-based models, AdamW remains the most commonly used choice.

Exercises¶

Try making the threshold max_norm of the gradient clipping method clip_grad_norm larger or smaller, and observe how training differs.