BlockPilot: Instanzadaptives Policy-Lernen für diffusionsbasiertes spekulatives Decodieren

Zusammenfassung

Spekulative Dekodierung beschleunigt die Inferenz, indem ein leichtes Entwurfsmodell parallel Kandidatentoken erzeugt, die anschließend vom Zielmodell verifiziert werden, was eine verlustfreie Beschleunigung ermöglicht. Die jüngst entwickelte diffusionsbasierte spekulative Dekodierung verbessert die Parallelität weiter, indem pro Vorwärtspass mehrere Token durch blockweise Diffusion generiert werden, und erreicht damit den aktuellen Spitzenwert (SOTA). Allerdings verwenden bestehende Methoden eine feste Inferenzblockgröße und gehen von einer einheitlichen optimalen Dekodierungsstrategie über alle Eingaben hinweg aus. In dieser Arbeit zeigen wir, dass diese Annahme suboptimal ist, da die optimale Blockgröße zwischen den Stichproben variiert und eine entscheidende Rolle für die Leistung der spekulativen Dekodierung spielt. Darüber hinaus weisen diese Werte eine klare lokale Struktur auf, die sich um die Trainingsblockgröße konzentriert, wodurch das Problem auf einen niedrigdimensionalen und strukturierten Entscheidungsraum reduziert wird. Basierend auf diesen Erkenntnissen schlagen wir BlockPilot vor, eine stichprobenadaptive Strategie, die die optimale Blockgröße aus der Repräsentation der Vorausfüllungsphase vorhersagt. Konkret formulieren wir die Blockgrößenauswahl als ein leichtgewichtiges Strategielernproblem und schlagen einen instanzadaptiven Entscheidungsmechanismus vor, der die optimale Blockgröße basierend auf der Repräsentation der Vorausfüllungsphase vorhersagt. Die Vorhersage erfolgt nur einmal nach der Vorausfüllung, was eine nahtlose Integration ermöglicht. Umfangreiche Experimente zeigen, dass unsere Methode plug-and-play ist, nur minimalen Overhead verursacht und durchgängig die Effizienz verbessert. Sie erreicht eine Akzeptanzlänge von 5,92 und eine 4,20-fache Beschleunigung auf Qwen3-4B bei einer Temperatur 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.