Advanced35 Minutes

Scaled Dot-Product & Multi-Head Self-Attention

Self-attention allows every token in an input sequence to dynamically route information from every other token in the sequence based on relevance, breaking the sequential bottlenecks of RNNs.

Prerequisites

Linear algebraMatrix multiplication

Learning Objectives

•Derive the Scaled Dot-Product Attention formula from first principles.
•Understand Query, Key, and Value linear projections.
•Explain why dividing by sqrt(d_k) prevents vanishing gradients in softmax.

In-Depth Architecture & Mechanics

Mathematical Formulation

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

Where: - $Q = X W_Q$ (Queries: What am I looking for?) - $K = X W_K$ (Keys: What content do I offer?) - $V = X W_V$ (Values: If matched, what information do I pass along?)

Why Scale by $\sqrt{d_k}$? For large projection dimensions $d_k$, the dot product $Q K^T$ grows large in magnitude. When input magnitudes to the softmax function are large, the softmax output pushes towards 0 or 1, causing extremely small gradients (vanishing gradient problem). Scaling by $\frac{1}{\sqrt{d_k}}$ preserves unit variance when keys and queries have zero mean and unit variance.

System Architecture Flow:

Tokens ──> [ Linear Projections: W_Q, W_K, W_V ]
                  │        │        │
                  ▼        ▼        │
                  Q        K        │
                  │        │        │
                  └──MatMul┘        │
                      │             │
                 Scale & Mask       │
                      │             │
                   Softmax          │
                      │             │
                      └───MatMul────┘
                             │
                      Output Vectors

Worked code example:

Scaled Dot-Product Attention in PyTorch

python

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    # Compute attention scores: (Batch, Heads, Seq_Q, Seq_K)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

Architectural Tradeoffs

Quadratic Complexity: Standard self-attention scales as O(N^2) with sequence length N. Techniques like FlashAttention and sliding-window attention alleviate memory bandwidth constraints.

Common Anti-Patterns

Forgetting causal masking during autoregressive decoder training, allowing tokens to cheat by looking at future tokens.

Staff & Principal Interview Questions

Confidence: