LSTM
Advanced 16 min read LSTM

LSTM Internals

The exact math of the gated memory cell

The gated math, formally: four gates (forget, input, candidate, output) combine linear maps over [h, x] with sigmoid/tanh activations. This concept nails the equations and shapes.

LSTM equationscandidatesigmoidpeepholehidden sizecell update
Prerequisites: LSTM Intuition
Before we start

What you will learn

  • Write the four LSTM update equations
  • Compute output shapes through a layer
  • Explain why cand-tanh on the belt is bounded
Introduction

What is this?

With input x_t (size d) and hidden h_(t-1) (size H), an LSTM computes a concatenated linear map z = W·[h_(t-1), x_t] + b (size 4H), splits it into forget/input/candidate/output, then: f = σ(f_gate), i = σ(i_gate), g = tanh(g_gate), o = σ(o_gate); c_t = f⊙c_(t-1) + i⊙g; h_t = o⊙tanh(c_t). All four gates share one matmul — that's why an LSTM layer is 4x an RNN's parameters.

Motivation

Why do you need it?

  • Reading the equations is how you tune init, inspect activations, and build custom cells.
  • The shared-matmul structure explains param counts (4·H·(H+d)) and why LSTMs cost 4x an RNN per step.
  • Knowing gate value ranges (sigmoid ∈ (0,1), tanh ∈ (−1,1)) is how you read cell-state norms.
  • It's the foundation for GRU and modern gated variants — difference in one line.
Think about it

Mental model

Imagine

A bank ledger updated by four clerks

Think of

The four LSTM gates

Clerk F stamps 'forgive this debit' (forget), clerk I approves new credits (input), clerk G drafts the amount (candidate), clerk O decides what's shown in the window (output). The ledger (c) is the authoritative record, the window (h_t) what customers see.

Mechanics

How it works

  • Precompute z = W·[h_(t-1), x_t] + b once: one (4H) logits vector.
  • Split: zf, zi, zg, zo = z[:H], z[H:2H], z[2H:3H], z[3H:].
  • f = σ(zf), i = σ(zi), g = tanh(zg), o = σ(zo).
  • c_t = f ⊙ c_(t-1) + i ⊙ g; h_t = o ⊙ tanh(c_t).
  • Bias trick: an extra +b_fg on the forget gate pre-logits biases 'remember'.

Data flow inside one LSTM timestep.

[h,x]
W·[...] + b
σ σ tanh σ
c_t belt
h_t
Real world

In production

  • Fine-grained LSTM layers in ASR/forecasting harnesses where the 4x cost is accepted for quality.
  • Custom cells with peepholes or layer-norm built from these equations power specialized speech models.
  • Distilled students sometimes replace big attention models with friendlier LSTMs at the serving tier.
Guidelines

Best practices

  • Count params: 4·H·(H+d) per layer — budget it before stacking.
  • Reuse one linear layer for all four gates (einsum) rather than four separate mats for speed and clarity.
  • Gate out of saturation: keep pre-activations moderate with init scaling and small LR.
  • Architect foresight: forget-bias initialization (+1) belongs in the constructor, not after.
Watch out

Common mistakes

  • Permuting the four gate slices wrong (a classic silent bug in hand-written cells).
  • Forgetting the tanh on the candidate — keeping it linear blows the belt unbounded.
  • Using i⊙g vs g⊙i convention mismatch across frameworks — keep it consistent.
  • Confusing the 4H split with 4 separate matmuls — params and cost are the same, correctness isn't.
Root cause analysis

Common issues

Error
Custom LSTM differs from torch's after a few steps
Root cause

Gate slice ordering or bias placement (torch uses coy bias_ih+bias_hh split)

Fix

Match the framework's exact ordering; check bias pattern explicitly.

Error
Belt norm → ±4 (out of tanh range)
Root cause

Linearly-mixed c path (candidate not tanh-clamped) or a ragged init

Fix

Verify g = tanh(zg) and c update; shift forget bias up; clip.

Error
Gradients still blow up near cell output
Root cause

Output gate + hidden path lacked clipping while c-path stayed stable

Fix

Clip global grad norm and normalize h rows; keep c-path guarded.

When things break

Debugging tips

  • Sanity: with all gates forced to extremes (f→1,i→0,o→1) the cell should behave like the identity/RNN baseline.
  • Print c_t and h_t norms: belt bounded by tanh; gates in (0,1); violations mean a math bug.
  • Gradient-check one timestep numerically before wiring a custom cell into training.
  • Compare framework vs your hand-rolled cell on identical weights — catches slice/init mismatches immediately.
Interviews

Interview checkpoint

Practice

Hands-on lab

Implement an LSTM Cell from Scratch
35 min · advanced
In progress
Objectives
  • Implement the four-gate LSTM math in numpy
  • Match a framework LSTMCell's output on identical weights
  • Prove the belt keeps near-constant scale with forget open

Implement an LSTM Cell from Scratch

35 min · advanced

0/4 steps
0%
sandbox · Linux sandbox
$
Try these commands
Test yourself

Quick quiz

0/2
1

Which activation bounds the cell-state candidate?

2

An LSTM layer performs four separate matmuls per step.

Go deeper

Official documentation

PyTorch LSTMCell docs
docs