AMTP는 검증 피드백으로 MTP 초안 길이를 매 라운드 적응적으로 고른다

무엇인가

이 논문은 추측 디코딩(speculative decoding)에서 한 라운드에 몇 개의 토큰을 초안할지, 즉 스페큘레이션 지평(horizon)을 어떻게 정할 것인가를 다룬다. MTP(다중 토큰 예측)는 공유 표현에서 여러 미래 토큰을 한 번의 순전파로 뽑고 타깃 모델 한 번의 패스로 검증한다. 지평이 길면 초안이 신뢰할 만할 때 검증 비용을 더 많은 토큰에 분산시켜 이득이지만, 앞에서 거절되면 버려지는 접미부가 지연만 더한다. 문제는 최적 지평이 문맥마다 달라진다는 것이다. 자기회귀 드래프터는 토큰을 하나 낼 때마다 계속할지 멈출지 정할 수 있지만, MTP는 모든 헤드가 같은 순전파에서 동시에 발화하므로 초안을 뽑기 전에 지평을 고정해야 한다. 저자들은 초안을 다 뽑은 뒤 표본을 보고 잘라내는 방식은 검증 비용만 아끼고 초안 비용은 그대로이며, 완벽한 드래프터에서도 정확성(exactness)을 깨뜨릴 수 있음을 보인다. 공정한 비트를 던지는 예시에서 두 번째 토큰이 1일 확률이 1/2이 아니라 3/4이 되어 버린다.

어떻게 동작하나

핵심 개념은 수용 생존율(acceptance survival)이다. 초안 접두부 Y1:j가 검증에서 전부 수용될 확률을 s(t,j)라 하면, 위치 j의 거절 위험률(rejection hazard) h(t,j)에 대해 s(t,j)는 (1-h)의 누적곱으로 표현되고, 지평 k의 기대 커밋 토큰 수는 G(t,k)=1+Σ s(t,j)가 된다. 즉 위험률 벡터 하나만 알면 모든 지평의 기대 이득을 한꺼번에 계산할 수 있다. 저자들은 이 위험률이 단순 평균 불일치가 아니라 생존 확률로 가중된 조건부 확률이어야 한다는 점을 강조한다. 이진 어휘 예시에서 가중치 없는 평균 불일치는 1/2이지만 실제 위험률은 1/3이고, 기대 이득은 2.25 대 2.125로 갈린다. 또한 토큰별 주변분포가 정확해도 다중 토큰 제안이 정확해지지는 않는다는 점, 독립 MTP 헤드는 미래 토큰 간 의존성을 놓칠 수 있다는 점도 지적한다.

무엇과 다른가

AMTP는 이 위험률을 사전 특징(pre-draft feature)과 위치에서 예측하는 가벼운 제어기다. 예측 위험률을 생존율과 기대 이득으로 변환하고, 라운드 비용 추정치와 함께 처리량 G/C를 계산해 매 라운드 최대화하는 지평을 O(K) 시간에 고른다. 모델 재학습이 없고 정확한 검증기(exact verifier)를 그대로 둔다. 학습 목표는 생존 가중 회귀 손실이며, 보정 항등식(Lemma 1)에 의해 이 손실의 초과분이 정확히 생존 가중 위험률 제곱오차가 된다. 문제는 짧은 지평을 고르면 뒤쪽 위치의 검증 피드백을 아예 관측하지 못한다는 점이다. AMTP는 확률 ε로 전체 윈도를 초안하는 탐색과 역확률 가중(inverse-probability weighting)으로 이 부분 피드백을 편향 없이 보정한다.

어떻게 쓰나

이론적으로는 어떤 온라인 학습기의 예측 오차든 각 문맥에서의 최적 지평 대비 처리량 후회(regret)로 번역되고, 다시 wall-clock 완료 시간 한계로 이어진다는 것을 보인다. 비교 대상이 고정 설정이 아니라 방문한 모든 문맥의 최적 지평이라는 점이 기존 밴딧 계열 연구보다 강한 기준이다. 온라인 경사하강법(OGD)으로 완전 피드백에서는 O(T^{3/4}(1+P_T)^{1/4}), 부분 피드백에서는 O(T^{4/5}(1+P_T)^{1/5})의 준선형 후회를 얻는다. 이로부터 속도향상 하한 S_N ≥ (ρ/r_base)/(1+K/N+c_max·B_N/N)도 유도한다.

전제와 한계

실험은 자기 추측(self-speculative) 3종(자기 증류 Qwen3-4B, Meta의 multi-token-prediction 7B, Llama-3.1-8B-Instruct에 Medusa 헤드를 붙인 구성)과 교차 모델 2종(자기 증류 4B→Qwen3-14B, Meta MTP 7B→Llama-2-13B)을 AIME 2025와 GSM8K에서 평가한다. KV 캐시를 끄고 조건당 단일 비시드 실행으로 측정했으며, 고정 윈도 베이스라인은 각 드래프트가 지원하는 최대 윈도(k=16, 4, 6)를 쓴다. 결과는 표준 디코딩 대비 최대 4.1배 빠른 디코딩, 10개 설정 중 8개에서 고정 윈도와 동등하거나 우수, 처리량 최대 10% 향상이다. 교차 모델 조합에서 이득이 컸다. Meta MTP 7B→Llama-2-13B는 최적 고정 윈도 대비 AIME에서 처리량 268%, GSM8K에서 147% 향상됐고, 자기 증류 4B→Qwen3-14B는 AIME 10% 향상, GSM8K에서는 최적 고정 윈도 대비 1% 이내를 유지했다. 자기 추측 조합은 드래프트와 타깃이 한 순전파를 공유해 지평을 늘리는 추가 비용이 작으므로 이득도 작아, AIME에서 2~7% 향상, GSM8K에서 3% 이내였다. 유일한 역행은 Meta MTP 7B의 AIME(−15%)로, 저자들은 해당 체크포인트가 답을 낸 뒤 안정적으로 멈추지 않아 긴 생성의 라운드 비용이 커진 탓으로 설명한다. 다만 그 설정에서도 수용률 자체는 적응 제어 쪽이 더 높았다.

개발자 관점에서 이 논문의 실용적 요지는 초안 길이를 상수로 두지 말라는 것이다. AMTP는 라운드마다 초안 토큰 수를 줄이면서 커밋하는 토큰 수는 비슷하게 유지해, 낭비되던 드래프트 연산을 완료 토큰으로 바꾼다. 재학습이 필요 없고 검증기를 건드리지 않으므로 기존 MTP/Medusa 스택 위에 제어기만 얹는 형태로 시도할 수 있다. 다만 지평 선택이 초안 표본을 보기 전 정보에만 의존해야 정확성이 유지되므로, 사후에 초안을 잘라내는 최적화와 혼용하면 분포가 틀어질 수 있다는 점을 확인해야 한다.

저자들이 밝힌 한계는 분명하다. 실험은 KV 캐시 없이 시퀀스를 하나씩 디코딩하는 설정이라, 캐시를 쓰는 배치 서빙에서는 지평별 비용 구조가 달라진다. 위험률 기반 예측과 생존율 직접 예측의 비교, 트리 구조 초안으로의 확장, 확산 기반 추측 기법으로의 확장은 향후 과제로 남겨져 있다.