RNNs, Transformers, Attention

Sequences in Machine Learning

Many real-world problems involve sequential data where order matters and inputs can be variable-length: text, audio, time series, code, genomics. Unlike tabular data, position relative to other elements is informative.

Challenges: long-range dependencies, variable length, sequential computation (hard to parallelize).

Recurrent Neural Networks (RNNs)

An RNN processes a sequence one element at a time, maintaining a hidden state hₜ that summarizes all past information:

hₜ = f(Wₕhₜ₋₁ + Wₓxₜ + b)

yₜ = g(Wᵧhₜ)

The same weight matrices W_h, W_x, W_y are reused at every step (parameter sharing across time).

import torch.nn as nn

rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=2,
             batch_first=True, dropout=0.3)
# Input: (batch, seq_len, input_size)
# Output: (batch, seq_len, hidden_size), final hidden state
output, h_n = rnn(x)

Vanilla RNN limitations: - Vanishing gradients through time: gradient of loss at step t with respect to step 1 involves T multiplications of W_h — shrinks exponentially - Cannot maintain information over 10+ steps in practice - Sequential computation prevents parallelization

LSTM (Long Short-Term Memory)

LSTM adds a cell state c_t (separate memory track) and gates that control information flow:

Forget gate: fₜ = σ(Wf · [hₜ₋₁, xₜ] + bf) — what to erase from cell Input gate: iₜ = σ(Wᵢ · [hₜ₋₁, xₜ] + bᵢ) — what new info to write Candidate: c̃ₜ = tanh(Wc · [hₜ₋₁, xₜ] + bc) — new candidate values Cell update: cₜ = fₜ ⊙ cₜ₋₁ + iₜ ⊙ c̃ₜ Output gate: oₜ = σ(Wo · [hₜ₋₁, xₜ] + bo) Hidden state: hₜ = oₜ ⊙ tanh(cₜ)

The cell state cₜ flows through the network with only element-wise operations — gradient highway that alleviates vanishing gradients. LSTMs can retain information over 100+ steps.

lstm = nn.LSTM(input_size=128, hidden_size=256, num_layers=2,
               batch_first=True, dropout=0.3, bidirectional=True)
# Bidirectional doubles hidden_size in output
output, (h_n, c_n) = lstm(x)   # h_n, c_n: final hidden and cell states

GRU (Gated Recurrent Unit)

Simplified LSTM with two gates instead of three — fewer parameters, often similar performance:

Reset gate: rₜ = σ(Wr · [hₜ₋₁, xₜ]) Update gate: zₜ = σ(Wz · [hₜ₋₁, xₜ]) Candidate: h̃ₜ = tanh(Wh · [rₜ ⊙ hₜ₋₁, xₜ]) Hidden state: hₜ = (1−zₜ) ⊙ hₜ₋₁ + zₜ ⊙ h̃ₜ

gru = nn.GRU(input_size=128, hidden_size=256, num_layers=2,
             batch_first=True, bidirectional=True)
output, h_n = gru(x)

RNN Architectures for Different Tasks

# Seq2Seq encoder-decoder (simplified)
class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.lstm  = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
    def forward(self, x):
        return self.lstm(self.embed(x))   # returns output, (h_n, c_n)

Attention Mechanism

The attention mechanism lets the decoder selectively focus on different encoder outputs when generating each output token, solving the bottleneck of compressing everything into a single vector.

Bahdanau Attention (additive):

eₜₛ = vᵀ tanh(Wₕhₛ + Wₛsₜ) (alignment score) αₜₛ = exp(eₜₛ) / Σₛ' exp(eₜₛ') (attention weights, sum to 1) cₜ = Σₛ αₜₛ hₛ (context vector for decoder step t)

Luong Attention (multiplicative):

eₜₛ = sₜᵀ Wₐ hₛ

Attention weights αₜₛ tell the decoder how much to attend to encoder position s when generating output position t.

Self-Attention

The key insight of Transformers: attention over the sequence itself rather than encoder-decoder. Each position attends to all other positions to compute a contextualized representation.

Given input matrix X (seq_len × d_model), project to Queries, Keys, Values:

Q = XWQ, K = XWK, V = XWV

Attention(Q, K, V) = softmax(QKᵀ/√d_k) V

• √d_k scaling prevents dot products from growing too large (which saturates softmax)
• The output is a weighted sum of Values, with weights determined by Query-Key similarity

Intuition: for each token (Query), compute how relevant each other token (Key) is, then aggregate those tokens' Values weighted by relevance.

Multi-Head Attention

Run h parallel attention heads, each with different projections:

MultiHead(Q, K, V) = Concat(head₁, …, headₕ) Wₒ

where headᵢ = Attention(QWᵢᵠ, KWᵢᴷ, VWᵢᵛ)

Each head can focus on different types of relationships (syntax, coreference, position, etc.).

Complexity: O(n² · d) per layer — quadratic in sequence length. This limits Transformers on very long sequences.

# PyTorch multi-head attention
attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, dropout=0.1,
                               batch_first=True)
output, attn_weights = attn(query, key, value, key_padding_mask=mask)

The Transformer Architecture

Introduced in "Attention Is All You Need" (Vaswani et al., 2017).

Encoder block:

x → LayerNorm → Multi-Head Self-Attention → + (residual)
  → LayerNorm → Feed-Forward Network (2 linear layers + ReLU/GELU) → + (residual)

Decoder block (autoregressive):

x → LayerNorm → Masked Self-Attention → + (residual)
  → LayerNorm → Cross-Attention (attends to encoder) → + (residual)
  → LayerNorm → Feed-Forward → + (residual)

Masked self-attention: prevents position i from attending to future positions j > i (causal masking), essential for autoregressive generation.

class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads,
                                           dropout=dropout, batch_first=True)
        self.ff   = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(),
            nn.Linear(d_ff, d_model), nn.Dropout(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_key_padding_mask=None):
        # Pre-norm variant
        attn_out, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x),
                                 key_padding_mask=src_key_padding_mask)
        x = x + self.dropout(attn_out)
        x = x + self.dropout(self.ff(self.norm2(x)))
        return x

Positional Encoding

Transformers have no built-in notion of position — all tokens attend in parallel. Positional encoding injects position information:

Sinusoidal (original): PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Properties: each dimension is a sinusoid at a different frequency; relative positions can be inferred from the pattern.

Learned absolute: trainable embeddings, one per position (BERT, GPT).

Rotary Positional Embedding (RoPE): encodes position as rotation of query/key vectors — supports extrapolation to longer sequences. Used in Llama, Mistral, and most modern LLMs.

ALiBi (Attention with Linear Biases): add a position-dependent penalty to attention scores — strong length generalization.

Comparing RNNs and Transformers

Efficient Attention Variants

FlashAttention is the practical standard — same results as full attention but 2–4× faster with much less memory.

Sequence-to-Sequence (seq2seq) Applications