LSTM Cell From Scratch (PyTorch)

This repository implements a single LSTM cell from scratch in PyTorch and explains the math behind it.
The aim is to make the mapping

$$(x_t, h_{t-1}, c_{t-1}) \rightarrow (h_t, c_t)$$

explicit and consistent with standard LSTM definitions.

Overview

• Custom LSTMCellCustom module in PyTorch
• Gate-by-gate mathematical formulation
• Minimal example that unrolls the cell across a sequence
• Optional tests (shape checks / comparison to nn.LSTMCell)

Problem Setup

We model an input sequence

$$x_1, x_2, \dots, x_T,\quad x_t \in \mathbb{R}^{d_x}$$

The LSTM maintains at each time step:

• hidden state: $h_t \in \mathbb{R}^{d_h}$
• cell state: $c_t \in \mathbb{R}^{d_h}$

Given $x_t$, $h_{t-1}$ and $c_{t-1}$, the LSTM cell outputs $h_t$ and $c_t$.

Mathematical Formulation

Notation and dimensions

• input dimension: $d_x$
• hidden / cell dimension: $d_h$

Gates:

• input gate: $i_t \in \mathbb{R}^{d_h}$
• forget gate: $f_t \in \mathbb{R}^{d_h}$
• output gate: $o_t \in \mathbb{R}^{d_h}$
• candidate cell (proposal): $g_t \in \mathbb{R}^{d_h}$

Nonlinearities:

• sigmoid: $\sigma(z) = \dfrac{1}{1 + e^{-z}}$
• hyperbolic tangent: $\tanh(z)$

Parameters (single layer, single direction):

• input-to-gates: $W_x \in \mathbb{R}^{4 d_h \times d_x}$
• hidden-to-gates: $W_h \in \mathbb{R}^{4 d_h \times d_h}$
• bias: $b \in \mathbb{R}^{4 d_h}$

All four gates share these matrices and are obtained by splitting the $4 d_h$-dimensional vector.

Forward pass (gate-by-gate)

1. Combined affine transform:

$$ z_t = W_x x_t + W_h h_{t-1} + b, \quad z_t \in \mathbb{R}^{4 d_h}. $$

2. Split into blocks (each $d_h$ long):

$$ z_t = \begin{bmatrix} z_t^{(i)} \\ z_t^{(f)} \\ z_t^{(o)} \\ z_t^{(g)} \end{bmatrix}, \qquad z_t^{(i)}, z_t^{(f)}, z_t^{(o)}, z_t^{(g)} \in \mathbb{R}^{d_h}. $$

3. Apply nonlinearities:

$$ \begin{aligned} i_t &= \sigma!\big(z_t^{(i)}\big), \\ f_t &= \sigma!\big(z_t^{(f)}\big), \\ o_t &= \sigma!\big(z_t^{(o)}\big), \\ g_t &= \tanh!\big(z_t^{(g)}\big). \end{aligned} $$

4. Cell state update:

$$ c_t = f_t \odot c_{t-1} + i_t \odot g_t , $$

where $\odot$ is elementwise multiplication.

5. Hidden state update:

$$ h_t = o_t \odot \tanh(c_t). $$

Compact form

Stack the gates:

$$ \gamma_t = \begin{bmatrix} i_t \ f_t \ o_t \ g_t \end{bmatrix} \in \mathbb{R}^{4 d_h}. $$

Then

$$ \gamma_t = \phi\big( W_x x_t + W_h h_{t-1} + b \big), $$

where $\phi$ applies $\sigma$ to the first three blocks and $\tanh$ to the last block.

Gradient Flow (Intuition)

The key recurrence is

$$ c_t = f_t \odot c_{t-1} + i_t \odot g_t. $$

For component $k$:

$$ \frac{\partial c_t^{(k)}}{\partial c_{t-1}^{(k)}} = f_t^{(k)}. $$

• $f_t^{(k)} \approx 1$ → information and gradients are kept over time.
• $f_t^{(k)} \approx 0$ → that component is reset and gradients vanish on this path.

A common trick is to initialize the forget-gate bias $b_f$ (the second $d_h$-block of $b$) to a positive value (e.g. 1.0) so that the initial forget gate values are biased toward keeping information.

PyTorch Implementation

The core cell is implemented as a single nn.Module.

src/lstm_cell.py

The module is called LSTMCellCustom(input_size, hidden_size) and implements the equations described above:

• input: x_t of shape (batch_size, input_size)
• hidden: h_{t-1} of shape (batch_size, hidden_size)
• cell: c_{t-1} of shape (batch_size, hidden_size)
• output: (h_t, c_t) with the same shapes as the hidden and cell states.

Sequence Example

A small script shows how to unroll the custom LSTM cell across a full sequence and collect all hidden states.

examples/simple_sequence_demo.py:

import torch
from src.lstm_cell import LSTMCellCustom


def run_demo():
    torch.manual_seed(0)

    batch_size = 2
    seq_len = 5
    input_size = 3
    hidden_size = 4

    # Input sequence: (batch, time, features)
    x = torch.randn(batch_size, seq_len, input_size)

    cell = LSTMCellCustom(input_size, hidden_size)

    # Initial states
    h_t = torch.zeros(batch_size, hidden_size)
    c_t = torch.zeros(batch_size, hidden_size)

    outputs = []

    for t in range(seq_len):
        x_t = x[:, t, :]
        h_t, c_t = cell(x_t, h_t, c_t)
        outputs.append(h_t.unsqueeze(1))

    # Stack over time: (batch, time, hidden)
    H = torch.cat(outputs, dim=1)
    print("Hidden sequence shape:", H.shape)


if __name__ == "__main__":
    run_demo()

Run the example

python examples/simple_sequence_demo.py

You should see:

Hidden sequence shape: torch.Size([2, 5, 4])

Tests

This repository includes a small test to check that the LSTM cell returns tensors with the expected shapes.

tests/test_lstm_cell_shapes.py:

import torch
from src.lstm_cell import LSTMCellCustom


def test_lstm_cell_shapes():
    batch_size = 3
    input_size = 5
    hidden_size = 7

    x_t = torch.randn(batch_size, input_size)
    h_prev = torch.randn(batch_size, hidden_size)
    c_prev = torch.randn(batch_size, hidden_size)

    cell = LSTMCellCustom(input_size, hidden_size)
    h_t, c_t = cell(x_t, h_prev, c_prev)

    assert h_t.shape == (batch_size, hidden_size)
    assert c_t.shape == (batch_size, hidden_size)

References

• S. Hochreiter and J. Schmidhuber, “Long Short-Term Memory,” Neural Computation, 1997.
• PyTorch documentation for torch.nn.LSTM and torch.nn.LSTMCell.