1. Descripción

En este tutorial, usaremos PyTorch para construir un modelo de transformador básico desde cero. El modelo Transformer presentado por Vaswani y otros en el artículo "Todo lo que necesita es atención" es una arquitectura de aprendizaje profundo diseñada para tareas de secuencia a secuencia, como la traducción automática y el resumen de texto. Se basa en el mecanismo de autoatención, que ha sido la base de muchos modelos de procesamiento de lenguaje natural de última generación, como GPT y BERT.

2. Actividades preparatorias

Para generar el modelo convertidor, seguiremos estos pasos:

Importar bibliotecas y módulos necesarios
Definición de bloques de construcción básicos: atención de múltiples cabezas, redes de retroalimentación posicional, codificación posicional
Construir las capas de codificador y decodificador
Combine capas de codificador y decodificador para crear un modelo de convertidor completo
Preparar datos de muestra
modelo de entrenamiento

Comencemos importando las bibliotecas y módulos necesarios.

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import math
import copy

Ahora definiremos los componentes básicos de un modelo convertidor.

3. Atención múltiple

Figura 2. Atención de múltiples cabezas (Fuente: Imagen creada por el autor)

El mecanismo de atención multicabezal calcula la atención entre cada par de posiciones en la secuencia. Consiste en múltiples "cabezas de atención" para capturar diferentes aspectos de la secuencia de entrada.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_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.W_o = nn.Linear(d_model, d_model)
        
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
        attn_probs = torch.softmax(attn_scores, dim=-1)
        output = torch.matmul(attn_probs, V)
        return output
        
    def split_heads(self, x):
        batch_size, seq_length, d_model = x.size()
        return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
        
    def combine_heads(self, x):
        batch_size, _, seq_length, d_k = x.size()
        return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
        
    def forward(self, Q, K, V, mask=None):
        Q = self.split_heads(self.W_q(Q))
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))
        
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
        output = self.W_o(self.combine_heads(attn_output))
        return output

El código MultiHeadAttention inicializa el módulo con parámetros de entrada y una capa de transformación lineal. Calcula las puntuaciones de atención, remodela los tensores de entrada en múltiples cabezas y combina las salidas de atención de todas las cabezas juntas. El enfoque directo calcula la autoatención de varios cabezales, lo que permite que el modelo se centre en algunos aspectos diferentes de la secuencia de entrada.

4. Posicionar la red de feedforward

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

La clase PositionWiseFeedForward amplía el nn de PyTorch. módulo e implementa una red feed-forward de posicionamiento. Esta clase se inicializa con dos capas de transformación lineal y una función de activación de ReLU. El método directo aplica estas transformaciones y funciones de activación secuencialmente para calcular la salida. Este proceso permite que el modelo tenga en cuenta la ubicación de los elementos de entrada al hacer predicciones.

5. Código de posición

Las codificaciones posicionales se utilizan para inyectar la información posicional de cada token en la secuencia de entrada. Utiliza funciones de seno y coseno de diferentes frecuencias para generar códigos de posición.

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length):
        super(PositionalEncoding, self).__init__()
        
        pe = torch.zeros(max_seq_length, d_model)
        position = torch.arange(0, max_seq_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)
        
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

La clase PositionalEncoding se inicializa con los parámetros de entrada d_model y max_seq_length, creando un tensor para almacenar los valores de codificación posicional. Esta clase calcula el seno y el coseno de los índices pares e impares respectivamente según el factor de escala div_term. El método directo calcula la codificación posicional agregando el valor de codificación posicional almacenado al tensor de entrada, lo que permite que el modelo capture la información posicional de la secuencia de entrada.

Ahora, construiremos las capas de codificador y decodificador.

6. Capa de codificador

Figura 3. La parte del codificador de la red del transformador (fuente: la imagen proviene del texto original)

La capa del codificador consta de una capa de atención de varios cabezales, una capa de avance de posición y dos capas de normalización.

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask):
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        return x

La clase se inicializa con parámetros y componentes de entrada, incluido un módulo de atención de varios cabezales, un módulo PositionWiseFeedForward, módulos de normalización de dos capas y una capa de abandono. El método feed-forward calcula la salida de la capa del codificador aplicando autoatención, agregando la salida de atención al tensor de entrada y normalizando el resultado. Luego calcula la salida de alimentación hacia adelante en cuanto a la posición, la combina con la salida de autoatención normalizada y normaliza el resultado final antes de devolver el tensor procesado.

7. Capa decodificadora

Figura 4. La parte del decodificador de la red del transformador (Fuente: imagen del artículo original)

La capa del decodificador consta de dos capas de atención de varios cabezales, una capa de avance de posición y tres capas de normalización.

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, enc_output, src_mask, tgt_mask):
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))
        attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_output))
        return x

La capa del decodificador se inicializa con parámetros y componentes de entrada, como un módulo de atención de varios cabezales para enmascarar la atención propia y cruzada, un módulo PositionWiseFeedForward, un módulo de normalización de tres capas y una capa de abandono.

El método directo calcula la salida de la capa del decodificador realizando los siguientes pasos:

Calcula la salida de autoatención enmascarada y la agrega al tensor de entrada, seguida de la normalización de capa y abandono.
La salida de atención cruzada entre las salidas del decodificador y del codificador se calcula y se suma a la salida de autoatención enmascarada normalizada, seguida de la normalización de caída y capa.
La salida de alimentación hacia adelante en cuanto a la posición se calcula y combina con la salida de atención cruzada normalizada, seguida de la caída y la normalización de la capa.
Devuelve el tensor procesado.

Estas operaciones permiten que el decodificador genere secuencias objetivo a partir de la entrada y la salida del codificador.

Ahora, combinemos las capas de codificador y decodificador para crear el modelo de convertidor completo.

8. Modelo de transformador

Figura 5. La red de transformadores (Fuente: la imagen proviene del texto original)

Combinarlos todos juntos:

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout):
        super(Transformer, self).__init__()
        self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)
        self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_seq_length)

        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])

        self.fc = nn.Linear(d_model, tgt_vocab_size)
        self.dropout = nn.Dropout(dropout)

    def generate_mask(self, src, tgt):
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
        tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
        seq_length = tgt.size(1)
        nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool()
        tgt_mask = tgt_mask & nopeak_mask
        return src_mask, tgt_mask

    def forward(self, src, tgt):
        src_mask, tgt_mask = self.generate_mask(src, tgt)
        src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
        tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))

        enc_output = src_embedded
        for enc_layer in self.encoder_layers:
            enc_output = enc_layer(enc_output, src_mask)

        dec_output = tgt_embedded
        for dec_layer in self.decoder_layers:
            dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)

        output = self.fc(dec_output)
        return output

Las clases combinan módulos previamente definidos para crear modelos convertidores completos. Durante la inicialización, el módulo Transformer establece los parámetros de entrada e inicializa varios componentes, incluidas las capas de incrustación para las secuencias de origen y de destino, los módulos de codificación de posición, los módulos de capa de codificador y capa de decodificador para crear capas apiladas y para proyectar la salida del decodificador La capa lineal y la caída capa.

El método generate_mask crea una máscara binaria para las secuencias de origen y de destino para ignorar los tokens de relleno y evitar que el decodificador procese tokens futuros. El método directo calcula la salida del modelo convertidor a través de los siguientes pasos:

Utilice el método generate_mask para generar máscaras de origen y destino.
Calcule incrustaciones de origen y destino, y aplique codificación posicional y abandono.
Procese la secuencia fuente a través de la capa del codificador, actualizando el tensor enc_output.
Procese la secuencia objetivo a través de la capa del decodificador, use enc_output y mask, y actualice el tensor dec_output.
Aplique una capa de proyección lineal a la salida del decodificador, tomando el logaritmo de salida.

Estos pasos permiten que el modelo Transformer procese una secuencia de entrada y genere una secuencia de salida basada en las funciones combinadas de sus componentes.

9. Preparar datos de muestra

En este ejemplo, crearemos un conjunto de datos de juguetes con fines de demostración. De hecho, utilizará un conjunto de datos más grande, preprocesará el texto y creará un mapa de vocabulario para los idiomas de origen y de destino.

src_vocab_size = 5000
tgt_vocab_size = 5000
d_model = 512
num_heads = 8
num_layers = 6
d_ff = 2048
max_seq_length = 100
dropout = 0.1

transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)

# Generate random sample data
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)

10. Modelo de formación

Ahora, entrenaremos el modelo utilizando los datos de ejemplo. En la práctica, utilizará un conjunto de datos más grande y lo dividirá en conjuntos de entrenamiento y validación.

criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

transformer.train()

for epoch in range(100):
    optimizer.zero_grad()
    output = transformer(src_data, tgt_data[:, :-1])
    loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:].contiguous().view(-1))
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

Podemos construir un transformador simple desde cero en Pytorch usando este enfoque. Todos los modelos de lenguajes grandes se entrenan usando estos bloques codificadores o decodificadores transformadores. Por lo tanto, es importante comprender la red que lo inicia todo. Espero que este artículo ayude a cualquiera que quiera entender LLM en profundidad.

[Transformers 03] Usa Pytorch para comenzar a construir transformadores