RNN vs LSTM vs GRU: trace states through a short sequence

An RNN updates a hidden state as it reads a sequence. LSTM adds a separate cell state and gates controlling information flow. GRU uses gates with a hidden state but no separately exposed cell state. These structural differences affect parameter counts and computation; they do not by themselves establish which model will perform best on a task.

We will pass a three-step, one-feature sequence through each PyTorch module. All parameters are deliberately set to the same small constant to make the calculation reproducible. These are untrained modules, so their numerical outputs are traces of their equations, not quality scores.

Track batch, time and hidden axes

With batch_first=True, the input is (batch, time, features) = (1, 3, 1). A single unidirectional layer with hidden size 2 produces output (1, 3, 2). The final hidden state has shape (layers×directions, batch, hidden) = (1, 1, 2).

LSTM also returns a cell state of shape (1, 1, 2) in this configuration. batch_first changes the input and output arrangement, but does not move the batch axis to the front of the final hidden-state tensor.

For input size 1 and hidden size 2, a basic RNN has 10 parameters with PyTorch's two bias vectors. GRU has three gate groups and 30 parameters. LSTM has four gate groups and 40. Projection, bidirectionality, additional layers and disabled biases change those counts.

Inspect states and reverse the input

import torch
torch.set_num_threads(1)
x = torch.tensor([[[1.], [0.], [-1.]]], dtype=torch.float64)
for cls in [torch.nn.RNN, torch.nn.LSTM, torch.nn.GRU]:
    module = cls(1, 2, batch_first=True).double()
    with torch.no_grad():
        for parameter in module.parameters():
            parameter.fill_(.1)
        output, state = module(x)
        reversed_output, _ = module(x.flip(1))
    hidden = state[0] if isinstance(state, tuple) else state
    print(cls.__name__, 'parameters=' + str(sum(p.numel() for p in module.parameters())),
          'output=' + str(tuple(output.shape)), 'hidden=' + str(tuple(hidden.shape)))
    print('step values:', [round(v, 6) for v in output[0, :, 0].tolist()])
    print('reversal changes final state:', not torch.allclose(output[:, -1], reversed_output[:, -1]))
    if isinstance(state, tuple):
        print('cell shape:', tuple(state[1].shape))
    assert torch.allclose(output[:, -1], hidden[0])

Checked output:

RNN parameters=10 output=(1, 3, 2) hidden=(1, 1, 2)
step values: [0.291313, 0.25267, 0.149407]
reversal changes final state: True
LSTM parameters=40 output=(1, 3, 2) hidden=(1, 1, 2)
step values: [0.095241, 0.116026, 0.093377]
reversal changes final state: True
cell shape: (1, 1, 2)
GRU parameters=30 output=(1, 3, 2) hidden=(1, 1, 2)
step values: [0.107199, 0.133299, 0.102356]
reversal changes final state: True

Interpret memory without inventing a benchmark

The hidden value at a later step depends on earlier state as well as the current input. Reversing the sequence changes that path, and the final states change in this fixture. This demonstrates order sensitivity; it does not show that a module learned a useful long-term dependency.

LSTM's cell update offers an additive route through time, controlled by gates. GRU combines state through its update and reset mechanisms. These designs address training and information-flow problems, but long-sequence performance still depends on data, initialization, optimization and architecture choices.

The constant initialization here makes hidden units symmetric. It is intentional for inspection and unsuitable as a general training recipe because identical units can remain redundant. Train with an appropriate initialization and evaluate a declared sequence task before comparing predictive quality.

Keep padding and state boundaries explicit

Real minibatches often contain variable-length sequences. Padding values can affect a recurrent computation unless lengths are handled correctly, for example through packed sequences or an appropriate masked workflow. The last padded output is not automatically the final meaningful state.

Decide whether state should reset between independent examples. Carrying a hidden state from one unrelated learner, sensor or document into another can leak information and change the task. Stateful streaming requires a clear sequence identity and reset policy.

Extend the trace

Use two sequences in the batch and verify the axes again. Then enable bidirectionality and derive the changed output width and hidden-state leading dimension before running the module. Finally, train each architecture on the same small prediction task with separate development and test sequences.

The Deep Learning syllabus includes sequence models; lessons are being published. Review matrix shapes for the weighted transformations, and train versus eval before interpreting dropout or recording gradients during evaluation.

References

PyTorch RNN — Basic recurrent equations, parameter layout and tensor shapes.

PyTorch LSTM — Gates, hidden and cell states and projection conventions.

PyTorch GRU — GRU equations and PyTorch-specific computation details.