En un estudio reciente, investigadores de Meta, Ecole des Ponts ParisTech y Université Paris-Saclay proponen un nuevo enfoque para mejorar la precisión y velocidad de los modelos de lenguaje grande (LLM) de IA, permitiendo la predicción de múltiples tokens simultáneamente. Esta innovación desafía el diseño tradicional de modelos autorregresivos, que predicen un token a la vez.
Beneficios de la Predicción Multi-Token
Aunque la predicción de múltiples tokens no es adecuada para todos los LLM o tareas lingüísticas, ofrece ventajas significativas en ciertos escenarios, como la aceleración de tareas generativas hasta tres veces más rápidas que los métodos convencionales. A pesar de que aún hay potencial para refinamiento, esta técnica podría ser una herramienta poderosa en aplicaciones específicas de LLM.
Desafíos de la Predicción del Siguiente Token
El método tradicional para entrenar LLM se llama "predicción del siguiente token." Esta técnica de aprendizaje autogestionado presenta al modelo una secuencia de tokens para que prediga el siguiente token, que se añade al input para más predicciones. Este proceso iterativo, aplicado a grandes corpus de texto, permite al modelo aprender a generar texto coherente. Sin embargo, los investigadores han identificado limitaciones en este enfoque, como el desarrollo de habilidades de procesamiento del lenguaje, adquisición de conocimiento y razonamiento. Al enfocarse en un token a la vez, los modelos corren el riesgo de ser demasiado sensibles a patrones locales y pueden pasar por alto razonamientos que requieren un contexto más amplio. Además, la predicción del siguiente token requiere vastos conjuntos de datos para alcanzar niveles de fluidez similares a los humanos, que logran con menos texto. El estudio de Meta afirma que "entrenar modelos de lenguaje para predecir múltiples tokens futuros a la vez resulta en una mayor eficiencia de muestra."
Explorando la Predicción Multi-Token
En contraste, la predicción multi-token guía al LLM para predecir varios tokens futuros en cada posición de los datos de entrenamiento simultáneamente. Los investigadores presentan una arquitectura de predicción multi-token sencilla que no impone requisitos adicionales de tiempo de entrenamiento o memoria. Este modelo se basa en la arquitectura Transformer, que es fundamental para la mayoría de los LLM, pero con modificaciones. En lugar de generar una única salida, incluye múltiples cabezas de salida independientes para cada predicción de token.
Implementación de la Predicción Multi-Token
Durante la inferencia, el modelo utiliza el método tradicional de predicción del siguiente token para cada cabeza de salida, empleando las cabezas adicionales para agilizar el proceso de decodificación. El marco aprovecha trabajos previos en el campo. "Mientras que es costo-efectivo y simple, la predicción multi-token mejora significativamente el entrenamiento de modelos Transformer más rápidos y potentes," afirman los investigadores.
Resultados y Observaciones
El equipo probó su estrategia de predicción multi-token con modelos que van desde 300 millones hasta 13 mil millones de parámetros. Sus hallazgos revelan patrones notables: los modelos más pequeños muestran menos beneficios de la predicción multi-token, que se vuelve cada vez más efectiva a medida que aumenta el tamaño del modelo. Por ejemplo, los modelos entrenados para predicciones de 4 tokens mostraron mejoras de rendimiento de varios puntos porcentuales sobre las predicciones de un solo token en el benchmark de codificación MBPP. Los investigadores concluyen que "es posible, utilizando los mismos recursos computacionales, lograr un mayor rendimiento de los modelos de lenguaje grande al emplear la predicción multi-token."
Además, la predicción multi-token mejora las velocidades de inferencia, haciendo los modelos hasta tres veces más rápidos en varios tamaños de lote. "El preentrenamiento con predicción multi-token mejora la precisión de las cabezas adicionales en comparación con simplemente ajustar un modelo de predicción del siguiente token, desbloqueando el potencial completo de la decodificación auto-especulativa," explican. El estudio también destaca que la predicción multi-token fomenta que el modelo aprenda patrones a más largo plazo, particularmente en experimentos con "tokenización a nivel de byte," donde cada byte se trata como un único token. En estos casos, la predicción de múltiples bytes superó significativamente a los modelos de un solo byte, lo cual es crucial para aplicaciones que carecen de un vocabulario predefinido.
Direcciones Futuras para la Investigación
A pesar de sus ventajas, la predicción multi-token no está exenta de desafíos. Determinar el número óptimo de tokens a predecir varía según la tarea y el tamaño del modelo. Los investigadores están explorando futuras opciones de investigación, incluyendo técnicas automatizadas para identificar el mejor número de tokens a predecir y la dinámica entre tamaños de vocabulario y estrategias multi-token. Esta investigación ofrece promesas para aplicaciones empresariales, potencialmente proporcionando velocidades de inferencia mejoradas y mayor precisión en tareas generativas como la finalización de código, sin grandes alteraciones en la arquitectura existente de LLM, garantizando compatibilidad con otras técnicas de optimización dentro del marco Transformer.