Putting all together

The four gates of the previous notes (forget, input, candidate update, and output) act on a single shared input pair $({\mathbf{x}}_t,{\mathbf{h}}_{t-1})$ and on a single shared memory ${\mathbf{c}}_{t-1}$. This note collects the complete forward pass of one LSTM cell in one place, then states the mini-batch form, the parameter count, and the initialization conventions used throughout the rest of the section.

One cell, one time step

Given ${\mathbf{x}}_t\in {\mathbb{R}}^{n_{\text{inputs}}}$, ${\mathbf{h}}_{t-1}\in {\mathbb{R}}^{n_{\text{neurons}}}$ and ${\mathbf{c}}_{t-1}\in {\mathbb{R}}^{n_{\text{neurons}}}$, the LSTM cell computes

$$\begin{array}{rlrl}{\mathbf{f}}_t& =\sigma ({\mathbf{W}}_{xf}{\mathbf{x}}_t+{\mathbf{W}}_{hf}{\mathbf{h}}_{t-1}+{\mathbf{b}}_f)& & \text{forget gate}\\ {\mathbf{i}}_t& =\sigma ({\mathbf{W}}_{xi}{\mathbf{x}}_t+{\mathbf{W}}_{hi}{\mathbf{h}}_{t-1}+{\mathbf{b}}_i)& & \text{input gate}\\ {\tilde{\mathbf{c}}}_t& =\tanh ({\mathbf{W}}_{xc}{\mathbf{x}}_t+{\mathbf{W}}_{hc}{\mathbf{h}}_{t-1}+{\mathbf{b}}_c)& & \text{candidate update}\\ {\mathbf{o}}_t& =\sigma ({\mathbf{W}}_{xo}{\mathbf{x}}_t+{\mathbf{W}}_{ho}{\mathbf{h}}_{t-1}+{\mathbf{b}}_o)& & \text{output gate}\\ {\mathbf{c}}_t& ={\mathbf{f}}_t\odot {\mathbf{c}}_{t-1}+{\mathbf{i}}_t\odot {\tilde{\mathbf{c}}}_t& & \text{cell-state update}\\ {\mathbf{h}}_t& ={\mathbf{o}}_t\odot \tanh ({\mathbf{c}}_t)& & \text{hidden-state read-out}\end{array}$$

and returns $({\mathbf{h}}_t,{\mathbf{c}}_t)$.

The first four lines are all the same affine-then-pointwise pattern, with separate parameters per gate and the same input pair going into each. The last two lines are the entire dynamics of the cell: a gated additive update on the memory line, followed by a gated, bounded read-out of the hidden state.

The whole architecture in two equations

Strip the gate definitions away and the LSTM reduces to

$${\mathbf{c}}_t={\mathbf{f}}_t\odot {\mathbf{c}}_{t-1}+{\mathbf{i}}_t\odot {\tilde{\mathbf{c}}}_t,{\mathbf{h}}_t={\mathbf{o}}_t\odot \tanh ({\mathbf{c}}_t).$$

Everything else (the four sigmoid/tanh layers, the shared input, the four times larger parameter count compared with a vanilla RNN) exists to make these two element-wise equations behave well. The first guarantees that the gradient survives across long gaps, because its Jacobian with respect to ${\mathbf{c}}_{t-1}$ is $\mathrm{diag}({\mathbf{f}}_t)$. The second guarantees that the hidden state is bounded and selectively informative. Reading the LSTM as “two equations with learned coefficients” makes the design intent unambiguous.

Mini-batch form

For a mini-batch of $B$ sequences processed in parallel, every per-example vector becomes a matrix with one column per example. The parameters are shared across the batch (and across the unrolled time axis), so their shapes are unchanged. The cell becomes

$$\begin{array}{rl}{\mathbf{F}}_t& =\sigma ({\mathbf{W}}_{xf}{\mathbf{X}}_t+{\mathbf{W}}_{hf}{\mathbf{H}}_{t-1}+{\mathbf{b}}_f),\\ {\mathbf{I}}_t& =\sigma ({\mathbf{W}}_{xi}{\mathbf{X}}_t+{\mathbf{W}}_{hi}{\mathbf{H}}_{t-1}+{\mathbf{b}}_i),\\ {\tilde{\mathbf{C}}}_t& =\tanh ({\mathbf{W}}_{xc}{\mathbf{X}}_t+{\mathbf{W}}_{hc}{\mathbf{H}}_{t-1}+{\mathbf{b}}_c),\\ {\mathbf{O}}_t& =\sigma ({\mathbf{W}}_{xo}{\mathbf{X}}_t+{\mathbf{W}}_{ho}{\mathbf{H}}_{t-1}+{\mathbf{b}}_o),\\ {\mathbf{C}}_t& ={\mathbf{F}}_t\odot {\mathbf{C}}_{t-1}+{\mathbf{I}}_t\odot {\tilde{\mathbf{C}}}_t,\\ {\mathbf{H}}_t& ={\mathbf{O}}_t\odot \tanh ({\mathbf{C}}_t).\end{array}$$

All element-wise operations apply column-wise, so the entire batch is processed in parallel as a small number of dense matrix multiplications. This is the form that actually runs on hardware.

Dimensions used in this note (master table for the section)

symbol	role	single example	mini-batch ($B$)
${\mathbf{x}}_t,{\mathbf{X}}_t$	input	$n_{\text{inputs}}$	$n_{\text{inputs}}\times B$
${\mathbf{h}}_{t-1},{\mathbf{h}}_t,{\mathbf{H}}_{t-1},{\mathbf{H}}_t$	hidden state	$n_{\text{neurons}}$	$n_{\text{neurons}}\times B$
${\mathbf{c}}_{t-1},{\mathbf{c}}_t,{\mathbf{C}}_{t-1},{\mathbf{C}}_t$	cell state	$n_{\text{neurons}}$	$n_{\text{neurons}}\times B$
${\mathbf{f}}_t,{\mathbf{i}}_t,{\mathbf{o}}_t,{\mathbf{F}}_t,{\mathbf{I}}_t,{\mathbf{O}}_t$	gates	$n_{\text{neurons}}$	$n_{\text{neurons}}\times B$
${\tilde{\mathbf{c}}}_t,{\tilde{\mathbf{C}}}_t$	candidate update	$n_{\text{neurons}}$	$n_{\text{neurons}}\times B$
${\mathbf{W}}_{x\bullet }$	input-to-hidden weights, $\bullet \in \{f,i,c,o\}$	$n_{\text{neurons}}\times n_{\text{inputs}}$	idem
${\mathbf{W}}_{h\bullet }$	hidden-to-hidden weights	$n_{\text{neurons}}\times n_{\text{neurons}}$	idem
${\mathbf{b}}_{\bullet }$	bias	$n_{\text{neurons}}$	$n_{\text{neurons}}$

Broadcasting: only in the bias addition

Across the entire LSTM cell, the only operation that requires broadcasting is the bias term in each affine map. In mini-batch form, ${\mathbf{b}}_{\bullet }\in {\mathbb{R}}^{n_{\text{neurons}}}$ is added to a matrix of shape $n_{\text{neurons}}\times B$; the bias is broadcast across the $B$ columns.

Every element-wise product $\odot$ in the cell acts on two operands of identical shape:

product	both operands are	shape (single / batch)
${\mathbf{f}}_t\odot {\mathbf{c}}_{t-1}$	gate $\odot$ cell state	$n_{\text{neurons}}$ / $n_{\text{neurons}}\times B$
${\mathbf{i}}_t\odot {\tilde{\mathbf{c}}}_t$	gate $\odot$ candidate	$n_{\text{neurons}}$ / $n_{\text{neurons}}\times B$
${\mathbf{o}}_t\odot \tanh ({\mathbf{c}}_t)$	gate $\odot$ bounded cell state	$n_{\text{neurons}}$ / $n_{\text{neurons}}\times B$

This is not a coincidence: the weight matrices ${\mathbf{W}}_{h\bullet }\in {\mathbb{R}}^{n_{\text{neurons}}\times n_{\text{neurons}}}$ are sized precisely so that every gate output matches the width of the cell state it modulates. The LSTM is designed so that each “switch” is a scalar attached to one slot of memory, and the $\odot$ operations therefore stay shape-aligned by construction.

The fused weight matrix used in practice

Modern frameworks (PyTorch, TensorFlow, cuDNN) store the eight weight matrices as a single fused tensor and perform all four affine maps with one matrix multiplication per source:

$$\left[\begin{array}{c}{\mathbf{Z}}_t^f\\ {\mathbf{Z}}_t^i\\ {\mathbf{Z}}_t^c\\ {\mathbf{Z}}_t^o\end{array}\right]=\left[\begin{array}{c}{\mathbf{W}}_{xf}\\ {\mathbf{W}}_{xi}\\ {\mathbf{W}}_{xc}\\ {\mathbf{W}}_{xo}\end{array}\right]{\mathbf{X}}_t+\left[\begin{array}{c}{\mathbf{W}}_{hf}\\ {\mathbf{W}}_{hi}\\ {\mathbf{W}}_{hc}\\ {\mathbf{W}}_{ho}\end{array}\right]{\mathbf{H}}_{t-1}+\left[\begin{array}{c}{\mathbf{b}}_f\\ {\mathbf{b}}_i\\ {\mathbf{b}}_c\\ {\mathbf{b}}_o\end{array}\right].$$

The four blocks of the result are then sliced apart and passed through their respective activations. This is purely an implementation optimization: it reduces kernel-launch overhead and exposes a single large GEMM to BLAS. Mathematically it is identical to the four separate affine maps written above. PyTorch’s nn.LSTMCell exposes this fused form through a single weight tensor named weight_ih (input-to-hidden) and weight_hh (hidden-to-hidden), each of shape $(4n_{\text{neurons}})\times (\cdot )$.

The fused matmul parallelizes the gates, not time

The fused weight matrix turns the four affine maps into a single GEMM, and the mini-batch turns the $B$ sequences into extra columns of that GEMM. Both of those axes are parallelized. The one axis that stays strictly sequential is time: ${\mathbf{c}}_t$ and ${\mathbf{h}}_t$ cannot be computed until ${\mathbf{c}}_{t-1}$ and ${\mathbf{h}}_{t-1}$ exist, so a length-$T$ sequence forces $T$ dependent steps that no amount of hardware can collapse into fewer.

This is the structural reason LSTMs are slow on long sequences, and the reason the architectures that followed were designed around it. Self-attention removes the recurrence entirely and processes all positions at once; linear state-space models keep a recurrence but choose one whose decay admits a parallel scan, trading the LSTM’s per-step nonlinearity for parallelism across the sequence. The LSTM’s inductive bias is excellent; its sequential dependency along time is the price it pays for it.

Parameter count and initialization

Each of the four gates has weight matrices ${\mathbf{W}}_{x\bullet }\in {\mathbb{R}}^{n_{\text{neurons}}\times n_{\text{inputs}}}$, ${\mathbf{W}}_{h\bullet }\in {\mathbb{R}}^{n_{\text{neurons}}\times n_{\text{neurons}}}$ and a bias ${\mathbf{b}}_{\bullet }\in {\mathbb{R}}^{n_{\text{neurons}}}$. The total parameter count of one LSTM cell is therefore

$$P_{\text{LSTM}}=4(n_{\text{neurons}}{n}_{\text{inputs}}+n_{\text{neurons}}^2+n_{\text{neurons}}).$$

Where the count comes from

One affine map of the form ${\mathbf{W}}_x{\mathbf{x}}_t+{\mathbf{W}}_h{\mathbf{h}}_{t-1}+\mathbf{b}$ contains:

• ${\mathbf{W}}_x$: a matrix of shape $n_{\text{neurons}}\times n_{\text{inputs}}$, i.e. $n_{\text{neurons}}\cdot n_{\text{inputs}}$ scalars;
• ${\mathbf{W}}_h$: a matrix of shape $n_{\text{neurons}}\times n_{\text{neurons}}$, i.e. $n_{\text{neurons}}^2$ scalars;
• $\mathbf{b}$: a vector of length $n_{\text{neurons}}$.

Summing, one affine map has $n_{\text{neurons}}{n}_{\text{inputs}}+n_{\text{neurons}}^2+n_{\text{neurons}}$ parameters. The LSTM cell has four such independent maps (one per gate, plus the candidate update), with no shared parameters between them, giving the factor of $4$.

For comparison, the vanilla RNN has a single such affine map and therefore $n_{\text{neurons}}{n}_{\text{inputs}}+n_{\text{neurons}}^2+n_{\text{neurons}}$ parameters, exactly one quarter of the LSTM count at the same hidden width.

The corresponding vanilla RNN has one quarter of this count. The factor of $4$ is the price of preservation discussed in the overview: three of the four gates exist solely to make the memory line behave as a stable additive register.

The initial states are conventionally fixed:

$${\mathbf{h}}_0=0,{\mathbf{c}}_0=0.$$

A simple but consequential initialization detail (Jozefowicz, Zaremba and Sutskever, 2015) deserves restating: the forget-gate bias ${\mathbf{b}}_f$ is initialized close to $1$ rather than $0$, so that $\sigma ({\mathbf{b}}_f)\approx 0.73$ at the start of training. Without this, a freshly initialized LSTM forgets roughly half of its memory at every step, and the cell-state path is effectively dead long before gradient descent can teach the gate to keep things. The justification is given at the end of the Forget gate note.

Unrolling and training

An LSTM layer processes a sequence by unrolling the cell across $T$ time steps, with the same parameters at every step. Forward, this is just the iterated application of the equations above, starting from ${\mathbf{h}}_0={\mathbf{c}}_0=0$. Backward, this is Backpropagation Through Time on the resulting computational graph, in any of its practical variants (full BPTT, truncated BPTT with detach() at chunk boundaries, or random-chunk training when long context is not informative).

The next note, Gradient in LSTM, works out the backward pass on this graph and shows explicitly why the cell-state line carries gradient across long gaps, while the hidden-state line, despite passing through all four gates, behaves benignly thanks to the gated additive structure underneath.