디코드는 토큰 하나를 만들 때마다 가중치 전체를 읽는다. 배치가 작으면 연산기는 거의 놀고 스텝 시간은 읽는 시간이 정한다. 그런데 같은 가중치를 한 번 읽어 토큰 여러 개를 검증할 수 있다면, 놀던 연산기로 스텝 수를 줄일 수 있다. 투기적 디코딩(speculative decoding)은 싼 방법으로 다음 토큰 몇 개를 추측하고, 대상 모델이 그 추측을 한 번의 실행으로 확인하는 방법이다.
동작 방식
초안 모델과 대상 모델이 번갈아 실행된다. 한 번의 반복은 세 단계로 이뤄진다.
- 초안: 작고 빠른 초안 모델이 다음 토큰 k개를 차례로 생성한다
- 검증: 대상 모델이 현재 문맥 뒤에 초안 k개를 붙여 한 번 실행한다. 프리필처럼 k+1개 위치의 다음 토큰 분포를 한꺼번에 얻는다
- 수락: 앞에서부터 초안 토큰을 하나씩 수락 규칙으로 확인하고, 처음 거절된 위치에서 멈춘다. 거절된 위치에는 대상 모델의 분포에서 새 토큰을 하나 뽑는다
초안이 모두 수락되면 대상 모델 한 번의 실행으로 토큰 k+1개를 얻는다. 첫 토큰부터 거절돼도 대상 모델이 그 위치의 토큰 하나를 내므로, 최악의 경우에도 보통의 디코드 스텝과 같은 토큰 하나는 얻는다.
검증 단계가 싼 이유는 디코드가 대역폭에 묶여 있어서다. 토큰 1개를 계산하든 5개를 계산하든 가중치를 읽는 양은 같으므로, 배치가 작을 때 검증 실행은 보통 디코드 스텝과 시간이 거의 같다.
출력 분포를 지키는 수락 규칙
투기적 디코딩은 근사가 아니다. Leviathan 등(ICML 2023)과 Chen 등이 제안한 수락 규칙은 대상 모델만으로 샘플링한 것과 같은 분포의 출력을 보장한다.
초안 모델이 어떤 토큰 x에 준 확률을 q(x), 대상 모델이 준 확률을 p(x)라고 하자. 초안 토큰 x는 확률 min(1, p(x)/q(x))로 수락한다. 대상 모델이 초안보다 그 토큰을 더 높게 보면 항상 수락하고 낮게 보면 그 비율만큼만 수락한다. 거절되면 max(0, p(x) − q(x))를 정규화한 분포에서 새 토큰을 뽑는다. 이 규칙을 따르면 최종 토큰의 분포가 정확히 p가 된다. 탐욕적 디코딩(온도 0)이라면 규칙이 단순해져서 초안 토큰이 대상 모델의 최댓값 토큰과 같을 때만 수락한다.
논문들이 보고한 속도 향상은 T5-XXL에서 23배(Leviathan 등), Chinchilla 70B에서 22.5배(Chen 등)다. 출력이 원래 모델과 같은 분포이므로 양자화와 달리 품질 평가를 다시 할 필요가 없다.
기대 토큰 수
이득은 초안이 얼마나 자주 수락되느냐에 달려 있다. 토큰마다 수락 확률이 α로 같다고 단순화하면 초안 k개로 한 번의 반복에서 얻는 토큰 수의 기댓값은 다음과 같다(Leviathan 등).
기대 토큰 수 = (1 − α^(k+1)) ÷ (1 − α)| 수락률 α | k = 2 | k = 4 | k = 8 |
|---|---|---|---|
| 0.6 | 1.96 | 2.31 | 2.47 |
| 0.8 | 2.44 | 3.36 | 4.33 |
| 0.9 | 2.71 | 4.10 | 6.13 |
k를 늘리면 기대 토큰 수는 늘지만 늘어나는 폭이 점점 줄고 초안 모델을 k번 실행하는 비용은 k에 비례해 늘어난다. 실제 속도 향상은 대략 다음과 같다.
속도 향상 ≈ 기대 토큰 수 ÷ (검증 1회 시간 + 초안 k회 시간) × 보통 디코드 스텝 시간수락률은 작업에 따라 크게 달라진다. 코드처럼 반복이 많고 예측하기 쉬운 출력, 입력을 많이 인용하는 요약은 수락률이 높고, 창작처럼 다음 토큰이 열려 있는 출력은 낮다. 온도가 높을수록 수락률이 떨어진다. 그래서 k는 서비스의 실제 요청으로 수락률을 재고 정한다.
초안을 만드는 방법
| 방법 | 초안을 만드는 것 | 특징 |
|---|---|---|
| 별도 초안 모델 | 같은 토크나이저를 쓰는 작은 모델 (예: 같은 계열의 1B 모델) | 구현이 단순하다. 초안 모델의 가중치와 KV 캐시가 메모리를 더 쓴다 |
| Medusa (ICML 2024) | 대상 모델 마지막 층 위에 붙인 여러 개의 디코딩 헤드가 앞쪽 여러 위치의 토큰을 한꺼번에 예측 | 별도 모델이 없다. 헤드를 학습해야 한다. 논문은 2.2배 이상(Medusa-1), 2.3~3.6배(Medusa-2)를 보고한다 |
| EAGLE (ICML 2024) | 대상 모델의 은닉 상태를 입력으로 받아 다음 은닉 상태를 예측하는 작은 층 | 토큰이 아니라 특징 수준에서 예측해 수락률이 높다. LLaMA2-Chat 70B에서 2.7~3.5배를 보고한다 |
| n-gram 조회 | 프롬프트나 이미 생성한 텍스트에서 현재 끝부분과 같은 구간을 찾아 그 뒤를 초안으로 쓴다 | 모델이 필요 없다. 입력을 많이 베끼는 작업(요약, 코드 수정)에서만 효과가 있다 |
별도 초안 모델은 대상 모델과 분포가 비슷할수록 수락률이 높다. 같은 계열의 작은 모델을 쓰고 대상 모델을 미세 조정했다면 초안 모델도 같은 데이터로 맞추는 편이 낫다.
배치가 커지면 줄어드는 이득
투기적 디코딩의 이득은 디코드가 대역폭에 묶여 연산기가 놀고 있다는 데서 나온다. 배치가 커지면 그 전제가 약해진다.
배치 B에 초안 k개를 붙이면 검증 실행은 토큰 B × (k+1)개를 처리한다. B가 작을 때는 이 토큰들이 놀던 연산기를 채우지만, B × (k+1)이 능선(H100 SXM, BF16 기준 약 295) 근처를 넘으면 검증 실행 시간이 토큰 수에 비례해 늘어나기 시작한다. 그때부터는 거절된 초안 토큰에 쓴 연산이 그대로 손해가 되고 그 연산으로 다른 요청의 토큰을 만들었다면 처리량이 더 높았다.
그래서 투기적 디코딩은 배치가 작은 상황에서 지연을 줄이는 수단이다. 동시 요청이 적은 시간대, 응답 속도가 중요한 소수의 요청, 큰 모델을 적은 사용자에게 제공하는 서비스에서 효과가 크다. 배치가 늘 큰 처리량 위주의 서버에서는 꺼 두거나, 부하에 따라 k를 줄이거나 끄는 기능이 있는지 엔진 문서를 확인한다. 켜고 끈 두 설정의 굿풋을 같은 부하로 재 보는 것이 가장 확실하다.