Fast Inference from Transformers via Speculative Decoding

논문 개요

대규모 언어 모델(LLM)의 자기회귀(autoregressive) 생성은 토큰을 하나씩 순차적으로 생성하기 때문에 추론 속도가 느리다. 모델 크기가 커질수록 단일 포워드 패스(forward pass)의 지연이 증가하며, 이는 실시간 대화형 애플리케이션에서 심각한 병목이 된다.

Leviathan et al.(2023)이 ICML 2023에서 발표한 추측 디코딩(Speculative Decoding)은 이 문제를 새로운 방식으로 해결한다. 핵심 통찰은 LLM 추론의 병목이 계산량(FLOPs)보다 메모리 대역폭과 순차적 실행에 있다는 점이다. GPU는 작은 배치나 큰 배치 모두 비슷한 시간 내에 처리할 수 있으므로, 여러 토큰을 동시에 검증하는 것이 토큰 하나를 검증하는 것과 거의 같은 시간이 걸린다. 추측 디코딩은 이 병렬성을 활용한다.

---

핵심 기여

1. 추측 디코딩 알고리즘: 초안 모델로 $\gamma$개 토큰을 생성하고 대상 모델이 이를 병렬로 검증하는 완전한 알고리즘을 제안한다.
2. 이론적 동등성 증명: 제안된 거부 샘플링 방식이 대상 모델에서 직접 샘플링한 것과 동일한 분포를 생성함을 수학적으로 증명한다.
3. 속도-품질 무손실 향상: 출력 분포가 보장되므로 모델 품질의 어떠한 저하도 없이 순수하게 추론 속도만 향상된다.

---

방법론 상세

문제 설정

대상 모델 $M_p$와 초안 모델 $M_q$를 정의한다. 두 모델은 동일한 어휘(vocabulary)를 사용하지만, $M_p\gg M_q$(파라미터 수 기준)이다. 주어진 프롬프트 $x_{1:n}$에 대해 $M_p$의 분포 $p(x_{n+1}|x_{1:n})$로부터 샘플링하되, $M_q$를 활용하여 속도를 향상하는 것이 목표이다.

추측 디코딩 알고리즘

단계 1: 초안 생성

초안 모델 $M_q$를 사용하여 자기회귀적으로 $\gamma$개의 토큰을 생성한다: ${\tilde{x}}_{n+1}\sim q(\cdot |x_{1:n}),{\tilde{x}}_{n+2}\sim q(\cdot |x_{1:n},{\tilde{x}}_{n+1}),\dots ,{\tilde{x}}_{n+\gamma }\sim q(\cdot |x_{1:n+\gamma -1})$

단계 2: 병렬 검증

대상 모델 $M_p$를 사용하여 프롬프트와 $\gamma$개의 초안 토큰을 한 번의 포워드 패스로 처리한다: $p(\cdot |x_{1:n}),p(\cdot |x_{1:n+1}),\dots ,p(\cdot |x_{1:n+\gamma })$

이 $\gamma +1$개의 분포를 단일 배치로 병렬 계산한다.

단계 3: 거부 샘플링으로 토큰 결정

각 초안 토큰 ${\tilde{x}}_{n+i}$에 대해 수용/거부를 결정한다. 토큰 $\tilde{x}$를 수용할 확률은:

${\alpha }_i=\min \left(1,\frac{p({\tilde{x}}_{n+i}|x_{1:n+i-1})}{q({\tilde{x}}_{n+i}|x_{1:n+i-1})}\right)$

즉, 대상 모델이 초안 모델보다 더 높은 확률을 부여하는 토큰($p>q$)은 항상 수용하고, 낮은 확률을 부여하는 토큰은 확률 $p/q$로 수용한다. 거부된 경우, 첫 거부 위치에서 다음 토큰을 수정된 분포로 샘플링한다:

$x^{\prime }\sim \text{norm}(\max (0,p(\cdot )-q(\cdot )))$

모든 $\gamma$개가 수용된 경우, 추가로 한 토큰을 $p(\cdot |x_{1:n+\gamma })$에서 샘플링한다.

이론적 동등성 증명

정리: 위 알고리즘으로 샘플링된 토큰의 분포는 대상 모델 $M_p$에서 직접 샘플링한 분포와 동일하다.

증명 스케치: 토큰 $x$가 최종적으로 선택될 확률을 계산하면:

수용 경로에서의 확률: $P(\text{수용, 토큰}=x)=q(x)\cdot \min \left(1,\frac{p(x)}{q(x)}\right)=\min (q(x),p(x))$

거부 경로에서의 확률: $P(\text{거부 후, 토큰}=x)=\left(1-{\sum }_x\min (p(x),q(x))\right)\cdot \frac{\max (0,p(x)-q(x))}{{\sum }_x\max (0,p(x)-q(x))}$

두 경로를 합산하면 $P(\text{토큰}=x)=p(x)$가 성립한다. $\square$

예상 수용 길이

한 라운드에서 수용되는 토큰 수의 기대값 $E[\text{수용 수}]$는 초안 모델과 대상 모델의 분포 유사도에 따라 결정된다:

$E[\text{수용 수}]=\frac{1-{\alpha }^{\gamma +1}}{1-\alpha }$

여기서 $\alpha ={\mathbb{E}}_{x\sim q}\left[\min \left(1,\frac{p(x)}{q(x)}\right)\right]$은 평균 수용률이다. 두 모델이 유사할수록 $\alpha o1$이 되어 거의 모든 초안 토큰이 수용된다.

---

실험 결과

속도 향상

대상 모델	초안 모델	속도 향상
Chinchilla 70B	Chinchilla 7B	2.1–3.0×
T5-XXL (11B)	T5-Small (60M)	2.5×
PaLM 540B	PaLM 62B	2.2×

품질 검증

대상 모델 직접 생성과 추측 디코딩을 통한 생성의 출력 분포가 수학적으로 동일하므로, BLEU, ROUGE, perplexity 등 어떤 지표에서도 차이가 없다. 실험적으로도 수백 개의 생성 샘플에서 통계적으로 유의미한 품질 차이가 관찰되지 않았다.

초안 길이 $\gamma$의 영향

$\gamma$가 커질수록 잠재적 속도 이득은 증가하지만 거부 확률도 증가한다. 실험에서 $\gamma =4\sim 8$ 범위가 대부분의 설정에서 최적으로 나타났다.

---

의의 및 한계

의의

• 무손실 가속: 기존의 양자화(quantization), 증류(distillation) 등은 품질 저하를 수반하는 반면, 추측 디코딩은 이론적으로 완전히 동등한 출력을 보장한다.
• 실용적 적용성: 기존 대상 모델과 초안 모델을 수정하지 않고 그대로 사용하므로, 배포된 시스템에 쉽게 적용할 수 있다.
• vLLM, TGI 통합: 이후 LLM 서빙 시스템에 표준 기능으로 통합되었다.

한계

• 초안 모델 의존성: 효과적인 초안 모델을 준비해야 하며, 초안 모델이 너무 약하면 수용률이 낮아 이득이 줄어든다.
• 메모리 오버헤드: 두 모델을 동시에 GPU 메모리에 올려야 하므로, 메모리 여유가 적은 환경에서는 적용이 어렵다.
• 추측 오류 비용: 거부율이 높으면 오히려 순수 대상 모델보다 느려질 수 있다.
• 배치 추론 한계: 추측 디코딩은 단일 시퀀스 생성에서 효과적이지만, 대용량 배치 추론 환경에서는 이득이 제한적이다.