Implementation

Architecture

Now we can assemble the modules from the previous chapter into the full transformer.

forward computes the output by calling each module in order Embedding -> Attention -> MLP -> Linear (output layer).
backward computes the gradient values by calling the same modules in the reverse order.
predict runs the forward pass and computes the probability distribution over the vocabulary for the last token.
train runs the full training loop calculating the loss first and then running the backward pass and updating the parameters.
Attention and the MLP each use a residual connection. They add the output back to their input ( $x + ...$ , $r + ...$ ). It gives the gradient a direct path back through the network, which keeps training stable.
Our training implementation runs one sequence at a time. Real LLMs process a batch of sequences at once to average the gradient over many examples so that the learning curve is smoother. It also parallelises better, but that doesn't matter for our small example.

Code

import numpy as np
from modules_implementation import Linear, Embedding, Attention, MLP
from helpers_implementation import softmax, cross_entropy, training_examples, Vocabulary

class Transformer:
    def __init__(self, vocab_size, d, hidden, block, seed=0):
        rng = np.random.default_rng(seed)
        self.block = block
        self.embed = Embedding(vocab_size, d, block, rng)
        self.attn = Attention(d, block, rng)
        self.mlp = MLP(d, hidden, rng)
        self.output = Linear(d, vocab_size, rng)
        self.mods = [self.embed, self.attn, self.mlp, self.output]

    def forward(self, X):
        x = self.embed.forward(X)
        r = x + self.attn.forward(x)
        r2 = r + self.mlp.forward(r)
        return self.output.forward(r2)

    def loss(self, X, Y):
        logits = self.forward(X)
        self.probs = softmax(logits)
        self.Y = Y
        return cross_entropy(self.probs, Y)

    def backward(self):
        dlogits = self.probs.copy()
        dlogits[np.arange(len(self.Y)), self.Y] -= 1
        dlogits /= len(self.Y)
        dr2 = self.output.backward(dlogits)
        dr = dr2 + self.mlp.backward(dr2)
        dx = dr + self.attn.backward(dr)
        self.embed.backward(dx)

    def params(self):
        return [pg for m in self.mods for pg in m.params()]

    def predict(self, context_ids):
        ids = list(context_ids)[-self.block:] or [0]
        logits = self.forward(np.array(ids))
        return softmax(logits[-1])

    def train(self, X, Y, steps, lr):
        n = X.shape[0]
        rng = np.random.default_rng(1)
        for _ in range(steps):
            i = rng.integers(0, n)
            for param, grad in self.params():
                grad[:] = 0
            self.loss(X[i], Y[i])
            self.backward()
            for param, grad in self.params():
                np.clip(grad, -5, 5, out=grad)
                param -= lr * grad
        return self

Our model has a single transformer block. If we want to scale it we will need to stack several blocks ( $N$ ).

The code is ready to be tested now. Let's train the model on our small corpus and run the generator predicting the next tokens for a given sequence.

Note that the corpus is very small in our example. The model will just memorize the exact order of the text.

The full code is available separately, so you can try it on a bigger corpus and at a larger scale.

Code

import numpy as np
from implementation_overview import d, hidden, block, vocab_size, steps, lr, corpus

def generate(model, vocab, prompt, n=40):
    ids = vocab.encode(prompt)
    for _ in range(n):
        probs = model.predict(ids)
        ids.append(int(probs.argmax()))
    return vocab.decode(ids)

vocab = Vocabulary(corpus, vocab_size)
ids = vocab.encode(corpus)

model = Transformer(len(vocab.itos), d=d, hidden=hidden, block=block)
X, Y = training_examples(ids, model.block)
model.train(X, Y, steps=steps, lr=lr)

print("Training")
print()
print(generate(model, vocab, "The little girl went to the", n=5))