Kangaroo: Decodificación Autospeculativa sin Pérdidas mediante Doble Salida Temprana

Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

April 29, 2024

Autores: Fangcheng Liu, Yehui Tang, Zhenhua Liu, Yunsheng Ni, Kai Han, Yunhe Wang

cs.AI

Resumen

La decodificación especulativa ha demostrado su eficacia para acelerar la inferencia de modelos de lenguaje grandes mientras se mantiene una distribución de muestreo consistente. Sin embargo, el enfoque convencional de entrenar un modelo borrador separado para lograr una tasa de aceptación de tokens satisfactoria puede resultar costoso. Inspirándonos en el concepto de salida temprana, proponemos un novedoso marco de decodificación auto-especulativa llamado Kangaroo, que utiliza una subred superficial fija como modelo borrador automático, mientras que las capas restantes funcionan como el modelo objetivo más grande. Entrenamos un módulo adaptador ligero y eficiente sobre la subred para cerrar la brecha entre la capacidad de representación de la subred y la del modelo completo. Es importante destacar que la latencia de inferencia del modelo borrador automático puede dejar de ser insignificante en comparación con el modelo grande, lo que requiere estrategias para aumentar la tasa de aceptación de tokens mientras se minimizan los pasos de generación del modelo pequeño. Para abordar este desafío, introducimos un mecanismo adicional de salida temprana para generar tokens borradores. Específicamente, detenemos la predicción subsiguiente del modelo pequeño durante la fase de generación una vez que el nivel de confianza para el token actual cae por debajo de un umbral determinado. Experimentos exhaustivos en Spec-Bench demuestran la efectividad de Kangaroo. Bajo verificación de secuencia única, Kangaroo logra aceleraciones de hasta 1.68 veces en Spec-Bench, superando a Medusa-1 con un 88.7% menos de parámetros adicionales (67M en comparación con 591M). El código de Kangaroo está disponible en https://github.com/Equationliu/Kangaroo.

English

Speculative decoding has demonstrated its effectiveness in accelerating the inference of large language models while maintaining a consistent sampling distribution. However, the conventional approach of training a separate draft model to achieve a satisfactory token acceptance rate can be costly. Drawing inspiration from early exiting, we propose a novel self-speculative decoding framework Kangaroo, which uses a fixed shallow sub-network as a self-draft model, with the remaining layers serving as the larger target model. We train a lightweight and efficient adapter module on top of the sub-network to bridge the gap between the sub-network and the full model's representation ability. It is noteworthy that the inference latency of the self-draft model may no longer be negligible compared to the large model, necessitating strategies to increase the token acceptance rate while minimizing the drafting steps of the small model. To address this challenge, we introduce an additional early exiting mechanism for generating draft tokens. Specifically, we halt the small model's subsequent prediction during the drafting phase once the confidence level for the current token falls below a certain threshold. Extensive experiments on the Spec-Bench demonstrate the effectiveness of Kangaroo. Under single-sequence verification, Kangaroo achieves speedups up to 1.68times on Spec-Bench, outperforming Medusa-1 with 88.7\% fewer additional parameters (67M compared to 591M). The code for Kangaroo is available at https://github.com/Equationliu/Kangaroo.