BlockPilot: Aprendizaje de Políticas Adaptativo por Instancia para Decodificación Especulativa Basada en Difusión

Resumen

La decodificación especulativa acelera la inferencia al emplear un modelo borrador ligero para generar tokens candidatos en paralelo, que luego son verificados por el modelo objetivo, permitiendo una aceleración sin pérdidas. Recientemente, la decodificación especulativa basada en difusión mejora aún más el paralelismo al generar múltiples tokens por paso forward mediante difusión a nivel de bloque, logrando un rendimiento de última generación (SOTA). Sin embargo, los métodos existentes adoptan un tamaño de bloque de inferencia fijo y asumen una estrategia de decodificación óptima uniforme para todas las entradas. En este artículo, demostramos que esta suposición es subóptima, ya que el tamaño de bloque óptimo varía entre muestras y juega un papel crítico en el rendimiento de la decodificación especulativa. Además, estos valores exhiben una clara estructura local, concentrándose alrededor del tamaño de bloque de entrenamiento, lo que reduce el problema a un espacio de decisión de baja dimensión y estructurado. Basándonos en estas observaciones, proponemos BlockPilot, una política adaptable a la muestra que predice el tamaño de bloque óptimo a partir de la representación de prellenado. Específicamente, formulamos la selección del tamaño de bloque como un problema de aprendizaje de políticas ligero y proponemos un mecanismo de decisión adaptable a la instancia que predice el tamaño de bloque óptimo basándose en la representación de la etapa de prellenado. La predicción se realiza solo una vez después del prellenado, permitiendo una integración fluida. Extensos experimentos demuestran que nuestro método es plug-and-play, introduce una sobrecarga mínima y mejora consistentemente la eficiencia, logrando una longitud de aceptación de 5.92 y un aumento de velocidad de 4.20 veces en Qwen3-4B bajo temperatura 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.