Speculative Decoding

작은 모델이나 보조 토큰 예측을 이용해 추론을 가속하는 원리는 무엇인가.

Scope

핵심 질문

Speculative Decoding은 작은 초안 모델이나 보조 예측을 이용해 큰 모델의 출력 분포를 유지하면서 왜 생성속도를 높일 수 있는가?

포함 범위

draft model
verification
acceptance
정확한 분포 유지
이득이 생기는 조건

제외 범위

특정 런타임의 설정법
최신 모델별 속도 벤치마크

Research

핵심 근거

autoregressive decoding은 기본적으로 한 번의 forward pass에서 한 토큰만 확정하므로 GPU 병렬성을 충분히 활용하기 어렵다.
speculative decoding은 저비용 draft가 여러 토큰을 제안하고 target model이 한 번에 검증한다.
rejection sampling 기반 방식은 target model의 원래 분포를 보존할 수 있다.
draft와 target의 예측이 자주 일치할수록 이득이 커진다.

참고 자료

Leviathan et al., Fast Inference from Transformers via Speculative Decoding: https://arxiv.org/abs/2211.17192
Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling: https://arxiv.org/abs/2302.01318

Draft

Speculative Decoding은 큰 언어 모델이 토큰을 하나씩 직접 생성하는 대신, 더 저렴한 방법으로 여러 토큰 후보를 미리 제안하고 큰 모델이 이를 한 번에 검증해 생성속도를 높이는 기법이다.

왜 기본 decoding이 느린가

Autoregressive LLM은 토큰 하나를 생성한 뒤 그 결과를 다시 입력에 붙여 다음 토큰을 생성한다.

따라서 100개의 토큰을 만들려면 기본적으로 100번의 순차적인 decoding step이 필요하다.

GPU는 대규모 병렬연산에 강하지만 한 토큰씩 순차적으로 확정하는 구조는 이 병렬성을 충분히 활용하기 어렵다.

Draft Model

Speculative Decoding에서는 작은 모델이나 별도의 predictor가 앞으로 나올 여러 토큰을 빠르게 제안한다.

예를 들어 작은 모델이 5개 토큰을 미리 만든 뒤 큰 target model이 이 5개를 한 번의 계산으로 검증할 수 있다.

큰 모델이 제안을 받아들이면 여러 토큰을 한꺼번에 확정한 효과를 얻는다.

틀리면 어떻게 되는가

Target model이 draft의 특정 토큰을 받아들이지 않으면 그 지점부터 target 분포에 맞는 토큰으로 교정한다.

정확한 rejection sampling 절차를 사용하면 최종적으로 생성되는 분포를 원래 target model의 autoregressive sampling과 동일하게 유지할 수 있다.

즉 단순히 작은 모델의 답을 믿고 품질을 희생하는 방식과는 다르다.

언제 빨라지는가

Draft model과 target model의 예측이 자주 일치할수록 좋다.

작은 모델이 제안한 토큰 대부분이 통과하면 target model 호출 한 번으로 여러 토큰을 확정할 수 있다.

반대로 두 모델의 예측이 자주 다르면 검증 후 버리는 토큰이 많아져 이득이 줄어든다.

Draft가 꼭 별도 모델이어야 하는가

아니다.

별도의 작은 모델을 사용할 수도 있고, 동일한 모델에 추가된 multi-token prediction head나 n-gram 기반 예측, 이전 문맥에서의 반복 패턴을 사용할 수도 있다.

핵심은 target model보다 싸게 미래 토큰 후보를 만들 수 있어야 한다는 것이다.

병목은 무엇인가

Speculative Decoding은 특히 batch가 작고 latency가 중요한 환경에서 효과가 클 수 있다.

서버가 이미 큰 batch로 가득 차 GPU 계산을 효율적으로 사용하고 있다면 추가 검증과 draft 실행의 이득이 상대적으로 작아질 수 있다.

핵심

Speculative Decoding은 작은 예측기가 미래 토큰을 먼저 추측하고 큰 모델이 여러 토큰을 병렬적으로 검증함으로써 autoregressive decoding의 순차성 병목을 줄이는 기법이다. 품질을 희생하지 않고 속도를 높일 수 있지만 draft의 정확도와 서빙 환경에 따라 실제 이득이 크게 달라진다.