Decoding 'Attention Is All You Need': Scaled Dot-Product & Multi-Head Mechanisms

1. The Core Motivation

Prior to Transformers, dominant sequence transduction models relied on complex recurrent or convolutional neural networks (LSTMs, GRUs). While effective, sequential computation creates a fundamental bottleneck:

Sequential Bottleneck: Computing hidden state $h_t$ strictly requires $h_{t-1}$. This precludes parallelization within training examples, becoming critical at longer sequence lengths $N$.

The Transformer discards recurrence entirely, drawing connections between all positions in constant $O(1)$ sequential operations.

2. Scaled Dot-Product Attention

The model maps a query $Q$ and a set of key-value pairs $(K, V)$ to an output vector. When dimensions are large, the dot product $QK^T$ grows large in magnitude, pushing the softmax function into regions with tiny gradients.

To counteract this, the inner product is scaled by the inverse square root of key dimension $d_k$:

PyTorch Implementation

Here is how scaled dot-product attention is implemented in code:

import torch
import torch.nn as nn
import math
 
def scaled_dot_product_attention(query, key, value, mask=None):
    d_k = query.size(-1)
    
    # Compute raw attention scores: (batch, heads, seq_len, seq_len)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
        
    # Softmax along key dimension
    p_attn = torch.softmax(scores, dim=-1)
    
    # Weight values by attention probabilities
    return torch.matmul(p_attn, value), p_attn

3. Multi-Head Attention

Instead of performing a single attention function with $d_{\text{model}}$-dimensional keys, queries, and values, Multi-Head Attention projects $Q, K, V$ linearly $h$ times with different learned linear projections to $d_k, d_k, d_v$ dimensions:

where $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$.

This allows the model to jointly attend to information from different representation subspaces at different positions.