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) ускоряет вывод больших языковых моделей, используя более компактную рабочую модель для предсказания нескольких токенов, которые затем проверяются параллельно более крупной целевой моделью. Однако ограниченная емкость рабочей модели часто требует использования деревянной выборки для улучшения точности прогнозирования, где на каждом шаге генерируется несколько кандидатов. Мы выявляем ключевое ограничение этого подхода: кандидаты на одном и том же шаге происходят из одного и того же представления, что ограничивает разнообразие и снижает общую эффективность. Для решения этой проблемы мы предлагаем Jakiro, использующий Смесь Экспертов (MoE), где независимые эксперты генерируют разнообразные прогнозы, эффективно разрывая корреляции между кандидатами. Кроме того, мы представляем гибридную стратегию вывода, объединяющую авторегрессионное декодирование для начальных токенов с параллельным декодированием для последующих этапов, и улучшаем последнее с помощью контрастного механизма в признаках для повышения точности. Наш метод значительно повышает точность прогнозирования и достигает более высоких ускорений вывода. Обширные эксперименты с различными моделями подтверждают эффективность и надежность нашего подхода, устанавливая новый 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.