Supercharge los grandes modelos de lenguaje con la predicción de múltiples tokens

Los grandes modelos de lenguaje (LLM) como GPT, LLaMA y otros han revolucionado el mundo con su capacidad para entender y generar texto similar al humano. Sin embargo, a pesar de sus capacidades impresionantes, el método estándar de entrenamiento de estos modelos, conocido como “predicción de token siguiente”, tiene algunas limitaciones inherentes.

En la predicción de token siguiente, el modelo se entrena para predecir el token siguiente en una secuencia dado el contexto previo. Aunque este enfoque ha demostrado ser exitoso, puede llevar a modelos que luchan con dependencias de largo alcance y tareas de razonamiento complejas. Además, la discrepancia entre el régimen de entrenamiento (fuerza del maestro) y el proceso de generación autoregresiva durante la inferencia puede resultar en un rendimiento subóptimo.

Un reciente artículo de investigación de Gloeckle et al. (2024) de Meta AI introduce un nuevo paradigma de entrenamiento llamado “predicción de múltiples tokens” que busca abordar estas limitaciones y mejorar los grandes modelos de lenguaje. En este artículo, profundizaremos en los conceptos básicos, detalles técnicos y posibles implicaciones de esta investigación innovadora.

¿Qué es la predicción de múltiples tokens?

La idea clave detrás de la predicción de múltiples tokens es entrenar a los modelos de lenguaje para predecir múltiples tokens futuros simultáneamente, en lugar de solo el token siguiente. Específicamente, durante el entrenamiento, el modelo se encarga de predecir los siguientes n tokens en cada posición del corpus de entrenamiento, utilizando n cabezales de salida independientes que operan sobre un tronco de modelo compartido.

Por ejemplo, con una configuración de predicción de 4 tokens, el modelo se entrenaría para predecir los siguientes 4 tokens a la vez, dado el contexto previo. Este enfoque fomenta que el modelo capture dependencias de largo alcance y desarrolle una mejor comprensión de la estructura y coherencia generales del texto.

Un ejemplo juguete

Para entender mejor el concepto de predicción de múltiples tokens, consideremos un ejemplo simple. Supongamos que tenemos la siguiente oración:

“El rápido zorro marrón salta sobre el perro perezoso.”

En el enfoque estándar de predicción de token siguiente, el modelo se entrenaría para predecir el token siguiente dado el contexto previo. Por ejemplo, dado el contexto “El rápido zorro marrón salta sobre el”, el modelo se encargaría de predecir el token siguiente, “perezoso”.

Con la predicción de múltiples tokens, sin embargo, el modelo se entrenaría para predecir múltiples tokens futuros a la vez. Por ejemplo, si configuramos n=4, el modelo se entrenaría para predecir los siguientes 4 tokens simultáneamente. Dado el mismo contexto “El rápido zorro marrón salta sobre el”, el modelo se encargaría de predecir la secuencia “perezoso ” (Nota el espacio después de “perezoso” para indicar el final de la oración).

Al entrenar al modelo para predecir múltiples tokens futuros a la vez, se fomenta que el modelo capture dependencias de largo alcance y desarrolle una mejor comprensión de la estructura y coherencia generales del texto.

Detalles técnicos

Los autores proponen una arquitectura simple pero efectiva para implementar la predicción de múltiples tokens. El modelo consiste en un tronco de transformador compartido que produce una representación latente del contexto de entrada, seguido de n capas de transformador independientes (cabezales de salida) que predicen los tokens futuros respectivos.

Durante el entrenamiento, las pasadas hacia adelante y hacia atrás se orquestan cuidadosamente para minimizar la huella de memoria de la GPU. El tronco compartido calcula la representación latente, y luego cada cabeza de salida realiza su pasada hacia adelante y hacia atrás secuencialmente, acumulando gradientes en el nivel del tronco. Este enfoque evita materializar todos los vectores de logit y sus gradientes simultáneamente, reduciendo el uso máximo de memoria de la GPU de O(nV + d) a O(V + d), donde V es el tamaño del vocabulario y d es la dimensión de la representación latente.

La implementación eficiente en términos de memoria

Uno de los desafíos en el entrenamiento de predictores de múltiples tokens es reducir su uso de memoria de la GPU. Dado que el tamaño del vocabulario (V) es típicamente mucho mayor que la dimensión de la representación latente (d), los vectores de logit se convierten en el cuello de botella del uso de memoria de la GPU.

Para abordar este desafío, los autores proponen una implementación eficiente en términos de memoria que adapta cuidadosamente la secuencia de operaciones hacia adelante y hacia atrás. En lugar de materializar todos los logit y sus gradientes simultáneamente, la implementación calcula las pasadas hacia adelante y hacia atrás para cada cabeza de salida independiente, acumulando gradientes en el nivel del tronco.

Este enfoque evita almacenar todos los vectores de logit y sus gradientes en memoria simultáneamente, reduciendo el uso máximo de memoria de la GPU de O(nV + d) a O(V + d), donde n es el número de tokens futuros que se están prediciendo.

Ventajas de la predicción de múltiples tokens

El artículo de investigación presenta varias ventajas convincentes del uso de la predicción de múltiples tokens para entrenar a los grandes modelos de lenguaje:

1. Mejora de la eficiencia de muestra: Al fomentar que el modelo prediga múltiples tokens futuros a la vez, la predicción de múltiples tokens conduce a una mejor eficiencia de muestra. Los autores demuestran mejoras significativas en el rendimiento en tareas de comprensión y generación de código, con modelos de hasta 13B de parámetros que resuelven alrededor de un 15% más de problemas en promedio.
2. Infancia más rápida: Las cabezales de salida adicionales entrenadas con la predicción de múltiples tokens se pueden aprovechar para la decodificación especulativa auto-especulativa, una variante de la decodificación especulativa que permite la predicción de tokens en paralelo. Esto resulta en tiempos de inferencia hasta 3 veces más rápidos en una amplia gama de tamaños de lote, incluso para modelos grandes.
3. Fomento de dependencias de largo alcance: La predicción de múltiples tokens fomenta que el modelo capture dependencias de largo alcance y patrones en los datos, lo que es particularmente beneficioso para tareas que requieren comprensión y razonamiento sobre contextos más amplios.
4. Razonamiento algorítmico: Los autores presentan experimentos en tareas sintéticas que demuestran la superioridad de la predicción de múltiples tokens en el desarrollo de cabezales de inducción y capacidades de razonamiento algorítmico, especialmente para tamaños de modelo más pequeños.
5. Coherencia y consistencia: Al entrenar al modelo para predecir múltiples tokens futuros simultáneamente, la predicción de múltiples tokens fomenta el desarrollo de representaciones coherentes y consistentes. Esto es particularmente beneficioso para tareas que requieren la generación de texto más largo y coherente, como la narración de historias, la escritura creativa o la generación de manuales instructivos.
6. Mejora de la generalización: Los experimentos de los autores en tareas sintéticas sugieren que los modelos de predicción de múltiples tokens exhiben una mejor capacidad de generalización, especialmente en configuraciones fuera de la distribución. Esto puede deberse a la capacidad del modelo para capturar patrones y dependencias de largo alcance, lo que puede ayudarlo a extrapolar más efectivamente a escenarios no vistos.

Ejemplos e intuiciones

Para proporcionar más intuición sobre por qué la predicción de múltiples tokens funciona tan bien, consideremos algunos ejemplos:

1. Generación de código: En el contexto de la generación de código, predecir múltiples tokens simultáneamente puede ayudar al modelo a entender y generar estructuras de código más complejas. Por ejemplo, al generar una definición de función, predecir solo el token siguiente puede no proporcionar suficiente contexto para que el modelo genere la firma de la función correctamente. Sin embargo, al predecir múltiples tokens a la vez, el modelo puede capturar mejor las dependencias entre el nombre de la función, los parámetros y el tipo de retorno, lo que lleva a una generación de código más precisa y coherente.
2. Razonamiento del lenguaje natural: Consideremos un escenario en el que un modelo de lenguaje se encarga de responder a una pregunta que requiere razonamiento sobre múltiples pasos o piezas de información. Al predecir múltiples tokens simultáneamente, el modelo puede capturar mejor las dependencias entre los diferentes componentes del proceso de razonamiento, lo que lleva a respuestas más coherentes y precisas.
3. Generación de texto a largo plazo: Al generar texto a largo plazo, como historias, artículos o informes, mantener la coherencia y la consistencia durante un período prolongado puede ser un desafío para los modelos de lenguaje entrenados con la predicción de token siguiente. La predicción de múltiples tokens fomenta que el modelo desarrolle representaciones que capturen la estructura y el flujo generales del texto, lo que puede llevar a generaciones más coherentes y consistentes a largo plazo.

Limitaciones y direcciones futuras

Aunque los resultados presentados en el artículo son impresionantes, hay algunas limitaciones y preguntas abiertas que requieren una investigación adicional:

1. Número óptimo de tokens: El artículo explora diferentes valores de n (el número de tokens futuros que se van a predecir) y encuentra que n=4 funciona bien para muchas tareas. Sin embargo, el valor óptimo de n puede depender de la tarea específica, el conjunto de datos y el tamaño del modelo. Desarrollar métodos principled para determinar el valor óptimo de n podría llevar a mejoras adicionales en el rendimiento.
2. Tamaño del vocabulario y tokenización: Los autores señalan que el tamaño óptimo del vocabulario y la estrategia de tokenización para los modelos de predicción de múltiples tokens pueden diferir de los utilizados para los modelos de predicción de token siguiente. Explorar este aspecto podría llevar a mejores compromisos entre la longitud de la secuencia comprimida y la eficiencia computacional.
3. Pérdidas de predicción auxiliares: Los autores sugieren que su trabajo podría despertar interés en el desarrollo de nuevas pérdidas de predicción auxiliares para los grandes modelos de lenguaje, más allá de la predicción de token siguiente estándar. Investigar pérdidas auxiliares alternativas y sus combinaciones con la predicción de múltiples tokens es una dirección de investigación emocionante.
4. Comprensión teórica: Aunque el artículo proporciona algunas intuiciones y evidencia empírica sobre la efectividad de la predicción de múltiples tokens, una comprensión teórica más profunda de por qué y cómo este enfoque funciona tan bien sería valiosa.

Conclusión

El artículo de investigación “Mejores y más rápidos grandes modelos de lenguaje a través de la predicción de múltiples tokens” de Gloeckle et al. introduce un nuevo paradigma de entrenamiento que tiene el potencial de mejorar significativamente el rendimiento y las capacidades de los grandes modelos de lenguaje. Al entrenar a los modelos para predecir múltiples tokens futuros simultáneamente, la predicción de múltiples tokens fomenta el desarrollo de dependencias de largo alcance, capacidades de razonamiento algorítmico y una mejor eficiencia de muestra.

La implementación técnica propuesta por los autores es elegante y computacionalmente eficiente, lo que la hace factible para aplicar a la formación de modelos de lenguaje a gran escala. Además, la capacidad de aprovechar la decodificación auto-especulativa para una inferencia más rápida es una ventaja práctica significativa.

Aunque todavía hay preguntas abiertas y áreas para explorar, esta investigación representa un paso emocionante hacia adelante en el campo de los grandes modelos de lenguaje. A medida que la demanda de modelos de lenguaje más capaces y eficientes continúa creciendo, la predicción de múltiples tokens podría convertirse en un componente clave en la próxima generación de estos sistemas de inteligencia artificial poderosos.