CourseTransformers: From Attention to Trained, Served and Interpreted Models · Module 1: The Road to Attention · part 8 of 111
Part 8 · Module 1: The Road to Attention

Lesson 2: Why recurrence needed replacing, measured

18 min read·13 Oct 2026
Module 1 · Lesson 2 · About 45 minutes · Needs: Lesson 1 (tensors, softmax, the chain rule)

The situation

Sol 400 at Aurora Station. You have just started tinyformer-lab, and before building a transformer you want proof that you need one. The station's older tools used recurrent networks, which read text one token at a time. Kofi Mensah, who reviews the lab's claims, asks: "What exactly is wrong with them?"

You answer with an experiment. You race a GRU, a recurrent network, against a small transformer on the shared copy task, then measure why the loser loses. The result becomes the lab's first report.

A real case, first

In 2018 Shaojie Bai, J. Zico Kolter and Vladlen Koltun published "An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling" (arXiv:1803.01271, revised April 2018, code at github.com/locuslab/TCN). They pitted a plain convolutional network (a TCN) against LSTMs, GRUs and plain recurrent networks.

One test was copy memory: ten digits, a long gap, then repeat the digits. Describing their Figure 4, they write that LSTMs and GRUs "simply converge to the same loss as predicting all zeros". In their Table 1, with a gap of about 1,000 steps, the recurrent models end near 0.02 and the TCN at 0.000035.

COPY MEMORY, T = 1000: TEST LOSS (LOWER IS BETTER) LSTM 0.0204 GRU 0.0197 plain RNN 0.0202 TCN (conv) 0.000035 Bai, Kolter, Koltun 2018, Table 1 (cited)
In Bai, Kolter and Koltun's copy memory test the three recurrent networks (red, failed) end at about the loss of guessing zeros everywhere, while the convolutional network (green) almost solves it.

They concluded that the "infinite memory" advantage of recurrent networks "is largely absent in practice". This lesson reproduces that shape on your copy task, then finds the cause.

What you will be able to do

  • Explain why a recurrent network must run its steps in order, and time the cost.
  • Race a GRU against a transformer of the same size and read the result honestly.
  • Measure vanishing and exploding gradients as a product of factors.
  • Say what LSTM and GRU gates fix, and what they leave broken.
  • Compute path length for recurrent, convolutional and attention layers.

1. Recurrence and the sequential dependency (Essential)

Every argument in this lesson goes back to one equation, and it explains why recurrent models were slow to train. Paste the torch setup cell and shared blocks 1 and 9 first.

python
import torch
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"   # every lesson runs on "cpu"; a GPU is optional
torch.manual_seed(11)                                     # the course seed (see course-guide/06)
print("torch", torch.__version__, "| device", DEVICE, "| CPU threads", torch.get_num_threads())

Code walkthrough

  • The course's torch setup cell: PyTorch, the CPU, seed 11. It printed torch 2.14.1 | device cpu | CPU threads 5 on the author's Mac.
python
import random
import torch

COURSE_DATE = "2026-10-09"   # versions, model names and prices checked on this day
SEED = 11                    # every generator in the course starts from this seed
SNAPSHOT_SOL = 400           # the corpus is the station record up to sol 399

def set_seed(seed=SEED):
    """Seed Python and PyTorch so a cell gives the same numbers every run."""
    random.seed(seed)
    torch.manual_seed(seed)

set_seed()
print("torch", torch.__version__, "| threads", torch.get_num_threads())

Code walkthrough

  • Shared block 1, as written: SEED = 11 and set_seed().
python
from typing import NamedTuple

class TaskData(NamedTuple):
    name: str
    vocab_size: int
    inputs: torch.Tensor    # (n, T) token ids
    targets: torch.Tensor   # (n, T) the next token at every position
    mask: torch.Tensor      # (n, T) 1 where the answer is scored, 0 elsewhere

def _shift(name, vocab, seq, scored_from):
    inputs, targets = seq[:, :-1], seq[:, 1:]
    mask = torch.zeros_like(inputs)
    mask[:, scored_from:] = 1
    return TaskData(name, vocab, inputs, targets, mask)

def make_task(name, n=4096, length=16, vocab=32, seed=SEED):
    """copy, reverse or induction. Token 0 is the separator; content tokens are 1 to vocab - 1."""
    g = torch.Generator().manual_seed(seed)
    body = torch.randint(1, vocab, (n, length), generator=g)
    sep = torch.zeros(n, 1, dtype=torch.long)
    if name == "copy":
        seq = torch.cat([body, sep, body], dim=1)
    elif name == "reverse":
        seq = torch.cat([body, sep, body.flip(1)], dim=1)
    elif name == "induction":                       # random tokens, then the same random tokens again
        seq = torch.cat([body, body], dim=1)
    return _shift(name, vocab, seq, scored_from=length)
TASKS = {name: make_task(name) for name in ["copy", "reverse", "induction"]}
for t in TASKS.values():
    print(t.name, tuple(t.inputs.shape), "scored positions per row:", int(t.mask[0].sum()))

Code walkthrough

  • Shared block 9, as written. Copy rows are 16 random tokens, the separator 0, then the same 16; _shift makes next-token targets and a mask scoring only the second half. Output: copy (4096, 32), 16 scored positions per row.
COPY TASK, ROW 0: 33 TOKENS 7 19 19 22 13 17 11 29 ... 0 separator 7 19 19 22 13 17 11 16 random tokens: not scored same 16 again: scored every answer = the token 16 positions back
The running example is row 0 of the shared copy task: grey marks the separator, amber marks the scored half, and every scored answer is the token exactly 16 positions earlier.

This is the row you follow all lesson. The first 16 tokens are random, so they are not scored. After the separator each answer is fully determined: the token exactly 16 positions back. A model passes only if it can reach that far.

python
copy = TASKS["copy"]
print("input :", copy.inputs[0].tolist())
print("target:", copy.targets[0].tolist())
print("scored:", copy.mask[0].tolist())

# every scored answer is the input token exactly 16 positions earlier
same = (copy.targets[:, 16:] == copy.inputs[:, :16]).all().item()
print("answer = input 16 steps back, in every row:", same)

Code walkthrough

  • Row 0 prints [7, 19, 19, 22, ...], 0, then the same tokens; targets are the row moved one place left; the mask is 16 zeros then 16 ones.
  • The comparison checks all 4,096 rows and prints True.
A RECURRENT NETWORK READS ONE TOKEN PER STEP h1 7 h2 19 h3 19 h4 22 ... h17 0 one state vector carries everything read so far
A recurrent network reads the copy row one token per step and passes a single state vector (blue, the normal data path) from each step to the next.

A recurrent network reads the row one token per step, keeping a running summary called the hidden state: h1 after token 1, h2 after token 2, up to h17 at the separator. Whatever it needs later must survive inside that one vector.

token x_t old state h_t-1 x_t W_x h_t-1 W_h h_t = tanh(sum) must wait for h_t-1
One recurrent step adds the token's part to the old state's part, each made with learned weights (amber), and squashes the sum with tanh; it cannot start until the old state exists.

One step is one line: h_t = tanh(x_t W_x + h_t-1 W_h). The second term is the catch: step t cannot start until step t - 1 has finished. This is the sequential dependency, and no clever code removes it, because the equation says so.

python
import time
torch.manual_seed(SEED)
B, T, d = 8, 512, 64                       # 8 sequences, 512 positions, 64 numbers per state
W_x = torch.randn(d, d) / d ** 0.5         # input weights
W_h = torch.randn(d, d) / d ** 0.5         # recurrent weights: state to state
x = torch.randn(B, T, d)                   # stand-in input vectors

def rnn_forward(x):
    h = torch.zeros(B, d)                  # h_0: the empty state
    for t in range(T):                     # step t needs h from step t - 1
        h = torch.tanh(x[:, t] @ W_x + h @ W_h)
    return h

def same_maths_no_dependency(x):           # same multiplications, no step waits for another
    flat = x.reshape(B * T, d)
    return torch.tanh(flat @ W_x + flat @ W_h)

def ms(f, runs=20):
    f(x); t0 = time.perf_counter()
    for _ in range(runs):
        f(x)
    return 1000 * (time.perf_counter() - t0) / runs

print(f"one step at a time: {ms(rnn_forward):.2f} ms   all at once: {ms(same_maths_no_dependency):.2f} ms")

Code walkthrough

  • W_x and W_h are random weights, scaled by one over the square root of the width to keep numbers moderate.
  • rnn_forward is the recurrence: a loop of 512 steps, each feeding h into the next.
  • same_maths_no_dependency does the same multiplications but gives each position its own input instead of the previous state. It is not a model; it shows the arithmetic with no waiting.
  • ms averages 20 runs, in milliseconds.
SAME MULTIPLICATIONS, 512 POSITIONS, BATCH 8 one step at a time 2.27 ms all at once 0.20 ms measured here, Apple silicon Mac, 5 threads
Doing the same multiplications one dependent step at a time took 2.27 ms against 0.20 ms all at once on the author's Mac (8 to 11 times across runs): the waiting, not the arithmetic, is the cost.

On the author's Mac the loop took 2.27 ms and the same arithmetic in one go 0.20 ms, 8 to 11 times faster across runs (measured here; Colab differs, the ratio is the point). Hardware is fastest on a few large multiplications, not 512 tiny ones that wait in line.

STEPS THAT MUST RUN IN ORDER, ONE COPY ROW GRU: 32 steps, one after another attention layer: 1 step, all 32 positions together
For one 32-token copy row a GRU layer must run 32 steps in order, while an attention layer handles all 32 positions in one step, so extra hardware can help it.

A GRU layer runs 32 ordered steps per copy row. An attention layer, built in Module 2, compares all 32 positions in one large multiplication: one ordered step per layer at any length. Vaswani and colleagues (2017) put it plainly: recurrence "precludes parallelization within training examples".

2. GRU against transformer on one task (Essential)

Speed is not what matters most. Kofi's real question is whether a recurrent network can learn the copy task at all, on a fair budget.

GRU embed 128, one GRU layer 107,296 parameters transformer width 64, 2 layers, 4 heads 106,144 parameters same data, optimiser, steps and seed
The race is fair on size: a one-layer GRU with 107,296 parameters against a two-layer transformer with 106,144, trained on identical batches.

The GRU (gated recurrent unit) is the recurrent racer; section 4 explains its gates. Its rival is a two-layer transformer from PyTorch's ready-made layer, which you rebuild yourself in Modules 2 to 5. Both have about 107,000 parameters.

python
import torch.nn as nn
import torch.nn.functional as F

class GRUModel(nn.Module):
    def __init__(self, vocab=32, d=128):
        super().__init__()
        self.embed = nn.Embedding(vocab, d)
        self.gru = nn.GRU(d, d, batch_first=True)          # one gated recurrent layer
        self.head = nn.Linear(d, vocab)

    def forward(self, ids):
        return self.head(self.gru(self.embed(ids))[0])

class AttentionModel(nn.Module):
    def __init__(self, vocab=32, d=64, layers=2, heads=4, context=32):
        super().__init__()
        self.embed, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)
        layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.0, batch_first=True, norm_first=True)
        self.blocks = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
        self.head = nn.Linear(d, vocab)

    def forward(self, ids):
        T = ids.shape[1]
        x = self.embed(ids) + self.pos(torch.arange(T))
        causal = nn.Transformer.generate_square_subsequent_mask(T)   # no looking ahead
        return self.head(self.blocks(x, mask=causal, is_causal=True))

count = lambda m: sum(p.numel() for p in m.parameters())
print(f"GRU {count(GRUModel()):,} parameters | transformer {count(AttentionModel()):,} parameters")

Code walkthrough

  • GRUModel: an embedding gives each token 128 numbers, nn.GRU runs the recurrence, head scores every possible next token. [0] keeps per-position outputs.
  • AttentionModel adds learned position vectors (self.pos), because attention alone does not know order, then stacks two nn.TransformerEncoderLayer blocks.
  • generate_square_subsequent_mask builds the causal mask: each position sees only itself and earlier positions.
  • Output: GRU 107,296 parameters | transformer 106,144 parameters.
python
N_VALID = 512                                        # last 512 rows are held out

def scored_loss(model, x, y, mask):
    losses = F.cross_entropy(model(x).flatten(0, 1), y.flatten(), reduction="none")
    return (losses * mask.flatten()).sum() / mask.sum()

@torch.no_grad()
def accuracy(model, task, by_position=False):
    model.eval()
    x, y, m = (t[-N_VALID:] for t in (task.inputs, task.targets, task.mask))
    right = (model(x).argmax(-1) == y).float()
    model.train()
    return right[:, 16:].mean(0) if by_position else (right * m).sum().item() / m.sum().item()

def race(model, task, steps=1000, every=50, lr=3e-3):
    set_seed(); g = torch.Generator().manual_seed(SEED)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    curve = []
    for step in range(1, steps + 1):
        i = torch.randint(0, len(task.inputs) - N_VALID, (64,), generator=g)
        loss = scored_loss(model, task.inputs[i], task.targets[i], task.mask[i])
        opt.zero_grad(); loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
        if step % every == 0:
            curve.append((step, accuracy(model, task)))
    return curve

Code walkthrough

  • The last 512 rows (N_VALID) are held out for measuring accuracy.
  • scored_loss is cross-entropy multiplied by the mask, so only the copied half counts.
  • accuracy compares the top-scoring token with the target on the scored half; by_position=True returns one accuracy per copied token.
  • race trains with AdamW for 1,000 steps of 64 rows from a seeded generator, so both models see identical batches; clip_grad_norm_ caps the gradient size (section 3). It logs accuracy every 50 steps.
python
results = {}
for name, make in [("GRU", GRUModel), ("transformer", AttentionModel)]:
    set_seed()
    model = make()
    t0 = time.perf_counter()
    curve = race(model, copy)
    secs = time.perf_counter() - t0
    first = next((s for s, a in curve if a >= 0.95), None)
    results[name] = (model, curve)
    print(f"{name:12} acc@100 {curve[1][1]:.3f}  acc@500 {curve[9][1]:.3f}  acc@1000 {curve[-1][1]:.3f}"
          f"  first >= 0.95 at step {first}  ({secs:.0f} s)")
print("chance:", round(1 / 31, 3))
for name, (model, _) in results.items():
    print(name, [round(a, 2) for a in accuracy(model, copy, by_position=True).tolist()])

Code walkthrough

  • Each model is built from the same seed, raced and timed; first is the first logged step at 0.95 accuracy. The last loop prints accuracy per copied position.
  • Training took 6 s (GRU) and 5 s (transformer) on the author's Mac; allow up to four times that on a free Colab CPU.
HELD-OUT ACCURACY ON THE SCORED HALF 0.0 0.5 1.0 0 500 1000 transformer 1.000 GRU 0.568 training step (chance 0.032, dashed)
On held-out rows the transformer (green) reached 0.998 by step 50 and 1.000 after, while the GRU (blue) crept to 0.568 in 1,000 steps (measured here).

The transformer reached 0.998 at step 50, its first check, and 1.000 after. The GRU reached 0.104 at step 100, 0.464 at step 500 and 0.568 at step 1,000, still creeping up; chance is about 0.032. This is one seed, and Module 5 adds a noise floor, but a gap of 0.43 over 8,192 scored answers is far beyond run-to-run wobble.

GRU ACCURACY AT EACH SCORED POSITION 0.71 0.44 0.71 transformer 1.00 1st 16th copied token
The GRU copies the first and last tokens best (0.71) and sags to 0.44 near the thirteenth, although every answer is 16 positions back; the transformer is perfect everywhere.

Every answer is the same 16 positions back, yet the GRU's accuracy is not flat: 0.71 on the first and last copied tokens, 0.44 around the thirteenth. Its one state must hold all 16 tokens while new ones arrive, and it holds some better than others. The transformer just looks 16 positions back.

3. Vanishing gradients, measured (Essential)

Why does the GRU learn so slowly? Learning is driven by gradients, and if almost none reaches the early tokens, the model cannot learn to keep them.

THE PROBE: LOSS ON THE LAST POSITION ONLY pos 0 pos 1 pos 2 ... pos T-1 loss how much gradient reaches position 0?
To measure vanishing gradients, put the loss (red) on the last position only and read how much gradient arrives back at each input position, starting with position 0.

The probe is small and untrained: an embedding, one layer that mixes positions, a head. Put the loss on the last position, backpropagate, and read the gradient arriving at each input. Position 0's number is all the signal a long-range dependency gets.

python
class Probe(nn.Module):
    """Embedding, ONE layer that mixes positions, a head. Never trained."""
    def __init__(self, kind, d=64, vocab=32):
        super().__init__()
        self.kind, self.embed, self.head = kind, nn.Embedding(vocab, d), nn.Linear(d, vocab)
        if kind == "attention":
            self.mix = nn.TransformerEncoderLayer(d, 4, 4 * d, dropout=0.0, batch_first=True, norm_first=True)
        else:
            self.mix = {"rnn": nn.RNN, "lstm": nn.LSTM, "gru": nn.GRU}[kind](d, d, batch_first=True)

    def forward(self, ids):
        self.x = self.embed(ids)
        self.x.retain_grad()                       # keep the gradient that arrives at each input position
        if self.kind == "attention":
            mask = nn.Transformer.generate_square_subsequent_mask(ids.shape[1], dtype=self.x.dtype)
            return self.head(self.mix(self.x, src_mask=mask, is_causal=True))
        return self.head(self.mix(self.x)[0])

def grad_by_position(model, T, dtype=torch.float64):
    model = model.to(dtype)
    g = torch.Generator().manual_seed(SEED)
    ids, target = torch.randint(1, 32, (16, T), generator=g), torch.randint(1, 32, (16,), generator=g)
    F.cross_entropy(model(ids)[:, -1], target).backward()     # loss on the LAST position only
    return model.x.grad.norm(dim=-1).mean(0)                   # one number per position

Code walkthrough

  • Probe picks its mixing layer by name. self.x.retain_grad() keeps the gradient of the embedded inputs, which PyTorch normally discards.
  • grad_by_position switches to float64 so tiny numbers stay visible, puts the loss on the last position, and returns the gradient size (.norm) per position, averaged over 16 rows.
BACKPROPAGATION MULTIPLIES ONE FACTOR PER STEP x 0.49 x 0.49 x 0.49 x 0.49 x 0.49 ... 255 0.49 x 0.49 x ... (255 times) = about 1e-80 measured: 5.7e-82 at T = 256 (tanh RNN)
Each step back through a tanh recurrence multiplies the gradient by a factor (amber), about 0.49 on average here, so 255 steps shrink it to around 1e-80.

Predict first. Backpropagating through one step multiplies the gradient by that step's Jacobian, the table of how each output number changes with each input number. Going back 255 steps multiplies 255 of them. If each roughly halves the gradient, the result is about 0.5 to the power 255, or 1.7e-77 (arithmetic).

python
KINDS = ["rnn", "lstm", "gru", "attention"]
def first_position(kind, T, dtype=torch.float64):
    torch.manual_seed(SEED)
    return grad_by_position(Probe(kind), T, dtype)[0].item()

print("    T  " + "".join(f"{k:>12}" for k in KINDS))
for T in [16, 32, 64, 128, 256]:
    print(f"{T:5d}  " + "".join(f"{first_position(k, T):12.1e}" for k in KINDS))
print("float32, T = 128:", [first_position(k, 128, torch.float32) for k in KINDS])
print("smallest normal float32:", torch.finfo(torch.float32).tiny)

torch.manual_seed(SEED)
g = grad_by_position(Probe("rnn"), 256)
print(f"tanh RNN: average shrink per step back = {(g[0] / g[-1]).item() ** (1 / 255):.3f}")

Code walkthrough

  • first_position returns a fresh probe's gradient at position 0; the loop prints one row per length.
  • The float32 line repeats length 128 in normal training precision; finfo(...).tiny is float32's smallest normal number.
  • The last line is the 255th root of the first-to-last gradient ratio: the average shrink per step.
GRADIENT AT POSITION 0 (LOG SCALE) 1e0 1e-40 1e-80 16 32 64 128 256 attention LSTM, GRU tanh RNN sequence length T
From length 16 to 256 the gradient reaching position 0 falls to 5.7e-82 for a tanh RNN (red) and about 1e-56 for LSTM and GRU, but only to 7.8e-5 for attention (green).

The prediction holds. For the tanh RNN the gradient at position 0 falls from 1.8e-07 at length 16 to 5.7e-82 at 256, a shrink of 0.488 per step. LSTM and GRU end near 1e-56. Attention falls only from 1.1e-03 to 7.8e-05, roughly as one over the length, since every position is one hop from the loss. This is the vanishing gradient problem that Pascanu, Mikolov and Bengio analysed in 2012.

float64, T = 128 RNN 2.7e-42 GRU 1.1e-29 float32, T = 128 RNN 0.0, LSTM 0.0, GRU 0.0 attention 1.4e-4 float32 cannot hold numbers below about 1.2e-38
In float32, the precision training normally uses, the recurrent gradients at length 128 are exactly zero (red): the first token gets no learning signal at all.

In float32 it is worse: at length 128 the three recurrent gradients print exactly 0.0, below float32's smallest normal number, about 1.2e-38. The first token gets no signal. The probes are untrained, and training can shape a GRU to keep what it needs, but this is where it starts.

Break it: make the recurrent weights bigger

If the factors are below 1 the product vanishes. Make them bigger.

python
for scale in [1, 2, 4]:
    torch.manual_seed(SEED)
    probe = Probe("rnn")
    with torch.no_grad():
        probe.mix.weight_hh_l0.mul_(scale)            # make the state-to-state weights bigger
    top = torch.linalg.matrix_norm(probe.mix.weight_hh_l0.double(), ord=2).item()
    g = grad_by_position(probe, 64)
    print(f"W_h x{scale}  largest stretch {top:.2f}  gradient at position 0: {g[0]:.1e}")

Code walkthrough

  • weight_hh_l0 is the RNN's W_h; mul_(scale) scales it in place without recording a gradient.
  • matrix_norm(..., ord=2) is the most W_h can stretch any vector, its largest singular value.
GRADIENT AT POSITION 0, T = 64 (LOG SCALE) W_h x1 2.6e-22 W_h x2 2.3e-07 W_h x4 6.7e+02 1
Break it: multiply the recurrent weights by 4 and the same product of factors grows instead of shrinking, reaching 670 (red, failure) at position 0; there is no safe middle setting.

At normal size (stretch 1.08) the gradient is 2.6e-22; doubled, 2.3e-07; times 4, it is 6.7e+02. That is the exploding gradient, where one update can wreck the weights. Pascanu and colleagues proposed clipping its size, which is why race calls clip_grad_norm_. Clipping does nothing for vanishing.

4. LSTM and GRU gating as partial fixes (Recommended)

LSTMs and GRUs were built to fight this decay. Knowing what their gates fix tells you when one is still a fair choice.

c_t-1 x f_t + c_t i_t x candidate forget gate, 0 to 1 gradient per step = f_t
An LSTM keeps a separate carry (blue) updated as c_t = f_t x c_t-1 + new input, so the gradient along it is multiplied by the forget gate (amber) each step, not by a weight matrix.

An LSTM keeps a second vector, the cell state: c_t = f_t x c_t-1 + i_t x candidate. The forget gate f_t is a number between 0 and 1 per slot. Along this carry path the gradient per step is just f_t, with no weight matrix and no tanh. A GRU does the same with one update gate and no separate cell.

python
def lstm_first_position(forget_bias, T):
    torch.manual_seed(SEED)
    probe = Probe("lstm")
    with torch.no_grad():                              # PyTorch packs gates as input, forget, cell, output
        probe.mix.bias_ih_l0[64:128].fill_(forget_bias)
        probe.mix.bias_hh_l0[64:128].zero_()
    return grad_by_position(probe, T)[0].item()

for T in [64, 256]:
    print(f"T = {T:3d}  forget bias 0: {lstm_first_position(0.0, T):.1e}"
          f"   forget bias +2: {lstm_first_position(2.0, T):.1e}")
gate = torch.sigmoid(torch.tensor([0.0, 2.0], dtype=torch.float64))
print("gate values", gate.tolist(), " gate ** 255:", (gate ** 255).tolist())

Code walkthrough

  • PyTorch packs an LSTM's four gate biases in the order input, forget, cell, output, so slots 64 to 127 are the forget gate. fill_ sets it; zero_() clears the second bias.
  • torch.sigmoid turns a bias into a gate value (0 gives 0.5, +2 about 0.88); the 255th power is the arithmetic for 255 steps.
forget bias 0, gate 0.5 T = 64: 3.8e-17 T = 256: 3.0e-58 forget bias +2, gate 0.88 T = 64: 5.7e-05 T = 256: 1.4e-09 a huge gain, and still a shrinking product
Opening the LSTM's forget gate with a +2 bias raised the gradient at position 0 from 3.0e-58 to 1.4e-9 at length 256 (measured here): a partial fix, still decaying.

With the default bias the gradient at length 256 is 3.0e-58; with +2 it is 1.4e-09, about 49 orders of magnitude better from one line. The arithmetic agrees: 0.5 to the 255th is 1.7e-77, 0.88 to the 255th is 8.8e-15. Better, and still a shrinking product.

16 TOKENS, ONE STATE 7 19 19 22 13 17 11 29 ... 8 more one 128-number state overwritten at every step
Gating does not remove the bottleneck: all 16 tokens must be squeezed into one 128-number state that every new step partly overwrites.

The second limit is invisible in the gradient table. The whole past must fit in one fixed-size state that every step partly overwrites; the GRU had 128 numbers for 16 tokens and reached 0.568. And gating leaves the sequential dependency untouched.

5. Path length and the convolutional alternative (Recommended)

Bai and colleagues' winner was convolutional, not attention. Comparing the two shows what attention actually buys.

POSITION 0 TO POSITION 16 recurrent 16 hops dilated conv 5 layers attention 1 hop
The copy answer at position 16 needs position 0: a recurrence (blue) gets there in 16 hops, dilated convolutions (grey) in 5 layers, attention (amber) in one direct hop.

Path length is the number of steps a signal travels between two positions. For the copy answer at position 16, a recurrence passes 16 states. A convolution mixes a small window of neighbours with no waiting; with dilation each layer skips further (1, 2, 4, 8, ...), so 5 layers cover 17 positions. Attention takes one hop.

python
def conv_stack(dilations, k=2, d=8):
    torch.manual_seed(SEED)
    return [nn.Conv1d(d, d, k, dilation=r) for r in dilations]

def run(layers, x):                         # causal: pad on the left only, so no output sees the future
    for conv in layers:
        x = torch.tanh(conv(F.pad(x, ((conv.kernel_size[0] - 1) * conv.dilation[0], 0))))
    return x

def reach(layers, T=64, d=8):
    x = torch.randn(1, d, T)
    poked = x.clone(); poked[0, :, 0] += 5.0     # change position 0 only
    with torch.no_grad():
        moved = (run(layers, x) - run(layers, poked)).abs().amax(dim=1)[0] > 0
    return int(moved.nonzero().max()) + 1       # how many positions position 0 can reach

for dil in [[1], [1, 1, 1, 1], [1, 2, 4, 8], [1, 2, 4, 8, 16]]:
    print(f"dilations {str(dil):18} formula 1 + sum = {1 + sum(dil):3d}   measured {reach(conv_stack(dil)):3d}")

Code walkthrough

  • conv_stack builds one nn.Conv1d with a window (kernel) of 2 per dilation.
  • run pads on the left only, so no output sees later positions: a causal convolution.
  • reach changes position 0 and counts how far the outputs moved, the receptive field; the formula for kernel 2 is 1 plus the sum of the dilations.
KERNEL 2, CAUSAL CONVOLUTIONS: REACH dilations formula measured [1] 2 2 [1, 1, 1, 1] 5 5 [1, 2, 4, 8] 16 16 [1, 2, 4, 8, 16] 32 32 formula: 1 + sum of dilations
Poking position 0 and watching which outputs move gives the same reach as 1 + sum of dilations in every row (green, passed): doubling dilations doubles the reach per layer.

Every row matches. Four plain layers reach 5 positions; dilations up to 16 reach 32, doubling per layer. That logarithmic growth let the TCN hold a 1,000-step gap that the recurrent networks lost.

python
import math
print(f"{'positions':>11} {'RNN steps':>10} {'conv layers':>12} {'dilated layers':>15} {'attention':>10}")
for T in [17, 1024, 1_000_000]:
    plain = T - 1                                # kernel 2, dilation 1: each layer reaches 1 further
    dilated = math.ceil(math.log2(T))            # kernel 2, dilations 1, 2, 4, ...
    print(f"{T:11,} {T - 1:10,} {plain:12,} {dilated:15} {1:10}")

Code walkthrough

  • For T positions a recurrence needs T - 1 steps, plain kernel-2 convolutions T - 1 layers, dilated ones ceil(log2(T)), attention 1.
TO CONNECT FIRST AND LAST POSITION positions RNN steps dilated layers attention 17 16 5 1 1,024 1,023 10 1 1,000,000 999,999 20 1
Arithmetic: joining the ends of a million positions takes 999,999 recurrent steps, 20 dilated convolution layers (kernel 2) or one attention layer.

For a million positions: 999,999 recurrent steps, 20 dilated layers, or one attention layer. Convolution removed the waiting and cut the path to a logarithm; attention cuts it to 1, with connections chosen by content.

VASWANI ET AL. 2017, TABLE 1 (ORDER OF GROWTH) layer work per layer in order max path self-attention n² d 1 1 recurrent n d² n n convolutional k n d² 1 log_k n n = length, d = width, k = kernel
The transformer paper's own trade (cited): self-attention (amber) pays n squared work per layer to get one sequential step and a maximum path length of 1.

Vaswani and colleagues tabulated this trade in Table 1 of "Attention Is All You Need". Attention's price is work that grows with the square of the length; Module 10 measures that bill.

In production

Before transformers, recurrent models ran production translation. Google's 2016 system (Wu and colleagues) used 8 LSTM layers in the encoder and 8 in the decoder, and the paper notes such systems "are known to be computationally expensive both in training and in translation inference". Its attention design was chosen "to improve parallelism and therefore decrease training time". The sequential dependency was felt first as cost.

Recurrence did not vanish: its fixed-size state suits long streams, where a transformer's memory grows with every token (Module 10). State-space models, which train in parallel but run like a recurrence, return in Module 14.

What you added to tinyformer-lab

  • The lab's first report: a matched-size GRU and transformer raced on the copy task, with curves and per-position accuracy.
  • A gradient probe that measures the signal reaching each input position, for any layer type.
  • A receptive-field check for convolution stacks against its formula.
  • The module lab's experiments/m01_rnn_vs_attention.py keeps the race, with this lesson's GRU as its first racer; the gradient probe and the receptive-field check stay in this lesson.

Check yourself

  1. Why can no implementation trick run step t of a recurrence before step t - 1?
  2. The GRU reached 0.568. Why is that not proof that GRUs can never copy?
  3. Why does the gradient shrink like a power through a recurrence, but only as one over the length through attention?
  4. What does a forget-gate bias of +2 change, and what does it leave?
  5. How many kernel-2 dilated layers connect the ends of 1,024 positions?
  1. h_t = tanh(x_t W_x + h_t-1 W_h) takes h_t-1 as an input, so it must exist first.
  2. It was still climbing, on one seed, size and budget: it learns slowly, which is all the run shows.
  3. A recurrence multiplies one factor per step; attention puts every position one hop from the loss.
  4. The gate opens to about 0.88, so the gradient decays far more slowly; it is still a product, the state is still fixed-size, steps still run in order.
  5. Ten (2 to the 10th is 1,024), against one attention layer.

Exercise

Double the copy length to 32. First write down your predicted accuracy for each model, then run the check. The longer task is added in this lesson; the shared TASKS stay unchanged.

python
long_copy = make_task("copy", length=32)                   # 32 tokens, separator, 32 tokens again
set_seed(); gru_curve = race(GRUModel(), long_copy)
set_seed(); att_curve = race(AttentionModel(context=64), long_copy)
gru_acc, att_acc = gru_curve[-1][1], att_curve[-1][1]
print(f"GRU {gru_acc:.3f}   transformer {att_acc:.3f}")
assert att_acc >= 0.95, "the transformer should still solve the longer copy"
assert gru_acc < accuracy(results["GRU"][0], copy), "the GRU should do worse at length 32 than at 16"
print("check passed")

Code walkthrough

  • Rows are now 65 tokens, so the transformer needs context=64. Both models are raced as before; the checks require the transformer above 0.95 and the GRU below its length-16 score.
  • On the author's Mac it printed GRU 0.342 transformer 1.000. Then explain the GRU's per-position accuracy using section 4.

Key terms

TermMeaning
Recurrent networkReads one token per step and carries a hidden state forward
Sequential dependencyStep t needs step t - 1, so steps run in order
JacobianHow each output number of a step changes with each input number
Vanishing gradientThe gradient shrinks geometrically going back through steps
Exploding gradientThe same product grows instead; clipping caps it
Forget gateAn LSTM gate between 0 and 1 scaling the carried state
Path lengthSteps a signal travels between two positions
Receptive fieldHow many input positions can affect one output

Sources

  • Shaojie Bai, J. Zico Kolter, Vladlen Koltun, "An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling", 2018. https://arxiv.org/abs/1803.01271
  • Razvan Pascanu, Tomas Mikolov, Yoshua Bengio, "On the difficulty of training Recurrent Neural Networks", 2012 (revised 2013). https://arxiv.org/abs/1211.5063
  • Ashish Vaswani et al., "Attention Is All You Need", 2017. https://arxiv.org/abs/1706.03762
  • Yonghui Wu et al. (Google), "Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation", 2016. https://arxiv.org/abs/1609.08144