Residual connection - From Neuron to GPT

Over 4 chapters, we completed the most central part of the GPT model: the self-attention mechanism.

Next, we will focus on the engineering implementation of GPT pre-training.

A distinctive feature of GPT models is that “scale works wonders” (a Chinese saying, literally “great force produces miracles”). It shows not only in the massive data used for training, but also in the huge parameter count and in deep network structures of dozens or even hundreds of layers. Specifically, the GPT-3 model has 96 Transformer blocks, and deeper models have more layers. The more Transformer blocks are stacked, the harder training gets. During backpropagation the gradient has to pass through dozens or hundreds of layers, and each layer “dilutes” it a little more, until it almost vanishes in the last few layers. This is the so-called vanishing gradient problem.

Residual connection¶

In 2015, Kaiming He and colleagues at Microsoft Research proposed the concept of the residual connection, which has greatly alleviated the challenges associated with training deep neural networks. This paper (ResNet) is one of the most cited papers in the history of deep learning.

The core idea of the residual connection is extremely simple: let the output of each layer = input + transformation.

\text{output} = x + f(x)

(1)

instead of the traditional $\text{output} = f(x)$ .

Why does the residual connection work¶

Forward pass: even if $f(x)$ is learned poorly, the output still contains $x$ at least, so no information is lost;
Backward pass: the gradient can flow back through $x$ and does not depend entirely on the gradient of $f(x)$ , so the vanishing gradient problem is greatly eased.

A residual connection is like building an extra highway for the gradient: however winding the $f(x)$ in between may be, the gradient can always flow back to the earlier layers unobstructed.

Where to add residual connections¶

In the GPT model, the depth of the network consists mainly of Transformer blocks. A Transformer block contains two sublayers: self-attention and the feed-forward network. Each sublayer needs a residual connection:

self-attention: $x = x + \text{attention}(x)$ ;
feed-forward network: $x = x + \text{feedforward}(x)$ .

In this way, no matter how many Transformer blocks are stacked, the gradient can reach the first layer smoothly, avoiding vanishing gradients.

GPT basic modules with residual connections: the attention layer and the feed-forward network layer each have a residual connection that bypasses the layer

Figure: GPT basic modules with residual connections

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})

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¶

We add a residual connection at the end of the GPT attention layer, changing the output from self.output(...) to x + self.output(...).

class GPTAttention(Composite):

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

        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.query,
                          self.key,
                          self.value,
                          self.mask,
                          self.softmax,
                          self.output])

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

GPT feed-forward network layer¶

Likewise, we add a residual connection at the end of the GPT feed-forward network layer, changing the output from self.output(...) to x + self.output(...).

class GPTFeedForward(Composite):

    def __init__(self, embedding_size):
        self.input = Linear(embedding_size, embedding_size * 4)
        self.gelu = GELU()
        self.output = Linear(embedding_size * 4, embedding_size)

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

    def forward(self, x: Tensor):
        h = self.gelu(self.input(x))
        # Residual connection
        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¶

class GPTOutput(Composite):

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

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

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

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.8085668785700706

Inference¶

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

ROMEO:
Well, I'll poisen, swearn thee, haprite,
Will me is is fatherior offenretcher rangershed is in tet Catie, there name the king,
Lets Comiou, who no: well I delieve what bow,
And
Berancionr her; and love stime and in him esceasured,
That try love.

SIY:
HtU

In this chapter, we added residual connections to both sublayers of the Transformer block. The gradient now has a direct channel, and one of the training difficulties of deep neural networks, the vanishing gradient, is solved with ease.

But the residual connection also has a side effect: the output of every layer is “input + transformation”, so as the number of layers grows, the values keep accumulating and getting larger, which makes training unstable. In the next chapter, we will set out to solve this problem.

Exercises¶

Both sublayers of the Transformer block now have residual connections. Try removing one of them and see how the training results change.