Jakiro: MoEを介したデカップルされたマルチヘッドによるスペキュレイティブデコーディングのブースティング

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

February 10, 2025

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

cs.AI

要旨

先読みデコーディング（SD）は、大規模言語モデルの推論を加速するために、より小さな下書きモデルを使用して複数のトークンを予測し、それらを後でより大きな対象モデルで並列に検証します。しかし、下書きモデルの容量が限られているため、予測精度を向上させるために木構造サンプリングがしばしば必要とされます。このアプローチの主要な制限を特定しました。すなわち、同じステップの候補は同じ表現から派生しており、多様性が制限され、全体的な効果が低下しています。これを解決するために、Mixture of Experts（MoE）を活用したJakiroを提案します。独立した専門家が多様な予測を生成し、候補者間の相関を効果的に分離します。さらに、初期トークンのための自己回帰デコーディングと、後続の段階のための並列デコーディングを組み合わせたハイブリッド推論戦略を導入し、後者を特徴の対比メカニズムで強化して精度を向上させます。当社の手法は予測精度を大幅に向上させ、推論の高速化を実現します。多様なモデルにわたる包括的な実験により、当社の手法の効果と堅牢性が確認され、先読みデコーディングにおける新たなSOTAが確立されました。当社のコードは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.