Weighted attention - From Neuron to GPT

In the previous chapter, we closed the loophole of peeking at the answers with the causal mask. But the current causal attention still has a hidden problem: a token wears many hats, and its different roles interfere with each other in the weights.

Let us review the computation of self-attention:

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

(1)

The same input $X$ plays three roles at the same time:

the asker: which token am I? What information do I need?
the answerer: which token am I? What information can I offer?
the information itself: which token am I? What content do I actually carry?

The same token has to ask, answer, and be the information itself. Playing several roles at once, it is hard to do all of them well.

Query, key, value¶

The solution is to separate the roles: use 3 different sets of weights to generate 3 different vectors for each token. Imagine giving the same token 3 different masks, each worn when playing a different role:

Query: plays the asker;
Key: plays the answerer;
Value: plays the information itself.

These three vectors are obtained by applying a linear transformation to the input $X$ separately:

Q = X \cdot W_Q, \quad K = X \cdot W_K, \quad V = X \cdot W_V

(2)

Here $W_Q$ , $W_K$ and $W_V$ are three trainable weight matrices, which model training can adjust toward whatever suits each role best. With the roles separated, the formula of self-attention becomes:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q \cdot K^T}{\sqrt{d}}\right) V

(3)

Weighted attention¶

Attention with the linear projections of $Q$ , $K$ and $V$ is called weighted 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 introduce the three linear projections $Q$ , $K$ and $V$ into the GPT attention layer. The forward pass steps are extended to:

generate $Q$ , $K$ and $V$ : obtain the query, key and value through three linear layers;
compute the attention score matrix: $\text{scores} = Q \cdot K^T$ ;
scale: multiply by $1 / \sqrt{d}$ to prevent the values from getting too large, where $d$ is the embedding size embedding_size;
causal mask: block the future positions with the lower triangular matrix layer;
compute the attention weights: use the Softmax activation function to turn each row into a probability distribution;
weighted sum: multiply the attention weights by $V$ to get a new attention vector that fuses the context.

class GPTAttention(Composite):

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

        # Query layer
        self.query = Linear(embedding_size, embedding_size)
        # Key layer
        self.key = Linear(embedding_size, embedding_size)
        # Value layer
        self.value = Linear(embedding_size, embedding_size)
        # Causal mask layer
        self.mask = Tril()
        # Activation function layer
        self.softmax = Softmax()
        # Output layer
        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):
        # Query
        query = self.query(x)
        # Key
        key = self.key(x)
        # Value
        value = self.value(x)
        # Scaling factor
        scale = Tensor(np.array(1.0 / np.sqrt(self.embedding_size)))
        # Attention score matrix (query, key)
        scores = query @ key.transpose((0, 2, 1)) * scale
        # Causal attention weights
        weights = self.softmax(self.mask(scores))
        # Scaled dot-product attention (value)
        return self.output(weights @ value)

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¶

class GPTTransformer(Composite):

    def __init__(self, embedding_size):
        self.attention = GPTAttention(embedding_size)
        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)
layer = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(layer.parameters, lr=LEARNING_RATE)
trainer = GPTTrainer(layer, 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.9071517813219405

Inference¶

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

ROMEO:
Where as his and andrieus
That end with love will so so be dever onot, as is bideiffutend, go stay.

TRANIO:
That he the lady! elty I'll will thee hom.

ANGELO:
Her davent.

MO' him busio's tis! seest.

PETHA:
Deel he, do mastion; poitrtys love:
Ay, my ye

In this chapter, we added a role separation mechanism to the self-attention mechanism: with the three linear projections $Q$ , $K$ and $V$ , a token has a distinct role when asking, answering and passing on information.

The current self-attention still has a limitation: all tokens use the same single set of $Q$ , $K$ and $V$ . It is like a person who always looks at a problem from one perspective, which is bound to be one-sided. Take the word “apple”: it means something completely different in “eat an apple” and in “Apple phone”, yet the current self-attention can only understand it in one way.

In the next chapter, we will introduce the multi-head mechanism, so that the model learns to understand the context from different angles.

Exercises¶

Try removing one of the three linear projections $Q$ , $K$ and $V$ and replacing it with $X$ directly. Observe how the training results change. Which projection has the biggest impact?