대규모 언어 모델을 위한 빠른 추론적 디코딩을 위한 순환 드래프터

Recurrent Drafter for Fast Speculative Decoding in Large Language Models

March 14, 2024

저자: Aonan Zhang, Chong Wang, Yi Wang, Xuanyu Zhang, Yunfei Cheng

cs.AI

초록

본 논문에서는 대규모 언어 모델 서빙 효율성을 향상시키기 위한 개선된 스펙티브 디코딩 접근법을 소개합니다. 우리의 방법은 기존의 두 가지 기술, 즉 고전적인 이중 모델 스펙티브 디코딩 접근법과 최근의 단일 모델 접근법인 Medusa의 장점을 활용합니다. Medusa에서 영감을 받아, 우리의 접근법은 단일 모델 전략을 스펙티브 디코딩에 적용합니다. 그러나 우리의 방법은 경량화된 단일 드래프트 헤드를 사용하며, 순환 의존성 설계를 채택한다는 점에서 차별화됩니다. 이는 고전적인 스펙티브 디코딩에서 사용되는 소형 드래프트 모델과 유사하지만, 전체 트랜스포머 아키텍처의 복잡성을 피합니다. 또한 순환 의존성 덕분에, 우리는 드래프트 헤드를 사용하여 원치 않는 후보들을 빠르게 걸러내는 빔 서치를 활용할 수 있습니다. 결과적으로, 이 방법은 단일 모델 설계의 단순성을 유지하면서도 Medusa에서 추론을 위해 데이터 의존적 트리 어텐션 구조를 생성할 필요를 피합니다. 우리는 여러 인기 있는 오픈 소스 언어 모델에서 제안된 방법의 효과를 실증적으로 입증하고, 이 접근법을 채택함에 따른 트레이드오프에 대한 포괄적인 분석을 제공합니다.

English

In this paper, we introduce an improved approach of speculative decoding aimed at enhancing the efficiency of serving large language models. Our method capitalizes on the strengths of two established techniques: the classic two-model speculative decoding approach, and the more recent single-model approach, Medusa. Drawing inspiration from Medusa, our approach adopts a single-model strategy for speculative decoding. However, our method distinguishes itself by employing a single, lightweight draft head with a recurrent dependency design, akin in essence to the small, draft model uses in classic speculative decoding, but without the complexities of the full transformer architecture. And because of the recurrent dependency, we can use beam search to swiftly filter out undesired candidates with the draft head. The outcome is a method that combines the simplicity of single-model design and avoids the need to create a data-dependent tree attention structure only for inference in Medusa. We empirically demonstrate the effectiveness of the proposed method on several popular open source language models, along with a comprehensive analysis of the trade-offs involved in adopting this approach.