BlockPilot: 확산 기반 투기적 디코딩 (Diffusion-based Speculative Decoding)을 위한 인스턴스 적응형
요약
BlockPilot은 확산 기반 투기적 디코딩의 효율성을 높이기 위해 입력 샘플마다 최적의 블록 크기를 예측하는 인스턴스 적응형 정책을 제안합니다. 프리필링 표현을 활용해 최적의 블록 크기를 결정함으로써 추가 오버헤드 없이 추론 속도를 크게 향상시킵니다.
핵심 포인트
- 기존 고정 블록 크기 방식의 비효율성 지적
- 프리필링 표현 기반의 샘플 적응형 블록 크기 예측 메커니즘 제안
- 플러그 앤 플레이 방식으로 기존 모델에 쉽게 통합 가능
- Qwen3-4B 모델에서 약 4.2배의 속도 향상 달성
투기적 디코딩 (Speculative decoding)은 경량화된 초안 모델 (draft model)을 사용하여 후보 토큰들을 병렬로 생성하고, 이를 타겟 모델이 검증함으로써 손실 없는 가속화를 가능하게 하여 추론을 가속화합니다. 최근에는 확산 기반 투기적 디코딩 (diffusion-based speculative decoding)이 블록 수준의 확산 (block-level diffusion)을 통해 한 번의 순전파 (forward pass) 당 여러 개의 토큰을 생성함으로써 병렬성을 더욱 향상시켜 최첨단 (SOTA) 성능을 달성하고 있습니다. 그러나 기존 방법들은 고정된 추론 블록 크기를 채택하며, 모든 입력에 대해 균일한 최적 디코딩 전략이 적용된다고 가정합니다. 본 논문에서는 이러한 가정이 최적이 아님을 보여줍니다. 최적의 블록 크기는 샘플마다 다르며 투기적 디코딩 성능에 결정적인 역할을 하기 때문입니다. 더욱이, 이러한 값들은 훈련된 블록 크기 주변에 집중되는 명확한 국소적 구조 (local structure)를 보이며, 이는 문제를 저차원의 구조화된 결정 공간 (decision space)으로 축소시킵니다. 이러한 통찰을 바탕으로, 우리는 프리필링 표현 (prefilling representation)으로부터 최적의 블록 크기를 예측하는 샘플 적응형 정책인 BlockPilot을 제안합니다. 구체적으로, 우리는 블록 크기 선택을 경량화된 정책 학습 (policy learning) 문제로 정식화하고, 프리필링 단계의 표현을 기반으로 최적의 블록 크기를 예측하는 인스턴스 적응형 결정 메커니즘을 제안합니다. 예측은 프리필링 이후 단 한 번만 수행되므로 원활한 통합이 가능합니다. 광범위한 실험을 통해 우리의 방법이 플러그 앤 플레이 (plug-and-play) 방식이며, 최소한의 오버헤드만을 도입하면서도 효율성을 지속적으로 개선함을 입증하였습니다. 온도 $T=1$ 조건의 Qwen3-4B 모델에서 5.92의 수락 길이 (acceptance length)와 4.20$ imes$의 속도 향상을 달성했습니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.CL (NLP)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기