🧠 LSTM — Long Short-Term Memory

Real-world scenario: Imagine reading a long mystery novel. When you reach the final chapter and the detective reveals the killer, you instantly recall a seemingly unimportant detail from chapter 3 — a red scarf mentioned in passing. You remembered it because something told you it mattered. You didn't remember every sentence equally; you selectively retained the important parts and discarded the filler. LSTM networks work exactly like this — they have explicit gates that decide what to remember, what to update, and what to output at each step.

🌟 Why LSTM?

RNNs fail on long sequences because gradients vanish during backpropagation — by the time the network tries to learn from something that happened 50 steps ago, the gradient connecting that event to the current step is essentially zero.

LSTM (introduced by Hochreiter & Schmidhuber, 1997) solved this by adding a cell state — a "conveyor belt" running through the entire sequence that gradients can flow through almost unchanged. Three learned gates control what information gets added, removed, or read from this conveyor.

The key insight: rather than forcing information to flow through the same weight matrix at every step (which causes gradient decay), LSTM uses additive updates to the cell state, which allow gradients to flow back hundreds of steps without vanishing.

Cell State

The cell state C_t is LSTM's long-term memory. It runs in parallel to the hidden state and can carry information across many time steps with minimal modification.

Unlike the hidden state (which is squashed through tanh at every step), the cell state is updated additively — new information is added to it, not multiplied through a non-linearity at each step. This is the structural reason why LSTM can learn long-range dependencies.

C_t = f_t ⊙ C_{t-1}  +  i_t ⊙ g_t

where:
  f_t = forget gate output  (how much of old cell state to keep)
  i_t = input gate output   (how much of new info to add)
  g_t = candidate values    (what new information to potentially add)
  ⊙   = element-wise multiplication

Visually: the cell state is a highway. The gates are on-ramps and off-ramps.

Hidden State

The hidden state h_t is LSTM's short-term, working memory — the output that's passed to the next time step AND used as the current time step's prediction.

Unlike the cell state, the hidden state is a filtered, compressed version of the cell state:

h_t = o_t ⊙ tanh(C_t)

where o_t is the output gate. The tanh squashes cell state values to [-1, 1], and the output gate decides how much of that to expose.

Think of it this way:

• Cell state = your long-term memory (what you've learned and remembered)
• Hidden state = what you're actively thinking about right now (a filtered read from long-term memory)

Forget Gate

The forget gate decides how much of the previous cell state to keep.

f_t = σ(W_f · [h_{t-1}, x_t] + b_f)

Output is between 0 and 1 (sigmoid). A value of 1 means "keep everything"; 0 means "forget everything."

import torch
import torch.nn as nn

# Toy example: forget gate in isolation
input_size, hidden_size = 4, 8
W_f = nn.Linear(input_size + hidden_size, hidden_size)

h_prev = torch.zeros(1, hidden_size)   # previous hidden state
x_curr = torch.randn(1, input_size)    # current input

combined = torch.cat([h_prev, x_curr], dim=1)
f_t = torch.sigmoid(W_f(combined))     # forget gate: values in [0, 1]

print(f"Forget gate: {f_t.squeeze().detach().numpy().round(2)}")
# e.g. [0.45  0.61  0.32  0.78  0.55  0.42  0.69  0.51]
# Each value says: keep this fraction of the corresponding cell state unit

In a trained model, the forget gate learns to output values near 1 for information that's still relevant (e.g., the subject of a sentence) and near 0 for information that's now irrelevant (e.g., the previous sentence's topic).

Input Gate

The input gate decides how much new information to write into the cell state. It works together with the candidate memory (see below):

i_t = σ(W_i · [h_{t-1}, x_t] + b_i)

Like the forget gate, output is in [0, 1]. A high value means "write a lot of this new information"; near zero means "don't update the cell state much this step."

W_i = nn.Linear(input_size + hidden_size, hidden_size)
i_t = torch.sigmoid(W_i(combined))     # input gate: how much to write

print(f"Input gate: {i_t.squeeze().detach().numpy().round(2)}")
# e.g. [0.52  0.33  0.71  0.29  0.48  0.65  0.38  0.57]

The forget and input gates work in tandem: the forget gate removes old information, and the input gate adds new information. A common but incorrect assumption is that they're opposites (forgetting = not writing). They're independent — the network can both forget old info AND write new info at the same step.

Candidate Memory

The candidate memory g_t (sometimes called c̃_t) is a proposal for what new information could be written to the cell state:

g_t = tanh(W_g · [h_{t-1}, x_t] + b_g)

Values in [-1, 1] (tanh). Think of it as "here's what's potentially important about this new input." The input gate then decides how much of this proposal to actually accept:

cell_update = i_t ⊙ g_t

W_g = nn.Linear(input_size + hidden_size, hidden_size)
g_t = torch.tanh(W_g(combined))        # candidate: what to potentially add

cell_update = i_t * g_t                # how much actually gets written
print(f"Cell update: {cell_update.squeeze().detach().numpy().round(3)}")

Output Gate

The output gate decides how much of the cell state to expose as the hidden state:

o_t  = σ(W_o · [h_{t-1}, x_t] + b_o)
h_t  = o_t ⊙ tanh(C_t)

Even if the cell state contains a lot of information, the output gate can choose to only expose a fraction of it at this time step. This is useful for tasks like language modeling, where certain information stored in the cell state is relevant for future steps but not the current prediction.

# Full LSTM step from scratch
W_o = nn.Linear(input_size + hidden_size, hidden_size)
C_prev = torch.zeros(1, hidden_size)

f_t_new = torch.sigmoid(W_f(combined))
i_t_new = torch.sigmoid(W_i(combined))
g_t_new = torch.tanh(W_g(combined))
o_t_new = torch.sigmoid(W_o(combined))

C_t = f_t_new * C_prev + i_t_new * g_t_new   # new cell state
h_t = o_t_new * torch.tanh(C_t)              # new hidden state

print(f"Cell state: {C_t.shape}")    # (1, 8)
print(f"Hidden state: {h_t.shape}")  # (1, 8)

LSTM Limitations

Despite solving the vanishing gradient problem, LSTMs still have significant limitations:

1. Sequential bottleneck LSTMs process tokens one at a time — step 5 can't be computed until step 4 is done. No parallelism means slow training on GPUs.

2. Fixed-size cell state Even with gating, a 512-dimensional cell state must compress an arbitrarily long document into 512 numbers. For very long texts, information still gets compressed and lost.

3. Long-range dependencies remain difficult LSTMs handle dependencies up to ~50-100 steps much better than RNNs. But for truly long-range dependencies (hundreds or thousands of steps), they still struggle. A character at position 1 and a character at position 800 have a very thin thread connecting them.

4. Attention does it better The Transformer's attention mechanism directly connects any two positions in a sequence in a single step — no recurrence required. For most NLP tasks, attention-based models now outperform LSTMs, train faster, and scale better.

LSTM vs Transformer comparison:
                    LSTM        Transformer
Sequential?         Yes ❌      No ✅
Parallelisable?     No  ❌      Yes ✅
Long-range deps?    Medium 🟡  Excellent ✅
Memory efficient?   Yes ✅      Less so ❌
Still used?         Rarely      Everywhere ✅

LSTMs are still found in:

• Edge/embedded devices where transformers are too large
• Time series forecasting where sequential processing is natural
• Legacy production systems

But for language understanding, LSTMs have been almost completely replaced by Transformers.