# Mecanismos de Atención en Transformers: Self-Attention, Cross-Attention y Multi-Head
Los mecanismos de atención permiten que los modelos de IA ponderen dinámicamente la importancia de los elementos en una secuencia de entrada. Para cada token, se calculan pesos mediante matrices Query (consulta), Key (clave) y Value (valor), generando representaciones contextualizadas. El self-attention analiza conexiones dentro de una misma secuencia, el cross-attention une el codificador y el decodificador, y el multi-head attention paraleliza el proceso para capturar distintos aspectos de los datos.
Teoría del Self-Attention
La secuencia de entrada se representa mediante la matriz $X \in \mathbb{R}^{n \times d}$, donde $n$ es el número de tokens y $d$ la dimensión de los embeddings. De $X$ proyectamos las matrices $Q = XW^Q$, $K = XW^K$, $V = XW^V$ con pesos $W^{Q,K,V} \in \mathbb{R}^{d \times d_k}$.
La matriz de puntuaciones de atención $A_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}}$ mide la relevancia del token $j$ para el $i$. La escala por $\sqrt{d_k}$ estabiliza los gradientes, evitando que desaparezcan en la softmax.
La matriz normalizada $N = \mathrm{softmax}(A)$ por filas da una distribución probabilística. Salida: $\mathrm{Attention}(Q, K, V) = N V$.
Ejemplo de Self-Attention
Considera la frase "Karina va a la tienda", tokenizada como ["karina", "va", "a", "la", "tienda"] con $n=5$. Embeddings (simplificados):
| Token | Embedding |
|--------|-----------|
| karina | [1, 0] |
| va | [0, 1] |
| a | [1, 1] |
| la | [1, 1] |
| tienda | [0.5, 1] |
Matriz $X_{5\times2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$. Por simplicidad, $W^Q = W^K = W^V = I$, por lo que $Q=K=V=X$.
$QK^T = \begin{bmatrix} 1 & 0 & 1 & 1 & 0.5 \\ 0 & 1 & 1 & 1 & 1 \\ 1 & 1 & 2 & 2 & 1.5 \\ 1 & 1 & 2 & 2 & 1.5 \\ 0.5 & 1 & 1.5 & 1.5 & 1.25 \end{bmatrix}$, dividida por $\sqrt{2} \approx 1.41$:
$A \approx \begin{bmatrix} 0.71 & 0 & 0.71 & 0.71 & 0.35 \\ 0 & 0.71 & 0.71 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.35 & 0.71 & 1.06 & 1.06 & 0.89 \end{bmatrix}$.
Tras softmax por fila: $N \approx \begin{bmatrix} 0.25 & 0.15 & 0.25 & 0.25 & 0.10 \\ 0.17 & 0.23 & 0.23 & 0.23 & 0.23 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.14 & 0.23 & 0.28 & 0.28 & 0.07 \end{bmatrix}$.
Salida $N \times V = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$. Los nuevos embeddings incorporan contexto: el vector de "va" refleja movimiento humano, no solo un tiempo abstracto.
import torch
import torch.nn as nn
class SelfAttention(nn.Module):
def __init__(self, embed_dim):
super().__init__()
self.embed_dim = embed_dim
self.q_linear = nn.Linear(embed_dim, embed_dim)
self.k_linear = nn.Linear(embed_dim, embed_dim)
self.v_linear = nn.Linear(embed_dim, embed_dim)
self.scale = embed_dim ** 0.5
def forward(self, x):
Q = self.q_linear(x)
K = self.k_linear(x)
V = self.v_linear(x)
scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
attn = torch.softmax(scores, dim=-1)
out = torch.bmm(attn, V)
return outCross-Attention: Uniendo Secuencias
En cross-attention, la Query proviene del decodificador ($Q = X_{dec}W^Q$), mientras que Key y Value del codificador ($K = X_{enc}W^K$, $V = X_{enc}W^V$). El cálculo es idéntico: $\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$.
Diferencias clave:
| Componente | Self-Attention | Cross-Attention |
|------------|----------------|------------------|
| Q | X | X_dec |
| K | X | X_enc |
| V | X | X_enc |
Esto permite que el decodificador se centre en partes relevantes del contexto del codificador.
import torch
import torch.nn as nn
class CrossAttention(nn.Module):
def __init__(self, embed_dim):
super().__init__()
self.embed_dim = embed_dim
self.q_linear = nn.Linear(embed_dim, embed_dim)
self.k_linear = nn.Linear(embed_dim, embed_dim)
self.v_linear = nn.Linear(embed_dim, embed_dim)
self.scale = embed_dim ** 0.5
def forward(self, x_dec, x_enc):
Q = self.q_linear(x_dec)
K = self.k_linear(x_enc)
V = self.v_linear(x_enc)
scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
attn = torch.softmax(scores, dim=-1)
out = torch.bmm(attn, V)
return outMulti-Head Attention: Perspectivas Paralelas
El multi-head attention ejecuta $h$ cálculos de atención independientes con proyecciones $Q^{(i)} = X W_Q^{(i)}$, $K^{(i)} = X W_K^{(i)}$, $V^{(i)} = X W_V^{(i)}$. Cada cabeza: $\mathrm{head}_i = \mathrm{Attention}(Q^{(i)}, K^{(i)}, V^{(i)})$.
Los resultados se concatenan y proyectan: $\mathrm{MultiHead}(X) = \mathrm{Concat}(head_1, \dots, head_h) W^O$. Tamaño de cabeza $d_k = d / h$, donde $d$ es embed_dim.
import torch
import torch.nn as nn
class MultiHeadSelfAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
assert embed_dim % num_heads == 0
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.q = nn.Linear(embed_dim, embed_dim)
self.k = nn.Linear(embed_dim, embed_dim)
self.v = nn.Linear(embed_dim, embed_dim)
self.out = nn.Linear(embed_dim, embed_dim)
def forward(self, x):
B, T, D = x.shape
Q = self.q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
K = self.k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
V = self.v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, T, D)
return self.out(out)Lecciones Clave
- Self-attention contextualiza tokens: Cada vector de salida depende de toda la secuencia.
- Escala por $\sqrt{d_k}$ es esencial: Estabiliza softmax y gradientes.
- Cross-attention para seq2seq: Une codificador y decodificador en transformers.
- Multi-head captura aspectos diversos: Gramática, semántica, dependencias a larga distancia.
- Implementaciones en PyTorch usan Linear y bmm: Para eficiencia en lotes.
— Editorial Team
Aún no hay comentarios.