BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding

Hao Zhang, Yiming Hu, Yong Wang, Mingqiao Mo, Xin Xiao, Xiangxiang Chu

arXiv:2606.31315 · 2026-07-01 공개 · arXiv · PDF

diffusion-models speculative-decoding policy-learning qwen3-4b token-generation efficiency-optimization instance-adaptive prefilling

Abstract

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.

한국어 요약

한 줄 요약

BlockPilot은 확산 기반 추론 가속 기법에서 입력에 따라 최적 블록 크기를 예측하여 4.20× 가속을 달성한 샘플 적응형 정책 학습 방법이다.

핵심 기여도

핵심 아이디어

기존 확산 기반 추론 가속 방법은 블록 크기를 고정값으로 사용하지만, 이는 입력별 예측 가능성 차이를 고려하지 못해 비효율적이다. BlockPilot은 블록 크기를 샘플에 따라 조정하는 정책 학습 문제로 정식화한다. prefilling 단계에서 생성된 마지막 토큰의 예측 분포를 현재 추론 상태의 표현으로 활용하여, 블록 크기 결정을 샘플별로 최적화한다. 이는 블록 크기 결정을 복잡한 온라인 최적화가 아닌, 작은 이산 공간에서의 분류 문제로 전환시켜 효율성을 높인다. 특히, 최적 블록 크기는 학습 시 설정값 주변에 집중되어 있어, 정책 학습 모델은 이 지역 내에서만 예측하면 충분하다는 점이 핵심 통찰이다.

기술적 접근법

주요 결과

의의 및 한계

BlockPilot은 확산 기반 추론 가속 기법에서 블록 크기 결정을 샘플에 따라 최적화함으로써, 기존 고정 블록 크기 방법의 한계를 극복한다. 이는 추론 효율성 향상에 기여하며, 추론 정책 자체를 학습 가능한 구성 요소로 다루는 새로운 접근을 제시한다. 그러나 BlockPilot은 블록 크기 결정을 지역 후보 집합 내에서만 수행하기 때문에, 매우 특이한 입력에 대해서는 최적 성능을 발휘하지 못할 수 있다. 또한, 정책 학습 모델의 정확도는 prefilling 단계의 예측 분포에 크게 의존하므로, 이 분포의 품질이 낮아지면 성능이 저하될 수 있다.

실용적 활용

BlockPilot은 대규모 언어 모델의 추론 가속에 적용할 수 있으며, 특히 실시간 대화, 코드 생성, 긴 텍스트 생성 등에서 유용하다. 기존 확산 기반 추론 프레임워크에 무해하게 통합 가능하므로, 기존 시스템에 최소한의 변경만으로도 성능 향상을 기대할 수 있다.