LSTM Cell: Input, Forget, and Output Gates
Learners will diagram the LSTM cell, trace information flow through each gate, and implement an LSTM text classifier using nn.LSTM.
LSTM Cell: Input, Forget, and Output Gates is a free Machine Learning Academy lesson on CoddyKit — lesson 3 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Machine Learning Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.
Why RNNs Need Gating Mechanisms
Standard Vanilla RNNs struggle with long-range dependencies because gradients vanish during backpropagation through time. The hidden state h_t = tanh(W_h * h_{t-1} + W_x * x_t) overwrites previous context with each new input, causing the network to forget information from many steps back.
The Long Short-Term Memory (LSTM) was designed in 1997 by Hochreiter and Schmidhuber to solve this problem. It introduces a cell state — a separate memory conveyor belt — alongside gating mechanisms that explicitly control what information is added, removed, or passed forward.
The LSTM Cell State Concept
The LSTM has two internal states: the cell state C_t and the hidden state h_t. The cell state runs like a conveyor belt through the entire sequence, with gates controlling minor linear interactions. This makes gradients flow more easily through time.
The hidden state h_t is the output at each timestep, computed from the cell state. Think of C_t as long-term memory and h_t as working memory that gets passed to the next layer and returned as output.
import torch
import torch.nn as nn
# An LSTM processes sequences and returns (output, (h_n, c_n))
lstm = nn.LSTM(input_size=10, hidden_size=32, batch_first=True)
print('LSTM parameters:', sum(p.numel() for p in lstm.parameters()))The Forget Gate
The forget gate is the first gate the input passes through. It reads the previous hidden state h_{t-1} and the current input x_t, and outputs a value between 0 and 1 for each element in the cell state.
A value of 1 means 'keep everything', while 0 means 'forget completely'. The formula is: f_t = sigmoid(W_f · [h_{t-1}, x_t] + b_f). For example, when processing a new sentence in an NLP task, the forget gate would reset any stored subject information from the previous sentence.
import torch
import torch.nn as nn
# Manual forget gate illustration
batch_size, hidden_size, input_size = 1, 4, 3
h_prev = torch.zeros(batch_size, hidden_size)
x_t = torch.randn(batch_size, input_size)
W_f = torch.randn(hidden_size, hidden_size + input_size)
b_f = torch.zeros(hidden_size)
combined = torch.cat([h_prev, x_t], dim=1)
f_t = torch.sigmoid(combined @ W_f.T + b_f)
print('Forget gate output:', f_t) # Values between 0 and 1The Input Gate and Candidate Cell
The input gate decides which new values to store in the cell state. It has two parts working together: the input gate i_t = sigmoid(W_i · [h_{t-1}, x_t] + b_i) determines how much to update, and the candidate cell g_t = tanh(W_g · [h_{t-1}, x_t] + b_g) creates candidate values to potentially add.
The sigmoid controls the gate (0 = closed, 1 = open) while tanh creates values in the range [-1, 1]. The new information that gets added to the cell state is i_t * g_t — the gate filtered by the candidate values.
# Input gate and candidate values
W_i = torch.randn(hidden_size, hidden_size + input_size)
b_i = torch.zeros(hidden_size)
W_g = torch.randn(hidden_size, hidden_size + input_size)
b_g = torch.zeros(hidden_size)
i_t = torch.sigmoid(combined @ W_i.T + b_i) # Input gate: what to update
g_t = torch.tanh(combined @ W_g.T + b_g) # Candidate values to add
print('Input gate:', i_t.detach())
print('Candidate cell:', g_t.detach())Updating the Cell State
Once we have the forget gate f_t, input gate i_t, and candidate cell g_t, updating the cell state is straightforward: C_t = f_t * C_{t-1} + i_t * g_t.
The first term f_t * C_{t-1} applies the forget gate — selectively erasing information from the previous cell state. The second term i_t * g_t adds new information selectively. This additive structure is what allows gradients to flow back through many timesteps without vanishing, since the gradient flows directly through the addition operation.
# Update cell state
C_prev = torch.zeros(batch_size, hidden_size) # Previous cell state
C_t = f_t * C_prev + i_t * g_t # Element-wise operations
print('Updated cell state C_t:', C_t.detach())
# The additive update is key: gradients flow through + easily
# Compare to vanilla RNN: h_t = tanh(W * h_{t-1} + U * x_t)
# where gradients must flow through the tanh compression each stepThe Output Gate
The output gate controls what part of the cell state gets exposed as the hidden state output h_t. First, the output gate decides which parts of the cell state to output: o_t = sigmoid(W_o · [h_{t-1}, x_t] + b_o).
Then, the cell state is passed through tanh (to push values between -1 and 1) and multiplied by the output gate: h_t = o_t * tanh(C_t). The hidden state h_t serves as both the output at this timestep and the input to the next LSTM step alongside the new cell state.
# Output gate and hidden state
W_o = torch.randn(hidden_size, hidden_size + input_size)
b_o = torch.zeros(hidden_size)
o_t = torch.sigmoid(combined @ W_o.T + b_o) # Output gate
h_t = o_t * torch.tanh(C_t) # New hidden state
print('Output gate:', o_t.detach())
print('New hidden state h_t:', h_t.detach())LSTM Parameter Count and Weight Matrices
An LSTM has four weight matrices (forget, input, candidate, output), each of size (hidden_size, hidden_size + input_size) plus bias. The total parameter count is 4 * hidden_size * (hidden_size + input_size) + 4 * hidden_size.
PyTorch's nn.LSTM packs all four gates into combined weight matrices weight_ih_l0 and weight_hh_l0 for efficiency. This is why LSTM training can be slower than a GRU (Gated Recurrent Unit), which uses only two gates and fewer parameters.
import torch.nn as nn
hidden = 64
input_s = 32
lstm = nn.LSTM(input_size=input_s, hidden_size=hidden, batch_first=True)
# PyTorch stores weight_ih (input-hidden) and weight_hh (hidden-hidden)
print('weight_ih_l0 shape:', lstm.weight_ih_l0.shape) # (4*hidden, input)
print('weight_hh_l0 shape:', lstm.weight_hh_l0.shape) # (4*hidden, hidden)
print('Total params:', sum(p.numel() for p in lstm.parameters()))Using nn.LSTM in PyTorch
PyTorch's nn.LSTM accepts a sequence tensor of shape (batch, seq_len, input_size) when batch_first=True, and optionally an initial state tuple (h_0, c_0). It returns the output tensor (all hidden states) and a tuple of the final states.
For sequence classification, you typically only need the last hidden state h_n from the final timestep. For sequence labelling tasks (like named entity recognition), you use all timestep outputs. Always initialise hidden states to zero unless you have a reason to pass context between batches.
import torch
import torch.nn as nn
batch_size, seq_len, input_size = 16, 20, 32
hidden_size = 64
lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, batch_first=True)
x = torch.randn(batch_size, seq_len, input_size)
# Forward pass
output, (h_n, c_n) = lstm(x)
print('Output shape (all timesteps):', output.shape) # (16, 20, 64)
print('h_n shape (final hidden):', h_n.shape) # (1, 16, 64)
print('c_n shape (final cell):', c_n.shape) # (1, 16, 64)Multi-Layer and Bidirectional LSTMs
A stacked LSTM feeds the output of one LSTM layer as input to the next, allowing the network to learn hierarchical temporal representations. Set num_layers=2 or more in nn.LSTM.
A bidirectional LSTM runs one LSTM forward through the sequence and another backward, concatenating their hidden states. This gives each timestep context from both the past and the future. Set bidirectional=True; the output dimension becomes 2 * hidden_size. Bidirectional LSTMs are very effective for NLP tasks where both left and right context matter.
import torch
import torch.nn as nn
# Bidirectional stacked LSTM
lstm = nn.LSTM(
input_size=32,
hidden_size=64,
num_layers=2,
batch_first=True,
bidirectional=True,
dropout=0.3 # Dropout between layers
)
x = torch.randn(16, 20, 32)
output, (h_n, c_n) = lstm(x)
print('Bidirectional output shape:', output.shape) # (16, 20, 128)
print('h_n shape:', h_n.shape) # (4, 16, 64): 2 layers * 2 directionsGRU: A Simpler Alternative to LSTM
The Gated Recurrent Unit (GRU) simplifies the LSTM by merging the forget and input gates into a single update gate, and combining the cell and hidden states. It uses only two gates: the reset gate (how much past to forget) and the update gate (how much to update).
GRU has fewer parameters and trains faster than LSTM while often achieving comparable performance. Use LSTM when you need maximum expressiveness on complex sequences, and GRU when speed and simplicity matter. In practice, empirical testing on your specific task determines the better choice.
import torch
import torch.nn as nn
# GRU is simpler: only 3 weight matrices instead of 4
gru = nn.GRU(input_size=32, hidden_size=64, batch_first=True)
lstm = nn.LSTM(input_size=32, hidden_size=64, batch_first=True)
print('GRU params:', sum(p.numel() for p in gru.parameters()))
print('LSTM params:', sum(p.numel() for p in lstm.parameters()))
# GRU has 25% fewer parameters than LSTM for same hidden sizeLSTM Intuition Through a Language Example
Consider the sentence: 'The author, who lived in Paris for many years, wrote a novel.' An LSTM needs to remember 'author' (singular) when it reaches 'wrote' many words later.
The forget gate keeps 'author' in the cell state throughout the long relative clause. The input gate adds new relevant information like 'wrote' to the cell state. The output gate uses the cell state at the right moment to produce the correct hidden state for downstream tasks like predicting the next word or classifying the sentence.
- Forget gate: Keep the subject through the relative clause
- Input gate: Store the verb information when encountered
- Output gate: Use stored information at prediction time
Quick Check
Test your understanding of LSTM gates from this lesson.
Lesson Recap
In this lesson you learned: the LSTM cell state acts as long-term memory that flows through the sequence with minimal modification, three gating mechanisms (forget, input, output) control information flow using sigmoid-gated multiplication, and the additive cell state update C_t = f_t * C_{t-1} + i_t * g_t allows gradients to flow without vanishing. Next up we apply LSTMs to a real sentiment analysis task, building an end-to-end text classifier.
Frequently asked questions
Is the “LSTM Cell: Input, Forget, and Output Gates” lesson free?
Yes — the full text of “LSTM Cell: Input, Forget, and Output Gates” is free to read here on the web, and the Machine Learning Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Machine Learning Academy course, upgrade to CoddyKit PRO.
What will I learn in “LSTM Cell: Input, Forget, and Output Gates”?
Learners will diagram the LSTM cell, trace information flow through each gate, and implement an LSTM text classifier using nn.LSTM. You practise Machine Learning Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.
Do I need any experience to start Machine Learning Academy?
No prior experience is required. Machine Learning Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 3 of 4, so you can start here or from the beginning and move at your own pace.
How long does the “LSTM Cell: Input, Forget, and Output Gates” lesson take?
Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.
Can I write and run code in this Machine Learning Academy lesson?
Yes. Every Machine Learning Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.
All lessons in this course
- Vanilla RNNs: Hidden State and Sequence Unrolling
- The Vanishing Gradient Problem in Deep Time Steps
- LSTM Cell: Input, Forget, and Output Gates
- Sequence-to-One: Sentiment Analysis with an LSTM