Capítulo 2: Auto-atención y Atención Multi-cabeza

Introducción

Basándonos en nuestra comprensión del mecanismo de atención, este capítulo explora cómo se aplica la atención cuando una secuencia se atiende a sí misma (auto-atención) y por qué el uso de múltiples cabezas de atención en paralelo mejora drásticamente la capacidad del modelo.

Auto-atención: El núcleo de los Transformers

¿Qué es la auto-atención?

En la auto-atención, una secuencia se atiende a sí misma. Cada posición en la entrada puede atender a cualquier otra posición en la misma entrada.

Input:  "The cat sat on the mat"

Self-Attention: Each word attends to ALL words
├─ "The" attends to ["The", "cat", "sat", "on", "the", "mat"]
├─ "cat" attends to ["The", "cat", "sat", "on", "the", "mat"]
├─ "sat" attends to ["The", "cat", "sat", "on", "the", "mat"]
└─ ... and so on

¿Por qué auto-atención?

La auto-atención permite:

1. Dependencias de largo alcance: Conexiones directas entre cualquier par de posiciones
2. Procesamiento en paralelo: Todas las posiciones se calculan simultáneamente
3. Contexto dinámico: Diferentes pesos para diferentes entradas

Formulación matemática

Para la auto-atención, Q, K y V provienen de la misma entrada:

$$\text{SelfAttention}(X) = \text{softmax}\left(\frac{XW_Q(XW_K)^T}{\sqrt{d_k}}\right)XW_V$$

Implementación

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    """
    Self-attention layer where Q, K, V all come from the same input.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        """
        x: Input tensor (batch, seq_len, d_model)
        mask: Optional attention mask
        """
        # All projections from the same input
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Compute attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Comparando tipos de atención

Auto-atención vs. Atención cruzada

Tipo	Fuente de Q	Fuente de K, V	Caso de uso
Auto-atención	Entrada X	Entrada X	Codificador, Decodificador
Atención cruzada	Decodificador	Codificador	Modelos codificador-decodificador

Ejemplo de atención cruzada

class CrossAttention(nn.Module):
    """
    Cross-attention where Q comes from one source, K and V from another.
    Used in encoder-decoder architectures.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k
    
    def forward(self, query_input, key_value_input, mask=None):
        """
        query_input: Source for Q (e.g., decoder state)
        key_value_input: Source for K, V (e.g., encoder output)
        """
        Q = self.W_Q(query_input)
        K = self.W_K(key_value_input)
        V = self.W_V(key_value_input)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Atención Multi-cabeza

La motivación

Una sola cabeza de atención solo puede captar un tipo de relación. La atención multi-cabeza ejecuta múltiples operaciones de atención en paralelo, permitiendo que el modelo atienda conjuntamente a información de diferentes subespacios de representación.

Intuición

Diferentes cabezas pueden aprender diferentes patrones:

• Cabeza 1: Relaciones sintácticas (sujeto-verbo)
• Cabeza 2: Relaciones posicionales (palabras cercanas)
• Cabeza 3: Relaciones semánticas (sinónimos)
• Cabeza 4: Correferencia (pronombres a sustantivos)

Definición matemática

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$$

Donde cada cabeza es:

$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Dimensiones de los parámetros

Para $h$ cabezas con dimensión del modelo $d_{model}$:

• Cada cabeza tiene dimensión $d_k = d_v = d_{model} / h$
• El total de parámetros se mantiene aproximadamente igual que en la atención de una sola cabeza

# Example: d_model = 512, h = 8
d_model = 512
num_heads = 8
d_k = d_model // num_heads  # 64 per head

Implementación completa

class MultiHeadAttention(nn.Module):
    """
    Multi-head attention mechanism.
    """
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Linear projections for Q, K, V (all heads combined)
        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)
        
        # Output projection
        self.W_O = nn.Linear(d_model, d_model, bias=False)
        
        self.dropout = nn.Dropout(dropout)
    
    def split_heads(self, x, batch_size):
        """
        Split the last dimension into (num_heads, d_k).
        Transpose to shape (batch, num_heads, seq_len, d_k)
        """
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Linear projections
        Q = self.W_Q(query)
        K = self.W_K(key)
        V = self.W_V(value)
        
        # Split into heads
        Q = self.split_heads(Q, batch_size)  # (batch, heads, seq, d_k)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)
        
        # Scaled dot-product attention for all heads in parallel
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            # Expand mask for heads dimension
            if mask.dim() == 3:
                mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        # Apply attention to values
        context = torch.matmul(attention_weights, V)
        
        # Concatenate heads
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)
        
        # Final linear projection
        output = self.W_O(context)
        
        return output, attention_weights

Uso de la atención multi-cabeza

# Example usage
d_model = 512
num_heads = 8
seq_len = 100
batch_size = 32

mha = MultiHeadAttention(d_model, num_heads)

# Input (same for self-attention)
x = torch.randn(batch_size, seq_len, d_model)

# Self-attention
output, attn_weights = mha(x, x, x)
print(f"Output shape: {output.shape}")  # (32, 100, 512)
print(f"Attention shape: {attn_weights.shape}")  # (32, 8, 100, 100)

Beneficios de la paralelización

Por qué los Transformers se entrenan rápido

A diferencia de las RNN, todas las posiciones en un Transformer pueden calcularse simultáneamente:

# RNN: Sequential (slow)
for t in range(seq_len):
    h_t = rnn_cell(x[t], h_{t-1})

# Transformer: Parallel (fast)
output = multi_head_attention(x, x, x)  # All positions at once

Eficiencia en GPU

La atención multi-cabeza es especialmente eficiente en GPU:

# All heads computed in a single batched matrix multiplication
# Shape: (batch, heads, seq, d_k) @ (batch, heads, d_k, seq)
# → (batch, heads, seq, seq)
scores = torch.matmul(Q, K.transpose(-2, -1))  # Parallel across all heads!

Visualización de múltiples cabezas

def visualize_multi_head_attention(attention_weights, tokens, num_heads_to_show=4):
    """
    Visualize attention patterns from multiple heads.
    attention_weights: (batch, heads, seq, seq)
    """
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(1, num_heads_to_show, figsize=(4*num_heads_to_show, 4))
    
    for head_idx in range(num_heads_to_show):
        ax = axes[head_idx]
        attn = attention_weights[0, head_idx].detach().numpy()  # First batch item
        
        im = ax.imshow(attn, cmap='Blues')
        ax.set_xticks(range(len(tokens)))
        ax.set_yticks(range(len(tokens)))
        ax.set_xticklabels(tokens, rotation=45, ha='right')
        ax.set_yticklabels(tokens)
        ax.set_title(f'Head {head_idx + 1}')
    
    plt.tight_layout()
    plt.show()

Patrones de atención en la práctica

La investigación ha demostrado que diferentes cabezas aprenden patrones especializados:

Tipo de patrón	Qué captura	Ejemplo
Diagonal	Misma posición / identidad	El token se atiende a sí mismo
Token previo	Relaciones secuenciales	Patrones de modelos de lenguaje
Primer token	Agregación global	Token CLS en BERT
Sintáctico	Estructura gramatical	Concordancia sujeto-verbo
Semántico	Relaciones de significado	Sinónimos, antónimos

Puntos clave

Summary

1. La auto-atención permite que cada posición atienda a todas las posiciones en la misma secuencia
2. La atención cruzada permite atender a una secuencia diferente (codificador-decodificador)
3. La atención multi-cabeza ejecuta múltiples operaciones de atención en paralelo
4. Diferentes cabezas aprenden distintos tipos de relaciones
5. La paralelización hace que los Transformers sean eficientes en GPU

Próximos pasos

En el siguiente capítulo, exploraremos cómo los Transformers manejan la información posicional mediante el codificado posicional, y cómo los tokens se representan a través de embeddings.

---

← Capítulo 1: Mecanismo de Atención

Capítulo 3: Codificado Posicional →