Speculative Decoding and MTP

LLM 추론의 가장 큰 병목은 시퀀셜성(sequential nature)입니다. 각 단계에서 한 토큰만 생성할 수 있고, 다음 토큰을 만들려면 앞의 모든 토큰이 완료되어야 합니다. 배치 처리로 여러 요청을 묶어도 각 요청 내에서는 이 순서를 피할 수 없습니다.

Speculative Decoding(투기적 디코딩)과 Multi-Token Prediction(MTP)은 이 병목을 우회하려는 두 가지 접근입니다.

Speculative Decoding: 작은 모델이 앞서 가다

투기적 디코딩의 기본 아이디어는 간단합니다. 큰 모델이 느리니까, 작은 모델(또는 시간 예측기)이 먼저 여러 토큰을 "추측"하고, 큰 모델이 그 추측을 검증하는 것입니다.

구체적으로:

1. 작은 모델(draft model) 이 다음 K개 토큰을 빠르게 샘플링합니다. 예: Phi-2 또는 TinyLlama 같은 작은 모델.
2. 큰 모델(target model) 이 K개 토큰을 한 번에 검증합니다. 즉, 원래 큰 모델이 각각 생성했을 확률과 작은 모델이 생성한 확률을 비교합니다.
3. 만약 작은 모델의 추측이 큰 모델의 확률과 유사하면, 그 토큰들을 수락합니다. 다르면 거기서 멈추고 큰 모델이 다시 생성합니다.

핵심은 병렬 검증(parallel verification) 입니다. 작은 모델이 K개를 하나씩 생성하는 시간이 짧으면, 큰 모델이 K개를 한 번에 검증하는 오버헤드(forward pass 1회)로 충분히 보상됩니다.

예시: Llama 2 7B로 Llama 2 70B 가속

OpenAI의 초기 논문(2023년 가을)에서는: - Draft model: Llama 2 7B (빠름) - Target model: Llama 2 70B (느림) - 결과: 2~3배 엔드-투-엔드 속도 향상 (전체 파이프라인, 불가피한 거부 포함)

거부율이 문제가 될 수 있습니다. 만약 작은 모델과 큰 모델의 예측이 자주 어긋나면, 많은 토큰이 거부되고 큰 모델이 다시 생성해야 합니다. 따라서 작은 모델의 선택이 중요합니다. 같은 계열의 작은 모델(예: Llama 계열에서 7B 선택)이 거부율을 낮춥니다.

MoGe, Lookahead Decoding 같은 변형들

최근에는 여러 변형이 나왔습니다. Lookahead Decoding(Google, 2024)은 별도의 작은 모델 없이, 큰 모델의 숨겨진 상태(hidden state)를 이용해서 앞당겨 계산합니다. 이렇게 하면 추가 모델 가중치가 필요 없으므로 메모리 오버헤드가 없습니다.

Multi-Token Prediction (MTP): 한 번에 여러 토큰 생성

MTP는 다른 접근입니다. 모델을 처음부터 K개 토큰을 동시에 예측하도록 학습시킵니다.

전통적으로는 다음 토큰 예측(next-token prediction)만 했습니다. 입력 "The cat sat on the" 다음에 "mat"을 한 토큰씩 예측합니다. MTP에서는 처음부터 "mat"의 3개 토큰을 한 번에 예측하도록 학습합니다.

학습 방식

가장 간단한 방법:

Input:  "The cat sat on the"
Targets: [mat, end, <pad>]  (3개 위치)

모델의 마지막 층이 3개의 출력 헤드를 가집니다. 각 헤드가 k번째 토큰의 확률 분포를 출력합니다. Cross-Entropy Loss는 3개 모두에 적용됩니다.

실제 추론

추론 단계에서는 여전히 자기회귀(autoregressive)입니다. 첫 번째 "mat" 토큰을 얻은 다음, 이를 입력에 추가해서 다음 1~3개 토큰을 예측합니다. 하지만 이론적으로는 1회 forward pass로 3개를 얻을 수 있으므로, 전체 생성 단계가 1/3로 줄어듭니다.

왜 지금 주목받나?

Speculative Decoding은 추가 모델이나 수정이 필요하지만, MTP는 기존 모델을 재학습하기만 하면 됩니다. 또한 FlexGen 같은 양자화 기법과 결합하면, 메모리도 덜 차지하면서 더 빠른 생성이 가능합니다.

최근 논문들에서: - Google의 "Decoding by Contrasting Layers"(2024): 추가 헤드만으로 2배 이상 속도 향상 - Together AI의 연구: 4개 토큰 동시 예측으로 3~4배 속도 향상 (배포 단계, 학습 비용 무시)

다만 MTP도 한계가 있습니다. 토큰들 사이의 상호작용을 무시합니다. "mat"을 예측할 때 " " (공백) 토큰과 "end"의 존재를 고려하지 않습니다. 따라서 첫 번째, 두 번째, 세 번째 토큰 예측의 정확도가 각각 떨어집니다.

Speculative Decoding vs MTP

실제 배포 사례

vLLM은 speculative decoding을 기본 옵션으로 제공합니다. 사용자가 draft model을 지정하면 자동으로 처리됩니다.

TensorRT-LLM(NVIDIA)도 speculative decoding을 지원하며, Llama 2 70B에서 약 2배 처리량 증가를 보고했습니다.

MTP는 아직 상용 배포가 제한적이지만, Llama 3 계열 모델의 일부 변형과 Mixtral 모델에서 실험 중입니다.

미래 방향

두 기법은 보완 관계입니다. Speculative Decoding은 '빠른 체크'에 강하고, MTP는 '미리 생성'에 강합니다. 이론적으로는 둘을 합칠 수 있습니다: 작은 모델이 4개를 추측하고, 큰 모델이 각각 2개씩 예측하는 식입니다.

또한 양쪽 모두 배치 처리와의 상호작용이 과제입니다. 여러 요청이 섞여 있을 때, speculative decoding의 거부율이 급상승할 수 있고, MTP도 요청마다 다른 시점에서 토큰이 소진됩니다. 이를 효율적으로 관리하는 것이 다음 세대 추론 엔진의 핵심 과제가 될 것 같습니다.