III · Forward

The forward pass, with its shapes

One modern, no-peephole LSTM layer. This is our batched pedagogical presentation, not a verbatim reproduction of the 1997 or 2013 papers. The translation page explains their architectural differences.

Source family: Graves 2013, §2.1, Eqs. 7–11, with all peepholes removed and the cell candidate named g. Fusing gates, batching and row-oriented storage are presentation/implementation choices. Sources and exact adaptations.

Axis identity	Meaning	Size
$B$	independent sequences in a batch	16
$T$	time steps	128
$E$	input embedding coordinates	512
$D$	recurrent state coordinates (source of the recurrent projection)	768
$R$	projection output coordinates, distinct from D; canonical coordinate-wise bijection R → D	768
$G$	gate labels, in order i, f, g, o	4
$K = G \cdot R$	fused gate-output axis; an atomic identity until unfused	3072 = 4·768
$J = E + D$	input and recurrent state concatenated, in that order	1280

Sizes are not axis identities. R and D both have 768 entries but cannot be silently matched. K has 3072 entries (often called “4D” in code); it does not share a free D or R axis with another tensor. A size-only numerical view is a check against this ledger, never a replacement for its identities.

Parameters

E \times K W, D \times K U, (B) \times K b, ∣ K ∣ = 4∣ R ∣ = 4∣ D ∣ = 3072

W has 1,572,864 entries; U has 2,359,296; the stored bias has 3072. K is a fused destination axis, not the source D. Unfused slices of U have D×R, so the two hidden directions cannot collide by name.

P = 4 D (E + D + 1) = 3072 (512 + 768 + 1) = 3, 935, 232

Here D and E in counts denote cardinalities. PyTorch stores two additive bias vectors: its one-layer parameter count is 4D(E+D+2)=3,938,304. Its stored weights have K×E and K×D layout, the reverse of the layout used here; matching named axes does not depend on their order. See the nn.LSTM equations and variables.

Read one slice of the input

B \times T \times E X, B \times E x_{t} = X_{t}, t = 1, \dots, T

X contains 16×128×512 = 1,048,576 entries. Each time slice is 16×512. Start with $B \times D h_{0} = 0, B \times D c_{0} = 0$ .

Project the input and the previous state

B \times K a_{t} = B \times E x_{t} E \cdot E \times K W

B \times K r_{t} = B \times D h_{t - 1} D \cdot D \times K U

In numbers: (16×512) times (512×3072), and (16×768) times (768×3072). Both yield 16×3072. The source E or D is summed; K survives. Only the second depends on the previous step.

B \times T \times K A = B \times T \times E X E \cdot W

The input projection can be computed across all times at once; this is an algebraic possibility, not a claim that every backend always does so.

Add the explicitly broadcast bias

B \times K z_{t} = a_{t} + r_{t} + (B) \times K b

The same 3072 bias numbers are added to each batch row. The equivalent concatenated calculation is:

B \times J v_{t} = cat_{E, D \to J} (x_{t}, h_{t - 1}), J \times K V = [W U], B \times K z_{t} = v_{t} J \cdot V + b

J=E+D=1280. This is (16×1280) times (1280×3072). Concatenation chooses the input-first ordering explicitly.

Unpack and rename the gate coordinates

(B \times D q_{t}^{i}, B \times D q_{t}^{f}, B \times D q_{t}^{g}, B \times D q_{t}^{o}) = unpack_{K \to G, R \to D} (z_{t})

Unpack means: unfuse K into G×R, select each gate (i,f,g,o), then apply the coordinate-wise bijection ρ from R to D. Before the rename, every gate slice is B×R. Afterwards it is B×D. This explicit map is required even though |R|=|D|=768; no D×D dense map is hiding here.

Apply the element-wise nonlinearities

B \times D i_{t} = σ (q_{t}^{i}), B \times D f_{t} = σ (q_{t}^{f}), B \times D g_{t} = tanh (q_{t}^{g}), B \times D o_{t} = σ (q_{t}^{o})

All four have 16×768 entries. i, f and o are in (0,1); g is in (−1,1).

Update the cell and read it out

B \times D c_{t} = f_{t} ⊙ c_{t - 1} + i_{t} ⊙ g_{t}, B \times D h_{t} = o_{t} ⊙ tanh (c_{t})

There is an additive, coordinate-wise route from c at one step to the next. That route can help preserve gradients when forget gates stay near one. It does not prevent vanishing gradients: repeated forget gates below one attenuate it. Gates also depend on h, which introduces other recurrent paths. The backward page distinguishes these derivatives precisely.

Repeat and stack

B \times T \times D H = stack_{T} (h_{1}, \dots, h_{T})

The output is 16×128×768 = 1,572,864 entries (6.29 MB in fp32). The two matrix products cost about 2BK(E+D)=125,829,120 FLOPs per batch step, or 16.1 GFLOPs over the whole batch sequence. This ignores nonlinearities and element-wise work.

The short form

All gate and preactivation symbols are declared as B×D families above. E1–E5 allow:

B \times K z_{t} = x_{t} \cdot W + h_{t - 1} \cdot U + b

(q_{t}^{i}, q_{t}^{f}, q_{t}^{g}, q_{t}^{o}) = unpack_{K \to G, R \to D} (z_{t})

(i_{t}, f_{t}, g_{t}, o_{t}) = (σ (q_{t}^{i}), σ (q_{t}^{f}), tanh (q_{t}^{g}), σ (q_{t}^{o}))

B \times D c_{t} = f_{t} ⊙ c_{t - 1} + i_{t} ⊙ g_{t}, B \times D h_{t} = o_{t} ⊙ tanh c_{t}

A = X @ W                              # B,T,K; K = 4*D
for t in range(T):
    z = A[:, t] + h @ U + b             # B,K; virtual B bias axis
    qi, qf, qg, qo = z.split(D, -1)     # K -> G,R; views R renamed D
    i, f, g, o = qi.sigmoid(), qf.sigmoid(), qg.tanh(), qo.sigmoid()
    c = f * c + i * g                   # B,D
    h = o * c.tanh()                    # B,D
    hs.append(h)
H = torch.stack(hs, dim=1)              # B,T,D