Modelos autorregresivos y arquitectura GPT

En este capítulo revisaremos las bases de los modelos de lenguaje actuales. Para esto, comenzaremos revisando la formulación autorregresiva de una secuencia de texto y concluiremos con la implementación de un modelo GPT [1] de juguete en español. El foco estará puesto en realizar implementaciones limpias y sencillas, al mismo tiempo que se formula cada elemento de manera precisa, lo cual muchas veces es pasado por alto en la literatura de este tipo de modelos.

Las librerías que se utilizarán son las siguientes:

import torch
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
from dataclasses import dataclass
import matplotlib.pyplot as plt
import tqdm
import re

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Formulación de un modelo de lenguaje

Tokenización

Una secuencia de tokens es un conjunto ordenado finito, $S=\left(w_1,\dots ,w_T\right)$, donde los elementos $w_1,\dots ,w_T$ son símbolos (e.g., letras o palabras), denominados tokens, que pertenecen a un vocabulario común, ${\mathcal{V}}_0=\left\{a_1,\dots ,a_K\right\}$ (e.g., letras del abecedario o palabras de un idioma). Por lo general, el vocabulario base ${\mathcal{V}}_0$ suele ser extendido a otro vocabulario $\mathcal{V}\supset {\mathcal{V}}_0$ para incluir algunos tokens especiales que son utilizados durante el entrenamiento e inferencia de un modelo generativo para texto. En este capítulo, se considerarán los siguientes tokens especiales:

• Begin of sentence (\mathrm{\text{<\BOS>}}): token especial agregado al comienzo de cada secuencia. Se utilizará como token inicial durante la generación.

• End of sentence (\mathrm{\text{<\EOS>}}): token especial agregado al final de cada secuencia. Se utilizará para saber cuándo detener la generación de texto.

• Unknown (\mathrm{\text{<\UNK>}}): token usado para representar tokens desconocidos (fuera de ${\mathcal{V}}_0$). Permitirá procesar símbolos desconocidos durante el entrenamiento e inferencia.

• Padding (\mathrm{\text{<\PAD>}}): token agregado al final de algunas secuencias para forzar que tengan un largo predefinido. Permitirá realizar el entrenamiento en batches de secuencias.

Por lo tanto, el vocabulario extendido que se utilizará en este capítulo es el siguiente:

Además, se denotará como ${\mathcal{V}}^{\ast }$ al conjunto de todas las secuencias finitas de tokens formadas por elementos del vocabulario $\mathcal{V}$, es decir:

donde la notación ${\mathcal{V}}^T:=\underset{T\text{ veces}}{\underbrace{\mathcal{V}\times \cdots \times \mathcal{V}}}$ representa el conjunto de secuencias de largo $T$ (i.e., ${\mathcal{V}}^T=\left\{S=\left(w_1,\dots ,w_T\right):w_1,\dots ,w_T\in \mathcal{V}\right\}$).

La siguiente clase Tokenizer define un vocabulario extendido, vocabulary, y utiliza el método encode para mapear un string a su respectiva secuencia de tokens (más precisamente, a su secuencia de índices de tokens según el orden del vocabulario). De forma inversa, el método decode transforma una lista de índices en el string de símbolos:

class Tokenizer:

    def __init__(self, corpus_vocabulary: list[str]) -> None:
        self.vocabulary = corpus_vocabulary + ['<BOS>', '<EOS>', '<UNK>', '<PAD>']
        self.vocab_size = len(self.vocabulary)

        self.token_to_id = {token: id for id, token in enumerate(self.vocabulary)}
        self.bos_id = self.token_to_id['<BOS>']
        self.eos_id = self.token_to_id['<EOS>']
        self.unk_id = self.token_to_id['<UNK>']
        self.pad_id = self.token_to_id['<PAD>']

    def encode(self, text: str) -> list[int]:
        seq_tokens = re.findall(r'\d|[^\w\s]|\w+|\s', text)
        seq_ids = [self.token_to_id.get(token, self.unk_id) for token in seq_tokens]
        return seq_ids

    def decode(self, seq_ids: list[int]) -> str:
        seq_tokens = ''.join(self.vocabulary[i] for i in seq_ids)
        return seq_tokens

• Ejemplo de uso.

  vocabulary = ['a', 'b', 'c', 'd', ' ']
   tokenizer = Tokenizer(vocabulary)
  
   print(f'<BOS> ID: {tokenizer.bos_id}.')
   print(f'<EOS> ID: {tokenizer.eos_id}.')
   print(f'<UNK> ID: {tokenizer.unk_id}.')
   print(f'<PAD> ID: {tokenizer.pad_id}.')
  
   token_seq = 'a b c d e'
   id_seq = tokenizer.encode(token_seq)
   print(f'encode("{token_seq}") = {id_seq}.')
   print(f'decode({id_seq}) = "{tokenizer.decode(id_seq)}".')
   

  ID: 5. ID: 6. ID: 7. ID: 8. encode(“a b c d e”) = [0, 4, 1, 4, 2, 4, 3, 4, 7]. decode([0, 4, 1, 4, 2, 4, 3, 4, 7]) = “a b c d “.

Para entrenar un modelo de lenguaje, es necesario contar con un conjunto de entrenamiento $\mathcal{D}=\left\{S^1,\dots ,S^N\right\}\subset {\mathcal{V}}^{\ast }$ formado por secuencias de tokens definidas sobre un mismo vocabulario $\mathcal{V}$. Además, para poder trabajar con batches de secuencias, asumiremos que todas las secuencias de entrenamiento tienen un mismo largo $T\in \mathbb{N}$:

• Las secuencias de largo mayor a $T$ serán truncadas hasta el largo máximo $T$.

• Las secuencias de largo menor a $T$ serán extendidas agregando el token especial \mathrm{\text{<\PAD>}} tantas veces como sea necesario para alcanzar el largo $T$.

En las implementaciones, el parámetro $T$ será representado por la el atributo seq_length.

El vocabulario ${\mathcal{V}}_0$ consistirá en todas las palabras distintas que se encuentran en el corpus contenido en el archivo data.txt, el cual consiste en un conjunto de cuentos en español, separados por un salto de línea. En consecuencia, cada secuencia del conjunto de entrenamiento, $S=\left(w_1,\dots ,w_T\right)\in \mathcal{D}$, consistirá en un cuento obtenido desde el corpus (truncado o extendido hasta un largo $T\in \mathbb{N}$), con los tokens especiales agregados donde corresponda. El dataset original fue descargado desde el repositorio karen-pal/borges y se preprocesó para eliminar todos los tildes y caracteres poco frecuentes. La siguiente clase TextDataset implementa esta lógica:

class TextDataset(Dataset):

    def __init__(self, filename: str, seq_length: int) -> None:

        with open(filename, 'r', encoding='utf-8') as file:
            corpus = file.read()

        corpus_vocabulary = sorted(set(re.findall(r'\d|[^\w\s]|\w+|\s', corpus)))
        self.tokenizer = Tokenizer(corpus_vocabulary)

        sentences = [sentence.strip() for sentence in corpus.split('\n') if sentence.strip()]
        self.data = [[self.tokenizer.bos_id] + self.tokenizer.encode(sentence) + [self.tokenizer.eos_id] for sentence in sentences]

        self.seq_length = seq_length

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, n: int) -> torch.Tensor:

        seq_ids = self.data[n][:]

        if len(seq_ids) > self.seq_length:
            seq_ids = seq_ids[:self.seq_length]
        else:
            seq_ids += [self.tokenizer.pad_id] * (self.seq_length - len(seq_ids))

        return torch.tensor(seq_ids)

Con esta clase se puede inicializar un dataset para el texto contenido en el archivo data.txt:

seq_length, batch_size = 256, 32

dataset = TextDataset('data.txt', seq_length)
dataloader = DataLoader(dataset, batch_size, shuffle=True, drop_last=True)

print(f'Tamaño del dataset: {len(dataset)} secuencias.')
print(f'Tamaño del dataloader: {len(dataloader)} batches.')
print(f'Tamaño del vocabulario: {dataset.tokenizer.vocab_size} tokens.')

Tamaño del dataset: 719 secuencias. Tamaño del dataloader: 22 batches. Tamaño del vocabulario: 60198 tokens.

Un ejemplo de secuencia contenida en este dataset es el siguiente (solo se mostrarán los primeros 40 tokens por espacio):

seq_ids = dataset[0]
seq_tokens = dataset.tokenizer.decode(seq_ids[:40].tolist())
print(seq_tokens)

Los enanos tienen una especie de sexto sentido que les permite reconocerse a primera vista.

Formulación autorregresiva

Dada una secuencia de tokens, $S=\left(w_1,\dots ,w_T\right)\in {\mathcal{V}}^{\ast }$, un modelo autorregresivo $p_{\theta }(S)$ descompone la probabilidad que le asigna a la secuencia $S$ de forma causal, expresando la probabilidad conjunta $p_{\theta }(S)=p_{\theta }\left(w_1,\dots ,w_T\right)$ como el producto de las probabilidades de cada token de $S$ dados los tokens que lo anteceden:

donde $p_{\theta }\left(w_1\right)$ es la distribución que el modelo le asigna al primer token de $S$, mientras que $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)$ es la distribución del $(t+1)$-ésimo token de $S$ dados los $t$ tokens anteriores. Además, dado que el primer token siempre es \mathrm{\text{<\BOS>}} \in \mathcal{V}, se puede fijar p_{\theta}\left( w_{1} \right) = p\left( w_{1} \right) = \delta_{\mathrm{\text{<\BOS>}}}(w_{1}), por lo que solo es necesario diseñar una red neuronal que aprenda (los parámetros de) la distribución condicional $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)$, la cual será utilizada posteriormente para generar nuevas secuencias de tokens. Notar que, a priori, la descomposición causal no realiza ninguna suposición de independencia sobre los tokens que componen la secuencia $S$ (como sí ocurre, por ejemplo, en las cadenas de Markov), aunque sí impone que el orden topológico del DAG asociado a esta red bayesiana sea el orden temporal de las variables.

Por otra parte, dado que cada distribución condicional $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)$ es una distribución de probabilidades sobre el vocabulario $\mathcal{V}$ (el cual es finito), se puede considerar, sin pérdida de generalidad, que $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)\sim \mathrm{Categ}\acute{\mathrm{o}}\mathrm{rica}(\mathcal{V};{\lambda }_{\theta }\left(w_1,\dots ,w_t\right))$ es una distribución categórica sobre $\mathcal{V}$ con vector de probabilidades ${\lambda }_{\theta }\left(w_1,\dots ,w_t\right)\in {\Delta }^{|\mathcal{V}|}$, es decir:

En consecuencia, para entrenar un modelo de lenguaje autorregresivo solo hace falta entrenar un modelo condicional que aprenda el vector de probabilidad ${\lambda }_{\theta }\left(w_1,\dots ,w_t\right)\in {\Delta }^{|\mathcal{V}|}$, los cuales son parametrizados mediante una red neuronal ${\lambda }_{\theta }:{\mathcal{V}}^{\ast }\to [0,1{]}^{|\mathcal{V}|}$, la cual recibe una secuencia de tokens $\left(w_1,\dots ,w_t\right)\in {\mathcal{V}}^{\ast }$ (llamada contexto) y entrega una distribución sobre $\mathcal{V}$ asociada al próximo token en el tiempo $t+1$. Por otro lado, dado que ${\lambda }_{\theta }\left(w_1,\dots ,w_t\right)\in {\Delta }^{|\mathcal{V}|}$ es un vector de probabilidades, se debe que cumplir que ${\sum }_{k=1}^{|\mathcal{V}|}{\lambda }_{\theta }{\left(w_1,\dots ,w_t\right)}_k=1$. Esto último se consigue aplicando la función $\mathrm{softmax}$ al final de la red neuronal.

Inferencia usando una red neuronal recurrente

Vectores de embedding y RNN

Dado que $\mathcal{V}$ es un conjunto abstracto sin ninguna estructura matemática definida, para implementar un modelo de lenguaje utilizando redes neuronales, es necesario obtener una representación continua para cada uno de los tokens del vocabulario discreto $\mathcal{V}$ ya que, de lo contrario, conceptos como derivada o gradiente no estarán bien definidos (al menos en su forma usual). Para esto, es usual asociar cada token $a\in \mathcal{V}$ con un vector $\mathrm{emb}(a)\in {\mathbb{R}}^D$, denominado vector de embedding. Más precisamente, el operador de embedding, $\mathrm{emb}:\mathcal{V}\to {\mathbb{R}}^D$, mapea el $k$-ésimo token del vocabulario $\mathcal{V}$ a la $k$-ésima fila de una matriz $E\in {\mathcal{M}}_{|\mathcal{V}|,D}\left(\mathbb{R}\right)$, denominada matriz de embedding, la cual es aprendida conjuntamente con el resto de los parámetros del modelo neuronal durante el entrenamiento. Notar que el operador de embedding asume que el vocabulario $\mathcal{V}$ tiene un orden fijo (i.e., se debería escribir $\mathcal{V}=\left(a_1,\dots ,a_{|\mathcal{V}|}\right)$ en vez de $\mathcal{V}=\left\{a_1,\dots ,a_{|\mathcal{V}|}\right\}$, pero se seguirá usando la notación de conjunto), lo cual se cumple en la clase Tokenizer, donde cada símbolo tiene asociado un índice, los cuales se almacenan en el atributo token_to_id. En PyTorch, se puede definir una capa de embedding utilizando la clase nn.Embedding, la cual implementa exactamente esta lógica. La dimensión de embedding $D\in \mathbb{N}$ será denotada como embedding_dim, aunque también se le suele llamar hidden dim.

• Ejemplo de uso.

  vocab_size, embedding_dim = dataset.tokenizer.vocab_size, 384
   emb = nn.Embedding(vocab_size, embedding_dim)
  
   x = torch.randint(vocab_size, size=[batch_size, seq_length])
   y = emb(x)
  
   assert y.shape == (batch_size, seq_length, embedding_dim)
   assert emb.weight.shape == (vocab_size, embedding_dim)
   

Por otro lado, aplicar el operador de embedding al $k$-ésimo token del vocabulario, $a_k\in \mathcal{V}$, es equivalente al realizar el producto matricial $E^{\top }{y}_k$, donde $y_k\in {\mathbb{R}}^{|\mathcal{V}|}$ es el vector one-hot asociado al token $a_k$ (i.e., es el $k$-ésimo vector de la base canónica de ${\mathbb{R}}^{|\mathcal{V}|}$). De esta forma, una capa de embedding puede verse como una capa lineal cuya entrada, $a_k\in \mathcal{V}$, es codificada mediante el vector $y_k\in {\mathbb{R}}^{|\mathcal{V}|}$.

Con respecto a la elección de la arquitectura neuronal a utilizar, es importante notar que el contexto $\left(w_1,\dots ,w_t\right)\in {\mathcal{V}}^{\ast }$ (entrada a la red neuronal) es de largo variable, por lo que no es posible utilizar una red fully connected estándar para parametrizar la distribución condicional $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)$ (aunque sí se podría utilizar una red convolucional 1D), por lo que hay que usar otro tipo de arquitectura. Una opción frecuente (al menos hasta antes de la aparición de la arquitectura Transformer) es utilizar una red neuronal recurrente (RNN), la cual procesa el contexto $\left(w_1,\dots ,w_t\right)\in {\mathcal{V}}^{\ast }$ comenzando con un vector inicial $h_0=0\in {\mathbb{R}}^D$ (o cualquier otro valor), y luego procesa iterativamente cada uno de los tokens del contexto $\left(w_1,\dots ,w_t\right)$, desde el primero hasta el último:

donde $f_{\theta }:\mathcal{V}\times {\mathbb{R}}^D\to {\mathbb{R}}^D$ es una red neuronal cuya entrada es un token del vocabulario $\mathcal{V}$ y el vector asociado a la iteración anterior. De este modo, se puede considerar la última iteración, $h_t\in {\mathbb{R}}^D$, para definir ${\lambda }_{\theta }\left(w_1,\dots ,w_t\right)=g_{\theta }\left(h_t\right)\in [0,1{]}^{|\mathcal{V}|}$, donde $g_{\theta }:{\mathbb{R}}^D\to [0,1{]}^{|\mathcal{V}|}$ es una función que transforma el último vector $h_t\in {\mathbb{R}}^D$ en un vector de probabilidades sobre $\mathcal{V}$ (a los bloques finales que cumplen esta función se les suele llamar cabezales de lenguaje).

Un ejemplo simple (que será usado como baseline) es la red de Elman clásica [2], la cual considera

• $f_{\theta }\left(w_s,h_{s-1}\right)=\tanh (A\mathrm{emb}(w_s)+B{h}_{s-1}+d)$.

• $g_{\theta }\left(h_t\right)=\mathrm{softmax}(C{h}_t+e)$.

Notar que se abusó de notación en $f_{\theta }\left(w_s,h_{s-1}\right)$, ya que la función de activación $\tanh :\mathbb{R}\to [-1,1]$ es aplicada coordenada a coordenada.

De esta forma, los parámetros entrenables del modelo son la capa de embedding $\mathrm{emb}$ (la cual aprende una matriz $E\in {\mathcal{M}}_{|\mathcal{V}|,D}\left(\mathbb{R}\right)$), las matrices $A,B\in {\mathcal{M}}_{D,D}\left(\mathbb{R}\right)$, $C\in {\mathcal{M}}_{|\mathcal{V}|,D}\left(\mathbb{R}\right)$, y los vectores de sesgo $d\in {\mathbb{R}}^D$, $e\in {\mathbb{R}}^{|\mathcal{V}|}$. Para efectos de implementación, no se agregará la capa $\mathrm{softmax}$ en la salida del modelo. Es decir, el modelo no retornará directamente un vector de probabilidad, si no que retornará un vector no normalizado ($C{h}_t+e\in {\mathbb{R}}^{|\mathcal{V}|}$ en este caso) denominado vector de logits. Esto permitirá calcular la función de pérdida (log-verosimilitud) de forma más eficiente ya que, como se verá en la sección de entrenamiento, la función logaritmo deshace la normalización realizada por la función $\mathrm{softmax}$.

class GenerativeRNN(nn.Module):

    def __init__(self, vocab_size: int, embedding_dim: int) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.rnn = nn.RNN(embedding_dim, embedding_dim, batch_first=True)
        self.lm_head = nn.Linear(embedding_dim, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_embedding = self.embedding(x.long())
        rnn_output, _ = self.rnn(seq_embedding)
        logits = self.lm_head(rnn_output)
        return logits

• Ejemplo de uso.

  vocab_size, embedding_dim = dataset.tokenizer.vocab_size, 384
   rnn = GenerativeRNN(vocab_size, embedding_dim)
  
   x = torch.randint(vocab_size, size=[batch_size, seq_length])
   y = rnn(x)
  
   assert y.shape == (batch_size, seq_length, vocab_size)
   

En este capítulo, se utilizará una red recurrente con $D=384$:

rnn = GenerativeRNN(dataset.tokenizer.vocab_size, embedding_dim=384)

n_params = sum(param.numel() for param in rnn.parameters()) / 1e6
print(f'Cantidad de parámetros: {n_params:.3} millones.')

Cantidad de parámetros: 46.6 millones.

Entrenamiento

Para entrenar una red neuronal ${\lambda }_{\theta }:{\mathcal{V}}^{\ast }\to [0,1{]}^{|\mathcal{V}|}$ que aprenda el vector de probabilidades de la distribución $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)\sim \mathrm{Categ}\acute{\mathrm{o}}\mathrm{rica}(\mathcal{V};{\lambda }_{\theta }\left(w_1,\dots ,w_t\right))$, el enfoque de optimización natural es el de máxima verosimilitud, el cual consiste en maximizar la cantidad $\log p_{\theta }\left(\mathcal{D}\right)$ definida sobre el conjunto de entrenamiento i.i.d. $\mathcal{D}=\left\{S^1,\dots ,S^N\right\}\subset {\mathcal{V}}^{\ast }$. Sin embargo, considerando que $\log p_{\theta }\left(\mathcal{D}\right)\in (-\infty ,0]$, es más usual formular este problema de manera equivalente como la minimización de $-\log p_{\theta }\left(\mathcal{D}\right)\ge 0$, es decir, la red neuronal es optimizada optimizando

De acuerdo a la factorización causal que realizan los modelos autorregresivos, para una secuencia $S=\left(w_1,\dots ,w_T\right)\in {\mathcal{V}}^{\ast }$, su log-verosimilitud se puede escribir como la suma de las log-probabilidades que el modelo le asigna a cada token de la secuencia, dados los tokens anteriores como contexto:

donde en la última igualdad, $\log p\left(w_1\right)=0$ ya que p\left( w_{1} \right) = \delta_{\mathrm{\text{<\BOS>}}}(w_{1}) = 1 (todas las secuencias en $\mathcal{D}$ comienzan con w_{1} = \mathrm{\text{<\BOS>}}). Sin embargo, una práctica usual es optimizar la log-verosimilitud normalizada por dos cantidades que influyen directamente en la función de costo $-\log p_{\theta }\left(\mathcal{D}\right)$:

• Por el tamaño del batch $\mathcal{D}$: esto permite que la cantidad $\log p_{\theta }\left(\mathcal{D}\right)$ sea independiente del tamaño del batch, por lo que no hay que ajustar la tasa de aprendizaje según esta cantidad.

• Por el largo de cada secuencia $S\in \mathcal{D}$: esto cambia levemente el problema para optimizar la log-verosimilitud promedio por token (en vez de la log-verosimilitud de las secuencias), lo cual es importante ya que $\log p_{\theta }\left(S^1\right)\ge \log p_{\theta }\left(S^2\right)$ si $|S^1|\le |S^2|$ (i.e., sin normalizar, el modelo le da mayor verosimilitud a secuencias cortas).

En resumen, si se cuenta con un dataset de entrenamiento $\mathcal{D}=\left\{S^1,\dots ,S^N\right\}\subset {\mathcal{V}}^{\ast }$, donde cada secuencia $S^n=\left(w_1^{(n)},\dots ,w_{T_n}^{(n)}\right)\in \mathcal{D}$ es de largo $T_n\in \mathbb{N}$, el problema de optimización estándar para el entrenamiento de un LLM es el siguiente:

Si bien en esta implementación todos los batches tendrán el mismo tamaño y todas las secuencias tendrán el mismo largo (gracias al truncamiento y padding), en entrenamientos más eficientes no siempre es así. Por ejemplo, para ahorrar computación y memoria durante el entrenamiento, se pueden formar batches de secuencias de largos similares para minimizar el uso de truncamiento y padding. Del mismo modo, los batches también pueden ser de diferente tamaño, sobretodo cuando se indica drop_last=False al instanciar el DataLoader, donde el último batch puede ser de un tamaño menor que el resto debido a la falta de más muestras en el dataset.

Por otra parte, es importante notar que, para una secuencia $S=\left(w_1,\dots ,w_T\right)\in \mathcal{D}$, esta función de pérdida no necesita conocer todo el vector de probabilidades ${\lambda }_{\theta }\left(w_1,\dots ,w_t\right)\in {\Delta }^{|\mathcal{V}|}$ asociado a la distribución $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)$, si no que es suficiente conocer su valor en la coordenada correspondiente al token $w_{t+1}\in \mathcal{V}$ correcto. En consecuencia, es suficiente implementar el cálculo del logaritmo y $\mathrm{softmax}$ del vector de logits solo para la coordenada que se necesita (y no para toda la salida de la red neuronal), lo cual resulta computacionalmente más eficiente.

La clase nn.CrossEntropyLoss de PyTorch permite calcular eficientemente esta función objetivo cuando la red neuronal está implementada para retornar el vector de logits (i.e., cuando no se aplica la función $\mathrm{softmax}$ en la salida de la red neuronal). En efecto, si logits representa la salida de la red neuronal para un batch de secuencias dado como input y targets es un tensor con los índices de los tokens que debería predecir el modelo, el valor obtenido con nn.CrossEntropyLoss equivale a la función de costo anterior:

batch_size, seq_length, vocab_size = 32, 256, dataset.tokenizer.vocab_size

logits = torch.randn(batch_size, seq_length, vocab_size)
targets = torch.randint(0, vocab_size, size=(batch_size, seq_length))

# Cálculo de la función de pérdida paso a paso:
probs = logits.softmax(dim=-1)
target_probs = probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)

log_likelihood = target_probs.log().sum(dim=-1)
normalized_log_likelihood = log_likelihood / seq_length
manual_loss = - normalized_log_likelihood.mean()

# Cálculo de la función de pérdida con CrossEntropyLoss:
loss_fn = nn.CrossEntropyLoss()
logits_reshaped = logits.view(batch_size * seq_length, vocab_size)
targets_reshaped = targets.view(batch_size * seq_length)
direct_loss = loss_fn(logits_reshaped, targets_reshaped)

assert torch.isclose(manual_loss, direct_loss)

Un último detalle importante al momento de optimizar este tipo de modelos es con respecto a los tokens de padding, \mathrm{\text{<\PAD>}} \in \mathcal{V}. Dado que este token solo se incluye en las secuencias para extenderlas hasta un cierto largo predefinido $T\in \mathbb{N}$, su predicción no debe incluirse en el cálculo de la función de pérdida ya que no forman parte de la secuencia original, ni se busca aprender la dinámica de extender las secuencias con tokens de padding luego de generar el token final, \mathrm{\text{<\EOS>}} \in \mathcal{V}. Esta omisión se consigue incluyendo el argumento ignore_index=pad_id al instanciar la clase nn.CrossEntropyLoss.

Con todo esto, el loop de entrenamiento de un modelo de lenguaje es el siguiente:

def train_model(
        model: nn.Module,
        optimizer: optim.Optimizer,
        dataloader: DataLoader,
        epochs: int,
        ckpt_filename: str) -> None:

    model.to(DEVICE)
    model.train()

    pad_id = dataloader.dataset.tokenizer.pad_id
    loss_fn = nn.CrossEntropyLoss(ignore_index=pad_id)

    training = {'losses': [], 'model': None}

    try:
        progressbar = tqdm.trange(epochs)
        for epoch in progressbar:

            for seq_batch in dataloader:

                seq_batch = seq_batch.to(DEVICE)
                x_batch, y_batch = seq_batch[:, :-1], seq_batch[:, 1:]

                logits = model(x_batch)

                batch_size, seq_length, vocab_size = logits.shape
                logits = logits.reshape(batch_size * seq_length, vocab_size)
                y_batch = y_batch.reshape(batch_size * seq_length)

                loss = loss_fn(logits, y_batch)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                training['losses'].append(loss.item())
                progressbar.set_postfix(loss=loss.item())

    except KeyboardInterrupt:
        print('Entrenamiento interrumpido.')

    training['model'] = model.state_dict()
    torch.save(training, ckpt_filename)

Se entrenará la RNN definida anteriormente durante 32 épocas utilizando AdamW [3], el cual es un optimizador usual para LLMs (aunque hoy en día también se utilizan otras opciones más modernas):

rnn_optimizer = optim.AdamW(rnn.parameters())
train_model(rnn, rnn_optimizer, dataloader, epochs=32, ckpt_filename='rnn_training.pt')

El entrenamiento toma aproximadamente 5 minutos en Google Colab (usando una GPU NVIDIA Tesla T4). Una vez concluido el entrenamiento, se puede cargar el modelo y ver la evolución de la función de costo a lo largo del entrenamiento:

rnn_training = torch.load('rnn_training.pt', DEVICE, weights_only=True)
rnn.load_state_dict(rnn_training['model'])

plt.figure(figsize=(10, 5))
plt.plot(rnn_training['losses'])
plt.xlabel('Iteración')
plt.ylabel('Entropía cruzada media')
plt.grid(alpha=0.3)
plt.show()

Notar que, en un comienzo, la entropía cruzada es cercana a $\log (|\mathcal{V}|)\approx 11$ nats, que es lo máximo que puede valer cuando el modelo asigna probabilidades uniformes a cada token.

Generación

Para la generación de nuevas secuencias de tokens, se puede utilizar el algoritmo de ancestral sampling revisado en la entrada de redes bayesianas. En el caso del enfoque autorregresivo, la generación comienza con el token inicial w_{1} = \mathrm{\text{<\BOS>}} (recordar que se fijó p\left( w_{1} \right) = \delta_{\mathrm{\text{<\BOS>}}}(w_{1})) y luego se debe evaluar iterativamente el modelo neuronal entrenado para predecir el siguiente token, utilizando como contexto la secuencia de tokens que ya han sido generados. La generación de nuevos tokens continúa hasta que se genera el token \mathrm{\text{<\EOS>}} o bien, hasta que se genera una cantidad máxima de tokens. Del mismo modo, se podría realizar generación condicional comenzando con una secuencia de tokens iniciales (llamada prompt), $S_{\text{prompt }}=\left(w_1,\dots ,w_t\right)$, lo cual es una tarea usual en aplicaciones como chatbots o agentes.

Dado que el modelo neuronal retorna (los logits de) la distribución $p_{\theta }\left(w_{t+1}|w_1,\dots ,w_t\right)\sim \mathrm{Categ}\acute{\mathrm{o}}\mathrm{rica}(\mathcal{V};{\lambda }_{\theta }\left(w_1,\dots ,w_t\right))$ asociada al siguiente token dado el contexto actual, se debe elegir un criterio para decidir cuál va a ser el próximo token a partir de la distribución entregada por el modelo. Si $S=\left(w_1,\dots ,w_t\right)\in {\mathcal{V}}^{\ast }$ son los tokens de entrada (contexto) y ${\lambda }_{\theta }(S)\in {\Delta }^{|\mathcal{V}|}$ es el vector de probabilidades (luego de aplicar $\mathrm{softmax}$) que entrega el modelo neuronal para el próximo token (i.e., $p_{\theta }\left(w_{t+1}=a_k|S\right)={\lambda }_{\theta }(S{)}_k$, para todo $k\in \left\{1,\dots ,|\mathcal{V}|\right\}$), una primera idea podría ser elegir el siguiente token $w_{t+1}\in \mathcal{V}$ como el token más probable según el modelo, es decir:

Una limitación de este enfoque, denominado greedy decoding, es que la elección del token más probable no es necesariamente una decisión óptima ya que, si bien la probabilidad de este token es alta, puede ocurrir que todas las futuras probabilidades sean bajas producto de haber elegido el token más probable en el momento (de ahí viene el adjetivo greedy). Para permitir más variabilidad en la generación, se puede elegir el siguiente token generando una muestra desde la distribución categórica $\mathrm{Categ}\acute{\mathrm{o}}\mathrm{rica}(\mathcal{V};{\lambda }_{\theta }(S))$ (e.g., utilizando torch.multinomial) y definir $w_{t+1}$ como la muestra generada. Esto garantiza que el modelo genere siempre secuencias nuevas, aunque reciba la misma secuencia inicial. Sin embargo, dado que ${\lambda }_{\theta }(S{)}_k>0$ para todo $k\in \left\{1,\dots ,|\mathcal{V}|\right\}$ (ya que este vector se determina usando $\mathrm{softmax}$ y $e^l>0,\forall l\in \mathbb{R}$), el sampling desde la distribución categórica usando ${\lambda }_{\theta }(S)$ puede generar, eventualmente, tokens muy poco probables, lo que muchas veces se traduce en secuencias de texto sin sentido debido a la propagación del error de generación. Una solución inmediata a este problema es permitir elegir solamente los $K\ge 1$ tokens más probables, redistribuyendo la masa de los $|\mathcal{V}|-K$ tokens restantes: si $\mathcal{I}=\left\{k_1,\dots ,k_K\right\}\subset \left\{1,\dots ,|\mathcal{V}|\right\}$ son las $K$ coordenadas donde ${\lambda }_{\theta }(S)\in {\Delta }^{|\mathcal{V}|}$ toma los valores más altos, se puede construir el vector de probabilidad truncado, ${\tilde{\lambda }}_{\theta }(S)\in {\Delta }^{|\mathcal{V}|}$, cuya $k$-esima coordenada es

Luego, la técnica de top-$K$ sampling elige el siguiente token generando una muestra desde la distribución $\mathrm{Categ}\acute{\mathrm{o}}\mathrm{rica}(\mathcal{V};{\tilde{\lambda }}_{\theta }(S))$. Notar que si $K=1$ se recupera la generación greedy, mientras que si $K=|\mathcal{V}|$ se recupera el sampling categórico usando el vector de probabilidades original, ${\lambda }_{\theta }(S)\in {\Delta }^{|\mathcal{V}|}$. Por otro lado, es importante mencionar que el escalar $K\in \mathbb{N}$ definido aquí no tiene relación con el escalar $K\in \mathbb{N}$ definido al comienzo del capítulo (usado para denotar la cardinalidad del vocabulario base, ${\mathcal{V}}_0=\left\{a_1,\dots ,a_K\right\}$), solo es un alcance de nombre.

Otra heurística comúnmente utilizada corresponde a aplicar la idea de temperatura, la cual modifica el vector de probabilidades ${\lambda }_{\theta }(S)$ para volverlo más determinista (baja temperatura) o más uniforme (alta temperatura). Más precisamente, si $l\in {\mathbb{R}}^{|\mathcal{V}|}$ corresponde al vector de logits entregado por la red neuronal (i.e., el vector previo a aplicar $\mathrm{softmax}$) y $T>0$ es un escalar, denominado temperatura, se puede considerar una versión escalada del vector de logits, $\tilde{l}=\frac{l}{T}\in {\mathbb{R}}^{|\mathcal{V}|}$, y aplicar $\mathrm{softmax}$ a esta nueva cantidad para obtener el vector de probabilidades desde el que se realiza el sampling del próximo token (ya sea usando todo el vector o usando top-$K$ sampling). El objetivo de esta ponderación es llevar todos los logits a una escala similar (si $T>1$) o destacar aún más el logit del token más probable (si $T<1$). Notar que el escalar de temperatura $T\in {\mathbb{R}}_{++}$ definido aquí no tiene relación con el escalar $T\in \mathbb{N}$ usado para definir el largo de las secuencias, solo es un alcance de nombre.

En la siguiente figura se observa cómo cambia el nuevo vector de probabilidades, $\mathrm{softmax}(\tilde{l})\in {\Delta }^{|\mathcal{V}|}$, para distintos valores de temperatura, donde $T=1$ corresponde al vector de probabilidades ${\lambda }_{\theta }(S)$ que se obtendría sin aplicar esta técnica:

• Código para generar la figura.

  vocab_size = 10
   vocabulary = [f'$a_{{{k}}}$' for k in range(1, vocab_size + 1)]
   logits = torch.randn(vocab_size)
   temperatures = [0.1, 0.5, 1, 2, 5, 10]
  
   fig, axes = plt.subplots(2, 3, figsize=(10, 4))
   axes = axes.flatten()
  
   for i, t in enumerate(temperatures):
       scaled_logits = logits / t
       prob = scaled_logits.softmax(-1)
       axes[i].bar(vocabulary, prob)
       axes[i].set_title(f'$T={t}$')
       axes[i].set_ylabel(f'Probabilidad')
  
   plt.tight_layout()
   plt.show()
   

Notar que si $T\to 0$ se recupera la generación greedy, mientras que si $T\to \infty$ el vector de probabilidades se vuelve uniforme (con una probabilidad $\frac{1}{|\mathcal{V}|}$ para cada token). Por lo tanto, el parámetro de temperatura $T>0$ permite controlar el trade-off entre versatilidad y confianza en la generación, lo cual es una herramienta muy útil en algunos contextos. Por ejemplo, para la generación de código es usual elegir una baja temperatura, mientras que para la generación de textos narrativos, puede ser preferible utilizar una temperatura más elevada.

Por último, otra heurística usual es penalizar la probabilidad de tokens que ya han sido generados o que estén en la secuencia inicial. Esta técnica se aplica debido a que es común que los modelos de lenguaje basados en redes neuronales entren en ciclos de repetición de un mismo token o secuencias de tokens. Si $c\left(a_k\right)\in \mathbb{N}$ es la cantidad de veces que el símbolo $a_k\in \mathcal{V}$ ha aparecido en la secuencia generada hasta el momento, se pueden escalar los logits $l\in {\mathbb{R}}^{|\mathcal{V}|}$ mediante

donde el hiperparámetro $\gamma \ge 1$ controla el nivel de penalización. Más específicamente, si $\gamma >1$, entonces los logits de tokens repetidos decrecen, reduciendo su probabilidad tras aplicar $\mathrm{softmax}$.

Todas estas técnicas pueden ser incluidas dentro de un mismo loop de generación:

def generate_tokens(
        model: nn.Module,
        context: str,
        tokenizer: Tokenizer,
        temperature: float = 0.8,
        top_k: int = 50,
        max_tokens: int = 256,
        repetition_penalty: float = 1.01) -> str:

    model.to(DEVICE)
    model.eval()

    seq_id = [tokenizer.bos_id] + tokenizer.encode(context)
    seq_id = torch.tensor(seq_id, device=DEVICE)

    with torch.no_grad():
        for _ in range(max_tokens):
            logits = model(seq_id.unsqueeze(0))[0, -1, :]

            token_counts = torch.bincount(seq_id, minlength=logits.size(0))
            for token_id, count in enumerate(token_counts):
                if count > 0:
                    if logits[token_id] > 0:
                        logits[token_id] /= (repetition_penalty ** count)
                    else:
                        logits[token_id] *= (repetition_penalty ** count)

            if temperature == 0:
                next_token = torch.argmax(logits, dim=-1, keepdim=True)
            else:
                logits = logits / temperature

                top_k_logits, top_k_indices = torch.topk(logits, top_k)
                probs = top_k_logits.softmax(dim=-1)
                next_token = top_k_indices[torch.multinomial(probs, num_samples=1)]

            seq_id = torch.cat((seq_id, next_token), dim=0)

            if next_token in (tokenizer.eos_id, tokenizer.pad_id):
                break

    return tokenizer.decode(seq_id.tolist())

Notar que la técnica de penalización por repetición se aplica antes de aplicar la de temperatura. Por otro lado, la implementación realizada no es muy eficiente: hay cálculos que podrían reutilizarse (p.g. usando programación dinámica), y en la parte de penalización por repetición, el loop con enumerate(token_counts) itera sobre todo el vocabulario (se podría hacer más eficiente operando solo sobre los tokens que efectivamente aparecen). En la entrada de LLMs y agentes se revisarán algunas heurísticas más avanzadas para un entrenamiento e inferencia más eficiente.

Usando la función anterior y el modelo neuronal ya entrenado, se puede realizar generación condicional continuando una secuencia de texto dada:

context = 'Habia una vez'

for n in range(10):
    new_tokens = generate_tokens(rnn, context, dataset.tokenizer)
    print(new_tokens)

Habia una vez por la luz del otro, de la puerta y de los detalles de los ojos, y el que no se acerco el caso. Habia una vez que me dijo que habia que en los ojos del mundo, y de la hora. Y es de otro. Habia una vez que como el dia los ojos de la tarde y su vida. Ya son menos. Pero, se sento en que, que lo pocos de los cuatro. Como el brazo, de los dedos, el sol, las cosas. En la ultima hora, habia? Habia una vez una manera de las yerbas que habia sido tan dias, y a la puerta, y el cuarto. ¿Que dio a el espacio de la mañana. Habia una vez un rio para el muro y la luz. No puede la mañana, al instante… ¿Es siempre por la costa. Sus noche, se llamaba a la tierra, sin la pequeña spicula…… Habia una vez que solo como a los dos años, no se ve sus amigos, es la casa de los dos del sol. Todo el mundo, despues, no, las pesadas mujer, ¡Era con una mesa,……. Habia una vez que no se me dijo nunca, de la puerta que en la muerte, como que? ¡No….. Habia una vez ese que no se ha que se encamino a quien estaba en el instante. Las calles y a la ciudad. Habia una vez por la costa del hombre del brazo, aunque no hace que su lado que se le habia veces se trata. Habia una vez su luz del corredor de los hombres para la puerta, y se no de las manos, pues, la calle. Para entre la puerta, a la cabeza, que?

En la generación anterior se observa que, si bien el modelo fue capaz de aprender a separar cada palabra por un espacio, el texto que genera es incoherente. Si bien es posible obtener mejores resultados cambiando algunos hiperparámetros de generación (o usando una RNN más profunda), en la segunda parte de este capítulo se estudiará la arquitectura GPT, la cual define el esqueleto común de prácticamente todos los LLMs modernos.

Otras heurísticas de generación

Pendiente: beam search, nucleus sampling, min-$p$ sampling, penalización por $n$-gramas, sampling especulativo, contrastive search.

Arquitectura GPT

A pesar de que las RNNs permiten procesar secuencias de largo variable, tienen limitaciones importantes que no permiten extenderlas fácilmente a escenarios más complejos como los que enfrentan los LLMs de hoy en día. Por ejemplo, omitiendo algunos componentes por simplicidad, una RNN está compuesta esencialmente por iteraciones de la forma $h_s=A\mathrm{emb}(w_s)+B{h}_{s-1}$, donde se observa que la matriz $B\in {\mathcal{M}}_{D,D}\left(\mathbb{R}\right)$ es multiplicada por sí misma varias veces, dependiendo del largo de la secuencia que se esté procesando. Esta multiplicación reiterada provoca que el gradiente $\frac{\partial L}{\partial B}$ (con $L$ la función de pérdida) crezca de manera irrestricta (exploding gradient problem) o bien, decrezca hasta hacerse cercano a cero (vanishing gradient problem), lo cual dependerá del radio espectral de la matriz $B$. Ambos problemas dificultan el entrenamiento de este tipo de modelos ya que un gradiente irrestricto provoca un entrenamiento inestable (que puede incluso divergir), mientras que un gradiente nulo no permite actualizar los parámetros de la red neuronal por algoritmos de gradiente. Más aún, el tamaño de estos gradientes no es uniforme a lo largo de las capas de la red neuronal, por lo que una tasa de aprendizaje pequeña (e.g., para aminorar la inestabilidad del exploding gradient) puede ser contraproducente para el desvenecimiento de gradiente observado en las primeras capas.

Si bien se han propuesto algunas modificaciones para mitigar estos problemas (e.g., gradient clipping para el exploding gradient problem o arquitecturas con compuertas como GRU o LSTM para el vanishing gradient problem), las RNNs siguen teniendo limitaciones intrínsecas debido a su forma recursiva. En efecto, su procesamiento secuencial ralentiza considerablemente el entrenamiento de este tipo de modelos ya que cada secuencia se debe procesar token a token (i.e., no se pueden procesar todos los tokens en paralelo). Por otro lado, toda la información de la secuencia que se está procesando debe poder almacenarse en un vector $h_t\in {\mathbb{R}}^D$, lo cual es inviable para secuencias largas, provocando que el modelo empiece a olvidar la información de los primeros tokens de la secuencia a medida que se avanza en la recursion (en particular, las RNNs tienen un sesgo hacia los tokens más recientes). La arquitectura Transformer [4], y en particular los modelos tipo GPT, solucionan (al menos en principio) las limitaciones anteriores:

• Las secuencias $S=\left(w_1,\dots ,w_T\right)\in \mathcal{V}$ se procesan en paralelo, es decir, se pueden procesar todos los tokens de todas las secuencias de un batch de entrenamiento en una sola llamada a la red neuronal. En particular, esto permite escalar el entrenamiento en tamaño del dataset y largo de las secuencias. Sin embargo, la generación sigue siendo secuencial (token a token).

• No se produce el problema de vanishing/exploding gradient gracias a que se insertan capas de normalización y conexiones residuales entre las capas de la red neuronal.

• No se olvida información al procesar la secuencia ya que cada token puede poner atención a la información que contienen todos los otros tokens de la secuencia.

En esencia, la arquitectura GPT consiste en mejorar una secuencia de embeddings iniciales, $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$ (asociados a una secuencia de tokens $S=\left(w_1,\dots ,w_T\right)\in {\mathcal{V}}^{\ast }$ dados como entrada), pasándolos a través de una serie de bloques neuronales, llamados bloques Transformer, los cuales buscan enriquecer cada uno de estos embeddings con información adicional obtenida de los tokens vecinos dentro de la secuencia. De este modo, al pasar por el último bloque Transformer, los embeddings obtenidos son lo suficientemente ricos en información como para poder pasarlos por un cabezal de lenguaje simple (similar al usado en la RNN) para realizar la predicción del próximo token. Cada uno de estos bloques Transformer, que son el elemento principal de la arquitectura GPT, está formado principalmente por dos componentes:

• Módulo de self-attention: en este sub-módulo cada token interactúa con los otros tokens de la secuencia para obtener una representación más acorde al contexto.

• Módulo feedforward: red fully connected de dos capas usada para mejorar la representación obtenida en el módulo de auto-atención. Este sub-módulo es aplicado de forma independiente a cada token.

Adicionalmente, la arquitectura GPT tiene otros componentes: al comienzo de la arquitectura se debe incluir una capa de embedding inicial (llamada token embedding), a la cual se le suma una capa de embedding posicional (llamada positional encoding), la cual permitirá inyectar la información acerca de la posición en la que se encuentra cada token. A la salida del último bloque Transformer, se anexa un cabezal de lenguaje similar al usado en la RNN implementada anteriormente. Adicionalmente, cada uno de los componentes que forman la arquitectura GPT incluye sub-módulos menores como capas de normalización (en este caso, layer normalization) y dropout.

En esta sección se explicará e implementará cada uno estos componentes, intentando dar una motivación natural para cada uno de ellos. La siguiente clase GPTConfig contiene todos los hiperparámetros que definen la arquitectura GPT (se utilizarán valores pequeños para poder realizar el entrenamiento en Google Colab):

@dataclass
class GPTConfig:
    vocab_size: int = dataset.tokenizer.vocab_size
    context_window: int = 256
    embedding_dim: int = 384
    num_layers: int = 6
    num_heads: int = 6
    head_dim: int = 64
    ff_factor: int = 4
    dropout: float = 0.1

config = GPTConfig()

Layer normalization

El primer módulo que se revisará será la capa de normalización, la cual será añadida al comienzo de cada componente de la arquitectura (aunque originalmente se aplicaba al final de cada componente).

Al igual como ocurre en todas las redes neuronales profundas, la aplicación consecutiva de varias capas neuronales suele afectar el entrenamiento del modelo (e.g., por vanishing/exploding gradient). Para mitigar estos problemas, es usual agregar capas de normalización entremedio de los bloques de la red. Por ejemplo, en redes convolucionales es usual utilizar la técnica de batch normalization [5] (o más frecuentemente hoy en día, group normalization [6]), la cual normaliza cada canal de un batch de imágenes para que tenga media nula y varianza unitaria.

Sin embargo, no es obvio cómo aplicar esta técnica en modelos secuenciales ya que, en general, las secuencias de tokens no suelen tener el mismo largo y muchas veces los batches de entrenamiento deben ser pequeños debido a limitaciones de memoria en GPU, lo cual provoca que los estadísticos calculados no sean precisos, perjudicando el entrenamiento.

Por estos motivos, en modelos de lenguaje y en arquitecturas tipo Transformer es usual utilizar otra técnica de normalización más natural para este tipo de modelos, conocida como layer normalization [7], la cual normaliza cada token dentro de una secuencia de forma individual, por lo que su normalización no depende de otras secuencias dentro de un batch de entrenamiento ni de otros tokens de la secuencia que se está procesando. Más precisamente, para un vector de embedding $x\in {\mathbb{R}}^D$, layer normalization primero calcula su media $\mu \in \mathbb{R}$ y varianza ${\sigma }^2\ge 0$ a lo largo de las dimensiones del embedding $x$:

Luego, normaliza el vector de embedding coordenada a coordenada usando los valores de $\mu$ y $\sigma$. Es decir, la $d$-ésima coordenada del embedding normalizado, $\tilde{x}\in {\mathbb{R}}^D$, viene dada por

donde $ϵ>0$ es un escalar fijo y pequeño utilizado para estabilidad numérica. Finalmente, de forma análoga a lo realizado por batch normalization, layer normalization aplica una transformación afín a cada coordenada del embedding normalizado:

donde $\alpha ,\beta \in {\mathbb{R}}^D$ son parámetros entrenables del módulo $\mathrm{LayerNorm}$ y $\odot$ es el producto de Hadamard (i.e., $\alpha \odot \tilde{x}\in {\mathbb{R}}^D$ con ${\left(\alpha \odot \tilde{x}\right)}_d={\alpha }_d{\tilde{x}}_d$). La motivación detrás de esta última transformación es poder permitirle al modelo deshacer la normalización anterior. El siguiente diagrama resume este proceso:

La siguiente clase LayerNorm implementa esta técnica de normalización:

class LayerNorm(nn.Module):

    def __init__(self, emb_dim: int, epsilon: float = 1e-5) -> None:
        super().__init__()

        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))
        self.epsilon = epsilon

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.epsilon)
        return self.scale * x_norm + self.shift

• Ejemplo de uso.

  ln = LayerNorm(config.embedding_dim)
   x = torch.randn(batch_size, seq_length, config.embedding_dim)
   y = ln(x)
  
   assert y.shape == (batch_size, seq_length, config.embedding_dim)
   assert y.mean(dim=-1).allclose(torch.zeros(batch_size, seq_length), atol=1e-5)
   assert y.var(dim=-1).allclose(torch.ones(batch_size, seq_length), atol=1e-2)
   

Notar que los parámetros $\alpha$ y $\beta$ son inicializados como el vector $1\in {\mathbb{R}}^D$ y el vector $0\in {\mathbb{R}}^D$ respectivamente, lo cual garantiza que, en un comienzo, estos parámetros no tengan ningún efecto posterior a la normalización, lo cual estabiliza el entrenamiento. Por otro lado, la transformación afín realizada luego de la normalización, $\alpha \odot \tilde{x}+\beta$, es aplicada sobre todos tokens de la secuencia (y sobre todas las secuencias de un batch de entrenamiento) usando los mismos parámetros $\alpha ,\beta \in {\mathbb{R}}^D$ ya que esta capa de normalización es local y se aplica de forma individual para cada token.

Es importante mencionar que la implementación realizada en la clase LayerNorm tiene el mismo comportamiento que la clase nn.LayerNorm de PyTorch cuando se utiliza para procesar secuencias, pero no tiene el mismo comportamiento cuando es utilizada para batches imágenes (donde los tensores suelen ser de orden 4). Sin embargo, como se mencionó anteriormente, en imágenes es usual considerar otras técnicas de normalización como nn.BatchNorm2d o nn.GroupNorm.

Self attention

El componente clave de la arquitectura GPT corresponde al mecanismo de auto-atención (self-attention), el cual está encargado de mejorar los embeddings de una secuencia utilizando la información contextual de cada token. Más precisamente, dada una secuencia de vectores de embeddings, $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$, asociada a una secuencia de tokens $S=\left(w_1,\dots ,w_T\right)\in {\mathcal{V}}^{\ast }$, el objetivo del mecanismo de auto-atención es obtener otra secuencia de embeddings, $\left(y_1,\dots ,y_T\right)\subset {\mathbb{R}}^P$, donde cada embedding $y_t\in {\mathbb{R}}^P$ corresponderá a una versión mejorada del respectivo embedding inicial, $x_t\in {\mathbb{R}}^D$, la cual incorporará información contextual de todos los otros tokens de la secuencia $S$. Para esto, el mecanismo de auto-atención comienza proyectando cada uno de los embeddings $x_1,\dots ,x_T\in {\mathbb{R}}^D$ sobre el espacio ${\mathbb{R}}^P$ usando 3 proyecciones distintas (en la implementación, el escalar $P\in \mathbb{N}$ se representa por el atributo head_dim), es decir, para cada $t\in \left\{1,\dots ,T\right\}$ se definen los vectores

donde $W_Q,W_K,W_V\in {\mathcal{M}}_{P,D}\left(\mathbb{R}\right)$ son las 3 matrices (parámetros entrenables) que realizan las 3 proyecciones lineales a ${\mathbb{R}}^P$ de los vectores de embedding $x_1,\dots ,x_T\in {\mathbb{R}}^D$. De esta forma, cada embedding mejorado, $y_i\in {\mathbb{R}}^P$, se calculará como una suma ponderada (más precisamente, una combinación convexa) de los vectores $v_1,\dots ,v_T\in {\mathbb{R}}^P$, donde los ponderadores (que serán obtenidos a partir de las proyecciones $q_1,\dots ,q_T\in {\mathbb{R}}^P$ y $k_1,\dots ,k_T\in {\mathbb{R}}^P$) indicarán el grado de importancia que tiene cada token de $S$ para el token $w_i\in S$:

Cada ponderador $a\left(q_i,k_j\right)\in [0,1]$ se interpreta como la atención que coloca el token $w_i$ sobre el token $w_j$, lo cual se determinará utilizando las proyecciones $q_i\in {\mathbb{R}}^P$ y $k_j\in {\mathbb{R}}^P$. Por otro lado, notar que para que la combinación lineal que define $y_i\in {\mathbb{R}}^P$ sea convexa, se debe cumplir que ${\sum }_{j=1}^{T}a\left(q_i,k_j\right)=1$.

La motivación de realizar 3 proyecciones distintas, $q_t$, $k_t$ y $v_t$, para cada embedding $x_t\in {\mathbb{R}}^D$ de la secuencia de embeddings $\left(x_1,\dots ,x_T\right)$, viene de permitir distintas representaciones que se ajusten al papel que está jugando el respectivo token en cada situación (lógica prestada desde information retrieval):

• El vector $q_t\in {\mathbb{R}}^P$ (denominado query) es usado cuando el embedding $x_t\in {\mathbb{R}}^D$ es el que busca poner atención sobre otros vectores de embedding de la secuencia.

• El vector $k_t\in {\mathbb{R}}^P$ (denominado key) es usado cuando un token cualquiera de la secuencia busca poner atención sobre el embedding $x_t\in {\mathbb{R}}^D$.

• El vector $v_t\in {\mathbb{R}}^P$ (denominado value) es usado para representar la información útil del embedding $x_t\in {\mathbb{R}}^D$, la cual se considerará en la suma que define cada uno de los embeddings finales $\left(y_1,\dots ,y_T\right)\subset {\mathbb{R}}^P$.

Dado que cada embedding mejorado, $y_t\in {\mathbb{R}}^D$, incluye información acerca de todos los demás tokens de la secuencia $S=\left(w_1,\dots ,w_T\right)$ (y no solo del token $w_t\in \mathcal{V}$), a estos nuevos embeddings se les suele llamar embeddings contextuales ya que permite, por ejemplo, la desambiguación de palabras homógrafas, lo cual no es posible a partir de embeddings obtenidos al mirar los tokens individualmente.

Por otra parte, para satisfacer las restricciones de convexidad, $a\left(q_i,k_j\right)\ge 0$ y ${\sum }_{j=1}^{T}a\left(q_i,k_j\right)=1$, es usual comenzar considerando escalares irrestrictos, $s\left(q_i,k_j\right)\in \mathbb{R}$, denominados scores de atención, los cuales luego son normalizados usando la función $\mathrm{softmax}$, de manera similar a lo que se realiza para transformar un vector de logits en un vector de probabilidad. Notar que el cálculo de la función $\mathrm{softmax}$ debe realizarse a lo largo de la segunda variable de la función de score $s:{\mathbb{R}}^D\times {\mathbb{R}}^D\to \mathbb{R}$ ya que es esa la coordenada donde se pide que la suma sea $1$:

Para la elección de los scores de atención, $s\left(q_i,k_j\right)\in \mathbb{R}$, si bien se podría utilizar una red neuronal más pequeña que aprenda buenos valores (como ocurre en la atención de Bahdanau [8]), un enfoque sencillo y que funciona bien corresponde al dot-product self-attention (propuesto en la arquitectura original de Transformer [4]), el cual considera

donde el producto punto indica el grado de similitud entre el par query-key (recordar la relación con la similitud coseno: $\cos \angle (x,y):=\frac{⟨x,y⟩}{\parallel x\parallel \parallel y\parallel }$), mientras que la división por $\sqrt{P}$ (con $P$ la dimensión de los vectores $q_i$ y $k_j$) busca mantener la magnitud de los scores controlada, lo cual evita que la función $\mathrm{softmax}$ concentre casi toda su masa en una única posición al momento de calcular los ponderadores de atención.

El siguiente diagrama resume el mecanismo de auto-atención para una secuencia de embeddings $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$:

Es importante notar que se pueden calcular todos los embeddings contextuales $y_1,\dots ,y_T\in {\mathbb{R}}^P$ en paralelo (hasta el momento, cada $y_t\in {\mathbb{R}}^P$ se debería computar de forma independiente, iterando sobre $t\in \left\{1,\dots ,T\right\}$). Para esto, se almacenarán las proyecciones realizadas dentro de matrices

donde las 3 matrices, las cuales pertenecen a ${\mathcal{M}}_{T,P}\left(\mathbb{R}\right)$, almacenan los respectivos embeddings proyectados, uno en cada fila, similar a como funciona una matriz de embedding. Luego, notando que $⟨q_i,k_j⟩=⟨{\text{fila}}_i(Q),{\text{ fila}}_j(K)⟩={\left(Q{K}^{\top }\right)}_{ij}$, se pueden obtener los scores de atención mediante $s\left(q_i,k_j\right)={\left(\frac{Q{K}^{\top }}{\sqrt{P}}\right)}_{ij}$. De esta forma, si $A\in {\mathcal{M}}_{T,T}\left(\mathbb{R}\right)$ es la matriz de atención, la cual almacena $A_{ij}=a\left(q_i,k_j\right)$, la observación anterior permite escribir $A={\mathrm{softmax}}_{\text{por filas }}\left(\frac{Q{K}^{\top }}{\sqrt{P}}\right)$, donde ${\mathrm{softmax}}_{\text{por filas}}$ indica que se debe aplicar la función $\mathrm{softmax}$ en cada una de las filas de la matriz de scores $S=\frac{Q{K}^{\top }}{\sqrt{P}}\in {\mathcal{M}}_{T,T}\left(\mathbb{R}\right)$ (recordar que cada fila debe sumar 1, no cada columna). De este modo, si $Y\in {\mathcal{M}}_{T,P}\left(\mathbb{R}\right)$ es la matriz que almacena los embeddings contextuales calculados por el mecanismo de atención, $\left(y_1,\dots ,y_T\right)\subset {\mathbb{R}}^P$ (un embedding en cada fila), se observa que $Y_{ij}={\sum }_{k=1}^T{a}_{ik}{v}_{kj}=(AV{)}_{ij}$, por lo que se puede escribir $Y=\mathrm{Attention}(Q,K,V)$, donde

es el mecanismo de atención actuando sobre los embeddings iniciales $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$, cuyas proyecciones en ${\mathbb{R}}^P$ se almacenan en las matrices $Q,K,V\in {\mathcal{M}}_{T,P}\left(\mathbb{R}\right)$.

Por otro lado, las proyecciones $Q,K,V\in {\mathcal{M}}_{T,P}\left(\mathbb{R}\right)$ también pueden ser calculadas mediante un producto matricial: si $X\in {\mathcal{M}}_{T,D}\left(\mathbb{R}\right)$ es la matriz que almacena en sus filas los embeddings iniciales, $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$, entonces se tiene que $q_{ij}={\left(W_Q{x}_i\right)}_j=⟨{\text{fila}}_j\left(W_Q\right),x_i⟩=⟨{\text{fila}}_j\left(W_Q\right),{\text{ fila}}_i(X)⟩={\left(X{W}_Q^{\top }\right)}_{ij}$ (análogo para $K$ y $V$). Por lo tanto, $Q=X{W}_Q^{\top }$, $K=X{W}_K^{\top }$ y $V=X{W}_V^{\top }$, donde $W_Q,W_K,W_V\in {\mathcal{M}}_{P,D}\left(\mathbb{R}\right)$ son las matrices de proyección definidas anteriormente.

Masked self-attention

Un detalle importante pasado por alto hasta el momento es la pérdida de causalidad al momento de computar los embeddings contextuales $y_1,\dots ,y_T\in {\mathbb{R}}^P$. En efecto, dado que el objetivo central sigue siendo construir un modelo neuronal para predecir la probabilidad del próximo token dados los tokens anteriores, cada embedding contextual $y_t$ solo debería poder tener acceso a la información asociada a los tokens $w_1,\dots ,w_t\in \mathcal{V}$ y no a la información de los tokens siguientes, $w_{t+1},\dots ,w_T\in \mathcal{V}$, por lo que se debe forzar que $a\left(x_i,x_j\right)=0$ para $j>i$.

Si bien se podría sustituir directamente $A_{ij}=0$ para $j>i$ en la matriz de atención, esto provocaría que ${\sum }_{j=1}^{T}a\left(q_i,k_j\right)<1$, por lo que la combinación lineal ya no sería convexa. Una solución que evita este problema es sustituir en la matriz de scores, $S=\frac{Q{K}^{\top }}{\sqrt{P}}$, cada entrada donde $j>i$ por un valor muy negativo ($-\infty$ en la implementación). De este modo, como $\underset{s\to -\infty }{\lim }{e}^s=0$, se tendrá que $a\left(q_i,k_j\right)=\frac{e^{s\left(q_i,k_j\right)}}{{\sum }_{t=1}^T{e}^{s\left(q_i,k_t\right)}}=0$ para $j>i$, y se seguirá cumpliendo la restricción ${\sum }_{j=1}^{T}a\left(q_i,k_j\right)=1$ ya que la función $\mathrm{softmax}$ es aplicada luego de esta sustitución.

Esta variante de auto-atención se conoce como masked self-attention, y es esencial para preservar la causalidad del lenguaje. En la práctica, este tipo de exclusión se implementa construyendo una matriz auxiliar $M\in {\mathcal{M}}_{T,T}\left(\mathbb{R}\right)$, llamada máscara, donde las entradas no nulas de su $i$-ésima fila indicarán los tiempos $\left\{1,\dots ,T\right\}$ sobre los que podrá poner atención el embedding contextual $x_i\in {\mathbb{R}}^D$. En este caso, para mantener la causalidad, $M_{ij}=\{\begin{array}{ll}1& \text{ si }j\le i\\ 0& \text{ si }j>i\end{array}$ (en particular, $M\in {\mathcal{M}}_{T,T}\left(\mathbb{R}\right)$ es una matriz triangular inferior). El siguiente diagrama muestra el mecanismo de auto-atención incluyendo el masking causal:

Notar que se incluyó una capa de dropout [9], la cual es aplicada a la matriz de atención, $A\in {\mathcal{M}}_{T,T}\left(\mathbb{R}\right)$. Recordar que para un vector cualquiera, $x\in {\mathbb{R}}^D$, ${\mathrm{Dropout}}_p(x):=\frac{1}{1-p}m\odot x$, donde el vector binario $m\in {\left\{0,1\right\}}^D$ determina sus coordenadas al azar desde $m_d\sim \mathrm{Bernoulli}(1-p)$. Es decir, ${\mathrm{Dropout}}_p(x)$ apaga el valor de ciertas coordenadas de $x$, impidiendo el flujo de información a través de las conexiones neuronales relacionadas con las componentes apagadas. En cambio, el factor de ponderación, $\frac{1}{1-p}\ge 1$ busca compensar el flujo perdido aumentando el valor de las coordenadas que sí quedaron encendidas. En el caso del dropout aplicado a la matriz de atención, el funcionamiento es el mismo: ciertas coordenadas de la matriz $A\in {\mathcal{M}}_{T,T}\left(\mathbb{R}\right)$ son sustituidas por $0$, mientras que las restantes son ponderadas por $\frac{1}{1-p}$, respetando (en esperanza) el cumplimiento de la restricción ${\sum }_{j=1}^{T}a\left(q_i,k_j\right)=1$. En el código, el parámetro de dropout $p\in [0,1]$ se representa por el atributo config.dropout, el cual se compartirá de manera global en todas las capas de dropout definidas en el modelo.

Para ejemplificar el uso de dropout, se aplicará dropout (con $p=0.2$) sobre una matriz cuadrada que simulará ser una matriz de atención (no se aplicará causalidad ni $\mathrm{softmax}$ por simplicidad). Notar que los valores que permanecen activos son multiplicados cada uno por $\frac{1}{1-p}=1.25$:

p = 0.2
dropout = nn.Dropout(p)
matrix = torch.ones(5, 5)
matrix = dropout(matrix)
print(matrix)

tensor([[1.2500, 1.2500, 1.2500, 1.2500, 0.0000], [0.0000, 1.2500, 1.2500, 0.0000, 1.2500], [1.2500, 0.0000, 1.2500, 1.2500, 0.0000], [1.2500, 1.2500, 0.0000, 1.2500, 1.2500], [0.0000, 1.2500, 1.2500, 1.2500, 1.2500]])

La motivación usual de la técnica de dropout es que el modelo aprenda a predecir la respuesta correcta aún cuando no tiene acceso a toda la información. En particular, para la matriz de atención, si $a\left(q_i,k_j\right)=0$, entonces el embedding contextual $y_i\in {\mathbb{R}}^P$ (asociado al token $w_i\in S$) no podrá obtener información contextual acerca del token $w_j\in S$, independientemente de si $j\le i$ o no. Dado que las posiciones donde ocurre el dropout se eligen de forma aleatoria en cada secuencia e iteración de entrenamiento, aplicar esta regularización en la matriz de atención fuerza a que el modelo tenga que aprender a utilizar la información de todas las posiciones del contexto para realizar una predicción. Durante la inferencia (i.e., cuando el modelo ya está entrenado y se invocó model.eval()), el dropout se deshabilita y todos los tokens del contexto son utilizados para el cálculo de los embeddings contextuales.

La siguiente clase implementa el mecanismo de masked self-attention. Además, aunque no es necesario hacerlo, se guardará la matriz de atención en el atributo attn_matrix para poder visualizarla posteriormente cuando el modelo esté en modo inferencia (recordar que el atributo training se hereda de la clase nn.Module, e indica si el modelo está en modo entrenamiento o en modo inferencia).

class SelfAttentionHead(nn.Module):

    def __init__(self, cfig: GPTConfig) -> None:
        super().__init__()

        self.query = nn.Linear(cfig.embedding_dim, cfig.head_dim, bias=False)
        self.key = nn.Linear(cfig.embedding_dim, cfig.head_dim, bias=False)
        self.value = nn.Linear(cfig.embedding_dim, cfig.head_dim, bias=False)

        self.dropout = nn.Dropout(cfig.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        Q = self.query(x)
        K = self.key(x)
        V = self.value(x)

        batch_size, seq_length, head_dim = Q.shape
        mask = torch.tril(torch.ones(seq_length, seq_length)).to(x.device)

        attn = Q @ K.transpose(-1, -2) / (head_dim ** 0.5)
        attn = attn.masked_fill(mask == 0, float('-inf'))
        attn = self.dropout(attn.softmax(dim=-1))

        if self.training == False:
            self.attn_matrix = attn.detach()

        return attn @ V

• Ejemplo de uso.

  sa = SelfAttentionHead(config)
   x = torch.randn(batch_size, seq_length, config.embedding_dim)
   y = sa(x)
  
   assert y.shape == (batch_size, seq_length, config.head_dim)
   

Atención multicabezal

El objetivo del mecanismo de auto-atención definido anteriormente es que cada embedding contextual $y_t\in {\mathbb{R}}^P$ pueda resumir de mejor forma la información contenida en el vector de embedding inicial, $x_t\in {\mathbb{R}}^D$, haciendo uso del contexto en el que se encuentra el token $w_t\in S$ asociado al embedding $x_t$ (i.e., usando información de los tokens vecinos). Este mecanismo puede repetirse varias veces de forma independiente sobre una misma secuencia de embeddings $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$, donde cada instancia del mecanismo de auto-atención, que se suele llamar cabezal de atención, tendrá sus propias matrices de proyección $W_Q,W_K,W_V\in {\mathcal{M}}_{P,D}\left(\mathbb{R}\right)$ y, en consecuencia, su propia matriz de atención y embeddings finales, $\left(y_1,\dots ,y_T\right)\subset {\mathbb{R}}^P$. La motivación de usar varios cabezales de atención es que cada cabezal puede aprender a extraer una información contextual diferente, de forma similar a como actúan los filtros en una red convolucional.

Para una secuencia de embeddings $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$ y $\mathrm{H}\ge 1$ cabezales de atención independientes actuando sobre esta secuencia, se obtendrán $\mathrm{H}$ nuevas secuencias de embeddings, $\left(y_1^{(h)},\dots ,y_T^{(h)}\right)\subset {\mathbb{R}}^P$, con $h\in \left\{1,\dots ,\mathrm{H}\right\}$. Con esto, se pueden combinar todos los cabezales para formar una secuencia de embeddings $\left({\tilde{y}}_1,\dots ,{\tilde{y}}_T\right)\subset {\mathbb{R}}^{P\cdot \mathrm{H}}$ de mayor dimensión concatenando los embeddings de todos los cabezales:

Además, para poder relacionar los embeddings obtenidos por cada cabezal, cada embedding extendido ${\tilde{y}}_t$ suele ser pasado por una capa lineal extra, lo que permite, además, recuperar la dimensión de los embeddings originales, $\left(x_1,\dots ,x_T\right)\subset {\mathbb{R}}^D$. Por lo tanto, los embeddings finales que retorna un bloque de atención multicabezal, $\left(u_1,\dots ,u_T\right)\in {\mathbb{R}}^D$, son de la forma $u_t=W_0{\tilde{y}}_t\in {\mathbb{R}}^D$, donde la matriz $W_0\in {\mathcal{M}}_{D,P\cdot \mathrm{H}}\left(\mathbb{R}\right)$ proyecta cada embedding concatenado, ${\tilde{y}}_t\in {\mathbb{R}}^{P\cdot \mathrm{H}}$, de vuelta al espacio ${\mathbb{R}}^D$. El siguiente diagrama resume esta idea:

En la práctica (y en la clase de hiperparámetros GPTConfig) es usual considerar como dimensión de cada cabezal a $P=\frac{D}{\mathrm{H}}$, donde se asume que $\mathrm{H}\mid D$ (i.e., $D$ es divisible por $\mathrm{H}$). La idea de esto es considerar que la dimensión de embedding $D\in \mathbb{N}$ se reparte entre los $\mathrm{H}\in \mathbb{N}$ cabezales de atención. En la siguiente implementación, $\mathrm{H}$ es representado por el atributo num_head:

class MultiHeadSelfAttention(nn.Module):

    def __init__(self, cfig: GPTConfig) -> None:
        super().__init__()
        self.heads = nn.ModuleList([SelfAttentionHead(cfig) for _ in range(cfig.num_heads)])
        self.projection = nn.Linear(cfig.num_heads * cfig.head_dim, cfig.embedding_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        output = torch.cat([head(x) for head in self.heads], dim=-1)
        output = self.projection(output)
        return output