BlockPilot: Адаптивное к экземплярам обучение политики для спекулятивного декодирования на основе диффузии

Аннотация

Спекулятивное декодирование ускоряет инференс за счет использования легковесной черновой модели для параллельной генерации токенов-кандидатов, которые затем проверяются целевой моделью, обеспечивая без потерь ускорение. Недавно спекулятивное декодирование на основе диффузии дополнительно улучшило параллелизм, генерируя несколько токенов за один прямой проход с помощью блочной диффузии, что позволило достичь передовой производительности (SOTA). Однако существующие методы используют фиксированный размер блока при инференсе и предполагают единую оптимальную стратегию декодирования для всех входных данных. В данной работе мы показываем, что это предположение не является оптимальным, поскольку оптимальный размер блока варьируется в зависимости от выборки и играет ключевую роль в производительности спекулятивного декодирования. Более того, эти значения демонстрируют четкую локальную структуру, концентрируясь вокруг размера блока, используемого при обучении, что сводит задачу к низкоразмерному и структурированному пространству решений. Основываясь на этих выводах, мы предлагаем BlockPilot — адаптивную к выборкам стратегию, которая предсказывает оптимальный размер блока на основе представления фазы префиллинга. В частности, мы формулируем выбор размера блока как задачу обучения легковесной политики и предлагаем инстансно-адаптивный механизм принятия решений, предсказывающий оптимальный размер блока на основе представления фазы префиллинга. Предсказание выполняется только один раз после префиллинга, что обеспечивает бесшовную интеграцию. Обширные эксперименты показывают, что наш метод является подключаемым ("plug-and-play"), вносит минимальные накладные расходы и последовательно повышает эффективность, достигая длины принятия 5.92 и ускорения в 4.20 раза на Qwen3-4B при температуре T=1.

English

Speculative decoding accelerates inference by using a lightweight draft model to generate candidate tokens in parallel, and are then verified by the target model, enabling lossless acceleration. Recently, diffusion-based speculative decoding further improves parallelism by generating multiple tokens per forward pass via block-level diffusion, achieving state-of-the-art (SOTA) performance. However, existing methods adopt a fixed inference block size and assume a uniform optimal decoding strategy across all inputs. In this paper, we show that this assumption is suboptimal, as the optimal block size varies across samples and plays a critical role in speculative decoding performance. Moreover, these values exhibit a clear local structure, concentrating around the training block size, which reduces the problem to a low-dimensional and structured decision space. Based on these insights, we propose BlockPilot, a sample-adaptive policy that predicts the optimal block size from the prefilling representation. Specifically, we formulate block size selection as a lightweight policy learning problem and propose an instance-adaptive decision mechanism that predicts the optimal block size based on the representation of the prefilling stage. The prediction is performed only once after prefilling, allowing for seamless integration. Extensive experiments demonstrate that our method is plug-and-play, introduces minimal overhead, and consistently improves efficiency, achieving an acceptance length of 5.92 and a 4.20times speedup on Qwen3-4B under temperature T=1.