Рекуррентный черновик для быстрого спекулятивного декодирования в больших языковых моделях

Recurrent Drafter for Fast Speculative Decoding in Large Language Models

March 14, 2024

Авторы: Aonan Zhang, Chong Wang, Yi Wang, Xuanyu Zhang, Yunfei Cheng

cs.AI

Аннотация

В данной статье мы представляем улучшенный подход к спекулятивному декодированию, направленный на повышение эффективности обслуживания больших языковых моделей. Наш метод основан на преимуществах двух устоявшихся техник: классического подхода к спекулятивному декодированию с двумя моделями и более нового подхода с одной моделью, Medusa. Вдохновляясь Medusa, наш метод принимает стратегию с одной моделью для спекулятивного декодирования. Однако наш метод отличается тем, что использует одну легкую концепцию черновика с рекуррентным дизайном зависимостей, сутью похожую на небольшую модель черновика, используемую в классическом спекулятивном декодировании, но без сложностей полной архитектуры трансформера. И благодаря рекуррентной зависимости, мы можем использовать жадный поиск для быстрого фильтрования нежелательных кандидатов с помощью концепции черновика. Результатом является метод, который объединяет простоту дизайна с одной моделью и избегает необходимости создания структуры внимания дерева, зависящей от данных, только для вывода в Medusa. Мы эмпирически демонстрируем эффективность предложенного метода на нескольких популярных открытых языковых моделях, а также проводим всесторонний анализ компромиссов, связанных с принятием этого подхода.

English

In this paper, we introduce an improved approach of speculative decoding aimed at enhancing the efficiency of serving large language models. Our method capitalizes on the strengths of two established techniques: the classic two-model speculative decoding approach, and the more recent single-model approach, Medusa. Drawing inspiration from Medusa, our approach adopts a single-model strategy for speculative decoding. However, our method distinguishes itself by employing a single, lightweight draft head with a recurrent dependency design, akin in essence to the small, draft model uses in classic speculative decoding, but without the complexities of the full transformer architecture. And because of the recurrent dependency, we can use beam search to swiftly filter out undesired candidates with the draft head. The outcome is a method that combines the simplicity of single-model design and avoids the need to create a data-dependent tree attention structure only for inference in Medusa. We empirically demonstrate the effectiveness of the proposed method on several popular open source language models, along with a comprehensive analysis of the trade-offs involved in adopting this approach.