LSTM Architecture › PyTorch nn.LSTM Module
Language AI HandbookAll chapters LSTM Architecture › PyTorch nn.LSTM ModuleOne section · Sections in this chapter7 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
PyTorch provides nn.LSTM as a built-in, highly optimized implementation. Understanding its interface is essential before building anything with LSTMs.
Creating an LSTM
import torch.nn as nn
# Single-layer LSTM: 10 input features, 20 hidden units
lstm = nn.LSTM(
input_size=10, # Number of features per timestep
hidden_size=20, # Dimensions of hidden state and cell state
num_layers=1, # Number of stacked LSTM layers
batch_first=True, # Input shape: (batch, seq_len, features)
)Total LSTM parameters: 2,560 Expected (4 gates x (input_size + hidden_size) x hidden_size + biases): 2,480
The parameter count formula reveals how LSTM scaling works. With input size and hidden size , the total parameters are . For a common configuration like and , that is roughly 1.6 million parameters per LSTM layer. This is a modest count compared to transformer layers with similar hidden sizes, though transformers achieve their expressiveness through attention rather than gated recurrence.
Input and Output Shapes
import torch
# Create a simple LSTM for shape demonstration
lstm_demo = nn.LSTM(input_size=10, hidden_size=20, batch_first=True)
batch_size = 3
seq_len = 5
# Input: (batch_size, seq_len, input_size)
x_demo = torch.randn(batch_size, seq_len, 10)
# Forward pass
# output: (batch, seq_len, hidden_size) - h_t at every timestep
# hidden: (num_layers, batch, hidden_size) - final h_T
# cell: (num_layers, batch, hidden_size) - final C_T
output_demo, (hidden_demo, cell_demo) = lstm_demo(x_demo)Input shape: [3, 5, 10] Output shape: [3, 5, 20] (h_t at every step) Hidden shape: [1, 3, 20] (final h_T only) Cell shape: [1, 3, 20] (final C_T only)
The output tensor contains at every timestep. This is what you use for per-token prediction tasks: named entity recognition, part-of-speech tagging, sequence-to-sequence encoding, or any task where you need a representation for each position in the input. The hidden and cell tuple contains only the final states, which is sufficient for sequence-level tasks like classification. If you only need the final summary representation, using hidden[-1] (the last layer's final hidden state) is both efficient and semantically appropriate.
Initializing Hidden State
By default, nn.LSTM initializes the hidden and cell states to zero when no initial state is provided. For many tasks, this default is fine. But there are important cases where explicit initialization matters.
When processing sequences in minibatches, sequences within a batch are unrelated to each other. You want the hidden state to be zero at the start of every new sequence, and PyTorch handles this automatically when you call lstm(x) without passing a hidden state.
When processing long sequences in chunks for memory efficiency, you need to carry the hidden state from one chunk to the next. Here you would save the (hidden, cell) output from one call and pass it as input to the next, carrying state across chunk boundaries.
When using the LSTM as a language model decoder initialized from an encoder, you would set the initial hidden state to the encoder's final hidden state, allowing the decoder to condition its generation on the encoded input. This encoder-decoder pattern was the dominant architecture for machine translation before transformers replaced it.