Part of Language AI Handbook
Covers LSTM architecture including cell state, gate mechanisms, and information flow. Explains how LSTMs solve vanishing gradients.
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
LSTM Architecture
In the previous chapters, we saw how vanilla RNNs struggle with long sequences. As we established in the Vanishing Gradients chapter, gradients either vanish into insignificance or explode beyond control, making it nearly impossible to learn dependencies that span more than a few dozen timesteps. This limitation motivated researchers to rethink the fundamental architecture of recurrent networks.
Long Short-Term Memory (LSTM) networks, introduced by Hochreiter and Schmidhuber in 1997, solve the vanishing gradient problem through a deceptively simple architectural insight: instead of forcing all information to pass through repeated nonlinear transformations at every timestep, create a dedicated pathway where information can flow relatively unchanged across many timesteps. This "memory highway" allows LSTMs to maintain and access information over hundreds or even thousands of steps, enabling applications that were previously impossible with vanilla RNNs. The key idea is so elegant that it can be stated in a single sentence: separate the cell's memory from its output, and let the memory update additively rather than through repeated multiplication.
This chapter focuses on the conceptual architecture and intuition behind LSTMs: what each component does, why the design works, and how the pieces fit together. The next chapter, LSTM Gate Equations, covers the precise mathematical derivations, parameter counts, and scratch implementations.
The Cell State: An Information Highway
The defining feature of an LSTM is its cell state, a separate memory channel that runs parallel to the hidden state. Think of it as a conveyor belt that carries information through time. Unlike the hidden state in vanilla RNNs, which gets completely rewritten at each timestep through a matrix multiplication and nonlinearity, the cell state can preserve information indefinitely by flowing forward with only element-wise operations applied to it.
To appreciate why this matters, consider what happens inside a vanilla RNN at each step. The hidden state is recomputed as . Both and tanh are involved in that transformation. Tanh squashes any value into the range , so the magnitude of the hidden state is always bounded. When you backpropagate through this squashing operation over many steps, the gradient of the loss with respect to an early hidden state gets multiplied by the derivative of tanh at each step. The derivative of tanh is bounded by 1, often much smaller for large inputs, so the gradient shrinks exponentially. After 50 or 100 steps, the gradient is effectively zero.
The LSTM cell state bypasses this problem in a fundamental way. Information stored in the cell state passes from one timestep to the next primarily through an additive operation: new information is added to old information, with each term selectively scaled by a gate value. Addition does not shrink gradients the way multiplication through nonlinearities does. The gradient of a sum is just the sum of the gradients, which means gradients can travel backward through the cell state pathway without multiplicative attenuation. This additive gradient path is the mathematical core of why LSTMs work.
It is worth pausing to appreciate how unusual this design was when it was proposed in 1997. The mainstream assumption in neural network research at the time was that computation should proceed through stacked layers of learned nonlinear transformations. Every layer or every timestep would apply a weight matrix and a squashing function, and the network would learn by adjusting those weights. Hochreiter and Schmidhuber challenged this by arguing that a memory system needed a way to copy information forward unchanged, not transform it at every step. Their cell state is essentially a write-once, read-selectively buffer: information can be injected into it, it can persist untouched for many steps, and it can later be read out when needed. This design philosophy would later resurface in the residual connections of deep CNNs and in the skip connections that are essential to modern Transformer architectures.
The two key vectors in an LSTM at timestep are:
- : the cell state, which stores long-term memory and flows through time with minimal modification
- : the hidden state, which is the output at each timestep and is derived from the cell state through the output gate

In a vanilla RNN, information at timestep must pass through every intermediate hidden state to reach timestep . Each passage involves a matrix multiplication and a tanh activation, which compresses values and causes gradients to shrink exponentially. After enough steps, the original signal is lost in noise.
The cell state sidesteps this problem. Information stored in the cell state can travel from the beginning to the end of a sequence with only element-wise additions and scalings applied to it. Neither of those operations compounds the gradient shrinkage the way repeated matrix multiplications through tanh do. That additive path through time is the core mathematical reason LSTMs can learn long-range dependencies.
Gate Mechanism Intuition
If the cell state flowed forward completely unchanged at every step, the network would have no way to update its memory. LSTMs need a principled way to selectively add new information, remove outdated information, and control what gets exposed as output. They accomplish this through gates, learned sigmoid layers that output values between 0 and 1.
Think of gates as dimmer switches rather than on/off toggles. A gate value of 0.0 means "block everything," 0.5 means "let half through," and 1.0 means "let everything through." Because gates use sigmoid activations, their outputs always fall in the range, making them perfect for scaling information flow. The sigmoid function is differentiable everywhere, which is critical for backpropagation. A hard threshold, where values below 0.5 are set to 0 and values above are set to 1, would create a non-differentiable decision boundary that prevents gradient-based learning.
The gate values are not fixed: they are computed fresh at every timestep by passing the current input and previous hidden state through a small learned network. This means the same LSTM cell can behave differently at different points in a sequence. When processing the beginning of a long document, the forget gate might learn to mostly preserve stored information. When a new named entity appears, the input gate might spike to write it into memory. When a question mark arrives, the output gate might activate to expose accumulated context for an answer prediction. All of this behavior is learned implicitly from training data, not handcrafted by the designer.


LSTMs have three gates, each serving a distinct purpose:
- Forget gate: Decides what information to discard from the cell state. When processing a new sentence, the network may want to forget who the previous sentence was about.
- Input gate: Decides what new information to store in the cell state. When the model encounters a new subject, it wants to record that subject.
- Output gate: Decides what portion of the cell state to expose as the hidden state. Not everything stored in memory is relevant to the current prediction.
Every gate takes the same two inputs: the previous hidden state and the current input . The general gate computation is:
where:
- : the sigmoid activation function, which outputs values in
- : the weight matrix for this gate (each gate has its own learned weights)
- : the concatenation of the previous hidden state and current input
- : the bias vector for this gate
Because each gate has its own weight matrix, they specialize during training. The forget gate might learn to fire when it sees punctuation marks that signal a sentence boundary; the input gate might learn to fire when it encounters a named entity that needs to be tracked. This specialization is not programmed in: it emerges from gradient descent. What makes the architecture powerful is that the parameterization is rich enough to learn such specialization, and the gating structure provides the mechanism through which that specialization can influence the cell state.
Why Sigmoid and Tanh? A Design Choice
It is worth asking why the LSTM specifically uses sigmoid for gates and tanh for memory computations. The answer is functional: each activation is chosen for what it computes in context.
Sigmoid maps any real-valued input to the range . When you multiply the cell state element-wise by a vector of sigmoid outputs, each dimension of the cell state gets scaled by a value between 0 and 1. A value of 0 completely erases that dimension; a value of 1 keeps it unchanged; anything in between produces a partial retention. This is exactly the "dimmer switch" behavior that gates require. No other common activation function produces values strictly in in a natural way that is also differentiable and trainable.
Tanh, by contrast, maps inputs to . This is the appropriate range for candidate values that will be added to the cell state, because a positive candidate increases a dimension's value while a negative candidate decreases it. If you used sigmoid for candidates, you could only ever add positive information to the cell state, never subtract it. The result would be a one-sided accumulation that could not represent the full range of patterns needed in practice. Using tanh allows the cell state to hold both positive and negative magnitudes, which is necessary for representing rich, varied information across different dimensions.
The combination of the two is what makes the full architecture work. The sigmoid gate decides how much to write; the tanh candidate provides what to write. Together they give the network full expressiveness over each cell state dimension.
LSTM Diagram Walkthrough
Now that the gates make sense individually, let us see how they work together to create a complete memory system. The LSTM processes information in four distinct stages per timestep: decide what to forget, decide what to add, update the cell state, and decide what to output. These four stages happen simultaneously in the equations, but understanding them sequentially builds the right intuition.
Stage 1: The Forget Gate Decides What to Discard
Before adding new information, the LSTM first filters out what is no longer relevant. The forget gate examines and , and produces a value between 0 and 1 for each element of the cell state. Elements scaled to near 0 get mostly erased; elements scaled to near 1 get preserved.
where:
- : the forget gate output, a vector with one value per cell state dimension
- : the forget gate's weight matrix
- : the forget gate's bias vector
Consider parsing the sentence "The cat sat on the mat. The dog barked." When the model encounters the period and then "The", the forget gate should fire strongly to clear information about "cat" from the cell state, making room for the upcoming "dog" to be stored as the new subject. The forget gate does this by producing values close to 0 for dimensions that are currently holding subject information, effectively zeroing out that memory slot.
In practice, the forget gate tends to learn values close to 1 for most dimensions most of the time. This is the correct behavior for tasks where context persists across many steps. If the forget gate habitually set values to 0, the LSTM would lose its long-term memory too quickly. The forget gate essentially asks: "Given what I just saw and what I remember, which parts of my memory are no longer needed?" and most of the time the answer is: almost nothing needs to be forgotten yet.
The bias in the forget gate is often initialized to a positive value, such as 1.0 or 2.0. The intuition is that a freshly initialized LSTM should err on the side of remembering rather than forgetting. A large positive bias pushes the sigmoid output close to 1, meaning the network starts with a strong tendency to preserve cell state values. During training, it then learns which contexts warrant forgetting. This initialization trick, popularized by Jozefowicz et al. (2015), improves training stability and long-sequence performance.
Stage 2: The Input Gate and Candidate Values
With old information selectively cleared, the LSTM decides what to write into memory. This stage involves two parallel computations that work together. The input gate decides which positions in the cell state to update, while a tanh layer computes candidate values that could be written.
where:
- : the input gate, controlling which cell state positions receive updates
- : candidate values that could be added to the cell state
- : separate weight matrices for the input gate and candidate computation
- : corresponding bias vectors
The tanh function is used for the candidates rather than sigmoid because tanh outputs values in , allowing the network to both increase and decrease cell state values. A positive candidate pushes that cell state element upward; a negative candidate pushes it downward. The input gate then scales these proposals, determining how much of each proposed change to apply.
The separation of the input gate from the candidate layer is subtle but important. Why not just compute new values and write them directly? The two-part design allows the network to decouple what kind of information to write from how much of it to write. The candidate layer might compute strong signals for certain dimensions, but the input gate might suppress them because the current context is not appropriate for that kind of update. Conversely, the input gate might activate for dimensions where the candidate values are small but consistent with a persistent pattern the network needs to build up incrementally. This decoupling gives the LSTM a richer vocabulary of memory operations than a single-component update would allow.
Consider what happens when the model encounters the word "not" in the phrase "not happy." The candidate layer might produce a negative value for dimensions associated with positivity, pushing the sentiment representation downward. The input gate activates strongly for those same dimensions to ensure the negation gets applied. Without the separate input gate, the network could not independently control where to update and how much to update.
Stage 3: The Cell State Update
The actual memory update combines the forgetting and remembering decisions:
where:
- : the new cell state at timestep
- : the previous cell state from timestep
- : the forget gate output, scaling old information
- : the input gate output, scaling new candidate information
- : the candidate values proposed by the tanh layer
- : element-wise (Hadamard) multiplication
This formula looks simple but does a lot. The first term, , selectively preserves old information. If for some element , we keep 90% of that element's previous value. The second term, , selectively writes new information.
The critical observation is that this update is additive, not purely multiplicative. In a vanilla RNN, information at timestep reaches timestep only through sequential matrix multiplications. In an LSTM, information stored in the cell state at timestep can reach timestep through a series of additions and element-wise scalings. The gradient that flows backward through addition does not get multiplied by a small number at each step the way it would through a tanh-saturated weight matrix. This additive pathway is the core mechanism behind LSTMs solving the vanishing gradient problem.
Consider the gradient of the loss with respect to the cell state at time :
This gradient flows backward by multiplication with , the forget gate value at the next timestep. If is close to 1, the gradient passes through almost unchanged. Compare this to a vanilla RNN, where the corresponding gradient involves the derivative of tanh applied to , a term that can be very small whenever weights or activations are in the saturating regions of tanh. The LSTM cell state gradient path is far gentler, and a well-trained forget gate that learns to keep values near 1 for important memories creates a near-constant gradient pathway. Hochreiter and Schmidhuber called this the "Constant Error Carousel" (CEC): the cell state can carry an error signal (a gradient) forward in time without it vanishing, as long as the forget gate cooperates.
Let us make this concrete with a numerical example using a 4-dimensional cell state:
# Concrete cell state update example with a 4-dimensional cell state
C_prev = np.array([0.8, -0.3, 0.5, 0.1]) # Previous cell state
f_t = np.array(
[0.9, 0.95, 0.2, 0.85]
) # Forget gate (mostly retain, except dim 2)
i_t = np.array([0.7, 0.1, 0.6, 0.8]) # Input gate
C_tilde = np.array([0.4, -0.2, 0.9, 0.5]) # Candidate values from tanh layer
# Stage 3 computation
forgotten = f_t * C_prev # What we keep from old memory
added = i_t * C_tilde # What we add from new information
C_new = forgotten + added # New cell stateCell State Update Breakdown: Dim C_prev f_t Retained i_t Cand Added C_new -------------------------------------------------------------------- 0 0.80 0.90 0.72 0.70 0.40 0.28 1.00 1 -0.30 0.95 -0.28 0.10 -0.20 -0.02 -0.30 2 0.50 0.20 0.10 0.60 0.90 0.54 0.64 3 0.10 0.85 0.09 0.80 0.50 0.40 0.49
Dimension 2 tells the most interesting story. The forget gate value of 0.2 means we discard 80% of the old value. The input gate value of 0.6 combined with a strong candidate of 0.9 adds 0.54. The new cell state for dimension 2 is dominated by new information, exactly the selective updating that enables an LSTM to track changing context.

Stage 4: The Output Gate Produces the Hidden State
After updating the cell state, the LSTM decides what portion of that memory to expose as the current output. The cell state may contain information from many previous steps, but only some of it is relevant to the current prediction. The output gate makes that selection.
where:
- : the output gate, deciding which cell state dimensions to expose
- : the hidden state, which is the output at this timestep
- : the cell state squashed to
The tanh applied to the cell state serves two purposes. First, it bounds the values to , preventing any single dimension from dominating the output. Without this squashing, cell state values could accumulate over many additive updates and grow large in magnitude, causing numerical instability and making the output hard for downstream layers to handle. Second, the tanh re-centers values around zero, which helps downstream computations that typically assume zero-mean inputs.
The output gate then selects which of these bounded values to include in the hidden state. This is where the LSTM distinguishes what it knows and what it reports. The cell state is the network's internal notebook, which may contain many things. The hidden state is what the network decides to share with the world at this particular moment. Suppose the LSTM is generating text and has been tracking the grammatical number (singular or plural) of the current subject. Most of the time that number-tracking information sits quietly in the cell state without being actively used. But when the model reaches a point where a verb needs to agree with the subject, the output gate activates for the dimensions storing number information, pulling it out into the hidden state where it can influence the verb prediction.
This hidden state serves two roles simultaneously: it is the output that downstream layers or classifiers use for predictions, and it is also fed back into the LSTM at the next timestep as . The cell state, by contrast, is internal. It acts as a memory buffer that influences the hidden state through the output gate but is not directly exposed to the rest of the network. This encapsulation is a design virtue: it means the long-term memory (cell state) and the immediate context signal (hidden state) are separate, and the network can develop them differently.
The Complete LSTM Step
Every timestep runs these six computations in order:
- Forget gate:
- Input gate:
- Candidate:
- Cell update:
- Output gate:
- Hidden state:
Each gate has its own weight matrix and bias, giving the network four parameter groups to learn: , , , and . During training, backpropagation adjusts these parameters so the gates learn to open and close at the right times for the task. The precise parameter count and a from-scratch implementation are covered in the LSTM Gate Equations chapter.
Notice that steps 1, 2, 3, and 5 all involve the same inputs ( and ) processed through different weight matrices. In practice, implementations often batch all four of these computations into a single large matrix multiplication, applying one weight matrix of shape at once, then splitting the result into four slices. This is both more efficient on modern hardware and mathematically equivalent to computing each gate separately. PyTorch's nn.LSTM uses this batched formulation internally.
Information Flow in LSTMs
Understanding how information flows through an LSTM helps build intuition about what these networks can learn. Consider processing the sentence: "The cat, which had been sleeping on the warm windowsill since early morning, finally stretched." The subject "cat" appears at the beginning, but the verb "stretched" does not arrive until the end. A vanilla RNN loses the subject information before reaching the verb. An LSTM stores the subject in the cell state and retrieves it when it becomes relevant.
The mechanics of this process illustrate each gate's role in a natural context. When "cat" is processed, the input gate fires strongly for dimensions in the cell state designated for tracking the grammatical subject. The candidate values for those dimensions are computed from the word embedding of "cat", encoding its syntactic and semantic properties. Meanwhile, the forget gate remains high: the network does not want to erase any previously accumulated context from the beginning of the passage.
As the long subordinate clause unfolds ("which had been sleeping on the warm windowsill since early morning"), the forget gate continues to protect the subject information. The input gate makes small writes for the content of the clause, but the cell state dimension holding the subject stays relatively stable. When "finally" appears as a signal that the main clause is about to continue, the output gate begins preparing to expose the subject information that was stored.
When "stretched" is reached, the output gate activates for the subject-tracking dimensions. The model can now see both the stored subject (the cat) and the current token (stretched), and it can use this to make predictions about what comes next. This is long-range dependency resolution in action, spanning roughly a dozen tokens.
# Simulated gate activation trace for the example sentence above.
# These are illustrative values, not from a real trained model,
# but representative of what a trained LSTM produces on such input.
words = [
"The",
"cat",
",",
"which",
"had",
"been",
"sleeping",
"...",
"finally",
"stretched",
".",
]
positions = np.arange(len(words))
forget_gates = [0.9, 0.3, 0.95, 0.95, 0.95, 0.95, 0.95, 0.95, 0.95, 0.4, 0.2]
input_gates = [0.2, 0.9, 0.1, 0.3, 0.2, 0.2, 0.3, 0.1, 0.3, 0.1, 0.1]
output_gates = [0.3, 0.5, 0.2, 0.4, 0.3, 0.3, 0.4, 0.2, 0.5, 0.9, 0.3]
cell_state_trace = [
0.1,
0.85,
0.82,
0.80,
0.78,
0.76,
0.74,
0.72,
0.70,
0.35,
0.1,
]

When "cat" appears, the input gate activates strongly, storing subject information in the cell state. The forget gate remains high through the long subordinate clause, preserving what was stored. When "stretched" arrives, the output gate fires, allowing the stored subject information to inform the verb-agreement prediction. This selective storage, preservation, and retrieval is the LSTM in action.
The Hidden State as a Working Memory
While the cell state provides long-term memory, the hidden state functions more like a working memory: the information currently active and available for computation. A useful mental model distinguishes the two as follows. The cell state is like a filing cabinet: it can hold a large amount of information over long periods, but accessing any particular piece requires an explicit retrieval action (the output gate). The hidden state is like what is currently on your desk: a small, immediately accessible summary of what is relevant right now.
This distinction explains why many sequence tasks use the final hidden state rather than the final cell state for downstream processing. The cell state at the end of a sequence may contain accumulated memories from early in the sequence that are no longer contextually active. The final hidden state, having been filtered through the output gate, contains a more focused representation of what was relevant at that last timestep. For classification tasks, this focused representation is often more useful as a fixed-size input to a classifier than the raw cell state would be.
It also explains a common practical observation: when fine-tuning or inspecting trained LSTMs, analysis of the hidden states is much more informative than analysis of the cell states. The hidden state is designed to be informative for the current prediction; the cell state is designed to be a stable reservoir. Probing tools that visualize LSTM representations typically examine rather than .
LSTMs for Long Sequences
The advantage of LSTMs becomes measurable when we compare gradient propagation to vanilla RNNs. Recall from the Vanishing Gradients chapter that vanilla RNN gradients decay exponentially: each timestep multiplies the gradient by a factor typically less than 1. For typical random initialization, this factor is around 0.7 to 0.8, making gradients negligible after 50 or so steps.
LSTMs change this substantially. The additive cell state update allows gradients to flow backward through the addition operation without multiplication by a small value. The forget gate introduces some scaling, but a well-trained forget gate learns values close to 1 for important information, creating a near-constant gradient highway.
The mathematical story is worth tracing explicitly. In a vanilla RNN, the gradient of the loss with respect to the hidden state at step , given the loss depends on step , involves a product of Jacobian matrices, each associated with the tanh nonlinearity and the weight matrix:
Each factor in this product contains the derivative of tanh, which saturates near 0 for large inputs. In the LSTM, the analogous quantity for the cell state path involves only the forget gates:
A product of forget gate values is much gentler than a product of tanh Jacobians. If the network learns forget gates close to 1, this product stays near 1 for many steps, preserving the gradient signal over arbitrarily long distances.
def simulate_gradient_magnitude(seq_length, architecture="vanilla"):
"""
Estimate gradient magnitude reaching the first timestep.
Vanilla RNN: exponential decay per step due to tanh saturation
and matrix multiplication. Typical factor around 0.75.
LSTM: much gentler decay because cell state gradients bypass
repeated matrix multiplications. Forget gate factor near 0.99.
"""
if architecture == "vanilla":
decay_factor = 0.75
else:
decay_factor = 0.992
return decay_factor**seq_length
seq_lengths = np.arange(1, 201)
vanilla_gradients = [
simulate_gradient_magnitude(l, "vanilla") for l in seq_lengths
]
lstm_gradients = [simulate_gradient_magnitude(l, "lstm") for l in seq_lengths]
threshold = 1e-6
vanilla_cutoff = next(
(l for l, g in zip(seq_lengths, vanilla_gradients) if g < threshold), None
)
lstm_cutoff = next(
(l for l, g in zip(seq_lengths, lstm_gradients) if g < threshold), None
)Gradient falls below 1e-06: Vanilla RNN: at sequence length ~49 LSTM: remains above threshold for all tested lengths

The difference is substantial. Vanilla RNN gradients become negligible after around 50 timesteps. LSTM gradients remain orders of magnitude larger and stay above the practical training threshold for sequences well over 100 steps. This gradient stability is why LSTMs enabled practical sequence modeling at scale: tasks like machine translation, speech recognition, and language modeling all involve dependencies that span many dozens of tokens.
Stacking LSTMs
For complex tasks, it is common to stack multiple LSTM layers on top of each other. The output hidden states of the first LSTM layer become the input to the second layer, and so on. Each layer operates on the sequence of hidden states produced by the layer below, allowing deeper layers to model longer-range or more abstract patterns.
A two-layer LSTM processes a sequence as follows. The first layer sees the raw input features and produces hidden states . These serve as the input to the second layer, which produces . The second layer's outputs can be used for prediction.
Stacking layers tends to improve performance up to a point. Two or three layers are common in many NLP applications. Beyond three layers, training stability can become an issue: the gradients that must flow backward through both time and depth face a compounding attenuation. Dropout between layers helps regularize stacked LSTMs and is almost always used in practice when stacking more than one layer. In PyTorch, the num_layers argument handles stacking automatically, and the dropout argument applies dropout between all layers except the last.
One subtlety with stacked LSTMs: each layer has its own cell state and hidden state. When you initialize the LSTM with hidden = (h_0, c_0), both tensors have shape (num_layers, batch, hidden_size) to accommodate one state per layer. This is easy to overlook when moving from single-layer to multi-layer configurations.
Bidirectional LSTMs
A standard LSTM processes sequences left to right. At each timestep, it has seen all previous tokens but none of the future ones. For tasks like named entity recognition, sentiment analysis, or document classification, the model has access to the entire input at inference time. In those cases, each token's representation can benefit from both past and future context.
A bidirectional LSTM (Bi-LSTM) addresses this by running two LSTMs over the same sequence: one in the forward direction (left to right) and one in the backward direction (right to left). At each timestep, the forward LSTM has seen all tokens up to position , and the backward LSTM has seen all tokens from position onward. The hidden states from both directions are concatenated to produce a representation that encodes bidirectional context.
For position in a bidirectional LSTM with hidden size , the combined representation has dimension :
where is the forward hidden state and is the backward hidden state.
The gains from bidirectionality can be substantial for understanding tasks. Consider labeling the word "bank" in a sentence. In "He walked to the river bank to fish," the word "bank" is a geographical feature. In "He visited the bank to deposit a check," it is a financial institution. A unidirectional LSTM would need to predict the label for "bank" before seeing the disambiguating words later in the sentence. A bidirectional LSTM incorporates both the context before and after "bank" simultaneously, making the correct label much easier to determine.
The key limitation of bidirectional LSTMs is that they cannot be used for autoregressive generation. When generating text one token at a time, there is no future context to pass to the backward layer. Bi-LSTMs are therefore primarily used in encoder roles: understanding and representing existing text, not producing new text. In PyTorch, setting bidirectional=True in nn.LSTM enables bidirectional processing, and the output has shape (batch, seq_len, 2 * hidden_size).
PyTorch nn.LSTM Module
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.
LSTM for Sequence Classification
One of the most common LSTM applications is classifying entire sequences: deciding the sentiment of a review, the intent of a query, or the topic of a document. The approach is to run the LSTM over the full sequence and use the final hidden state as a fixed-size summary vector for a linear classifier. The LSTM does the hard work of compressing a variable-length sequence into a fixed-size representation; the classifier on top is just a linear layer applied to that representation.
Building the Classifier
class LSTMClassifier(nn.Module):
"""LSTM sequence classifier using the final hidden state."""
def __init__(
self, input_size, hidden_size, num_classes, num_layers=1, dropout=0.0
):
super().__init__()
self.lstm = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
dropout=dropout if num_layers > 1 else 0.0,
)
self.classifier = nn.Linear(hidden_size, num_classes)
def forward(self, x):
# x: (batch, seq_len, input_size)
_, (hidden, _) = self.lstm(x)
# hidden[-1]: final layer's hidden state (batch, hidden_size)
return self.classifier(hidden[-1])The simplicity of this architecture is one of LSTMs' strengths for classification. You stack an LSTM on top of an embedding layer, apply the LSTM to get a sequence of hidden states, take the last hidden state, and pass it through a linear layer. The LSTM does all the work of summarizing variable-length sequences into fixed-size vectors.
A variation that often works better in practice uses the mean of all hidden states rather than just the last one. The final hidden state represents the network's processing state after the last token, which gives disproportionate weight to the end of the sequence. Averaging over all positions gives equal weight to information throughout the sequence. Max pooling across the time dimension is another common alternative: take the element-wise maximum across time, which captures the strongest activation for each hidden state dimension regardless of when it occurred. In practice, mean pooling and max pooling both tend to outperform using only the final hidden state on classification tasks, especially for long sequences.
Training and Evaluating
import torch.optim as optim
torch.manual_seed(42)
model_cls = LSTMClassifier(input_size=8, hidden_size=32, num_classes=2)
optimizer_cls = optim.Adam(model_cls.parameters(), lr=0.01)
criterion_cls = nn.CrossEntropyLoss()
def make_cls_batch(batch_size=32, seq_len=20, input_size=8):
"""Synthetic binary classification: class determined by sequence mean."""
x = torch.randn(batch_size, seq_len, input_size)
labels = (x.mean(dim=(1, 2)) + 0.3 * torch.randn(batch_size) > 0).long()
return x, labels
cls_losses = []
for step in range(100):
model_cls.train()
xb, yb = make_cls_batch()
optimizer_cls.zero_grad()
logits = model_cls(xb)
loss = criterion_cls(logits, yb)
loss.backward()
optimizer_cls.step()
cls_losses.append(loss.item())
model_cls.eval()
with torch.no_grad():
x_test, y_test = make_cls_batch(batch_size=200)
preds = model_cls(x_test).argmax(dim=1)
test_acc = (preds == y_test).float().mean().item()Initial loss: 0.7037 Final loss: 0.6739 Test accuracy: 58.5%

Character-Level Language Model with LSTM
Language modeling is among the most natural applications for LSTMs. The model reads characters one at a time and learns to predict the next character. Because individual characters carry minimal meaning on their own, the model must internalize word-level, phrase-level, and sentence-level patterns to make accurate predictions.
Character-level modeling was historically significant because it required no tokenization decisions and could handle any input text without preprocessing. The famous 2015 blog post by Andrej Karpathy, "The Unreasonable Effectiveness of Recurrent Neural Networks," demonstrated that LSTM character models trained on large corpora could generate surprisingly coherent text, structured code, and even plausible mathematical notation. This work helped establish LSTMs as the dominant sequence modeling architecture before the transformer era and inspired many researchers to explore what kinds of structure an LSTM could learn to represent internally.
Dataset Preparation
# Sample text corpus for demonstration
corpus = (
"To be, or not to be, that is the question: "
"Whether 'tis nobler in the mind to suffer "
"The slings and arrows of outrageous fortune, "
"Or to take arms against a sea of troubles "
"And by opposing end them."
)
# Build character-level vocabulary
chars = sorted(set(corpus))
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for c, i in char_to_idx.items()}
vocab_size = len(chars)
encoded = [char_to_idx[c] for c in corpus]Corpus length: 197 characters Vocabulary size: 30 unique characters First 30 chars: 'To be, or not to be, that is t'
Model Architecture
The character model uses a small embedding layer before the LSTM. Although character indices could be passed directly as one-hot vectors, an embedding layer allows the model to learn a dense representation for each character. In a well-trained model, characters with similar roles (all lowercase vowels, all punctuation marks) tend to cluster together in the embedding space, making the downstream LSTM's job easier.
class CharLSTM(nn.Module):
"""Character-level language model using LSTM."""
def __init__(self, vocab_size, embed_dim, hidden_size, num_layers=1):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.lstm = nn.LSTM(
input_size=embed_dim,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
)
self.output_layer = nn.Linear(hidden_size, vocab_size)
def forward(self, x, hidden=None):
# x: (batch, seq_len) integer character indices
embeds = self.embedding(x) # (batch, seq_len, embed_dim)
lstm_out, hidden = self.lstm(embeds, hidden)
logits = self.output_layer(lstm_out) # (batch, seq_len, vocab_size)
return logits, hiddenNote that the LSTM returns the hidden state, which is passed back into the model during generation. This stateful generation pattern, where the hidden state is carried from one call to the next, is what enables the model to generate coherent text: each new character prediction takes into account the most recent character and the full history encoded in the recurrent state.
Preparing Training Batches
def make_char_batches(encoded, seq_len=25):
"""Create (input, target) pairs where target is input shifted by 1 position."""
inputs, targets = [], []
for i in range(0, len(encoded) - seq_len - 1, seq_len):
inputs.append(encoded[i : i + seq_len])
targets.append(encoded[i + 1 : i + seq_len + 1])
return torch.tensor(inputs), torch.tensor(targets)
seq_len_lm = 25
inputs_lm, targets_lm = make_char_batches(encoded, seq_len=seq_len_lm)Number of training sequences: 7 Sequence length: 25 Example: Input: 'To be, or not to be, that' Target: 'o be, or not to be, that '
The input-target offset by one position is the standard autoregressive training setup. The model sees a character and must predict the next one. By training on all consecutive pairs in the corpus, the model learns the statistical regularities of the text: which characters commonly follow which others, which word patterns are frequent, and which punctuation conventions appear.
Training the Character LM
torch.manual_seed(42)
char_model = CharLSTM(
vocab_size=vocab_size, embed_dim=32, hidden_size=64, num_layers=2
)
char_optimizer = optim.Adam(char_model.parameters(), lr=0.005)
char_criterion = nn.CrossEntropyLoss()
char_losses = []
for epoch in range(150):
char_model.train()
epoch_loss = 0.0
for inp, tgt in zip(inputs_lm, targets_lm):
inp_b = inp.unsqueeze(0) # (1, seq_len)
tgt_b = tgt.unsqueeze(0)
char_optimizer.zero_grad()
logits, _ = char_model(inp_b)
loss = char_criterion(logits.view(-1, vocab_size), tgt_b.view(-1))
loss.backward()
# Gradient clipping prevents occasional exploding gradients in stacked LSTMs
torch.nn.utils.clip_grad_norm_(char_model.parameters(), max_norm=1.0)
char_optimizer.step()
epoch_loss += loss.item()
char_losses.append(epoch_loss / max(len(inputs_lm), 1))Initial loss: 3.2606 Final loss: 0.0108 Improvement: 99.7%

Generating Text
def generate_text(model, seed_text, num_chars=100, temperature=0.8):
"""
Generate characters autoregressively from a seed string.
Temperature scales the logits before softmax sampling.
Lower values produce more predictable output; higher values more varied.
"""
model.eval()
generated = seed_text
hidden = None
with torch.no_grad():
# Prime the LSTM hidden state with the seed characters
for char in seed_text[:-1]:
idx = torch.tensor([[char_to_idx[char]]])
_, hidden = model(idx, hidden)
# Generate new characters one at a time
next_char = seed_text[-1]
for _ in range(num_chars):
idx = torch.tensor([[char_to_idx[next_char]]])
logits, hidden = model(idx, hidden)
probs = torch.softmax(logits[0, -1] / temperature, dim=0)
next_idx = torch.multinomial(probs, 1).item()
next_char = idx_to_char[next_idx]
generated += next_char
return generatedSeed: 'To be' Generated: To be, or not to be, that a s eo thet aks againd take arms againd a sea of troubles And againd the be, th
Even on a small corpus with brief training, the LSTM generates plausible character sequences. The temperature parameter plays an large effect in generation quality. Low temperatures (0.3 to 0.5) make the model more conservative: it selects characters with high probability more consistently, producing repetitive but well-formed text. High temperatures (1.2 to 1.5) make the model more exploratory: it sometimes picks lower-probability characters, producing more varied but potentially less coherent output. Temperature 1.0 corresponds to unmodified sampling from the model's distribution. In practice, values around 0.7 to 0.9 often produce a good balance of coherence and diversity.
Given more data and training, the model would learn word boundaries, punctuation conventions, and longer structural patterns. This is the foundation underlying language models that operate at the subword level, though modern systems have replaced LSTMs with transformers as the underlying architecture.
Key Parameters
When configuring nn.LSTM:
- input_size: Dimensionality of the input at each timestep. For embedded text, this is the embedding dimension.
- hidden_size: Dimensionality of both the hidden state and cell state. Larger values (128 to 1024) provide more capacity but increase computation.
- num_layers: Number of stacked LSTM layers. More layers can model more complex patterns but require more gradient clipping and careful tuning.
- batch_first: When
True, expects input shape(batch, seq_len, features). UsingTruealigns with standard data loaders. - dropout: Applied between LSTM layers when
num_layers > 1. Values in the range 0.1 to 0.5 reduce overfitting. - bidirectional: When
True, processes the sequence in both directions simultaneously, doubling the effective hidden size. Useful for classification but inapplicable for generation.
LSTM Variants and Alternatives
The original LSTM architecture spawned a family of related designs that modify the gating structure in different ways. Understanding these variants illuminates which parts of the LSTM architecture are necessary and which admit alternatives.
Peephole Connections
Standard LSTM gates receive only and as inputs. A variant proposed by Gers and Schmidhuber (2000) adds "peephole connections" that also feed the cell state directly into the gates:
The intuition is that the gates should be able to look at the memory contents in addition to the previously exposed hidden state. Consider a clock-like memory: the gate deciding when to reset a counter would benefit from seeing the current count directly, rather than inferring it through the hidden state. Peephole connections add a small number of parameters per gate and can improve performance on tasks requiring precise timing, such as speech recognition with long silence gaps. In practice, the improvement is often modest, and peephole connections are not included in PyTorch's standard nn.LSTM.
Gated Recurrent Units
The Gated Recurrent Unit (GRU), introduced by Cho et al. (2014), simplifies the LSTM by merging the cell state and hidden state into a single vector and replacing the three-gate system with two gates: a reset gate and an update gate. The GRU has fewer parameters than an LSTM of the same hidden size and tends to train faster. On many tasks, GRUs and LSTMs perform comparably. The choice between them is often empirical. The next chapter covers GRUs in depth.
Coupled Forget and Input Gates
The LSTM's forget gate and input gate are independent, which means it is possible for both to have high values simultaneously. A coupled variant forces them to be complementary:
This ensures that the total information in each cell state dimension remains roughly constant: if the forget gate retains 70% of the old value, the input gate contributes exactly 30% of the new candidate. This hard coupling reduces the parameter count and enforces a conservation-like property on the cell state. It works well on some tasks but can be too restrictive for others.
Practical Training Tips
Training LSTMs effectively requires attention to several details beyond the architecture itself. These practical considerations can mean the difference between a model that trains smoothly and one that fails to converge.
Gradient Clipping
Despite the LSTM's protection against vanishing gradients, exploding gradients remain a risk, particularly in deep or wide LSTMs. A single unusually large gradient update can destabilize the weights irreversibly. Gradient clipping addresses this by rescaling the gradient vector if its norm exceeds a threshold. This is standard practice for any RNN training:
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()The threshold max_norm=1.0 is a common default, though values up to 5.0 are used in some applications. If training loss is erratic or shows sudden spikes, increasing the clipping threshold (or checking that clipping is enabled) is usually the first debugging step.
Weight Initialization
The LSTM's weight matrices are typically initialized with small random values, often from a uniform distribution over . PyTorch's nn.LSTM uses this initialization by default. More specialized initializations, such as orthogonal initialization for the recurrent weight matrix , can improve training on tasks with very long sequences by preserving the norm of the gradient through time.
The forget gate bias initialization is one of the most reliable practical tricks. Setting or at initialization costs nothing and often measurably improves convergence on tasks with long-range dependencies. The intuition is that a freshly initialized LSTM should err on the side of remembering rather than forgetting, and a large positive forget gate bias achieves exactly that. PyTorch does not do this by default, but it is easy to apply after construction by directly modifying the bias parameter.
Truncated Backpropagation Through Time
Training on very long sequences requires storing all intermediate hidden and cell states for the backward pass. For sequences of thousands of tokens, this can exhaust GPU memory. Truncated BPTT addresses this by processing the sequence in fixed-length chunks, backpropagating only within each chunk while carrying the hidden state forward across chunk boundaries.
The trade-off is explicit: truncated BPTT prevents the model from learning dependencies that span chunk boundaries. A chunk size of 35 tokens means the model cannot directly learn from dependencies longer than 35 tokens via gradients. In practice, many of the dependencies that matter in language are shorter than common chunk sizes, and the cell state's ability to carry information forward across chunk boundaries provides some implicit long-range memory even when gradients do not flow backward across them. Truncated BPTT is the standard training regime for language model training on modern hardware.
Limitations and Impact
The LSTM architecture has clear limitations that are worth understanding before choosing it for a task.
The most significant practical limitation is sequential computation. Computing requires , which requires , and so on. This causal dependency means that no matter how many processors are available, an LSTM must take sequential steps for a sequence of length . Modern GPUs designed for massively parallel computation cannot be fully utilized on long sequences under this constraint. Transformers, by contrast, process all positions simultaneously in a single attention pass, making them dramatically faster to train when sequences are long and batches are large. This sequential bottleneck is the primary reason LSTMs have been largely displaced by transformers in large-scale language modeling, despite their excellent gradient properties.
Memory requirements during training also scale linearly with sequence length. Backpropagation through time must retain every intermediate hidden and cell state so gradients can be computed on the backward pass. For sequences of thousands of tokens, this memory cost becomes prohibitive on GPU. Truncated BPTT mitigates the cost by limiting how many steps backward gradients are propagated, at the price of sacrificing the ability to learn the very longest dependencies.
LSTMs also do not provide a natural mechanism for direct attention between arbitrary timesteps. If the model needs to compare the token at position 10 with the token at position 100, it must carry the position 10 information forward through 90 hidden states before making the comparison. Attention mechanisms, which allow any two positions to interact directly in a single computation step, provide a fundamentally richer interface for long-range relationship modeling. This is one of the key architectural advantages of the transformer that drove its widespread adoption.
Despite these limitations, LSTMs fundamentally changed sequence modeling. Before LSTMs, no neural architecture could reliably learn dependencies spanning more than a few dozen tokens. After LSTMs, machine translation, speech recognition, handwriting recognition, and language modeling at scale all became tractable. The Hochreiter and Schmidhuber (1997) paper remains among the most cited in all of machine learning, with over 70,000 citations as of 2024.
LSTMs also introduced a new design primitive: the gated memory cell. Before LSTMs, most network layers transformed information through learned nonlinear functions. LSTMs showed that networks could also learn to dynamically route and preserve information, with the routing decisions driven by the data itself. This gating idea directly inspired the GRU (Gated Recurrent Unit), covered in the next chapter, and more distantly influenced the attention mechanisms that underpin Transformer models. The residual connections essential to training deep transformers also echo the LSTM cell state's strategy of enabling identity paths for information to flow without being transformed at every step.
Historical Context
Understanding the historical context of LSTMs helps appreciate the significance of the contribution. In the mid-1990s, several groups were working on the vanishing gradient problem, but most approaches involved either gradient clipping, careful initialization, or modifications to the training algorithm rather than changes to the architecture itself. Hochreiter's 1991 diploma thesis, written before the LSTM paper, contained a detailed analysis of why gradients vanish in standard RNNs, laying the theoretical foundation for the solution.
The 1997 LSTM paper was initially received with limited enthusiasm, partly because the architecture seemed complicated relative to simpler RNNs and partly because the computing hardware of the era was not capable of training LSTMs on the scale needed to demonstrate their advantages clearly. It was not until the 2010s, when GPU-based deep learning made large-scale sequence modeling practical, that LSTMs became the dominant approach. The breakthrough applications, including Google's translation system, Apple's Siri, Amazon's Alexa, and numerous speech recognition systems, all relied heavily on LSTMs trained at scales that would have been unimaginable in 1997.
The story of LSTMs is a story about an idea that was correct and important, but required the right computational context to be fully appreciated. It also illustrates a pattern common in deep learning: architectural innovations often precede by years or decades the practical context in which they prove decisive. Understanding the architecture deeply positions you to recognize when a similar conceptual advance becomes practically relevant in the future.
Summary
This chapter introduced the LSTM architecture and the core innovations that allow it to learn long-range dependencies.
The central idea is the cell state: an information highway that accumulates information through additive updates rather than repeated matrix multiplications. This additive structure keeps gradients from vanishing over long sequences, allowing the network to propagate learning signals across hundreds of timesteps. The mathematical reason is direct: gradients flowing back through the additive cell state update are multiplied by the forget gate value at each step, rather than by the derivative of a saturating nonlinearity, keeping the gradient signal alive over long distances.
Three gates coordinate access to the cell state:
- The forget gate selectively discards information from the previous cell state, controlled by a learned sigmoid layer that examines the current input and previous hidden state. Initializing its bias to a positive value at the start of training is a simple trick that improves long-range learning.
- The input gate and candidate layer together control what new information to write into memory: the candidate layer proposes values using tanh (enabling both positive and negative updates), while the input gate scales how much of each proposal is applied. This two-part design decouples what to write from how much to write.
- The output gate filters the cell state to produce the current hidden state output, creating a fundamental distinction between what the network knows (cell state) and what it currently reports (hidden state).
These gates are learned, not fixed. They specialize during training based on the task, learning when to remember, when to forget, and what to output. The PyTorch nn.LSTM module packages all six equations into a single optimized call, with output providing per-timestep hidden states and hidden, cell providing the final states.
LSTMs introduced the gating abstraction into neural network design, a contribution whose influence extends far beyond recurrent networks. The GRU simplifies the LSTM by merging the cell and hidden states and using two gates instead of three. Transformer attention generalizes the idea of selective access to memory through a fundamentally different mechanism that operates without sequential computation. Both of these advances owe a conceptual debt to the LSTM's demonstration that learned routing of information can be more powerful than learned transformation alone.
In the LSTM Gate Equations chapter, we will work through the mathematics precisely: derive the parameter count, implement all six equations from scratch in NumPy, and verify that the result matches PyTorch. That mathematical grounding will prepare you for LSTM Gradient Flow, where we analyze exactly why the cell state path prevents vanishing gradients.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about LSTM architecture.
LSTM Architecture Quiz
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of Language AI Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore Language AI HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.

Comments
No comments yet. Be the first to share your thoughts!