Tutorial TensorFlow: Transformers

Fundamentos de los Transformers

Los Transformers han revolucionado el campo del aprendizaje profundo, especialmente en el procesamiento del lenguaje natural (PLN). Se presentaron como una alternativa a las redes neuronales recurrentes (RNN) y las redes neuronales convolucionales (CNN) para manejar secuencias de datos de manera más eficiente.

A diferencia de las RNN, que procesan secuencias de forma secuencial, los Transformers permiten el procesamiento en paralelo de los datos de entrada. Esto es posible gracias al enfoque en el mecanismo de atención, que permite al modelo identificar y ponderar las relaciones entre los elementos de la secuencia sin necesidad de procesarlos en orden.

Las limitaciones de las RNN en capturar dependencias a largo plazo y su ineficiencia en el procesamiento paralelo llevaron al desarrollo de los Transformers. Estas arquitecturas eliminan la recurrencia y utilizan únicamente mecanismos de atención y capas totalmente conectadas, mejorando así la paralelización durante el entrenamiento y reduciendo los tiempos de cómputo.

Un elemento clave en los Transformers es la incorporación de codificaciones posicionales (positional encodings), que añaden información sobre la posición de cada elemento en la secuencia. Esto es esencial para que el modelo pueda entender el orden y la estructura de los datos, dado que no se procesan de forma secuencial.

Los Transformers también han impulsado el avance en el aprendizaje transferido, permitiendo el afinamiento fino (fine-tuning) de modelos preentrenados en grandes conjuntos de datos. Este enfoque ha facilitado la adaptación de modelos a tareas específicas con menos recursos y datos etiquetados.

Además de su impacto en el PLN, los Transformers se han expandido a otras áreas como la visión por computadora, donde han demostrado ser efectivos en tareas de reconocimiento y generación de imágenes. 

La arquitectura de los Transformers representa un cambio paradigmático en cómo se abordan las tareas de aprendizaje con secuencias, ofreciendo una combinación de eficiencia y rendimiento que supera a las arquitecturas tradicionales en muchas aplicaciones.

Mecanismo de atención (Attention Mechanism)

El mecanismo de atención es el elemento central en la arquitectura de los Transformers, permitiendo al modelo enfocarse en distintas partes de la entrada para generar cada elemento de la salida. Este mecanismo supera las limitaciones de las redes neuronales recurrentes al capturar dependencias a largo plazo sin necesidad de procesar las secuencias de forma secuencial.

En el corazón del mecanismo de atención se encuentran los conceptos de Consultas (Queries), Claves (Keys) y Valores (Values). Cada token de entrada se transforma en estos tres vectores mediante matrices de proyección aprendidas. Las Consultas representan la información que se busca, las Claves indican dónde está esa información y los Valores contienen la información real que se necesita extraer.

La atención de producto escalar escalado (scaled dot-product attention) es el cálculo fundamental que permite al modelo determinar la importancia de cada elemento en relación con otro. Se calcula el producto escalar entre las consultas y las claves, se divide por la raíz cuadrada de la dimensión de las claves para estabilizar los gradientes y se aplica una función softmax para obtener los pesos de atención.

Matemáticamente, se define de la siguiente manera:

import tensorflow as tf

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Producto escalar entre consultas y claves
    matmul_qk = tf.matmul(Q, K, transpose_b=True)
    
    # Escalado
    d_k = tf.cast(tf.shape(K)[-1], tf.float32)
    scaled_scores = matmul_qk / tf.math.sqrt(d_k)
    
    # Aplicación de máscara (opcional)
    if mask is not None:
        scaled_scores += (mask * -1e9)
    
    # Cálculo de pesos de atención
    attention_weights = tf.nn.softmax(scaled_scores, axis=-1)
    
    # Producto de pesos por valores
    output = tf.matmul(attention_weights, V)
    
    return output, attention_weights

En esta implementación, Q (Consultas), K (Claves) y V (Valores) son tensores, y la función devuelve tanto la salida ponderada como los pesos de atención. La máscara permite bloquear ciertas posiciones, lo que es útil en tareas como el modelado de lenguaje.

El concepto de atención multi-cabezas (multi-head attention) amplía este mecanismo al permitir que el modelo atienda a diferentes subespacios de representación. Se utilizan múltiples cabezas de atención que operan en paralelo, cada una con sus propias matrices de proyección para Q, K y V. Esto permite al modelo capturar diversas relaciones y patrones en los datos.

A continuación se muestra una implementación simplificada de la atención multi-cabezas en TensorFlow:

class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model debe ser divisible entre num_heads"
        
        self.num_heads = num_heads
        self.depth = d_model // num_heads
        
        self.Wq = tf.keras.layers.Dense(d_model)
        self.Wk = tf.keras.layers.Dense(d_model)
        self.Wv = tf.keras.layers.Dense(d_model)
        self.linear = tf.keras.layers.Dense(d_model)
    
    def split_heads(self, x, batch_size):
        # Divide la última dimensión en (num_heads, depth)
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        # Transpone para obtener (batch_size, num_heads, seq_len, depth)
        return tf.transpose(x, perm=[0, 2, 1, 3])
    
    def call(self, v, k, q, mask=None):
        batch_size = tf.shape(q)[0]
        
        # Proyecciones lineales
        Q = self.Wq(q)  # (batch_size, seq_len_q, d_model)
        K = self.Wk(k)  # (batch_size, seq_len_k, d_model)
        V = self.Wv(v)  # (batch_size, seq_len_v, d_model)
        
        # División en cabezas
        Q = self.split_heads(Q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
        K = self.split_heads(K, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
        V = self.split_heads(V, batch_size)  # (batch_size, num_heads, seq_len_v, depth)
        
        # Atención de producto escalar escalado
        scaled_attention, attention_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # Concatenación de cabezas
        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
        concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))
        
        # Proyección final
        output = self.linear(concat_attention)
        
        return output, attention_weights

Esta clase personalizada MultiHeadAttention extiende la funcionalidad de una capa de TensorFlow, encapsulando el proceso de atención multi-cabezas. La utilización de múltiples cabezas permite al modelo enfocarse en diferentes posiciones y aspectos de la secuencia simultáneamente.

El mecanismo de atención es crucial para capturar dependencias a larga y corta distancia, permitiendo que cada posición en la secuencia pueda acceder directamente a todas las demás posiciones sin tener que pasar por un camino secuencial. Esto mejora significativamente la capacidad del modelo para comprender contextos complejos y relaciones entre elementos.

Además, la atención facilita la paralelización del procesamiento, ya que no depende de estados previos como en las RNN tradicionales. Esto resulta en una mayor eficiencia computacional y acelera el entrenamiento de modelos grandes en TensorFlow.

La flexibilidad del mecanismo de atención también permite su adaptación a diferentes tipos de datos y tareas. Por ejemplo, en visión por computador, la atención puede aplicarse a características espaciales, mientras que en audio puede enfocarse en características temporales.

En el contexto de TensorFlow, el uso del mecanismo de atención se integra de manera fluida con las APIs de Keras, permitiendo construir modelos personalizados y aprovechar al máximo las capacidades de esta arquitectura avanzada. La comprensión profunda de este mecanismo es esencial para desarrollar soluciones de vanguardia en el campo del aprendizaje profundo.

Arquitectura del Transformer

La arquitectura del Transformer es una innovación que permite procesar secuencias de datos en su totalidad, sin necesidad de recurrencia o convolución. Se compone de una serie de bloques modulares que facilitan el aprendizaje de relaciones complejas en los datos. El Transformer está formado principalmente por dos componentes: el codificador (encoder) y el decodificador (decoder).

El codificador es una pila de N capas idénticas, cada una de las cuales contiene dos subcapas principales: una capa de atención multi-cabezas y una red feed-forward totalmente conectada. Cada subcapa está rodeada por conexiones residuales y seguida de una normalización de capa (Layer Normalization). Esta arquitectura permite al modelo capturar dependencias en diferentes niveles y escalas.

Matemáticamente, la salida de cada subcapa se puede definir como:

import tensorflow as tf

def sublayer_output(x, sublayer):
    # Aplicación de conexión residual y normalización
    return tf.keras.layers.LayerNormalization(epsilon=1e-6)(x + sublayer(x))

En este código, la función sublayer_output aplica una conexión residual sumando la entrada original x al resultado de la subcapa, seguida de una normalización. Este enfoque ayuda a mitigar el problema del desvanecimiento del gradiente y facilita el entrenamiento de arquitecturas profundas.

La capa de atención multi-cabezas en el codificador permite que el modelo atienda a diferentes posiciones en la secuencia de entrada simultáneamente. Esto es especialmente útil para capturar relaciones contextuales complejas. La salida de esta capa se formula como:

attention_output, _ = MultiHeadAttention(d_model, num_heads)(q=x, k=x, v=x)

Aquí, la misma entrada x se utiliza para las consultas, claves y valores, ya que el codificador utiliza auto-atención (self-attention). La salida de la atención se procesa luego a través de la red feed-forward, que es una red neuronal aplicada de manera independiente a cada posición de la secuencia.

La red feed-forward utilizada en el codificador es una MLP (Perceptrón Multicapa) aplicada a cada posición de forma individual y compartiendo parámetros a lo largo de las posiciones. Se compone de dos capas lineales con una activación no lineal intermedia. En TensorFlow, se puede implementar así:

def point_wise_feed_forward_network(d_model, dff):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(dff, activation='relu'),  # Capa intermedia
        tf.keras.layers.Dense(d_model)  # Capa de salida
    ])

En este fragmento, dff es la dimensión de la capa oculta, que suele ser mayor que la dimensión del modelo d_model. Esta expansión en la dimensión permite al modelo aprender representaciones más complejas.

El decodificador también es una pila de N capas, similar al codificador, pero con algunas diferencias clave. Cada capa del decodificador contiene tres subcapas:

1. Atención máscara multi-cabezas (Mecanismo de atención con máscara) que previene que las posiciones futuras influyan en las actuales durante el entrenamiento, crucial para tareas de secuencia a secuencia.
2. Atención multi-cabezas que toma la salida del codificador como claves y valores, y la salida de la subcapa anterior como consultas.
3. Red feed-forward similar a la del codificador.

La atención con máscara se implementa aplicando una máscara que impide que el modelo atienda a posiciones futuras. En TensorFlow, esto se puede lograr de la siguiente manera:

def create_look_ahead_mask(size):
    # Máscara triangular superior
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # Tensor de forma (size, size)

Esta máscara se aplica durante el cálculo de la atención para asegurar que el modelo solo considere las posiciones anteriores o actuales.

Las codificaciones posicionales son esenciales en la arquitectura del Transformer debido a la ausencia de recurrencia o convolución. Estas codificaciones aportan información sobre la posición relativa o absoluta de los tokens en la secuencia, permitiendo al modelo incorporar el orden de la secuencia. Se utilizan funciones seno y coseno de diferentes frecuencias para generar estas codificaciones:

import numpy as np

def positional_encoding(position, d_model):
    # Cálculo de posiciones
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                            np.arange(d_model)[np.newaxis, :],
                            d_model)
    # Aplicación de seno a las dimensiones pares
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
    # Aplicación de coseno a las dimensiones impares
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    
    pos_encoding = angle_rads[np.newaxis, ...]
    return tf.cast(pos_encoding, dtype=tf.float32)

En este código, la función positional_encoding genera una matriz de codificaciones posicionales que se suma a los embeddings de entrada.

La combinación de conexiones residuales, normalización de capa y codificaciones posicionales permite que la arquitectura del Transformer aprenda eficientemente relaciones complejas. Las conexiones residuales ayudan a una mejor propagación del gradiente, la normalización de capa estabiliza las activaciones y las codificaciones posicionales introducen información de orden.

El flujo de datos a través del Transformer sigue un patrón bien definido. En el codificador, los embeddings de entrada se suman con las codificaciones posicionales y pasan por las capas de atención y feed-forward. El decodificador recibe los embeddings de salida, también combinados con codificaciones posicionales, y aplica atención máscara antes de interactuar con la representación del codificador.

La implementación completa de un Transformer en TensorFlow Keras implica ensamblar estos componentes de manera modular. A continuación, se muestra cómo podría definirse una clase para el codificador:

class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super().__init__()
        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = point_wise_feed_forward_network(d_model, dff)
        
        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        
        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
        
    def call(self, x, training, mask):
        # Capa de atención multi-cabezas
        attn_output, _ = self.mha(x, x, x, mask)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)  # Conexión residual
        
        # Capa feed-forward
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)  # Conexión residual
        
        return out2

Este código define una capa del codificador con los componentes esenciales: atención, feed-forward, normalización y conexiones residuales. El parámetro rate controla la tasa de dropout para regularización.

De manera similar, se puede definir el decodificador:

class DecoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super().__init__()
        self.mha1 = MultiHeadAttention(d_model, num_heads)
        self.mha2 = MultiHeadAttention(d_model, num_heads)
        
        self.ffn = point_wise_feed_forward_network(d_model, dff)
        
        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        
        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
        self.dropout3 = tf.keras.layers.Dropout(rate)
        
    def call(self, x, enc_output, training, look_ahead_mask, padding_mask):
        # Capa de atención máscara
        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)
        attn1 = self.dropout1(attn1, training=training)
        out1 = self.layernorm1(x + attn1)
        
        # Capa de atención con salida del codificador
        attn2, attn_weights_block2 = self.mha2(out1, enc_output, enc_output, padding_mask)
        attn2 = self.dropout2(attn2, training=training)
        out2 = self.layernorm2(out1 + attn2)
        
        # Capa feed-forward
        ffn_output = self.ffn(out2)
        ffn_output = self.dropout3(ffn_output, training=training)
        out3 = self.layernorm3(out2 + ffn_output)
        
        return out3, attn_weights_block1, attn_weights_block2

El decodificador integra la atención hacia la entrada del codificador, permitiendo que cada paso de la generación de salida considere la información de la secuencia de entrada.

Finalmente, se pueden construir las clases completas del codificador y decodificador agrupando múltiples capas:

class Encoder(tf.keras.layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, rate=0.1):
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        
        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
        self.pos_encoding = positional_encoding(1000, d_model)
        
        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate)
                           for _ in range(num_layers)]
        
        self.dropout = tf.keras.layers.Dropout(rate)
        
    def call(self, x, training, mask):
        seq_len = tf.shape(x)[1]
        
        # Embedding y codificación posicional
        x = self.embedding(x)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]
        
        x = self.dropout(x, training=training)
        
        # Aplicación de las capas del codificador
        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)
        
        return x  # Salida final del codificador

La modularidad de esta arquitectura facilita la adaptación y extensión del Transformer para diversas tareas. La implementación en TensorFlow Keras permite aprovechar las capacidades de aceleración por GPU y TPU, esenciales para entrenar modelos de gran escala en aprendizaje profundo.

La arquitectura del Transformer representa un avance significativo en el procesamiento de secuencias, combinando mecanismos de atención eficientes con principios de ingeniería de redes profundas. Su diseño ha sentado las bases para modelos avanzados como BERT, GPT-3 y otros, que han alcanzado resultados sobresalientes en múltiples tareas.

Aplicaciones de los Transformers en NLP y más allá

Los Transformers han transformado significativamente el campo del procesamiento del lenguaje natural (PLN), estableciendo nuevos estándares en diversas tareas. Una de las aplicaciones más destacadas es en la traducción automática. Modelos como el Transformer de Vaswani et al. han demostrado superar a las arquitecturas tradicionales, proporcionando traducciones más precisas y fluidas. Esto se debe a la capacidad de los Transformers para capturar dependencias a largo plazo en secuencias de texto, gracias al mecanismo de atención.

Otra área donde los Transformers han tenido un impacto notable es en la generación de lenguaje natural. Modelos como GPT-3 y GPT-4 utilizan arquitecturas basadas en Transformers para generar texto coherente y contextualmente relevante. Estos modelos han sido entrenados en grandes conjuntos de datos y pueden realizar tareas como completar frases, responder preguntas y incluso escribir código.

En la clasificación de texto, los Transformers permiten analizar y categorizar documentos con alta precisión. Modelos como BERT (Bidirectional Encoder Representations from Transformers) se han utilizado para tareas de clasificación de sentimientos, detección de spam y análisis de temas. La naturaleza bidireccional de BERT permite al modelo comprender el contexto tanto anterior como posterior a una palabra, mejorando la comprensión contextual.

Implementar estos modelos en TensorFlow es accesible gracias a la integración de Keras y las herramientas proporcionadas por la biblioteca. Por ejemplo, para utilizar un modelo preentrenado de BERT para clasificación de texto:

import tensorflow as tf
from transformers import TFBertForSequenceClassification, BertTokenizer

# Carga del tokenizador y modelo preentrenado
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased')

# Preparación de los datos de entrada
texts = ["Este curso es excelente.", "No me gusta este producto."]
inputs = tokenizer(texts, return_tensors='tf', padding=True, truncation=True)

# Inferencia
outputs = model(inputs)
logits = outputs.logits

# Predicciones
predictions = tf.nn.softmax(logits, axis=-1)

En este ejemplo, se emplea el modelo TFBertForSequenceClassification para realizar una clasificación binaria de secuencias de texto. La librería transformers facilita la integración con modelos preentrenados, permitiendo un afinamiento fino para tareas específicas.

Más allá del PLN, los Transformers han sido adaptados para la visión por computadora mediante los Vision Transformers (ViT). Estos modelos tratan imágenes como una secuencia de parches, aplicando el mismo mecanismo de atención utilizado en texto. Los ViT han demostrado rendir de forma competitiva en tareas de clasificación de imágenes, a veces superando a las redes convolucionales tradicionales.

Para implementar un Vision Transformer en TensorFlow:

import tensorflow as tf

# Definición de la función para dividir la imagen en parches
def extract_patches(images, size, stride):
    patches = tf.image.extract_patches(
        images=images,
        sizes=[1, size, size, 1],
        strides=[1, stride, stride, 1],
        rates=[1, 1, 1, 1],
        padding='VALID'
    )
    return patches

# Supongamos que tenemos un lote de imágenes de entrada
images = tf.random.normal([batch_size, height, width, channels])

# Extracción de parches
patch_size = 16
patches = extract_patches(images, patch_size, patch_size)

# Reorganización de los parches para preparar la entrada al Transformer
num_patches = (height // patch_size) * (width // patch_size)
patch_dim = patch_size * patch_size * channels
patches = tf.reshape(patches, [batch_size, num_patches, patch_dim])

# Añadir codificaciones posicionales y continuar con el modelo Transformer

En este fragmento, las imágenes se dividen en parches que son aplanados y tratados como tokens en una secuencia, replicando el enfoque de los Transformers en texto. Esto permite al modelo aprender relaciones espaciales entre diferentes partes de la imagen.

Los Transformers también han sido aplicados en series temporales y datos secuenciales en campos como la meteorología, finanzas y control de sistemas. Su capacidad para atender a diferentes partes de una secuencia facilita el pronóstico y la detección de anomalías en datos temporales. Por ejemplo, en el análisis financiero, los Transformers pueden mejorar la predicción de tendencias del mercado al captar patrones complejos.

En bioinformática, los Transformers han sido empleados para analizar secuencias de aminoácidos y ácidos nucleicos. Modelos como AlphaFold han utilizado principios similares para predecir la estructura tridimensional de las proteínas, aportando avances significativos en la comprensión de procesos biológicos.

Otra área emergente es la aplicación de Transformers en sistemas de recomendación. Al tratar las interacciones de los usuarios como secuencias, los modelos pueden proporcionar recomendaciones más personalizadas y contextuales. Esto mejora la experiencia del usuario al ofrecer contenido relevante basado en su historial de interacciones.

Además, los Transformers han encontrado usos en el campo del audio. En tareas como el reconocimiento de voz y la transformación de texto a voz, los modelos basados en Transformers han mejorado la calidad y precisión, permitiendo una mejor transcripción y síntesis de voz más natural.

En cuanto a la generación de contenido multimodal, los Transformers han sido clave en modelos que combinan texto e imágenes, como DALL·E. Estos modelos pueden generar imágenes a partir de descripciones textuales, abriendo nuevas posibilidades en diseño, arte y publicidad.

Es importante destacar que, aunque los Transformers ofrecen ventajas significativas, también requieren una cantidad considerable de recursos computacionales para entrenar. Sin embargo, con herramientas como TensorFlow y el soporte para entrenamiento distribuido en GPU y TPU, es posible manejar modelos grandes y complejos de manera más eficiente.