Learning rate scheduler - From Neuron to GPT

In the previous chapter, Adam adjusted the effective step size of each parameter individually. Another hyperparameter related to the step size is the learning rate, which we have still been tuning by hand.

So far, the learning rate has been fixed once training starts: whatever it was at the beginning, it stays the same until the end. But this is not reasonable:

Early in training: the parameters are still randomly initialized and the gradient direction is unreliable. If we rush ahead with a large learning rate from the start, we can easily go off course;
Late in training: the model is already close to the optimum and needs careful fine-tuning rather than more big steps. Otherwise it swings back and forth around the optimum, or even jumps out of it.

In this chapter, we introduce automatic tuning: a learning rate scheduler changes the learning rate dynamically during training. We will use the scheme most commonly used to train large models: warmup and cosine annealing.

Warmup: early in training, the learning rate starts from 0 and gradually climbs to its maximum;
Cosine annealing: after warmup, the learning rate falls smoothly from the maximum to the minimum along a cosine curve.

import math
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¶

Abstract class¶

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)

Adaptive moment estimation optimizer¶

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
        self.m = [None] * len(parameters)
        self.v = [None] * len(parameters)
        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)

                self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * p.grad
                self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (p.grad ** 2)
                m = self.m[i] / (1 - self.b1 ** self.t)
                v = self.v[i] / (1 - self.b2 ** self.t)
                self._apply_weight_decay(p)
                p.data -= self.lr * m / (np.sqrt(v) + self.eps)

    def _apply_weight_decay(self, p):
        pass

    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¶

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

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

Learning rate scheduler (warmup cosine annealing)¶

To create a learning rate scheduler, we need to provide four initialization parameters:

lr (learning rate): the learning rate we have been using all along, which is also the maximum learning rate reached after warmup;
total_steps (total steps): the number of steps in the whole training process;
warmup_steps (warmup steps): the number of steps in the warmup phase;
min_lr (minimum learning rate): the minimum value the learning rate falls to late in training.

Warmup¶

In the first warmup_steps steps of training, the learning rate climbs linearly from 0 to the maximum lr:

lr_t = lr \cdot \frac{t}{warmup\_steps}

(1)

Where:

$lr_t$ : the current learning rate;
$t$ : the current step count (counted from 1, which corresponds to current_step + 1 in the code).

Cosine annealing¶

After warmup, the learning rate has reached its maximum, and then falls smoothly along a cosine curve from the maximum to the minimum $min\_lr$ :

lr_t = min\_lr + (lr - min\_lr) \cdot \frac{1 + \cos(\pi \cdot \frac{t - warmup\_steps}{total\_steps - warmup\_steps})}{2}

(2)

Where:

$lr_t$ : the current learning rate;
$t$ : the current step count (which corresponds to current_step in the code).

The benefit of annealing along a cosine curve is that the step size goes slow, then fast, then slow again. So over the whole training run, the learning rate is like driving a car: a slow start (warmup), acceleration to cruising speed in the middle, and a smooth slowdown as you near the destination (cosine annealing). This is the standard training rhythm of large models such as GPT.

The learning rate scheduler provides only one method:

step (step method): computes the new learning rate from the current step count.

class WarmupCosineScheduler:

    def __init__(self, lr, total_steps, warmup_steps, min_lr=0.0):
        self.lr = lr
        self.total_steps = max(total_steps, 1)
        self.warmup_steps = max(min(warmup_steps, self.total_steps), 0)
        self.min_lr = min_lr

    # Step method
    def step(self, current_step):
        # Warmup
        if self.warmup_steps > 0 and current_step < self.warmup_steps:
            return self.lr * (current_step + 1) / self.warmup_steps

        # Hold the minimum learning rate
        if current_step >= self.total_steps:
            return self.min_lr

        # Cosine annealing
        progress = (current_step - self.warmup_steps) / max(self.total_steps - self.warmup_steps, 1)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return self.min_lr + (self.lr - self.min_lr) * cosine

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, scheduler=None):
        pass

    @abstractmethod
    def test(self, dataset):
        pass

RNN trainer¶

In the trainer, every step first calls the step method step of the learning rate scheduler to get the latest learning rate.

class RNNTrainer(Trainer):

    def train(self, dataset, epochs, scheduler=None):
        dataset.train()

        steps = 0
        for epoch in range(epochs):
            for i in range(len(dataset)):
                # Update the learning rate
                if scheduler is not None:
                    self.optimizer.lr = scheduler.step(steps)

                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()
                self.optimizer.clip_grad_norm()
                self.optimizer.step()
                steps += 1

    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¶

With a learning rate scheduler, LEARNING_RATE is in fact used as the maximum learning rate. The scheduler lowers the average learning rate, so we can raise it a little, and set the minimum learning rate to 1/10 of the maximum.

LEARNING_RATE = 0.005

Batch size¶

BATCH_SIZE = 4

Number of epochs¶

EPOCHS = 2

Context size¶

CONTEXT_SIZE = 32

Embedding size¶

EMBEDDING_SIZE = 16

Modeling¶

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = RNN(dataset.vocab_size, EMBEDDING_SIZE)
loss_fn = CELoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = RNNTrainer(model, loss_fn, optimizer)

Training¶

Before training the model, we first create an instance of the learning rate scheduler and pass it to the training method. The initialization parameters here mean:

LEARNING_RATE: the maximum learning rate;
EPOCHS * len(dataset): the total number of training steps;
100: the number of warmup steps;
LEARNING_RATE / 10: the minimum learning rate.

# Learning rate scheduler
scheduler = WarmupCosineScheduler(LEARNING_RATE, EPOCHS * len(dataset), 100, LEARNING_RATE / 10)
trainer.train(dataset, EPOCHS, scheduler)

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(64.7635956614507)

Inference¶

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

ROMEO:
Hat that, he you!
You spelr gexssigean I'red
ho,
That his whirovellm to waind:
Ye Gosis
The rest patiet,
Pervunded mantlacem! But do stay,
Hath muthood;
I wore not hanted gaing, you knes is not thought be but what you a you geeden Hartion as bearter a Pro

In this chapter, we implemented automatic tuning of the learning rate: warmup lets training start smoothly, and cosine annealing lets it converge finely at the end. With that, the three practical workhorses of model training, the Adam optimizer, gradient clipping, and learning rate scheduling, are all in place. This is also the standard setup for training large models.

Looking back over chapters 11 to 16, we started from digitizing text and implemented an RNN, creating a model that can read text.

But we must point out that the RNN has inherent flaws, limited memory and no parallelism, which keep it from carrying the weight of building a real large language model. In the next part, we take a completely new approach, the attention mechanism, and build a real GPT model step by step.

Exercises¶

Try changing the initialization parameters of the learning rate scheduler. For example, set the warmup steps to 0 (no warmup) or to a very large number, or set the minimum learning rate to 0 so that the learning rate finally falls to 0. Observe whether training gets faster or slower and whether the results get better or worse.