스펙큘러티브 디코딩 — 작은 모델이 먼저 쓰고 큰 모델이 검사한다

실행: python NN_46_speculative.py (검증 환경: torch 2.8.0+cu129, RTX 5080) 이 글의 수치는 전부 그 스크립트를 돌려 얻은 것이다. 이어지는 글: 추론 서빙 편

추론 서빙 편에서 디코드가 계산이 아니라 가중치를 읽어 오는 데 묶여 있다는 것을 봤다(0.3 TFLOP/s). 가중치 전체를 한 번 읽어 오면서 토큰을 딱 하나 얻는 것은 아깝다.

그래서 이런 생각을 한다 — 작은 모델에게 앞으로 몇 글자를 미리 쓰게 하고, 큰 모델은 그것을 한 번에 검사만 하면 어떨까. 검사는 이미 쓰인 토큰들을 한 번에 통과시키는 일이라 프리필처럼 싸다.

1. 어떻게 공짜가 되는가 — 그림으로 먼저

핵심은 ②다. 큰 모델은 초안을 한 번에 통과시키면서 각 자리에서 “내가 원하는 분포”를 동시에 얻는다. 토큰 하나를 위해 부를 때와 같은 한 번이다. 그래서 초안이 받아들여진 만큼이 공짜로 벌린 토큰이다.

검사 규칙은 이렇다. 초안 모델이 $p$ 로 토큰 $t$ 를 뽑았고 큰 모델이 원하는 분포가 $q$ 라면,

\text{확률 } \min\!\Bigl(1,\ \frac{q(t)}{p(t)}\Bigr) \text{ 로 받아들이고,}\quad \text{거절하면 } \frac{\max(0,\ q - p)}{\lVert \max(0,\ q-p)\rVert_1} \text{ 에서 다시 뽑는다.}

기호	뜻
$p$	초안(작은) 모델의 분포
$q$	큰 모델이 그 자리에서 원하는 분포
$t$	초안이 실제로 뽑은 토큰
$k$	한 번에 쓰는 초안의 길이

이 규칙이 왜 중요한지는 4절에서 잰다 — 결과가 근사가 아니라 정확히 같아진다.

2. 직접 재 보기 A — 몇 개나 받아 주나

문자 단위 언어모델을 두 개 학습시켰다. 큰 모델은 4층(388K), 초안 모델은 1층(108K)으로 파라미터가 3.6배 차이다. 초안 길이 $k$ 만 바꿔 가며 큰 모델이 몇 개를 받아 주는지 300번씩 쟀다.

스펙큘러티브 디코딩

초안 길이 $k$	1	2	3	4	6	8
평균 채택 수	0.91	1.66	2.56	2.95	4.10	4.85
전부 통과한 비율	91.3%	78.0%	76.0%	60.0%	53.7%	35.0%

길게 쓸수록 전부 통과할 확률은 떨어지지만, 평균적으로 확정되는 토큰은 늘어난다. $k=8$ 이면 8개 중 4.85개가 살아남고, 나머지는 버린다.

버려지는 초안이 아깝게 느껴지지만 그렇지 않다. 초안 모델은 싸고, 큰 모델을 한 번 부른 값은 어차피 치른 것이다. 실제로 어디까지 받아들여지는지 몇 번 찍어 봤다.

초안이 어디까지 받아들여지나

3. 직접 재 보기 B — 가장 좋은 초안 길이

이득은 확정 토큰 수와 드는 시간의 비다. 한 단계에 확정되는 토큰은 채택 수 $+1$ 이다 (거절된 자리는 큰 모델이 직접 채우므로 최소 한 개는 나온다).

한 번 호출에 큰 모델 1.243 ms, 초안 모델 0.507 ms였다(비 2.45).

초안 길이 $k$	1	2	3	4	6	8
한 단계에 확정되는 토큰	1.91	2.66	3.56	3.95	5.10	5.85
드는 시간	1.750 ms	2.257 ms	2.764 ms	3.271 ms	4.286 ms	5.300 ms
속도 향상	1.36배	1.46배	1.60배	1.50배	1.48배	1.37배

$k=3$ 에서 1.60배로 가장 빠르다. 더 길게 쓰면 초안 비용이 쌓이는데 채택률은 떨어져 손해다.

여기서 초안 모델이 얼마나 싸야 하는지가 나온다. 이 실험의 비는 2.45배인데, 실제 서비스에서는 큰 모델과 초안 모델의 차이가 수십 배라 훨씬 유리하다. 반대로 초안 모델이 너무 작아 채택률이 낮으면 그것대로 손해다 — “충분히 싸면서 충분히 비슷한” 모델을 고르는 문제다.

4. 직접 재 보기 C — 근사가 아니다

가장 중요한 부분이다. 빨라지는 대신 답이 조금 나빠지는 것 아닌가?

먼저 두 모델이 실제로 얼마나 다른지 400자리에서 쟀다 — 평균 총변동거리 0.115, 가장 많이 다른 자리는 0.657이었다. 그 “가장 많이 다른 자리”에서 6만 번을 뽑아 분포를 비교했다.

큰 모델의 참 분포와 얼마나 떨어져 있나	총변동거리
초안 모델의 분포 자체	0.6571
스펙큘러티브로 뽑은 6만 표본	0.0018
큰 모델로 직접 뽑은 6만 표본	0.0035

초안 모델이 0.657만큼 다른 자리인데, 스펙큘러티브 표본은 0.0018로 참 분포와 붙어 있다. 큰 모델로 직접 6만 번 뽑았을 때의 표본 잡음(0.0035)보다도 작다. 즉 구별할 수 없다.

이것이 1절 검사 규칙의 값어치다. 단순히 “비슷하면 받아들인다”가 아니라, 받아들일 확률과 거절했을 때 다시 뽑는 분포를 정확히 맞춰 두어 전체 분포가 보존되게 만든 것이다. 거절 표집(rejection sampling)의 고전적인 논법이 그대로 쓰인다.

그래서 스펙큘러티브 디코딩은 품질을 한 톨도 내주지 않는 최적화다. 양자화나 증류처럼 “조금 나빠지지만 빨라지는” 것들과 성격이 다르다.

5. 흔한 오해와 한계

“초안이 틀리면 품질이 떨어진다” — 아니다(4절). 틀린 초안은 버려지고 큰 모델이 채운다.
“초안을 길게 쓸수록 좋다” — 3절에서 $k=3$ 이 최적이었다. 채택률과 초안 비용의 균형이다.
“작은 모델이면 아무거나 된다” — 채택률이 곧 이득이라, 큰 모델과 비슷하게 쓰는 초안이어야 한다. 그래서 같은 데이터로 학습한 같은 계열 모델을 쓴다.
“항상 빨라진다” — 배치가 크면 얘기가 다르다. 추론 서빙 편에서 봤듯 배치가 차면 계산이 병목이 되는데, 스펙큘러티브는 버리는 계산을 늘리므로 이득이 줄어든다.
이 글의 실험 — 문자 단위 장난감 모델이고 두 모델의 크기 차이가 3.6배뿐이다. 1.60배·0.0018 같은 수는 이 설정의 값이고, 요점은 검사는 싸다, 채택률이 이득을 정한다, 분포는 정확히 보존된다는 구조다.

6. 한 문단 요약

디코드가 가중치를 읽어 오는 데 묶여 있다면, 한 번 읽어 올 때 토큰을 여러 개 확정하면 된다. 작은 초안 모델이 $k$ 개를 미리 쓰고 큰 모델이 한 번에 통과시켜 검사하는 것이 스펙큘러티브 디코딩이다. 초안 3개를 쓰면 평균 2.56개가 받아들여져 한 단계에 3.56토큰이 확정됐고, 초안 비용을 빼고도 1.60배 빨라졌다(초안이 너무 길면 버리는 것이 늘어 다시 나빠진다). 무엇보다 이것은 근사가 아니다. 초안 모델의 분포가 0.657만큼 다른 자리에서도, 받아들일 확률과 거절 시 재표집 분포를 정확히 맞춰 두었기 때문에 최종 표본은 참 분포와 0.0018 — 큰 모델로 직접 뽑은 표본의 잡음보다도 가깝다. 품질을 한 톨도 내주지 않고 얻는 속도라는 점에서, 양자화나 증류와는 성격이 다른 최적화다.