RNNs & LSTMs — Sequences & the Vanishing Gradient

Why this session exists

Recurrence applies one transition repeatedly while maintaining a fixed-width state. It is useful for streaming, but learning distant dependencies involves long derivative paths. LSTMs add gated additive memory; they mitigate rather than abolish that difficulty. We will derive BPTT, inspect the exact PyTorch LSTM/GRU conventions, train a small character model on an inline corpus and separate those mechanics from claims about realistic language generation.

A feedforward model can predict a missing word from a fixed context window. An RNN's distinction is state reuse and variable-length processing, not a unique ability to model sequences. There is no universal 20-step RNN limit, 500-step LSTM limit, or rule that Transformers made every recurrent deployment obsolete.

Prerequisites

S098 backpropagation, S100 PyTorch and S102 shared filters. We use column-vector equations and batch-first code; identify the convention before comparing Jacobians.

(a) Intuition · 5 min

Reading left to right suggests a running summary, but a finite hidden vector does not contain everything seen so far. With $h_t=\tanh(W_xx_t+W_hh_{t-1}+b)$, each new token changes the same state using the same parameters. Unlike a CNN's spatial sharing, the next state depends on the previous computation. Fixed parameter count permits arbitrary input length in principle; runtime, memory and representational capacity still limit the useful length.

What RNNs unlocked that MLPs couldn't

Recurrence supplies a natural variable-length interface, incremental memory and an autoregressive generator. A fixed-window MLP or temporal CNN can also generate by feeding its predictions back. The tradeoff is what history remains available: a finite window drops older inputs, whereas recurrence compresses history into state. Compression may preserve task-relevant information or lose it. Generation is sequential for both a recurrent model and an autoregressive Transformer, even though their teacher-forced training dependencies differ.

A short history of sequence modelling

Recurrent networks predate modern LSTMs. The original 1997 LSTM and later forget-gate variants should not be conflated into a claim that the first paper already specified today's complete cell. Seq2seq (2014) demonstrated encoder-decoder LSTMs; GNMT (2016) combined deep recurrence, residual connections, attention, wordpieces and inference engineering. The Transformer (2017) replaced recurrence within its translation architecture, not every sequence model overnight. Gates, attention and autoregressive factorization have related but distinct histories; not every Transformer idea originated in LSTM.

An RNN unrolled through time

Each state receives a local-loss gradient and future-state gradients. Shared parameters collect contributions from all uses. Unrolling is a graph explanation; nn.LSTM can use optimized kernels rather than a Python loop.

Why gradients vanish — the math

Define $D_t=\operatorname{diag}(1-h_t^2)$, so $J_t=\partial h_t/\partial h_{t-1}=D_tW_h$ in the column convention. For a final loss, the derivative to an earlier state contains an ordered product of time-varying Jacobians. For losses at every step, let $g_t$ be the local derivative with respect to $h_t$ and $\delta_t$ the total derivative with respect to the preactivation:

The terminal future contribution is zero. For the final-loss path, $\|J_T\cdots J_{s+1}\|\le\prod_t\|J_t\|$. If every factor's operator norm is bounded above by the same constant $q<1$, this path's norm is at most $q^{T-s}$ and decays with path length. A scalar constant factor 0.9 gives $0.9^{30}\approx0.0424$ and $0.9^{100}\approx2.66\times10^{-5}$, but a general nonlinear RNN's gradient path is an ordered product of different $D_tW_h$ factors, not a power of the recurrent weight matrix alone. Eigenvalues alone do not control finite products of nonnormal/time-varying matrices; singular values, activation saturation and alignment matter. Feedforward deep networks can suffer the same issue without shared weights. Orthogonal initialization and normalization can help conditioning but do not guarantee retention.

Clipping caps a finite aggregate gradient norm. It cannot recover a vanished signal, reverse earlier overflow, or guarantee the gradient direction is useful. With mixed precision, unscale before clipping. Truncated BPTT detaches carried state between windows: the forward value persists while earlier graph paths do not. Shuffled unrelated windows should not inherit each other's states.

The LSTM cell — three gates and a cell state

There is no peephole connection from the previous cell directly into the gates in this standard cell. The candidate $g_t$, omitted by many sketches, is essential.

Anatomy of an LSTM cell

For $a_t=[h_{t-1};x_t]$:

The conveyor-belt metaphor means selective retain/write/read, not immortal memory. Holding gates fixed, the direct cell path derivative is $\operatorname{diag}(f_t)$ and over steps is a product of forget gates. At 0.99 for 100 steps it is about 0.366; at 0.9 it is about 0.0000266. The total derivative also includes gate dependence through hidden states. Sigmoid gates mathematically lie strictly between zero and one for finite logits, though floating arithmetic can saturate. A copy-delay sweep is an experiment, not a prewritten guarantee that an LSTM will achieve 100% at a particular delay.

RNN vs GRU vs LSTM vs Transformer — when to use what

PyTorch's GRU uses reset $r$, update $z$, candidate $n$ and $h_t=(1-z_t)\odot n_t+z_t\odot h_{t-1}$. Its reset multiplies the hidden affine candidate after that affine map, unlike formulations applying reset before the matrix. It has one state tensor, whereas LSTM returns (h,c). For equal input/hidden widths, standard GRU has three affine gate groups versus LSTM's four, so its recurrent-module parameter count is 3/4, not necessarily the full model's count or runtime.

No row is strictly better. Some state-space recurrences admit parallel scans; that does not remove the data dependence of an arbitrary nonlinear LSTM. Full dense attention pairwise work is quadratic in length; materialized score memory is quadratic, but fused algorithms can avoid storing it. Fixed-width recurrent inference has bounded state, whereas training through a long sequence stores or recomputes intermediate states.

Bidirectional and stacked variants

Bidirectional outputs concatenate forward/backward states: useful for a fully observed tagging/encoder input, inappropriate for a causal next-token predictor that would see its target. Stacking feeds one layer's sequence to another. batch_first=True gives output (B,T,directions*H) but hidden/cell states remain (layers*directions,B,H) without projections. With padding, pack valid lengths or mask appropriately; the last padded output is not generally the final real token's state. GNMT used a bidirectional bottom encoder layer with unidirectional higher layers, not eight bidirectional encoder layers.

(c) Hands-on · 25 min

The program trains both cells on an inline toy character corpus with contiguous train/validation/test partitions formed before windows. Its repetitive text is a mechanics fixture, not Shakespeare or a language benchmark. Validation chooses checkpoints; printed samples are unscored illustrations. Replace the inline text with a permitted local corpus for a real experiment, keep document-level splits, and report corpus/version/vocabulary and token-weighted held-out loss.

import copy
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader,TensorDataset
 
torch.set_num_threads(1);torch.manual_seed(11)
class CharModel(nn.Module):
    def __init__(self,vocab,kind='lstm',hidden=24,layers=2):
        super().__init__();self.kind=kind
        self.embed=nn.Embedding(vocab,hidden)
        cls={'lstm':nn.LSTM,'gru':nn.GRU}[kind]
        self.rnn=cls(hidden,hidden,num_layers=layers,batch_first=True,
                     dropout=.2 if layers>1 else 0.)
        self.head=nn.Linear(hidden,vocab)
    def forward(self,x,state=None):
        out,state=self.rnn(self.embed(x),state)
        return self.head(out),state
 
def windows(ids,length=16):
    n=(len(ids)-1)//length
    if n<1:raise ValueError('partition too short')
    return TensorDataset(ids[:n*length].reshape(n,length),
                         ids[1:n*length+1].reshape(n,length))
 
def epoch(m,loader,opt=None):
    m.train(opt is not None);total=count=0
    with torch.set_grad_enabled(opt is not None):
        for x,y in loader:
            if opt:opt.zero_grad(set_to_none=True)
            logits,_=m(x) # independent windows start with zero state
            loss=F.cross_entropy(logits.flatten(0,1),y.flatten())
            if opt:
                loss.backward()
                torch.nn.utils.clip_grad_norm_(m.parameters(),1.,error_if_nonfinite=True)
                opt.step()
            total+=loss.item()*y.numel();count+=y.numel()
    return total/count
 
@torch.no_grad()
def sample(m,stoi,itos,seed,n=30,temperature=.8):
    if not seed or temperature<=0:raise ValueError('nonempty seed and positive temperature required')
    m.eval();ids=torch.tensor([[stoi.get(c,0) for c in seed]])
    logits,state=m(ids);result=seed
    for _ in range(n):
        nxt=torch.multinomial((logits[:,-1]/temperature).softmax(-1),1)
        result+=itos[nxt.item()];logits,state=m(nxt,state)
    return result
 
text=('the cat sat. the dog ran. a cat can run.\n'*55)
a,b=int(.6*len(text)),int(.8*len(text));parts=[text[:a],text[a:b],text[b:]]
itos=['?']+sorted(set(parts[0])-{'?'});stoi={c:i for i,c in enumerate(itos)}
ids=[torch.tensor([stoi.get(c,0) for c in p],dtype=torch.long) for p in parts]
assert sum(map(len,ids))==len(text)
for kind in ['lstm','gru']:
    torch.manual_seed(11);m=CharModel(len(itos),kind)
    tr=DataLoader(windows(ids[0]),batch_size=8,shuffle=True,
        generator=torch.Generator().manual_seed(9))
    va=DataLoader(windows(ids[1]),batch_size=8);te=DataLoader(windows(ids[2]),batch_size=8)
    opt=torch.optim.Adam(m.parameters(),lr=.015);best=float('inf')
    for e in range(6):
        train=epoch(m,tr,opt);valid=epoch(m,va)
        if valid<best:best=valid;state=copy.deepcopy(m.state_dict())
    m.load_state_dict(state);print(kind,'toy held-out CE',epoch(m,te),sample(m,stoi,itos,'the '))
    m.eval();x=ids[0][:12][None];full,s=m(x)
    left,s1=m(x[:,:5]);right,s2=m(x[:,5:],s1)
    assert torch.allclose(full,torch.cat([left,right],1),atol=1e-6)
    assert isinstance(s,tuple)==(kind=='lstm')
    if kind=='lstm':assert s[0].shape==(2,1,24) and s[1].shape==(2,1,24)
    else:assert s.shape==(2,1,24)
    _,reset=m(x[:,5:])
    # Compare outputs, not differently structured states.
    wrong,_=m(x[:,5:]);assert not torch.allclose(right,wrong)
    print(kind,'state carry/reset negative control passed')
try:windows(torch.tensor([1,2]),16);raise AssertionError('empty windows accepted')
except ValueError:pass
# Exact gate implementation parity with PyTorch, using its i,f,g,o packing.
cell=nn.LSTMCell(3,2).double();x=torch.randn(1,3,dtype=torch.double)
h=torch.randn(1,2,dtype=torch.double);c=torch.randn_like(h)
i,f,g,o=(F.linear(x,cell.weight_ih,cell.bias_ih)+F.linear(h,cell.weight_hh,cell.bias_hh)).chunk(4,-1)
c1=f.sigmoid()*c+i.sigmoid()*g.tanh();h1=o.sigmoid()*c1.tanh()
h2,c2=cell(x,(h,c));assert torch.allclose(h1,h2) and torch.allclose(c1,c2)
print('LSTM gate oracle passed')

Anatomy of the script

Embedding IDs become (B,T,H), recurrent outputs become (B,T,H), and the head makes (B,T,V) logits. Flattening the first two axes matches one next-character target per input position. This is teacher forcing: training receives true preceding characters; free generation receives sampled ones, so errors can compound. The model may still encounter mistaken-looking contexts in real data; scheduled sampling is a separate method, not a universal LLM training step.

Inter-layer dropout is used only for more than one recurrent layer, not as guaranteed recurrent-state dropout. Sampling sets eval and no-grad independently. Temperature divides logits: below one sharpens, above one flattens; neither means objectively more coherent or creative. Empty seeds and nonpositive temperature are rejected. Unknown characters map to a declared fallback fitted from training vocabulary. The GRU swap retains valid multilayer dropout and changes state type, not an arbitrary instruction to remove dropout.

Gradient and memory practice

For a real copy-delay experiment, generate random symbols, reveal the first symbol, insert independently sampled fillers, and score its delayed recovery against the class-prior baseline. Sweep delays using matched initialization budgets, hidden sizes, optimizer searches and seeds. Keep train/test generated examples independent. Retain gradients on nonleaf states before examining them; reading .grad without retain_grad() is not a working diagnostic. A direct-path forget-gate calculation tests the equation, not successful copy-task training. The supplemental checks compare analytical BPTT and finite differences on a tiny tanh RNN, including a wrong-Jacobian negative control.

(d) Production reality · 15 min

GNMT is a concrete source for deep recurrent translation, wordpieces and quantization-aware constraints. It does not support the earlier invented launch latency, overnight worldwide switch or universal requirement that recurrent models use special hardware. Quantization has device/range/quality constraints and must be evaluated end to end.

For demand forecasting, compare last-valid-state readout, masked attention pooling and explicit calendar features under rolling-origin evaluation. Separate seasonal shift from a state-readout bug; no Uber Monday failure or fabricated MAPE reduction is needed to teach that diagnostic.

Learning to Generate Reviews and Discovering Sentiment (2017) studied byte-level language models, including a multiplicative LSTM trained on Amazon reviews and a salient sentiment unit. That is not a 2016 proof of perfect sentiment classification or that next-token prediction inevitably creates intelligence. Features can be useful beyond the training loss, but performance is dataset-specific.

Where this shows up in the rest of the plan

Attention offers queries over retained states; it does not erase recurrent state costs. Forecasting and speech require task-specific evaluation rather than blanket LSTM histories. Positional information in nonrecurrent models replaces some ordering supplied by recurrence. Serving compares bounded recurrent state with cache growth and actual sequential decoding latency.

(e) Recall + stretch · 10 min

1. Weight sharing? The same transition parameters at every step; parameter count is independent of sequence length, not training memory.
2. BPTT? Sum local and future loss contributions through $D_tW_h$; inspect time-varying products, not one eigenvalue or a 20-step cutoff.
3. Forget gate? Multiplies old cell content; its direct derivative product may still decay.
4. GRU versus LSTM? Three versus four gate affine groups at matched widths, one versus two state tensors; accuracy and speed require measurements.
5. Parallelism? Input projections can batch, but nonlinear recurrent transitions depend on prior states. Transformer teacher forcing parallelizes positions; autoregressive sampling remains sequential.
6. Teacher forcing risk? Conditioning mismatch between true histories and generated histories; check held-out loss and free-running behavior separately.
7. Bidirectionality? Valid when the whole source is available; future information leaks in causal generation or forecasting.
8. Clip? Rescale a finite total gradient norm to a chosen cap; not a repair for NaNs or missing gradients.
9. Temperature 0.5, 1, 1.5? Respectively sharpened, unchanged and flattened probabilities, not promised quality levels.
10. CNN–LSTM hybrid? A shared CNN encodes video frames; a masked/packed LSTM integrates temporal order before classification. Split by video/person, not frame, and compare 3D conv/attention baselines without claiming a historical winner.

Explain-out-loud test

Write BPTT for three losses, draw all six LSTM equations, and explain a case where a windowed model or fixed-state recurrent model is preferable to dense attention.

What comes next

Embeddings explains how the character lookup in this code becomes a learned representation. Learning-plan hub.

Primary references

• Pascanu et al. (2013): recurrent gradient products and clipping.
• PyTorch 2.9 tagged recurrent implementation: LSTM/GRU gates, shapes and dropout.
• GNMT and sentiment-unit study: bounded primary case studies.