Attention Mechanism

The attention mechanism allows a model to dynamically focus on relevant parts of its input when producing each element of its output. At its core, scaled dot-product attention computes a weighted sum of values, where the weights are derived from a compatibility function between queries and keys.

Motivation

Before attention, sequence models (RNNs, LSTMs) processed tokens sequentially, creating a bottleneck for long-range dependencies. Attention removes this sequential constraint: every output position can directly attend to every input position, enabling parallel computation and better gradient flow.

import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)      # (seq_len, seq_len)
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)
    weights = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
    return weights @ V

# Demo: 3 tokens, 4-dim embeddings
Q = K = V = np.random.randn(3, 4)
out = scaled_dot_product_attention(Q, K, V)
print(f"Output shape: {out.shape}")  # (3, 4)

Multi-Head Attention

Instead of a single attention function, multi-head attention runs several attention operations in parallel on different learned linear projections of Q, K, and V. Each head can attend to different aspects of the input (syntax, semantics, positional relationships), and the results are concatenated and projected:

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O

where \text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V).