Kangaroo: Decodificação Autosspeculativa sem Perdas via Dupla Saída Antecipada

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

Resumo

A decodificação especulativa demonstrou sua eficácia na aceleração da inferência de modelos de linguagem grandes, mantendo uma distribuição de amostragem consistente. No entanto, a abordagem convencional de treinar um modelo de rascunho separado para alcançar uma taxa de aceitação de tokens satisfatória pode ser custosa. Inspirados pelo conceito de saída antecipada, propomos uma nova estrutura de decodificação autoespeculativa chamada Kangaroo, que utiliza uma sub-rede superficial fixa como um modelo de rascunho automático, com as camadas restantes servindo como o modelo alvo maior. Treinamos um módulo adaptador leve e eficiente sobre a sub-rede para preencher a lacuna entre a capacidade de representação da sub-rede e a do modelo completo. É importante notar que a latência de inferência do modelo de rascunho automático pode não ser mais desprezível em comparação com o modelo grande, exigindo estratégias para aumentar a taxa de aceitação de tokens enquanto minimiza as etapas de rascunho do modelo pequeno. Para enfrentar esse desafio, introduzimos um mecanismo adicional de saída antecipada para gerar tokens de rascunho. Especificamente, interrompemos a previsão subsequente do modelo pequeno durante a fase de rascunho assim que o nível de confiança para o token atual cai abaixo de um certo limite. Experimentos extensivos no Spec-Bench demonstram a eficácia do Kangaroo. Sob verificação de sequência única, o Kangaroo alcança acelerações de até 1,68 vezes no Spec-Bench, superando o Medusa-1 com 88,7% menos parâmetros adicionais (67M comparados a 591M). O código do Kangaroo está disponível em 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.