Explain the role of each LSTM gate (forget, input, output) and the cell state, using both equations and plain language.
Derive how the LSTM cell state update provides an additive gradient pathway that mitigates vanishing gradients.
Compare LSTM and GRU architectures — parameter counts, gating structure, and typical empirical tradeoffs.
Implement LSTM and GRU models in PyTorch and diagnose their behaviour on a controlled long-range dependency task.
Justify the choice between vanilla RNN, LSTM, and GRU for a given sequence modelling setting using evidence from learning curves and validation performance.
In a vanilla RNN, the hidden state \(\mathbf{h}_t \in \mathbb{R}^H\) must:
🤔 Remember anything from the past that might be useful later (e.g., the subject of a sentence for agreement checking 20 words away).
💭 Forget everything that is no longer relevant.
🧮 Represent the current input well enough to produce a useful output now.
These three goals conflict:
If \(W_{hh}\) is tuned to preserve old information → gradient of current prediction w.r.t. parameters gets diluted.
If \(W_{hh}\) is tuned for current responsiveness → old information is overwritten.
Key design question
Can we separate long-term memory from short-term working representation? LSTM says: yes — give them different variables with different update rules.
LSTM: the cell state and its gates
The Long Short-Term Memory(Hochreiter and Schmidhuber 1997) adds a cell state\(\mathbf{c}_t \in \mathbb{R}^H\) alongside the hidden state \(\mathbf{h}_t\).
This is a diagonal matrix. When \(\mathbf{f}_t \approx \mathbf{1}\) (forget gate open), the gradient flows back through \(c\) with roughly no shrinkage.
Compare to vanilla RNN where the Jacobian is \(\operatorname{diag}(1 - \mathbf{h}_{t-1}^2) \cdot W_{hh}\) — a full matrix multiply at each step.
Takeaway
The cell state provides a near-identity gradient path when the forget gate is near 1.
The model learns when to open or close this path — it is not hardcoded.
🔬 Gradient flow — formal sketch
For a loss \(\mathcal{L}\) at the final step \(T\), the gradient w.r.t. cell state at time \(t\) is:
Because each factor is diagonal and elementwise-bounded by \([0, 1]\):
When \(f_k^{(j)} \approx 1\) for all \(k\): the \(j\)-th gradient component survives all \(T - t\) steps undiluted.
When \(f_k^{(j)} \approx 0\) for some \(k\): the \(j\)-th gradient is intentionally killed — the network decided this memory slot is irrelevant beyond step \(k\).
Each factor is a full matrix, and powers of matrices with spectral radius \(< 1\) shrink exponentially.
The structural difference
LSTM replaces full-matrix products (multiplicative, potentially shrinking) with element-wise products (controllable, potentially identity-like). This is the core of why LSTMs can learn dependencies across hundreds of steps while vanilla RNNs cannot.
GRU: two gates, no separate cell state
The Gated Recurrent Unit(Cho et al. 2014) simplifies the LSTM by merging the cell state and hidden state into one, with two gates:
Reset gate\(\mathbf{r}_t\): how much of the past hidden state the candidate should see. \(\mathbf{r}_t \approx 0\) forces \(\tilde{\mathbf{h}}_t\) to ignore the past.
Update gate\(\mathbf{z}_t\): interpolation coefficient between old hidden state and new candidate. \(\mathbf{z}_t \approx 0\) → “carry old”; \(\mathbf{z}_t \approx 1\) → “write new”.
Long sequences, tasks needing fine-grained control
Rule of thumb
Start with LSTM. If training time matters and sequences are not extremely long, try GRU. Vanilla RNNs are mainly used for pedagogical or very constrained settings today.
❌ Anti-example: using LSTM when vanilla RNN would suffice
Scenario:
Task: classify 3-word product titles as “electronics” or “clothing”.
Sequences are short; no long-range dependency exists.
A student reaches for nn.LSTM by default and reports slightly worse validation accuracy than a nn.RNN trained for the same number of steps.
Issues:
Unnecessary complexity: LSTM has \(4\times\) the parameters of a vanilla RNN for the same hidden size. With very little data and short sequences, this extra capacity increases overfitting risk.
Slower convergence: More parameters → more gradient updates needed before the gates settle into useful configurations.
Misdiagnosis: The student suspects “LSTMs are worse” and may over-generalise. In reality, LSTMs are solving a problem (long-range gradients) that simply does not exist here.
Takeaway:
Match the architectural complexity to the task. The key question is always: how far back in the sequence does the relevant context actually reach?
For sequences of length ≤ 10–15, a well-regularized vanilla RNN or even a simple bag-of-words approach is often competitive.
Worked example 1
Long-range memory: T = 100
In M48 we showed a vanilla RNN can retrieve a seed value from 30 steps ago.
Now we push to T = 100 steps.
Task:
Input: a random walk \(x_0, x_1, \ldots, x_{T-1}\) of length \(T = 100\).
Target: predict the seed value\(x_0\) from the final hidden state.
Why this is a clean test:
There is exactly one relevant timestep: step 0.
A model that generalises must carry \(x_0\) in its state for 100 steps without corrupting it.
The theoretical noise floor (for an MLP or any model that can’t see step 0) is \(\text{Var}(x_0) \approx 0.083\).
We compare three models with the same hidden size (\(H = 32\)): vanilla RNN, LSTM, GRU.
At T=100, the vanilla RNN typically fails to cross the naive floor — it cannot reliably remember \(x_0\).
Both LSTM and GRU cross it and approach the noise floor, demonstrating that gated architectures can carry information across 100 steps while vanilla RNNs cannot.
Binary task: classify a movie review snippet as 👍positive or 👎negative.
Average review length: ~20 tokens — moderate, not trivially short.
Vocabulary: top 5 000 words; fixed 30-token padded sequences.
Why this task?
Sentiment often depends on negation and context that spans several words: “not at all bad” — the word “not” needs to influence the interpretation of “bad” 3 steps later.
Long-range dependencies are real but not extreme, making this a good intermediate test.
Goal: compare vanilla RNN, LSTM, and GRU on accuracy and training stability.
Load and preprocess SST-2 (subset)
Code
from collections import Counterimport refrom datasets import load_datasetimport random# ── Real SST-2 subsample (5 000 train + 400 validation) ───────────────────────# Loads from HuggingFace Hub; requires network access.# The official SST-2 validation split has no labels publicly, so we carve our# own val set from the training split.N_SAMPLES =5000VAL_FRAC =0.2MAX_LEN =30VOCAB_SIZE =5000EMBED_DIM =64random.seed(3151)np.random.seed(3151)torch.manual_seed(3151)torch.cuda.manual_seed_all(3151)torch.backends.cudnn.deterministic =Truetorch.backends.cudnn.benchmark =Falseraw = (load_dataset("stanfordnlp/sst2", split="train") .shuffle(seed=3151) .select(range(N_SAMPLES)))# ── Build a simple word-level vocabulary from this subsample ──────────────────def tokenize(text):return re.findall(r"[a-z]+", text.lower())all_tokens = [tok for ex in raw for tok in tokenize(ex["sentence"])]vocab_counts = Counter(all_tokens)# token→index: 0=PAD, 1=UNK, then top (VOCAB_SIZE-2) wordsidx2tok = ["<PAD>", "<UNK>"] + [w for w, _ in vocab_counts.most_common(VOCAB_SIZE -2)]tok2idx = {w: i for i, w inenumerate(idx2tok)}def encode(text, max_len=MAX_LEN): ids = [tok2idx.get(t, 1) for t in tokenize(text)][:max_len] ids += [0] * (max_len -len(ids)) # right-pad with PAD=0return idsX_sst = np.array([encode(ex["sentence"]) for ex in raw], dtype=np.int64)y_sst = np.array([ex["label"] for ex in raw], dtype=np.int64)split =int((1- VAL_FRAC) * N_SAMPLES)Xtr_s, Xvl_s = X_sst[:split], X_sst[split:]ytr_s, yvl_s = y_sst[:split], y_sst[split:]g = torch.Generator()g.manual_seed(3151)sst_train_loader = DataLoader( TensorDataset(torch.from_numpy(Xtr_s), torch.from_numpy(ytr_s)), batch_size=64, shuffle=True, generator=g)sst_val_loader = DataLoader( TensorDataset(torch.from_numpy(Xvl_s), torch.from_numpy(yvl_s)), batch_size=256)print(f"Train: {Xtr_s.shape} Val: {Xvl_s.shape}")print(f"Class balance — train: {np.bincount(ytr_s)}, val: {np.bincount(yvl_s)}")print(f"Vocab size used: {len(tok2idx)}")
For sequences of moderate length (~30 tokens), LSTM and GRU often outperform the vanilla RNN — but the gap shrinks as sequences get shorter. Always compare empirically rather than assuming “LSTM is always better.”
The train–val gap is a useful proxy for how much the model is overfitting given the small dataset.
Summary
Choosing your recurrent cell
Situation
Recommendation
Reason
Short sequences (< 15 tokens)
Vanilla RNN or even bag-of-words
No long-range dependency; extra gates add cost without benefit
Moderate sequences (15–100)
GRU first
Fewer params than LSTM, fast to train, competitive accuracy
Long sequences (100+) or tasks requiring fine-grained memory control
LSTM
Cell state highway helps; output gate decouples memory from representation
Compute-constrained / mobile
GRU
75 % of LSTM params; often matches performance
Cutting-edge NLP (2024+)
Transformer
Attention has replaced recurrence for most large-scale tasks
Quick diagnostic
Does your validation loss plateau while training loss keeps falling? → likely overfitting; reduce capacity or add dropout.
Do both losses plateau early at a high value? → likely underfitting; increase hidden size or train longer.
Does LSTM outperform GRU substantially? → the task probably requires fine-grained cell control.
LSTM/GRU can remember and condition on early context for hundreds of steps.
If the training corpus encodes stereotypes (e.g., gender–profession correlations in text), gated models can leverage that context more persistently than a vanilla RNN would.
Example: a sentiment model trained on review text may associate certain names or dialects with positive or negative sentiment and carry that context across the full review.
2. Memorisation of sensitive sequences:
Long hidden states in LSTMs trained on personal data (medical notes, chat logs, location traces) can encode individual-specific patterns.
These can be extracted via model inversion or membership-inference attacks.
Mitigation: differential privacy training (e.g., DP-SGD), careful data minimisation, and evaluation for memorisation before deployment.
3. Opaque temporal reasoning:
Unlike a bag-of-words model, a gated RNN’s decision is influenced by a hidden state that is hard to inspect.
When deployed in high-stakes settings (medical, legal, financial), temporal models should be paired with interpretability tools (attention visualization, SHAP for time series, or integrated gradients).
Simple check: Before deploying a gated RNN, ask: What information from the past could the model be weighting heavily? Could that information encode protected attributes or be used to discriminate?
Quick checks
Q1. An LSTM has input size \(d = 10\), hidden size \(H = 20\).
How many trainable parameters does this LSTM have (excluding any output head)?
\(20 \times (10 + 20) + 20 = 620\)
\(4 \times [20 \times (10 + 20) + 20] = 2480\)
\(4 \times [20 \times 10 + 20] = 960\)
\(3 \times [20 \times (10 + 20) + 20] = 1860\)
Q2. Short answer: In the LSTM cell update \(\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \tilde{\mathbf{c}}_t\), what happens when the forget gate \(\mathbf{f}_t \approx \mathbf{0}\) and the input gate \(\mathbf{i}_t \approx \mathbf{1}\)? Describe the effect on memory and identify a natural use case.
Model answer: The cell state is almost entirely replaced by the new candidate \(\tilde{\mathbf{c}}_t\). This is a “hard reset” — the old long-term memory is erased and rewritten with the current input. A natural use case: the beginning of a new sentence or document in a streaming text task, where the previous context is irrelevant and the model should start fresh.
Q3. In the GRU, the candidate hidden state is computed as \(\tilde{\mathbf{h}}_t = \tanh(W_h[\mathbf{r}_t \odot \mathbf{h}_{t-1};\, \mathbf{x}_t] + \mathbf{b}_h)\).
What does setting the reset gate \(\mathbf{r}_t \approx \mathbf{0}\) effectively do?
Forces the GRU to fully copy the old hidden state \(\mathbf{h}_{t-1}\) to \(\mathbf{h}_t\).
Makes the candidate \(\tilde{\mathbf{h}}_t\) ignore \(\mathbf{h}_{t-1}\), so the update depends only on \(\mathbf{x}_t\).
Sets \(\mathbf{h}_t \approx \mathbf{h}_{t-1}\) regardless of \(\mathbf{x}_t\).
Doubles the effective learning rate for this time step.
Q4. Numeric / reasoning: Consider a sequence classification task where the relevant label signal is contained only in the last 3 tokens of a length-50 sequence (e.g., “…, nothing, wrong, here → positive”).
A student trains a vanilla RNN for 200 epochs and a GRU for 100 epochs (same hidden size, same data). The GRU converges faster to higher validation accuracy.
Provide one architectural explanation for the GRU’s advantage, and one training-setup reason that might make the comparison unfair.
References
Cho, Kyunghyun, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. 2014. “Learning Phrase Representations Using RNN Encoder–Decoder for Statistical Machine Translation.” In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), 1724–34. Association for Computational Linguistics. https://doi.org/10.3115/v1/D14-1179.
