Auto-Atención & Transformers

S2: Atención Multi-Cabeza y Codificación Posicional

Prof. Francisco Suárez

Universidad Católica Boliviana

2026-04-15

Agenda de Hoy

  1. 🔁 Repaso: auto-atención de un solo cabezal
  2. 🎭 Multi-Head Attention: múltiples perspectivas
  3. 📐 Implementación de MHA desde cero
  4. 📍 El problema del orden: ¿por qué importa la posición?
  5. 🌊 Codificación posicional sinusoidal
  6. 🧪 Experimento: MHA + posicional encoding sobre texto

Objetivo

Entender por qué un solo cabezal de atención es insuficiente, cómo Multi-Head Attention captura múltiples tipos de relaciones simultáneamente, y por qué los Transformers necesitan codificación posicional explícita.

Prerequisitos: Scaled Dot-Product Attention (S1)

Repaso: Un Solo Cabezal

Limitación del Cabezal Único

Lo que aprendimos en S1

Un cabezal de atención aprende una proyección \(\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V\) y calcula:

\[\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}}\right)\mathbf{V}\]

El problema

Una sola proyección solo puede capturar un tipo de relación a la vez.

En la oración “El banco cerca del río se inundó”:

  • ¿“banco” se refiere a institución financiera o a orilla?
  • Un solo cabezal debe elegir una perspectiva

Multi-Head Attention

La Idea: Múltiples Perspectivas Simultáneas

En lugar de una proyección, usar \(h\) proyecciones

Cada cabezal \(i\) aprende sus propias matrices \(\mathbf{W}_Q^{(i)}, \mathbf{W}_K^{(i)}, \mathbf{W}_V^{(i)}\) y puede especializarse en un tipo diferente de relación:

\[\text{head}_i = \text{Attention}(\mathbf{Q}\mathbf{W}_Q^{(i)},\; \mathbf{K}\mathbf{W}_K^{(i)},\; \mathbf{V}\mathbf{W}_V^{(i)})\]

Sus salidas se concatenan y se proyectan:

\[\text{MultiHead}(\mathbf{Q},\mathbf{K},\mathbf{V}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\,\mathbf{W}^O\]

Dimensiones

  • \(\mathbf{W}_Q^{(i)}, \mathbf{W}_K^{(i)} \in \mathbb{R}^{d_\text{model} \times d_k}\), con \(d_k = d_\text{model}/h\)
  • La concatenación produce \(h \cdot d_k = d_\text{model}\) → proyección \(\mathbf{W}^O \in \mathbb{R}^{d_\text{model} \times d_\text{model}}\)

Tip

En Vaswani et al. (2017): \(d_\text{model}=512\), \(h=8\), \(d_k=d_v=64\).

Implementación: MultiHeadAttention

Code
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    """
    Atención multi-cabeza (Vaswani et al., 2017).
    d_model: dimensión del modelo
    num_heads: número de cabezales h
    d_k = d_v = d_model // num_heads
    """
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % num_heads == 0, "d_model debe ser divisible por num_heads"

        self.d_model    = d_model
        self.num_heads  = num_heads
        self.d_k        = d_model // num_heads   # dimensión por cabezal

        # Una proyección lineal grande equivale a h proyecciones pequeñas
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)

    def split_heads(self, x):
        """(B, T, d_model) → (B, h, T, d_k)"""
        B, T, _ = x.shape
        x = x.view(B, T, self.num_heads, self.d_k)
        return x.transpose(1, 2)   # (B, h, T, d_k)

    def forward(self, x, mask=None):
        """x: (B, T, d_model)"""
        B, T, _ = x.shape

        # Proyectar y dividir en cabezales
        Q = self.split_heads(self.W_Q(x))   # (B, h, T, d_k)
        K = self.split_heads(self.W_K(x))
        V = self.split_heads(self.W_V(x))

        # Atención escalada por cabezal
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)  # (B, h, T, T)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Contexto: (B, h, T, d_k)
        context = torch.matmul(attn_weights, V)

        # Reensamblar cabezales: (B, T, d_model)
        context = context.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_O(context), attn_weights


# Prueba
torch.manual_seed(0)
B, T, d_model, h = 2, 10, 64, 8
mha = MultiHeadAttention(d_model=d_model, num_heads=h)
x = torch.randn(B, T, d_model)
out, w = mha(x)
print(f"Entrada:      {x.shape}")
print(f"Salida:       {out.shape}          ← misma forma")
print(f"Pesos MHA:    {w.shape}  ← (batch, heads, T, T)")
npar = sum(p.numel() for p in mha.parameters())
print(f"Parámetros:   {npar:,}  (4 matrices de {d_model}x{d_model})")
Entrada:      torch.Size([2, 10, 64])
Salida:       torch.Size([2, 10, 64])          ← misma forma
Pesos MHA:    torch.Size([2, 8, 10, 10])  ← (batch, heads, T, T)
Parámetros:   16,384  (4 matrices de 64x64)

Visualización: Cada Cabezal Aprende Algo Diferente

Note

Estos patrones son ilustrativos. En modelos reales (BERT), Clark et al. (2019) identificaron cabezales especializados en: tokens directamente anteriores/siguientes, tokens relacionados sintácticamente, el token [SEP], y más.

El Problema del Orden

La Auto-Atención es Invariante a Permutaciones

Demostración

Code
torch.manual_seed(1)
B, T, d = 1, 5, 32
mha_test = MultiHeadAttention(d_model=d, num_heads=4)
mha_test.eval()

x = torch.randn(B, T, d)

# Salida original
with torch.no_grad():
    out_orig, _ = mha_test(x)

# Permutar tokens
perm = torch.tensor([2, 0, 4, 1, 3])
x_perm = x[:, perm, :]

with torch.no_grad():
    out_perm, _ = mha_test(x_perm)

# ¿La salida permutada coincide con permutar la salida original?
diff = (out_perm - out_orig[:, perm, :]).abs().max().item()
print(f"Diferencia máxima (out_perm vs out_orig[perm]): {diff:.2e}")
print("→ La MHA no distingue el orden de tokens ⚠️")
Diferencia máxima (out_perm vs out_orig[perm]): 5.96e-08
→ La MHA no distingue el orden de tokens ⚠️

Por qué esto es un problema

Para MHA, estas dos oraciones son idénticas:

  • “El perro muerde al hombre”
  • “El hombre muerde al perro”

Tienen exactamente los mismos tokens — diferente orden.

El modelo nunca verá la diferencia si no le decimos explícitamente qué posición ocupa cada token.

Solución: añadir un vector de posición a cada embedding antes de la atención:

\[\mathbf{x}_i' = \mathbf{e}_i + \text{PE}(i)\]

Codificación Posicional

Codificación Posicional Sinusoidal

Definición (Vaswani et al., 2017)

\[\text{PE}(pos, 2i) = \sin\!\left(\frac{pos}{10000^{2i/d_\text{model}}}\right)\]

\[\text{PE}(pos, 2i+1) = \cos\!\left(\frac{pos}{10000^{2i/d_\text{model}}}\right)\]

  • \(pos\): posición del token en la secuencia (\(0, 1, \ldots, T-1\))
  • \(i\): índice de la dimensión del embedding (\(0, 1, \ldots, d/2-1\))
  • Dimensiones pares: seno, dimensiones impares: coseno

¿Por qué funciona?

  • Cada posición tiene un vector único
  • El modelo puede inferir distancia relativa entre posiciones: \(\text{PE}(pos+k)\) es una transformación lineal de \(\text{PE}(pos)\)
  • No requiere entrenamiento — es determinista

Implementación: PositionalEncoding

Code
import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """
    Codificación posicional sinusoidal (Vaswani et al., 2017).
    Suma un vector posicional fijo al embedding de cada token.
    """
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # Matriz PE: (max_len, d_model)
        PE = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)    # (T, 1)
        div = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float) *
            (-math.log(10000.0) / d_model)
        )                                                                   # (d_model/2,)
        PE[:, 0::2] = torch.sin(pos * div)
        PE[:, 1::2] = torch.cos(pos * div)

        # Registrar como buffer (no es parámetro aprendible)
        self.register_buffer('PE', PE.unsqueeze(0))   # (1, max_len, d_model)

    def forward(self, x):
        """x: (B, T, d_model) → suma PE para las primeras T posiciones"""
        return self.dropout(x + self.PE[:, :x.size(1)])


# Prueba
torch.manual_seed(0)
d_model = 64
pe = PositionalEncoding(d_model=d_model, dropout=0.0)

x_emb = torch.randn(2, 10, d_model)      # (B=2, T=10, d=64)
x_enc = pe(x_emb)

print(f"Embedding sin PE:  {x_emb.shape}")
print(f"Embedding con PE:  {x_enc.shape}  ← mismas dimensiones")
print(f"\nPrimeras 5 dims del PE, posiciones 0-4:")
print(pe.PE[0, :5, :5].numpy().round(3))
print("\n→ Cada fila es única: el modelo puede distinguir posiciones ✅")
Embedding sin PE:  torch.Size([2, 10, 64])
Embedding con PE:  torch.Size([2, 10, 64])  ← mismas dimensiones

Primeras 5 dims del PE, posiciones 0-4:
[[ 0.     1.     0.     1.     0.   ]
 [ 0.841  0.54   0.682  0.732  0.533]
 [ 0.909 -0.416  0.997  0.071  0.902]
 [ 0.141 -0.99   0.778 -0.628  0.993]
 [-0.757 -0.654  0.142 -0.99   0.778]]

→ Cada fila es única: el modelo puede distinguir posiciones ✅

¿Por Qué Ondas de Diferentes Frecuencias?

Tip

La similitud coseno decae con la distancia → el modelo puede inferir que posiciones cercanas son más similares que posiciones lejanas, sin haberlo sido entrenado explícitamente para ello.

Juntando Todo: MHA + Positional Encoding

El Bloque de Entrada de un Transformer

Pipeline completo (input processing)

Tokens  →  Embedding  →  + PE  →  MHA  →  Salida contextualizada

Formalmente:

\[\mathbf{X}' = \text{Embedding}(\mathbf{w}) + \text{PE}\] \[\mathbf{Z} = \text{MultiHead}(\mathbf{X}', \mathbf{X}', \mathbf{X}')\]

Note

En el Transformer completo (S3), hay más componentes: - Add & Norm (conexiones residuales + Layer Norm) - Feed-Forward (MLP por token) - Múltiples capas apiladas

Pero MHA + PE son el corazón.

Code
class TransformerInputBlock(nn.Module):
    """
    Bloque de entrada: Embedding + PE + Multi-Head Attention.
    (Versión simplificada — sin residual ni FFN)
    """
    def __init__(self, vocab_size, d_model, num_heads,
                 max_len=512, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pe         = PositionalEncoding(d_model, max_len, dropout)
        self.mha        = MultiHeadAttention(d_model, num_heads, dropout)

    def forward(self, x, mask=None):
        """x: (B, T) token ids"""
        emb = self.embedding(x)      # (B, T, d_model)
        emb = self.pe(emb)           # + posicional encoding
        out, attn = self.mha(emb, mask)    # multi-head attention
        return out, attn


# Ejemplo
torch.manual_seed(42)
block = TransformerInputBlock(
    vocab_size=10000, d_model=64, num_heads=8
)
tokens = torch.randint(1, 10000, (3, 20))   # batch=3, seq_len=20
out, attn = block(tokens)
print(f"Tokens entrada:     {tokens.shape}")
print(f"Salida contextual:  {out.shape}")
print(f"Pesos atención:     {attn.shape}   ← (B, heads, T, T)")
npar = sum(p.numel() for p in block.parameters())
print(f"Total parámetros:   {npar:,}")
Tokens entrada:     torch.Size([3, 20])
Salida contextual:  torch.Size([3, 20, 64])
Pesos atención:     torch.Size([3, 8, 20, 20])   ← (B, heads, T, T)
Total parámetros:   656,384

Experimento: ¿Importa el Orden Con PE?

Code
torch.manual_seed(7)

d_model, n_heads = 32, 4
vocab_sz = 500

# Sin PE: solo embedding + MHA
class BlockSinPE(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(vocab_sz, d_model, padding_idx=0)
        self.mha = MultiHeadAttention(d_model, n_heads)
    def forward(self, x):
        e = self.emb(x)
        out, _ = self.mha(e)
        return out

# Con PE: embedding + PE + MHA
class BlockConPE(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(vocab_sz, d_model, padding_idx=0)
        self.pe  = PositionalEncoding(d_model, dropout=0.0)
        self.mha = MultiHeadAttention(d_model, n_heads)
    def forward(self, x):
        e = self.pe(self.emb(x))
        out, _ = self.mha(e)
        return out

# Oración y su versión permutada
frase    = torch.tensor([[10, 25, 87, 43, 62, 91, 15, 37]])   # (1, 8)
perm_idx = torch.tensor([3, 0, 6, 1, 7, 4, 2, 5])
frase_p  = frase[:, perm_idx]

blk_sin = BlockSinPE(); blk_sin.eval()
blk_con = BlockConPE(); blk_con.eval()

with torch.no_grad():
    out_sin_orig = blk_sin(frase)
    out_sin_perm = blk_sin(frase_p)
    out_con_orig = blk_con(frase)
    out_con_perm = blk_con(frase_p)

diff_sin = (out_sin_perm - out_sin_orig[:, perm_idx]).abs().mean().item()
diff_con = (out_con_perm - out_con_orig[:, perm_idx]).abs().mean().item()

print("Sin PE — diferencia media entre original y permutado:", f"{diff_sin:.2e}")
print("Con PE — diferencia media entre original y permutado:", f"{diff_con:.2e}")
print()
print("→ Sin PE: permutación NO cambia la salida (invariante al orden)")
print("→ Con PE: permutación SÍ cambia la salida (sensible al orden) ✅")
Sin PE — diferencia media entre original y permutado: 1.41e-08
Con PE — diferencia media entre original y permutado: 3.19e-02

→ Sin PE: permutación NO cambia la salida (invariante al orden)
→ Con PE: permutación SÍ cambia la salida (sensible al orden) ✅

Visualización: PE en el Espacio de Atención

Resumen y Cierre

Lo Que Aprendimos Hoy

Multi-Head Attention

\[\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1,\ldots,\text{head}_h)\mathbf{W}^O\]

  • Un cabezal = una perspectiva sobre las relaciones entre tokens
  • \(h\) cabezales en paralelo → múltiples perspectivas
  • \(d_k = d_\text{model}/h\) para mantener el coste computacional constante
  • Vaswani et al.: \(h=8\), \(d_k=64\), \(d_\text{model}=512\)

Codificación Posicional

\[\text{PE}(pos, 2i) = \sin\!\left(\frac{pos}{10000^{2i/d}}\right)\!, \quad \text{PE}(pos, 2i+1) = \cos(\cdots)\]

  • La MHA sola es invariante al orden — sin PE, “perro muerde hombre” = “hombre muerde perro”
  • PE añade información posicional sin parámetros aprendibles
  • Similitud coseno decae con distancia → el modelo puede inferir proximidad

Lo que queda por completar

Componente Descripción ¿Visto?
Scaled Dot-Product Attention Q·K·V ✅ S1
Multi-Head Attention \(h\) cabezales ✅ Hoy
Positional Encoding PE sinusoidal ✅ Hoy
Add & Norm Residual + LayerNorm ❌ S3
Feed-Forward MLP por token ❌ S3
Stacking \(N\) bloques apilados ❌ S3

→ La próxima sesión completa el Transformer.

Para la Próxima Sesión

S3: El Bloque Transformer Completo

Note

Ya tenemos todos los componentes de atención. La sesión S3 añade los elementos que hacen el entrenamiento estable y profundo:

  1. Conexiones residuales (\(x + \text{sublayer}(x)\)) — evitan el gradiente que desaparece en redes profundas
  2. Layer Normalization — estabiliza las activaciones
  3. Feed-Forward Network por token — capacidad no-lineal
  4. Encoder completo\(N\) bloques apilados
  5. Decoder — con cross-attention (Q del decoder, K/V del encoder)

Arquitectura completa (preview)

Input Tokens
    ↓
Embedding + PE
    ↓
┌─────────────────┐ ×N
│  Multi-Head SA  │
│  Add & Norm     │
│  Feed-Forward   │
│  Add & Norm     │
└─────────────────┘
    ↓
Output (encoder)

Lecturas recomendadas

  • Vaswani et al. (2017) — Secciones 3.3–3.5
  • “The Illustrated Transformer” — Jay Alammar (blog)

Referencias

  • Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017.
  • Clark, K., Khandelwal, U., Levy, O., & Manning, C. D. (2019). What Does BERT Look at? An Analysis of BERT’s Attention. ACL 2019 Workshop BlackboxNLP.
  • Vig, J. (2019). A Multiscale Visualization of Attention in the Transformer Model. ACL 2019.
  • Alammar, J. (2018). The Illustrated Transformer. https://jalammar.github.io/illustrated-transformer/

¿Preguntas? 🙋

¡Gracias!

📧 fsuarez@ucb.edu.bo

🔗 Materiales: github.com/fjsuarez/ucb-nlp