0gcvkiopdnmm1xi5z.jpeg

Construyamos un modelo LSTM simple y entrenémoslo para predecir el siguiente token dado un prefijo de tokens. Ahora quizás te preguntes qué es un token.

Tokenización

Normalmente, para los modelos de lenguaje, un token puede significar

  1. Un solo carácter (o un solo byte)
  2. Una palabra completa en el idioma de destino.
  3. Algo entre 1 y 2. Esto generalmente se llama subpalabra.

Asignar un solo carácter (o byte) a un token es muy restrictivo ya que estamos sobrecargando ese token para contener mucho contexto sobre dónde ocurre. Esto se debe a que el carácter «c», por ejemplo, aparece en muchas palabras diferentes, y para predecir el siguiente carácter después de ver el carácter «c», es necesario que analicemos detenidamente el contexto principal.

Asignar una sola palabra a un token también es problemático, ya que el inglés tiene entre 250.000 y 1 millón de palabras. Además, ¿qué sucede cuando se agrega una nueva palabra al idioma? ¿Necesitamos volver atrás y volver a entrenar todo el modelo para tener en cuenta esta nueva palabra?

La tokenización de subpalabras se considera el estándar de la industria en el año 2023. Asigna subcadenas de bytes que aparecen con frecuencia juntos a tokens únicos. Normalmente, los modelos de lenguaje tienen desde unos pocos miles (digamos 4000) hasta decenas de miles (digamos 60 000) de tokens únicos. El algoritmo para determinar qué constituye un token está determinado por el Algoritmo BPE (codificación de pares de bytes)metro.

Para elegir la cantidad de tokens únicos en nuestro vocabulario (llamado tamaño de vocabulario), debemos tener en cuenta algunas cosas:

  1. Si elegimos muy pocas fichas, volvemos al régimen de una ficha por personaje y es difícil para el modelo aprender algo útil.
  2. Si elegimos demasiados tokens, terminamos en una situación en la que las tablas de incrustación del modelo eclipsan el resto del peso del modelo y resulta difícil implementar el modelo en un entorno restringido. El tamaño de la tabla de incrustación dependerá de la cantidad de dimensiones que usemos para cada token. No es raro usar un tamaño de 256, 512, 786, etc. Si usamos una dimensión de incrustación de tokens de 512 y tenemos 100k tokens, terminamos con una tabla de incrustación que usa 200MiB en memoria.

Por lo tanto, debemos lograr un equilibrio al elegir el tamaño del vocabulario. En este ejemplo, seleccionamos 6600 tokens y entrenamos nuestro tokenizador con un tamaño de vocabulario de 6600. A continuación, echemos un vistazo a la definición del modelo en sí.

El modelo PyTorch

El modelo en sí es bastante sencillo. Tenemos las siguientes capas:

  1. Incrustación de tokens (tamaño de vocabulario = 6600, incrustación tenue = 512), para un tamaño total de aproximadamente 15 MiB (asumiendo 4 bytes float32 como tipo de datos de la tabla de incrustación)
  2. LSTM (número de capas = 1, dimensión oculta = 786) para un tamaño total de aproximadamente 16 MiB
  3. Perceptrón multicapa (de 786 a 3144 a 6600 dimensiones) para un tamaño total de aproximadamente 93 MiB

El modelo completo tiene alrededor de 31 millones de parámetros entrenables para un tamaño total de aproximadamente 120 MiB.

Aquí está el código PyTorch para el modelo.

class WordPredictionLSTMModel(nn.Module):
def __init__(self, num_embed, embed_dim, pad_idx, lstm_hidden_dim, lstm_num_layers, output_dim, dropout):
super().__init__()
self.vocab_size = num_embed
self.embed = nn.Embedding(num_embed, embed_dim, pad_idx)
self.lstm = nn.LSTM(embed_dim, lstm_hidden_dim, lstm_num_layers, batch_first=True, dropout=dropout)
self.fc = nn.Sequential(
nn.Linear(lstm_hidden_dim, lstm_hidden_dim * 4),
nn.LayerNorm(lstm_hidden_dim * 4),
nn.LeakyReLU(),
nn.Dropout(p=dropout),

nn.Linear(lstm_hidden_dim * 4, output_dim),
)
#

def forward(self, x):
x = self.embed(x)
x, _ = self.lstm(x)
x = self.fc(x)
x = x.permute(0, 2, 1)
return x
#
#

Aquí está el resumen del modelo usando torchinfo.

Resumen del modelo LSTM

=================================================================
Layer (type:depth-idx) Param #
=================================================================
WordPredictionLSTMModel -
├─Embedding: 1–1 3,379,200
├─LSTM: 1–2 4,087,200
├─Sequential: 1–3 -
│ └─Linear: 2–1 2,474,328
│ └─LayerNorm: 2–2 6,288
│ └─LeakyReLU: 2–3 -
│ └─Dropout: 2–4 -
│ └─Linear: 2–5 20,757,000
=================================================================
Total params: 30,704,016
Trainable params: 30,704,016
Non-trainable params: 0
=================================================================

Interpretando la exactitud: Después de entrenar este modelo con 12 millones de oraciones en inglés durante aproximadamente 8 horas en una GPU P100, logramos una pérdida de 4.03, una precisión entre los primeros 1 del 29 % y una precisión entre los 5 primeros del 49 %. Esto significa que el 29 % de las veces, el modelo pudo predecir correctamente el siguiente token, y el 49 % de las veces, el siguiente token en el conjunto de entrenamiento fue una de las 5 predicciones principales del modelo.

¿Cuál debería ser nuestra métrica de éxito? Si bien los números de precisión de los primeros 1 y 5 primeros de nuestro modelo no son impresionantes, no son tan importantes para nuestro problema. Nuestras palabras candidatas son un pequeño conjunto de palabras posibles que se ajustan al patrón de deslizamiento. Lo que queremos de nuestro modelo es poder seleccionar un candidato ideal para completar la oración de manera que sea sintáctica y semánticamente coherente. Dado que nuestro modelo aprende el naturaleza del lenguaje a través de los datos de entrenamiento, esperamos que asigne una mayor probabilidad a oraciones coherentes. Por ejemplo, si tenemos la frase “El jugador de béisbol” y posibles candidatos de finalización (“corrió”, “nadó”, “se escondió”), entonces la palabra “corrió” es una mejor palabra de seguimiento que las otras dos. Entonces, si nuestro modelo predice la palabra corrió con mayor probabilidad que el resto, nos funciona.

Interpretando la pérdida: Una pérdida de 4,03 significa que la probabilidad logarítmica negativa de la predicción es 4,03, lo que significa que la probabilidad de predecir correctamente el siguiente token es e^-4,03 = 0,0178 o 1/56. Un modelo inicializado aleatoriamente normalmente tiene una pérdida de aproximadamente 8,8, que es -log_e(1/6600), ya que el modelo predice aleatoriamente 1/6600 tokens (siendo 6600 el tamaño del vocabulario). Si bien una pérdida de 4.03 puede no parecer muy grande, es importante recordar que el modelo entrenado es aproximadamente 120 veces mejor que un modelo no entrenado (o inicializado aleatoriamente).

A continuación, veamos cómo podemos usar este modelo para mejorar las sugerencias desde nuestro teclado deslizante.

Usar el modelo para eliminar sugerencias no válidas

Veamos un ejemplo real. Supongamos que tenemos una oración parcial. «Creo»y el usuario realiza el patrón de deslizamiento que se muestra en azul a continuación, comenzando en «o», pasando entre las letras «c» y «v» y terminando entre las letras «e» y «v».

Algunas posibles palabras que podrían representarse con este patrón de deslizamiento son

  1. Encima
  2. Oct (abreviatura de octubre)
  3. Hielo
  4. He (con el apóstrofe implícito)

De estas sugerencias, la más probable será «He». Introduzcamos estas sugerencias en nuestro modelo y veamos qué arroja.

[I think] [I've] = 0.00087
[I think] [over] = 0.00051
[I think] [ice] = 0.00001
[I think] [Oct] = 0.00000

El valor después del signo = es la probabilidad de que la palabra complete válidamente el prefijo de la oración. En este caso, vemos que a la palabra “tengo” se le ha asignado la mayor probabilidad. Por lo tanto, es la palabra más probable que siga al prefijo de oración «Creo».

La siguiente pregunta que quizás tengas es cómo podemos calcular estas probabilidades de la siguiente palabra. Vamos a ver.

Calcular la probabilidad de la siguiente palabra

Para calcular la probabilidad de que una palabra sea una finalización válida de un prefijo de oración, ejecutamos el modelo en modo de evaluación (inferencia) e ingresamos el prefijo de oración tokenizado. También tokenizamos la palabra después de agregar un prefijo de espacio en blanco a la palabra. Esto se hace porque el pretokenizador de HuggingFace divide las palabras con espacios al principio de la palabra, por lo que queremos asegurarnos de que nuestras entradas sean consistentes con la estrategia de tokenización utilizada por los tokenizadores de HuggingFace.

Supongamos que la palabra candidata se compone de 3 tokens T0, T1 y T2.

  1. Primero ejecutamos el modelo con el prefijo de oración tokenizado original. Para el último token, verificamos la probabilidad de predecir el token T0. Agregamos esto a la lista de «problemas».
  2. A continuación, ejecutamos una predicción en el prefijo+T0 y verificamos la probabilidad del token T1. Agregamos esta probabilidad a la lista de «problemas».
  3. A continuación, ejecutamos una predicción en el prefijo+T0+T1 y verificamos la probabilidad del token T2. Agregamos esta probabilidad a la lista de «problemas».

La lista de «problemas» contiene las probabilidades individuales de generar los tokens T0, T1 y T2 en secuencia. Dado que estos tokens corresponden a la tokenización de la palabra candidata, podemos multiplicar estas probabilidades para obtener la probabilidad combinada de que el candidato complete el prefijo de la oración.

El código para calcular las probabilidades de finalización se muestra a continuación.

 def get_completion_probability(self, input, completion, tok):
self.model.eval()
ids = tok.encode(input).ids
ids = torch.tensor(ids, device=self.device).unsqueeze(0)
completion_ids = torch.tensor(tok.encode(completion).ids, device=self.device).unsqueeze(0)
probs = []
for i in range(completion_ids.size(1)):
y = self.model(ids)
y = y[0,:,-1].softmax(dim=0)
# prob is the probability of this completion.
prob = y[completion_ids[0,i]]
probs.append(prob)
ids = torch.cat([ids, completion_ids[:,i:i+1]], dim=1)
#
return torch.tensor(probs)
#

Podemos ver algunos ejemplos más a continuación.

[That ice-cream looks] [really] = 0.00709
[That ice-cream looks] [delicious] = 0.00264
[That ice-cream looks] [absolutely] = 0.00122
[That ice-cream looks] [real] = 0.00031
[That ice-cream looks] [fish] = 0.00004
[That ice-cream looks] [paper] = 0.00001
[That ice-cream looks] [atrocious] = 0.00000

[Since we're heading] [toward] = 0.01052
[Since we're heading] [away] = 0.00344
[Since we're heading] [against] = 0.00035
[Since we're heading] [both] = 0.00009
[Since we're heading] [death] = 0.00000
[Since we're heading] [bubble] = 0.00000
[Since we're heading] [birth] = 0.00000

[Did I make] [a] = 0.22704
[Did I make] [the] = 0.06622
[Did I make] [good] = 0.00190
[Did I make] [food] = 0.00020
[Did I make] [color] = 0.00007
[Did I make] [house] = 0.00006
[Did I make] [colour] = 0.00002
[Did I make] [pencil] = 0.00001
[Did I make] [flower] = 0.00000

[We want a candidate] [with] = 0.03209
[We want a candidate] [that] = 0.02145
[We want a candidate] [experience] = 0.00097
[We want a candidate] [which] = 0.00094
[We want a candidate] [more] = 0.00010
[We want a candidate] [less] = 0.00007
[We want a candidate] [school] = 0.00003

[This is the definitive guide to the] [the] = 0.00089
[This is the definitive guide to the] [complete] = 0.00047
[This is the definitive guide to the] [sentence] = 0.00006
[This is the definitive guide to the] [rapper] = 0.00001
[This is the definitive guide to the] [illustrated] = 0.00001
[This is the definitive guide to the] [extravagant] = 0.00000
[This is the definitive guide to the] [wrapper] = 0.00000
[This is the definitive guide to the] [miniscule] = 0.00000

[Please can you] [check] = 0.00502
[Please can you] [confirm] = 0.00488
[Please can you] [cease] = 0.00002
[Please can you] [cradle] = 0.00000
[Please can you] [laptop] = 0.00000
[Please can you] [envelope] = 0.00000
[Please can you] [options] = 0.00000
[Please can you] [cordon] = 0.00000
[Please can you] [corolla] = 0.00000

[I think] [I've] = 0.00087
[I think] [over] = 0.00051
[I think] [ice] = 0.00001
[I think] [Oct] = 0.00000

[Please] [can] = 0.00428
[Please] [cab] = 0.00000

[I've scheduled this] [meeting] = 0.00077
[I've scheduled this] [messing] = 0.00000

Estos ejemplos muestran la probabilidad de que la palabra complete la oración anterior. Los candidatos se clasifican en orden decreciente de probabilidad.

Dado que Transformers está reemplazando lentamente los modelos LSTM y RNN para tareas basadas en secuencias, echemos un vistazo a cómo sería un modelo Transformer para el mismo objetivo.