[논문 리뷰] Online Speculative Decoding

Xiaoxuan Liu, Lanxiang Hu외 5명|arXiv (Cornell University)|2023. 10. 11.

Topic ModelingComputer Science인용 수 3

한 줄 요약

이 논문은 LLM 서빙 클러스터의 유휴 계산 자원을 활용해 실시간 사용자 쿼리 데이터를 기반으로 드래フト 모델을 지속적으로 재학습하는 온라인 사전 추론 기법을 제안한다. 온라인 지식 이행을 적용함으로써 드래프트 모델은 실시간으로 쿼리 분포에 적응하게 되어 토큰 수용률을 10–65% 향상시키고 추론 지연을 1.22×에서 3.06×까지 감소시킨다.

ABSTRACT

Speculative decoding is a pivotal technique to accelerate the inference of large language models (LLMs) by employing a smaller draft model to predict the target model's outputs. However, its efficacy can be limited due to the low predictive accuracy of the draft model, particularly when faced with diverse text inputs and a significant capability gap between the draft and target models. We introduce online speculative decoding to address this challenge. The main idea is to continuously update the (multiple) draft model(s) on observed user query data. Adapting to query distribution mitigates the shifts between the training distribution of the draft model and the query distribution, enabling the draft model to more accurately predict the target model's outputs. We develop a prototype of online speculative decoding based on knowledge distillation and evaluate it using both synthetic and real query data. The results show a substantial increase in the token acceptance rate by 0.1 to 0.65, bringing 1.42x to 2.17x latency reduction. Our code is available at https://github.com/LiuXiaoxuanPKU/OSD.

연구 동기 및 목표

• 분포 이탈과 능력 격차로 인해 드래프트 모델의 사전 추론 정확도가 낮아지는 문제를 해결한다.
• LLM 서빙 클러스터 내에서 활용도가 낮은 계산 자원(여유 FLOPs)을 활용해 드래프트 모델의 지속적 재학습을 수행한다.
• 추론 비용을 증가시키지 않고도 드래프트 모델이 변화하는 쿼리 분포에 실시간으로 적응할 수 있도록 한다.
• 사용자 쿼리 데이터 기반의 온라인 지식 이행을 통해 온라인 LLM 서빙에서 토큰 수용률을 향상시키고 지연을 감소시킨다.
• 제한된 데이터로도 쿼리에 특화된 데이터 기반의 온라인 재학습이 정적 오프라인 피팅보다 우수한 성능을 내는지 입증한다.

제안 방법

• LLM 서빙 클러스터의 유휴 계산 능력(여유 FLOPs)을 활용해 드래프트 모델의 온라인 재학습을 수행한다.
• 일반화된 지식 이행(GKD)을 적용하여 드래프트 모델과 타겟 모델 출력 간의 차이를 기반으로 드래프트 모델을 업데이트한다.
• 틀린 추론 결과를 버퍼에 유지하고 여유 FLOPs가 확보될 때만 기울기 업데이트를 수행함으로써 오버헤드를 최소화한다.
• 특정 언어나 주제 클러스터에 맞게 최적화된 여러 전용 드래프트 모델을 훈련시켜 분포 기반 정확도를 향상시킨다.
• 온라인 이행 중에 드래프트 모델과 타겟 모델의 분포 간 거리 측도를 계산하여 손실를 산정한다.
• 기회적 훈련을 통해 추론 지연에 거의 영향을 주지 않으면서도 지속적인 모델 적응을 가능하게 한다.

Figure 1: Online speculative decoding overview. For each prompt, the draft model suggests multiple tokens in a single step. The target model then verifies these tokens, accepting some and rejecting others. If the student proposes incorrect tokens, both the draft and target distributions are stored i

실험 결과

연구 질문

• RQ1 실제 사용자 쿼리 데이터를 기반으로 한 드래프트 모델의 온라인 재학습이 사전 추론 정확도를 향상시킬 수 있는가?
• RQ2 소량의 쿼리 데이터만 이용 가능한 상황에서 온라인 사전 추론은 오프라인 피팅보다 어떻게 성능을 내는가?
• RQ3 LLM 서빙 클러스터 내 여유 FLOPs를 드래프트 모델 훈련에 재활용할 수 있는 정도는 어느 정도이며, 추론 성능에 영향을 주지 않는가?
• RQ4 쿼리 분포의 변화에 실시간으로 적응함으로써 토큰 수용률 향상이 지속적으로 이루어지는가?
• RQ5 각각 특정 언어나 주제에 맞게 훈련된 다수의 전용 드래프트 모델이 단일 일반 목적 드래프트 모델보다 우수한 성능을 내는가?

주요 결과

• 다양한 데이터셋에서 온라인 사전 추론은 토큰 수용률을 0.1에서 0.65(10–65%)까지 향상시켜 사전 추론 효율성을 크게 향상시켰다.
• 이 방법은 추론 지연을 1.22×에서 3.06×까지 감소시켰으며, 실제 LMSYS-chat 데이터에서 가장 높은 성능 향상을 보였다.
• 실제 LLM 채팅 대화에서 OSD는 언어당 2,000건 미만의 데이터 포인트로도 수용률을 0.1에서 0.2까지 향상시켰다.
• 주제 기반 대화에서는 모든 주제에서 수용률이 0.6를 초과했으며, 사회학 및 컴퓨터 과학 대화는 거의 0.8에 도달했다.
• SELECT, FROM, WHERE, +와 같은 고빈도 토큰은 온라인 이행 후 정밀도와 재현율 모두 뚜렷한 향상을 보였다.
• 전체 쿼리 데이터의 70–100%를 사용해 피팅한 모델과 비교해도 OSD는 동등하거나 그 이상의 성능를 보였다. 다만 사용한 데이터의 일부분에 불과했다.

Figure 3: Online acceptance rate ( $\alpha$ ) across different datasets. The x-axis represents the number of records that OSD has processed. Alpha is averaged over the most recent 50 records.