드롭아웃 기반 토큰 검증으로 가속하는 Speculative Decoding

초록

DropMatch는 대형 언어 모델의 LM 헤드에만 Monte‑Carlo 드롭아웃을 적용해 다중 샘플을 생성하고, 이 샘플들의 경험적 토큰 분포와 초안 모델이 제안한 토큰을 비교해 수용 여부를 결정한다. 훈련·데이터·캘리브레이션이 필요 없으며, 기존 Speculative Decoding 파이프라인에 그대로 삽입해 1.09×~1.33× 정도의 속도 향상을 달성한다.

상세 분석

본 논문은 Speculative Decoding의 핵심 병목인 “수용 길이(acceptance length)”를 확대하기 위해, 목표 모델의 출력 분포를 직접 샘플링하는 새로운 메커니즘을 제안한다. 기존 방식은 초안 모델(draft model)과 목표 모델(target model)의 확률 분포를 정량적으로 비교해 토큰을 수용하거나 거부했지만, 토큰 수준에서의 미세 차이까지도 거부로 이어져 속도 향상에 한계가 있었다. DropMatch는 목표 모델의 마지막 레이어인 LM 헤드에만 Monte‑Carlo 드롭아웃을 적용함으로써, 전체 네트워크의 KV‑cache를 그대로 유지하면서도 K개의 독립적인 확률 분포(p(i)ₜ)를 동시에 얻는다. 이는 기존의 전체 모델에 대한 다중 전방패스와 달리 연산 비용이 미미하고, KV‑cache 정렬 문제를 회피한다는 장점이 있다.

드롭아웃 마스크 m(i)∈{0,1}ᵈ는 Bernoulli(1‑p_drop)로 샘플링되고, 역스케일링을 통해 기대값이 원본 hidden state hₜ와 일치하도록 설계된다. 각 마스크에 대해 동일한 LM‑head 가중치 W를 사용해 로그잇 l(i)ₜ와 토큰 확률 p(i)ₜ=Softmax(l(i)ₜ)를 계산한다. 이렇게 얻은 K개의 샘플은 “클러스터” 형태의 경험적 분포를 형성한다.

수용 기준은 두 단계로 구성된다. ① Naïve 토큰 매칭: 초안 토큰 ŷₜ가 K개의 top‑1 토큰 중 하나와 일치하면 수용한다. 이는 드롭아웃 확률이 낮아 샘플 간 변동이 적을 때 효과적이다. ② Jensen‑Shannon(JS) 발산 기반 기준: 각 샘플의 로그잇 평균을 중심 분포 \bar{p}ₜ로 정의하고, 초안 모델의 분포 \hat{p}ₜ와 \bar{p}ₜ 사이의 JS 발산이 샘플들 간 최대 JS 발산 이하이면 수용한다. 이때 샘플이 고도로 집중(concentrated)된 경우, 다수결(majority) 기준을 추가해 동일 토큰이 다수 헤드에서 나타나면 수용한다.

실험에서는 Llama‑3.1‑70B‑Instruct를 대상으로 K=5, p_drop∈{0.1,0.2,0.3,0.4,0.5}를 변동시켰다. 낮은 p_drop(0.1)에서는 샘플 간 코사인 유사도와 의미 일관성이 높아, 대부분의 초안 토큰이 클러스터 내에 포함되는 것을 확인했다. 또한 HumanEval에서 각 헤드별 Pass@1 점수가 원본 모델에 근접함을 보여, 드롭아웃이 모델 성능을 크게 저하시키지 않음을 입증했다.

DropMatch는 기존 Speculative Decoding(Leviathan 등) 및 최신 가속 기법(EAGLE3)과 조합 가능하며, 추가적인 검증 비용은 LM‑head에 한정된 드롭아웃 연산으로 매우 경량이다. 결과적으로 평균 수용 길이가 증가하고, 전체 추론 속도가 1.09×~1.33×, EAGLE3와 병합 시 추가 1.09×까지 향상된다.