A Recurrent Net from Scratch

You've used L S T M layers. But what's really happening inside a network with memory? The surprising truth is that the core of it is one line of maths. So let's write it by hand, with raw tensors, and teach it to write text one letter at a time.

• What's inside a network with memory?
• One line of maths
• Write it by hand

The hidden state

We've got three weight tables, and that's all of it. For each letter, this one line does the whole job. The new memory is a mix of the current letter and the old memory, squashed by tanh. Then we guess the next letter from that memory. And we add up the loss, across forty letters. When we call backward, gradients flow back through all forty steps. That's called backpropagation through time. Run it. The loss falls from about two point four down to almost nothing. And look at the text that it writes. The cat ate. The cat sat on the mat. It's generating sentences, one letter at a time, using nothing but memory, from a few lines of Python.

import torch

torch.manual_seed(0)
text = "the cat sat on the mat. the cat ate. " * 30
chars = sorted(set(text))
ix = {c: i for i, c in enumerate(chars)}
data = torch.tensor([ix[c] for c in text])
n, size = len(chars), 64

Wx = torch.nn.Parameter(torch.randn(n, size) * .1)
Wh = torch.nn.Parameter(torch.eye(size))
Wy = torch.nn.Parameter(torch.randn(size, n) * .1)
opt = torch.optim.Adam([Wx, Wh, Wy], lr=.01)

for step in range(401):
    start = torch.randint(0, len(data) - 41, (1,)).item()
    h, loss = torch.zeros(size), 0
    for t in range(40):
        h = torch.tanh(Wx[data[start + t]] + h @ Wh)
        loss = loss + torch.nn.functional.cross_entropy(h @ Wy, data[start + t + 1])
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 200 == 0:
        print(f"Step {step}: loss {loss.item() / 40:.2f}")

h, c, out = torch.zeros(size), ix["t"], "t"
for _ in range(40):
    h = torch.tanh(Wx[c] + h @ Wh)
    c = (h @ Wy).argmax().item()
    out += chars[c]
print(out)

Step 0: loss 2.41
Step 200: loss 0.01
Step 400: loss 0.03
the cat ate. the cat sat on the mat. the

Why this matters

Now that you've seen it, some mysteries disappear. Vanishing gradients happen because that one line gets applied over and over. L S T M adds gates to protect the memory. And notice the identity matrix that we used to begin with. Starting the memory weights close to the identity helps a lot, because the state then passes through unchanged at first.

Summary