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.

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.

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.

MLP

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