Sequential Models: LSTMs, GRUs & Recurrent Networks

1. Why Vanilla RNNs Fail on Long Sequences

In a vanilla RNN, the hidden state update formula is h_t = tanh(W_hh * h_(t-1) + W_xh * x_t + b). During backpropagation through time (BPTT), computing the gradient of the loss at step T with respect to step 1 requires multiplying the weight matrix W_hh repeatedly T-1 times:

If the largest eigenvalue of W_hh < 1, the gradient decays exponentially to zero (Vanishing Gradient). If > 1, it explodes uncontrollably to NaN. Thus, vanilla RNNs cannot maintain memories past 10-15 time steps.

2. LSTM Cell Architecture: The Additive Constant Error Carousel

Hochreiter & Schmidhuber solved this with the Long Short-Term Memory (LSTM) cell, which separates the fast-changing hidden state $h_t$ from a slow, persistent Cell State $C_t$ regulated by three sigmoid gates:

3. PyTorch nn.LSTM Tensor Dimensions

By default, PyTorch expects sequential inputs in [Seq_Len, Batch, Feature_Dim] layout. For modern pipelines, setting batch_first=True allows standard[Batch, Seq_Len, Feature_Dim] tensors:

import torch
import torch.nn as nn

# Initializing an LSTM layer
lstm = nn.LSTM(
    input_size=10,       # Number of features per time step
    hidden_size=64,      # Number of features in hidden state h
    num_layers=2,        # Stack 2 LSTM layers on top of each other
    batch_first=True,    # Input tensor shape: [Batch, Seq_Len, Input_Size]
    bidirectional=False, # If True, doubles hidden_size with backward pass
    dropout=0.2          # Inter-layer dropout (applied between layer 1 and 2)
)

# Input: Batch of 16 sequences, each 50 time steps long, with 10 features per step
batch_x = torch.randn(16, 50, 10)

# Forward pass: returns (output_sequence, (h_n, c_n))
output, (h_n, c_n) = lstm(batch_x)

print(f"Output Sequence Shape: {output.shape}")  # [16, 50, 64]
print(f"Final Hidden State Shape: {h_n.shape}")    # [2, 16, 64] (num_layers, batch, hidden)
print(f"Final Cell State Shape:   {c_n.shape}")    # [2, 16, 64]

4. Production Time-Series Forecasting Architecture

Below is a complete, modular LSTM sequence-to-one forecasting architecture that extracts the final temporal state and maps it to future predicted targets:

class LSTMForecaster(nn.Module):
    """
    Modular LSTM network for sequence regression or classification.
    """
    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int, output_dim: int):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=0.2 if num_layers > 1 else 0.0
        )
        
        # Fully connected projection head
        self.head = nn.Sequential(
            nn.Linear(hidden_dim, 32),
            nn.ReLU(),
            nn.Linear(32, output_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: [batch_size, seq_len, input_dim]
        # output shape: [batch_size, seq_len, hidden_dim]
        lstm_out, (h_n, c_n) = self.lstm(x)
        
        # Take the hidden state of the very last time step: [batch_size, hidden_dim]
        last_step_features = lstm_out[:, -1, :]
        
        # Predict target values
        predictions = self.head(last_step_features)
        return predictions

# Verify model execution
model = LSTMForecaster(input_dim=5, hidden_dim=64, num_layers=2, output_dim=1)
mock_seq = torch.randn(32, 24, 5)  # 32 batches, 24 hours history, 5 sensor metrics
forecast = model(mock_seq)
print(f"Forecast shape: {forecast.shape}")  # torch.Size([32, 1])