LSTM Cell Mathematics

What §3.3 Established

Section 3.3 introduced the LSTM cell as a working-memory analogue: three gates (input / forget / output) plus a candidate state, with the cell-state update $c_t = f_t \odot c_{t-1} + i_t \odot g_t$ being the magic line that lets gradients flow through long sequences. We worked through a SCALAR cell to keep the math visible. This section shows the same algebra in vector form - the production-grade calculation that nn.LSTM does internally.

From Scalar Cell to Vector Cell

With $x_t \in \mathbb{R}^D$ and $h_t, c_t \in \mathbb{R}^H$, every gate becomes a linear-plus-non-linearity from $(x_t, h_{t-1})$:

$i_t = \sigma\!\bigl(W_{ix} x_t + W_{ih} h_{t-1} + b_i\bigr)$

$f_t = \sigma\!\bigl(W_{fx} x_t + W_{fh} h_{t-1} + b_f\bigr)$

$g_t = \tanh\!\bigl(W_{gx} x_t + W_{gh} h_{t-1} + b_g\bigr)$

$o_t = \sigma\!\bigl(W_{ox} x_t + W_{oh} h_{t-1} + b_o\bigr)$

Then $c_t = f_t \odot c_{t-1} + i_t \odot g_t$ and $h_t = o_t \odot \tanh(c_t)$. The only difference from §3.3 is that everything is now a coordinate-wise vector op of length H, and the W matrices have shape (H, D) or (H, H).

Interactive: Step-Trace (Recap)

The 8-cycle pulse trace from §3.3, reproduced. The vector form does the same thing per coordinate that the scalar form did once.

Python: Vector LSTM Cell From Scratch

The class below packs all 4 gate weights into one matrix and does the gate computation in one matmul (matching nn.LSTM's internal layout). Run on a 30-cycle, 64-D input with hidden size 256.

1import numpy as np
2
3# A vector LSTM cell — generalises the scalar version from §3.3.
4# Inputs and states are now vectors of length H; weights are matrices.
5class LSTMCellVector:
6    def __init__(self, input_size: int, hidden_size: int, seed: int = 0):
7        rng = np.random.default_rng(seed)
8        H, D = hidden_size, input_size
9        # PyTorch packs all 4 gates into one weight matrix of shape (4H, ·).
10        # Layout: rows 0..H-1 = i, H..2H-1 = f, 2H..3H-1 = g, 3H..4H-1 = o.
11        self.W_x = rng.standard_normal((4 * H, D)).astype(np.float32) / np.sqrt(D)
12        self.W_h = rng.standard_normal((4 * H, H)).astype(np.float32) / np.sqrt(H)
13        self.b   = np.zeros(4 * H, dtype=np.float32)
14        self.b[H:2*H] = 1.0          # forget bias = 1, like §3.3
15        self.H = H
16
17    def step(self, x: np.ndarray, h_prev: np.ndarray, c_prev: np.ndarray):
18        sig = lambda z: 1.0 / (1.0 + np.exp(-z))
19        H = self.H
20
21        # ALL 4 gates in one matmul
22        z = self.W_x @ x + self.W_h @ h_prev + self.b   # (4H,)
23
24        i = sig(z[:H])
25        f = sig(z[H:2*H])
26        g = np.tanh(z[2*H:3*H])
27        o = sig(z[3*H:4*H])
28
29        c = f * c_prev + i * g
30        h = o * np.tanh(c)
31        return h, c, (i, f, g, o)
32
33
34# ----- Run on a 30-cycle window -----
35np.random.seed(0)
36T, D, H = 30, 64, 256
37cell = LSTMCellVector(input_size=D, hidden_size=H, seed=0)
38seq  = np.random.randn(T, D).astype(np.float32) * 0.5
39h, c = np.zeros(H, dtype=np.float32), np.zeros(H, dtype=np.float32)
40
41for t in range(T):
42    h, c, gates = cell.step(seq[t], h, c)
43
44print("h.shape :", h.shape)         # (256,)
45print("c.shape :", c.shape)         # (256,)
46print("|h|     :", np.linalg.norm(h).round(3))
47print("|c|     :", np.linalg.norm(c).round(3))

PyTorch: nn.LSTMCell vs nn.LSTM

Two APIs do almost the same thing. nn.LSTMCell does ONE timestep at a time and is useful for custom recurrent loops. nn.LSTM processes a whole sequence and invokes the cuDNN fused kernel - always prefer it for normal sequence-to-sequence work.

1import torch
2import torch.nn as nn
3
4# nn.LSTMCell: ONE timestep at a time
5torch.manual_seed(0)
6cell = nn.LSTMCell(input_size=64, hidden_size=256)
7
8x_t   = torch.randn(2, 64)            # (B, D) at one timestep
9h     = torch.zeros(2, 256)
10c     = torch.zeros(2, 256)
11h, c  = cell(x_t, (h, c))
12print("after cell h.shape:", tuple(h.shape))    # (2, 256)
13print("after cell c.shape:", tuple(c.shape))    # (2, 256)
14
15
16# nn.LSTM: WHOLE sequence in one call
17seq   = torch.randn(2, 30, 64)        # (B, T, D)
18lstm  = nn.LSTM(input_size=64, hidden_size=256, num_layers=1, batch_first=True)
19out, (h_n, c_n) = lstm(seq)
20
21print("\nseq                :", tuple(seq.shape))
22print("LSTM out           :", tuple(out.shape))     # (2, 30, 256)
23print("LSTM h_n           :", tuple(h_n.shape))     # (1, 2, 256)

Parameter Accounting

Per LSTM layer with hidden size $H$ and input dim $D$:

$P = 4H(D + H + 1).$

For our backbone $D = 64, H = 256$: $4 \cdot 256 \cdot (64 + 256 + 1) = 4 \cdot 256 \cdot 321 = 328{,}704$ parameters per direction per layer. Two directions x two layers = ~1.3M, but the second layer's input dim is 2H=512 (it receives the bidirectional concat from layer 1) so the accounting balloons to ~2.1M total. Section 9.4 walks through the exact numbers.

Two Math Pitfalls

The point. The LSTM cell is six lines of vector math. Production code packs the gate weights into one (4H, D) matrix and does the four gate pre-activations in a single matmul. Everything else - bidirectional, multi-layer - is wrapping around this six-line core.

Takeaway

• Vector LSTM cell = scalar LSTM cell, vectorised. Same algebra; H-dim element-wise ops.
• Per-layer parameter count is 4H(D+H+1). For our backbone this is ~328k per direction per layer.
• nn.LSTM >> loop over nn.LSTMCell. The cuDNN kernel is 5-10× faster.
• Pack the 4 gate weights into one matrix. One matmul; slice to recover individual gates.