Self-Attention — Architecture Concept

In One Sentence

Self-attention allows every token in a sequence to dynamically query and aggregate information from all other tokens, bypassing sequential bottlenecks to capture long-range dependencies in parallel.

Intuition

Imagine reading the sentence: "The bank of the river was muddy."

When your eyes land on the word "bank", how do you know if it refers to a financial institution or the side of a river? You instantly look at the surrounding context—specifically the word "river".

Self-attention formalizes this process mathematically using three roles for every token:

Query (Q): "What am I looking for?" (e.g., *I am the word "bank", looking for words that clarify my meaning*).
Key (K): "What information do I offer?" (e.g., *I am the word "river", offering geographical context*).
Value (V): "What is my actual content?" (e.g., *The actual semantic representation of "river" to be mixed into "bank"*).

By multiplying Queries with Keys, the model computes an **attention map** (a matrix of compatibility scores). It scales these scores, applies a softmax to convert them into probabilities, and uses them to compute a weighted sum of the Values.

Formal Definition

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

Where:

$Q, K, V$ are matrices of shape $(N, d_k)$ (or $d_v$), representing the Queries, Keys, and Values projected from input sequence $X \in \mathbb{R}^{N \times d_{model}}$ via learned weight matrices $W_Q, W_K, W_V$.
$QK^T$ computes the raw pairwise similarity scores between all tokens.
$\sqrt{d_k}$ is the **scaling factor** (dimension of the keys). It prevents the dot products from growing extremely large in high dimensions, which would push the softmax function into regions with dangerously small gradients.
$\text{softmax}$ normalizes each row so that the attention weights sum to 1.

Multi-Head Attention (MHA)

To allow the model to jointly attend to information from different representation subspaces at different positions, Multi-Head Attention projects the queries, keys, and values $h$ times with different, learned linear projections:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$$ $$\text{where } \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Interactive Walkthrough: Scaling Factor ($\sqrt{d_k}$)

Select a query token below to see its attention distribution over the sentence. Adjust the Scaling Factor ($\sqrt{d_k}$) to see how it prevents the softmax from collapsing into a hard one-hot distribution (which causes vanishing gradients during training).

1. Select Query Token:

Scaling Factor ($\sqrt{d_k}$) 8

Lower scaling factors amplify raw score differences (sharper distribution); higher scaling factors smooth them out.

2. Resulting Attention Weights ($\text{softmax}$):

When to use & When not to

Use When

You need to model complex, long-range dependencies in sequential data (text, audio, time-series).
You want to maximize parallelization during training (unlike RNNs, which must be processed sequentially).
You require dynamic, context-dependent representations of inputs.

Avoid When

You are working with extremely long sequences (e.g., >10k tokens) where the $O(N^2)$ memory and computational footprint becomes prohibitive.
You are deploying on hardware with strict memory bandwidth limitations (unless using optimized kernels like FlashAttention).

History & Evolution

June 2017

The Birth of the Transformer

Vaswani et al. publish "Attention Is All You Need", proposing a new simple network architecture based solely on attention mechanisms, dispensing with recurrence and convolutions entirely.

April 2021

Rotary Position Embeddings (RoPE)

Su et al. introduce RoFormer, encoding absolute positions with a rotation matrix to naturally incorporate relative position dependency directly into the self-attention formulation.

May 2022

FlashAttention: IO-Aware Acceleration

Dao et al. propose FlashAttention, an exact attention algorithm that uses tiling to minimize memory reads/writes between GPU HBM and SRAM, dramatically speeding up the quadratic bottleneck.

Pseudocode

scaled_dot_product_attention.py PyTorch

import torch
import torch.nn as nn
import math

def self_attention(Q, K, V, mask=None):
    """
    Computes Scaled Dot-Product Attention.
    Q, K, V shape: [batch_size, num_heads, seq_len, d_k]
    """
    d_k = Q.size(-1)
    
    # 1. Compute similarity scores (QK^T)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. Apply optional causal/padding mask
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
        
    # 3. Softmax to obtain attention weights
    attention_weights = torch.softmax(scores, dim=-1)
    
    # 4. Weighted sum of values
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights

Common Variants

Causal Attention

Masks out future tokens (setting scores to $-\infty$ before softmax) to prevent information leakage in autoregressive language generation.

Rotary Embeddings (RoPE)

Rotates Query and Key vectors in the complex plane to inject relative distance information directly into the dot product.

FlashAttention

An IO-aware exact attention implementation that uses tiling to avoid materializing the massive $N \times N$ attention matrix in slow GPU memory.

Representative Flash Papers

yudgctsnk7

Attention Is All You Need — Explaining the Transformer

The foundational paper establishing self-attention as a complete replacement for recurrent and convolutional layers.

66asdqe6cn

BERT Explained: Deep Bidirectional Transformers

Demonstrates how bidirectional self-attention can be pre-trained using masked language modeling to achieve state-of-the-art representation learning.