Jakiro: Steigerung der spekulativen Dekodierung mit entkoppeltem Multi-Head über MoE

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

February 10, 2025

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

cs.AI

Zusammenfassung

Spekulatives Decodieren (SD) beschleunigt die Inferenz großer Sprachmodelle, indem ein kleinerer Entwurf eines Modells verwendet wird, um mehrere Token vorherzusagen, die dann parallel vom größeren Zielmodell überprüft werden. Die begrenzte Kapazität des Entwurfsmodells erfordert jedoch häufig baumbasiertes Sampling, um die Vorhersagegenauigkeit zu verbessern, wobei mehrere Kandidaten in jedem Schritt generiert werden. Wir identifizieren eine wesentliche Einschränkung dieses Ansatzes: Die Kandidaten im selben Schritt stammen aus derselben Repräsentation, was die Vielfalt einschränkt und die Gesamteffektivität verringert. Um dies zu lösen, schlagen wir Jakiro vor, das Mixture of Experts (MoE) nutzt, bei dem unabhängige Experten vielfältige Vorhersagen generieren, um Korrelationen zwischen Kandidaten effektiv zu entkoppeln. Darüber hinaus führen wir eine hybride Inferenzstrategie ein, die autoregressives Decodieren für anfängliche Token mit parallelem Decodieren für nachfolgende Stufen kombiniert und letzteres mit einem kontrastiven Mechanismus in den Merkmalen zur Verbesserung der Genauigkeit verstärkt. Unsere Methode steigert die Vorhersagegenauigkeit erheblich und erzielt höhere Beschleunigungen bei der Inferenz. Umfangreiche Experimente mit verschiedenen Modellen bestätigen die Wirksamkeit und Robustheit unseres Ansatzes und etablieren einen neuen SOTA im spekulativen Decodieren. Unser Code ist verfügbar unter 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.