Module 4

Long Short-Term Memory (LSTM)

Gating mechanisms, LSTM cells, and character-level language modelling. Understand why LSTMs matter and how Transformers replaced them.

Learning Objectives

  • Understand gating mechanisms in sequence models
  • Implement LSTM cells and networks
  • Train on sequence prediction problems
  • Recognise LSTM's role in early language models

4.1 The LSTM Cell

Key Concepts

  • Forget gate — what to remove from memory
  • Input gate — what new information to add
  • Output gate — what to expose as the hidden state
  • Cell state — the "constant error carousel": protected long-term memory
LSTMs were dominant for sequence modelling before Transformers. Understanding them shows the evolution of sequence models and why each component was added.

LSTM Equations

f_t = σ(W_f·[h_{t-1}, x_t] + b_f) # Forget gate i_t = σ(W_i·[h_{t-1}, x_t] + b_i) # Input gate C̃_t = tanh(W_c·[h_{t-1}, x_t] + b_c) # Cell candidate C_t = f_t ⊙ C_{t-1} + i_t ⊙ C̃_t # Cell state update o_t = σ(W_o·[h_{t-1}, x_t] + b_o) # Output gate h_t = o_t ⊙ tanh(C_t) # Hidden state where σ = sigmoid, ⊙ = element-wise multiply
Haskell — LSTM step
data LSTMCell = LSTMCell
  { weightsInput  :: Matrix Double      -- Input weight for all gates + candidate
  , weightsHidden :: Matrix Double     -- Recurrent weight for all gates + candidate
  , biases        :: Vector Double     -- Biases for all gates + candidate
  , hiddenSize    :: Int
  }

lstmStep
  :: LSTMCell
  -> Vector Double                          -- Current input x_t
  -> (Vector Double, Vector Double)         -- (prev hidden h_{t-1}, prev cell c_{t-1})
  -> ((Vector Double, Vector Double), Vector Double)  -- ((new h, new c), output)
lstmStep cell input (prevHidden, prevCell) =
  let allGates = matmul (weightsInput cell) input
               + matmul (weightsHidden cell) prevHidden
               + biases cell

      gateSize       = hiddenSize cell
      forgetGate     = fmap sigmoid (slice allGates 0 gateSize)
      inputGate      = fmap sigmoid (slice allGates gateSize (2*gateSize))
      outputGate     = fmap sigmoid (slice allGates (2*gateSize) (3*gateSize))
      cellCandidate  = fmap tanh    (slice allGates (3*gateSize) (4*gateSize))

      newCell   = elementwiseMultiply forgetGate prevCell
                + elementwiseMultiply inputGate cellCandidate
      newHidden = elementwiseMultiply outputGate (fmap tanh newCell)

  in ((newHidden, newCell), newHidden)

4.2 Training LSTMs on Sequences

Implementation Task — Character-Level Language Model

Haskell — character-level data preparation
type CharData = String

-- Training: predict next character given previous characters
prepareSequenceData :: Int -> CharData -> [(String, Char)]
prepareSequenceData seqLen text =
  let sequences = [take seqLen (drop i text) | i <- [0..length text - seqLen - 1]]
      targets   = [text !! (i + seqLen)       | i <- [0..length text - seqLen - 1]]
  in zip sequences targets

-- Simple LSTM language model structure
data LSTMLanguageModel = LSTMLanguageModel
  { embedding  :: Matrix Double
  , lstmLayers :: [LSTMCell]
  , outputLayer :: (Matrix Double, Vector Double)
  , charToIdx  :: Char -> Int
  , idxToChar  :: Int  -> Char
  }
Haskell — gradient clipping
-- Use truncated BPTT (max 35–50 steps)
-- Clip gradients to prevent explosion:
clipGradientByNorm :: Double -> [Double] -> [Double]
clipGradientByNorm maxNorm grads =
  let norm  = sqrt (sum (map (^2) grads))
      scale = min 1.0 (maxNorm / norm)
  in map (*scale) grads

4.3 Why Transformers Displaced LSTMs

  • Parallelisation — RNNs must process sequentially (hidden state depends on previous). Transformers process entire sequences in parallel.
  • Long-term dependencies — despite LSTMs, still struggles beyond ~1000 tokens. Transformers handle this better via direct attention.
  • Scaling — Transformers scale better with both model size and data.

Reading & Viewing

Essential papers

Long Short-Term Memory (1997) Hochreiter & Schmidhuber — Neural Computation 9(8) The foundational paper — read this. MIT Press link
Generating Sequences with Recurrent Neural Networks (2013) Alex Graves Character-level generation with LSTMs. arxiv.org/abs/1308.0850
LSTM: A Search Space Odyssey (2015) Jozefowicz, Zaremba, Sutskever Comprehensive empirical study of LSTM variants. arxiv.org/abs/1503.04069

Blog

Understanding LSTM Networks Chris Olah The best visual introduction to LSTMs available. colah.github.io/posts/2015-08-Understanding-LSTMs/
The Unreasonable Effectiveness of Recurrent Neural Networks Andrej Karpathy karpathy.github.io/2015/05/21/rnn-effectiveness/