Recurrent Neural Networks, LSTMs & Sequence Modeling

1. Recurrent Neural Network (RNN) Formalism

Recurrent Neural Networks model sequential dependencies by maintaining an internal hidden state memory across time steps $t \in \{1, 2, \dots, T\}$. Given an input sequence $\mathbf{x}_t \in \mathbb{R}^{d_x}$, the hidden state $\mathbf{h}_t \in \mathbb{R}^{d_h}$ updates recurrently:

            Recurrent Unrolling Across Time Steps
            
                h_0               h_1               h_2               h_T
                 │                 │                 │                 │
                 ▼                 ▼                 ▼                 ▼
             ┌───────┐         ┌───────┐         ┌───────┐         ┌───────┐
  x_1 ──────►│  RNN  │───h_1──►│  RNN  │───h_2──►│  RNN  │──...───►│  RNN  │──► y_T
             └───────┘         └───────┘         └───────┘         └───────┘
                 │                 │                 │                 │
                 ▼                 ▼                 ▼                 ▼
                y_1               y_2               y_3               y_T

2. Backpropagation Through Time (BPTT) & Vanishing Gradient Proof

Let $\mathcal{L} = \sum_{t=1}^T \mathcal{L}_t$ be the total sequence loss. Computing the gradient of $\mathcal{L}_T$ with respect to the recurrent weight matrix $\mathbf{W}_{hh}$ requires unrolling the chain rule across all preceding temporal states $k \le T$:

The temporal Jacobian product $\frac{\partial \mathbf{h}_T}{\partial \mathbf{h}_k}$ expands into a sequential matrix product:

Let $\gamma = \|\mathbf{W}_{hh}\|_2$ be the spectral norm of the recurrent transition matrix. Since $|1 - \tanh^2(z)| \le 1$:

• If $\gamma < 1$, the gradient decays exponentially with temporal distance $T-k$, preventing vanilla RNNs from capturing long-range dependencies ($T > 20$).
• If $\gamma > 1$, gradients explode exponentially, requiring gradient clipping.

3. Long Short-Term Memory (LSTM) Architecture

To preserve long-range gradient propagation, Long Short-Term Memory (LSTM) networks introduce an additive linear Cell State $\mathbf{c}_t \in \mathbb{R}^{d_h}$ regulated by three multiplicative gating mechanisms.

       LSTM Cell Internal Gating Architecture
       
   c_{t-1} ───────────────────(×)────────────────────(+)──────────────► c_t
                               ▲                      ▲
                               │ f_t                  │ i_t * g_t
                             ┌───┐                  ┌───┐
                             │ σ │                  │ σ │
                             └───┘                  └───┘
                               ▲                      ▲
                               │                      │
             h_{t-1}, x_t ─────┴──────────────────────┴───────► o_t ──► h_t

Exact LSTM Gating Equations

Given concatenated input $\mathbf{v}_t = [\mathbf{h}_{t-1}, \mathbf{x}_t]^T$:

Constant Error Carousel (CEC)

Taking the derivative of the cell state $\mathbf{c}_t$ with respect to $\mathbf{c}_{t-1}$:

When the forget gate $\mathbf{f}_t \approx \mathbf{1}$, the Jacobian $\frac{\partial \mathbf{c}_t}{\partial \mathbf{c}_{t-1}} = \mathbf{I}$, guaranteeing zero gradient attenuation along the cell state highway across hundreds of time steps.

4. Gated Recurrent Units (GRU) Formalism

GRUs streamline LSTMs by merging the cell state and hidden state into a single representation $\mathbf{h}_t$ regulated by Reset ($\mathbf{r}_t$) and Update ($\mathbf{z}_t$) gates:

5. PyTorch Implementation: Attention-Augmented BiLSTM Sequence Classifier

Below is a production-grade canonical PyTorch implementation of a Bidirectional LSTM with Additive Attention pooling suitable for general sequence classification and sequential embedding representation.

import torch
import torch.nn as nn
import torch.nn.functional as F

class TemporalAdditiveAttention(nn.Module):
    """Additive attention mechanism pooling sequence hidden states into a single context vector."""
    def __init__(self, hidden_dim: int):
        super().__init__()
        self.attn_proj = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.Tanh(),
            nn.Linear(hidden_dim // 2, 1, bias=False)
        )
        
    def forward(self, lstm_outputs: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        # lstm_outputs shape: (Batch, SequenceLength, HiddenDim)
        attn_scores = self.attn_proj(lstm_outputs).squeeze(-1)  # (Batch, SeqLen)
        
        if mask is not None:
            attn_scores = attn_scores.masked_fill(~mask, -1e9)
            
        attn_weights = F.softmax(attn_scores, dim=-1).unsqueeze(-1)  # (Batch, SeqLen, 1)
        context_vector = torch.sum(lstm_outputs * attn_weights, dim=1)
        return context_vector, attn_weights

class CanonicalBiLSTMClassifier(nn.Module):
    """Bidirectional LSTM Sequence Classifier with Attention Pooling."""
    def __init__(self, vocab_size: int = 20000, embed_dim: int = 128, hidden_dim: int = 128,
                 num_classes: int = 5, num_layers: int = 2, dropout: float = 0.2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.bilstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0.0
        )
        # Bidirectional doubles output feature dimension
        bilstm_out_dim = hidden_dim * 2
        self.attention = TemporalAdditiveAttention(bilstm_out_dim)
        self.classifier = nn.Sequential(
            nn.Linear(bilstm_out_dim, 64),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(64, num_classes)
        )
        
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        # input_ids shape: (Batch, SequenceLength)
        mask = (input_ids != 0)
        embedded = self.embedding(input_ids)
        lstm_out, _ = self.bilstm(embedded)
        
        context_vec, _ = self.attention(lstm_out, mask=mask)
        logits = self.classifier(context_vec)
        return logits

# Verification: Forward pass on benchmark sequence token indices
if __name__ == "__main__":
    synthetic_sequences = torch.randint(1, 1000, (16, 50))  # Batch=16, SeqLen=50
    model = CanonicalBiLSTMClassifier(vocab_size=1000, num_classes=5)
    
    logits = model(synthetic_sequences)
    print("Input Sequence Batch Tensor Shape:", synthetic_sequences.shape)
    print("Output Class Logits Matrix Shape:", logits.shape)