Neural language models

Here's the same recipe as one formula. Read it from the inside out.

The symbols, one at a time:

• $t$ is the position of the character you're predicting.
• $x_{t-1}$ is the character just before it. $x_{t-n}$ is the character $n$ places back.
• $e$ stands for an embedding. So $e_{x_{t-1}}$ is the embedding of the character just before.
• The square brackets with semicolons mean "place side by side". The dots stand for the characters in between.
• $W_1$ is the hidden layer's table of weights.
• $b_1$ is the hidden layer's bias, a list of numbers added on just before tanh.
• $W_2$ is the output layer's table of weights.
• $b_2$ is the output layer's bias, added on at the very end.

In plain words: look up the embeddings of the last few characters, line them up, pass them through one hidden layer, and give every possible next character a score.

Why this beats the table

It shares what it learns. Characters (or, in real models, tokens) that behave alike get similar embeddings. So whatever the network learns about one context also helps with similar ones. That's how it generalizes to contexts it has never seen.

It grows slowly. Only the first layer gets bigger when the window grows. A model like the one you're about to train has 13,112 weights, counting the biases, with 6 characters of context. Double the window to 12 characters and it has 20,024 weights. A counting table for 6 characters would already need about 31 billion rows. The network grows in step with the context length. The table multiplies.

V, d, H = 56, 12, 96          # characters, numbers per embedding, hidden units
for block in [6, 12]:
    params = V * d + (block * d) * H + H + H * V + V
    print(f"{block} characters of context: {params:,} weights")
print(f"a counting table for 6 characters: {V ** 6:,} rows")

import math

emb = {"a": [1, 0], "b": [0, 1], "c": [1, 1]}       # each character → 2 numbers
x = emb["a"] + emb["b"]                               # context "ab", lined up side by side
W1, b1 = [[1, 0, 0, 1], [0, 1, 1, 0]], [0, 0]        # hidden layer: 2 units
h = [math.tanh(sum(w * v for w, v in zip(row, x)) + b) for row, b in zip(W1, b1)]
W2, b2 = [[1, 0], [0, 1], [-1, 0]], [0, 0, 0]        # output layer: one row per character
logits = [sum(w * v for w, v in zip(row, h)) + b for row, b in zip(W2, b2)]
probs = [math.exp(z) / sum(math.exp(s) for s in logits) for z in logits]
print("lined up:", x)
print("hidden:", [round(v, 2) for v in h])
print("logits:", [round(v, 2) for v in logits])
print("probs:", [round(p, 2) for p in probs])

block, emb_dim = 6, 12
print(block * emb_dim)

Training one on the catalog

The code below builds the real thing and trains it. Here's what each part does.

• windows slides along the text. Each row of X holds 6 characters that sit next to each other in the text. The matching entry of y is the character that comes after them.
• The model is the four steps. nn.Embedding gives each character 12 numbers. nn.Flatten lines the 6 embeddings up into 72 numbers. Then come a hidden layer of 96 units with tanh, and an output layer with one score per character.
• Training runs 800 steps. Each step picks 64 random windows and measures the cross-entropy loss. Then the Adam optimizer nudges every weight.
• Afterwards the code measures the loss on held-out text. It writes 110 new characters, starting from "The Ju". Last, it lists the 5 characters whose embeddings point most nearly the same way as the one for "a".

import numpy as np
import requests
import torch
import torch.nn as nn

text = requests.get("https://api.gradient.test/v1/corpus").json()["text"]
docs = text.split("\n")
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
V, block = len(chars), 6

def windows(s):
    ids = torch.tensor([stoi[c] for c in s])
    X = torch.stack([ids[i:len(ids) - block + i] for i in range(block)], dim=1)    # 6 previous characters
    return X, ids[block:]                                                          # the character after them

X_train, y_train = windows("\n".join(docs[:270]))
X_val, y_val = windows("\n".join(docs[270:]))

torch.manual_seed(0)
model = nn.Sequential(
    nn.Embedding(V, 12),              # each character → 12 numbers
    nn.Flatten(),                     # 6 embeddings side by side → 72 numbers
    nn.Linear(block * 12, 96), nn.Tanh(),
    nn.Linear(96, V),                 # logits for the next character
)
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
for step in range(800):
    idx = torch.randint(0, len(X_train), (64,))
    loss = nn.functional.cross_entropy(model(X_train[idx]), y_train[idx])
    opt.zero_grad()
    loss.backward()
    opt.step()

with torch.no_grad():
    val_loss = nn.functional.cross_entropy(model(X_val), y_val).item()
    print(f"held-out loss: {val_loss:.3f}   (bigram: about 2.17)")
    ctx, out = [stoi[c] for c in "\nThe Ju"[-block:]], "The Ju"
    torch.manual_seed(3)
    for _ in range(110):
        probs = torch.softmax(model(torch.tensor([ctx[-block:]])), dim=-1)[0]
        nxt = torch.multinomial(probs, 1).item()
        out += chars[nxt]
        ctx.append(nxt)
print(repr(out))

E = model[0].weight.detach().numpy()
E = E / np.linalg.norm(E, axis=1, keepdims=True)
sims = E @ E[stoi["a"]]
print("characters whose embeddings are closest to 'a':", [chars[j] for j in np.argsort(-sims)[1:6]])

Six characters of context take the held-out loss from about 2.17 (the bigram's score) to well under 1. Lower is better. The samples now contain real catalog phrases ("made from", "Built") mixed with near-misses.

The embeddings are still rough after so little training. Even so, "a" already sits near "A" and other vowels. With more data, the patterns in learned embeddings become much clearer. (As before, torch is the course's in-browser PyTorch subset.)

Watch out: this catalog repeats itself a lot, which makes it easier to predict. Don't expect a loss this low on ordinary text.

What it still can't do

Imagine reading a book through a slot that shows six letters at a time. You'd do well on spelling. You'd lose track of the story. This model has that problem, plus two more.

• A fixed window. The model sees exactly six characters. Anything earlier is invisible, such as which product a description is about. So it can't keep a sentence consistent with how it began.
• Separate weights for each slot. Each position in the window has its own weights. So a pattern learned at position 2 has to be learned again at position 5.
• A bigger window is costly. Doubling the context doubles the first layer. And every slot is still treated separately.

Recurrent networks, from Phase 4, tried to carry information forward without limit. They ran into vanishing gradients, and they're slow because they work one step at a time. The fix that won is attention. It lets every position look straight at every earlier position. Then it decides, based on content, which ones matter. That's the next lesson.