Practical build 5

Lab: train a gated sequence predictor

Introduction

In this generated sequence task, the first token gives an instruction that matters at the end. A memory-free two-choice baseline has loss about 0.693, so held-out performance can test whether the LSTM carried that instruction.

Learning goal: Train a gated sequence predictor end to end and compare held-out loss with a meaningful memory-free baseline.
Before you start: Embeddings, recurrent state, sigmoid and tanh, cross-entropy, training loops, and held-out evaluation.

Lesson plan

Generate the sequence task and calculate one gate-controlled memory update by hand.
Build the LSTM classifier, trace state shapes, and train with gradient clipping.
Compare held-out results with the stronger baseline, then investigate controlled failures.

Prerequisites

Read Lessons 3, 4, 5, 8, and 9, plus the learned-embeddings build. You should know how integer IDs enter nn.Embedding, how cross-entropy compares logits with class IDs, and why training and validation data must stay separate. This page uses a model class. Recall the Lesson 3 bridge: __init__ stores layers in a reusable object; forward defines the tensor calculation. PyTorch registers a layer when it is stored as self.layer, so its parameters appear in model.parameters().

Plain-language start: carry one instruction

Each sequence begins with a control token. ID 0 means “count up.” ID 1 means “count down.” The next token is a number from zero to seven, stored as IDs 2 to 9. Later values wrap around. If the control says up, 6 is followed by 7 and then 0. If it says down, 1 is followed by 0 and then 7.

At a later position, the current number alone is not enough. After value 3, the correct next value may be 4 or 2. The answer depends on the first control token. An LSTM must carry information about that early token through its hidden and cell states. This is a small but genuine sequence-memory problem.

Concrete input, target, output, and success

Consider the full token sequence [0, 4, 5, 6, 7]. ID 0 says count up. IDs 4, 5, 6, 7 represent values 2, 3, 4, 5. We make next-step examples by shifting:

input:  [0, 4, 5, 6]
target: [4, 5, 6, 7]

For B = 64 examples and T = 10 positions, input and target tensors both have shape (64, 10). Both use torch.long. The model produces logits of shape (64, 10, 10): batch, time, vocabulary. At every position there are ten unnormalized scores.

The first target is the random starting number, so the model cannot predict it from the control token. Later targets are predictable if the model knows the direction. Success means low held-out cross-entropy and high token accuracy, while honestly retaining that irreducible first-position uncertainty.

Build the generated dataset

inputs, targets = make_sequence_data(
    count=512,
    sequence_length=10,
    seed=30,
)

assert inputs.shape == (512, 10)
assert targets.shape == (512, 10)
assert torch.equal(inputs[:, 1:], targets[:, :-1])

The final assertion checks the next-step shift. Training uses seed 30. Validation uses seed 40 and 256 new generated sequences. The sets may contain the same possible rule combinations because the task has only two directions and eight starts, but validation examples are generated independently and never used for parameter updates. The goal is to recover the known rule, not memorize unique natural sentences.

Why an LSTM has gates

Predict one carry: an old cell value of 3 with forget gate 0.25 contributes 0.75 to the next cell before the input contribution. The gate scales memory; it does not label it as simply kept or deleted.

A basic recurrent network repeatedly rewrites one hidden state. An LSTM introduces a cell state c, which is a longer memory path, and a hidden state h, which is the exposed output. At time t, the layer combines current input vector x_t with previous hidden state h_(t-1). One affine calculation creates four raw vectors. Sigmoid and tanh turn them into gates:

i_t = sigmoid(W_i x_t + U_i h_(t-1) + b_i)   input gate
f_t = sigmoid(W_f x_t + U_f h_(t-1) + b_f)   forget gate
g_t = tanh(   W_g x_t + U_g h_(t-1) + b_g)   candidate
o_t = sigmoid(W_o x_t + U_o h_(t-1) + b_o)   output gate

c_t = f_t * c_(t-1) + i_t * g_t
h_t = o_t * tanh(c_t)

The star is element-wise multiplication. Sigmoid values lie from zero to one. A forget value near one preserves old memory; near zero removes it. The input gate controls new candidate information. The output gate controls which cell information appears in the hidden state. These are learned soft controls, not human-written decisions and not guarantees of perfect memory.

Implement one gate update before using the layer

def manual_lstm_cell(x, previous_hidden, previous_cell,
                     input_weight, hidden_weight, bias):
    gates = x @ input_weight.T + previous_hidden @ hidden_weight.T + bias
    input_raw, forget_raw, candidate_raw, output_raw = gates.chunk(4, dim=-1)

    input_gate = torch.sigmoid(input_raw)
    forget_gate = torch.sigmoid(forget_raw)
    candidate = torch.tanh(candidate_raw)
    output_gate = torch.sigmoid(output_raw)

    cell = forget_gate * previous_cell + input_gate * candidate
    hidden = output_gate * torch.tanh(cell)
    return hidden, cell

Suppose the batch is 3, input size is 4, and hidden size is 5. Then x is (3, 4), each old state is (3, 5), input weight is (20, 4), recurrent weight is (20, 5), and bias is (20,). Why 20? Four gates times hidden size five. The combined gates tensor is (3, 20). chunk(4) produces four (3, 5) tensors. The new hidden and cell states are each (3, 5).

The complete script returns the gates too, which lets tests assert that sigmoid gates remain between zero and one. Production code uses PyTorch’s optimized nn.LSTM, but writing this one update removes the mystery from the layer.

The sequence model class

class NextStepLSTM(nn.Module):
    def __init__(self, embedding_size=8, hidden_size=24):
        super().__init__()
        self.embedding = nn.Embedding(10, embedding_size)
        self.lstm = nn.LSTM(embedding_size, hidden_size, batch_first=True)
        self.output = nn.Linear(hidden_size, 10)

    def forward(self, token_ids):
        vectors = self.embedding(token_ids)       # (B, T, 8)
        states, _ = self.lstm(vectors)             # (B, T, 24)
        return self.output(states)                  # (B, T, 10)

batch_first=True makes input and output use batch, time, feature order. It does not change the final-state order. If we kept the returned state pair, h_n and c_n would each have shape (1, B, 24): layers, batch, hidden features. The underscore means this build does not need the final pair because it predicts from every time position in states.

The output linear layer acts on the final axis. It processes each of the B * T hidden vectors independently and changes 24 hidden features into 10 vocabulary logits. It does not mix positions.

Loss, optimization, and gradient clipping

logits = model(inputs)  # (64, 10, 10)
loss = F.cross_entropy(
    logits.reshape(-1, 10),   # (640, 10)
    targets.reshape(-1),      # (640,)
)

optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

Cross-entropy expects one class-score row per target. Reshaping joins batch and time while preserving the vocabulary axis. Gradient clipping scales an overly large combined gradient norm down to one. It is a safety measure for recurrent training; it does not fix wrong data, wrong targets, or a bad learning rate.

The script samples deterministic batches of 64 from 512 training sequences and runs 300 Adam steps. It sets one CPU thread so this tiny workload avoids thread overhead. Validation uses model.eval() and torch.no_grad().

A stronger baseline than random guessing

The memoryless baseline should receive the same current token but no earlier state. If the LSTM does not beat it on held-out sequences, we have no evidence that memory helped. Check this comparison before inspecting attractive generated examples.

Uniform guessing across ten tokens has loss ln(10), about 2.303. That is too weak. Our baseline knows the sequence grammar but forgets the first control token. At position zero it assigns probability 1/8 to each possible starting number. At later positions it knows the next value is either one step up or one step down, so it assigns probability 1/2 to each.

baseline loss = [ln(8) + (T - 1) * ln(2)] / T
              = 0.8318 when T = 10

Beating this baseline is evidence that the model uses direction information from earlier in the sequence. It is more informative than claiming victory over random ten-way predictions.

Run the complete program

Download build_lstm.py. It includes generated data, the manual gate calculation, model, baseline, training, evaluation, assertions, and command-line controls.

python examples/pytorch/build_lstm.py
python examples/pytorch/build_lstm.py --steps 120
python examples/pytorch/build_lstm.py --smoke-test

The smoke command runs at most three updates and skips learning thresholds. It only checks the path through the program. On the tested CPU runtime, the full 300-step command produced:

input and target shapes: (256, 10) (256, 10)
logit shape: (4, 10, 10)
memoryless baseline loss: 0.8318
held-out loss: 0.2156
held-out token accuracy: 0.913

Nine of ten targets are predictable. At the first position, the best predictor has a one-in-eight chance. The expected ceiling is therefore (9 + 1/8) / 10 = 91.25%. A finite sample can lie slightly above or below that value. The main result is that held-out loss is far below the memoryless baseline.

Failure drills

Drill 1: remove the control token. Replace the first input with a constant after generating data. After seeing only one number, the model cannot know whether to count up or down. After seeing two consecutive numbers, it can infer the direction again. Measure loss separately at each position. Do not expect the whole sequence to remain ambiguous or the total loss to reach the fully memoryless baseline. To test a truly memoryless model, reset its hidden and cell states before every token.

Drill 2: forget batch_first=True. The default LSTM expects time first. A tensor shaped (B, T, D) may be interpreted as (T, B, D). The program can run while mixing your intended axes. Fix the flag or transpose explicitly, then assert shapes.

Drill 3: compare unshifted targets. Set targets equal to inputs. The model learns to copy the current token and reports a misleadingly easy score. Always inspect one input-target pair and assert the shift.

Drill 4: apply softmax before cross-entropy. Cross-entropy expects raw logits and performs a stable log-softmax itself. Passing probabilities changes the calculation and weakens gradients. Remove the extra softmax.

Solved exercises

1. What are the shapes for B=16, T=7, embedding 12, hidden 20?

IDs are (16, 7). Embeddings are (16, 7, 12). LSTM outputs are (16, 7, 20). One-layer final hidden and cell states are each (1, 16, 20). Vocabulary logits remain (16, 7, 10).

2. Why does the baseline use ln(2) after the first position?

Without remembering direction, it sees two equally plausible next values: one step up and one step down. The negative log-likelihood of the correct event with probability one half is -ln(1/2) = ln(2).

3. How would two LSTM layers change state shapes?

Set num_layers=2. The sequence output still has final feature size 24. Both final state tensors become (2, B, 24). The first axis now stores one final state per layer.

4. Why not report only training loss?

A model can memorize training examples or exploit an accidental pattern. Held-out data checks the learned rule on examples excluded from optimizer updates. Comparing it with a relevant baseline gives the number context.

Limits and next step

This experiment proves that a small LSTM can carry one control signal across a short generated sequence and beat a baseline that cannot remember it. It does not prove performance on natural language, very long dependencies, variable-length padding, or noisy real data. Longer tasks can still cause optimization problems, and recurrent processing is sequential. The next build changes direction: instead of predicting a sequence, it learns a probability distribution in a latent space.

← Previous build: embeddings Next build: VAE →