Recurrent Neural Networks

Before attention, a sequence model was a loop. Keep a hidden state, update it with each token, and read the answer off the end.

Everything appealing about the design is in that equation. The state is a fixed size, so memory does not grow with the sequence. The weights are shared across steps, so the model handles any length. Decoding is $O(1)$ per token, which is the property the whole recurrence shelf is trying to recover sixty years later.

Why it does not train

Backpropagation through time unrolls the loop and applies the chain rule down it. The gradient from step $t$ to step $t-k$ is a product of $k$ Jacobians.

A product of $k$ matrices is governed by the largest singular value. Below one it decays geometrically and the gradient at distance 20 is numerically zero; above one it grows geometrically and overflows. The knife-edge in between is not a region you can train in, because the value drifts as the weights update.

Bengio et al. proved the trade-off is not an artefact of a bad initialisation: a recurrence that stores information robustly must have its Jacobian’s spectral radius below one, and that is precisely the condition under which gradients vanish. Stability and trainability pull in opposite directions.

What the fixes have in common

Everything after this attacks the product in eq. 2.

Gating — LSTM and GRU insert an additive path with a Jacobian near the identity, so the product does not decay.

Attention — remove the recurrence altogether. Every position reaches every other in one step, so there is no product to control. This is what won, and the cost is the quadratic.

Linearity — state space models drop the nonlinearity from the recurrence, which makes the Jacobian a constant diagonal and its product analytically controllable — and, unexpectedly, makes the whole sequence computable by parallel scan rather than a loop.

That third route is the interesting one, because the nonlinearity looked essential. It turns out that a linear recurrence with nonlinearities between layers is expressive enough, and it is the only formulation that keeps constant- memory decoding while training as fast as attention.

import torch
from torch import Tensor, nn


class ElmanRNN(nn.Module):
    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.wx = nn.Linear(dim, hidden, bias=False)
        self.wh = nn.Linear(hidden, hidden, bias=True)
        self.hidden = hidden

    def forward(self, x: Tensor) -> Tensor:                   # [B, T, D]
        h = x.new_zeros(x.size(0), self.hidden)
        out = []
        # Sequential and unavoidably so: h_t needs h_{t-1}.
        for t in range(x.size(1)):
            h = torch.tanh(self.wx(x[:, t]) + self.wh(h))
            out.append(h)
        return torch.stack(out, dim=1)


def spectral_radius(w: Tensor) -> float:
    """The number that decides whether gradients survive. Log it."""
    return torch.linalg.eigvals(w).abs().max().item()

The loop is the other half of the problem, and it is the half that made attention inevitable. $h_t$ depends on $h_{t-1}$, so training cannot parallelise over the sequence — a GPU that could process 4096 tokens at once processes them one at a time. Attention’s quadratic cost bought perfect parallelism, and on the hardware of 2017 that was a trade worth making by a wide margin.