RNN, LSTM, GRU Networks
1. Recurrent Neural Network Foundations
1.1 Basic RNN Equations
A Recurrent Neural Network (RNN) processes sequential data by maintaining a hidden state that evolves over time:
where:
- : input at time
- : hidden state
- : hidden-to-hidden weight
- : input-to-hidden weight
- : activation function (typically )
1.2 Unrolled Computation
For a sequence of length , the unrolled network computes:
The output at time depends on all inputs through the chain of hidden states.
1.3 Gradient Flow Analysis
Backpropagation Through Time (BPTT):
The gradient of the loss at time with respect to the hidden state at time is:
where
with .
2. Vanishing and Exploding Gradients
2.1 The Vanishing Gradient Problem
For the gradient to vanish, we need:
Using the spectral norm (largest singular value):
For activation, (at ). In practice, is often much less than 1.
Gradient magnitude after steps:
If the maximum singular value , then:
2.2 Exploding Gradients
Conversely, if :
Gradient clipping is the standard mitigation:
where is the clipping threshold.
2.3 Mathematical Analysis of Gradient Flow
The long-term dependencies problem can be quantified using the BPTT length:
For vanilla RNNs, is typically 10-100 timesteps, making it impossible to learn dependencies spanning hundreds or thousands of steps.
3. Long Short-Term Memory (LSTM)
3.1 LSTM Cell Architecture
The LSTM addresses vanishing gradients through a cell state that acts as a "conveyor belt" with carefully controlled gates:
3.2 Gradient Flow in LSTM
The key insight is the additive update of the cell state:
The gradient flows through:
Since (sigmoid output), the gradient magnitude is:
Key property: If the forget gate , then , allowing gradients to flow indefinitely.
3.3 LSTM Gradient Magnitude
For the gradient from time to time :
If for all , the gradient is preserved. The LSTM can learn to set to maintain long-term dependencies.
3.4 LSTM Variants
Peephole Connections:
Coupled Forget-Input Gate:
Projection LSTM:
4. Gated Recurrent Unit (GRU)
4.1 GRU Equations
The GRU simplifies the LSTM by merging the cell state and hidden state:
4.2 GRU vs LSTM Comparison
| Property | LSTM | GRU |
|---|---|---|
| Parameters | ||
| Gates | 3 (forget, input, output) + cell state | 2 (update, reset) |
| Memory | Separate cell state | Merged with hidden state |
| Training speed | ~1.5× slower | ~1× faster |
| Performance | Slightly better on long sequences | Competitive, simpler |
4.3 Gradient Flow in GRU
The update gate controls the trade-off between old and new:
When , the gradient , preserving gradients.
5. Bidirectional RNNs
5.1 Bidirectional Formulation
A bidirectional RNN processes the sequence in both directions:
Forward pass:
Backward pass:
Combined representation:
5.2 Bidirectional Gradient Analysis
The forward pass has gradient flow length while the backward pass has length . The effective gradient path is:
The bidirectional architecture effectively halves the maximum gradient path length.
6. Sequence-to-Sequence (Seq2Seq)
6.1 Encoder-Decoder Architecture
Encoder: Processes input sequence to produce context vector :
Decoder: Generates output sequence :
6.2 Attention Mechanism (Bahdanau et al., 2015)
The attention mechanism allows the decoder to focus on different parts of the input:
where
The context vector is:
6.3 Teacher Forcing
During training, the decoder receives the ground truth as input:
Scheduled Sampling (Bengio et al., 2015) gradually transitions from teacher forcing to free running:
7. Code Examples
7.1 LSTM Cell Implementation
import torch
import torch.nn as nn
class LSTMCell(nn.Module):
"""LSTM cell with optional peephole connections."""
def __init__(self, input_size, hidden_size, use_peepholes=False):
super().__init__()
self.hidden_size = hidden_size
self.use_peepholes = use_peepholes
# Input-to-hidden weights
self.W_i = nn.Linear(input_size, hidden_size)
self.W_f = nn.Linear(input_size, hidden_size)
self.W_c = nn.Linear(input_size, hidden_size)
self.W_o = nn.Linear(input_size, hidden_size)
# Hidden-to-hidden weights
self.U_i = nn.Linear(hidden_size, hidden_size, bias=False)
self.U_f = nn.Linear(hidden_size, hidden_size, bias=False)
self.U_c = nn.Linear(hidden_size, hidden_size, bias=False)
self.U_o = nn.Linear(hidden_size, hidden_size, bias=False)
# Peephole connections
if use_peepholes:
self.P_f = nn.Parameter(torch.randn(hidden_size) * 0.1)
self.P_i = nn.Parameter(torch.randn(hidden_size) * 0.1)
self.P_o = nn.Parameter(torch.randn(hidden_size) * 0.1)
self._init_weights()
def _init_weights(self):
"""Orthogonal initialization for recurrent weights."""
for name, param in self.named_parameters():
if 'U' in name:
nn.init.orthogonal_(param)
elif 'weight' in name:
nn.init.xavier_uniform_(param)
elif 'bias' in name:
nn.init.zeros_(param)
def forward(self, x_t, h_prev, c_prev):
"""
Forward pass for one timestep.
Parameters
----------
x_t : Tensor (batch, input_size)
h_prev : Tensor (batch, hidden_size)
c_prev : Tensor (batch, hidden_size)
Returns
-------
h_t : Tensor (batch, hidden_size)
c_t : Tensor (batch, hidden_size)
"""
# Gates
i_t = torch.sigmoid(self.W_i(x_t) + self.U_i(h_prev))
f_t = torch.sigmoid(self.W_f(x_t) + self.U_f(h_prev))
o_t = torch.sigmoid(self.W_o(x_t) + self.U_o(h_prev))
# Peephole connections
if self.use_peepholes:
f_t = torch.sigmoid(self.W_f(x_t) + self.U_f(h_prev) + self.P_f * c_prev)
i_t = torch.sigmoid(self.W_i(x_t) + self.U_i(h_prev) + self.P_i * c_prev)
# Candidate cell state
c_tilde = torch.tanh(self.W_c(x_t) + self.U_c(h_prev))
# Cell state update
c_t = f_t * c_prev + i_t * c_tilde
# Output gate with peephole
if self.use_peepholes:
o_t = torch.sigmoid(self.W_o(x_t) + self.U_o(h_prev) + self.P_o * c_t)
# Hidden state
h_t = o_t * torch.tanh(c_t)
return h_t, c_t
def compute_gradient_norm(self, h_T, c_T, target):
"""Compute gradient norm to analyse vanishing gradients."""
loss = nn.MSELoss()(h_T, target)
loss.backward()
grad_norms = {}
for name, param in self.named_parameters():
if param.grad is not None:
grad_norms[name] = param.grad.norm().item()
return grad_norms
# Example: Analyse gradient flow
lstm = LSTMCell(input_size=32, hidden_size=64)
x = torch.randn(10, 32) # sequence of 10 timesteps
h, c = torch.zeros(1, 64), torch.zeros(1, 64)
# Forward pass
for t in range(10):
h, c = lstm(x[t:t+1], h, c)
# Compute gradients
target = torch.randn(1, 64)
grad_norms = lstm.compute_gradient_norm(h, c, target)
print("Gradient norms:")
for name, norm in grad_norms.items():
print(f" {name}: {norm:.6f}")
7.2 GRU Implementation
class GRUCell(nn.Module):
"""GRU cell with reset and update gates."""
def __init__(self, input_size, hidden_size):
super().__init__()
self.hidden_size = hidden_size
# Combined weights for efficiency
self.W_z = nn.Linear(input_size + hidden_size, hidden_size)
self.W_r = nn.Linear(input_size + hidden_size, hidden_size)
self.W_n = nn.Linear(input_size + hidden_size, hidden_size)
self._init_weights()
def _init_weights(self):
for name, param in self.named_parameters():
if 'weight' in name:
nn.init.orthogonal_(param)
elif 'bias' in name:
nn.init.zeros_(param)
def forward(self, x_t, h_prev):
"""
Forward pass for one timestep.
Parameters
----------
x_t : Tensor (batch, input_size)
h_prev : Tensor (batch, hidden_size)
Returns
-------
h_t : Tensor (batch, hidden_size)
"""
combined = torch.cat([x_t, h_prev], dim=-1)
# Update gate
z_t = torch.sigmoid(self.W_z(combined))
# Reset gate
r_t = torch.sigmoid(self.W_r(combined))
# Candidate hidden state
combined_r = torch.cat([x_t, r_t * h_prev], dim=-1)
n_t = torch.tanh(self.W_n(combined_r))
# Hidden state update
h_t = (1 - z_t) * h_prev + z_t * n_t
return h_t
class GRUModel(nn.Module):
"""Full GRU model for sequence classification."""
def __init__(self, input_size, hidden_size, num_classes, num_layers=2):
super().__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.cells = nn.ModuleList([
GRUCell(input_size if i == 0 else hidden_size, hidden_size)
for i in range(num_layers)
])
self.classifier = nn.Linear(hidden_size, num_classes)
def forward(self, x):
"""
Forward pass through the full GRU.
Parameters
----------
x : Tensor (batch, seq_len, input_size)
Returns
-------
output : Tensor (batch, num_classes)
"""
batch_size, seq_len, _ = x.shape
# Initialize hidden states
h = [torch.zeros(batch_size, self.hidden_size).to(x.device)
for _ in range(self.num_layers)]
# Process sequence
for t in range(seq_len):
inp = x[:, t, :]
for layer, cell in enumerate(self.cells):
h[layer] = cell(inp, h[layer])
inp = h[layer]
# Use last hidden state for classification
output = self.classifier(h[-1])
return output
7.3 Bidirectional LSTM with Attention
class BiLSTMAttention(nn.Module):
"""Bidirectional LSTM with self-attention."""
def __init__(self, input_size, hidden_size, num_classes):
super().__init__()
self.hidden_size = hidden_size
# Bidirectional LSTM
self.lstm = nn.LSTM(
input_size, hidden_size,
num_layers=2,
batch_first=True,
bidirectional=True
)
# Attention
self.attention = nn.Linear(hidden_size * 2, 1)
# Classification
self.classifier = nn.Linear(hidden_size * 2, num_classes)
def forward(self, x, mask=None):
"""
Forward pass.
Parameters
----------
x : Tensor (batch, seq_len, input_size)
mask : Tensor (batch, seq_len), optional
Returns
-------
output : Tensor (batch, num_classes)
attention_weights : Tensor (batch, seq_len)
"""
# LSTM
lstm_out, _ = self.lstm(x) # (batch, seq_len, hidden*2)
# Attention
attn_scores = self.attention(lstm_out).squeeze(-1) # (batch, seq_len)
if mask is not None:
attn_scores = attn_scores.masked_fill(~mask.bool(), float('-inf'))
attn_weights = torch.softmax(attn_scores, dim=-1)
# Weighted sum
context = torch.bmm(attn_weights.unsqueeze(1), lstm_out).squeeze(1)
# Classification
output = self.classifier(context)
return output, attn_weights
def compute_gradient_flow(self, x):
"""Analyse gradient flow through the network."""
output, attn = self.forward(x)
loss = output.sum()
loss.backward()
# Check gradient norms at different layers
grad_norms = {
'lstm_ih_l0': self.lstm.weight_ih_l0.grad.norm().item(),
'lstm_hh_l0': self.lstm.weight_hh_l0.grad.norm().item(),
'lstm_ih_l1': self.lstm.weight_ih_l1.grad.norm().item(),
'attention': self.attention.weight.grad.norm().item(),
}
return grad_norms
# Example: Test bidirectional LSTM
model = BiLSTMAttention(input_size=64, hidden_size=128, num_classes=10)
x = torch.randn(32, 50, 64) # batch=32, seq_len=50, features=64
output, attn_weights = model(x)
print(f"Output shape: {output.shape}")
print(f"Attention shape: {attn_weights.shape}")
print(f"Attention sum: {attn_weights[0].sum():.4f}") # should be ~1.0
8. Summary
-
Vanishing gradients in vanilla RNNs limit the ability to learn long-range dependencies to ~10-100 timesteps.
-
LSTM solves this through additive cell state updates and forget gates, enabling gradients to flow for 1000+ steps.
-
GRU simplifies LSTM by merging cell and hidden states, achieving comparable performance with fewer parameters.
-
Bidirectional RNNs halve the effective gradient path length by processing in both directions.
-
Seq2seq with attention enables the decoder to focus on relevant parts of the input, breaking the information bottleneck.
References
- Hochreiter, S., & Schmidhuber, J. (1997). Long short-term memory. Neural Computation, 9(8), 1735-1780.
- Cho, K., et al. (2014). Learning phrase representations using RNN encoder-decoder for statistical machine translation. EMNLP.
- Bahdanau, D., Cho, K., & Bengio, Y. (2015). Neural machine translation by jointly learning to align and translate. ICLR.
- Bengio, S., et al. (2015). Scheduled sampling for sequence prediction with recurrent neural networks. NeurIPS.
- Pascanu, R., Mikolov, T., & Bengio, Y. (2013). On the difficulty of training recurrent neural networks. ICML.