BlockPilot: 확산 기반 추측 디코딩을 위한 인스턴스 적응형 정책 학습

초록

추론적 디코딩은 경량 초안 모델을 사용하여 후보 토큰을 병렬로 생성하고, 이를 목표 모델이 검증함으로써 무손실 가속을 가능하게 한다. 최근 확산 기반 추론적 디코딩은 블록 수준 확산을 통해 단일 순방향 패스에서 여러 토큰을 생성함으로써 병렬성을 더욱 개선하여 최첨단 성능을 달성한다. 그러나 기존 방법들은 고정된 추론 블록 크기를 채택하고 모든 입력에 대해 균일한 최적 디코딩 전략을 가정한다. 본 논문에서는 최적 블록 크기가 샘플에 따라 달라지며 추론적 디코딩 성능에 중요한 역할을 하기 때문에 이러한 가정이 차선임을 보인다. 또한, 이러한 값들은 훈련 블록 크기 주변에 집중되는 명확한 국소 구조를 나타내며, 이는 문제를 저차원의 구조화된 결정 공간으로 축소한다. 이러한 통찰을 바탕으로, 우리는 프리필링 표현으로부터 최적 블록 크기를 예측하는 샘플 적응형 정책인 BlockPilot을 제안한다. 구체적으로, 블록 크기 선택을 경량 정책 학습 문제로 정식화하고, 프리필링 단계의 표현에 기반하여 최적 블록 크기를 예측하는 인스턴스 적응형 결정 메커니즘을 제안한다. 예측은 프리필링 이후에 한 번만 수행되므로 원활한 통합이 가능하다. 광범위한 실험을 통해 우리의 방법이 플러그 앤 플레이 방식이며, 최소한의 오버헤드를 도입하고, 지속적으로 효율성을 개선하여, 온도 T=1에서 Qwen3-4B에 대해 수용 길이 5.92와 4.20배의 속도 향상을 달성함을 입증한다.

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.