BlockPilot: 확산 기반 스펙큘레이티브 디코딩을 위한 샘플 적응형 정책 학습

대형 언어 모델의 추론은 토큰 단위의 순차적 특성 때문에 병렬화에 한계가 있고, speculative decoding은 이 병목을 해소해 추론 속도를 높이는 유력한 접근법이다. 기존 확산 기반 블록 제안 방법은 추론 시 고정된 블록 크기를 사용해 입력별로 다른 예측 난이도를 반영하지 못했다. BlockPilot은 prefilling 단계의 마지막 토큰 예측분포를 이용해 샘플별 최적 블록 크기를 예측함으로써 병렬성 활용을 개선하고 실사용 환경에서 유의미한 속도 향상을 달성했다.

스펙큘레이티브 디코딩의 효율성은 한 번의 검증 단계에서 수락되는 평균 토큰 수인 acceptance length에 의해 지배된다. diffusion 기반의 블록 제안은 한 forward로 여러 토큰을 생성해 병렬성을 크게 높이지만, 블록 크기 B가 커질수록 예측 오차가 누적되어 acceptance가 떨어질 가능성이 있다. 따라서 효율은 draft/verify의 비용과 τ(B)의 상호작용에 의해 결정되므로 적절한 B 선택이 중요하다.

논문은 우선 블록 크기 후보 집합을 훈련에 사용한 블록 크기 B를 중심으로 한 지역 이웃 {B−k,...,B+k}로 제한해 문제를 정식화했다. 각 샘플에 대해 prefilling 단계의 마지막 위치에서 예측 분포 p(x)를 추출하고 오프라인 열거(enumeration)로 각 후보 b에 대한 실제 acceptance length τ(b;x)를 측정해 최적 블록 크기 B*(x)를 레이블로 확보했다. 이후 이 (p(x),B*(x)) 쌍을 사용해 얕은 MLP 분류기를 학습해 추론 시 한 번의 호출로 인스턴스별 블록 크기를 예측하도록 했다.

광범위한 실험에서 BlockPilot은 DFlash 등 고정 블록 기법을 일관되게 능가했다. Qwen3-4B 기준으로 온도 1에서 평균 acceptance length 5.92와 4.20× 속도향상을 기록했으며, Qwen3-8B에서도 유의한 향상을 보였다. 어블레이션에서 예측기 깊이·폭·입력 전처리·반경 k의 선택이 성능에 영향을 주었고, 기본 설정(L=2, hidden=2048, k=2, 원본 분포 입력)이 균형점으로 채택되었다.

전체 아키텍처는 대상 LLM의 prefilling 단계에서 마지막 토큰의 예측 확률 분포 p(x)를 추출하고 이를 입력으로 받는 얕은 n-layer MLP 분류기로 구성된다. 이 MLP는 입력 확률 벡터에 대해 softmax 기반의 후보 블록 크기 집합 ℬ에 대한 logits를 출력하며 argmax로 최종 블록 크기를 선택한다. 수식적으로 예측기는 P(b|x)=exp(o_b)/Σ_{b'∈ℬ}exp(o_{b'})를 계산하고 교차엔트로피 손실을 최소화하며 학습된다.

논문은 후보 반경 k를 확장하면 후보 커버리지는 늘어나지만 분류 난이도도 증가해 예측 성능이 저하될 수 있음을 실험적으로 보였다. 오프라인 레이블 생성은 후보별 speculative decoding을 병렬화하면 실용적이나 초기 데이터 수집 비용이 존재한다는 점이 언급되었다. 또한 diffusion 기반 drafter 자체가 많은 denoising 단계 또는 대형 초안 모델을 요구하는 기존 접근의 한계를 완전히 제거하지는 못한다는 배경이 기술되었다.