Implementation

Modules

In this chapter we will implement small modules that will be the building blocks of our final model's architecture.

Each module will have three main functions:

forward will be used in the forward pass. It takes the input from the previous module and produces the output for the next.
backward will be used in the backward pass. It takes the gradient of the loss with respect to this module's output, and returns the gradient with respect to its input.
params provides the parameter variables so the training loop can update them.

This will help us in the next chapter to build the whole architecture by stacking these modules, which share the same interface.

Embedding

Embedding stores the positional encodings and the embedding table. The values for the given tokens are summed up in the forward pass ( $x = e + p$ ).

The backward pass accumulates gradients into dtok/dpos and does not return anything as it is the first module in the architecture.

The same token can appear several times in one sequence, so np.add.at sums its gradient across those positions.

Code

import numpy as np

class Embedding:
    def __init__(self, vocab_size, d, block, rng):
        self.tok = rng.normal(0, 0.05, (vocab_size, d))
        self.dtok = np.zeros_like(self.tok)
        self.pos = rng.normal(0, 0.05, (block, d))
        self.dpos = np.zeros_like(self.pos)

    def forward(self, X):
        self.X = X
        T = len(X)
        return self.tok[X] + self.pos[:T]

    def backward(self, dout):
        np.add.at(self.dtok, self.X, dout)
        self.dpos[:len(self.X)] += dout

    def params(self):
        return [(self.tok, self.dtok), (self.pos, self.dpos)]

Linear layer

Linear computes the function $z = w_0 + w_1 \cdot x_1 + ... + w_n \cdot x_n$ . Note that $w_0$ parameters are defined separately as $b$ (bias) as it has a different shape and is not multiplied by the input.

Code

import numpy as np

class Linear:
    def __init__(self, nin, nout, rng, bias=True):
        s = 1 / np.sqrt(nin)
        self.W = rng.normal(0, s, (nin, nout))
        self.dW = np.zeros_like(self.W)
        self.b = np.zeros(nout) if bias else None
        self.db = np.zeros(nout) if bias else None

    def forward(self, x):
        self.x = x
        y = x @ self.W
        return y + self.b if self.b is not None else y

    def backward(self, dy):
        self.dW += self.x.T @ dy
        if self.b is not None:
            self.db += dy.sum(0)
        return dy @ self.W.T

    def params(self):
        return [(self.W, self.dW)] + ([(self.b, self.db)] if self.b is not None else [])

Attention

Attention computes contextual embeddings with the attention mechanism formulas.

The mask matrix ensures each token attends only to previous tokens in the sequence, by pushing the masked scores to a large negative value before the softmax.

Code

import numpy as np
from helpers_implementation import softmax

class Attention:
    def __init__(self, d, block, rng):
        self.Wq = Linear(d, d, rng, bias=False)
        self.Wk = Linear(d, d, rng, bias=False)
        self.Wv = Linear(d, d, rng, bias=False)
        self.Wo = Linear(d, d, rng, bias=False)
        self.mask = np.triu(np.full((block, block), -1e9), 1)
        self.d = d

    def forward(self, x):
        T = x.shape[0]
        Q = self.Wq.forward(x)
        K = self.Wk.forward(x)
        V = self.Wv.forward(x)
        self.Q, self.K, self.V = Q, K, V
        scores = Q @ K.T / np.sqrt(self.d)
        scores = scores + self.mask[:T, :T]
        self.attn = softmax(scores)
        return self.Wo.forward(self.attn @ V)

    def backward(self, dout):
        dctx = self.Wo.backward(dout)
        dattn = dctx @ self.V.T
        dV = self.attn.T @ dctx
        dscores = self.attn * (dattn - (dattn * self.attn).sum(-1, keepdims=True))
        dscores = dscores / np.sqrt(self.d)
        dQ = dscores @ self.K
        dK = dscores.T @ self.Q
        return self.Wq.backward(dQ) + self.Wk.backward(dK) + self.Wv.backward(dV)

    def params(self):
        return self.Wq.params() + self.Wk.params() + self.Wv.params() + self.Wo.params()

Note that this implementation is single-head attention only. If we want to scale, we need to extend it to handle more heads in the attention block, each with its own separate parameters.

MLP

MLP is a two-layer network. For the first layer we apply the sigmoid activation function.

Code

import numpy as np
from helpers_implementation import sigmoid

class MLP:
    def __init__(self, d, hidden, rng):
        self.layer1 = Linear(d, hidden, rng)
        self.layer2 = Linear(hidden, d, rng)

    def forward(self, x):
        self.h = sigmoid(self.layer1.forward(x))
        return self.layer2.forward(self.h)

    def backward(self, dout):
        dh = self.layer2.backward(dout)
        return self.layer1.backward(dh * self.h * (1 - self.h))

    def params(self):
        return self.layer1.params() + self.layer2.params()

As we discussed previously the sigmoid function for the hidden layers might not be the best choice if we want to scale the model. But we keep it as it works well for our small example.