추측적 디코딩을 위한 검증 인식 훈련

초록

추측적 디코딩(speculative decoding)은 드래프트 모델(draft model)이 후보 토큰을 생성하면 타깃 모델(target model)이 단일 순전파(forward pass)로 이를 검증함으로써 대규모 언어 모델 추론을 가속화한다. 검증은 순차적으로 진행되며 첫 번째 거부가 발생한 지점부터 이후의 모든 위치를 폐기한다. 그러나 기존의 드래프트 학습은 위치별 고정 가중치를 사용하여 타깃을 토큰 수준에서 모방할 뿐, 이러한 두 가지 특성을 모두 반영하지 못한다. 우리는 모든 학습 단계에서 검증 과정을 시뮬레이션하고 그 결과 나타나는 수용(accept) 및 거부(reject) 패턴을 지도 신호로 변환하는 플러그인 프레임워크인 검증 인식 학습(VAT, Verification-Aware Training)을 제안한다. VAT는 두 가지 구성 요소로 이루어진다. (i) 검증 헤드(verification head)는 각 위치가 순차 검증에서 생존하는지 여부에 대해 드래프트 모델을 지도하는 경량의 공동 학습 이진 분류기이다. (ii) 검증 적응 가중치 부여(verification-adaptive weighting)는 각 샘플의 첫 번째 거부 지점까지는 전체 가중치를 유지하고 감쇠가 해당 지점에서 시작하도록 재설정함으로써, 고정된 가중치 스케줄을 대체한다. VAT는 학습 목표만 수정하므로 드래프트 아키텍처, 타깃 모델, 또는 추론 절차를 변경하지 않고도 기존 방법 위에 적용할 수 있다. VAT를 Qwen3-4B, Qwen3-8B, LLaMA-3.1-8B 기반의 EAGLE-3 및 DFlash에 적용한 결과, 평균 수용 길이(acceptance length)는 최대 11.4%, wall-clock 시간 기준 속도 향상은 최대 8.7% 개선되었으며, 수학, 코드, 채팅 벤치마크 전반에서 일관된 성능 향상을 보였다. 코드는 https://github.com/naver-ai/vat 에서 공개될 예정이다.

English

Speculative decoding accelerates large language model inference by using a draft model to generate candidate tokens, which are verified by the target model in a single forward pass. Verification proceeds sequentially and discards every position from the first rejection onward, yet existing draft training relies on token-level imitation of the target with a fixed per-position weighting that reflects neither property. We introduce Verification-Aware Training (VAT), a plug-in framework that simulates verification at every training step and turns the resulting accept and reject patterns into supervision. VAT consists of two components: (i) a verification head, a lightweight jointly trained binary classifier that supervises the draft model on whether each position survives sequential verification; (ii) verification-adaptive weighting, which replaces the fixed weighting schedule by keeping full weight up to each sample's first rejection point and re-anchoring the decay to start there. VAT modifies only the training objective, so it can be layered on top of existing methods without changing the draft architecture, the target model, or the inference procedure. Applied to EAGLE-3 and DFlash on Qwen3-4B, Qwen3-8B, and LLaMA-3.1-8B, VAT improves average acceptance length by up to 11.4% and wall-clock speedup by up to 8.7%, with consistent gains across math, code, and chat benchmarks. Code will be available at https://github.com/naver-ai/vat