Capítulo 1: Mecanismo de Atención

Introducción

El mecanismo de atención es la innovación central que hace posibles los Transformers. Permite que los modelos se enfoquen dinámicamente en las partes relevantes de la entrada al producir cada salida, superando las limitaciones de los vectores de contexto de tamaño fijo.

La intuición detrás de la atención

Una analogía humana

Piense en cómo lee una oración compleja:

"El animal no cruzó la calle porque estaba demasiado cansado."

Para entender a qué se refiere "estaba", usted mentalmente presta atención a "el animal" en lugar de "la calle". Esto es exactamente lo que hace la atención computacionalmente.

La idea central

La atención calcula una suma ponderada de valores, donde los pesos se determinan según cuán relevante es cada valor para una consulta dada:

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

Query, Key, Value: Las tres proyecciones

Comprendiendo Q, K, V

Componente	Rol	Analogía
Query (Q)	Lo que estamos buscando	Una consulta de búsqueda
Key (K)	Lo que cada posición ofrece	Entradas de un índice
Value (V)	El contenido real	Los documentos

Definición matemática

Dada una entrada $X \in \mathbb{R}^{n \times d}$ (secuencia de n tokens, cada uno con d dimensiones):

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

Donde:

• $W_Q \in \mathbb{R}^{d \times d_k}$
• $W_K \in \mathbb{R}^{d \times d_k}$
• $W_V \in \mathbb{R}^{d \times d_v}$

Implementación

import torch
import torch.nn as nn
import math

class AttentionProjections(nn.Module):
    def __init__(self, d_model, d_k, d_v):
        super().__init__()
        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_v, bias=False)
    
    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        Returns Q, K, V projections
        """
        Q = self.W_Q(x)  # (batch, seq_len, d_k)
        K = self.W_K(x)  # (batch, seq_len, d_k)
        V = self.W_V(x)  # (batch, seq_len, d_v)
        return Q, K, V

Atención de producto punto escalado

Cálculo paso a paso

Paso 1: Calcular los puntajes de atención

$$\text{scores} = QK^T$$

Esto nos da una matriz donde la entrada $(i, j)$ representa cuánto la posición $i$ debe atender a la posición $j$.

# Q: (batch, seq_len, d_k)
# K: (batch, seq_len, d_k)
scores = torch.matmul(Q, K.transpose(-2, -1))
# scores: (batch, seq_len, seq_len)

Paso 2: Escalar los puntajes

$$\text{scaled\_scores} = \frac{QK^T}{\sqrt{d_k}}$$

Why Scale?

Sin escalado, para valores grandes de $d_k$, los productos punto pueden crecer mucho en magnitud, llevando al softmax a regiones con gradientes extremadamente pequeños. Escalar por $\sqrt{d_k}$ mantiene la varianza estable.

d_k = Q.size(-1)
scaled_scores = scores / math.sqrt(d_k)

Paso 3: Aplicar Softmax

$$\text{attention\_weights} = \text{softmax}(\text{scaled\_scores})$$

attention_weights = torch.softmax(scaled_scores, dim=-1)
# Each row sums to 1

Paso 4: Suma ponderada de valores

$$\text{output} = \text{attention\_weights} \cdot V$$

output = torch.matmul(attention_weights, V)
# output: (batch, seq_len, d_v)

Implementación completa

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Compute scaled dot-product attention.
    
    Args:
        Q: Query tensor (batch, seq_len, d_k)
        K: Key tensor (batch, seq_len, d_k)
        V: Value tensor (batch, seq_len, d_v)
        mask: Optional mask tensor
    
    Returns:
        output: Attention output (batch, seq_len, d_v)
        attention_weights: Attention weights (batch, seq_len, seq_len)
    """
    d_k = Q.size(-1)
    
    # Step 1 & 2: Compute scaled scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Optional: Apply mask (for causal attention)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    # Step 3: Softmax
    attention_weights = torch.softmax(scores, dim=-1)
    
    # Step 4: Weighted sum
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights

Visualización de la atención

Matriz de atención

Para una oración como "The cat sat":

         The   cat   sat
The    [ 0.7   0.2   0.1 ]
cat    [ 0.3   0.5   0.2 ]
sat    [ 0.2   0.4   0.4 ]

Cada fila muestra cuánto esa posición atiende a todas las demás posiciones.

Código para visualizar

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(attention_weights, tokens):
    """
    Visualize attention weights as a heatmap.
    """
    plt.figure(figsize=(10, 8))
    sns.heatmap(
        attention_weights.detach().numpy(),
        xticklabels=tokens,
        yticklabels=tokens,
        cmap='Blues',
        annot=True,
        fmt='.2f'
    )
    plt.xlabel('Keys')
    plt.ylabel('Queries')
    plt.title('Attention Weights')
    plt.show()

Atención causal (enmascarada)

Para modelos de lenguaje autorregresivos, cada posición solo debe atender a posiciones anteriores.

La máscara causal

def create_causal_mask(seq_len):
    """
    Create a causal mask for autoregressive attention.
    Position i can only attend to positions <= i.
    """
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# Example for seq_len=4:
# [[1, 0, 0, 0],
#  [1, 1, 0, 0],
#  [1, 1, 1, 0],
#  [1, 1, 1, 1]]

¿Por qué enmascaramiento causal?

Durante el entrenamiento, queremos que el modelo prediga el siguiente token basándose solo en los tokens anteriores, no en los futuros (que no tendrá en el momento de inferencia).

def causal_attention(Q, K, V):
    """
    Attention with causal masking.
    """
    seq_len = Q.size(1)
    mask = create_causal_mask(seq_len).to(Q.device)
    return scaled_dot_product_attention(Q, K, V, mask)

Complejidad de la atención

Complejidad temporal

$$O(n^2 \cdot d)$$

Donde:

• $n$ = longitud de la secuencia
• $d$ = dimensión del modelo

Esta complejidad cuadrática en la longitud de la secuencia es la razón por la que los modelos de contexto largo son un desafío.

Complejidad de memoria

$$O(n^2)$$

Para la propia matriz de atención (almacenando todos los puntajes por pares).

Técnicas de optimización

Existen varios enfoques para abordar la complejidad cuadrática de la atención:

1. Flash Attention: Implementación eficiente en memoria
2. Linear Attention: Aproximación de la atención en O(n)
3. Sparse Attention: Solo atiende a un subconjunto de posiciones
4. Sliding Window: Atención local con tokens globales

Resumiendo todo

class Attention(nn.Module):
    """
    Complete attention module.
    """
    def __init__(self, d_model, d_k=None, d_v=None, causal=False):
        super().__init__()
        d_k = d_k or d_model
        d_v = d_v 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_v, bias=False)
        self.W_O = nn.Linear(d_v, d_model, bias=False)
        
        self.causal = causal
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # Project to Q, K, V
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Compute attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # Apply causal mask if needed
        if self.causal:
            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
            scores = scores.masked_fill(causal_mask == 0, float('-inf'))
        
        # Apply additional mask if provided
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Softmax and weighted sum
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)
        
        # Output projection
        output = self.W_O(output)
        
        return output, attn_weights

Puntos clave

Summary

1. La atención calcula sumas ponderadas basadas en la relevancia
2. Q, K, V son tres proyecciones de la entrada con roles distintos
3. El escalado por $\sqrt{d_k}$ estabiliza el entrenamiento
4. El enmascaramiento causal previene mirar tokens futuros
5. La complejidad es $O(n^2)$ en la longitud de la secuencia

Mirando hacia adelante

En el próximo capítulo, exploraremos cómo se utiliza la atención dentro de una sola secuencia (self-attention) y cómo se combinan múltiples mecanismos de atención (multi-head attention).

---

← Descripción general del módulo

Capítulo 2: Self & Multi-Head Attention →