Long Short-Term Memory (LSTM)

LSTMs are a specialised type of RNN designed to reduce the vanishing gradient / short-term memory problem. They introduce a separate cell state which is like a "memory highway" that can carry information across many timesteps with minimal degradation.

---

Why use LSTM for longer dependencies?

In a standard RNN, the hidden state at step $t$ is:

$$h_t=\tanh (W_{xh}{x}_t+W_{hh}{h}_{t-1}+b_h)$$

RNN Mechanism	Problem
Words --> vectors	Fine
Process one-by-one, passing hidden state	Hidden state gets overwritten each step
tanh squeezes values to $(-1,1)$	Repeated squashing causes vanishing gradients
Early inputs far from the output	The network "forgets" them

LSTMs provide an additive path through a separate cell state $C_t$. Gates control how much earlier information remains on this path.

---

LSTM Architecture

Main concept is that the LSTM cell has two parallel flows:

• Cell state $C_t$: the long-term memory highway (modified by gates, not a full overwrite)
• Hidden state ($h_t$): the short-term output passed to the next timestep and to the prediction head

---

Cell State

The cell state acts as a transport highway through the sequence:

C_{t-1}  →  [Forget Gate]  →  [Input Gate adds new info]  →  C_t

Information can be added or removed from the cell state at each step. A forget gate near one retains earlier information. A forget gate near zero can remove almost all of it. The cell state therefore supports long-term memory when the learned gates retain useful information.

---

Gates

LSTM uses three gates, each implemented as a small neural network with a sigmoid activation:

$$\text{gate output}=\sigma (W_x{x}_t+W_h{h}_{t-1}+b)\in (0,1)$$

• Output near 0 means gate is closed (block information)
• Output near 1 means gate is open (allow information through)

---

Gate 1: Forget Gate

"What should we erase from the cell state?"

$$f_t=\sigma (W_{xf}{x}_t+W_{hf}{h}_{t-1}+b_f)$$

1. Takes previous hidden state ($h_{t-1}$) + current input ($x_t$)
2. Outputs a value in $(0,1)$ per cell-state dimension
3. Multiplied element-wise with ($C_{t-1}$) where values near $0$ are erased, while values near $1$ are kept

---

Gate 2: Input Gate

"What new information should we write to the cell state?"

Two parallel operations:

$$i_t=\sigma (W_{xi}{x}_t+W_{hi}{h}_{t-1}+b_i)\text{(what to update)}$$ $${\tilde{C}}_t=\tanh (W_{xc}{x}_t+W_{hc}{h}_{t-1}+b_c)\text{(candidate values)}$$

Combined cell state update:

$$C_t=f_t\odot C_{t-1}+i_t\odot {\tilde{C}}_t$$

• $f_t\odot C_{t-1}$: old memory, partially forgotten
• $i_t\odot {\tilde{C}}_t$: new information, selectively written

---

Gate 3: Output Gate

"What part of the cell state should we expose as the hidden state?"

$$o_t=\sigma (W_{xo}{x}_t+W_{ho}{h}_{t-1}+b_o)$$ $$h_t=o_t\odot \tanh (C_t)$$

1. $tanh(C_t)$ squashes the cell state to $(-1,1)$
2. The output gate $o_t$ selectively exposes parts of it as the new hidden state
3. Both $C_t$ and $h_t$ are passed to the next timestep

---

Final Prediction, y

$$y_t=W_{hy}{h}_t+b_y$$

This is the same as RNN.

Summary

Gate	Controls	Key operation
Forget $f_t$	What to erase from $C_{t-1}$	$f_t\odot C_{t-1}$
Input $i_t$	What new info to write to $C_t$	$i_t\odot {\tilde{C}}_t$
Output $o_t$	What portion of $C_t$ becomes $h_t$	$o_t\odot \tanh (C_t)$

Candidate state and updated state

The candidate ${\tilde{C}}_t$ proposes new content. The forget gate and input gate combine earlier memory with that candidate to produce the updated cell state $C_t$. The output gate then controls how much of $\tanh (C_t)$ becomes the hidden state.

---

Code

PyTorch: two returned states

import torch
from torch import nn
 
ids = torch.tensor([[1, 2, 3], [3, 2, 1]], dtype=torch.long)
embedding = nn.Embedding(10, 8)
lstm = nn.LSTM(8, 16, batch_first=True)
head = nn.Linear(16, 2)
 
output, (h_n, c_n) = lstm(embedding(ids))
logits = head(h_n[-1])
print(tuple(output.shape))  # (2, 3, 16)
print(tuple(h_n.shape))     # (1, 2, 16)
print(tuple(c_n.shape))     # (1, 2, 16)
print(tuple(logits.shape))  # (2, 2)

Read the returned values:

• output contains the hidden state at each of the three positions
• h_n contains the final hidden state for each layer and direction
• c_n contains the final cell state, which is separate from the hidden state
• h_n[-1] selects the single layer's final hidden state for both sequences
• head maps each 16-value state to two class scores

For token-level predictions, use head(output) to obtain shape (2, 3, 2). batch_first changes the input/output layout; the layer axis remains first in h_n and c_n. PyTorch LSTM reference.

TensorFlow/Keras: select sequences and states

import tensorflow as tf
 
ids = tf.constant([[1, 2, 3], [3, 2, 1]])
embedded = tf.keras.layers.Embedding(10, 8)(ids)
lstm = tf.keras.layers.LSTM(16, return_sequences=True, return_state=True)
output, h_n, c_n = lstm(embedded)
logits = tf.keras.layers.Dense(2)(h_n)
print(tuple(output.shape))  # (2, 3, 16)
print(tuple(h_n.shape))     # (2, 16)
print(tuple(c_n.shape))     # (2, 16)
print(tuple(logits.shape))  # (2, 2)

return_sequences=True keeps every position. return_state=True returns the two final states as additional values. Keras returns the final states directly as (batch, units) for this layer.

In a Sequential model, replace SimpleRNN(16) with LSTM(16). Its default output is one final hidden state per sequence. The gate calculations are performed inside the layer.