BlockPilot: Aprendizagem de Política Adaptativa por Instância para Decodificação Especulativa Baseada em Difusão

Resumo

A decodificação especulativa acelera a inferência ao utilizar um modelo rascunho leve para gerar tokens candidatos em paralelo, que são então verificados pelo modelo alvo, possibilitando uma aceleração sem perdas. Recentemente, a decodificação especulativa baseada em difusão aprimora ainda mais o paralelismo ao gerar múltiplos tokens por passagem direta por meio de difusão em nível de bloco, alcançando desempenho de ponta (SOTA). No entanto, os métodos existentes adotam um tamanho de bloco de inferência fixo e assumem uma estratégia de decodificação ótima uniforme para todas as entradas. Neste artigo, mostramos que essa suposição é subótima, pois o tamanho de bloco ótimo varia entre as amostras e desempenha um papel crítico no desempenho da decodificação especulativa. Além disso, esses valores exibem uma estrutura local clara, concentrando-se em torno do tamanho de bloco de treinamento, o que reduz o problema a um espaço de decisão de baixa dimensão e estruturado. Com base nessas percepções, propomos o BlockPilot, uma política adaptativa por amostra que prevê o tamanho de bloco ótimo a partir da representação de pré-preenchimento. Especificamente, formulamos a seleção do tamanho de bloco como um problema de aprendizado de política leve e propomos um mecanismo de decisão adaptativo por instância que prevê o tamanho de bloco ótimo com base na representação do estágio de pré-preenchimento. A previsão é realizada apenas uma vez após o pré-preenchimento, permitindo integração contínua. Experimentos extensos demonstram que nosso método é plug-and-play, introduz uma sobrecarga mínima e melhora consistentemente a eficiência, alcançando um comprimento de aceitação de 5,92 e uma aceleração de 4,20 vezes no Qwen3-4B sob 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.