Build Your Own LLM, Lesson 4: Embeddings and Your First Neural Language Model

📎 This article includes 1 downloadable practice file ↓

🧠 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
8. Continue learning

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

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.

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.

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

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

@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.

Continue learning

All Build Your Own LLM guides →

Build Your Own LLM, Lesson 12: Quantize and Run Your Model LocallyShrink models with quantization: fp32 to int8 with PyTorch dynamic quantization, size and speed…Build Your Own LLM, Lesson 11: Instruction Tuning and Chat TemplatesTurn a text predictor into an assistant: instruction-response data, chat templates with role tokens,…Build Your Own LLM, Lesson 10: Fine-Tuning a Pre-Trained Model With LoRAFine-tune efficiently with LoRA: freeze the base model, add low-rank adapter matrices to attention…Build Your Own LLM, Lesson 9: Evaluating a Model — Loss, Perplexity and Test PromptsMeasure a language model: cross-entropy loss, perplexity and what it means, validation sets, fixed…

📎 Practice files for this article

• 🐍
  All code from this post5 examples in one file, ready to open in your editor.
  ⬇ PY · 2 KB

Free to use for learning. Files with macros (.bas) are plain text — import them with Alt+F11 → File → Import File, and always test on a copy.