Jakiro : Amélioration du décodage spéculatif avec des têtes multiples découplées via MoE

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

February 10, 2025

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

cs.AI

Résumé

Le décodage spéculatif (SD) accélère l'inférence des grands modèles de langage en utilisant un modèle initial plus petit pour prédire plusieurs jetons, qui sont ensuite vérifiés en parallèle par le modèle cible plus grand. Cependant, la capacité limitée du modèle initial nécessite souvent un échantillonnage basé sur les arbres pour améliorer la précision des prédictions, où plusieurs candidats sont générés à chaque étape. Nous identifions une limitation clé dans cette approche : les candidats à la même étape sont dérivés de la même représentation, limitant la diversité et réduisant l'efficacité globale. Pour y remédier, nous proposons Jakiro, en tirant parti du Mélange d'Experts (MoE), où des experts indépendants génèrent des prédictions diverses, dissociant efficacement les corrélations entre les candidats. De plus, nous introduisons une stratégie hybride d'inférence, combinant le décodage autorégressif pour les jetons initiaux avec un décodage parallèle pour les étapes suivantes, et améliorons ce dernier avec un mécanisme contrastif dans les caractéristiques pour améliorer la précision. Notre méthode améliore significativement la précision des prédictions et atteint des accélérations d'inférence plus élevées. Des expériences approfondies sur divers modèles valident l'efficacité et la robustesse de notre approche, établissant un nouveau record en matière de décodage spéculatif. Nos codes sont disponibles sur 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.