Attention Mechanisms Explained

Attention mechanisms allow neural networks to dynamically focus on relevant parts of an input sequence.

Scaled Dot-Product Attention

import torch
import torch.nn.functional as F
 
def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

Multi-Head Attention

Transformers use multiple attention heads in parallel, each learning different relationship types.

Key Takeaways

1. Attention is permutation-equivariant without positional encoding
2. Self-attention has O(n²) complexity in sequence length
3. Flash Attention reduces memory bandwidth bottlenecks significantly