QKV Attention Mechanism in Transformers
Scaled dot-product attention matrices (Softmax), multi-head attention, and contextual representations
Analytical visualizer of the algorithmic core powering modern Large Language Models (LLMs). Explore how Query, Key, and Value vector matrices interact, computing contextual attention weights token by token.
Attention Is All You Need: Mathematical Anatomy of the Transformer
{lang === 'es' ? 'En 2017, Vaswani et al. (Google Brain & Research) publicaron el artículo fundacional «Attention Is All You Need», demoliendo la dependencia hegemónica de la recurrencia secuencial ($O(n)$) y las convoluciones locales. Al reemplazar las redes recurrentes por el mecanismo de Atención Escalar Producto Punto (Scaled Dot-Product Attention) y cabezas múltiples en espacios proyectivos ortogonales, el Transformer redujo la longitud máxima del camino entre cualquier par de tokens a un asombroso $O(1)$.' : 'In 2017, Vaswani et al. (Google Brain & Research) published the groundbreaking paper "Attention Is All You Need", dismantling the hegemony of sequential recurrence ($O(n)$) and local convolutions. By substituting recurrence with Scaled Dot-Product Attention and multi-head orthogonal projections, the Transformer collapsed the maximum path length between any arbitrary token pair down to $O(1)$.'}
Binds pronouns ("it", "he", "she") to their true noun antecedents across long temporal sequence spans.
{lang === 'es' ? 'Si se desactiva, los productos escalares crecen desmesuradamente en dimensiones altas ($d_k=64$), saturando el softmax y anulando los gradientes.' : 'When disabled, dot products explode in higher dimensions ($d_k=64$), saturating the softmax into regions with vanishing gradients.'}
💡 Hover or click any row or cell to inspect the attention probability distribution for that query token.
"it"
Distribution over Key Tokens (kⱼ):
Why Divide Exactly by √dₖ?
{lang === 'es' ? 'Asumiendo que las componentes de los vectores de consulta $q$ y clave $k$ son variables aleatorias independientes con media $\\mathbb{E}[q_i] = 0$ y varianza unitaria $\\text{Var}(q_i) = 1$, el producto escalar cartesiano es la suma de $d_k$ variables independientes:' : 'Assuming the components of query $q$ and key $k$ vectors are independent random variables with zero mean $\\mathbb{E}[q_i] = 0$ and unit variance $\\text{Var}(q_i) = 1$, the dot product is the sum of $d_k$ independent random variables:'}
{lang === 'es' ? 'Por tanto, la desviación estándar es $\\sigma = \\sqrt{d_k}$. Para valores típicos como $d_k = 64$ o $128$, los productos escalares adquieren valores de gran magnitud ($\pm 25$), empujando a la función $\\text{softmax}$ hacia regiones asintóticas donde el gradiente $\\frac{\\partial \\text{softmax}_i}{\\partial z_j} = s_i(\\delta_{ij} - s_j)$ es infinitesimal. Dividir por $\\sqrt{d_k}$ restablece la varianza unitaria $\\text{Var}\\left(\\frac{q \\cdot k}{\\sqrt{d_k}}\\right) = 1$, preservando el flujo de gradiente durante el entrenamiento.' : 'Hence, the standard deviation is $\\sigma = \\sqrt{d_k}$. For typical dimensions like $d_k = 64$ or $128$, dot products reach extreme magnitudes ($\pm 25$), pushing the $\\text{softmax}$ function into saturated plateaus where the Jacobian gradient $\\frac{\\partial \\text{softmax}_i}{\\partial z_j} = s_i(\\delta_{ij} - s_j)$ is virtually zero. Scaling by $\\sqrt{d_k}$ recovers unit variance $\\text{Var}\\left(\\frac{q \\cdot k}{\\sqrt{d_k}}\\right) = 1$, preventing gradient vanishing.'}
Multi-Head Attention: Subspace Projections
{lang === 'es' ? 'En lugar de calcular una sola atención con dimensión $d_{\\text{model}} = 512$, el Transformer proyecta linealmente $Q, K, V$ en $h = 8$ cabezas paralelas con matrices aprendidas $W_i^Q, W_i^K \\in \\mathbb{R}^{d_{\\text{model}} \\times d_k}$ y $W_i^V \\in \\mathbb{R}^{d_{\\text{model}} \\times d_v}$ ($d_k = d_v = 64$):' : 'Rather than computing a single attention function with dimension $d_{\\text{model}} = 512$, the Transformer linearly projects $Q, K, V$ into $h = 8$ parallel heads using learned projections $W_i^Q, W_i^K \\in \\mathbb{R}^{d_{\\text{model}} \\times d_k}$ and $W_i^V \\in \\mathbb{R}^{d_{\\text{model}} \\times d_v}$ ($d_k = d_v = 64$):'}
This allows the network to jointly attend to information from distinct representation subspaces at distinct positions. One head tracks local grammatical modifiers, another captures distant coreference, and another registers subject-verb-object alignments without destructive interference.
Sinusoidal Positional Encoding (Givens Rotations)
{lang === 'es' ? 'El mecanismo de autoatención es equivariante ante permutaciones: no tiene conocimiento intrínseco del orden temporal de las palabras. Para dotar a la secuencia de geometría espacial, se inyectan ondas sinusoidales de frecuencias geométricas decrecientes:' : 'Self-attention is strictly permutation equivariant: it possesses no intrinsic awareness of sequence order. To embed temporal order, sinusoidal wave patterns of geometrically decreasing frequencies are added to input embeddings:'}
{lang === 'es' ? 'La elegancia matemática de esta formulación radica en que para cualquier desplazamiento relativo fijo $k$, existe una transformación ortogonal lineal $R_k \\in \\text{SO}(2)$ tal que $PE_{pos+k} = R_k PE_{pos}$. Esto permite que el modelo aprenda a interpolar y extrapolar distancias relativas fácilmente mediante productos escalares aprendidos.' : 'The mathematical beauty of this scheme is that for any fixed relative offset $k$, there exists a linear orthogonal rotation $R_k \\in \\text{SO}(2)$ such that $PE_{pos+k} = R_k PE_{pos}$. This allows the model to learn relative position offsets purely through learned linear projections.'}
Asymptotic Complexity Comparison (Vaswani et al., Table 1)
Where n is the sequence length, d is the representation dimension, and k is the convolutional kernel size.
| Layer Type | Complexity per Layer | Sequential Operations | Maximum Path Length | Architectural Impact |
|---|---|---|---|---|
| Self-Attention (Transformer) | O(n^2 · d) | O(1) | O(1) | 100% parallelizable on GPU/TPU. Any token pair interacts in a single step with zero vanishing gradient across path length. |
| Recurrente (RNN / LSTM / GRU) | O(n · d^2) | O(n) | O(n) | Sequential unparallelizable bottleneck. Token n strictly depends on step n-1, suffering from vanishing long-term gradients. |
| Convolucional (CNN 1D / Dilatada) | O(k · n · d^2) | O(1) | O(log_k(n)) | Parallelizable but requires stacking multiple dilated layers to connect distant tokens across the sequence. |
| Self-Attention con Ventana Restringida | O(r · n · d) | O(1) | O(n / r) | Variant for ultra-long context windows (Longformer, BigBird). Trades global receptive horizon for linear O(n) scaling. |
# Implementación formal del mecanismo de atención de Vaswani et al. (2017)
import torch
import torch.nn.functional as F
import math
def scaled_dot_product_attention(Q, K, V, mask=None, tau=1.0):
"""
Q, K: (Batch, Heads, Seq_Len, d_k)
V: (Batch, Heads, Seq_Len, d_v)
"""
d_k = Q.size(-1)
# 1. Producto escalar cartesiano Q K^T escalado por sqrt(d_k)
scores = torch.matmul(Q, K.transpose(-2, -1)) / (tau * math.sqrt(d_k))
# 2. Máscara autoregresiva causal (opcional para decodificadores)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# 3. Normalización probabilística Softmax sobre la dimensión de claves
attn_weights = F.softmax(scores, dim=-1)
# 4. Agregación ponderada de los vectores de valor V
output = torch.matmul(attn_weights, V)
return output, attn_weights