5.3: Attention Mechanism/Transformers

Core Idea: Dynamically focus on relevant parts of input sequences when generating output

Motivation: Overcome limitations of fixed-length context vectors in seq2seq models (e.g., information bottleneck)

Key Components:

• Query (Q): Represents the current decoding step (e.g., target word)

• Key (K): Represents input words to be matched against the query

• Value (V): Represents the actual content of input words

Mathematical Formulation:

• Attention Score: Computed as compatibility between query and keys
  $$\text{score}(Q, K_i) = \frac{Q \cdot K_i^T}{\sqrt{d_k}}$$
  where d_k = dimension of keys (scaling factor for stable gradients)

• Attention Weights: Softmax over scores
  $$\alpha_i = \frac{\exp(\text{score}(Q, K_i))}{\sum_j \exp(\text{score}(Q, K_j))}$$

• Context Vector: Weighted sum of values $$ \text{context} = \sum_i \alpha_i V_i $$

Definition: Attention mechanism where queries, keys, and values come from the same sequence

Purpose: Model intra-sequence dependencies (e.g., word relationships within a sentence)

Mathematical Steps:

• Input: Sequence X = [x₁, x₂, ..., x_n] (each x_i ∈ ℝ^d_model)

• Linear Projections: $$ Q = X W^Q, \quad K = X W^K, \quad V = X W^V $$
  where W^Q, W^K, W^V ∈ ℝ^{d_model × d_k} (learnable weights)

• Scaled Dot-Product Attention:
  $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

• Output: Sequence of context vectors [z₁, z₂, ..., z_n], where each z_i ∈ ℝ^d_v

Concept: Run multiple self-attention mechanisms in parallel and concatenate outputs

Advantage: Capture diverse relationships (e.g., syntactic, semantic) from different representation subspaces

Mathematical Formulation:

• Projections for h Heads: $$  Q_i = X W_i^Q, \quad K_i = X W_i^K, \quad V_i = X W_i^V \quad (i = 1, ..., h) $$
  where W_i^Q, W_i^K, W_i^V ∈ ℝ^{d_model × d_k} (and d_k = d_model/h)

• Head Output: $$ \text{head}_i = \text{Attention}(Q_i, K_i, V_i) $$

• Concatenation & Final Projection: $$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h) W^O $$
  where W^O ∈ ℝ^{h ⋅ d_v × d_model} (output projection matrix)

Problem: Self-attention is permutation-invariant (lacks sequence order information)

Solution: Add positional embeddings to input embeddings

Formulation:

• Sinusoidal Encoding (fixed):
  $$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$$
  
  $$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$$
  where pos = position index, i = dimension index

• Learned Embeddings: Alternative trainable position vectors

Input to Transformer: $$  X_{\text{input}} = \text{Embedding}(X) + \text{PositionalEncoding} $$

Overall Structure: Encoder-decoder framework with stacked self-attention and feed-forward layers

Encoder:

• Layers: N identical blocks (e.g., N = 6)

• Sublayers:

  1. Multi-Head Self-Attention: $$  Z = \text{MultiHead}(X, X, X) $$

  2. Add & Norm: Residual connection + layer normalization: $$ X_{\text{sub}} = \text{LayerNorm}(X + Z) $$

  3. Position-wise Feed-Forward Network (FFN): $$ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 $$
     (applied independently to each position)

  4. Add & Norm: $$  X_{\text{out}} = \text{LayerNorm}(X_{\text{sub}} + \text{FFN}(X_{\text{sub}})) $$

Decoder:

• Layers: N identical blocks with two self-attention sublayers and one cross-attention

• Sublayers:

  1. Masked Multi-Head Self-Attention: Prevents attending to future tokens (autoregressive decoding)

     • Masking: Set scores of future positions to  − ∞ before softmax

  2. Add & Norm

  3. Encoder-Decoder Attention (Cross-Attention):

     • Query: From decoder output

     • Key/Value: From encoder output $$  Z = \text{MultiHead}(Q_{\text{dec}}, K_{\text{enc}}, V_{\text{enc}})  $$

  4. Add & Norm

  5. FFN + Add & Norm

Purpose: Stabilize training by normalizing activations

Formula:
$$\text{LayerNorm}(x) = \gamma \left(\frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}\right) + \beta$$
where μ = mean, σ² = variance over features, γ, β = learnable parameters

Loss Function: Cross-entropy loss for token prediction

Key Techniques:

• Label Smoothing: Prevents overconfidence (e.g., replace 1-hot targets with 1 − ϵ for correct token, ϵ/(V − 1) for others)

• Learning Rate Schedule: Warmup + decay,  e.g., $$ \text{lr} = d_{\text{model}}^{-0.5} \cdot \min(\text{step}^{-0.5}, \text{step} \cdot \text{warmup}^{-1.5}) $$

• Regularization: Dropout applied to embeddings, attention weights, and FFN outputs

Parallelization: Self-attention enables \( O(1) \) sequential dependency (vs. \( O(n) \) in RNNs)

Long-Range Dependencies: Direct path between any two positions (vs. vanishing gradients in RNNs)

Interpretability: Attention weights visualize alignment (e.g., word-to-word relationships)

BERT: Bidirectional encoder (masked language modeling)

GPT: Autoregressive decoder-only (next-token prediction)

T5/BART: Encoder-decoder fusion (sequence-to-sequence transformation for translation/summarization)

Mixture of Experts (MoE): Sparse routing decoder (conditional activation of specialist sub-networks per token)

Hybrid Mamba-Transformer: State-space linear decoder (interleaved attention and recurrent blocks for efficient long contexts)

Large Reasoning Models (LRM): Test-time compute decoder (reinforcement learning-tuned generation of internal chain-of-thought traces)

Efficient Transformers:

• Sparse Attention (e.g., Longformer, BigBird)

• Linformer: Low-rank projection of keys/values

• Performer: Kernelized attention

Self-Attention: $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V  $$

Multi-Head Attention: $$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h) W^O  $$

FFN: $$ \text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2 $$

Positional Encoding:
$$PE_{(pos, i)} = \begin{cases} \sin(pos / 10000^{2i/d_{\text{model}}}) & \text{if } i \text{ even} \\ \cos(pos / 10000^{2i/d_{\text{model}}}) & \text{if } i \text{ odd} \end{cases}$$