Jakiro: Speculatieve Decodering Verbeteren met Ontkoppelde Multi-Head 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

Samenvatting

Speculatieve decodering (SD) versnelt de inferentie van grote taalmodelmodellen door gebruik te maken van een kleiner conceptmodel om meerdere tokens te voorspellen, die vervolgens parallel worden geverifieerd door het grotere doelmodel. De beperkte capaciteit van het conceptmodel vereist echter vaak op boomgebaseerde steekproeven om de voorspellingsnauwkeurigheid te verbeteren, waarbij meerdere kandidaten bij elke stap worden gegenereerd. Wij identificeren een belangrijke beperking in deze aanpak: de kandidaten in dezelfde stap zijn afgeleid van dezelfde representatie, wat de diversiteit beperkt en de algehele effectiviteit vermindert. Om dit aan te pakken, stellen wij Jakiro voor, dat gebruikmaakt van Mixture of Experts (MoE), waarbij onafhankelijke experts diverse voorspellingen genereren, waardoor de correlaties tussen kandidaten effectief worden ontkoppeld. Bovendien introduceren wij een hybride inferentiestrategie, waarbij autoregressieve decodering voor initiële tokens wordt gecombineerd met parallelle decodering voor latere fasen, en versterken wij de laatste met een contrastief mechanisme in kenmerken om de nauwkeurigheid te verbeteren. Onze methode verbetert de voorspellingsnauwkeurigheid aanzienlijk en behaalt hogere inferentieversnellingen. Uitgebreide experimenten met diverse modellen valideren de effectiviteit en robuustheid van onze aanpak, waarmee een nieuwe SOTA in speculatieve decodering wordt gevestigd. Onze codes zijn beschikbaar op 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.