Implementación de un Modelo Transformer Basado en el Paper “Attention is All You Need”

En el mundo del aprendizaje profundo, los modelos Transformer han revolucionado la inteligencia artificial al lograr un rendimiento excepcional en tareas de procesamiento de lenguaje natural y posteriormente a otras como la visión por computadora. Veremos la implementación de un modelo Transformer basada en el influyente paper “Attention is All You Need”, que dio origen posteriormente a ChatGPT. Además, es interesante explicar cómo adaptar este modelo para abordar aplicaciones más avanzadas, como el procesamiento de lenguaje natural.

Implementación de la atención escalonada de producto escalado (scaled dot-product attention)

Iniciamos nuestra implementación con la piedra angular del transformer, la capa de atención de producto escalado. Esta capa permite al modelo calcular interacciones entre elementos en una secuencia. En nuestro código, la clase `ScaledDotProductAtt` realiza esta tarea y se basa en multiplicaciones escalares de matrices. A continuación, proporcionamos explicaciones detalladas de los parámetros y las operaciones:

– `query`, `key`, `value`: Representan las consultas, claves y valores utilizados en el cálculo de la atención. Serán las entradas a nuestro modelo.

– `dropout`: Es un parámetro que controla la regularización al aplicar eliminación evitando el sobre entrenamiento (overfitting).

La atención de producto escalado calcula puntuaciones de atención mediante el producto escalar de `query` y `key`, escalado por la raíz cuadrada de la dimensión de las claves (`key`). Luego, aplicamos la función softmax para obtener los pesos de atención. Si se proporciona una máscara, se aplica para evitar que ciertos elementos se consideren en la atención.

# Scaled Dot-Product Attention
class ScaledDotProductAtt(nn.Module):
    def __init__(self, dropout=0.1):
        super(ScaledDotProductAtt, self).__init__()
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        attScores = torch.matmul(query, key.transpose(-2, -1)) / np.sqrt(key.size(-1))
        if mask is not None:
            attScores = attScores.masked_fill(mask == 0, -1e10)

        attention = torch.softmax(attScores, dim=-1)
        attention = self.dropout(attention)
        return torch.matmul(attention, value), attention

Atención multicabeza

El siguiente paso en nuestra implementación es la atención multicabeza. Esta capa permite al modelo enfocarse en diferentes aspectos de la entrada de manera simultánea. La clase `MultiHeadAttention` se encarga de esto, donde:

– `d_model`, `nhead`, `dropout`: Estos parámetros controlan las dimensiones del modelo y la cantidad de cabezas de atención, así como la regularización.

– `d_k`, `d_v`: Representan las dimensiones de las claves y los valores en cada cabeza de atención.

– `linear_q`, `linear_k`, `linear_v`: Son capas lineales que proyectan las consultas, claves y valores en espacios de dimensiones adecuadas.

– `scaledDotProductAttention`: Es una instancia de la capa de atención de producto escalado que utilizamos.

– `linearLayer`: Otra capa lineal para transformar las salidas de las cabezas de atención.

– `dropout`: Controla la regularización.

Esta capa divide la entrada en múltiples cabezas, cada una de las cuales se procesa individualmente. Luego, las salidas se concatenan y se transforman para obtener la salida final.

# Multi-Head Attention
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.1):
        super(MultiHeadAttention, self).__init()
        self.d_model = d_model
        self.nhead = nhead
        self.d_k = d_model // nhead
        self.d_v = d_model // nhead

        self.linear_q = nn.Linear(d_model, d_model)
        self.linear_k = nn.Linear(d_model, d_model)
        self.linear_v = nn.Linear(d_model, d_model)

        self.scaledDotProductAttention = ScaledDotProductAtt(dropout)
        self.linearLayer = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

Codificación posicional

En el modelo Transformer, la codificación posicional se utiliza para que el modelo comprenda la posición de las palabras en la secuencia. La clase `PositionalEncoding` se encarga de esto. Algunos parámetros y operaciones clave son:

– `d_model`, `dropout`, `max_length`: Controlan la dimensión del modelo, la regularización y la longitud máxima de la secuencia.

– `pe`: Es una matriz que contiene la codificación posicional.

– `position`, `div_term`: Variables auxiliares utilizadas en el cálculo de la codificación posicional.

La codificación posicional se suma a las representaciones de entrada, lo que permite al modelo tener en cuenta la posición de las palabras.

# Positional Encoding
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_length=100):
        super(PositionalEncoding, self).__init()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_length, d_model)
        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)
        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)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

Capas de feed-forward y normalización

Las capas de feed-forward y normalización son componentes esenciales en cada capa del modelo:

– `d_model`, `d_mlp`, `dropout`: controlan las dimensiones del modelo, la dimensión de las capas de Feed-Forward y la regularización.

– `linear1`, `linear2`: son capas lineales utilizadas en la capa de Feed-Forward.

– `gamma`, `beta`, `epsilon`: parámetros y constantes utilizadas en la capa de normalización.

La capa de Feed-Forward transforma las representaciones intermedias, mientras que la Normalización mantiene la estabilidad del entrenamiento.

# Feed-Forward Layer
class FeedForward(nn.Module):
    def __init__(self, d_model, d_mlp=1024, dropout=0.1):
        super(FeedForward, self).__init()
        self.linear1 = nn.Linear(d_model, d_mlp)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_mlp, d_model)

# Normalization Layer
class NormalizationLayer(nn.Module):
    def __init__(self, d_model, epsilon=1e-5):
        super(NormalizationLayer, self).__init()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.epsilon = epsilon

Modelo de lenguaje

Finalmente, ensamblamos todas estas piezas en un modelo de lenguaje completo en la clase `LanguageModel`. Aquí están los detalles:

– `d_model`, `nhead`, `nEncoder`, `d_mlp`, `vocab_size`, `dropout`: Controlan las dimensiones del modelo, la cantidad de cabezas de atención, el número de capas del codificador, la dimensión de las capas de Feed-Forward, el tamaño del vocabulario y la regularización.

– `encoderLayer`, `encoderNorm`: Representan una capa del codificador y una capa de normalización.

– `posEncoder`, `inputEmbed`, `outputLayer`: Son la codificación posicional, la capa de embedding de entrada y la capa lineal de salida.

Este modelo toma una secuencia de entrada, aplica la codificación posicional y utiliza la atención multi-cabeza y las capas de Feed-Forward para aprender representaciones significativas. Luego, una capa lineal proyecta las representaciones en las probabilidades de las palabras del vocabulario.

# Language Model
class LanguageModel(nn.Module):
    def __init__(self, d_model, nhead, nEncoder, d_mlp, vocab_size, dropout=0.1):
        super(LanguageModel, self).__init__()
        self.d_model = d_model
        encoderLayer = Encoder(d_model, nhead, d_mlp, dropout)
        encoderNorm = NormalizationLayer(d_model)
        self.encoder = nn.TransformerEncoder(encoderLayer, nEncoder, encoderNorm)
        self.posEncoder = PositionalEncoding(d_model, dropout, max_length)
        self.inputEmbed = nn.Embedding(vocab_size, d_model)
        self.outputLayer = nn.Linear(d_model, vocab_size)

Ampliando la funcionalidad: GPT-3, visión y modelos multi-modales

Una vez que hemos implementado el modelo Transformer básico, podemos considerar cómo ampliar su funcionalidad. En un mundo en constante evolución de la IA, es esencial estar al tanto de las últimas tendencias. A continuación, algunas formas de ampliar tu modelo:

• 1. Incorporar GPT-3 o ChatGPT: Puedes combinar tu implementación con modelos de lenguaje pre-entrenados como GPT-3 o ChatGPT para dotar a tu modelo de habilidades de generación de lenguaje natural de vanguardia.

• 2. Incorporar Visión por Computadora: Para tareas de visión por computadora, agrega una rama de red neuronal convolucional (CNN) a tu modelo para procesar datos de imágenes. Esto habilita aplicaciones que requieren la comprensión de texto y visión.

• 3. Explorar Modelos Multi-Modales: Considera explorar modelos multi-modales que pueden manejar datos de texto, imágenes y otros tipos de entrada en un solo modelo.

Posibles modificaciones para un modelo de lenguaje (LLM)

Un Modelo de Lenguaje (LLM) tiene como objetivo principal predecir la probabilidad de la palabra siguiente en una secuencia de palabras. Para adaptar nuestro modelo actual a un LLM, debemos realizar las siguientes modificaciones:

1. Cambio en la función de pérdida

En lugar de utilizar la función de pérdida de entropía cruzada categórica, que se utiliza en tareas de clasificación, como la traducción automática, necesitamos una función de pérdida que sea adecuada para la tarea de predicción de palabras. Comúnmente, se utiliza la pérdida de entropía cruzada de palabras, que mide la discrepancia entre las distribuciones de probabilidad del modelo y las distribuciones reales de palabras.

# Pérdida para un Modelo de Lenguaje (Cross-Entropy Loss)
criterion = nn.CrossEntropyLoss()

2. Modificación de los datos de entrada y salida

En un LLM, los datos de entrada y salida se organizan de manera ligeramente diferente. La entrada consiste en una secuencia de palabras hasta un punto de corte, y la salida es la palabra siguiente a predecir.

input_seq = tokens[:-1]  # Secuencia de entrada
target_seq = tokens[1:]   # Secuencia de salida (palabra siguiente)

3. Entrenamiento con datos de texto

El modelo se entrena con datos de texto sin procesar en lugar de datos clasificados. Esto significa que el preprocesamiento de los datos, como la tokenización y la construcción del vocabulario, se debe realizar antes del entrenamiento.

4. Evaluación de la generación de texto

Una vez que el modelo está entrenado, se puede utilizar para generar texto autónomamente. La generación de texto implica alimentar una palabra inicial y permitir que el modelo prediga palabras sucesivas de manera autónoma.

# Generación de texto autónoma
def generate_text(model, starting_word, max_length):
    current_word = starting_word
    generated_text = [current_word]
    for _ in range(max_length):
        input_seq = [current_word]
        input_seq = torch.tensor(input_seq)
        output = model(input_seq)
        predicted_word = torch.argmax(output, dim=-1)[-1].item()
        generated_text.append(predicted_word)
        current_word = predicted_word
    return generated_text

Estas modificaciones permiten que el modelo funcione como un Modelo de Lenguaje (LLM) que puede generar texto de manera autónoma y predecir palabras siguientes en una secuencia. La adaptación de un modelo Transformer a un LLM es un paso importante para aplicaciones como la generación de texto creativo, corrección gramatical y completado automático de texto.

En resumen, construir un modelo Transformer desde cero basado en el paper original es un desafío emocionante que nos brinda una base sólida para comprender cómo funcionan estos modelos. Además, ampliar la funcionalidad del modelo nos permite aprovechar todo su potencial en aplicaciones del mundo real.

La IA es una tecnología, que si bien no es nueva, está generando muchísimo hype y  avanza constantemente, estar al tanto de las últimas tendencias y modelos, como los modelos multi-modales, es esencial para mantenerse actualizado en este campo en constante evolución.