Causal attention - From Neuron to GPT

In the previous chapter, we implemented the self-attention mechanism, and tokens can finally see each other. But we also found a fatal loophole: every token can see all positions, including the tokens after it.

What does this mean?

Suppose the model is processing the 3rd token and is about to predict the 4th. But now it can see not only the 1st and 2nd tokens, but also the 4th, 5th and 6th, that is, the answer itself. You can imagine that during training the model cheats with ease and the loss looks great; but at inference the answer does not exist yet, and the model gives itself away at once.

It is like a student who peeks at the standard answers while doing homework and gets good grades. But in the exam room there are no answers to copy, and the truth comes out.

Causal mask¶

The solution is actually straightforward: cover up the future information.

Specifically, after computing the attention score matrix, replace the values at all future positions of the matrix with a very large negative number (for example -10⁹). After Softmax, the attention weights at these positions approach 0, which means they are completely ignored and do not take part in the prediction.

This treatment is called the causal mask: each position can see only itself and the positions before it, not the positions after it. This matches exactly the causal nature of language generation: the next word can only be determined by the words before it.

Causal attention weights: only the diagonal and the part below it of the weight matrix have values, and everything above the diagonal is 0

Figure: causal attention weights

Lower triangular matrix¶

How do we implement the causal mask?

Look at the shape of the attention score matrix: (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:

when $j \leq i$ : the $j$ -th token is before the $i$ -th token, or is the token itself, and should be kept;
when $j > i$ : the $j$ -th token is after the $i$ -th token and should be masked.

That is, we only need to keep the diagonal and the lower triangular part below it, and replace the whole upper triangular part above the diagonal with -10⁹. NumPy provides a ready-made lower-triangle function np.tril, which produces a matrix with 1 in the lower triangle and 0 in the upper triangle, and we can use it directly as the mask.

Causal attention¶

Self-attention with a causal mask is called causal attention, or masked self-attention. It is the standard configuration of decoder-only architectures such as GPT.

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¶

We design the causal mask as a separate lower triangular matrix layer (Tril layer):

Forward pass: generate the lower triangular mask with np.tril and replace the values at the upper triangular positions with -10⁹;
Backward pass: the gradient flows back only through the lower triangular positions. The upper triangular positions did not take part in the forward pass, so their gradient is 0.

class Tril(Layer):

    def __init__(self, value=-1e9):
        super().__init__()
        self.value = value

    def forward(self, x: Tensor):
        # Causal mask
        keep = np.tril(np.ones(x.shape[-2:]))
        # Create the output value tensor
        p = Tensor(np.where(keep, x.data, self.value))

        # Gradient (closure) function of the output value tensor
        def gradient_fn():
            x.grad += p.grad * keep

        # Set the gradient (closure) function and the parent node list of the output value tensor
        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¶

In the GPT attention layer, we use the lower triangular mask to block the future positions of the attention score matrix scores.

weights = self.softmax(self.mask(scores))

class GPTAttention(Composite):

    def __init__(self, embedding_size):
        self.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.mask,
                          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
        # Causal attention weights
        weights = self.softmax(self.mask(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¶

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)
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.0406119945949857

Inference¶

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

ROMEO:
Famtany at!--will and, though this him at.

ISTARIO:
Hen ifthat not mid'dses mepion.

LERMIONA:
Sordecer, wor I tontlouce alsucitend waseres, It what you sor ande bust trune do's to whave with Me
In keal dies the is my antas;
I fill wak selse that of this

In this chapter, we closed the loophole of self-attention peeking at the answers with the causal mask. With this, the core of the self-attention mechanism is complete. From the inference result, we can already see obvious sentence patterns and dramatic structure.

Exercises¶

Think about why the lower triangular matrix layer fills the future positions in the upper triangle with -10⁹. Try changing it to 0 and observe how the training results change.