Building a Transformer from Scratch

import math
import numpy as np
import pylab as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

Building a Transformer from Scratch#

Extra Credit

This page is extra credit for the machine learning unit. It builds directly on Masked Diffusion, so finish that page first.

On the masked diffusion page we handed you a neural network called TinyDenoiser and told you not to worry about what was inside of it. All you needed to know was that it took in a sentence with some masks in it and returned a guess for every position.

TinyDenoiser is a transformer. So is the ModernBERT network that you chatted with, and so are the networks behind chatGPT and Claude. Our goal on this page is to open the box. You are going to build a transformer yourself, one piece at a time, and at each step we will check what the new piece buys you. By the end you will

  • have your own transformer which you can drop into the training loop that you wrote on the masked diffusion page in place of TinyDenoiser

  • train it on stories and have it write (using your Undiffuse)

  • change one line of it and turn it into a model which writes left to right the way that chatGPT does

We are not going to assume that you know much pytorch. Whenever we need something new from it we will show it to you first.

Getting Set Up#

Packages. You need torch, requests, numpy and matplotlib. On Google Colab all of these are already installed. On your own machine run pip install torch requests numpy matplotlib.

Downloads. The only thing that gets downloaded is the same 4 MB piece of TinyStories that you used before. The code we give you fetches it. (There is no pre-trained network on this page. You are building the network.)

Hardware. Everything up to Stories runs in a few seconds on any laptop. The two training runs on stories at the end take a few minutes each on a GPU (on Google Colab select Runtime > Change runtime type > T4 GPU) and something like ten times longer without one.

Things to bring with you. Copy the following from your masked diffusion notebook into this one:

  • Undiffuse (and UndiffuseGreedy and your version with a temperature, if you’d like to use them)

  • DiffusionLoss

  • your two-sentence SamplePInit and GuessX0 (the ones where the cat is alive has probability 0.8 and the box is empty has probability 0.2)

Here is the rest of what you need from that page.

# GIVEN: the same vocabulary, device and MakeGuesser as on the masked diffusion page
vocab = ["the", "cat", "box", "is", "alive", "empty", "nobody", "knows", "if", "or", "until", "somebody", "looks", "_"]
MASK = 13   # the index of "_"

def Encode(sentence):
    return torch.tensor([vocab.index(word) for word in sentence.split()])

def Decode(x):
    return " ".join(vocab[i] for i in x.tolist())

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

def MakeGuesser(net, forbidden=()):
    """Wrap a network as a guessing function for Undiffuse: x_t of shape (L,) -> probabilities (L, vocab_size).
    `forbidden` is a list of tokens the guess is never allowed to be."""
    def Guess(x_t):
        with torch.no_grad():
            logits = net(x_t[None].to(next(net.parameters()).device))[0].cpu()
        logits[:, list(forbidden)] = -math.inf
        return logits.softmax(dim=-1)
    return Guess

The rules. You are allowed to use the basic building blocks of pytorch: nn.Linear, nn.Embedding, nn.LayerNorm, nn.GELU, the optimizers, and the automatic differentiation (loss.backward()). You are not allowed to use the pieces of pytorch which already are a transformer: nn.MultiheadAttention, nn.TransformerEncoder (and its relatives) and F.scaled_dot_product_attention. (We will use the last one once, to check your answer.)

The Plan#

Let’s remind ourselves what the network has to do. It is handed a batch of sentences, which is a tensor of integers of shape (batch, L). It has to return a tensor of shape (batch, L, vocab_size): for every position of every sentence, one number (a logit) for every token in the vocabulary. The softmax of those logits is the network’s guess for what the token at that position is. Remember that a logit is just minus an energy.

A transformer does this in three stages.

tokens                  (batch, L)               integers
   |
   |  1. embed:   turn every token into a vector of d numbers
   v
vectors                 (batch, L, d)
   |
   |  2. blocks:  the vectors at different positions talk to each other (attention)
   |              and then each one thinks about what it heard (a small ordinary neural network).
   |              Repeat this a few times.
   v
vectors                 (batch, L, d)
   |
   |  3. output:  turn every vector into vocab_size logits
   v
logits                  (batch, L, vocab_size)

The picture to have in your head is a chain of \(L\) sites. On every site there sits a vector of \(d\) numbers which is everything that the network currently “knows” at that position. Stages 1 and 3 work on every site separately. The only place where one site can find out anything about another is the attention in stage 2.

We are going to build this in the order 1, 3, 2. That is, we will first build the network with nothing in the middle (which is very easy) and see exactly what it can’t do. That will tell us what the middle is for.

A Little Pytorch#

Before we start, here is the small amount of pytorch that we are going to need. Run each of these cells and make sure the output is what you expect. There is nothing to hand in for this section.

Shapes. A tensor is an array of numbers and its shape tells you how many indices it has and how big each one is. Nearly every bug you will have on this page is going to be a shape that isn’t what you thought it was, so get used to printing them. Our convention is going to be that the vectors sitting on our sites are stored as a tensor of shape (batch, L, d): which sentence, which position, which component of the vector.

# GIVEN
x = torch.randn(2, 5, 8)        # 2 sentences, 5 positions, a vector of 8 random numbers at each position
print(x.shape)
print(x[0].shape)               # the first sentence: (5, 8)
print(x[0, 3].shape)            # the vector at position 3 of the first sentence: (8,)
print(x[..., 0].shape)          # "..." means "all the indices I didn't mention": the first component everywhere, (2, 5)

A linear layer. nn.Linear(8, 3) is a matrix \(W\) (of shape 3 x 8) and a vector \(b\) (of length 3) which pytorch will later be able to train. When you call it on a vector \(x\) of length 8 it returns \(Wx+b\), which has length 3. The important thing is that if you hand it a tensor with more indices it does this to the last index and leaves all the others alone. So we can apply the same matrix at every position of every sentence with one call.

# GIVEN
layer = nn.Linear(8, 3)
print(layer(x).shape)                                   # (2, 5, 8) -> (2, 5, 3)
print(layer.weight.shape, layer.bias.shape)             # the matrix W and the vector b
no_bias = nn.Linear(8, 3, bias=False)                   # just the matrix
print(torch.allclose(no_bias(x), x @ no_bias.weight.T)) # ...so this is the same as multiplying by it

Matrix multiplication. a @ b multiplies matrices. If a and b have more than two indices it treats the last two as the matrix and does a separate multiplication for each value of the others. a.transpose(-2, -1) swaps the last two indices - i.e. it is the transpose of each of those matrices. (Negative numbers count indices from the end: -1 is the last index and -2 is the one before it.)

# GIVEN
a = torch.randn(2, 5, 8)
b = torch.randn(2, 5, 8)
print(a.transpose(-2, -1).shape)              # (2, 8, 5)
dots = a @ b.transpose(-2, -1)                # (2, 5, 8) @ (2, 8, 5) -> (2, 5, 5)
print(dots.shape)
# dots[s, i, j] is the dot product of the vector at position i of a with the vector at position j of b (in sentence s)
print(dots[0, 1, 3].item(), (a[0, 1] * b[0, 3]).sum().item())

Softmax. torch.softmax(a, dim=-1) turns numbers into probabilities: it exponentiates everything and then divides by the sum over the last index, \(p_j = e^{a_j}/\sum_{j'} e^{a_{j'}}\). In other words, it is a Boltzmann distribution where \(-a_j\) is the energy.

# GIVEN
p = torch.softmax(dots, dim=-1)
print(p.shape)                 # still (2, 5, 5)
print(p.sum(dim=-1))           # every row adds up to one

Making your own network. In pytorch a neural network is a class that has two functions in it.

  • __init__ is run once, when you make the network. This is where you create all of the pieces which have trainable numbers in them (like nn.Linear) and store them as self.something. The first line always has to be super().__init__(), which is what lets pytorch keep track of those pieces.

  • forward is what happens when you call the network on some input.

Here is a network with one linear layer in it. Notice that we never call forward ourselves. Writing net(x) does it for us.

# GIVEN
class Example(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()                      # always the first line
        self.layer = nn.Linear(d_in, d_out)     # create the trainable pieces here...

    def forward(self, x):
        return self.layer(x)                    # ...and use them here

net = Example(8, 3)
print(net(x).shape)                                      # calls net.forward(x)
print(sum(p.numel() for p in net.parameters()))          # how many trainable numbers: 3*8 + 3 = 27

That last line is worth remembering. net.parameters() is a list of every trainable tensor that pytorch found inside your network and it is what you hand to the optimizer. If a piece of your network doesn’t show up in that count, pytorch doesn’t know about it and it isn’t going to get trained.

Step 1: Words to Vectors#

The first thing we have to do is turn a token (an integer) into a vector of \(d\) numbers. There is no clever way to do this. We just keep a table with one row for each token in the vocabulary and look up the row we want. This is called an embedding.

If table is a tensor of shape (len(vocab), d) and x is a tensor of integers, then table[x] looks up all of the rows at once: where x had an integer there is now a vector.

Pytorch’s nn.Embedding(len(vocab), d) is exactly this. It holds a table (it is called .weight) which starts out random and it looks things up when you call it. The only difference from doing it yourself is that pytorch knows that the table is something it is allowed to train. The network will get to decide what vector each word should have.

Grading

  • Make a random table table = torch.randn(len(vocab), 8) and let x = Encode("the cat is the box"). Print the shape of table[x]. Which rows of it are identical? Why?

  • Make an embed = nn.Embedding(len(vocab), 8) and verify (using torch.allclose) that embed(x) is the same thing as embed.weight[x].

  • What is the shape of embed(x) if x is a batch of 64 sentences of 4 tokens each? Check your answer. (torch.stack([SamplePInit() for b in range(64)]) will make you such a batch.)

Step 2: A Network Where Nobody Talks#

We can already make a network which has the right inputs and outputs. Embed each token into a vector (stage 1) and then use a linear layer to turn that vector into len(vocab) logits (stage 3). There is nothing in the middle.

# GIVEN: the simplest network with the right shape. Each position makes its guess all on its own.
class NoTalking(nn.Module):
    def __init__(self, vocab_size, dim=32):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)     # stage 1: a table with one vector for each token
        self.out = nn.Linear(dim, vocab_size)          # stage 3: a vector -> one logit for each token

    def forward(self, x):                              # x:  (batch, L) integers
        h = self.embed(x)                              # h:  (batch, L, dim)
        return self.out(h)                             #     (batch, L, vocab_size)

net = NoTalking(len(vocab))
x_0 = torch.stack([SamplePInit() for b in range(64)])
print(x_0.shape, "->", net(x_0).shape)

Let’s see how well this does on our two sentences. Here is the training loop from the masked diffusion page. Nothing in it has changed except the first line.

net = NoTalking(len(vocab))
opt = torch.optim.Adam(net.parameters(), lr=1e-3)

losses = []
for step in range(0, 1500):
  opt.zero_grad()
  x_0 = torch.stack([SamplePInit() for b in range(64)])   # a batch of 64 sentences
  loss = DiffusionLoss(net, x_0, MASK)
  loss.backward()
  opt.step()
  losses.append(loss.item())

We are going to be training quite a few different networks this way, so it is worth putting this loop into a function def TrainTwoSentences(net, steps=1500) which returns the list of losses.

Recall that the best that any guesser can do on this problem is a loss of 0.50 (the entropy of our distribution).

Grading

Train a NoTalking network.

  • Plot the loss as a function of training step (with a window average over about 100 steps as before). Draw a horizontal line at 0.50. Where does your loss end up?

  • Make it into a guesser with GuessNet = MakeGuesser(net). For each of _ _ _ _, _ _ is alive and the _ _ empty print out the four most probable tokens at position 2 and their probabilities. (p[1].topk(4) gives you the .values and the .indices of the four largest entries of p[1].)

  • Explain what you see. Why is the guess the same for all three? Where do the numbers come from?

You should have found that your network did really badly. Its loss gets stuck around 6.5 and at a masked position it always makes the same guess: the or is about 25% of the time, cat or alive about 20% of the time and box or empty about 5% of the time.

Think about what this network gets to see when it is making its guess at a masked position. The vector on that site is the embedding of the mask token. It is the same vector whether the mask is at position 1 or at position 4, and it is the same vector no matter what the rest of the sentence is. So the network has to give the same answer every time, and the best single answer is how often each word shows up anywhere in a sentence.

There are two things this network is missing:

  1. It doesn’t know where it is. If it knew that it was at position 1 it could say the every time.

  2. It doesn’t know what the other words are. If it knew that position 4 was alive it could say cat.

Let’s fix them one at a time.

Where does 6.5 come from?

The best this network can do at a masked position is to guess the frequencies \(p(v)\) of the words. Each time it does, it pays (on average) the entropy of that distribution

\[-\sum_v p(v) \log p(v) = -2(0.25 \log 0.25) -2(0.2 \log 0.2) - 2(0.05\log 0.05) = 1.64\]

Our loss was \(\frac{T}{t}\sum_{i \textrm{ masked}} -\log p_\theta(x_0^i|x_t)\). On average \(4t/T\) of the four positions are masked so the \(T/t\) cancels and the loss is four times the cost of one position: \(4 \times 1.64 = 6.55\).

Knowing where you are#

To tell the network where it is we are going to use exactly the same trick that we used for telling it what token it has. We keep a second table. This one has a row for each position (instead of each token) and we add the vector for the position to the vector for the token.

\[h_i = \textrm{embed}[x_i] + \textrm{position}[i]\]

Now the vector sitting on site \(i\) carries both pieces of information: the mask at position 1 and the mask at position 4 are different vectors. Again, the network gets to learn what the position vectors should be.

Write a class WithPosition(nn.Module) which is NoTalking with this added. You will need to

  • add an argument max_length=64 to __init__ and make a second table there: self.position = nn.Embedding(max_length, dim). (max_length is the longest sentence that the network will ever be able to deal with.)

  • in forward, make the list of positions 0, 1, ..., L-1 with torch.arange(x.shape[-1], device=x.device), look them up in your new table, and add the result to self.embed(x)

Two things about that last step. First, the device=x.device makes sure that if x is living on a GPU then your list of positions is too. It doesn’t matter yet but it will when we get to the stories, and forgetting it is the most common bug on this page. Second, your token vectors have shape (batch, L, dim) and your position vectors have shape (L, dim). When you add two tensors like this, pytorch lines the shapes up from the right and adds the smaller one to every sentence in the batch, which is exactly what we want.

Grading

Train a WithPosition network on the two sentences.

  • Add its loss to your plot. Where does it end up now?

  • Print out its guesses at position 2 for the same three sentences as before.

  • Use your Undiffuse with this network as the guesser to generate 1000 sentences with \(T=20\) and make the bar chart of what you get (the one from the masked diffusion page, with the four bars for the cat is alive, the box is empty, the cat is empty, the box is alive). What fraction of your sentences are nonsense? If your network guesses cat with probability \(p\) at position 2 and alive with probability \(p\) at position 4, what fraction would you predict? (Compare this with what the exact guesser gave you with \(T=1\) on the masked diffusion page.)

  • On the masked diffusion page the cure for nonsense sentences was to use more steps. Try \(T=100\). Does it help? Why not?

Much better. Your loss should now be about 1.0, positions 1 and 3 are always right, and at position 2 your network says cat 80% of the time and box 20% of the time.

But it says that no matter what. Even when the sentence in front of it is the _ is empty, it still thinks the cat is the better bet. And so somewhere between a quarter and a third of the sentences it writes are nonsense, and (unlike on the masked diffusion page) taking more steps doesn’t fix it. Undiffusing slowly only helps if your guess changes after a word has been written down. The sites of this network don’t talk to each other so there is no way for the guess at position 2 to find out what got written at position 4.

This network has learned everything that you can learn about each position separately. The whole gap between its loss of 1.0 and the best possible loss of 0.50 is the cost of not knowing what the other words are. To close that gap, the sites have to talk.

Step 3: Attention#

We need a way for the vector on site \(i\) to collect information from the vectors on the other sites. The simplest thing you could imagine is to replace each vector with the average of all of them. That gets information from everywhere, but an unweighted average treats every site the same and what we want is for position 2 to be able to listen to position 4 in particular.

So we are going to do a weighted average, and we are going to let each site choose its own weights depending on what is sitting on all of the sites. This is called attention and here is how it works. From the vector \(x_i\) on each site we make three new vectors, each by multiplying by a matrix (that the network will learn):

  • a query \(q_i = Q x_i\). This is site \(i\) describing what it is looking for.

  • a key \(k_i = K x_i\). This is site \(i\) describing what it has to offer.

  • a value \(v_i = V x_i\). This is what site \(i\) will actually say to anybody who listens.

Site \(i\) decides how much to listen to site \(j\) by how well its query matches \(j\)’s key. We measure the match with a dot product, and a physicist would turn this into an energy

\[E_{ij} = -\frac{q_i \cdot k_j}{\sqrt{d}}\]

where \(d\) is the length of the query and key vectors. A good match is a low energy. The weights are then the Boltzmann distribution for these energies

\[w_{ij} = \frac{e^{-E_{ij}}}{\sum_{j'} e^{-E_{ij'}}}\]

(i.e. a softmax over \(j\)) and what site \(i\) hears is the average of the values using these weights

\[o_i = \sum_j w_{ij} v_j\]

That is all attention is: each site computes a thermal average over all of the sites of what they have to say, where the energies come from how well queries match keys. If one match is much better than the others then site \(i\) listens (nearly) only to that one site. If all of the matches are about the same you get back the unweighted average.

The \(\sqrt{d}\) is there to set the temperature. The dot product of two random vectors of length \(d\) has a typical size which grows like \(\sqrt{d}\). If we didn’t divide by it, bigger vectors would mean bigger energies, which is the same thing as a lower temperature, and at low temperature the Boltzmann distribution freezes onto a single site before the network has learned anything.

Write a function

def Attention(q, k, v):
  # do stuff
  return out, weights

which takes the queries, keys and values (each a tensor of shape (L, d)) and returns both the outputs \(o_i\) (of shape (L, d)) and the weights \(w_{ij}\) (of shape (L, L)). It should only take about four lines:

  • q @ k.transpose(-2, -1) gives you the (L, L) matrix of all the dot products \(q_i \cdot k_j\)

  • q.shape[-1] is \(d\)

  • torch.softmax(..., dim=-1) does the Boltzmann distribution over \(j\) (the last index)

  • weights @ v does the sum over \(j\)

If you write it using -1 and -2 like this (instead of 0 and 1) it will also work without any changes on a whole batch of sentences, where everything has shape (batch, L, d).

Grading

Test your Attention using random queries, keys and values q, k, v = torch.randn(3, 5, 8) (this makes three tensors of shape (5, 8)).

  • Verify that every row of your weights adds up to one.

  • Verify that your output agrees with pytorch’s: torch.allclose(out, F.scaled_dot_product_attention(q, k, v), atol=1e-6)

  • Verify that your function also works on a batch: with q, k, v = torch.randn(3, 64, 5, 8) the output should have shape (64, 5, 8) and the weights (64, 5, 5).

  • Temperature. Plot the matrix of weights (with plt.matshow(weights, vmin=0, vmax=1)) for Attention(q, k, v), for Attention(10 * q, k, v) and for Attention(0.1 * q, k, v). Multiplying the queries by a number multiplies all of the energies by it. Which of these is hot and which is cold?

  • Shuffling. Pick a random reordering of the five sites with perm = torch.randperm(5) and compare Attention(q[perm], k[perm], v[perm]) to your original output. You should find that the new output is just the old output reordered in the same way (i.e. it is out[perm]). Attention has no idea what order the sites are in! This is why we needed to put the position into the vectors ourselves.

A head#

Now we need to package this up with the three matrices \(Q\), \(K\) and \(V\) so that pytorch can train them. One set of these three matrices along with the attention that uses them is called a head.

Write a class Head(nn.Module). Here is its skeleton:

class Head(nn.Module):
    def __init__(self, dim, head_dim):
        super().__init__()
        self.Q = nn.Linear(dim, head_dim, bias=False)
        # make self.K and self.V the same way

    def forward(self, x):
        # 1. make q, k and v by applying self.Q, self.K and self.V to x
        # 2. call your Attention on them
        self.weights = weights.detach()     # remember the weights so that we can look at them later
        return out

The vectors coming in have length dim. The queries, keys, values (and so the output) have length head_dim, which doesn’t have to be the same. You’ll see why in a second. (The .detach() tells pytorch that this saved copy of the weights is just for us to look at and isn’t part of anything it needs to differentiate.)

Many heads#

A head produces one set of weights, so on each site it can pay attention to one kind of thing. But there might be several different things worth looking for at the same time (in real text: the subject of the sentence, the previous word, whether we are inside quotation marks, …). So transformers use several heads side by side, each with its own \(Q\), \(K\) and \(V\), and stick their outputs together.

To keep the total size of things fixed, if there are heads heads we give each of them a head_dim of dim // heads (the // is division which rounds down to an integer). When we stick the heads outputs of length dim // heads end to end we are back to a vector of length dim. Finally we apply one more linear layer to the result to mix together what the different heads found.

Write a class MultiHead(nn.Module).

class MultiHead(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = nn.ModuleList([Head(dim, dim // heads) for h in range(heads)])
        self.mix = nn.Linear(dim, dim)

    def forward(self, x):
        # 1. make a (python) list of the output of each head: [head(x) for head in self.heads]
        # 2. stick them together along the last index with torch.cat(that_list, dim=-1)
        # 3. apply self.mix

There is one new piece of pytorch here. nn.ModuleList is a list which pytorch knows about. If you were to store your heads in an ordinary python list, pytorch wouldn’t find them, their matrices wouldn’t show up in net.parameters(), and they would silently never get trained.

Grading

Use x = torch.randn(64, 5, 32) (a batch of 64 sentences with 5 positions and vectors of length 32).

  • Make a head = Head(32, 8). Verify that head(x) has shape (64, 5, 8) and that head.weights has shape (64, 5, 5).

  • Make a multi = MultiHead(32, 4). Verify that multi(x) has shape (64, 5, 32).

  • Count the number of trainable parameters in multi using sum(p.numel() for p in multi.parameters()). Work out (with pencil and paper) what this number should be and check that you get it.

  • Change your nn.ModuleList([...]) to just [...] and count the parameters again. (Then change it back!)

Step 4: The Block#

Attention is how sites talk. Once a site has heard from the others it needs to do something with what it heard, and for that we use the most ordinary kind of neural network (the kind you used for continuous diffusion): a linear layer which makes the vector four times bigger, a nonlinear function, and a linear layer which brings it back down.

mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

nn.Sequential just means “do these one after another” and nn.GELU is a smooth function which is close to zero for negative numbers and close to \(x\) for positive ones. Remember that a linear layer acts on the last index, so this network is applied to the vector on each site separately. Nobody talks during this part.

One round of talking (attention) followed by one round of thinking (the mlp) is called a block. There are two more ingredients in it which aren’t deep but which matter a great deal for getting a stack of blocks to train.

Adding instead of replacing. A block doesn’t replace the vector on each site with what the attention says. It adds what the attention says to what was already there

\[x \rightarrow x + \textrm{attention}(x)\]

and the same for the mlp. So each piece only has to compute a correction to what the site already knows. In particular the token and the position that we put in at the beginning don’t have to be carefully carried through each layer. They stay there unless something changes them.

Normalizing. Every time we add something to our vectors they can get bigger, and after a few blocks they can be a very different size from what the next layer was expecting. So before the vectors go into the attention or the mlp we rescale them. nn.LayerNorm(dim) takes the vector on each site, subtracts its mean (over its dim components) and divides by its standard deviation. (It then multiplies by a trainable scale and adds a trainable shift, which start out as one and zero.)

Putting these together, a block does

x = x + attention(norm1(x))
x = x + mlp(norm2(x))

Write a class Block(nn.Module) which takes dim and heads. In __init__ you need to make four things: two nn.LayerNorm(dim) (they each have their own trainable scale and shift so you can’t use the same one twice), your MultiHead(dim, heads) (call it self.attention), and the mlp. In forward you need the two lines above.

Grading

  • Let x = torch.randn(64, 5, 32) and norm = nn.LayerNorm(32). Verify that norm(x) is the same (to about atol=1e-4) as subtracting the mean and dividing by the standard deviation yourself:

    (x - x.mean(dim=-1, keepdim=True)) / x.std(dim=-1, keepdim=True, unbiased=False)

    (The keepdim=True leaves the mean with shape (64, 5, 1) instead of (64, 5) so that it lines up with x when you subtract.)

  • Make a block = Block(32, 4) and verify that block(x) has the same shape as x.

Step 5: The Whole Thing#

We now have all of the pieces. Go back and look at the picture in The Plan.

Write a class MyTransformer(nn.Module) whose __init__ takes (self, vocab_size, dim=128, heads=4, layers=4, max_length=64). It is your WithPosition network with a stack of blocks in the middle.

In __init__ make

  • the token table self.embed and the position table self.position (exactly as in WithPosition)

  • self.blocks, an nn.ModuleList of layers blocks

  • one last self.norm = nn.LayerNorm(dim)

  • the output layer self.out = nn.Linear(dim, vocab_size)

  • and then, as the very last line, call SmallStart(self) (see below)

In forward

  1. add together the token vectors and the position vectors (exactly as in WithPosition)

  2. send the result through each block in turn: for block in self.blocks:

  3. apply self.norm and then self.out

About SmallStart: pytorch fills in a new nn.Embedding table with random numbers of size about one. For a transformer that is too big. It makes the network start off very confident about its (random) guesses, and it then has to spend a long time unlearning that before it can learn anything. The following function goes through all of the tables and matrices of a network and makes them start off small instead.

# GIVEN: call this as the last line of your __init__
def SmallStart(net):
    for piece in net.modules():                              # every piece inside the network (and inside those...)
        if isinstance(piece, (nn.Linear, nn.Embedding)):
            nn.init.normal_(piece.weight, std=0.02)          # gaussian random numbers with standard deviation 0.02

Grading

  • Make a small transformer net = MyTransformer(len(vocab), dim=32, heads=2, layers=2) and verify that if you hand it a batch of 64 of our four word sentences you get back something of shape (64, 4, 14).

  • How many trainable parameters does a default sized MyTransformer(2002) have? (2002 is going to be the size of our vocabulary when we get to the stories.) Which piece of the network has the most of them?

  • Verify that net(x)[0] (the logits for the first sentence in a batch) doesn’t depend on what the other sentences in the batch are.

Using Your Transformer#

Two sentences#

Let’s see what talking has bought us. We will start with the smallest transformer there is: one block with one head in it.

Grading

Train a MyTransformer(len(vocab), dim=32, heads=1, layers=1) on the two sentences.

  • Plot the loss of all three of your networks (NoTalking, WithPosition and MyTransformer) on the same plot using a log scale for the loss. Draw horizontal lines at 6.55, 1.0 and 0.50.

  • Print out its guesses at position 2 for _ _ _ _, _ _ is alive and the _ _ empty and compare them to the exact answer from your GuessX0.

  • Generate 1000 sentences with \(T=20\) and make the bar chart. What fraction are nonsense now?

That’s the whole gap closed. With a single round of attention your network gets down to (nearly) the best possible loss, it changes its mind about the cat when it sees the word empty, and your undiffusion has stopped writing nonsense.

Since you kept a copy of the attention weights we can go and look at who is listening to whom. After you call your network on a sentence

x = Encode("the _ is alive")
net(x[None])          # x[None] adds an index at the front:  it turns one sentence of shape (4,) into a batch of shape (1, 4)
weights = net.blocks[0].attention.heads[0].weights[0]       # block 0, head 0, sentence 0 of the batch:  an (L, L) matrix

row \(i\) of weights tells you how much site \(i\) listened to each of the sites.

Grading

For your trained one block, one head transformer plot the matrix of attention weights (with plt.matshow(weights, vmin=0, vmax=1)) for the sentences the _ is alive, the _ is empty and _ _ _ _. Label the ticks with the words of the sentence.

  • Look at the row for position 2 (the masked one) in the first two. To get its answer right it needs to find out what is at position 4. Is it listening there?

  • What about when everything is masked?

Don’t expect this to look perfectly clean, and expect it to look somewhat different if you train again. All the training asks of the network is that it gets the right answer. It doesn’t have to do it in the way that is easiest for us to read!

Stories#

Now for the real test. On the masked diffusion page you trained TinyDenoiser on a small piece of TinyStories and had it write stories. Let’s do that again with a network that you built.

Here is the same code as before for downloading the stories.

# GIVEN: a tiny corpus (the same one as on the masked diffusion page)
import re
import requests
from collections import Counter

url = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt"
raw = requests.get(url, headers={"Range": "bytes=0-3999999"}).content[:4_000_000].decode("utf-8", errors="ignore")
words = re.findall(r"[a-z']+|[.,!?\"]", raw.lower().replace("<|endoftext|>", " "))

story_vocab = [w for w, n in Counter(words).most_common(2000)] + ["<unk>", "_"]
STORY_UNK, STORY_MASK = len(story_vocab) - 2, len(story_vocab) - 1
index = {w: i for i, w in enumerate(story_vocab)}

def StoryTokens(text):
    return torch.tensor([index.get(w, STORY_UNK) for w in re.findall(r"[a-z']+|[.,!?\"_]", text.lower())])

def StoryText(x):
    return " ".join(story_vocab[i] for i in x.tolist())

corpus = torch.tensor([index.get(w, STORY_UNK) for w in words])

def GetBatch(batch=64, L=32):
    start = torch.randint(len(corpus) - L, (batch,))
    return torch.stack([corpus[s:s + L] for s in start])

print(f"{len(corpus)} tokens of text, {len(story_vocab)} tokens in the vocabulary")
print(StoryText(GetBatch(1)[0]))

Train a default sized MyTransformer on it. This is the same thing that you did on the masked diffusion page:

  • The network is net = MyTransformer(len(story_vocab)).to(device). The .to(device) moves all of its tables and matrices onto the GPU (if you have one).

  • Your data is x_0 = GetBatch().to(device) and your mask is STORY_MASK

  • Use 5000 steps. (As before, if you don’t have a GPU you can get away with 1500.)

  • Keep track of the loss per token, loss.item() / 32, and print it out every few hundred steps so that you know it’s working.

If you get an error complaining that two tensors are on different devices, the culprit is almost certainly the torch.arange in your forward. It needs that device=x.device.

Grading

  • Plot the loss per token as a function of training step (with a window average). Draw horizontal lines at 7.6 (guessing every token with equal probability) and 5.4 (knowing which words are common but nothing else). How does your final loss compare to what you got from TinyDenoiser?

  • Make a guesser with GuessStory = MakeGuesser(net, forbidden=[STORY_MASK, STORY_UNK]) and write five stories of 32 tokens with your favorite undiffuser and settings from the masked diffusion page.

  • Have it write three stories which start one day , lily went to the park . and end and they were all happy . (Recall that StoryTokens turns a _ into a mask.)

  • Pick one sentence from the stories, mask one word in it, and plot the attention weights of each of the four heads in the first block and in the last block. You don’t need to be able to explain what you see. Do the different heads look at different things?

Flipping one switch: writing left to right#

The models that you are used to talking to (chatGPT, Claude, …) don’t undiffuse. They write one token at a time from left to right, each time guessing what the next token is given all of the tokens so far. Their network is also a transformer and it differs from the one you just built in exactly one way:

When site \(i\) is listening, it is only allowed to listen to sites \(j \leq i\). It can’t look to its right.

You already have the tool for forbidding something: give it an infinite energy. If we set \(E_{ij} = \infty\) whenever \(j > i\) then the Boltzmann weight \(e^{-E_{ij}}\) of those sites is exactly zero and the weights of the sites that are left still add up to one. A transformer with this change made to its attention is called causal (because what comes out at site \(i\) can only be affected by things at or before \(i\)).

Add an argument causal=False to your Attention function. When it is True, set the energies with \(j>i\) to infinity before you do the softmax. Here is how to do that in pytorch:

L = q.shape[-2]
future = torch.triu(torch.ones(L, L, dtype=torch.bool, device=q.device), diagonal=1)
energy = energy.masked_fill(future, math.inf)

The first line gets the number of sites. The second makes an (L, L) matrix which is True above the diagonal (where \(j>i\)) and False everywhere else. (torch.triu keeps the upper triangle of a matrix and diagonal=1 says to start one above the diagonal.) The third sets energy to infinity wherever future is True. If you are doing your softmax on minus the energy (i.e. on the dot products) then you want to set those entries to -math.inf instead.

Then add a causal=False argument to the __init__ of each of your Head, MultiHead, Block and MyTransformer and pass it down the chain until it gets to the place where Attention is called. (In Head you will need to save it as self.causal so that forward can use it.) Nothing else in your network changes.

Grading

  • Repeat the plot of the attention weights for random q, k, v = torch.randn(3, 5, 8) with causal=True. Verify that the rows still add up to one.

  • Make a net = MyTransformer(len(vocab), dim=32, heads=2, layers=2, causal=True). Call it on Encode("the cat is alive") and on Encode("the cat is empty"). These only differ in the last word. Verify that the logits at the first three positions are identical and the logits at the last position are not.

  • Do the same thing with causal=False. Now all four positions should differ.

Now we need to train it. A left to right model has a different (and simpler) job than our diffusion model did. There are no masks. We show the network real text, and at every position we ask it to guess the next token:

what the network sees:        once   upon   a      time   there   was
what it should guess:         upon   a      time   there  was     a

So the input is the sentence with its last token chopped off, the right answer is the sentence with its first token chopped off, and the loss is the same cross entropy as before: the average of \(-\log p_\theta(\textrm{the correct next token})\). Here it is.

# GIVEN: the loss for a left to right model
def NextTokenLoss(net, x_0):
    logits = net(x_0[:, :-1])                                     # everything but the last token goes in
    return F.cross_entropy(logits.transpose(1, 2), x_0[:, 1:])    # everything but the first token is the answer

You can now see why the network has to be causal. The answer that we want at position \(i\) is the token that is sitting at position \(i+1\) of the input. If the sites were allowed to listen to their right, the network would just copy it. It would get a perfect score and it would have learned nothing. With causal attention the only way to do well is to actually predict.

Notice also that we are getting a lot out of each sentence. Position 1 is being trained to continue a one token story, position 2 to continue a two token story, and so on, all at the same time and all with one call to the network. This only works because of causality: the output at position 5 is exactly what it would have been if the sentence had stopped there.

To write with this model you start with some text and then repeat:

  1. call the network on everything you have so far: logits = net(x[None].to(device))[0, -1].cpu() gives you the logits at the last position, which are the network’s guess for the token that comes next

  2. set the logits of the tokens you never want (STORY_MASK and STORY_UNK) to -math.inf

  3. turn the logits into probabilities with a softmax. (If you’d like a temperature \(\tau\), divide the logits by \(\tau\) first. A logit is minus an energy so this is \(e^{-E/\tau}\).)

  4. sample one token from them: torch.multinomial(p, 1) returns a tensor with one sampled token in it

  5. stick it on the end: x = torch.cat([x, new_token])

Write a function def WriteLeftToRight(net, start, L=32, temperature=1.0) which does this starting from StoryTokens(start) until the story is L tokens long, and returns StoryText(x). Put a @torch.no_grad() on the line above your def (this tells pytorch that it doesn’t need to keep track of derivatives inside this function, which makes it faster).

Grading

Train a new network net = MyTransformer(len(story_vocab), causal=True).to(device) on the stories for 5000 steps using NextTokenLoss(net, x_0) as the loss. (It is already a loss per token so you don’t need to divide by 32.)

  • Plot the loss as a function of training step on the same plot as the loss per token of your diffusion model.

  • Have it write five stories which start with once upon a time. Try a couple of different temperatures.

  • Have your diffusion model from the last section write five stories which start the same way. How do the two compare?

  • How many times does each model have to call its network to write a 32 token story? Could either of them do it in fewer?

  • On the masked diffusion page your model could write a story which started one day , lily went to the park . and ended and they were all happy . Can this one?

What You Built#

You now have a transformer that you wrote yourself, you know what each piece of it is for, and you have used the same code both as a diffusion model and as a (very small) GPT.

It is worth going back to the masked diffusion page and reading the TinyDenoiser code that we handed you. You should be able to recognize every line. There are two differences from yours.

  • All of the heads are done at once. You made a python list of heads and called them one at a time. TinyDenoiser uses one big matrix to make the queries, keys and values for all of the heads together and then reshapes the result so that “which head” is one more index of the tensor. It is the same calculation. It is just faster.

  • Position is put in differently. You added a position vector to the token vector at the very beginning. TinyDenoiser also rotates each query and key by an angle which is proportional to its position (the Rotate function) before taking the dot products. The dot product of two rotated vectors only depends on the difference of their angles, so the energies \(E_{ij}\) get to depend directly on how far apart \(i\) and \(j\) are. This works a bit better, which is why your loss on the stories came out slightly higher.

And here is how the networks that you have met compare. The first three rows use attention where every site listens to every site and the last two use causal attention. Other than that (and their size) they are all the thing that you just built.

attention

blocks

length of the vectors

parameters

trained to

your MyTransformer

everybody

4

128

1 million

fill in masks (diffusion)

ModernBERT (the chat model)

everybody

28

1024

400 million

fill in masks (diffusion)

Qwen3-0.6B diffusion (the optional bigger model)

everybody

28

1024

600 million

fill in masks (diffusion)

your causal MyTransformer

only to the left

4

128

1 million

guess the next token

GPT-2

only to the left

12

768

120 million

guess the next token

The Qwen3 model has an interesting history. It started out life as a causal, left to right model. The people who made the diffusion version took those trained weights, removed the triangle of infinite energies (that was the see_everything in the code we gave you), and trained it some more with the diffusion loss. You could try this yourself: you have a trained causal network and you have DiffusionLoss.

Further reading

  • Vaswani et al., Attention Is All You Need (2017). The paper that introduced the transformer.

  • Andrej Karpathy, Let’s build GPT: from scratch, in code, spelled out (a video lecture) and the accompanying nanoGPT code. Builds the left to right version at about the level of this page.

  • Ramsauer et al., Hopfield Networks is All You Need (2020). Shows that the attention you wrote is the update rule of a modern version of the Hopfield network from the first page of this unit.

  • Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021). The rotation trick that TinyDenoiser uses for position.