Multi-head attention - From Neuron to GPT

In the previous chapter, we added the $Q$ , $K$ , $V$ role separation to self-attention. But one limitation remains: all tokens share the same single set of $Q$ , $K$ and $V$ .

It is like a person who can look at a problem from only one perspective, and cannot understand that the same word can mean different things in different contexts. Take the word “apple”:

in “eat an apple”, it is a fruit;
in “Apple phone”, it is a brand.

The current model can understand it in only one way, and cannot capture these two different meanings at the same time.

Multi-head attention¶

The solution is multi-head attention: split the embedding vector of a token into several parts, called heads; generate a separate set of $Q$ , $K$ , $V$ linear projections for each part; let each develop its own independent attention; and finally concatenate them together:

each head has its own $Q$ , $K$ , $V$ linear projections and attends to a different feature subspace;
each head develops independently and picks up a different logic in the data;
the heads work in parallel and are combined at the end, so the model can understand the context from several angles.

Split and merge¶

Suppose the embedding size is 64 and the number of heads is 2. Then:

the embedding size of each head is 64 / 2 = 32;
the input shape changes from (batch size, context size, 64) to (batch size, context size, 2, 32);
transpose it to (batch size, 2, context size, 32), so that each head can compute its attention independently;
after the computation, transpose back to (batch size, context size, 2, 32), then concatenate into (batch size, context size, 64).

Scaling factor¶

The embedding size of each head is EMBEDDING_SIZE / HEADS, so the $d$ in the scaling factor $1 / \sqrt{d}$ is now EMBEDDING_SIZE / HEADS.

Number of parameters¶

The number of parameters of multi-head attention is exactly the same as that of single-head attention: the weight matrices of the three linear projections $Q$ , $K$ and $V$ are still (embedding size, embedding size). The only difference is that the computation splits the vector into several parts and computes each independently.

💡 Tip: computing multi-head attention is not a stack of several attentions, but a parallel computation after expanding the dimensions inside the same attention.

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 multi-head support to the GPT attention layer.

The forward pass gains two steps, split and merge:

generate $Q$ , $K$ and $V$ : obtain the query, key and value through three linear layers;
split the heads: divide the embedding vector into several parts;
compute attention: each head independently computes the attention score matrix, scales, masks, computes the weights, and takes the weighted sum;
merge the heads: concatenate and combine the multi-head attention.

How the shape of the attention changes during split and merge:

\text{split}: (\text{batch}, \text{context}, \text{embedding}) \xrightarrow{\text{split}} (\text{batch}, \text{context}, \text{heads}, \text{per-head embedding}) \xrightarrow{\text{transpose}} (\text{batch}, \text{heads}, \text{context}, \text{per-head embedding})

(1)

\text{merge}: (\text{batch}, \text{heads}, \text{context}, \text{per-head embedding}) \xrightarrow{\text{transpose}} (\text{batch}, \text{context}, \text{heads}, \text{per-head embedding}) \xrightarrow{\text{merge}} (\text{batch}, \text{context}, \text{embedding})

(2)

Therefore, the multi-head mechanism is handled only inside the GPT attention layer. The shape of the output does not change, and the later steps need no modification.

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):
        # Embedding size of each head
        head_embedding_size = self.embedding_size // self.heads
        # Multi-head shape
        multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
        # Shape after merging
        shape = (-1, x.shape[1], self.embedding_size)

        # Query (multi-head)
        query = self.query(x).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        # Key (multi-head), transposed in advance
        key = self.key(x).reshape(multi_head_shape).transpose((0, 2, 3, 1))
        # Value (multi-head)
        value = self.value(x).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        # Scaling factor
        scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
        # Attention score matrix (query, key)
        scores = query @ key * scale
        # Causal attention weights
        weights = self.softmax(self.mask(scores))
        # Scaled dot-product attention (value, after merging)
        return self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape))

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¶

heads needs to be passed in when creating it.

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¶

heads needs to be passed in when creating it.

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¶

This is the hyperparameter that defines the number of heads in multi-head attention.

HEADS = 2

Number of Transformer blocks¶

BLOCKS = 2

Modeling¶

When creating the model, we need to pass HEADS in as an additional model parameter.

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

Inference¶

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

ROMEO:
Where as her, mise. what I shall bestill tranter is it fate,
Why onot untock: papon you have for you wife?

HORTA:
I what a the know
You will the'll cof Pefier there and which:
Chold we man use to love, thing the mall geread; thou reavedss lovela, by off

In this chapter, we added the multi-head mechanism to attention: split the embedding vector into several parts, give each part its own $Q$ , $K$ , $V$ linear projections, compute attention independently, and finally concatenate and combine the results. The model can therefore understand the context from several angles.

With this, we have implemented the complete self-attention mechanism and built, step by step, the core component of the GPT model:

self-attention: tokens see each other;
causal attention: mask the future to prevent peeking at the answers;
weighted attention: separate the roles of $Q$ , $K$ and $V$ ;
multi-head attention: understand the context in parallel from several angles.

Exercises¶

Try changing HEADS from 2 to 4 or 8 and observe how the training speed and results change. Is more heads always better?