OmniDraft: Кросс-словарный, адаптивный онлайн-драфтер для спекулятивного декодирования на устройстве

OmniDraft: A Cross-vocabulary, Online Adaptive Drafter for On-device Speculative Decoding

July 3, 2025

Авторы: Ramchalam Kinattinkara Ramakrishnan, Zhaocong Yuan, Shaojie Zhuo, Chen Feng, Yicheng Lin, Chenzheng Su, Xiaopeng Zhang

cs.AI

Аннотация

Спекулятивное декодирование, как правило, предполагает использование небольшой и эффективной черновой модели, которая либо предварительно обучена, либо дистиллирована оффлайн для конкретной серии целевых моделей, например, моделей Llama или Qwen. Однако в условиях онлайн-развертывания возникают две основные проблемы: 1) использование целевой модели, несовместимой с черновой моделью; 2) ожидание улучшения задержек в процессе использования и со временем. В данной работе мы предлагаем OmniDraft — унифицированную структуру, которая позволяет одной черновой модели работать с любой целевой моделью и динамически адаптироваться к пользовательским данным. Мы вводим онлайн n-граммный кэш с гибридной дистилляционной тонкой настройкой для устранения несоответствий в словарях между черновой и целевой моделями; а также дополнительно улучшаем скорость декодирования за счет адаптивных техник чернового прогнозирования. OmniDraft особенно подходит для приложений LLM на устройствах, где стоимость модели, эффективность и пользовательская настройка являются основными точками разногласий. Это подчеркивает необходимость решения вышеуказанных проблем и мотивирует парадигму «один черновик для всех». Мы демонстрируем эффективность структуры OmniDraft, выполняя онлайн-обучение на задачах математического рассуждения, программирования и генерации текста. В частности, OmniDraft позволяет одной модели Llama-68M работать с различными целевыми моделями, включая Vicuna-7B, Qwen2-7B и Llama3-8B, для спекулятивного декодирования; а также обеспечивает ускорение до 1,5–2 раз.

English

Speculative decoding generally dictates having a small, efficient draft model that is either pretrained or distilled offline to a particular target model series, for instance, Llama or Qwen models. However, within online deployment settings, there are two major challenges: 1) usage of a target model that is incompatible with the draft model; 2) expectation of latency improvements over usage and time. In this work, we propose OmniDraft, a unified framework that enables a single draft model to operate with any target model and adapt dynamically to user data. We introduce an online n-gram cache with hybrid distillation fine-tuning to address the cross-vocabulary mismatch across draft and target models; and further improve decoding speed by leveraging adaptive drafting techniques. OmniDraft is particularly suitable for on-device LLM applications where model cost, efficiency and user customization are the major points of contention. This further highlights the need to tackle the above challenges and motivates the ``one drafter for all'' paradigm. We showcase the proficiency of the OmniDraft framework by performing online learning on math reasoning, coding and text generation tasks. Notably, OmniDraft enables a single Llama-68M model to pair with various target models including Vicuna-7B, Qwen2-7B and Llama3-8B models for speculative decoding; and additionally provides up to 1.5-2x speedup.