Building an LSTM Classifier

The Goal: From a Sentence to a Verdict

In section 1 we turned raw text into a sequence of embedding vectors. Beautiful numbers, but still useless on their own — a sequence of vectors is not a verdict. The job of a text classifier is to read a whole sentence and emit a single symbol: "positive", "spam", "urgent", "French". The question of this section is: how do we compress a variable-length sequence of embeddings into one fixed-size prediction?

The recipe we are going to build is the canonical one that dominated NLP from about 2014 until the transformer era:

1. Each token is looked up in an embedding table and becomes a dense vector $x_t \in \mathbb{R}^{d_\text{embed}}$.
2. An LSTM reads the sequence left to right, updating its hidden state $h_t$ and cell state $c_t$ at every step.
3. A pooling step collapses the sequence of hidden states $(h_0, \ldots, h_{T-1})$ into a single summary vector.
4. A linear layer + softmax turns that summary into class probabilities.

Why This Matters: Almost every "classical" NLP benchmark — sentiment, topic, spam, language identification — was solved to super-human accuracy by exactly this architecture. Today it has been replaced at the top of the leaderboard by transformers, but it is still deployed in countless real products because it is cheap, fast, and good enough for the majority of tasks.

Anatomy of an LSTM Classifier

Conceptually, the classifier is a three-stage pipeline. Data flows left-to-right, dimensions change at each stage, and nothing flows back to earlier stages during inference.

B is the batch size, T is the sequence length, d_embed is the embedding dimension, d_hidden is the LSTM hidden size, and n_classes is the number of labels. Every one of those numbers is a hyperparameter you will tune.

Hover any box in the diagram below to see the learnable parameter matrices that live at that stage. The counts use the production defaults we will train in section 3: $B=64,\ T=200,\ d_\text{embed}=128,\ d_\text{hidden}=256$.

The Mathematical Recipe, End to End

Let us write the whole classifier as a composition of functions. Given a sentence of token ids $(\text{id}_0, \ldots, \text{id}_{T-1})$:

1. Embedding lookup

$x_t = E[\text{id}_t] \in \mathbb{R}^{d_\text{embed}}$. We met this operation in section 1 — a single row-lookup that runs in $O(T \cdot d_\text{embed})$ time.

2. LSTM recurrence

This is the gated recurrence introduced by Hochreiter & Schmidhuber (1997) that solves the vanishing-gradient problem of vanilla RNNs. For $t = 0, 1, \ldots, T-1$, starting from $h_{-1} = c_{-1} = 0$:

Gate pre-activations (one matmul computes all four gates): $\begin{bmatrix} a_i \\ a_f \\ a_g \\ a_o \end{bmatrix} = W_x \, x_t + W_h \, h_{t-1} + b$

Gate activations: $i_t = \sigma(a_i),\; f_t = \sigma(a_f),\; g_t = \tanh(a_g),\; o_t = \sigma(a_o)$. The three sigmoids produce soft switches in $(0, 1)$; tanh produces the candidate content in $(-1, 1)$.

State updates — the one piece of the LSTM that is purely additive, and therefore the reason gradients do not vanish: $c_t = f_t \odot c_{t-1} + i_t \odot g_t$ (cell) and $h_t = o_t \odot \tanh(c_t)$ (hidden). $\odot$ is element-wise multiplication.

3. Pooling

Choose one of three canonical recipes. The simplest is $\bar{h} = h_{T-1}$ — just take the final hidden state. We will compare the alternatives in the next section.

4. Classifier head

$z = W_\text{cls} \, \bar{h} + b_\text{cls} \in \mathbb{R}^{n_\text{classes}}$, then $p = \text{softmax}(z)$. The predicted label is $\hat{y} = \arg\max_k p_k$.

5. Loss

Training uses cross-entropy against the true label $y$: $\mathcal{L} = -\log p_y = -\log \frac{e^{z_y}}{\sum_k e^{z_k}}$. Section 3 of this chapter will walk through the training loop; here we focus exclusively on the forward pass.

Before we roll the four gates over a real sentence, anchor each one in isolation. Drag any slider; the three sigmoids produce soft switches in $(0, 1)$, while tanh produces candidate content in $(-1, 1)$. The cell update $c_t = f_t \odot c_{t-1} + i_t \odot g_t$ multiplies these outputs together elementwise.

Now zoom out to the cell-state highway. Each of the $d_\text{hidden}$ channels has its own lane; memory from the previous step persists through the forget gate, and new content rides in via the input gate. Drag $f_t$ close to 1 and the original memory survives every step; drop it near 0 and the lane clears. This additive path is why LSTM gradients do not vanish.

Pooling: Last, Mean, or Max Hidden State?

The pooling step is where a surprising amount of modelling judgment lives. Given the full trajectory $H = (h_0, h_1, \ldots, h_{T-1}) \in \mathbb{R}^{T \times d_\text{hidden}}$, we need one vector.

The comparison below pins the idea to the math. The left panel is a canned hidden-state trajectory $H \in \mathbb{R}^{T \times d_\text{hidden}}$; the right panel shows the pooled vector for whichever strategy you select. Notice how max-pool captures a spike that mean-pool dilutes, and how attention-pool biases toward the end of the sequence for this particular query.

The attention pooling row is the crack through which transformers crawled into NLP. An attention pool is literally a single-head attention block with a learned query — once you accept that the sum $\sum_t \alpha_t h_t$ is a good idea, you are two steps away from dropping the LSTM entirely and letting every token attend to every other token.

Interactive Visualization: A Sentence Through an LSTM

The animation below lets you scrub through a sentence token by token and watch the hidden state $h_t$ evolve. Two sentences are pre-wired — an obviously positive one and an obviously negative one. The 4-bar histogram above the active cell is the current hidden vector; the yellow text below is the cell state $c_t$. When you reach the last token, the classifier head appears on the right and softmax turns $h_T$ into a probability over positive and negative.

Three things to look for while you play with the slider:

• Slow start. At $t = 0$ both sentences share the same token ("I") and therefore the same tiny hidden activations. The network has almost no opinion yet.
• The decisive token. At $t = 2$ ("love" or "hate") the trajectories split dramatically — the sign of the hidden bars flips. This is the cell-state additive highway at work: one strong token can rewrite memory in a single step.
• Noisy coda. The last two tokens ("this movie") are shared again, yet the two trajectories stay separated. The LSTM has committed. A transformer would let those later tokens re-weight the verdict via attention.

Building the Classifier From Scratch in NumPy

Before we hand the problem to PyTorch we are going to build the full classifier in pure NumPy. You will see that there is nothing magical inside nn.LSTM — just the four-gate recipe from the previous section, rolled over time with a Python for loop, followed by a single linear projection. Every matrix, every gate, every state update is visible.

The example uses $d_\text{embed} = 2$, $d_\text{hidden} = 2$, and a four-word sentence so every intermediate value fits comfortably on one line. Click any line number on the right to see what that line computes.

          d0    d1
<pad>   0.00  0.00
this    0.10  0.20
movie   0.40  0.30
is      0.20  0.10
great   0.90  0.70
awful  -0.80 -0.60

          d0    d1
this    0.10  0.20
movie   0.40  0.30
is      0.20  0.10
great   0.90  0.70

         d0    d1
i[0]    0.5   0.3
i[1]    0.1  -0.2
f[0]    0.4   0.6
f[1]    0.2   0.1
g[0]    0.7   0.5
g[1]    0.3   0.4
o[0]    0.2  -0.1
o[1]    0.0   0.3

all zeros

           h[0]    h[1]
positive  -1.00   1.20
negative   1.10  -0.90

1import numpy as np
2
3# --- 1. A tiny vocabulary and three embeddings (d_embed = 2) ---
4vocab = {'<pad>': 0, 'this': 1, 'movie': 2, 'is': 3, 'great': 4, 'awful': 5}
5E = np.array([
6    [0.00, 0.00],   # <pad>
7    [0.10, 0.20],   # this
8    [0.40, 0.30],   # movie
9    [0.20, 0.10],   # is
10    [0.90, 0.70],   # great  — pulls activations positive
11    [-0.80, -0.60], # awful  — pulls activations negative
12])
13
14# --- 2. Input sentence, tokenised + mapped to ids ---
15sentence = "this movie is great"
16input_ids = np.array([vocab[w] for w in sentence.split()])
17X = E[input_ids]                       # (T=4, d_embed=2)
18
19# --- 3. LSTM parameters (d_hidden = 2) ---
20d_embed, d_hidden = 2, 2
21# 4 gates stacked: [input, forget, cell, output] → rows = 4 * d_hidden
22W_x = np.array([                       # (4*d_hidden, d_embed) = (8, 2)
23    [ 0.5,  0.3], [ 0.1, -0.2],        # input gate    rows 0-1
24    [ 0.4,  0.6], [ 0.2,  0.1],        # forget gate   rows 2-3
25    [ 0.7,  0.5], [ 0.3,  0.4],        # cell cand.    rows 4-5
26    [ 0.2, -0.1], [ 0.0,  0.3],        # output gate   rows 6-7
27])
28W_h = np.zeros((4 * d_hidden, d_hidden))   # simplified — no recurrent bias here
29b   = np.zeros(4 * d_hidden)
30
31# --- 4. Classifier head: W_cls (2, d_hidden) + bias ---
32W_cls = np.array([[-1.0,  1.2],        # logit for 'positive'
33                  [ 1.1, -0.9]])        # logit for 'negative'
34b_cls = np.array([0.0, 0.0])
35
36def sigmoid(x: np.ndarray) -> np.ndarray:
37    return 1.0 / (1.0 + np.exp(-x))
38
39def lstm_step(x_t: np.ndarray, h_prev: np.ndarray, c_prev: np.ndarray):
40    """One LSTM cell step — the four-gate recipe."""
41    gates = W_x @ x_t + W_h @ h_prev + b     # (8,)
42    i = sigmoid(gates[0:2])                  # input gate   (2,)
43    f = sigmoid(gates[2:4])                  # forget gate  (2,)
44    g = np.tanh(gates[4:6])                  # cell cand.   (2,)
45    o = sigmoid(gates[6:8])                  # output gate  (2,)
46    c_t = f * c_prev + i * g                 # cell state   (2,)
47    h_t = o * np.tanh(c_t)                   # hidden state (2,)
48    return h_t, c_t
49
50# --- 5. Roll the LSTM over the sentence ---
51T = X.shape[0]
52h = np.zeros(d_hidden)
53c = np.zeros(d_hidden)
54hidden_history = []
55for t in range(T):
56    h, c = lstm_step(X[t], h, c)
57    hidden_history.append(h.copy())
58
59# --- 6. Take the final hidden state and classify ---
60h_T     = hidden_history[-1]             # (d_hidden,)
61logits  = W_cls @ h_T + b_cls             # (2,)
62probs   = np.exp(logits - logits.max())
63probs  /= probs.sum()                     # softmax → (2,)
64pred    = int(np.argmax(probs))
65label   = ['positive', 'negative'][pred]
66
67print("h_T      :", h_T)
68print("logits   :", logits)
69print("probs    :", probs)
70print("predicted:", label)

Quick Check: The untrained network above predicts "negative" for "this movie is great". That is not a bug — random weights give random verdicts. Training (chapter 18 §3) is what shapes $W_x, W_h, W_\text{cls}$ so that $h_T$ aligns with the correct class direction.

The Same Model in PyTorch

Now we wrap the same recipe in an nn.Module. Three things to notice as you read the PyTorch version:

1. The shapes are identical to the NumPy version — we simply trade np.array for torch.Tensor.
2. The entire LSTM recurrence collapses into one line: self.lstm(x). PyTorch uses a fused cuDNN kernel that runs on GPU at roughly 100× the speed of the NumPy loop.
3. We return logits, not probabilities, because nn.CrossEntropyLoss expects raw logits and performs a numerically-stable log-softmax internally.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class LSTMClassifier(nn.Module):
6    def __init__(self, vocab_size: int, d_embed: int,
7                 d_hidden: int, n_classes: int, pad_idx: int = 0):
8        super().__init__()
9        self.embedding = nn.Embedding(vocab_size, d_embed, padding_idx=pad_idx)
10        self.lstm      = nn.LSTM(d_embed, d_hidden, batch_first=True)
11        self.classifier = nn.Linear(d_hidden, n_classes)
12
13    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
14        x            = self.embedding(input_ids)          # (B, T, d_embed)
15        outputs, (h_n, c_n) = self.lstm(x)                # outputs: (B, T, d_hidden)
16        h_T          = outputs[:, -1, :]                  # last-step pooling
17        logits       = self.classifier(h_T)                # (B, n_classes)
18        return logits
19
20model  = LSTMClassifier(vocab_size=6, d_embed=2, d_hidden=2, n_classes=2)
21ids    = torch.tensor([[1, 2, 3, 4]])                     # 'this movie is great'
22logits = model(ids)
23probs  = F.softmax(logits, dim=-1)
24print("logits:", logits)
25print("probs :", probs)

Making it Better: Bidirectional and Stacked LSTMs

The classifier above has a structural weakness: every hidden state $h_t$ only depends on tokens $0, 1, \ldots, t$. The phrase "not great" is read in order — by the time the LSTM sees "great" it has already committed to a representation of "not". It cannot go back.

The fix is mechanically trivial: run a second LSTM in the reverse direction and concatenate the two hidden state streams. This is a bidirectional LSTM (BiLSTM).

$h_t^\rightarrow = \text{LSTM}_\rightarrow(x_{0:t})$, $h_t^\leftarrow = \text{LSTM}_\leftarrow(x_{t:T})$, and $h_t = [\,h_t^\rightarrow ;\, h_t^\leftarrow\,]$. The classifier head now receives a vector of size $2 d_\text{hidden}$.

In PyTorch the change is a single keyword argument:

1self.lstm = nn.LSTM(d_embed, d_hidden, batch_first=True, bidirectional=True)
2self.classifier = nn.Linear(2 * d_hidden, n_classes)

Stacking is the other standard upgrade. Passing num_layers=2 to nn.LSTM runs a second LSTM on top of the first — the second layer reads the hidden states of the first as its inputs. Two or three layers is typical; more than that rarely helps and routinely hurts because of vanishing gradients across layers (not time).

The Fundamental Bottleneck

Take a second look at the interactive animation. At every time step the LSTM rewrites one fixed-size vector $h_t \in \mathbb{R}^{d_\text{hidden}}$. Regardless of sentence length — 4 tokens or 4,000 — the entire memory of what was read must fit into those few hundred floats.

Empirically, three things go wrong once sequences get long:

• Forgetting. Tokens read 200 steps ago have almost certainly been overwritten. Even with the additive cell highway, gradients of length-200 paths are numerically tiny during training.
• Position blur. "The cat sat on the mat" and "The mat sat on the cat" produce different $h_T$, but not by much — the LSTM has no explicit notion of position, only of order.
• Sequential latency. Step $t$ cannot begin until step $t-1$ has finished. GPUs hate this. For a transformer, all timesteps go through the matmul in parallel.

The original motivation for attention was exactly this bottleneck — Bahdanau et al.'s 2014 paper literally let the decoder "peek" back at the whole encoder trajectory instead of relying on one final state. Eight years and one transformer paper later, that same idea would replace the entire recurrence.

From LSTM Classifier to Transformer Classifier

Conceptually the transformer classifier is the same shape: embed, encode, pool, classify. The only stage that changes is the encoder.

Self-attention is what replaces the recurrence. Instead of threading a single hidden vector through time, every token attends to every other token via queries, keys, and values: $\text{Attn}(Q, K, V) = \text{softmax}(QK^\top / \sqrt{d_k})\, V$. Crucially, all T positions are processed in parallel; there is no sequential loop. That is the superpower that made transformers feasible on modern GPUs.

Multi-head attention runs $h$ such attention functions in parallel on different linear projections of Q, K, V. Each head can specialise — one looks at syntactic parents, another at coreference, another at position. The concatenated heads are projected back to $d_\text{model}$. In an LSTM the hidden state mixes all of those roles into one vector; in multi-head attention they live in disjoint subspaces.

Flash Attention and the Memory Story

Self-attention fixes the LSTM's information bottleneck, but introduces a new one: memory. The naive formulation materialises the $T \times T$ attention matrix $S = QK^\top / \sqrt{d_k}$. For a sequence of 8,192 tokens that is 67 million floats — 256 MB per head per layer, just for intermediates. On a 32-head, 80-layer model you run out of GPU HBM long before you run out of compute.

Flash Attention (Dao et al., 2022) is the fix that made long-context transformers practical. The key insight: we never actually need S in full — we only need its softmax output times V. So we tile Q, K, V into blocks that fit in GPU SRAM, compute the softmax incrementally (online softmax), and accumulate the output without ever writing the full $T \times T$ matrix to HBM.

An LSTM classifier has no analogue of this problem because its state size is $O(d_\text{hidden})$ independent of sequence length. That same fact is why LSTMs cannot reach back: constant state is a feature and the bug.

Positional Encodings and the KV-Cache

A transformer layer is permutation-equivariant — swap any two input tokens and the output swaps along with them. Self-attention sees a bag of vectors, not a sequence. That is a deal-breaker for language, where "dog bites man" is very different from "man bites dog".

Positional encodings are the patch. The original transformer added a fixed sinusoidal vector $\text{PE}(t) \in \mathbb{R}^{d_\text{model}}$ to every token embedding. Modern variants — RoPE (rotary), ALiBi (additive linear biases), learned absolute, T5-relative — all share one goal: inject order information that an LSTM got for free from its recurrence.

Trade-off: an LSTM gets order structure from its sequential computation — free in bits, expensive in wall-clock time. A transformer gets parallelism for free but must pay a parameter budget for positional encodings. Rotary and ALiBi minimise that cost and generalise to sequences longer than training.

KV-cache — reclaiming incremental generation

There is one thing LSTMs do effortlessly that transformers struggle with: autoregressive generation. An LSTM, at decoding time, only needs to keep the last $(h_t, c_t)$ pair in memory. A transformer, naively, re-computes attention over the entire prefix at every new token — an $O(T^2)$ blowup.

The KV-cache is the standard fix. Observe that once a token's key and value vectors are computed, they never change — they depend only on that token's embedding and position. So we cache them. At generation step $t$, we compute only the new token's $q_t, k_t, v_t$, append $k_t, v_t$ to the cache, and attend $q_t$ over the full cache.

For a 70-billion-parameter model at 8k context, the KV-cache can itself be 40 GB — which is why production systems use techniques like multi-query attention (one K/V head shared across all Q heads), grouped-query attention, and paged attention (a vLLM innovation) to keep the cache tractable.

The Scaling Story: Why We Left LSTMs Behind

If LSTMs are so elegant, why is almost every state-of-the-art language system a transformer? The answer is one word: parallelism.

On a modern accelerator, wall-clock time is dominated by how well the computation matches the hardware's parallelism. A transformer layer is essentially two giant matmuls per block; an A100 or H100 isdesigned to do exactly those matmuls at peak throughput. An LSTM layer is a tight sequential loop — the GPU's thousands of cores sit mostly idle waiting on a data dependency.

This is the deep reason the field migrated. Not that LSTMs cannot learn the task — for sentiment classification they still do, and well — but that once you want to scale to billions of parameters and hundreds of billions of tokens, an architecture that wastes 70% of your hardware is not tenable. Transformers were the architecture that hugged the hardware.

But LSTMs are not dead. On-device keyboard suggestion, streaming speech recognition, small-footprint sensor models — anywhere latency and energy matter more than peak quality — still use LSTMs or GRUs. And in 2023-2024 the SSM/Mamba family revived constant-state recurrence with transformer-level quality on several benchmarks. The story is not over.

Summary

In this section we built the canonical LSTM text classifier end-to-end:

1. Derived the four-stage architecture (embedding → LSTM → pool → linear + softmax) and wrote it as a single differentiable function.
2. Implemented the entire forward pass in pure NumPy, tracing every gate, cell update and hidden state to numerical values so nothing remains mysterious.
3. Wrapped the same recipe in a PyTorch nn.Module, contrasted its parameter layout, and explained every argument of nn.Embedding, nn.LSTM and nn.Linear.
4. Compared four pooling strategies and upgraded to the standard bidirectional, stacked workhorse.
5. Exposed the fixed-state bottleneck that drove the field toward attention, and traced the line from it to multi-head attention, positional encodings, Flash Attention and the KV-cache.

In section 3 we close the loop: feeding this classifier a labelled dataset, training with cross-entropy and Adam, monitoring F1, and watching the hidden-state trajectories in the visualization above actually learn to separate positive from negative.

References

• Hochreiter, S. & Schmidhuber, J. (1997). Long Short-Term Memory. Neural Computation 9(8), 1735–1780.
• Bahdanau, D., Cho, K. & Bengio, Y. (2014). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015 / arXiv:1409.0473.
• Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS 2017 / arXiv:1706.03762.
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A. & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022 / arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Ainslie, J. et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023 / arXiv:2305.13245.