Jakiro: Mejora de la Decodificación Especulativa con Múltiples Cabeceras Desacopladas a través de MoE

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

February 10, 2025

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

cs.AI

Resumen

La decodificación especulativa (SD) acelera la inferencia de grandes modelos de lenguaje mediante el uso de un modelo preliminar más pequeño para predecir múltiples tokens, los cuales son luego verificados en paralelo por el modelo objetivo más grande. Sin embargo, la capacidad limitada del modelo preliminar a menudo requiere muestreo basado en árboles para mejorar la precisión de la predicción, donde se generan múltiples candidatos en cada paso. Identificamos una limitación clave en este enfoque: los candidatos en el mismo paso se derivan de la misma representación, limitando la diversidad y reduciendo la efectividad general. Para abordar esto, proponemos Jakiro, aprovechando la Mezcla de Expertos (MoE), donde expertos independientes generan predicciones diversas, desacoplando efectivamente las correlaciones entre candidatos. Además, introducimos una estrategia de inferencia híbrida, combinando decodificación autoregresiva para tokens iniciales con decodificación paralela para etapas posteriores, y mejoramos esta última con un mecanismo contrastivo en características para mejorar la precisión. Nuestro método aumenta significativamente la precisión de la predicción y logra mayores aceleraciones en la inferencia. Experimentos extensos en diversos modelos validan la efectividad y robustez de nuestro enfoque, estableciendo un nuevo estado del arte en decodificación especulativa. Nuestros códigos están disponibles en 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.