Jakiro: Potenziare la decodifica speculativa con Multi-Head disaccoppiato tramite MoE

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

February 10, 2025

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

cs.AI

Abstract

Il decoding speculativo (SD) accelera l'inferenza dei grandi modelli linguistici utilizzando un modello di bozza più piccolo per prevedere più token, che vengono successivamente verificati in parallelo dal modello target più grande. Tuttavia, la capacità limitata del modello di bozza spesso richiede il campionamento basato su alberi per migliorare l'accuratezza delle previsioni, generando più candidati ad ogni passaggio. Identifichiamo una limitazione chiave in questo approccio: i candidati allo stesso passaggio derivano dalla stessa rappresentazione, limitando la diversità e riducendo l'efficacia complessiva. Per affrontare ciò, proponiamo Jakiro, sfruttando il Mixture of Experts (MoE), dove esperti indipendenti generano previsioni diverse, separando efficacemente le correlazioni tra i candidati. Inoltre, introduciamo una strategia di inferenza ibrida, combinando il decoding autoregressivo per i token iniziali con il decoding parallelo per le fasi successive, potenziando quest'ultimo con un meccanismo contrastivo nelle caratteristiche per migliorare l'accuratezza. Il nostro metodo aumenta significativamente l'accuratezza delle previsioni e ottiene maggiori accelerazioni nell'inferenza. Estesi esperimenti su diversi modelli convalidano l'efficacia e la robustezza del nostro approccio, stabilendo un nuovo SOTA nel decoding speculativo. I nostri codici sono disponibili su 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.