Dropout - From Neuron to GPT

In the previous chapter, we added layer normalization to the GPT model, which keeps the values and the training stable after the residual connections. We also mentioned a possible problem of large-scale training: overfitting.

The model performs well on the training set but may perform poorly on the test set. It has memorized too many details of the training data instead of truly learning the patterns of language. It is like a student who just memorizes all the practice problems by rote and is at a loss when facing new problems in the exam.

Dropout is a classic technique against overfitting: during training, randomly drop the outputs of a part of the neurons (set them to 0), so that the model cannot depend too much on any single neuron:

During training: randomly set the outputs of a certain proportion $p$ (for example 0.1) of the neurons to 0, and divide the remaining outputs by $1 - p$ to keep the expected value unchanged;
During inference: do not drop any neuron, and use all the outputs.

GPT basic modules with dropout: dropout comes after the embedding layer and before the residual connections of the attention layer and the feed-forward network layer

Figure: GPT basic modules with dropout

import numpy as np

from src.core import *

np.random.seed(42)

Dataset¶

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)

Lower triangular matrix layer¶

class Tril(Layer):

    def __init__(self, value=-1e9):
        super().__init__()
        self.value = value

    def forward(self, x: Tensor):
        keep = np.tril(np.ones(x.shape[-2:]))
        p = Tensor(np.where(keep, x.data, self.value))

        def gradient_fn():
            x.grad += p.grad * keep

        return p.attach(gradient_fn, {x})

Normalization layer¶

class LayerNorm(Layer):

    def __init__(self, normalized_size, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = Tensor(np.ones(normalized_size, dtype=DTYPE))
        self.bias = Tensor(np.zeros(normalized_size, dtype=DTYPE))

    def forward(self, x: Tensor):
        mean = np.mean(x.data, axis=-1, keepdims=True)
        var = np.var(x.data, axis=-1, keepdims=True, ddof=0)
        norm = (x.data - mean) / np.sqrt(var + self.eps)
        p = Tensor(self.weight.data * norm + self.bias.data)

        def gradient_fn():
            axis = tuple(range(p.grad.ndim - 1)) if p.grad.ndim > 1 else None
            self.weight.grad += np.sum(p.grad * norm, axis=axis)
            self.bias.grad += np.sum(p.grad, axis=axis)
            grad = p.grad * self.weight.data
            grad_mean = np.mean(grad, axis=-1, keepdims=True)
            norm_mean = np.mean(grad * norm, axis=-1, keepdims=True)
            x.grad += (grad - grad_mean - norm * norm_mean) / np.sqrt(var + self.eps)

        return p.attach(gradient_fn, {x})

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

Dropout layer¶

The dropout layer (the Dropout layer) randomly drops the outputs of a part of the neurons during training:

Forward pass:
- During training: generate a random mask with the same shape as the input, keep each value with probability $1 - p$ (with a scaling factor of $1 / (1 - p)$ ) and drop it with probability $p$ (the value becomes 0);
- During inference: return the input directly;
Backward pass: the gradient flows back only through the neurons that were kept, and the gradient of the dropped neurons is 0.

💡 Tip: dividing the kept input values by $1 - p$ keeps the overall expected value of the output unchanged, so that the output scale is the same in training and in inference.

class Dropout(Layer):

    def __init__(self, prob=0.1):
        super().__init__()
        self.prob = prob

    def forward(self, x: Tensor):
        # No dropout at inference
        if not self.training or self.prob == 0:
            return x

        # Random mask
        keep_prob = 1.0 - self.prob
        mask = (np.random.rand(*x.shape) < keep_prob).astype(DTYPE) / keep_prob
        # Create the output value tensor
        p = Tensor(x.data * mask)

        # Gradient (closure) function of the output value tensor
        def gradient_fn():
            x.grad += p.grad * mask

        # Set the gradient (closure) function and the parent node list of the output value tensor
        return p.attach(gradient_fn, {x})

Activation functions¶

Gaussian error linear unit (GELU)¶

class GELU(Layer):

    def __init__(self):
        super().__init__()
        self.c = np.sqrt(2.0 / np.pi)

    def forward(self, x: Tensor):
        tanh = np.tanh(self.c * (x.data + 0.044715 * x.data ** 3))
        a = Tensor(0.5 * x.data * (1.0 + tanh))

        def gradient_fn():
            grad = 0.5 * (1.0 + tanh) + 0.5 * x.data * (1.0 - tanh ** 2) * self.c * (1.0 + 3.0 * 0.044715 * x.data ** 2)
            x.grad += a.grad * grad

        return a.attach(gradient_fn, parents=[x])

Loss function¶

GPT cross-entropy loss function¶

class GPTLoss(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)

        flat_softmax = softmax.reshape(-1, softmax.shape[-1])
        flat_y = y.data.reshape(-1).astype(np.int64)
        n = len(flat_y)
        rows = np.arange(n)

        log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
        ce = Tensor(0 - np.sum(log) / n)

        def gradient_fn():
            flat_grad = flat_softmax.copy()
            flat_grad[rows, flat_y] -= 1
            p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n

        return ce.attach(gradient_fn, parents=[p])

Model¶

GPT embedding layer¶

After the word vector and the position vector are combined, we add dropout.

class GPTEmbedding(Composite):

    def __init__(self, vocab_size, context_size, embedding_size):
        self.embedding = Embedding(vocab_size, embedding_size)
        self.positional_embedding = Embedding(context_size, embedding_size)
        # Dropout layer
        self.dropout = Dropout()

        super().__init__([self.embedding,
                          self.positional_embedding,
                          self.dropout])

    def forward(self, x: Tensor):
        token = self.embedding(x)
        position = self.positional_embedding(Tensor(range(x.shape[1])))
        # Dropout
        return self.dropout(token + position)

GPT attention layer¶

In the attention layer, dropout is added before the residual connection.

class GPTAttention(Composite):

    def __init__(self, embedding_size, heads=1):
        self.embedding_size = embedding_size
        self.heads = heads

        self.normalize = LayerNorm(embedding_size)
        self.query = Linear(embedding_size, embedding_size)
        self.key = Linear(embedding_size, embedding_size)
        self.value = Linear(embedding_size, embedding_size)
        self.mask = Tril()
        self.softmax = Softmax()
        self.output = Linear(embedding_size, embedding_size)
        # Dropout layer
        self.dropout = Dropout()

        super().__init__([self.normalize,
                          self.query,
                          self.key,
                          self.value,
                          self.mask,
                          self.softmax,
                          self.output,
                          self.dropout])

    def forward(self, x: Tensor):
        norm = self.normalize(x)
        head_embedding_size = self.embedding_size // self.heads
        multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
        shape = (-1, x.shape[1], self.embedding_size)

        query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
        value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
        scores = query @ key * scale
        weights = self.softmax(self.mask(scores))
        # Dropout, then residual connection
        return x + self.dropout(self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape)))

GPT feed-forward network layer¶

Likewise, in the feed-forward network layer, dropout is also added before the residual connection.

class GPTFeedForward(Composite):

    def __init__(self, embedding_size):
        self.normalize = LayerNorm(embedding_size)
        self.input = Linear(embedding_size, embedding_size * 4)
        self.gelu = GELU()
        self.output = Linear(embedding_size * 4, embedding_size)
        # Dropout layer
        self.dropout = Dropout()

        super().__init__([self.normalize,
                          self.input,
                          self.gelu,
                          self.output,
                          self.dropout])

    def forward(self, x: Tensor):
        norm = self.normalize(x)
        h = self.gelu(self.input(norm))
        # Dropout, then residual connection
        return x + self.dropout(self.output(h))

GPT Transformer block¶

class GPTTransformer(Composite):

    def __init__(self, embedding_size, heads):
        self.attention = GPTAttention(embedding_size, heads)
        self.feed_forward = GPTFeedForward(embedding_size)

        super().__init__([self.attention,
                          self.feed_forward])

    def forward(self, x: Tensor):
        x = self.attention(x)
        return self.feed_forward(x)

GPT output layer¶

class GPTOutput(Composite):

    def __init__(self, embedding_size, vocab_size):
        self.normalize = LayerNorm(embedding_size)
        self.output = Linear(embedding_size, vocab_size)

        super().__init__([self.normalize,
                          self.output])

    def forward(self, x: Tensor):
        norm = self.normalize(x)
        return self.output(norm)

GPT model¶

class GPT(Composite):

    def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
        self.embedding = GPTEmbedding(vocab_size, context_size, embedding_size)
        self.transformers = [GPTTransformer(embedding_size, heads) for _ in range(blocks)]
        self.output = GPTOutput(embedding_size, vocab_size)

        super().__init__([self.embedding] + self.transformers + [self.output])

    def forward(self, x: Tensor):
        x = self.embedding(x)
        for layer in self.transformers:
            x = layer(x)
        return self.output(x)

GPT trainer¶

class GPTTrainer(Trainer):

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

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

                feature, label = dataset[i]

                self.optimizer.zero_grad()
                prediction = self.model(feature)
                loss = self.loss_fn(prediction, label)
                loss.backward()
                self.optimizer.clip_grad_norm()
                self.optimizer.step()
                steps += 1

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

        predictions = []
        total_loss = 0.0
        with Tensor.no_grad():
            for i in range(len(dataset)):
                feature, label = dataset[i]
                prediction = self.model(feature)
                loss = self.loss_fn(prediction, label)
                predictions.append(prediction)
                total_loss += float(loss.data)

        loss = total_loss / len(dataset)
        dataset.train()
        self.model.train()
        return predictions, loss

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

        with Tensor.no_grad():
            for _ in range(steps):
                feature = Tensor([tokens[-dataset.context_size:]])
                logits = self.model(feature)

                last_logits = logits.data[0, -1]
                exp = np.exp(last_logits - np.max(last_logits))
                probs = exp / np.sum(exp)
                token = np.random.choice(len(probs), p=probs)
                tokens.append(token)

        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 = 64

Number of heads¶

HEADS = 2

Number of Transformer blocks¶

BLOCKS = 2

Modeling¶

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = GPTTrainer(model, loss_fn, optimizer)

Training¶

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}\nloss:\t{loss}')

prediction:	1742 steps, each (4, 32, 65)
loss:	1.8636507043608568

Inference¶

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

ROMEO:
O:
O deseep would my lethis this they-anio, you hale should, gentis,
I'll I
Boy, this I have of del mogriess.

Secondself faic, not lifect were Clains; i' be So thour man I complo, and do hand
Mistrue such to which a city his old vice frieunts not man,
An

In this chapter, we added dropout to the GPT model: during training, it randomly drops the outputs of a part of the neurons to prevent overfitting.

With this, the three major engineering optimization techniques for pre-training are all in place:

residual connection: lets the gradient flow smoothly and prevents vanishing gradients;
layer normalization: keeps the values stable, and training stable;
dropout: prevents overfitting.

With these engineering optimization techniques, the model can be stacked higher, trained more stably, and generalize better.

Exercises¶

Try changing the dropout probability $p$ from 0.1 to 0.3 or even 0.5 and observe how the training results change. Is a larger probability always better?