BlockPilot : Apprentissage de politiques adaptatif aux instances pour le décodage spéculatif basé sur la diffusion

Résumé

Le décodage spéculatif accélère l'inférence en utilisant un modèle de brouillon léger pour générer des tokens candidats en parallèle, qui sont ensuite vérifiés par le modèle cible, permettant une accélération sans perte. Récemment, le décodage spéculatif basé sur la diffusion a encore amélioré le parallélisme en générant plusieurs tokens par passage avant via une diffusion au niveau des blocs, atteignant des performances de pointe (SOTA). Cependant, les méthodes existantes adoptent une taille de bloc d'inférence fixe et supposent une stratégie de décodage optimale uniforme pour toutes les entrées. Dans cet article, nous montrons que cette hypothèse est sous-optimale, car la taille de bloc optimale varie selon les échantillons et joue un rôle critique dans les performances du décodage spéculatif. De plus, ces valeurs présentent une structure locale claire, se concentrant autour de la taille de bloc d'entraînement, ce qui réduit le problème à un espace de décision structuré et de faible dimension. Sur la base de ces observations, nous proposons BlockPilot, une politique adaptative à l'échantillon qui prédit la taille de bloc optimale à partir de la représentation du préremplissage. Plus précisément, nous formulons la sélection de la taille de bloc comme un problème d'apprentissage de politique léger et proposons un mécanisme de décision adaptatif à l'instance qui prédit la taille de bloc optimale à partir de la représentation de la phase de préremplissage. La prédiction n'est effectuée qu'une seule fois après le préremplissage, permettant une intégration transparente. Des expériences approfondies démontrent que notre méthode est prête à l'emploi, introduit une surcharge minimale et améliore constamment l'efficacité, atteignant une longueur d'acceptation de 5,92 et une accélération de 4,20 fois sur Qwen3-4B à température 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.