Recurrent networks

A character-level language model

Now a real one. You'll train an RNN to predict the next character of a short text. Then you'll let it write. This is a language model in the same sense as an LLM, only tiny. It works on characters instead of tokens. And it learns from a few hundred characters instead of trillions of tokens.

The training text is a few short lines about outdoor gear, repeated four times. Here is what the code does, in order:

1. Turn characters into numbers. chars lists every different character in the text. stoi ("string to integer") gives each one a number.
2. Build the model. nn.Embedding turns each character's number into a list of 48 learned numbers. Wx and Wh are the two weight tables from the formula. out turns the hidden state into one score per character. These scores are called logits. A higher score means "more likely to come next".
3. Train it. Each of the 300 steps picks 16 random snippets of 20 characters. The targets are the same snippets moved along by one character. So at every position, the right answer is the next character. The loss is cross-entropy, from Phase 2.
4. Let it write. Start from the text "good ". Turn the scores for the next character into chances, pick one at random, and add it to the text. Then feed it back in. Repeat 80 times.

import torch
import torch.nn as nn

text = ("great gear for great trails. light layers for long days. "
        "good boots for wet trails. warm layers for cold nights. ") * 4
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
V = len(chars)

class CharRNN(nn.Module):
    def __init__(self, vocab, hidden):
        super().__init__()
        self.hidden = hidden
        self.embed = nn.Embedding(vocab, hidden)          # each character → a vector
        self.Wx = nn.Linear(hidden, hidden)
        self.Wh = nn.Linear(hidden, hidden, bias=False)
        self.out = nn.Linear(hidden, vocab)               # hidden state → next-character logits

    def forward(self, idx, h=None):                       # idx: (batch, time)
        h = torch.zeros(idx.shape[0], self.hidden) if h is None else h
        x = self.embed(idx)
        logits = []
        for t in range(idx.shape[1]):                     # one step at a time, same weights
            h = torch.tanh(self.Wx(x[:, t]) + self.Wh(h))
            logits.append(self.out(h))
        return torch.stack(logits, dim=1), h

torch.manual_seed(0)
model = CharRNN(V, 48)
opt = torch.optim.Adam(model.parameters(), lr=0.01)
T = 20
for step in range(300):
    starts = torch.randint(0, len(data) - T - 1, (16,)).tolist()
    xb = torch.stack([data[i:i + T] for i in starts])           # 16 snippets of 20 characters
    yb = torch.stack([data[i + 1:i + T + 1] for i in starts])   # the same snippets shifted by one
    logits, _ = model(xb)
    loss = nn.functional.cross_entropy(logits.view(-1, V), yb.view(-1))
    opt.zero_grad()
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    if step % 100 == 0 or step == 299:
        print(f"step {step:>3}  loss {loss.item():.3f}")

with torch.no_grad():
    out = "good "
    logits, h = model(torch.tensor([[stoi[c] for c in out]]))
    for _ in range(80):
        probs = torch.softmax(logits[:, -1] / 0.7, dim=-1)       # temperature 0.7
        nxt = torch.multinomial(probs, 1)
        out += chars[nxt.item()]
        logits, h = model(nxt, h)                                # feed the sample back in
print(repr(out))

Everything you met in Phase 1 is here in miniature. Next-token prediction is the training target. Cross-entropy is the loss. Writing works by sampling with a temperature and feeding each output back in. On this little text the model mostly memorizes. Sometimes it mixes phrases it saw into new ones.

Two details are worth a note. Dividing the scores by 0.7 sets the temperature. A value below 1 makes likely characters even more likely, so the writing is less random. And clip_grad_norm_ caps how big the gradient can get at each step. RNNs are prone to exploding gradients, and this stops one bad step from wrecking training. (torch here is the course's in-browser PyTorch subset.)

Why RNNs forget

In the worked example, a memory faded a little at every step. Training has the same problem, and it's worse.

Training uses backpropagation. Through an unrolled RNN, the gradient travels backward one time step at a time. At each step it gets multiplied by one factor, the same way it does through the layers of a deep network. The difference is that it's the same $W_h$ every time. If those factors are below 1, the gradient shrinks step after step. You met this in the lesson on vanishing gradients.

The code below measures it. It runs a fresh, untrained RNN over a 40-character sequence. The loss depends only on the last step. Then it checks how much gradient reaches the hidden state at earlier steps. retain_grad() asks PyTorch to keep those in-between gradients so you can look at them. .norm() gives the size of each gradient as one number.

import torch
import torch.nn as nn

torch.manual_seed(1)
H = 48
embed, Wx, Wh, out = nn.Embedding(20, H), nn.Linear(H, H), nn.Linear(H, H, bias=False), nn.Linear(H, 20)
x = embed(torch.randint(0, 20, (1, 40))).detach()
h, states = torch.zeros(1, H), []
for t in range(40):
    h = torch.tanh(Wx(x[:, t]) + Wh(h))
    h.retain_grad()                         # keep the gradient of this intermediate value
    states.append(h)
out(h).sum().backward()                     # a loss that depends only on the last step
for t in [39, 35, 30, 20, 10, 0]:
    print(f"gradient reaching step {t:>2}: {states[t].grad.norm().item():.1e}")

The numbers print in short form. A value like 2.5e-04 means 2.5 with the decimal point moved 4 places to the left: 0.00025.

The signal shrinks by roughly 10× every few steps. By the start of the sequence, it's close to zero. So the network can't learn that something at step 0 matters for step 39. A plain RNN has a short memory, in training as much as in use.

g = 1.0
for step in range(10):
    g = g * 0.5
print(g)

LSTMs and GRUs

The LSTM (long short-term memory, 1997) was the fix. Alongside the hidden state, it keeps a second memory called the cell state. Think of it as a notebook. A plain RNN rewrites its whole memory at every step, through tanh. An LSTM mostly adds new notes to its notebook and crosses a few old ones out.

Small learned switches called gates decide what happens at each step. Each gate gives numbers between 0 and 1. A 0 means "block" and a 1 means "let through".

• The forget gate decides what to cross out.
• The input gate decides what new notes to write.
• The output gate decides what to read out as the hidden state.

The cell state is updated mostly by addition. So gradients can flow back along it without being squashed at every step. It works like a residual connection through time: a shortcut path the signal can travel along mostly unchanged. As a result, LSTMs can learn from much longer sequences. The GRU (gated recurrent unit) is a simpler variant with fewer gates.

In PyTorch they're one line each:

rnn = nn.LSTM(input_size=64, hidden_size=128, num_layers=2, batch_first=True)
outputs, (h_n, c_n) = rnn(x)            # x: (batch, time, 64) → outputs: (batch, time, 128)

Here x is a batch of sequences, with 64 numbers per step. hidden_size=128 sets the size of the hidden state. num_layers=2 stacks two LSTMs, the second reading the first one's outputs. Besides the outputs, you get back the last hidden state h_n and the last cell state c_n.

LSTMs drove the big advances in translation and speech around 2014–2017.

The limits that led to transformers

Even LSTMs kept two problems.

1. One step at a time. Step $t$ needs the result of step $t-1$. So a long sequence can't be processed in parallel. Think of a long queue at a single checkout. GPUs are built to do thousands of sums at once, so during training they sit underused.
2. A bottleneck. Everything about the past must squeeze into one fixed-size hidden vector. It's like summing up a whole book on one index card. And distant information still fades.

The transformer (2017) dropped recurrence. It uses attention: every position looks directly at every earlier position, all in one parallel step. There's no long chain of multiplications in between. That's Phase 5. (Some recent "state-space" models bring recurrence back, with clever ways around both limits. So the idea isn't gone.)