BiTA: Sintonización Bidireccional para la Aceleración sin Pérdidas en Modelos de Lenguaje a Gran Escala

BiTA: Bi-Directional Tuning for Lossless Acceleration in Large Language Models

January 23, 2024

Autores: Feng Lin, Hanling Yi, Hongbin Li, Yifan Yang, Xiaotian Yu, Guangming Lu, Rong Xiao

cs.AI

Resumen

Los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) suelen emplear generación autoregresiva durante la inferencia, lo que genera una alta demanda de ancho de banda de memoria y, en consecuencia, una latencia prolongada. Para mitigar esta ineficiencia, presentamos Bi-directional Tuning for Lossless Acceleration (BiTA), un método innovador que acelera los LLMs mediante una generación semi-autoregresiva optimizada y verificación de borradores. Inspirados por el concepto de ajuste de prompts, mejoramos los LLMs con un diseño eficiente en parámetros llamado ajuste bidireccional para potenciar la capacidad de generación semi-autoregresiva. Utilizando decodificación basada en árboles eficiente, los modelos generan candidatos de borradores y los verifican en paralelo, asegurando salidas idénticas a sus contrapartes autoregresivas bajo muestreo codicioso. BiTA funciona como un módulo complementario ligero, impulsando sin problemas la eficiencia de inferencia de los LLMs existentes sin requerir modelos de asistencia adicionales ni incurrir en costos significativos de memoria extra. Al aplicar el BiTA propuesto, LLaMA-2-70B-Chat logra una aceleración de 2.7 veces en el benchmark MT-Bench. Experimentos extensos confirman que nuestro método supera las técnicas de aceleración más avanzadas.

English

Large language models (LLMs) commonly employ autoregressive generation during inference, leading to high memory bandwidth demand and consequently extended latency. To mitigate this inefficiency, we present Bi-directional Tuning for lossless Acceleration (BiTA), an innovative method expediting LLMs via streamlined semi-autoregressive generation and draft verification. Inspired by the concept of prompt tuning, we enhance LLMs with a parameter-efficient design called bi-directional tuning for the capability in semi-autoregressive generation. Employing efficient tree-based decoding, the models perform draft candidate generation and verification in parallel, ensuring outputs identical to their autoregressive counterparts under greedy sampling. BiTA serves as a lightweight plug-in module, seamlessly boosting the inference efficiency of existing LLMs without requiring additional assistance models or incurring significant extra memory costs. Applying the proposed BiTA, LLaMA-2-70B-Chat achieves a 2.7times speedup on the MT-Bench benchmark. Extensive experiments confirm our method surpasses state-of-the-art acceleration techniques.