BlockPilot: Instantie-Adaptief Policy-Leren voor Diffusie-Gebaseerd Speculatief Decoderen

Samenvatting

Speculatieve decodering versnelt inferentie door een lichtgewicht conceptmodel te gebruiken om parallel kandidaat-tokens te genereren, die vervolgens worden geverifieerd door het doelmodel, wat verliesvrije versnelling mogelijk maakt. Recentelijk verbetert diffusie-gebaseerde speculatieve decodering verder de parallelliteit door meerdere tokens per forward pass te genereren via blokniveau-diffusie, waarmee state-of-the-art (SOTA) prestaties worden behaald. Echter, bestaande methoden hanteren een vaste inferentieblokgrootte en nemen een uniforme optimale decoderingstrategie aan voor alle invoer. In dit artikel tonen we aan dat deze aanname suboptimaal is, aangezien de optimale blokgrootte varieert per sample en een cruciale rol speelt in de prestatie van speculatieve decodering. Bovendien vertonen deze waarden een duidelijke lokale structuur, geconcentreerd rond de trainingsblokgrootte, wat het probleem reduceert tot een laagdimensionale en gestructureerde beslissingsruimte. Op basis van deze inzichten stellen we BlockPilot voor, een sample-adaptief beleid dat de optimale blokgrootte voorspelt uit de prefilling-representatie. Specifiek formuleren we blokgrootte-selectie als een lichtgewicht beleidsleerprobleem en stellen we een instantie-adaptief beslissingsmechanisme voor dat de optimale blokgrootte voorspelt op basis van de representatie van de prefilling-fase. De voorspelling wordt slechts eenmaal uitgevoerd na prefilling, wat naadloze integratie mogelijk maakt. Uitgebreide experimenten tonen aan dat onze methode plug-and-play is, minimale overhead introduceert en consistent de efficiëntie verbetert, met een acceptatielengte van 5,92 en een 4,20 keer versnelling op Qwen3-4B bij temperatuur 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.