Jakiro: Impulsionando a Decodificação Especulativa com Múltiplas Cabeças Desacopladas via MoE

Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE

February 10, 2025

Autores: Haiduo Huang, Fuwei Yang, Zhenhua Liu, Yixing Xu, Jinze Li, Yang Liu, Xuanwu Yin, Dong Li, Pengju Ren, Emad Barsoum

cs.AI

Resumo

A descodificação especulativa (SD) acelera a inferência de grandes modelos de linguagem usando um modelo preliminar menor para prever múltiplos tokens, que são então verificados em paralelo pelo modelo alvo maior. No entanto, a capacidade limitada do modelo preliminar muitas vezes requer amostragem baseada em árvore para melhorar a precisão da previsão, onde múltiplos candidatos são gerados em cada etapa. Identificamos uma limitação chave nessa abordagem: os candidatos na mesma etapa são derivados da mesma representação, limitando a diversidade e reduzindo a eficácia geral. Para lidar com isso, propomos o Jakiro, aproveitando a Mistura de Especialistas (MoE), onde especialistas independentes geram previsões diversas, desacoplando efetivamente as correlações entre os candidatos. Além disso, introduzimos uma estratégia de inferência híbrida, combinando descodificação autoregressiva para tokens iniciais com descodificação paralela para estágios subsequentes, e aprimoramos esta última com mecanismo contrastivo em características para melhorar a precisão. Nosso método aumenta significativamente a precisão da previsão e alcança maiores acelerações na inferência. Experimentos extensos em diversos modelos validam a eficácia e robustez de nossa abordagem, estabelecendo um novo estado da arte na descodificação especulativa. Nosso código está disponível em https://github.com/haiduo/Jakiro.

English

Speculative decoding (SD) accelerates large language model inference by using a smaller draft model to predict multiple tokens, which are then verified in parallel by the larger target model. However, the limited capacity of the draft model often necessitates tree-based sampling to improve prediction accuracy, where multiple candidates are generated at each step. We identify a key limitation in this approach: the candidates at the same step are derived from the same representation, limiting diversity and reducing overall effectiveness. To address this, we propose Jakiro, leveraging Mixture of Experts (MoE), where independent experts generate diverse predictions, effectively decoupling correlations among candidates. Furthermore, we introduce a hybrid inference strategy, combining autoregressive decoding for initial tokens with parallel decoding for subsequent stages, and enhance the latter with contrastive mechanism in features to improve accuracy. Our method significantly boosts prediction accuracy and achieves higher inference speedups. Extensive experiments across diverse models validate the effectiveness and robustness of our approach, establishing a new SOTA in speculative decoding. Our codes are available at https://github.com/haiduo/Jakiro.