LSTM and GRU

Long-term memory.

⏱ 2 min read▶ Video at 57:38🎓 Free

Watch this lesson in the Deep Learning courseWatch on YouTube ↗Full Deep Learning course

Try this sentence. I grew up in France, and after many years abroad, working in many countries, I still speak fluent. You know that the next word is French, because you remember France from the start. A simple R N N would've forgotten. The vanishing gradient strikes again, this time across time.

I grew up in France...
A simple RNN forgets
Vanishing gradients across time

Gates that control memory

In nineteen ninety seven, two researchers invented the L S T M, or long short term memory. It adds a separate memory lane that runs straight through time, along with three gates. The forget gate decides what to drop. The input gate decides what to store. And the output gate decides what to reveal. Because information can travel along this lane unchanged, gradients survive. A G R U is a simpler version, with two gates, and often it works just as well.

Cell statea memory lane through time

Forget gatewhat to drop

Input gatewhat to store

Output gatewhat to reveal

The memory test

Here's a task that needs real memory. In a sequence of thirty random numbers, two of them are marked. The model must add those two numbers, ignoring the rest. One marker is right at the start, so the model has to hold on to that across many steps. We train all three kinds of network. Run it. The plain R N N stays stuck, close to just guessing. L S T M and G R U do much better. That's the gain from having gates.

memory_test.pypython

import torch
from torch import nn

torch.manual_seed(0)

def data(n, length):
    x = torch.zeros(n, length, 2)
    x[:, :, 0] = torch.rand(n, length)
    x[:, 0, 1] = x[:, length // 2, 1] = 1
    y = x[:, 0, 0] + x[:, length // 2, 0]
    return x, y.unsqueeze(1)

class Net(nn.Module):
    def __init__(self, kind):
        super().__init__()
        self.rnn = kind(2, 32, batch_first=True)
        self.out = nn.Linear(32, 1)
    def forward(self, x):
        return self.out(self.rnn(x)[0][:, -1])

for name, kind in [("RNN", nn.RNN), ("LSTM", nn.LSTM), ("GRU", nn.GRU)]:
    torch.manual_seed(0)
    net = Net(kind)
    opt = torch.optim.Adam(net.parameters(), lr=0.003)
    for _ in range(500):
        x, y = data(64, 30)
        loss = ((net(x) - y) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"{name:<5} final error {loss.item():.4f}")

Output

RNN   final error 0.1468
LSTM  final error 0.1583
GRU   final error 0.0020

A network that writes

Now for some fun. We train a small G R U to predict the next character, on a few repeated sentences. Each target is simply the same text, shifted by one place. After training, we give it a single letter. We sample the next one, feed that back in, and keep going. Run it. It writes new text, one character at a time, that looks like its training data. That's a language model. A tiny one, but the same basic idea as the ones that you've heard of.

char_rnn.pypython

import torch
from torch import nn

text = "the cat sat on the mat. the dog sat on the log. " * 40
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])

class CharRNN(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.emb = nn.Embedding(n, 32)
        self.gru = nn.GRU(32, 64, batch_first=True)
        self.out = nn.Linear(64, n)
    def forward(self, x, h=None):
        y, h = self.gru(self.emb(x), h)
        return self.out(y), h

torch.manual_seed(0)
model = CharRNN(len(chars))
opt = torch.optim.Adam(model.parameters(), lr=0.01)
for step in range(150):
    i = torch.randint(0, len(data) - 41, (16,))
    x = torch.stack([data[j:j + 40] for j in i])
    y = torch.stack([data[j + 1:j + 41] for j in i])
    loss = nn.functional.cross_entropy(model(x)[0].transpose(1, 2), y)
    opt.zero_grad(); loss.backward(); opt.step()

ids, h = [stoi["t"]], None
for _ in range(48):
    logits, h = model(torch.tensor([[ids[-1]]]), h)
    ids.append(torch.multinomial(logits[0, -1].softmax(0), 1).item())
print("".join(chars[i] for i in ids))

Output

the mat. the dog sat on the log. the cat sat on t

Summary

Gates let networks remember, and they powered translation and speech for twenty years. Aria has now written her first words. But recurrent networks read one word at a time, which is slow, and they still forget long passages. In twenty seventeen, a single paper changed everything. Next, attention.

Key takeaways

LSTM adds gates and a memory lane
GRU is a simpler cousin
Predict the next character, repeat

Next up: Attention — look at everything at once

❮ Tokens and Embeddings Attention ❯