Atención y Transformers desde Cero: Implementando Redes Neuronales Auto-Atentivas en TensorFlow y PyTorch

🚀 Introducción a los Mecanismos de Atención y Transformers

El campo de la Inteligencia Artificial, y más específicamente el Procesamiento del Lenguaje Natural (PLN), ha experimentado una revolución con la llegada de los mecanismos de atención y las arquitecturas Transformer. Antes de su aparición, las redes neuronales recurrentes (RNNs) y las redes de memoria a corto y largo plazo (LSTMs) eran el estado del arte para tareas secuenciales. Sin embargo, tenían limitaciones inherentes, como la dificultad para manejar dependencias a largo plazo y la falta de paralelización en el entrenamiento.

Los Transformers, introducidos en el artículo seminal "Attention Is All You Need" (Vaswani et al., 2017), resolvieron muchos de estos problemas al depender exclusivamente de mecanismos de atención, desechando por completo la recurrencia. Esto permitió no solo capturar dependencias a distancias mucho mayores, sino también entrenar modelos de manera significativamente más eficiente. Hoy en día, son la base de modelos gigantes como BERT, GPT-3, T5, y muchos otros que dominan el panorama de la IA.

En este tutorial, desglosaremos la magia detrás de los mecanismos de atención y la arquitectura Transformer. Partiremos de los conceptos básicos y avanzaremos hacia una implementación práctica de la auto-atención, tanto en TensorFlow como en PyTorch, dos de las bibliotecas de aprendizaje profundo más populares.

🧐 ¿Qué es la Atención y Por Qué es Importante?

Imagina que estás leyendo un libro y te encuentras con una oración compleja. Naturalmente, tu cerebro presta atención a las palabras más relevantes para entender el significado general de la oración o el contexto de una palabra específica. Los mecanismos de atención en las redes neuronales funcionan de manera análoga: permiten que el modelo ponga mayor foco en las partes más relevantes de la entrada al procesar una porción específica de ella.

La Necesidad de Atención en Secuencias Largas

En los modelos Encoder-Decoder tradicionales sin atención (por ejemplo, basados en RNNs), el encoder comprimía toda la información de la secuencia de entrada en un único vector de contexto de tamaño fijo. Luego, el decoder utilizaba este vector para generar la secuencia de salida. El problema es que para secuencias de entrada muy largas, este vector de contexto a menudo resultaba ser un cuello de botella, incapaz de retener toda la información necesaria. Esto se conocía como el problema de la 'información perdida' o 'cuello de botella del contexto'.

El mecanismo de atención resuelve esto permitiendo que el decoder no solo acceda al vector de contexto final, sino que también "mire" directamente a todas las salidas del encoder en cada paso de tiempo. Crucialmente, asigna un peso o una 'puntuación de atención' a cada una de estas salidas, indicando cuánto debe "atender" a cada una para generar el siguiente elemento de la secuencia de salida.

🔍 Auto-Atención (Self-Attention): El Corazón de los Transformers

Mientras que la atención tradicional conecta un encoder con un decoder, la auto-atención (o self-attention) permite que una secuencia atienda a sí misma. Es decir, cada elemento en una secuencia de entrada puede "pesar" la importancia de todos los demás elementos de la misma secuencia para computar su propia representación. Esto es fundamental para entender cómo los Transformers capturan relaciones contextuales complejas sin recurrencia.

Componentes Clave de la Auto-Atención: Query, Key y Value

Para calcular la auto-atención, cada elemento de la secuencia de entrada (por ejemplo, cada palabra en una oración) se transforma en tres representaciones diferentes:

1. Query (Consulta) Q: Representa la pregunta o el elemento que estamos intentando entender o para el cual estamos buscando información relevante.
2. Key (Clave) K: Representa la información que cada elemento puede ofrecer.
3. Value (Valor) V: Contiene el contenido real del elemento que se agregará a la representación ponderada si su clave es relevante para la consulta.

Estas Q, K y V se obtienen a partir de la misma representación de entrada, multiplicándolas por matrices de pesos entrenables (W_Q, W_K, W_V).

El proceso de auto-atención se puede resumir en los siguientes pasos:

1. Cálculo de similitud: Para cada Query, se calcula su similitud con todas las Keys. Esto se suele hacer mediante un producto escalar (Q @ K.T).
2. Escalado: Las puntuaciones de similitud se escalan dividiéndolas por la raíz cuadrada de la dimensión de las Keys (sqrt(d_k)). Esto ayuda a estabilizar el gradiente durante el entrenamiento.
3. Normalización (Softmax): Se aplica una función softmax a las puntuaciones escaladas para obtener los pesos de atención. Estos pesos suman 1 y representan la importancia relativa de cada Value para la Query actual.
4. Combinación ponderada: Los pesos de atención se multiplican por los Values y se suman para obtener la nueva representación del elemento actual.

La fórmula compacta de la auto-atención escalada por producto escalar es:

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

🛠️ Implementando Auto-Atención en TensorFlow

Vamos a construir una capa de auto-atención desde cero usando TensorFlow. Primero, definiremos las transformaciones lineales para Q, K y V, y luego la lógica para calcular los pesos y la salida ponderada.

import tensorflow as tf

class SelfAttention(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads=8, **kwargs):
        super(SelfAttention, self).__init__(**kwargs)
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}"
            )
        self.proj_dim = embed_dim // num_heads
        
        self.query_dense = tf.keras.layers.Dense(embed_dim, use_bias=False)
        self.key_dense = tf.keras.layers.Dense(embed_dim, use_bias=False)
        self.value_dense = tf.keras.layers.Dense(embed_dim, use_bias=False)
        self.combine_heads = tf.keras.layers.Dense(embed_dim, use_bias=False)

    def attention(self, query, key, value):
        score = tf.matmul(query, key, transpose_b=True)
        dim_key = tf.cast(tf.shape(key)[-1], tf.float32)
        scaled_score = score / tf.math.sqrt(dim_key)
        weights = tf.nn.softmax(scaled_score, axis=-1)
        output = tf.matmul(weights, value)
        return output, weights # Retornamos weights para visualización opcional

    def separate_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.proj_dim))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs):
        batch_size = tf.shape(inputs)[0]

        query = self.query_dense(inputs)  # (batch_size, seq_len, embed_dim)
        key = self.key_dense(inputs)      # (batch_size, seq_len, embed_dim)
        value = self.value_dense(inputs)  # (batch_size, seq_len, embed_dim)

        query = self.separate_heads(query, batch_size)  # (batch_size, num_heads, seq_len, proj_dim)
        key = self.separate_heads(key, batch_size)      # (batch_size, num_heads, seq_len, proj_dim)
        value = self.separate_heads(value, batch_size)  # (batch_size, num_heads, seq_len, proj_dim)

        attention_output, _ = self.attention(query, key, value)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3]) # (batch_size, seq_len, num_heads, proj_dim)
        concat_attention = tf.reshape(attention_output, (batch_size, -1, self.embed_dim)) # (batch_size, seq_len, embed_dim)

        output = self.combine_heads(concat_attention) # (batch_size, seq_len, embed_dim)
        return output

Explicación del Código TensorFlow:

1. __init__: Inicializa las capas densas (Dense) que proyectarán la entrada a las dimensiones de Query, Key y Value. También define num_heads para la atención multi-cabeza, una extensión de la auto-atención que veremos a continuación.
2. attention: Esta es la función central que implementa la fórmula de auto-atención: producto escalar de Q y K transpuesta, escalado por sqrt(d_k), softmax y multiplicación con V.
3. separate_heads: Divide el vector de embed_dim en num_heads cabezas más pequeñas de proj_dim (dimensión proyectada), permitiendo que cada cabeza aprenda un subespacio de atención diferente. Luego transpone las dimensiones para que las cabezas queden agrupadas y el cálculo de atención pueda hacerse en paralelo.
4. call: Este método orquesta todo el proceso. Transforma la entrada en Q, K, V, las divide en cabezas, aplica la función attention y luego concatena las salidas de todas las cabezas antes de proyectarlas de nuevo a embed_dim.

💡 La Importancia de la Atención Multi-Cabeza (Multi-Head Attention)

La auto-atención simple es poderosa, pero la atención multi-cabeza la lleva un paso más allá. En lugar de calcular una única capa de atención, la atención multi-cabeza divide las Q, K y V en num_heads partes, y luego ejecuta la auto-atención en paralelo para cada una de estas "cabezas".

Después de calcular la atención para cada cabeza de forma independiente, las salidas de todas las cabezas se concatenan y se pasan a través de una capa lineal final. Esto permite que el modelo aprenda a procesar la información de una secuencia desde diferentes "perspectivas" o "subespacios de representación", enriqueciendo enormemente su capacidad para capturar relaciones complejas.

🛠️ Implementando Auto-Atención en PyTorch

Ahora, implementaremos la misma lógica de auto-atención multi-cabeza usando PyTorch. La estructura es similar, pero la sintaxis y el manejo de tensores difieren ligeramente.

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

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads=8):
        super(SelfAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}"
            )
        
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)

    def forward(self, x):
        batch_size, seq_len, _ = x.size()

        # Proyecciones lineales para Q, K, V
        q = self.q_proj(x) # (batch_size, seq_len, embed_dim)
        k = self.k_proj(x) # (batch_size, seq_len, embed_dim)
        v = self.v_proj(x) # (batch_size, seq_len, embed_dim)

        # Dividir Q, K, V en múltiples cabezas
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim)

        # Calcular puntuaciones de atención (Q @ K^T)
        scores = torch.matmul(q, k.transpose(-2, -1)) # (batch_size, num_heads, seq_len, seq_len)
        
        # Escalar puntuaciones
        scores = scores / (self.head_dim ** 0.5)

        # Aplicar Softmax para obtener pesos de atención
        attention_weights = F.softmax(scores, dim=-1) # (batch_size, num_heads, seq_len, seq_len)

        # Multiplicar pesos por Values
        context_layer = torch.matmul(attention_weights, v) # (batch_size, num_heads, seq_len, head_dim)

        # Concatenar cabezas y proyectar la salida
        context_layer = context_layer.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim) # (batch_size, seq_len, embed_dim)
        output = self.out_proj(context_layer) # (batch_size, seq_len, embed_dim)

        return output

Explicación del Código PyTorch:

1. __init__: Similar a TensorFlow, se inicializan las capas lineales (nn.Linear) para las proyecciones de Q, K, V y la proyección final de la salida multi-cabeza. head_dim es el equivalente a proj_dim.
2. forward: Este método es el equivalente al call de TensorFlow. Realiza las siguientes operaciones:
   • Proyección: Aplica las capas lineales a la entrada x para obtener q, k, v.
   • Separación de cabezas: Utiliza view y transpose para reconfigurar los tensores de Q, K, V, dividiéndolos en num_heads y moviendo la dimensión de la cabeza para el cálculo en paralelo.
   • Producto escalar y escalado: Calcula Q @ K.T y escala por sqrt(head_dim).
   • Softmax: Aplica F.softmax a lo largo de la dimensión correcta para obtener los pesos de atención.
   • Ponderación con V: Multiplica los pesos de atención por v para obtener el contexto ponderado de cada cabeza.
   • Concatenación y proyección final: Reconfigura el tensor para concatenar las salidas de las cabezas y aplica la proyección lineal final.

🧪 Probando Nuestra Implementación

Para verificar que nuestras capas de auto-atención funcionan, podemos pasar un tensor de entrada simulado y observar las dimensiones de la salida.

Prueba en TensorFlow:

# Crear una instancia de la capa SelfAttention
embed_dim = 128
num_heads = 4
attention_layer_tf = SelfAttention(embed_dim, num_heads)

# Simular datos de entrada: (batch_size, sequence_length, embedding_dimension)
batch_size = 2
seq_len = 10
random_input_tf = tf.random.normal((batch_size, seq_len, embed_dim))

# Pasar la entrada a la capa de atención
output_tf = attention_layer_tf(random_input_tf)

print(f"TensorFlow - Input shape: {random_input_tf.shape}")
print(f"TensorFlow - Output shape: {output_tf.shape}")
# Expected output shape: (2, 10, 128)

Prueba en PyTorch:

# Crear una instancia de la capa SelfAttention
embed_dim = 128
num_heads = 4
attention_layer_pt = SelfAttention(embed_dim, num_heads)

# Simular datos de entrada: (batch_size, sequence_length, embedding_dimension)
batch_size = 2
seq_len = 10
random_input_pt = torch.randn(batch_size, seq_len, embed_dim)

# Pasar la entrada a la capa de atención
output_pt = attention_layer_pt(random_input_pt)

print(f"PyTorch - Input shape: {random_input_pt.shape}")
print(f"PyTorch - Output shape: {output_pt.shape}")
# Expected output shape: (2, 10, 128)

En ambos casos, la forma de la salida debe ser la misma que la de la entrada (batch_size, seq_len, embed_dim), ya que la auto-atención procesa cada token en relación con los demás, pero mantiene la dimensión de la secuencia y la dimensión de embedding.

🏗️ Integrando la Auto-Atención en un Modelo Transformer Básico

Aunque hemos implementado solo la capa de auto-atención, es útil entender dónde encajaría dentro de una arquitectura Transformer completa. Un bloque Transformer típico consiste en:

1. Capa de Auto-Atención Multi-Cabeza: La que acabamos de implementar.
2. Capa de Normalización (Layer Normalization): Aplica normalización a la salida de la atención.
3. Conexión Residual: Suma la entrada original a la salida normalizada de la atención (x + attention_output).
4. Capa Feed-Forward (FFN): Una red neuronal densa con activaciones ReLU (dos capas lineales con una activación intermedia).
5. Capa de Normalización (Layer Normalization): Después de la FFN.
6. Conexión Residual: Suma la entrada del bloque FFN a la salida normalizada del FFN.

Estos bloques se apilan para formar los encoders y decoders de un Transformer completo.

📈 Futuras Extensiones y Aplicaciones

Habiendo implementado la auto-atención, has sentado las bases para explorar arquitecturas Transformer completas. Aquí hay algunas direcciones para continuar:

• Atención Enmascarada (Masked Self-Attention): Esencial para los decoders Transformer, donde la predicción de un token solo puede depender de los tokens anteriores en la secuencia, no de los futuros.
• Atención Encoder-Decoder: La atención cruzada que permite que el decoder atienda a la salida del encoder.
• Embeddings Posicionales: Los Transformers no tienen recurrencia, por lo que necesitan una forma de codificar la posición de los tokens en la secuencia. Esto se hace añadiendo un embedding posicional a los embeddings de los tokens.
• Construcción de un Modelo Transformer Completo: Apilar múltiples bloques encoder y/o decoder para construir un modelo para traducción, resumen, generación de texto, etc.
• Modelos Pre-entrenados: Explorar y ajustar modelos Transformer pre-entrenados como BERT, GPT, T5, que han demostrado un rendimiento excepcional en una amplia gama de tareas.

El conocimiento adquirido sobre la auto-atención es la clave para entender la mayoría de los avances recientes en PLN y más allá, ya que los Transformers también se están aplicando con éxito en visión por computadora (Vision Transformers) y otras áreas.

✅ Conclusión

Los mecanismos de atención y las arquitecturas Transformer han transformado el panorama del aprendizaje profundo, ofreciendo soluciones más eficientes y potentes para el procesamiento de secuencias. Al entender y ser capaz de implementar la auto-atención desde cero, ya sea con TensorFlow o PyTorch, has dado un paso fundamental para comprender cómo funcionan estos modelos de vanguardia.

Este tutorial te ha proporcionado no solo la teoría sino también la implementación práctica, permitiéndote construir una de las piezas más críticas de los Transformers. Con esta base, estás bien equipado para explorar arquitecturas más complejas y contribuir al emocionante campo de la Inteligencia Artificial.