The series so far built models that take a fixed-size input: a question’s bag of words, four routes out. Language isn’t fixed-size. A sentence can be three tokens or three thousand, and meaning depends on order: “the model trains the data” isn’t “the data trains the model”. This part of the series is about models that read sequences, and it starts with the design that dominated neural sequence modeling in the years before the Transformer: the recurrent neural network. The idea fits in one sentence: read one token at a time, and keep one fixed-size vector that summarizes everything read so far. Both halves of that sentence turn out to be problems.
- One at a time means the computation is serial, and the gradient has to travel back through every step, where it shrinks or explodes.
- One fixed-size vector means everything the model will ever know about the past has to fit in it.
This post builds an RNN by hand, measures both problems, and shows the two fixes that came before the Transformer: gates (LSTM, GRU) for the first, and attention for the second. Attention then turns out to be the whole answer, which is where The Transformer’s Three Attention Patterns picks up.
The two designs this post compares. An RNN squeezes everything it has read into one vector; attention lets the reader look back at every state.
- An RNN reads one token at a time into one fixed-size state. Backpropagation through time is Backpropagation, by Hand’s backprop through the unrolled steps.
- The gradient to an early step is multiplied once per step, so it vanishes or explodes geometrically: a vanilla RNN passed back of it over 50 steps.
- LSTM and GRU gates add to the state instead of overwriting it, giving the gradient a path with no matrix in it. With a forget bias of 3, an untrained LSTM passes back 0.46 of it over 50 steps.
- A fixed-size state is a bottleneck. Reversing 40 digits through one vector got at most 2% of sequences right; with attention, 93–95%.
- Transformers won on path length, memory and parallelism. On a 1,024-token sequence, attention did more arithmetic than a GRU layer and finished about six times sooner, 1.5 ms against 9.5 ms.
Table of contents
Open Table of contents
A state that carries the past
An Elman RNN (Elman, 1990) keeps a hidden state , a vector, and updates it at every step from the previous state and the new input :
The same weights , and are used at every step. That’s what lets one network read a sequence of any length, and it’s also where the trouble starts. A prediction at step reads , through an output layer like the router’s. Here are three steps with a 2-number state, every number shown, and the gradient computed backward through all three:
# An Elman RNN with a 2-number state, run for 3 steps:
# h_t = tanh(x_t W_x + h_{t-1} W_h + b), and a prediction y = h_3 w_out at the end.
W_x = np.array([[0.5, -0.3]]) # one input number per step
W_h = np.array([[0.8, 0.2], [-0.4, 0.6]]) # the same matrix at every step
b, w_out = np.array([0.1, 0.0]), np.array([1.0, -1.0])
xs = [1.0, 0.0, -1.0]
def forward(W_h):
h, hs = np.zeros(2), [np.zeros(2)]
for x in xs:
h = np.tanh(x * W_x[0] + h @ W_h + b)
hs.append(h)
return hs, hs[-1] @ w_out
hs, y = forward(W_h)
for t, h in enumerate(hs[1:], 1):
print(f"h_{t} = {np.round(h, 4)}")
print(f"prediction y = h_3 . w_out = {y:.4f}")
# Backprop through time: run "Backpropagation, by Hand"'s rules backward through the 3 steps.
# The same W_h was used at every step, so its gradient is the sum over steps.
dW_h, dh = np.zeros_like(W_h), w_out.copy() # dy/dh_3
for t in range(3, 0, -1):
dz = dh * (1 - hs[t] ** 2) # back through tanh
dW_h += np.outer(hs[t - 1], dz) # this step's share of the gradient
dh = dz @ W_h.T # pass back to h_{t-1}
print(f"gradient reaching h_{t - 1}: {np.round(dh, 4)}")
eps, numeric = 1e-6, np.zeros_like(W_h)
for i in np.ndindex(W_h.shape):
E = np.zeros_like(W_h); E[i] = eps
numeric[i] = (forward(W_h + E)[1] - forward(W_h - E)[1]) / (2 * eps)
print(f"dy/dW_h = {np.round(dW_h, 4).tolist()}; largest gap to finite differences "
f"{np.abs(dW_h - numeric).max():.1e}")
h_1 = [ 0.537 -0.2913]
h_2 = [ 0.5691 -0.0673]
h_3 = [0.082 0.357]
prediction y = h_3 . w_out = -0.2750
gradient reaching h_2: [ 0.6201 -0.9208]
gradient reaching h_1: [ 0.1521 -0.7177]
gradient reaching h_0: [-0.0448 -0.4374]
dy/dW_h = [[0.7904, -0.9888], [-0.189, 0.3257]]; largest gap to finite differences 4.6e-11
The forward pass is three applications of the same layer. The backward pass is Backpropagation, by Hand’s rules applied through the unrolled chain, from back to . This is backpropagation through time, and it has two properties worth noticing:
- The shared weight collects a gradient from every step. was used three times, so its gradient is the sum of three contributions (the first is zero, since ), the rule from PyTorch from the Inside’s autograd for a tensor used more than once.
- The gradient reaching earlier states is multiplied by and tanh’s slope once per step. It went from at to at in three steps. Over fifty, that product decides everything.
The vanishing gradient, measured
Initialization and Optimizers showed a signal passing through 20 layers shrinking or exploding by a factor per layer. An RNN is the same thing, sideways: steps is applications of the same layer, so the gradient reaching step 0 from step is multiplied by the same matrix times. If shrinks vectors, the gradient vanishes; if it stretches them, it explodes. For with every direction scaled by the same spectral radius :
# How much of the gradient at step 50 survives back to step 0?
# Backprop multiplies by W_h^T (times tanh's slope) once per step, 50 times.
rng = np.random.default_rng(0)
d = 64
Q, _ = np.linalg.qr(rng.normal(size=(d, d))) # a rotation: every direction the same size
def run(radius, act, input_size=0.5):
W = radius * Q # every singular value equals `radius`
h, hs = rng.normal(size=d) * 0.5, []
for _ in range(50): # forward, keeping the states
z = h @ W + rng.normal(size=d) * input_size # plus a fresh input each step
h = np.tanh(z) if act == "tanh" else z
hs.append(h)
g, sizes = np.ones(d) / np.sqrt(d), {}
for t in range(49, -1, -1): # backward from step 50
if act == "tanh":
g = g * (1 - hs[t] ** 2)
g = g @ W.T
if 50 - t in (1, 10, 25, 50):
sizes[50 - t] = np.linalg.norm(g)
return sizes
The script runs it at radii 0.9, 1.0, 1.1 and 2.0, linear and with tanh, and once more at 1.1 with no inputs.
radius 0.9, linear: gradient size after 1 steps 9.0e-01, 10 steps 3.5e-01, 25 steps 7.2e-02, 50 steps 5.2e-03
radius 0.9, tanh : gradient size after 1 steps 6.6e-01, 10 steps 1.6e-02, 25 steps 1.2e-04, 50 steps 1.6e-08
radius 1.0, linear: gradient size after 1 steps 1.0e+00, 10 steps 1.0e+00, 25 steps 1.0e+00, 50 steps 1.0e+00
radius 1.0, tanh : gradient size after 1 steps 7.4e-01, 10 steps 6.1e-02, 25 steps 1.1e-03, 50 steps 1.8e-06
radius 1.1, linear: gradient size after 1 steps 1.1e+00, 10 steps 2.6e+00, 25 steps 1.1e+01, 50 steps 1.2e+02
radius 1.1, tanh : gradient size after 1 steps 8.4e-01, 10 steps 1.8e-01, 25 steps 5.3e-03, 50 steps 2.8e-05
radius 2.0, tanh : gradient size after 1 steps 1.1e+00, 10 steps 2.9e+00, 25 steps 1.3e+01, 50 steps 1.1e+02
radius 1.1, tanh, no inputs: gradient size after 1 steps 1.0e+00, 10 steps 9.8e-01, 25 steps 1.0e+00, 50 steps 8.7e-01
- Linear, : , exactly as the product predicts. explodes to 120. Only holds steady.
- With tanh and real inputs it’s worse. tanh’s slope is at most 1, and inputs push units into its flat regions where the slope is near 0. So even vanishes, to in 50 steps, and , which exploded without tanh, now vanishes too. It’s the inputs that do it: the same with no inputs keeps 87% of the gradient.
- tanh doesn’t prevent explosion. At the tanh RNN explodes as well, a hundredfold over 50 steps.
(Here every direction is scaled by the same , so the spectral radius and the singular values coincide; in general it’s the largest singular value that matters: below 1 is enough for vanishing, above 1 is needed for explosion, as Pascanu, Mikolov and Bengio, 2013, show.) Bengio, Simard and Frasconi (1994) put the dilemma precisely: for an RNN to store a piece of information robustly against noise, its dynamics must be contracting, and contracting dynamics are exactly the ones whose gradients vanish over long distances. The model gets essentially no learning signal about how a word 50 steps back should have changed its state, so it never learns to use long-range context. Exploding gradients are the easier half to fix, by clipping the gradient’s norm (Initialization and Optimizers; Pascanu et al. introduced the standard version); vanishing ones need a different architecture.
Gates: LSTM and GRU
The LSTM (Hochreiter and Schmidhuber, 1997) adds a second state, the cell , and changes how it’s updated. Instead of rewriting the whole state every step through a squashing function, it adds to it, under the control of gates, sigmoids between 0 and 1 computed from the input and the previous state:
(each gate with its own weights). The original 1997 LSTM kept the cell with no forget gate at all, a fixed “constant error carousel”; the forget gate came three years later (Gers, Schmidhuber and Cummins, 2000).
The forget gate decides how much of the old cell to keep, the input gate how much of the new candidate to add, and the output gate how much of the cell to show as . The key is the first term. Along the cell path, the gradient from to is multiplied by and nothing else: no matrix, no tanh slope. (The other paths, through and the gates, still have both; the cell path is the one that can survive.) When is near 1, the gradient can survive many steps along it. (It’s the same idea as the residual connection of The Residual Stream: an update that adds instead of overwrites.) The GRU (Cho et al., 2014) is a lighter version with two gates and no separate cell.
The LSTM’s cell is a highway: along it, the gradient from to is multiplied by the forget gate and nothing else.
Here’s how much gradient reaches the input 50 steps back in untrained cells of the same width, at initialization. How much a gated cell “starts out remembering” is set by one bias, the LSTM’s forget gate or the GRU’s update gate, so both are shown with and without a head start:
# How much does the input at step 0 affect the state at step T? The size of
# d h_T / d x_0 for untrained cells of the same width, at initialization.
# The gates' biases decide how much a cell "starts out remembering": the
# LSTM's forget gate (PyTorch order: input, forget, cell, output) and the
# GRU's update gate (order: reset, update, new) keep the old state when near 1.
T, d = 50, 64
def gradient_back(cls, gate_bias=None):
sizes = []
for seed in range(5):
torch.manual_seed(seed)
cell = cls(d, d, batch_first=True)
if gate_bias is not None:
with torch.no_grad():
cell.bias_ih_l0[d:2 * d].fill_(gate_bias); cell.bias_hh_l0[d:2 * d].fill_(0.0)
x = torch.randn(1, T, d, requires_grad=True)
out, _ = cell(x)
out[0, -1].sum().backward()
sizes.append([x.grad[0, t].norm().item() for t in (T - 1, T - 10, T - 25, 0)])
return torch.tensor(sizes).median(0).values
It is run for a plain RNN, a GRU with its default and a raised update bias, and an LSTM with forget biases 0, 1 and 3.
RNN : gradient reaching the input 1 / 10 / 25 / 50 steps back: 3.5e+00 / 2.4e-03 / 1.4e-08 / 3.5e-17
GRU, default : gradient reaching the input 1 / 10 / 25 / 50 steps back: 2.2e+00 / 1.6e-02 / 1.3e-05 / 6.2e-11
GRU, update bias +1 : gradient reaching the input 1 / 10 / 25 / 50 steps back: 1.3e+00 / 9.9e-02 / 1.7e-03 / 3.5e-06
LSTM, forget bias 0 : gradient reaching the input 1 / 10 / 25 / 50 steps back: 1.0e+00 / 6.7e-03 / 2.1e-06 / 8.1e-12
LSTM, forget bias 1 : gradient reaching the input 1 / 10 / 25 / 50 steps back: 1.0e+00 / 8.6e-02 / 3.1e-03 / 1.8e-05
LSTM, forget bias 3 : gradient reaching the input 1 / 10 / 25 / 50 steps back: 8.9e-01 / 4.1e-01 / 4.3e-01 / 4.6e-01
How much gradient survives the trip back. A plain RNN loses it geometrically (left). A gated cell’s survival is set by one bias: with a forget bias of 3, almost half reaches 50 steps back (right).
Data: the two outputs above- The vanilla RNN passes back after 50 steps, effectively nothing.
- A gated cell’s head start is set by one bias. With a forget bias of 0, the LSTM is no better than the default GRU at 50 steps; with a bias of 1, it passes back , and with 3, almost half the gradient survives. A GRU given the matching update-gate bias gets , close to the LSTM with bias 1. The gap between cells is mostly in how they start, not which cell they are; the standard LSTM trick of a forget bias of 1 (from Gers et al., popularized by Jozefowicz et al., 2015) is exactly this head start.
- None of it is the trained behavior. These are untrained cells; training changes the gates. But they show the mechanism: a gate near 1 lets the gradient through, and one bias decides where the gate starts.
Gates don’t remove the decay, but they slow it enough that a model can learn from context dozens of steps away. After about fifteen years of development, LSTMs became the standard sequence model from about 2013, in speech recognition and then in translation, until the Transformer.
A character model of this blog
The running example gets its first language model: read the blog one
character at a time and predict the next. The training text is the blog’s
own English posts as of 2026-09-30, frozen into two files so every later
post can train on exactly the same text, with two whole posts held out for
validation so the model is scored on text it has never seen.
corpus.py builds them from the posts, removing frontmatter, imports and
component tags and keeping prose, math and code. The folder links the
frozen files from the companion repo’s shared data, so the script only
reports on them:
train 593,616 characters, validation 56,786; 147 distinct characters
validation characters never seen in training: 123 of 56,786
Before any network, a baseline: predict the next character from counts of what followed the previous characters in the training text.
import math
from collections import Counter
# Baselines with no neural network at all: predict the next character from
# the previous k characters by counting, with add-one smoothing.
train, val = open("blog_train.txt").read(), open("blog_val.txt").read()
V = len(set(train)) + 1 # + an "unseen" character
for k in (0, 1, 2, 3, 4):
ctx, pair = Counter(), Counter()
for i in range(k, len(train)):
ctx[train[i - k:i]] += 1
pair[train[i - k:i + 1]] += 1
nll = -sum(math.log((pair[val[i - k:i + 1]] + 1) / (ctx[val[i - k:i]] + V))
for i in range(k, len(val))) / (len(val) - k)
print(f"previous {k} characters: validation {nll:.3f} nats/char = {nll / math.log(2):.3f} bits/char")
previous 0 characters: validation 3.468 nats/char = 5.003 bits/char
previous 1 characters: validation 2.849 nats/char = 4.110 bits/char
previous 2 characters: validation 2.481 nats/char = 3.580 bits/char
previous 3 characters: validation 2.628 nats/char = 3.791 bits/char
previous 4 characters: validation 3.112 nats/char = 4.490 bits/char
Counting two characters back gives 3.58 bits per character. Counting
further makes it worse, because of how the counts are smoothed: most
4-character contexts in the validation posts never appeared in training,
and add-one smoothing spreads their probability evenly over 148
characters. A better smoother backs off to a shorter context instead.
Kneser–Ney smoothing (Kneser and Ney, 1995) does that, and also estimates
the shortest context by how many different contexts a character follows,
not how often it occurs. In the code, the discount D is 0.75, and V
counts the training characters plus one “unseen” character:
def build(order):
count = [Counter() for _ in range(order + 1)] # count[n][context + char], n = context length
for n in range(order + 1):
for i in range(n, len(train)):
count[n][train[i - n:i + 1]] += 1
# For lower orders, KN uses continuation counts: in how many distinct contexts does it appear?
cont = [Counter() for _ in range(order + 1)]
for n in range(order):
for gram in count[n + 1]:
cont[n][gram[1:]] += 1
tables = []
for n in range(order + 1):
c = count[n] if n == order else cont[n]
total, types = defaultdict(int), defaultdict(int)
for gram, k in c.items():
total[gram[:-1]] += k
types[gram[:-1]] += 1
tables.append((c, total, types))
return tables
def prob(tables, ctx, ch):
p = 1 / V # the uniform floor
for n in range(len(tables)):
c, total, types = tables[n]
h = ctx[len(ctx) - n:] if n else ""
if total[h] == 0:
continue # context never seen: keep the shorter estimate
p = max(c[h + ch] - D, 0) / total[h] + D * types[h] / total[h] * p
return p
The rest of the file scores context lengths 1 to 7 on the held-out posts, and checks that the probabilities after “the” sum to one.
previous 1 characters: 4.124 bits/char
previous 2 characters: 3.370 bits/char
previous 3 characters: 2.844 bits/char
previous 4 characters: 2.650 bits/char
previous 5 characters: 2.600 bits/char
previous 6 characters: 2.594 bits/char
previous 7 characters: 2.606 bits/char
probabilities over all characters sum to 1.000000 after 'the'
With good smoothing, counting keeps improving up to 5–6 characters of context, to 2.59 bits per character (choosing 6 by looking at these results, so it’s a hindsight number). That’s a strong baseline, and the recurrent models below have to be measured against it, not against add-one. What counts can’t do is use context beyond a few characters, or share what they learned between contexts that differ in one letter: that’s what a neural network’s vectors are for.
Now three recurrent models of the same width, trained the same way, three
seeds each. Each training window is 128 characters with a fresh state, so
the gradient never travels further back than that (truncated
backpropagation through time); at evaluation the state is carried across
windows. All of this post’s training runs on a laptop CPU: on an Apple M4,
this post’s scripts, which train about thirty small recurrent models,
take about 100 minutes. The model, the training windows and the
evaluation that carries the state along, from charlm.py:
class CharRNN(torch.nn.Module):
def __init__(self, cell, hidden=256):
super().__init__()
self.embed = torch.nn.Embedding(V, 64)
self.rnn = {"RNN": torch.nn.RNN, "GRU": torch.nn.GRU, "LSTM": torch.nn.LSTM}[cell](
64, hidden, batch_first=True)
self.out = torch.nn.Linear(hidden, V)
def forward(self, x, state=None):
h, state = self.rnn(self.embed(x), state)
return self.out(h), state
def batch(data, rng, B=64, T=128):
i = torch.randint(0, len(data) - T - 1, (B,), generator=rng)
x = torch.stack([data[j:j + T] for j in i])
return x, torch.stack([data[j + 1:j + T + 1] for j in i])
@torch.no_grad()
def evaluate(model, data, T=128):
"""Average loss per character over the whole text, carrying the state along."""
total, state = 0.0, None
for s in range(0, len(data) - 1 - T, T):
logits, state = model(data[s:s + T][None], state)
total += torch.nn.functional.cross_entropy(logits[0], data[s + 1:s + T + 1], reduction="sum").item()
return total / ((len(data) - 1 - T) // T * T)
The rest of charlm.py trains one model: 64 windows of 128 characters
per step, AdamW at a learning rate of , gradients clipped
at 1.0, for 3,000 steps unless told otherwise. It takes the cell and the
seed on the command line, and char_models.py runs it for each cell with
seeds 0, 1 and 2:
RNN seed 0, 3,000 steps: 129,940 parameters, 58 s; train 1.941, validation 2.839 bits/char
RNN seed 1, 3,000 steps: 129,940 parameters, 52 s; train 1.946, validation 2.829 bits/char
RNN seed 2, 3,000 steps: 129,940 parameters, 52 s; train 1.947, validation 2.832 bits/char
GRU seed 0, 3,000 steps: 294,804 parameters, 141 s; train 1.547, validation 2.827 bits/char
GRU seed 1, 3,000 steps: 294,804 parameters, 173 s; train 1.575, validation 2.842 bits/char
GRU seed 2, 3,000 steps: 294,804 parameters, 168 s; train 1.615, validation 2.812 bits/char
LSTM seed 0, 3,000 steps: 377,236 parameters, 263 s; train 1.714, validation 2.694 bits/char
LSTM seed 1, 3,000 steps: 377,236 parameters, 268 s; train 1.729, validation 2.696 bits/char
LSTM seed 2, 3,000 steps: 377,236 parameters, 257 s; train 1.781, validation 2.690 bits/char
The recurrent models against counting. Well-smoothed counts of the last six characters beat every recurrent model on this small blog.
Data: the outputs above- Every network beats add-one counting, and none beats Kneser–Ney. The RNN and GRU, at 2.81–2.84 bits per character, are 0.7 bits better than add-one’s best but 0.2 worse than Kneser–Ney’s 2.59.
- At 3,000 steps, the LSTM is ahead of the other two. All three seeds land at 2.69–2.70, against 2.81–2.84 for the RNN and GRU, a gap much larger than the spread between seeds. It also has nearly three times the RNN’s parameters (377,000 to 130,000), so “same width” isn’t “same size”.
- All three are overfitting. On training text they’re much better, 0.9 to 1.3 bits lower. Part of that gap is the held-out posts covering different topics; the next check shows the rest is overfitting. The GRU, the best fit to its training text, is no better than the plain RNN on new text. The blog is small: 600,000 characters, seen about 41 times each in 3,000 steps.
Is 3,000 steps a fair place to compare them? Trained four times as long,
one seed each, by char_models_long.py:
Full script: char_models_long.py
RNN seed 0, 12,000 steps: 129,940 parameters, 326 s; train 1.793, validation 3.017 bits/char
GRU seed 0, 12,000 steps: 294,804 parameters, 737 s; train 1.351, validation 3.182 bits/char
LSTM seed 0, 12,000 steps: 377,236 parameters, 761 s; train 1.211, validation 3.369 bits/char
Longer training, opposite directions. The model that fits its training text best, the LSTM, does worst on new text.
Data: the outputs above, seed 0Longer training makes every model worse on new text, and the LSTM worst of all, 3.37 bits per character, while its training loss drops to 1.21. So the 3,000-step numbers are a snapshot of models that are all past their best on this little text, not their ceilings; What “Generative” Means tracks the held-out loss during training and finds the character LSTM’s best point near 2,000 steps. Which model is best depends on when training stops: 3,000 steps is a step count fixed in advance, not each model’s best, and there’s no separate test post to choose a stopping point on without peeking. The lesson that holds up best is about data: on a corpus this size, a well-smoothed count model is as good as a small recurrent one, and the model that generalizes is the one that’s stopped in time.
The bottleneck, and attention
The second problem is the fixed-size state. It shows most clearly in sequence-to-sequence models (Sutskever, Vinyals and Le, 2014), which translate one sequence into another: an encoder RNN reads the input, and its last state is handed to a decoder RNN that writes the output. Everything the decoder knows about the input is in that one vector, however long the input was.
Attention (Bahdanau, Cho and Bengio, 2014; published at ICLR 2015) was invented to remove that bottleneck. At every output step, the decoder scores each of the encoder’s states against its own current state, turns the scores into weights with a softmax (Cross-Entropy and the Softmax Gradient), and reads a weighted average of the encoder states: a context vector chosen fresh for each output. Instead of remembering the whole input in one vector, the decoder looks back at it.
A test that isolates the bottleneck: reverse a sequence of random digits. There’s nothing to learn about language, only whether the input survives the trip:
class Seq2Seq(torch.nn.Module):
def __init__(self, attention):
super().__init__()
self.attention = attention
self.embed = torch.nn.Embedding(V, 32)
self.encoder = torch.nn.GRU(32, H, batch_first=True)
self.decoder = torch.nn.GRUCell(32 + (H if attention else 0), H)
self.W_dec, self.W_enc = torch.nn.Linear(H, H), torch.nn.Linear(H, H)
self.v = torch.nn.Linear(H, 1, bias=False)
self.out = torch.nn.Linear(H, 10)
def forward(self, src, tgt_in):
enc, h = self.encoder(self.embed(src)) # enc: every step; h: the last one
h, logits = h[0], []
for t in range(tgt_in.shape[1]):
x = self.embed(tgt_in[:, t])
if self.attention: # additive attention over enc
score = self.v(torch.tanh(self.W_dec(h)[:, None] + self.W_enc(enc)))[..., 0]
context = (torch.softmax(score, dim=1)[..., None] * enc).sum(1)
x = torch.cat([x, context], dim=1)
h = self.decoder(x, h)
logits.append(self.out(h))
return torch.stack(logits, 1)
The rest of seq2seq.py trains with Adam on fresh random batches of 64
and scores 1,000 new sequences by greedy decoding. Its arguments are the
length, attention or none, the seed, and optionally the number of steps
(2,000) and the learning rate (). reverse.py runs every
length, both models, two seeds:
length 5, no attention, seed 0: digits right 0.990, whole sequences right 0.972
length 5, no attention, seed 1: digits right 0.994, whole sequences right 0.979
length 5, attention , seed 0: digits right 0.998, whole sequences right 0.996
length 5, attention , seed 1: digits right 1.000, whole sequences right 1.000
length 10, no attention, seed 0: digits right 0.962, whole sequences right 0.832
length 10, no attention, seed 1: digits right 0.971, whole sequences right 0.853
length 10, attention , seed 0: digits right 0.987, whole sequences right 0.958
length 10, attention , seed 1: digits right 0.984, whole sequences right 0.958
length 20, no attention, seed 0: digits right 0.784, whole sequences right 0.148
length 20, no attention, seed 1: digits right 0.818, whole sequences right 0.189
length 20, attention , seed 0: digits right 0.938, whole sequences right 0.840
length 20, attention , seed 1: digits right 0.944, whole sequences right 0.869
length 40, no attention, seed 0: digits right 0.374, whole sequences right 0.000
length 40, no attention, seed 1: digits right 0.379, whole sequences right 0.000
length 40, attention , seed 0: digits right 0.479, whole sequences right 0.106
length 40, attention , seed 1: digits right 0.250, whole sequences right 0.000
At length 5 both models reverse almost every sequence. Without attention, accuracy then collapses as the input grows: 83–85% of whole 10-digit sequences right, 15–19% at 20 digits, none at 40. A rerun on an NVIDIA RTX 3090 (its saved outputs) got 88–92% and 18–22%, and none at 40. With attention it degrades far more slowly: 96–97% at 10 digits and 84–87% at 20, on both machines.
At 40 digits, training at this setup is unstable, and the reason is
worth a detour. On the M4 both models look broken. On the rerun the model
without attention is broken again, while one attention seed got 74% of
digits and 47% of whole sequences right and the other 35% and 1%. At 40
digits and 3,000 steps at the original learning rate (below), the
attention model does worse than the one without it on both machines,
0.17–0.20 of digits right against 0.48–0.51, a sign that the learning rate
that suits the short sequences is too large here, the instability of
Initialization and Optimizers. So compare both models at 40 digits and 3,000
steps, at the original learning rate and a smaller one (reverse_40_longer.py
and reverse_40_lr.py):
Full script: reverse_40_longer.py
length 40, no attention, seed 0, 3,000 steps: digits right 0.490, whole sequences right 0.000
length 40, no attention, seed 1, 3,000 steps: digits right 0.513, whole sequences right 0.000
length 40, attention , seed 0, 3,000 steps: digits right 0.195, whole sequences right 0.000
length 40, attention , seed 1, 3,000 steps: digits right 0.177, whole sequences right 0.000
length 40, no attention, seed 0, 3,000 steps, learning rate 0.001: digits right 0.224, whole sequences right 0.000
length 40, no attention, seed 1, 3,000 steps, learning rate 0.001: digits right 0.238, whole sequences right 0.000
length 40, attention , seed 0, 3,000 steps, learning rate 0.001: digits right 0.971, whole sequences right 0.927
length 40, attention , seed 1, 3,000 steps, learning rate 0.001: digits right 0.975, whole sequences right 0.940
and the model without attention, at its better rate, trained twice as long
(reverse_40_6000.py):
Full script: reverse_40_6000.py
length 40, no attention, seed 0, 6,000 steps: digits right 0.701, whole sequences right 0.006
length 40, no attention, seed 1, 6,000 steps: digits right 0.745, whole sequences right 0.015
The bottleneck, measured. Without attention, accuracy collapses as the input grows; with it, 93% of 40-digit sequences come out right. The 40-digit points use each model’s better setting: learning rate for attention, 6,000 steps without it.
Data: the runs above, means of two seeds(The test set is the same 1,000 random sequences for a given seed at every setting, and the better learning rate was chosen by looking at these results, so it’s a hindsight choice; with 1,000 sequences and a gap this size, that can’t explain the difference. Two seeds each.)
- Each model has a different best learning rate. At the original rate, the model without attention gets 48–51% of digits right, and the attention model diverges, to 17–20%. At the smaller rate the attention model trains cleanly and the other one learns more slowly, 22–24%.
- At its better setting, attention wins by a mile. 93–95% of whole 40-digit sequences right with attention; without it, none at 3,000 steps and at most 2% after 6,000.
- The model without attention is still improving. Its digit accuracy climbs from 48–51% at 3,000 steps to 70–75% at 6,000, and a 128-number state can in principle hold 40 digits. So this isn’t proof the task can’t be learned through one state. It’s a measure of how much harder it is: attention makes the task easy, not merely possible.
The digit accuracy tells the rest. The model without attention gets most
digits of a 40-digit sequence right but almost never all 40. Two things
compound. The last encoder state holds a blurred summary of the whole
input. And at test time the decoder is fed its own previous outputs
(greedy decoding) instead of the true ones it saw in training (teacher
forcing, the tgt_in in seq2seq.py), so one wrong digit makes the next
harder. Reversal is also the kindest case for the model without attention,
since its first output is the last digit it read; Sutskever et al.
reversed their source sentences for exactly that reason. With attention,
each output digit reads the one input digit it needs, directly. (This is a
simplified Bahdanau attention: the paper’s encoder is bidirectional, and
its output layer also reads the context.)
Why Transformers won
Attention fixed the bottleneck but kept the recurrence underneath: the encoder and decoder still read one step at a time. The Transformer (Vaswani et al., 2017) asked what happens if attention does all the work, every token attending to every other, with no recurrence at all. Three consequences decided the matter.
- Path length. In an RNN, information from token 1 reaches token 1,000 through 999 steps of the same layer, and the gradient comes back the same way, shrinking. With attention, any token reaches any other in one step per layer (Vaswani et al.’s Table 1 compares the two).
- Memory. An RNN’s knowledge of the past is one fixed vector. A Transformer keeps every previous token’s keys and values and reads them directly, which is also why it needs a KV cache (The KV Cache) and an RNN doesn’t.
- Parallelism. An RNN’s step needs step first. Attention processes all positions of a training sequence at once, as big matrix multiplications, exactly the kind of work How a GPU Runs Deep Learning’s GPU is built for:
# The same 1,024-token sequence at width 256: a GRU must take 1,024 steps one
# after another; self-attention does it as a few big matrix multiplies.
torch.manual_seed(0)
T, d = 1024, 256
x = torch.randn(1, T, d)
gru = torch.nn.GRU(d, d, batch_first=True)
Wq, Wk, Wv = (torch.randn(d, d) / d ** 0.5 for _ in range(3))
def attention(x):
q, k, v = x @ Wq, x @ Wk, x @ Wv
mask = torch.ones(T, T).triu(1).bool() # each token sees only the past
s = (q @ k.transpose(-1, -2) / d ** 0.5).masked_fill(mask, float("-inf"))
return torch.softmax(s, -1) @ v
def timed(fn, trials=10):
fn(); times = []
for _ in range(trials):
t0 = time.perf_counter(); fn(); times.append(time.perf_counter() - t0)
return sorted(times)[trials // 2]
with torch.no_grad():
t_rnn, t_att = timed(lambda: gru(x)), timed(lambda: attention(x))
flops_rnn = T * 3 * 2 * 2 * d * d # 3 gates, input and hidden matmuls
flops_att = 3 * 2 * T * d * d + 2 * 2 * T * T * d # projections, then scores and weighted sum
print(f"GRU: {T} sequential steps, {flops_rnn / 1e9:.2f} GFLOPs, {t_rnn * 1e3:6.1f} ms")
print(f"attention: 1 parallel pass, {flops_att / 1e9:.2f} GFLOPs, {t_att * 1e3:6.1f} ms")
GRU: 1024 sequential steps, 0.81 GFLOPs, 9.5 ms
attention: 1 parallel pass, 1.48 GFLOPs, 1.5 ms
Parallelism in one comparison: attention does more arithmetic and finishes six times sooner, because none of it waits.
Data: the measurement above(On this laptop’s CPU, a batch of one sequence, and one attention sublayer against one GRU layer.) The GRU does fewer FLOPs, 0.81 billion against 1.48, and takes six times as long: 9.5 ms against 1.5 ms. Its 1,024 steps each wait for the last, and each is a matrix-vector product too small to keep the hardware busy, the small-operation problem of How a GPU Runs Deep Learning. (Only the part is truly serial; the input side can be computed for all steps at once.) Attention does more arithmetic, all of it in a few large matrix multiplications that run at full speed. On a GPU we’d expect the gap to grow, since each of the 1,024 steps pays a kernel launch.
The trade-off didn’t disappear, it moved. Attention’s cost grows with the square of the sequence length, and generating text is still one token at a time. Each new token costs an RNN the same, a fixed-size state update; attention with a KV cache has to read every earlier token’s keys and values, so each step costs more as the text grows. Some of the newest architectures, linear attention (Katharopoulos et al., 2020, “Transformers are RNNs”) and state-space models such as Mamba (Gu and Dao, 2023), are RNNs again: a fixed-size state updated per token, but designed so training can run in parallel. Attention Alternatives covers them, and they make much more sense after this one.
Where this leads
- Embeddings and Contrastive Learning turns tokens into vectors: embeddings, and the contrastive losses that make “similar” mean “close”.
- What “Generative” Means is the language-model view of this post’s character model: autoregressive generation, teacher forcing, and perplexity.
- The Transformer’s Three Attention Patterns builds attention properly, and Build a Tiny GPT, End to End, and Train It on This Blog trains a tiny GPT on the same blog text, against the same validation posts, so its number can be compared directly with the RNNs here.
The assistant now has a language model of its own blog: a character-level LSTM, 2.69 bits per character at 3,000 steps on two posts it never saw. It’s small and serial, it was trained on 128-character windows, and a well-smoothed 6-character count model (2.59) still beats it. Those two numbers are the baselines the tiny GPT of Build a Tiny GPT, End to End, and Train It on This Blog has to beat, on the same validation text.
Takeaways
- An RNN reads one token at a time and keeps one fixed-size state, updated with the same weights every step. Backpropagation through time is Backpropagation, by Hand’s backprop through the unrolled steps; shared weights sum their gradients.
- The gradient to an early step is multiplied by the recurrent matrix and the activation’s slope once per step, so it vanishes or explodes geometrically. With tanh and real inputs, a vanilla RNN passed back of the gradient over 50 steps; tanh doesn’t stop explosion when the recurrent weights are large enough.
- LSTM and GRU gates add to the state instead of overwriting it, giving the gradient a path along the cell with no matrix in it. How much gets through at the start is set by one gate bias; with the same head start, the GRU and LSTM are close.
- Add-one n-gram counts stop improving after two characters; Kneser–Ney counts keep improving to six (2.59 bits per character) and beat every recurrent model here. On 600,000 characters, all of them overfit, and which is best depends on when training stops.
- A fixed-size state is a bottleneck: reversing 40 digits through one encoder state, a model got at most 2% of sequences right; with attention over the encoder’s states, 93–95%.
- Transformers won on path length, memory and parallelism. The serial trade-off comes back in generation and in state-space models.
References
Start with Karpathy’s The Unreasonable Effectiveness of Recurrent Neural Networks for intuition, then Olah’s Understanding LSTM Networks for the cell diagrams. Bahdanau et al. is the paper that introduced attention, and the step from it to Vaswani et al. is short.
- Elman, Finding Structure in Time (Cognitive Science, 1990): the simple recurrent network.
- Bengio, Simard and Frasconi, Learning Long-Term Dependencies with Gradient Descent is Difficult (IEEE Transactions on Neural Networks, 1994): the vanishing-gradient result.
- Pascanu, Mikolov and Bengio, On the Difficulty of Training Recurrent Neural Networks (ICML, 2013): singular-value conditions for vanishing and exploding gradients, and gradient clipping.
- Hochreiter and Schmidhuber, Long Short-Term Memory (Neural Computation, 1997): the LSTM.
- Gers, Schmidhuber and Cummins, Learning to Forget: Continual Prediction with LSTM (Neural Computation, 2000): the forget gate.
- Jozefowicz, Zaremba and Sutskever, An Empirical Exploration of Recurrent Network Architectures (ICML, 2015): among other things, the forget-gate bias of 1.
- Kneser and Ney, Improved Backing-Off for M-gram Language Modeling (ICASSP, 1995): Kneser–Ney smoothing.
- Cho et al., Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation (EMNLP, 2014): the GRU and the encoder–decoder.
- Sutskever, Vinyals and Le, Sequence to Sequence Learning with Neural Networks (NeurIPS, 2014): sequence-to-sequence with LSTMs.
- Bahdanau, Cho and Bengio, Neural Machine Translation by Jointly Learning to Align and Translate (ICLR, 2015): attention over encoder states.
- Vaswani et al., Attention Is All You Need (NeurIPS, 2017): the Transformer.
- Katharopoulos, Vyas, Pappas and Fleuret, Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (ICML, 2020): linear attention as a recurrence.
- Gu and Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023): a state-space model trained in parallel.