🧠 Lesson 4 of 12 · Build Your Own LLM course

In this article

1. Prepare the data
2. A bigram model as a neural network
3. What one training step does
4. Adding real embeddings
5. Looking inside the embeddings
6. Generate text
7. Check your understanding

In Lesson 3 our model was a table of counts. Now we get the same behaviour from a neural network trained with gradient descent — the exact training loop that scales up to GPT. Install PyTorch first (pip install torch; the CPU version is fine).

Prepare the data

“`python
import torch
import torch.nn as nn
import torch.nn.functional as F

text = open(“input.txt”, encoding=”utf-8″).read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
data = torch.tensor([stoi[c] for c in text])
V = len(chars)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]
“`

A bigram model as a neural network

nn.Embedding(V, V) is a trainable V×V table. Row i holds scores (logits) for which character follows character i — the neural version of our count table.

“`python
class Bigram(nn.Module):
def __init__(self, V):
super().__init__()
self.table = nn.Embedding(V, V)

def forward(self, idx):
return self.table(idx) # (batch, V) logits

def batch(split, size=256):
d = train_data if split == “train” else val_data
ix = torch.randint(0, len(d) – 1, (size,))
return d[ix], d[ix + 1] # inputs, next characters

model = Bigram(V)
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
for step in range(3001):
x, y = batch(“train”)
loss = F.cross_entropy(model(x), y)
opt.zero_grad()
loss.backward()
opt.step()
if step % 500 == 0:
print(step, round(loss.item(), 3))
“`

The loss starts near log(V) (random) and settles around the same value the counting model reached. That is the point: gradient descent rediscovered the count table. F.cross_entropy is the exact loss we computed by hand in Lesson 3.

What one training step does

1. Forward: look up logits for 256 random characters.
2. Loss: how much probability went to the true next character.
3. Backward: loss.backward() computes, for every number in the table, which direction would reduce the loss.
4. Update: the optimiser nudges each number a little in that direction.

Adding real embeddings

Now the important step. Instead of jumping straight from a character to V scores, we first map each character to a short vector — its embedding — then use a linear layer to produce scores.

“`python
class EmbedModel(nn.Module):
def __init__(self, V, d=32):
super().__init__()
self.emb = nn.Embedding(V, d) # each token -> 32 numbers
self.head = nn.Linear(d, V) # 32 numbers -> V scores

def forward(self, idx):
return self.head(self.emb(idx))
“`

Why bother? Because tokens that behave alike end up with similar vectors. The model can then share what it learns: if it learns something about a, it partly applies to e. In a real LLM, the embedding for “Monday” ends up near “Tuesday”, and “VLOOKUP” near “XLOOKUP”.

Looking inside the embeddings

“`python
emb_model = EmbedModel(V)
# … train it with the same loop as above (swap `model` for `emb_model`) …

E = F.normalize(emb_model.emb.weight.detach(), dim=1)
def similar(c, k=5):
sims = E @ E[stoi[c]]
best = sims.argsort(descending=True)[1:k + 1]
return [(itos[i.item()], round(sims[i].item(), 2)) for i in best]

print(similar(“a”)) # often other vowels
print(similar(“5”)) # often other digits
“`

Nobody told the model what a vowel or a digit is. It discovered that these characters appear in similar places, so their vectors drifted together.

Generate text

“`python
@torch.no_grad()
def generate(model, start=”T”, n=300):
idx = torch.tensor([stoi[start]])
out = [start]
for _ in range(n):
probs = F.softmax(model(idx)[-1], dim=-1)
idx = torch.multinomial(probs, 1)
out.append(itos[idx.item()])
return “”.join(out)

print(generate(emb_model))
“`

Still gibberish — each prediction sees only one character. To use more context, tokens need a way to look at each other. That mechanism is attention, next lesson.

Check your understanding

• Validation loss: compute F.cross_entropy on a batch("val") inside torch.no_grad(). It should be close to training loss; if it is much higher, the model is memorising.
• Parameter count: sum(p.numel() for p in emb_model.parameters()). With V=65 and d=32 it is only a few thousand. GPT-3 has 175 billion.

Keep reading

Build Your Own LLMBuild Your Own LLM, Lesson 1: How a Language Model Actually WorksBuild Your Own LLMBuild Your Own LLM, Lesson 2: Write a BPE Tokenizer in PythonBuild Your Own LLMBuild Your Own LLM, Lesson 3: A Language Model Made of Counting