Consistency LLM: convertir un LLM en un decodificador paralelo mejora 3.5x la velocidad de inferencia

Consistency Large Language Models (CLLM)

• Tradicionalmente, los LLM se han considerado decodificadores secuenciales que funcionan decodificando un token de manera continua a la vez
• Este blog muestra que los LLM preentrenados pueden convertirse fácilmente en decodificadores paralelos eficientes
• Consistency Large Language Models (CLLM) es una nueva familia de decodificadores paralelos capaz de decodificar eficientemente secuencias de n tokens en cada paso de inferencia para reducir la latencia
• CLLM imita el proceso cognitivo humano de formar una oración completa en la mente antes de pronunciar palabra por palabra, y puede entrenarse de forma efectiva con solo hacer fine-tuning a un LLM preentrenado
• CLLM se entrena realizando decodificación paralela para mapear una secuencia de n tokens inicializada aleatoriamente al mismo resultado que la decodificación autoregresiva (autoregressive, AR), en la menor cantidad de pasos posible
• Según los resultados experimentales, CLLM mostró mejoras de velocidad de generación de entre 2.4x y 3.4x frente a los decodificadores AR existentes, con un rendimiento comparable o superior al de técnicas de inferencia rápida como Medusa2 o Eagle
• CLLM puede lograr estas mejoras de rendimiento sin costos adicionales de memoria

Antecedentes y limitaciones de la decodificación Jacobi

• Un LLM genera tokens uno por uno con decodificación AR, por lo que las respuestas largas provocan una latencia elevada
• La decodificación Jacobi se deriva de los métodos Jacobi y Gauss-Seidel para resolver ecuaciones no lineales, y se ha demostrado que es equivalente a la generación AR con greedy sampling
• La decodificación Jacobi reestructura el proceso secuencial de generación como un sistema de n ecuaciones no lineales con n variables basado en iteraciones de Jacobi, lo que permite el procesamiento en paralelo
• En cada paso iterativo, puede predecir uno o más tokens correctos, por lo que potencialmente puede acelerar la decodificación AR
• Sin embargo, en la práctica, como los LLM entrenados con AR casi no pueden generar el token correcto cuando hay errores en los tokens previos, en la mayoría de las iteraciones de Jacobi solo se obtiene una corrección para la secuencia de n tokens, lo que produce trayectorias (trajectories) de Jacobi más largas
• La decodificación lookahead o speculative decoding busca aliviar esta ineficiencia de la decodificación Jacobi, pero implica costos adicionales de memoria durante la inferencia, mientras que CLLM no

Método de entrenamiento de CLLM

• El entrenamiento de CLLM se compone en gran medida de dos partes: preparación de trayectorias de Jacobi y optimización de las pérdidas de consistency y AR
• En la etapa de preparación de trayectorias de Jacobi, se realiza decodificación Jacobi secuencial recortando de n en n hasta que se generan los l tokens de la secuencia completa de respuesta, y la secuencia generada en cada trayectoria se considera un elemento de datos
• Durante el entrenamiento, se optimizan conjuntamente la pérdida de consistency y la pérdida AR; la pérdida de consistency garantiza la predicción de varios tokens a la vez y la pérdida AR evita que CLLM se desvíe del LLM objetivo, para mantener la calidad de generación
• La pérdida de global consistency (GC) minimiza la distancia entre un punto arbitrario de la trayectoria de Jacobi y el punto fijo, alentando a CLLM a predecir el punto fijo desde cualquier punto de la trayectoria
• La pérdida de local consistency (LC) induce a que estados adyacentes de la trayectoria de Jacobi produzcan la misma salida
• La pérdida AR incluye la pérdida AR tradicional basada en los resultados de generación del LLM objetivo, y su objetivo es evitar que CLLM se desvíe de la distribución del LLM objetivo

Resultados experimentales

• Los experimentos incluyeron tareas de dominios especializados como Spider (text-to-SQL), Human-Eval (completado de código en Python) y GSM8k (matemáticas), así como desafíos amplios de diálogo abierto como MT-bench
• CLLM mostró la mayor mejora de velocidad frente al modelo objetivo, y logró mejoras similares o mejores que Medusa2 sin costo adicional durante la inferencia
• En MT-bench, CLLM logró una mejora de velocidad casi igual a la obtenida al combinarlo con Medusa2, pero ofreciendo mayor adaptabilidad y eficiencia de memoria
• El costo de fine-tuning de CLLM es moderado y, incluso cuando el tamaño del dataset es grande, usar alrededor del 10% del dataset para generar trayectorias de Jacobi permitió obtener una mejora de velocidad de aproximadamente 2.5x
• CLLM puede predecir correctamente varios tokens consecutivos en una sola iteración de Jacobi mediante el fenómeno de fast forwarding
• CLLM muestra una capacidad de tokens stationary, manteniendo sin cambios tokens correctamente predichos por adelantado incluso si hay errores en los tokens previos
• A través del entrenamiento, CLLM aprende conceptos lingüísticos clave como las colocaciones (collocation), lo que le permite inferir la estructura desde cualquier punto de la trayectoria de Jacobi y predecir varias palabras simultáneamente para minimizar los pasos iterativos

La opinión de GN⁺

• CLLM parece resolver de forma efectiva el problema de la alta latencia del esquema de decodificación AR en los LLM existentes mediante el uso de decodificación Jacobi. En particular, resulta llamativo que haya logrado acelerar la velocidad mediante decodificación paralela sin costos adicionales de memoria

• El método de entrenamiento de CLLM parece relativamente simple, ya que consiste en hacer fine-tuning a un LLM existente mediante pérdida de consistency, pero resulta significativo que con ello haya aprendido colocaciones (collocation), una de las propiedades importantes del lenguaje, y que eso haya mejorado considerablemente el rendimiento de la decodificación paralela

• Aun así, como CLLM asume greedy sampling, parece necesario investigar más si también puede funcionar bien con estrategias de decoding más diversas. Además, por ahora los resultados experimentales están limitados al inglés, por lo que también haría falta validar la posibilidad de generalización a distintos idiomas

• CLLM parece un enfoque práctico para acelerar la velocidad de respuesta de los LLM. Da la impresión de que podría aplicarse bien a tareas que requieren respuesta en tiempo real, como búsqueda web o chatbots

• En lo personal, da expectativa que el enfoque de entrenamiento con consistency de CLLM pueda aplicarse no solo a LLM como GPT, sino también a otros modelos generativos como modelos de generación de imágenes o de síntesis de voz. Ojalá la idea de CLLM pueda contribuir en adelante a mejorar la eficiencia de diversos modelos generativos

[GN⁺]

Opiniones de Hacker News

• En clases de dibujo libre, practicar repetidamente con límite de tiempo mejora mucho la habilidad para dibujar. Practicar dibujando rápido incluso ayuda a desarrollar la capacidad de captar con precisión las proporciones y los contornos.
• Jacobi decoding es igual que greedy autoregressive decoding, pero en la práctica hay que establecer la temperatura de muestreo por encima de 0 para evitar repeticiones y respuestas demasiado genéricas.
• Pronto nos daremos cuenta de que no hace falta entrenar modelos. Solo se necesita un buen indexado y muestreo. Los LLM son básicamente equivalentes a una BD de datasets con una excelente interfaz de NLP.
• La mejora en rendimiento de inferencia del paper de CLLM es notable. El costo de fine-tuning es razonable (alrededor de 0.01% del costo de preentrenamiento) y la mejora de rendimiento también es bastante consistente.
• Como las Jacobi trajectories excluyen rutas de alta temperatura, pueden ser positivas para la recuperación de datos, pero negativas para maximizar la creatividad.
• Estaría bien tener una función de "preguntar a un experto en IA" sobre por qué los LLM no responden de forma determinista al mismo prompt.
• Me pregunto si la mejora de velocidad de Groq (https://groq.com/) se debe a este método o si usan otro distinto.