60 - LSTM Put to the Test: The Adding Problem, Stacked Layers and Padding

In the previous article we saw that the LSTM opens a road for the gradient. Here we check that the optimiser really walks it, with a classic test bench: the adding problem, proposed by Hochreiter and Schmidhuber in their 1997 paper. Then we look at how nn.LSTM is used in everyday practice: several layers, dropout, sequences of different lengths. Every number comes from runs of the code shown (PyTorch 2.12.0 on CPU, two runs with identical output, a single seed per configuration).

The words of this article

A benchmark is a task built on purpose to measure a single capability: its value lies in the fact that you know the right answer and can compute what a model that has not understood the problem will get. A long-term dependency is the case in which the information needed to answer at a given step sits many steps earlier. A plateau is the stretch of the loss curve where the error stays flat, only to (sometimes) drop suddenly: it is not a crash, the network is looking for the way.

A layer is a block of the network: a two-layer LSTM has two LSTMs stacked, and the second reads, as a sequence, the states produced by the first. Dropout is a regularisation technique that, during training, switches off a random fraction of the neurons at every step to stop the network from leaning too heavily on a few of them. Padding is the dummy values (here zeros) appended to the end of short sequences to bring them to the length of the longest one so that they can be stacked into a single tensor.

The task: adding two numbers among many

The network receives a sequence of T steps, with two channels per step: a random number in [0, 1] and a marker that is 1 in exactly two positions (chosen at random) and 0 everywhere else. The goal is the sum of the two marked numbers. To succeed, the network must ignore almost the whole sequence, notice the first marker, keep that number in memory until the second one shows up, and add. The distance between the two numbers can be almost T: it is a direct test of long-term memory.

There is an honest baseline: if the network always predicts the mean of the sum (1.0), the mean squared error equals the variance of the sum of two uniforms, that is 2/12 ≈ 0.167. A model that has not understood the task sits at 0.167; one that has understood it goes close to zero.

import torch
import torch.nn as nn

def adding_batch(batch, T):
    """Hochreiter & Schmidhuber's adding problem.
    Input: 2 channels. Channel 0: random numbers in [0,1]. Channel 1: a marker with exactly two 1s.
    Target: the sum of the two marked numbers."""
    values = torch.rand(batch, T)
    markers = torch.zeros(batch, T)
    for b in range(batch):
        idx = torch.randperm(T)[:2]
        markers[b, idx] = 1.0
    x = torch.stack([values, markers], dim=-1)          # (batch, T, 2)
    y = (values * markers).sum(dim=1)                   # (batch,)
    return x, y

class Net(nn.Module):
    def __init__(self, kind, hidden=64):
        super().__init__()
        cls = {"rnn": nn.RNN, "lstm": nn.LSTM}[kind]
        self.rec = cls(2, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)
        if kind == "lstm":                              # forget-gate bias = 1
            with torch.no_grad():
                self.rec.bias_ih_l0[hidden:2 * hidden] = 1.0
    def forward(self, x):
        out, _ = self.rec(x)
        return self.head(out[:, -1]).squeeze(-1)

def run(kind, T, steps=4000, seed=0):
    torch.manual_seed(seed)
    net = Net(kind)
    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
    for _ in range(steps):
        x, y = adding_batch(64, T)
        loss = nn.functional.mse_loss(net(x), y)
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
    with torch.no_grad():
        x, y = adding_batch(1000, T)
        return nn.functional.mse_loss(net(x), y).item()

# Reference: always predicting the mean (1.0) gives MSE = Var(sum of two U(0,1)) = 2/12 = 0.167
for T in (20, 50, 100):
    print(f"T={T:3d}  RNN {run('rnn', T):.4f}   LSTM {run('lstm', T):.4f}")

Reading the code. adding_batch builds a batch of examples: values are the random numbers, markers is a tensor of zeros in which, for each sequence, two randomly chosen positions (torch.randperm(T)[:2]) are set to 1. The input is the stacking of the two channels (torch.stack(..., dim=-1), shape (batch, T, 2)) and the target is (values * markers).sum(dim=1): it zeroes the unmarked values and sums the rest. The Net class is an RNN or an LSTM plus a linear head on the last state; for the LSTM, the forget-gate bias is set to 1 with the same [hidden : 2*hidden] slice seen in the previous lesson. run trains with Adam and clipping and returns the MSE on 1000 fresh sequences: since every training batch is generated on the fly, there is no chance of memorisation and the final number measures the real ability to generalise.

Results

Mean squared error on 1000 new sequences, after 4000 training steps with Adam, clipping at 1.0 and (for the LSTM) forget-gate bias at 1:

         RNN       LSTM
T=20   0.0199    0.0001
T=50   0.1491    0.0005
T=100  0.1785    0.0002

The result reads in three points. The plain RNN learns at distance 20 and collapses beyond it, staying at the level of "predict the mean" (0.167), in line with the article on the gradient but with a single seed. The LSTM solves it at every length tried, with errors a thousand times smaller than the baseline. The initial plateau is normal: with few steps the LSTM at T=100 stays stuck on the baseline and only with more training does it reach tiny values, so if you stop too early you conclude the model does not work when it only needs time.

The general lesson: a loss that does not go down does not necessarily mean the model is wrong. Before changing architecture, check the baseline of the task and train for longer.

Looking at the plateau up close

The sentence "the LSTM at T=100 after 4000 steps is at 0.0002" hides the most instructive part: how it got there. Let us record the mean training loss every 250 steps, with the same setup (T = 100, one seed).

def loss_curve(kind, T, steps=4000, every=250, seed=0):
    """Same training as run(), but returns the training loss every `every` steps."""
    torch.manual_seed(seed)
    net = Net(kind)
    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
    curve, acc = [], 0.0
    for step in range(1, steps + 1):
        x, y = adding_batch(64, T)
        loss = nn.functional.mse_loss(net(x), y)
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
        acc += loss.item()
        if step % every == 0:
            curve.append(acc / every); acc = 0.0           # average loss of the last `every` steps
    return curve

curve = loss_curve("lstm", T=100)
for i, v in enumerate(curve, start=1):
    print(f"steps {(i-1)*250+1:4d}-{i*250:4d}  mean train loss {v:.4f}")

steps    1- 250  mean train loss 0.1817
steps  251- 500  mean train loss 0.1669
steps  501- 750  mean train loss 0.1675
steps  751-1000  mean train loss 0.1668
steps 1001-1250  mean train loss 0.1654
steps 1251-1500  mean train loss 0.1703
steps 1501-1750  mean train loss 0.1455
steps 1751-2000  mean train loss 0.0087
steps 2001-2250  mean train loss 0.0023
...
steps 3751-4000  mean train loss 0.0002

How to read the table: loss_curve is the usual training, but it accumulates the loss of every step and writes the mean every 250. From steps 251 to 1500 the loss sits pinned at about 0.167: exactly the "always predict the mean" baseline. The network has learned the easiest thing, the mean output, and has not yet found the structure of the problem. Then, between steps 1500 and 2000, the loss collapses by two orders of magnitude (from 0.1455 to 0.0087), and keeps falling.

This pattern has an intuitive explanation: to improve on the mean, the network must first learn to do two things together (look at the marker channel and keep the marked value); until it does both, neither reduces the error, so the gradient towards each is weak. When one of the two appears, the other becomes useful and progress plummets. For you, training a real model, the practical lesson is: do not stop for a flat loss if the task baseline is exactly that value, but do not wait forever either: compare with the baseline, and if after a generous budget you do not break away, the problem is elsewhere (learning rate, initialisation, architecture).

Using nn.LSTM in practice

The three parameters that really change the behaviour are num_layers, dropout and bidirectional (the latter has a lesson of its own). With num_layers=2 the second LSTM reads the whole sequence of states produced by the first: more capacity for abstraction, but also more cost and more difficulty in training. PyTorch's dropout is applied between layers and not inside the time loop, so with a single layer it does nothing (PyTorch warns you).

The other practical topic is variable length. The sentences in a batch have different lengths: you pad with zeros up to the longest. But if you then read the last state of a sequence that was short, you are reading the state after a string of zeros, not after the last real element. pack_padded_sequence tells the network to stop at the true length of each sequence.

import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_sequence

torch.manual_seed(0)
seqs = [torch.randn(5, 2), torch.randn(3, 2), torch.randn(1, 2)]      # three sequences, different lengths
lengths = torch.tensor([len(s) for s in seqs])
padded = pad_sequence(seqs, batch_first=True)                          # (3, 5, 2), zeros at the end

lstm = nn.LSTM(2, 8, num_layers=2, dropout=0.2, batch_first=True)    # dropout acts BETWEEN layers only
packed = pack_padded_sequence(padded, lengths, batch_first=True, enforce_sorted=False)
_, (h_n, c_n) = lstm(packed)           # h_n: (num_layers, batch, hidden) = last REAL step of each sequence

print(h_n.shape)                       # torch.Size([2, 3, 8])
last_layer = h_n[-1]                   # (3, 8): one summary vector per sequence, padding ignored

# Compare with the naive way: last position of the padded tensor (reads zeros for the short sequences)
lstm.eval()
naive, _ = lstm(padded)
_, (h_ok, _) = lstm(pack_padded_sequence(padded, lengths, batch_first=True, enforce_sorted=False))
print((naive[:, -1] - h_ok[-1]).abs().max(dim=1).values)    # ~0 for the longest sequence, large for the others

Let us follow the code. pad_sequence brings three sequences of lengths 5, 3 and 1 to a (3, 5, 2) matrix, filling with zeros. pack_padded_sequence packs the tensor together with the true lengths, so that the LSTM skips the padding steps (enforce_sorted=False saves you from having to sort the sequences). After the LSTM, h_n[-1] is the state of the last layer at the true last step of each sequence: three vectors of 8. The final comparison reads naive[:, -1], that is the last position of the padded tensor, and compares it with h_ok[-1]: where there was padding (second and third sequence), the two readings differ.

The last line measures the difference between the two ways of reading the last state: zero (2.98·10⁻⁸, numerical noise) for the longest sequence, which has no padding, and up to 0.047 and 0.143 for the two shorter sequences. It is not a huge difference on random untrained data, but it is a systematic distortion that the model has to learn to compensate for, and that in production pollutes every prediction on short sentences.

The shape of h_n is (num_layers, batch, hidden), with the batch in second position even if you set batch_first=True (that flag concerns only the input and the out output). The last layer is h_n[-1]. This is a frequent shape mistake, and it raises no error as long as the dimensions happen to coincide.

More layers and dropout: do they really help?

A common reflex is "deeper = better". Let us try it on the adding task at T = 50, with a one-layer LSTM, a two-layer one, and a two-layer one with dropout 0.2, all with the forget-gate bias at 1 in every layer.

class DeepNet(nn.Module):
    def __init__(self, layers, dropout, hidden=64):
        super().__init__()
        self.rec = nn.LSTM(2, hidden, num_layers=layers, dropout=dropout if layers > 1 else 0.0, batch_first=True)
        self.head = nn.Linear(hidden, 1)
        with torch.no_grad():
            for l in range(layers):
                getattr(self.rec, f"bias_ih_l{l}")[hidden:2 * hidden] = 1.0
    def forward(self, x):
        out, _ = self.rec(x)
        return self.head(out[:, -1]).squeeze(-1)

def run_deep(layers, dropout, T=50, steps=3000, seed=0):
    torch.manual_seed(seed)
    net = DeepNet(layers, dropout)
    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
    for _ in range(steps):
        x, y = adding_batch(64, T)
        loss = nn.functional.mse_loss(net(x), y)
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
    net.eval()
    with torch.no_grad():
        x, y = adding_batch(1000, T)
        return nn.functional.mse_loss(net(x), y).item(), sum(p.numel() for p in net.parameters())

for layers, dropout in ((1, 0.0), (2, 0.0), (2, 0.2)):
    mse, n_par = run_deep(layers, dropout)
    print(f"layers={layers} dropout={dropout}: MSE {mse:.4f}  ({n_par} parameters)")

layers=1 dropout=0.0: MSE 0.0001  (17473 parameters)
layers=2 dropout=0.0: MSE 0.0003  (50753 parameters)
layers=2 dropout=0.2: MSE 0.0003  (50753 parameters)

Explaining the code: DeepNet creates an nn.LSTM with the desired number of layers; dropout is passed only if layers > 1 (otherwise PyTorch warns that it would be ignored); the loop sets the forget-gate bias of every layer to 1 (bias_ih_l0, bias_ih_l1…). After training we switch to eval(), which turns dropout off, and count the parameters with numel. Result: the single layer is already perfect (0.0001) with 17,473 parameters; two layers use 50,753 (about three times as many) without improving (0.0003), and dropout changes nothing here because the task is generated on the fly and there is no overfitting to regularise. Rule: add complexity only when a measurement justifies it. Every extra layer roughly triples the parameters, slows down training and introduces another source of instability.

Which knobs matter? Measuring "steps to solve"

When two configurations both end at 0.0002, the final value does not tell them apart. A more informative metric is how many steps it takes to solve the task: the first step at which the moving average of the loss (last 100 steps) drops below 0.05, that is below a third of the baseline. The function steps_to_solve computes it, and returns None if the budget runs out before it gets there. We vary one knob at a time, with T = 50 and one seed.

def steps_to_solve(lr=3e-3, forget_bias=1.0, hidden=64, T=50, max_steps=4000, threshold=0.05, seed=0):
    """First training step at which the running mean of the loss (last 100 steps) drops below `threshold`.
    Returns None if it never does within max_steps."""
    torch.manual_seed(seed)
    net = DeepNet(1, 0.0, hidden)
    with torch.no_grad():
        net.rec.bias_ih_l0[hidden:2 * hidden] = forget_bias
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    recent = []
    for step in range(1, max_steps + 1):
        x, y = adding_batch(64, T)
        loss = nn.functional.mse_loss(net(x), y)
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
        recent.append(loss.item()); recent = recent[-100:]
        if step >= 100 and sum(recent) / 100 < threshold:
            return step
    return None

print("learning rate :", {lr: steps_to_solve(lr=lr) for lr in (1e-3, 3e-3, 1e-2)})
print("forget bias   :", {fb: steps_to_solve(forget_bias=fb) for fb in (0.0, 1.0, 2.0)})
print("hidden size   :", {h: steps_to_solve(hidden=h) for h in (8, 32, 128)})

learning rate : {0.001: 1645, 0.003: 1027, 0.01: 604}
forget bias   : {0.0: 887, 1.0: 1027, 2.0: 788}
hidden size   : {8: 1497, 32: 851, 128: 928}

How to read it, with due caution (one seed, differences of 10-20% are noise). Learning rate: it is the clearest knob, and monotonic: 0.001 takes 1645 steps, 0.003 takes 1027, 0.01 takes 604. A higher learning rate speeds up the plateau; the price is instability, which here the clipping kept under control. Forget-gate bias: at T = 50 and with this task, 0, 1 and 2 give results of the same order (887, 1027, 788), with no clear winner. This does not contradict the earlier lessons: the bias had a clear effect on the flow of the untrained gradient and on the recall task at T = 20 with a standard cell; here the optimiser, with 4000 steps at its disposal, manages anyway. A sign that the same knob weighs more in harder conditions (longer sequences, fewer steps). State size: 8 units are too few (1497 steps), 32 and 128 are equivalent (851 and 928). You do not need a huge state for a task that requires remembering a single number.

The methodology is the part to take away: measure the learning time as well as the final result, because two models that reach the same point can differ by a factor of two or three in the cost of getting there, and that is a cost you pay at every training run.

Practical trade-offs

• A single layer is often enough. Add layers only if validation justifies it; every layer is another source of instability.
• Hidden size: 32 is enough for the adding task. For real tasks you start from 128-256 and check with validation.
• LSTM versus GRU: the GRU is lighter; at equal size the LSTM has more parameters.
• Very long lengths (hundreds or thousands of steps): even the LSTM struggles and costs memory in BPTT, and you turn to truncated BPTT or attention models.

The four items share a logic: every extra capability costs parameters, time or stability, and must be justified by a measurement. The reasonable starting point is the single layer with a medium state; from there you go up only if validation asks for it, and for very long sequences the choice shifts to other architectures.

Real traps

The first is concluding too early: the plateau can last thousands of steps, so set a sensible training budget and compare with the baseline of the task. The second is using a benchmark that does not require memory: if the task can be solved by looking at the last elements, a plain RNN "wins" and you have measured nothing; the adding problem is useful precisely because the position of the two numbers is random. The third is reading the last state of padded sequences: use pack_padded_sequence or select the index lengths − 1 by hand. The fourth is confusing batch_first with the shape of h_n, which always has the batch in position 1. The last is dropout with a single layer, which does nothing; to regularise inside the loop you need techniques such as variational dropout, not included in nn.LSTM.

Checked with four calculations

The baseline of the adding problem is 2/12 because the sum of two uniforms in [0, 1] has mean 1 and variance 1/12 + 1/12 = 1/6 ≈ 0.167, and always predicting the mean gives a mean squared error equal to the variance. If the loss stays at 0.167 for 1000 steps the model is not necessarily broken: it can be a pre-learning plateau, and before changing architecture it is worth training longer and checking learning rate, clipping and forget-gate bias. With sequences of lengths 3, 7 and 10 you avoid reading the state after the padding by using pack_padded_sequence with the true lengths, or by selecting by hand the index length − 1 of the out states. The shape of h_n for a 3-layer LSTM, batch 16, state 64 is (3, 16, 64): layers, batch, state.

A worked example by hand: the adding problem over eight steps

Take a sequence of T = 8 steps (as in the drawing): values 0.3, 0.9, 0.1, 0.7, 0.5, 0.2, 0.8, 0.4 and markers 0, 1, 0, 0, 0, 0, 1, 0. The target is the sum of the two marked values: 0.9 + 0.8 = 1.7. What must the network do, step by step? At steps 1, 3, 4, 5, 6 and 8 the marker is 0: the cell must leave the memory unchanged (input gate close to 0, forget close to 1). At step 2 the marker is 1: it must write 0.9 into memory (input open). In steps 3-6 it must keep it without being disturbed by the values 0.1, 0.7, 0.5, 0.2. At step 7 it must also write 0.8, adding it to the one already there (additive memory: c = f·c + i·g is precisely a sum). At step 8 it must not spoil anything. At the output, the linear head reads 1.7.

It is an almost literal description of the behaviour of the LSTM gates, and it is the reason this task is a good test: within a single sequence it requires writing, keeping, adding and ignoring. The baseline is 1/6 ≈ 0.167 because each uniform value in [0, 1] has variance 1/12 and the variances of two independent variables add up; a mean squared error of 0.0002 corresponds to a typical deviation of √0.0002 ≈ 0.014 on the sum, against √0.167 ≈ 0.41 for someone who predicts the mean: about 29 times smaller.

The padding too can be counted by hand. With three sequences of lengths 5, 3 and 1, the padded tensor has 3 × 5 = 15 positions, but only 5 + 3 + 1 = 9 are real: 40% of the computation, without packing, works on zeros (and, worse, produces a final state that has read those zeros).

Frequently asked questions

Why does the loss stay flat at 0.167 and then collapse? Before discovering that the marked value matters, the best move is to predict the mean; discovering it requires looking at the marker channel and keeping the value at the same time, and until that happens the gradient towards each of the two skills is weak. The collapse comes when one of the two is "hooked" and makes the other useful.

Wouldn't clipping be enough to avoid the plateau? No: clipping limits gradients that are too large, while the plateau has small gradients. You shorten it with a higher learning rate (604 steps at 0.01 against 1645 at 0.001), a better initialisation or more data per step.

How many layers do I need? Here one is enough; two did not help (0.0003 against 0.0001) and tripled the parameters. Add layers when validation justifies it, not by default.

How do I choose the batch size? Here 64. Larger batches give less noisy gradients but cost more per step; small batches are noisy and, in the plateau, may help to get out of it. We did not vary it: we cannot say which is better for this task.

Does the result hold for text and audio? It is a result on a synthetic memory task: it indicates that the LSTM can keep a value for tens of steps. For real language many other things matter (vocabulary, data, tokenisation) that we did not measure here.

Exercises with solutions

Before moving on, try to solve these exercises with pen and paper; the solutions (checked with the code) are right below each one.

Exercise 1. How many parameters does nn.LSTM(2, 64) have with PyTorch's two bias vectors? Solution: 4·64·(64 + 2 + 2) = 17,408 (four transformations, each with a 64×(64+2) matrix and two biases of 64). The adding network adds a linear head of 65 parameters: 17,473, the number printed in the experiment.

Exercise 2. Four sequences have lengths 10, 6, 4 and 2: how many positions are real and how many are padding? Solution: the tensor has 4·10 = 40 positions, of which 10 + 6 + 4 + 2 = 22 are real and 18 are padding, 45% of the computation without packing.

Exercise 3. A model has MSE 0.04 on the adding problem (baseline 1/6): what fraction of the variance does it explain? Solution: 1 − 0.04/0.1667 = 0.76, 76% of the variance of the target; an MSE of 0.0002 explains 99.9% of it. One more question, without calculations: why can the same problem give 0.167 at 1500 steps and 0.0002 at 4000? The answer is in the section on the plateau: the loss has a flat stretch, so the number of steps is part of the experiment.

What you can do now

You can use the adding problem as a test bench for memory, comparing RNN and LSTM against the 2/12 baseline and recognising the initial plateau for what it is. You can use nn.LSTM with several layers and dropout, handle sequences of different lengths with pack_padded_sequence, read the shape of h_n and measure the knobs not only on the final result but on the steps needed to solve the task. The next step is the lighter cell, the GRU.