추측 디코딩 드래프터 학습에서 TVP가 수용 길이를 높인다

무엇인가

추측 디코딩은 작은 드래프터가 여러 토큰을 제안하고 큰 타깃 모델이 병렬 검증해 자기회귀 생성의 병목을 줄인다. 손실 없는 speculative sampling에서 각 드래프트 토큰의 수용 확률은 타깃 분포 p와 드래프트 분포 q의 total variation 거리로 1−TV(p,q)가 되지만, 실제 드래프터는 주로 cross-entropy나 KL로 학습된다. 이 논문은 어떤 목표가 실제 수용 길이에 유리한지 통제 비교하고, TV의 두 약점인 부족 커버 토큰에서의 그래디언트 소실과 수용 길이 목표의 prefix product로 인한 뒤쪽 위치 감독 약화를 해결하는 TVP 손실을 제안한다.

어떻게 동작하나

TVP는 expected acceptance length(EAL)에서 출발한다. EAL은 각 위치의 수용 확률을 곱한 prefix product의 합인데, 이 곱 때문에 뒤쪽 위치의 학습 신호가 지수적으로 약해진다. 저자들은 로그와 Jensen 부등식으로 이를 위치별 항의 합 L0로 바꾸고, prefix가 살아남을 확률에 따라 뒤쪽 패널티를 줄이는 합성 규칙 φ(xy)=φ(x)+d(x)φ(y)를 만족하는 φβ(r)=(1−r^β)/β를 유도한다. 이로부터 Dβ(p,q)=1/β Σv p(v)[1−min{1,q(v)/p(v)}^β]와 Lβ를 얻는다. β=0이면 D0는 −log r 기반 커버리지 손실, β=1이면 D1=TV다. β를 0에서 1로 선형으로 어닐링해(주 실험은 학습 앞 80%) 처음에는 부족 커버리지를 강하게 교정하고 후반에는 수용 지향으로 이동한다. 그래디언트 계수는 p(v)^{1−β}q(v)^β로, β=0에서 p(v), β=1에서 q(v) 쪽으로 보간되며 부족 커버 토큰을 (p/q)^{1−β}만큼 증폭한다.

무엇과 다른가

이론적으로 TVP의 모든 구성원은 EAL 하한을 준다. EAL/K ≥ [1−β C Lβ,h]+^{1/β}이고, β=0이면 EAL/K ≥ e^{−C L0,h}다. 또한 TV ≤ Dβ ≤ min(D0, TV/β)이고 1−EAL/K ≤ L1 ≤ Lβ2 ≤ Lβ1이며, 그래디언트는 ||gβ−g1||1 ≤ 2(1−β)TV로 β가 1에 갈수록 TV 그래디언트에 회복된다. β=1의 TVP는 정확한 EAL 결손이 아니라 가산적 TV 목표라서, prefix가 거의 수용되지 않는 뒤쪽 위치에도 감독을 남긴다.

어떻게 쓰나

실험은 DSpark와 DFlash2, Eagle3 드래프터를 Qwen3-8B, Qwen3-14B, Qwen3.8-27B 타깃에 붙여 진행했다. Open-PerfectBlend에서 평가 프롬프트를 제거하고 타깃이 T=1로 생성한 thinking-on 완성을 썼으며, 수학 GSM8K·MATH·AIME25, 코드 MBPP·HumanEval·LiveCodeBench, 채팅 MT-Bench·Alpaca·Arena-Hard 아홉 벤치마크에서 SGLang speculative sampling의 τ(보정 토큰 포함 검증 단계당 평균 생성 토큰 수)를 측정했다. 파인튜닝은 1 epoch, peak LR 2e-4, 사전학습은 6e-4로 DSpark 5 epoch 또는 DFlash2 2 epoch다. Qwen3-8B DSpark 파인튜닝에서 TVP는 TV 대비 +1.5%, KL 대비 +8.2%, CE 대비 +9.8%, 공개된 head 대비 +25.7%로 최고 τ를 냈고 Qwen3-14B도 같은 순서로 TV 대비 +1.4%였다. Eagle3에서는 TVP가 TV보다 +0.2% 앞섰고, DFlash2 파인튜닝은 TV나 TVP가 공개 head를 2.3~2.9% 개선한 반면 공개 CE 레시피는 동률이거나 악화됐다. 공개 DFlash2 학습 레시피 대비 TVP 개선은 6.2~12.9%로 보고됐다.

전제와 한계

사전학습에서는 TVP가 다음 best인 LK보다 +1.2%, TV보다 +4.2%였다. DFlash2 path selector가 top-16만 보는데도 전체 어휘로 학습한 TVP가 top-16 TVP보다 +1.0%였고, TV는 top-16 제한 영향이 작았다. 드래프트 길이를 B=7에서 B=15로 늘리면 TVP의 TV 대비 이득이 +1.5~+1.9%로 커졌고, likelihood 목표는 더 뒤처졌다(KL −6.1%→−9.7%, CE −7.5%→−11.2%). Block verification에서는 TVP가 BV보다 B=7에서 +1.2%, B=15에서 +0.6%였고 TV 대비 이득은 +1.8%, +2.5%로 커졌다. Greedy decoding에서도 TVP가 TV보다 파인튜닝 +1.6%, 사전학습 +3.6% 앞섰고, 시도한 모든 학습률에서 최고 τ였으며 1~5 epoch 순위도 유지됐다. 세 seed 차이는 최대 0.27%였다.

흥미롭게도 파인튜닝에서는 공개된 어떤 베이스라인도 plain TV를 노이즈 이상으로 이기지 못했고 BV가 +0.1%로 가장 가까웠다. TV와 가까운 LK, TV+KL, TV+CE, BV만 근접했고, LK-log는 adaptive LK보다 2.9% 낮았으며 D-PACE는 TV보다 3.4% 낮았다. Qwen3-8B에서 KL은 TV보다 6.1%, Eagle3에서 7.4% 낮았고, KL 학습 드래프터는 모든 엔트로피 구간에서 forward KL이 가장 낮았지만 드래프트 엔트로피가 TV의 최대 4배였고 타깃 엔트로피 1 nat 미만인 86% 위치에서 수용률이 TV보다 7% 낮았다. 사전학습에서는 TV가 순수 likelihood 목표보다 강했지만 KL을 섞은 모든 베이스라인에는 못 미쳤고, 90% TV+10% KL이 LK(+3.0%)와 노이즈 수준인 +2.8%를 냈다. TVP가 왜 작동하나면 β를 고정한 버전은 어닐링보다 항상 나빴고, 특히 β=0 고정은 파인튜닝에서 TV보다 6.7% 낮아 KL에 가까웠다. 사전학습에서 어닐링 시작이 이를수록 τ가 단조 증가해 고정 β=0.5 +2.3%, 0.5 시작 +2.9%, 0.25 시작 +3.8%, 0 시작 +4.2%였고, 어닐링 방향을 뒤집으면 8.3% 손해였다. TVP는 β가 1에 도달하기 전에 TV+KL과 LK를 추월해 커버리지 단계가 더 나은 해를 만든다고 저자들은 해석한다. 타깃 엔트로피가 가장 낮은 48% 위치에서는 TVP와 TV 수용률이 같고, 이득은 엔트로피가 높을수록 커져 약 +0.02였다. KL은 2 nats 초과에서만 앞서므로 TVP는 KL의 고엔트로피 이점 일부를 저엔트로피 비용 없이 가져온다.

개발자 입장에서 이 논문은 speculative decoding 드래프터를 학습할 때 CE/KL만 보고 선택하면 실제 서빙 수용 길이를 놓칠 수 있음을 보여준다. 파인튜닝 단계에서는 잘 튜닝한 TV 베이스라인을 반드시 비교하고, 무작위 초기화나 크게 어긋난 드래프터를 처음부터 학습할 때는 TV 단독보다 소량의 KL을 섞거나 TVP처럼 커버리지에서 수용으로 넘어가는 스케줄을 검토할 만하다. 다만 이득은 드래프터 구조, 타깃 모델, 드래프트 길이, block verification, greedy/샘플링 설정에 따라 달라지므로 자사 워크로드에서 τ와 수용률을 직접 측정해야 한다. 또한 공개 head와의 비교는 thinking-on 데이터 재학습 여부 때문에 일부 이득이 데이터 분포 차이일 수 있으니, 동일 데이터·동일 옵티마이저·동일 스케줄의 통제 실험으로 확인해야 한다.

저자들이 밝힌 한계는 다음과 같다. Eagle3 실험은 DSpark보다 훨씬 적은 학습 샘플을 봐서 TVP와 TV의 +0.2% 근접이 완전히 학습된 Eagle3를 반영하지 않을 수 있다. TVP는 학습 초반에 더 느리게 올라가 사전학습 첫 1/3 동안 TV+KL과 LK에 뒤지므로 낮은 학습 예산에서는 불리할 수 있고, β 어닐링은 학습 예산을 미리 안다고 가정하므로 adaptive schedule은 열린 문제다. 공개 목표들은 하나의 공유 드래프터와 학습 레시피 안에서 비교해 손실 자체를 분리했을 뿐 각 논문의 전체 시스템을 재현한 것은 아니다. 평가는 T=1과 greedy decoding에 한정됐고 중간 온도는 다루지 않았으며, 속도 향상은 생산 워크로드를 대표하지 않는 고정 concurrency에서 측정됐다.