Mecanismo de Atención en PyTorch

El mecanismo de atención (Attention Mechanism) es uno de los conceptos más importantes en el aprendizaje profundo.

El mecanismo de atención permite que el modelo aprenda a "prestar atención" a las partes más relevantes de la entrada, logrando un gran éxito en áreas como el procesamiento del lenguaje natural y la visión por computadora.

Esta sección detalla los principios fundamentales del mecanismo de atención, su implementación en PyTorch y varias variantes de atención.

Versión aplicable:El código de este artículo está basado en PyTorch 2.0+.nn.MultiheadAttentiondebatch_firstEste parámetro se introdujo en PyTorch 1.9; en versiones anteriores es necesario manejar las dimensiones manualmente.

1. Fundamentos del mecanismo de atención

1.1 ¿Por qué se necesita la atención?

Los modelos tradicionales de secuencia a secuencia (Seq2Seq) tienen un problema fundamental:El Encoder necesita comprimir toda la información en un vector de longitud fija. Para secuencias largas, este vector se convierte en un cuello de botella de información: cuanto más larga es la oración, más grave es la pérdida de información.

La idea central del mecanismo de atención es: permitir que el Decoder, al generar cada salida, pueda "ver" todos los estados ocultos del Encoder y asignar dinámicamente diferentes pesos de atención según el contexto actual. Es como cuando un traductor humano, al traducir cada palabra, vuelve a mirar la parte correspondiente del texto original.

1.2 La esencia del mecanismo de atención

El mecanismo de atención puede considerarse como una operaciónde suma ponderada. Dados Query, Key y Value, se asignan pesos calculando la similitud entre Query y cada Key, y luego se realiza una suma ponderada de Value:

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

Donde:

Q（Query）: vector de consulta, representa "qué información estoy buscando"
K（Key）: vector de clave, representa "qué información tengo aquí" (para emparejamiento)
V（Value）: vector de valor, representa "el contenido que realmente puedo proporcionar"
$$d_k$$: la dimensión de Key, $$\sqrt{d_k}$$ se utiliza para escalar, evitando que los valores del producto punto sean demasiado grandes y provoquen la desaparición del gradiente en softmax.

¿Por qué dividir entre $$\sqrt{d_k}$$? Cuando $$d_k$$ es grande, la varianza del producto punto entre Q y K crece linealmente con la dimensión, lo que hace que los valores de entrada de softmax sean demasiado grandes y el gradiente se acerque a cero. Después del escalado, la varianza se restablece a 1, asegurando un flujo normal del gradiente.

Ejemplo

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
"""
Atención de producto punto escalada

Parámetros:
Q: tensor de consulta [batch, n_heads, seq_len_q, d_k]
K: tensor de clave [batch, n_heads, seq_len_k, d_k]
V: tensor de valor [batch, n_heads, seq_len_v, d_v]
mask: tensor de máscara, las posiciones enmascaradas se establecen en False/0

Retorna:
output: salida de atención [batch, n_heads, seq_len_q, d_v]
attention_weights: pesos de atención [batch, n_heads, seq_len_q, seq_len_k]
"""
d_k = Q.size(-1)

# 1. Calcular QK^T / sqrt(d_k)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

# 2. Aplicar la máscara (opcional)
# Las posiciones enmascaradas se establecen con un valor mínimo, tras softmax se aproximan a 0
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)

# 3. Normalización Softmax → pesos de atención
attention_weights = F.softmax(scores, dim=-1)

# 4. Suma ponderada
output = torch.matmul(attention_weights, V)

return output, attention_weights

# Prueba
batch_size, n_heads = 2, 4
seq_len_q, seq_len_k = 5, 6
d_k, d_v = 8, 16

Q = torch.randn(batch_size, n_heads, seq_len_q, d_k)
K = torch.randn(batch_size, n_heads, seq_len_k, d_k)
V = torch.randn(batch_size, n_heads, seq_len_k, d_v)

output, attn_weights = scaled_dot_product_attention(Q, K, V)

print(f"Forma de salida: {output.shape}") # [2, 4, 5, 16]
print(f"Forma de los pesos de atención: {attn_weights.shape}") # [2, 4, 5, 6]
print(f"Suma de fila de pesos (debería ser 1.0): {attn_weights[0, 0, 0].sum().item():.4f}")

La esencia de los pesos de atención es una distribución de probabilidad: después de softmax, cada fila suma 1. Cuanto mayor es el peso, más "atención" presta el modelo a esa posición.

2. Módulos de atención en PyTorch

2.1 Multi-Head Attention

Multi-Head Attention (Atención de múltiples cabezas)Permite que el modelo preste atención simultáneamente a información de diferentes subespacios de representación en diferentes posiciones. Proyecta Q, K, V a múltiples subespacios mediante diferentes proyecciones lineales, calcula la atención en cada subespacio de forma independiente y, finalmente, concatena los resultados y aplica una transformación lineal.

Es como hacer que varias personas examinen el mismo texto desde diferentes ángulos al mismo tiempo: alguien se fija en la estructura sintáctica, alguien en las relaciones semánticas, alguien en las dependencias de larga distancia. El mecanismo de múltiples cabezas permite que el modelo capture patrones más ricos.

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

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

Ejemplo

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0, "d_model debe ser divisible por n_heads"

self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads # Dimensión de cada cabeza

# Proyección lineal de Q, K, V (calcula todas las cabezas de una vez, más eficiente)
self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)

# Proyección de salida
self.w_o = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)

def forward(self, query, key, value, mask=None):
"""
Parámetros:
query: [batch, seq_len_q, d_model]
key: [batch, seq_len_k, d_model]
value: [batch, seq_len_k, d_model]
mask: [batch, 1, 1, seq_len_k] o [batch, 1, seq_len_q, seq_len_k]
"""
batch_size = query.size(0)

# 1. Proyección lineal y luego división en cabezas
# [batch, seq_len, d_model] → [batch, seq_len, n_heads, d_k]
# → [batch, n_heads, seq_len, d_k]
Q = self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

# 2. Atención de producto punto escalada
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)

attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)

# 3. Suma ponderada
context = torch.matmul(attn_weights, V)

# 4. Combinar cabezas: [batch, n_heads, seq_len, d_k] → [batch, seq_len, d_model]
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

# 5. Proyección de salida
output = self.w_o(context)

return output, attn_weights

# Prueba
d_model, n_heads = 128, 8
seq_len, batch = 10, 4

layer = MultiHeadAttention(d_model, n_heads)

query = torch.randn(batch, seq_len, d_model)
key = torch.randn(batch, seq_len, d_model)
value = torch.randn(batch, seq_len, d_model)

output, attn_weights = layer(query, key, value)

print(f"Forma de salida: {output.shape}") # [4, 10, 128]
print(f"Forma de los pesos de atención: {attn_weights.shape}") # [4, 8, 10, 10]
print(f"Dimensión por cabeza d_k = {d_model // n_heads}")

2.2 MultiheadAttention integrado en PyTorch

PyTorch proporciona un nn.MultiheadAttention altamente optimizadonn.MultiheadAttentioncuyo backend utiliza operadores fusionados (fused kernels), por lo que en la mayoría de los casos es más rápido que una implementación manual. Se recomienda su uso en producción.

Ejemplo

import torch
import torch.nn as nn

class TransformerAttention(nn.Module):
"""Capa de autoatención usando el MultiheadAttention integrado de PyTorch"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
embed_dim=d_model,
num_heads=n_heads,
dropout=dropout,
batch_first=True, # El formato de entrada es [batch, seq, features]
# PyTorch 2.0+ puede habilitar el backend Flash Attention:
# attn_implementation="flash_attention_2" # requiere instalar flash-attn
)
self.layernorm = nn.LayerNorm(d_model)

def forward(self, x, key_padding_mask=None):
"""
Autoatención: Q, K, V son todos x
Parámetros:
x: [batch, seq_len, d_model]
key_padding_mask: [batch, seq_len], True indica que la posición es padding
"""
attn_output, attn_weights = self.attention(
x, x, x, # self-attention
key_padding_mask=key_padding_mask
)

# Conexión residual Pre-Norm (entrenamiento más estable que Post-Norm)
output = self.layernorm(x + attn_output)

return output, attn_weights

# Prueba
d_model, n_heads = 128, 8
seq_len, batch = 10, 4

model = TransformerAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)

# Crear máscara de padding: True indica que la posición debe ignorarse
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 7:] = True # Las últimas 3 posiciones de la primera muestra son padding
key_padding_mask[1, 5:] = True # Las últimas 5 posiciones de la segunda muestra son padding

output, attn_weights = model(x, key_padding_mask=key_padding_mask)

print(f"Forma de salida: {output.shape}") # [4, 10, 128]
print(f"Forma de los pesos de atención: {attn_weights.shape}") # [4, 10, 10]

Nota sobre el tipo de mask:nn.MultiheadAttentiondekey_padding_maskSe usa la convenciónTrue = ignorarque es lo contrario de la convención 0 = ignorar en algunas implementaciones personalizadas. Asegúrate de verificarlo al usarla.

3. Variantes del mecanismo de atención

3.1 Autoatención (Self-Attention)

AutoatenciónEs una forma especial del mecanismo de atención: Q, K y V provienen todos de la misma secuencia de entrada. Permite que cada posición de la secuencia atienda directamente a todas las demás, capturando dependencias globales. Este es el componente central de Transformer y la clave de su superioridad sobre RNN: RNN necesita transmitir información paso a paso, mientras que la autoatención lo hace de una sola vez.

Ejemplo

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SelfAttention(nn.Module):
"""Capa de autoatención multi-cabeza con soporte para máscara causal (para generación autorregresiva)"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads

self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)

def _split_heads(self, x, batch_size):
"""[batch, seq, d_model] → [batch, n_heads, seq, d_k]"""
return x.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

def forward(self, x, causal=False):
"""
Autoatención: Q, K, V provienen todos del mismo x de entrada
Parámetros:
x: [batch, seq_len, d_model]
causal: si se usa máscara causal (para evitar ver posiciones futuras)
"""
batch_size, seq_len, _ = x.size()

# Proyección + división en cabezas
Q = self._split_heads(self.w_q(x), batch_size)
K = self._split_heads(self.w_k(x), batch_size)
V = self._split_heads(self.w_v(x), batch_size)

# Calcular puntuaciones de atención
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

# Máscara causal: la matriz triangular superior se establece en -inf para evitar ver tokens futuros
if causal:
causal_mask = torch.triu(
torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
diagonal=1
)
scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))

attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)

# Suma ponderada + combinación de cabezas
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

return self.out_proj(context), attn_weights

# Prueba
d_model, n_heads = 256, 8
seq_len, batch = 20, 2

attn = SelfAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)

# Sin máscara causal (escenario Encoder)
output_enc, weights_enc = attn(x, causal=False)
print(f"[Encoder] Salida: {output_enc.shape}, Pesos: {weights_enc.shape}")

# Con máscara causal (escenario Decoder / generación autorregresiva)
output_dec, weights_dec = attn(x, causal=True)
print(f"[Decoder] Salida: {output_dec.shape}, Pesos: {weights_dec.shape}")

# Verificar la causalidad: el primer token no debe atender a los tokens posteriores
print(f"Con causal=True, weights[0,0,0,5:] (deberían ser todos 0): {weights_dec[0, 0, 0, 5:].tolist()[:5]}")

3.2 Atención cruzada (Cross-Attention)

Atención cruzadaEn ella, Q proviene de una secuencia (por ejemplo, Decoder), y K y V de otra secuencia (por ejemplo, Encoder). Es el puente de la arquitectura Encoder-Decoder: permite que el Decoder, al generar cada token, consulte dinámicamente todas las salidas del Encoder.

Ejemplo

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CrossAttention(nn.Module):
"""Atención cruzada: Q proviene de la secuencia objetivo, K/V de la secuencia fuente"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads

self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)

def forward(self, target, source, mask=None):
"""
Parámetros:
target: secuencia objetivo [batch, target_len, d_model] → proporciona Q
source: secuencia fuente [batch, source_len, d_model] → proporciona K, V
mask: máscara opcional
"""
batch_size = target.size(0)
target_len = target.size(1)
source_len = source.size(1)

# Proyección + división en cabezas
Q = self.w_q(target).view(batch_size, target_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)

# Cálculo de atención
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'))

attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)

# Suma ponderada + combinación de cabezas
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, target_len, self.d_model)

return self.out_proj(context), attn_weights

# Prueba
d_model, n_heads = 256, 8
batch = 2
target_len, source_len = 10, 15

cross_attn = CrossAttention(d_model, n_heads)
target = torch.randn(batch, target_len, d_model) # Salida del Decoder
source = torch.randn(batch, source_len, d_model) # Salida del Encoder

output, weights = cross_attn(target, source)

print(f"Forma de la secuencia objetivo: {target.shape}")
print(f"Forma de la secuencia fuente: {source.shape}")
print(f"Forma de la salida: {output.shape}") # [2, 10, 256]
print(f"Forma de los pesos: {weights.shape}") # [2, 8, 10, 15]
# Significado de los pesos: grado de atención de cada posición target hacia cada posición source

Autoatención vs Atención cruzada:En la autoatención, las formas de Q/K/V son iguales (matriz de pesos $$seq \times seq$$); en la atención cruzada, la longitud de Q y la de K pueden diferir (matriz de pesos $$target \times source$$).

3.3 Codificación posicional (Positional Encoding)

El mecanismo de atención en sí tieneinvariancia a la permutación— desordenar el orden de entrada no cambia la salida. Pero el lenguaje y las imágenes son ordenados, por lo que se necesita inyectar explícitamente la información posicional mediante codificación posicional. El Transformer original utiliza funciones seno/coseno para generar codificaciones posicionales fijas:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right), \quad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

Ejemplo

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
"""Codificación posicional seno/coseno (esquema original de Transformer)"""
def __init__(self, d_model, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)

# Precalcular la matriz de codificación posicional [max_len, d_model]
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

# div_term: controla la velocidad de decaimiento de las diferentes dimensiones de frecuencia
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)

pe[:, 0::2] = torch.sin(position * div_term) # Dimensiones pares
pe[:, 1::2] = torch.cos(position * div_term) # Dimensiones impares

pe = pe.unsqueeze(0) # [1, max_len, d_model]
self.register_buffer('pe', pe) # No participa en el entrenamiento

def forward(self, x):
"""
x: [batch, seq_len, d_model]
Retorna: x + positional_encoding
"""
x = x + self.pe[:, :x.size(1), :]
return self.dropout(x)

# Prueba
d_model = 128
seq_len = 50
batch = 4

pe = PositionalEncoding(d_model)
x = torch.randn(batch, seq_len, d_model)
output = pe(x)

print(f"Forma de entrada: {x.shape}")
print(f"Forma de salida: {output.shape}")
print(f"Forma de la codificación posicional: {pe.pe.shape}") # [1, 5000, 128]

4. Aplicaciones visuales del mecanismo de atención

El mecanismo de atención no solo se aplica a datos secuenciales, también brilla en la visión por computadora. A continuación se presentan tres módulos clásicos de atención visual:

4.1 Atención de canal (Channel Attention)

Atención de canal(como SENet) permite que la red aprenda "qué canales son más importantes". Realiza un pooling global en cada canal para comprimir la información espacial y luego aprende las dependencias entre canales mediante una red totalmente conectada de dos capas.

Ejemplo

import torch
import torch.nn as nn

class SEBlock(nn.Module):
"""Squeeze-and-Excitation Block (atención de canal)"""
def __init__(self, channels, reduction=16):
super().__init__()
self.squeeze = nn.AdaptiveAvgPool2d(1) # Pooling promedio global
self.excitation = nn.Sequential(
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid() # Salida de pesos de canal entre 0 y 1
)

def forward(self, x):
# x: [batch, channels, H, W]
b, c, _, _ = x.size()

# Squeeze: [B, C, H, W] → [B, C, 1, 1] → [B, C]
y = self.squeeze(x).view(b, c)

# Excitation: aprender pesos de canal [B, C]
y = self.excitation(y).view(b, c, 1, 1)

# Ponderación de canales: cada canal se multiplica por su peso correspondiente
return x * y.expand_as(x)

# Prueba
se = SEBlock(channels=256, reduction=16)
x = torch.randn(4, 256, 32, 32)
output = se(x)
print(f"Entrada: {x.shape} → Salida: {output.shape}")

4.2 Atención espacial (Spatial Attention)

Atención espacialpermite que la red aprenda "qué posiciones son más importantes". Comprime a lo largo de la dimensión de canales (tomando la media y el máximo) y luego usa una convolución para aprender el mapa de pesos espaciales.

Ejemplo

import torch
import torch.nn as nn

class SpatialAttention(nn.Module):
"""Módulo de atención espacial"""
def __init__(self, kernel_size=7):
super().__init__()
padding = kernel_size // 2
# Entrada con 2 canales (media + máximo), salida con 1 canal (mapa de atención)
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()

def forward(self, x):
# x: [batch, channels, H, W]

# Compresión de canales: tomar media y máximo → uno [B, 1, H, W] cada uno
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)

# Concatenación → [B, 2, H, W]
scale = torch.cat([avg_out, max_out], dim=1)

# Convolución + Sigmoid → mapa de pesos espaciales [B, 1, H, W]
scale = self.sigmoid(self.conv(scale))

# Ponderación espacial
return x * scale

# Prueba
spatial_attn = SpatialAttention(kernel_size=7)
x = torch.randn(4, 256, 32, 32)
output = spatial_attn(x)
print(f"Entrada: {x.shape} → Salida: {output.shape}")

4.3 CBAM（Convolutional Block Attention Module）

CBAM combina la atención de canal y la atención espacialen serieAplicación: primero se seleccionan los canales importantes mediante atención de canal, luego se seleccionan las posiciones importantes mediante atención espacial. Son complementarios y el resultado es mejor que usarlos por separado.

Ejemplo

import torch
import torch.nn as nn

class CBAM(nn.Module):
"""CBAM: Atención de canal + Atención espacial"""
def __init__(self, channels, reduction=16, kernel_size=7):
super().__init__()
# Atención de canal
self.channel_attn = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid()
)
# Atención espacial
self.spatial_attn = nn.Sequential(
nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False),
nn.Sigmoid()
)

def forward(self, x):
# 1. Atención de canal
b, c, _, _ = x.size()
ca = self.channel_attn(x).view(b, c, 1, 1)
x = x * ca

# 2. Atención espacial
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
sa = self.spatial_attn(torch.cat([avg_out, max_out], dim=1))
x = x * sa

return x

# Prueba
cbam = CBAM(channels=256)
x = torch.randn(4, 256, 32, 32)
output = cbam(x)
print(f"Entrada: {x.shape} → Salida: {output.shape}")

5. Aplicaciones del mecanismo de atención en NLP

5.1 Transformer Encoder

El Encoder de Transformer se compone de N capas idénticas apiladas, cada capa contiene dos subcapas:Autoatención multi-cabezayred feed-forwardCada subcapa utiliza conexión residual y normalización de capa.

Ejemplo

import torch
import torch.nn as nn
import math

class TransformerEncoderLayer(nn.Module):
"""Encoder de Transformer de una capa"""
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
# Autoatención multi-cabeza
self.self_attn = nn.MultiheadAttention(
d_model, n_heads, dropout=dropout, batch_first=True
)
# Red feed-forward (dos capas lineales + activación)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(), # GELU es más común que ReLU
nn.Dropout(dropout),
nn.Linear(d_ff, d_model)
)
# Normalización de capas (estilo Pre-Norm, entrenamiento más estable)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)

def forward(self, x, key_padding_mask=None):
# Sublayer 1: Auto-atención
residual = x
x = self.norm1(x) # Pre-Norm
attn_out, _ = self.self_attn(x, x, x, key_padding_mask=key_padding_mask)
x = residual + self.dropout(attn_out)

# Sublayer 2: Red feed-forward
residual = x
x = self.norm2(x) # Pre-Norm
x = residual + self.dropout(self.ffn(x))

return x

class TransformerEncoder(nn.Module):
"""Encoder Transformer multicapa"""
def __init__(self, d_model, n_heads, d_ff, num_layers, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList([
TransformerEncoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(num_layers)
])
self.final_norm = nn.LayerNorm(d_model)

def forward(self, x, key_padding_mask=None):
for layer in self.layers:
x = layer(x, key_padding_mask)
return self.final_norm(x)

# Prueba
d_model, n_heads, d_ff = 128, 4, 512
num_layers, seq_len, batch = 3, 20, 4

encoder = TransformerEncoder(d_model, n_heads, d_ff, num_layers)
x = torch.randn(batch, seq_len, d_model)

# Máscara de padding: True = ignorar esta posición
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 15:] = True

output = encoder(x, key_padding_mask=key_padding_mask)
print(f"Entrada: {x.shape} → Salida: {output.shape}")

5.2 Visualización de la atención

Visualizar los pesos de atención es un medio importante para comprender y depurar los modelos Transformer. Mediante mapas de calor se puede ver de forma intuitiva qué información atiende el modelo en cada posición.

Ejemplo

import torch
import matplotlib.pyplot as plt
import numpy as np

def visualize_attention(attention_weights, tokens=None, save_path=None):
"""
Visualizar los pesos de atención multi-cabeza

Parámetros:
attention_weights: [n_heads, seq_len, seq_len]
tokens: lista de tokens (opcional)
save_path: ruta de guardado (opcional)
"""
n_heads = attention_weights.shape[0]
seq_len = attention_weights.shape[1]

fig, axes = plt.subplots(1, n_heads, figsize=(n_heads * 3, 3.5))
if n_heads == 1:
axes = [axes]

for head in range(n_heads):
ax = axes[head]
attn = attention_weights[head].detach().cpu().numpy()

im = ax.imshow(attn, cmap='Blues', aspect='auto', vmin=0, vmax=attn.max())
ax.set_title(f'Head {head + 1}', fontsize=10, fontweight='bold')
ax.set_xlabel('Key')
ax.set_ylabel('Query')

if tokens is not None:
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=90, fontsize=8)
ax.set_yticklabels(tokens, fontsize=8)

plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.show()

# Simular los pesos de atención tras el entrenamiento
n_heads, seq_len = 4, 8
attention_weights = torch.softmax(torch.randn(n_heads, seq_len, seq_len), dim=-1)
tokens = ['The', 'cat', 'sat', 'on', 'the', 'mat', '.', '[PAD]']

visualize_attention(attention_weights, tokens)

6. Clasificador completo con atención

6.1 Clasificación de texto basada en atención

A continuación se muestra un modelo completo de clasificación de texto: utiliza el Encoder Transformer para extraer características, luego comprime secuencias de longitud variable en vectores de longitud fija mediante atención de pooling (Attention Pooling) y finalmente las envía al clasificador.

Ejemplo

import torch
import torch.nn as nn
import math

class AttentionClassifier(nn.Module):
"""Modelo de clasificación de texto basado en Transformer + atención de pooling"""
def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes,
n_heads=4, num_layers=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim

# Embedding de palabras
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.pos_encoding = PositionalEncoding(embed_dim, dropout=dropout)

# Encoder Transformer
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=n_heads,
dim_feedforward=hidden_dim,
dropout=dropout,
batch_first=True,
activation='gelu'
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)

# Atención de pooling: aprende el peso de importancia de cada posición
self.attn_pool = nn.Sequential(
nn.Linear(embed_dim, 1) # Cada posición genera un peso escalar
)

# Cabeza de clasificación
self.classifier = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, num_classes)
)

def forward(self, x):
"""
x: [batch, seq_len] — token ids
Retorna: logits [batch, num_classes], attn_weights [batch, seq_len]
"""
# Embedding de palabras + codificación posicional
mask = (x == 0) # padding mask
x = self.embedding(x) * math.sqrt(self.embed_dim)
x = self.pos_encoding(x)

# Codificación Transformer
x = self.transformer(x, src_key_padding_mask=mask)

# Atención de pooling: calcula pesos para cada posición y suma ponderada
scores = self.attn_pool(x).squeeze(-1) # [batch, seq_len]
scores = scores.masked_fill(mask, float('-inf')) # Ignorar padding
attn_weights = torch.softmax(scores, dim=1) # [batch, seq_len]
pooled = torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1) # [batch, embed_dim]

# Clasificación
logits = self.classifier(pooled)
return logits, attn_weights

# Prueba
model = AttentionClassifier(
vocab_size=10000, embed_dim=128, hidden_dim=256, num_classes=5
)

# Entrada simulada (con padding)
x = torch.randint(1, 10000, (16, 50))
x[:, -10:] = 0 # Las últimas 10 posiciones son padding

logits, attn_weights = model(x)
print(f"Entrada: {x.shape}")
print(f"Logits: {logits.shape}") # [16, 5]
print(f"Pesos de atención: {attn_weights.shape}") # [16, 50]

6.2 Comparación de variantes comunes de atención

Tipo	Origen de Q	Origen de K/V	Uso principal	Aplicaciones típicas
Auto-atención (Self-Attention)	Propio	Propio	Interacción global dentro de la secuencia	Transformer、BERT、ViT
Atención cruzada (Cross-Attention)	Secuencia objetivo	Secuencia fuente	Fusión de información entre secuencias	Traducción automática, descripción de imágenes, Stable Diffusion
Atención de canal (Channel Attention)	Características de pooling global	Igual que el anterior	Aprender la importancia de los canales	SENet、CBAM
Atención espacial (Spatial Attention)	Características comprimidas por canal	Igual que el anterior	Aprender la importancia de las posiciones espaciales	CBAM、Spatial Transformer
Atención dispersa (Sparse Attention)	Algunas posiciones	Algunas posiciones	Reducir la complejidad computacional en secuencias largas	Longformer、BigBird

7. Mejores prácticas y preguntas frecuentes

7.1 Consejos de uso

Consistencia de la máscara: Asegúrate de que la dirección de la máscara de padding y la máscara causal sean consistentes (True = ignorar vs 1 = conservar). Las convenciones de las distintas API pueden variar.
Selección del número de cabezas: Normalmente se establece como factor de $$d_{model}$$, comúnmente 8 o 16. Demasiadas cabezas reducen la dimensión $$d_k$$ de cada cabeza y pueden perjudicar el rendimiento.
Conexión residual: Asegúrate de utilizar conexión residual + normalización de capas antes y después de la capa de atención; esto es una condición necesaria para entrenar Transformers profundos.
Calentamiento de la tasa de aprendizaje: Al inicio del entrenamiento de Transformers, los gradientes fluctúan considerablemente; usar un programador de warmup + decaimiento de coseno puede mejorar significativamente la estabilidad.
Flash Attention: PyTorch 2.0+ permite, mediantenn.functional.scaled_dot_product_attentioninvocar automáticamente Flash Attention, con grandes optimizaciones en memoria y velocidad.

7.2 Preguntas frecuentes

Problema	Causa	Solución
Los pesos de atención tienen distribución uniforme	Entrada sin normalizar o tasa de aprendizaje demasiado alta	Revisar la escala de la entrada, usar LayerNorm, reducir la tasa de aprendizaje
La pérdida de entrenamiento no disminuye	Codificación posicional ausente o máscara incorrecta	Confirmar que se añadió la codificación posicional y verificar la dirección de la máscara
Desbordamiento de memoria (secuencias largas)	Complejidad $$O(n^2)$$ de la matriz de atención	Usar Flash Attention, puntos de control de gradiente o atención dispersa
Degradación de la atención en inferencia	Máscara inconsistente entre entrenamiento/inferencia	Asegurarse de que la máscara de padding use la misma lógica en entrenamiento e inferencia

La atención es la tecnología central del deep learning moderno. Desde NLP hasta CV, desde voz hasta predicción de estructura de proteínas, casi todos los modelos de vanguardia se basan en la atención. Comprender sus principios e implementación es el primer paso clave para dominar modelos como Transformer, ViT, Stable Diffusion y GPT.

Otras extensiones