Self-attention - From Neuron to GPT

In the previous chapter, position embedding let the GPT model know where each token is. Next, we need to make the model interact with the neighboring tokens on this basis.

An RNN model passes the information of earlier tokens to the current token through the hidden state. A GPT model, in contrast, processes all tokens in parallel, and no token knows anything about the others. Just as we told the model the position of a token with the position vector, we now need to tell the model explicitly how the current token is related to all the other tokens.

For example, how do humans read and understand? Take this sentence:

Ming bought apples, not bananas, because they are fresh.

When we reach “they”, we look back at what came before and put more of our attention on “apples” rather than “bananas”. This ability to look back and weigh the context according to the current content is what we want to simulate. It is called the attention mechanism.

Attention mechanism¶

The core idea of the attention mechanism is: when processing each token, first look back at all the tokens in the context, and assign them different weights according to how relevant they are:

tokens with high relevance get large weights and contribute more;
tokens with low relevance get small weights and contribute less.

Self-attention¶

Self-attention is a kind of attention mechanism: a token attends not to another sequence, but to the very sequence it belongs to. Every token plays two roles at the same time:

the attender: how do the other tokens affect me?
the attended: how do I affect the other tokens?

Attention weights¶

So how do we measure the relevance of two tokens? A simple and effective way is to compute the dot product of their word vectors:

\text{score}_{ij} = x_i \cdot x_j

(1)

The geometric meaning of the dot product is the agreement in direction of two vectors: the closer the directions, the larger the dot product, and the more relevant the two tokens are.

Computing the dot product for every pair of tokens gives an attention score matrix (attention scores) with the shape (context size, context size). The value in row $i$ and column $j$ is the relevance of the $i$ -th token to the $j$ -th token.

The values in the attention score matrix are not bounded, and using them directly would be unstable. We apply the Softmax activation function to turn each row into a probability distribution, which gives the attention weights: all the weights in each row sum to 1.

Self-attention weights: the weights of each token over all tokens in the context; each row passes through Softmax and sums to 1

Figure: self-attention weights

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)

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 GPT attention layer to give the embedding vector of each token (word vector + position vector) the attention from the other tokens:

compute the attention score matrix: $\text{scores} = X \cdot X^T$ ;
scale: multiply by $1 / \sqrt{d}$ to prevent the values from getting too large, where $d$ is the embedding size embedding_size;
compute the attention weights: use the Softmax activation function to turn each row into a probability distribution;
weighted sum: multiply the attention weights by the embedding vectors to get a new attention vector that fuses the context.

Finally, we get the scaled dot-product attention:

\text{Attention}(X) = \text{softmax}\left(\frac{X \cdot X^T}{\sqrt{d}}\right) X

(2)

This is the core formula of the Transformer paper.

The code has one more step beyond the formula: the result of the weighted sum then passes through a linear layer self.output (the output projection) before it becomes the final output of the attention layer.

class GPTAttention(Composite):

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

        # Activation function layer
        self.softmax = Softmax()
        # Output layer
        self.output = Linear(embedding_size, embedding_size)

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

    def forward(self, x: Tensor):
        # Scaling factor
        scale = Tensor(np.array(1.0 / np.sqrt(self.embedding_size)))
        # Attention score matrix
        scores = x @ x.transpose((0, 2, 1)) * scale
        # Attention weights
        weights = self.softmax(scores)
        # Scaled dot-product attention
        return self.output(weights @ x)

GPT feed-forward network layer¶

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))
        return self.output(h)

GPT Transformer block¶

Now we have the complete Transformer block: attention layer + feed-forward network layer.

Attention layer: lets tokens communicate with each other and fuse context information;
Feed-forward network layer: further refines the features of the fused attention vectors.

GPT basic modules with self-attention: the Transformer block consists of an attention layer and a feed-forward network layer

Figure: GPT basic modules with self-attention

class GPTTransformer(Composite):

    def __init__(self, embedding_size):
        # Attention layer
        self.attention = GPTAttention(embedding_size)
        # Feed-forward network layer
        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, blocks):
        self.embedding = GPTEmbedding(vocab_size, context_size, embedding_size)
        self.transformers = [GPTTransformer(embedding_size) 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 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, 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:	2.872809682593143

Inference¶

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

ROMEO:::::M::MM:MMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMM

In this chapter, we preliminarily implemented the most central mechanism of the GPT model: self-attention. Tokens can finally see each other, and every position fuses the relevant information of the whole context. The Transformer block has also been upgraded from a plain feed-forward network to the complete structure of “attention + feed-forward network”.

But the inference result is surprising: it fails completely. This is because the current self-attention has a fatal loophole: every token can see all positions, including the tokens after it. It is as if the model peeks at the answers during training, and gives itself away at inference.

In the next chapter, we will explore how to close this loophole.

Exercises¶

What does the output projection of the GPT attention layer (the self.output linear layer) do? Try removing it and observe how the training results change.