Normalization - From Neuron to GPT

In the previous chapter, we added residual connections to the two sublayers of the Transformer block, which gave the gradient a fast channel. But the residual connection also brings a side effect: the output of each layer is “input + transformation”, so as the number of layers grows, the values keep accumulating and getting larger, and training becomes unstable.

Layer normalization¶

The solution is layer normalization: pull the output of each layer back to a standard distribution with mean 0 and variance 1.

Compute the mean and variance:

\mu = \frac{1}{d} \sum_{i=1}^{d} x_i, \quad \sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2

(1)

Standardize:

\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}

(2)

Scale and shift:

\text{output} = \gamma \odot \hat{x} + \beta

(3)

Here:

$d$ : the embedding size;
$\epsilon$ : a very small number (for example 10^-5) that keeps the divisor from being 0;
$\gamma$ (weight): the learnable scaling factor (weight);
$\beta$ (bias): the learnable shifting factor (bias).

No matter how large the values are, layer normalization can pull them back to the standard range and prevent them from exploding. It also keeps the distribution of the values stable, makes gradient changes smoother, and speeds up convergence.

The learnable scaling and shifting factors let the magnitude of the normalization adjust itself automatically, so that the one-size-fits-all normalization formula fits the dataset better.

GPT basic modules with normalization: the inputs of the attention layer, the feed-forward network layer and the output layer all pass through a normalization layer first

Figure: GPT basic modules with normalization

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¶

We add a normalization layer (the LayerNorm layer):

Forward pass: compute the mean and variance over the last dimension (the feature dimension). After standardizing, scale and shift;
Backward pass: the gradients of the input $x$ , the scaling factor $\gamma$ and the shifting factor $\beta$ need to be computed separately.

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
        mean = np.mean(x.data, axis=-1, keepdims=True)
        # Variance
        var = np.var(x.data, axis=-1, keepdims=True, ddof=0)
        # Standardize
        norm = (x.data - mean) / np.sqrt(var + self.eps)
        # Create the output value tensor
        p = Tensor(self.weight.data * norm + self.bias.data)

        # Gradient (closure) function of the output value tensor
        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)

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

    # List of trainable model parameters
    @property
    def parameters(self):
        return [self.weight, self.bias]

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¶

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)

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

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

GPT attention layer¶

In the GPT attention layer, we first normalize the input, then compute attention, and finally add the residual connection.

class GPTAttention(Composite):

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

        # Normalization layer
        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)

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

    def forward(self, x: Tensor):
        # Normalize
        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))
        return x + self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape))

GPT feed-forward network layer¶

Likewise, in the GPT feed-forward network layer we first normalize the input, then apply the feed-forward network, and finally add the residual connection.

class GPTFeedForward(Composite):

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

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

    def forward(self, x: Tensor):
        # Normalize
        norm = self.normalize(x)
        h = self.gelu(self.input(norm))
        return x + 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¶

The output layer likewise first normalizes its input, to make sure that the data entering the linear layer has a stable distribution.

class GPTOutput(Composite):

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

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

    def forward(self, x: Tensor):
        # Normalize
        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.7808665809757263

Inference¶

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

ROMEO:
Well.

SABELLA:
Therefore thee go by the tranted is is fastely, and is but more him, where's for you wife?

HORFOMIO:
I am as oncove whomostion on holve all me there and when yalk?

DUKE VINCENTIO:
You speech that be, being suppy,
That try love.

SAVINCO:

In this chapter, we added layer normalization to the modules of the GPT model, to make sure that after the residual connection the data entering the next layer is back in the standard range, which keeps training stable.

With this, we have basically solved the problem of how to stack a GPT model deeper while keeping training stable, clearing the obstacles to training large-scale GPT models. In the next chapter, we will look at a problem that large-scale training easily runs into: overfitting. The model performs well on the training set but may perform poorly on the test set. It “memorizes” too many details of the training data instead of truly “learning” the patterns of language.

Exercises¶

Root mean square normalization (RMSNorm) is an effective simplified version of layer normalization. RMSNorm directly assumes that the mean is 0 and computes only the root mean square. It also keeps only the learnable scaling factor $\gamma$ and drops the shifting factor $\beta$ . Try replacing the normalization layer with RMSNorm and see how the training results change.