[논문 리뷰] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

Yuhui Li, Fangyun Wei외 2명|arXiv (Cornell University)|2024. 01. 26.

Natural Language Processing Techniques인용 수 5

한 줄 요약

EAGLE은 예측 샘플링의 불확실성 문제를 해결하기 위해 한 단계 앞선 토큰 정보를 갖춘 특징 레벨 자동회귀 초안을 도입하고, 파인튜닝 없이 대상 LLM의 출력 분포를 보존하면서 MT-bench에서 2.1배~3.8배의 지연 시간 속도 향상을 달성한다.

ABSTRACT

Autoregressive decoding makes the inference of Large Language Models (LLMs) time-consuming. In this paper, we reconsider speculative sampling and derive two key observations. Firstly, autoregression at the feature (second-to-top-layer) level is more straightforward than at the token level. Secondly, the inherent uncertainty in feature (second-to-top-layer) level autoregression constrains its performance. Based on these insights, we introduce EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency), a simple yet highly efficient speculative sampling framework. By incorporating a token sequence advanced by one time step, EAGLE effectively resolves the uncertainty, enabling precise second-to-top-layer feature prediction with minimal overhead. We conducted comprehensive evaluations of EAGLE, including all models from the Vicuna and LLaMA2-Chat series, the MoE model Mixtral 8x7B Instruct, and tasks in dialogue, code generation, mathematical reasoning, and instruction following. For LLaMA2-Chat 70B, EAGLE achieved a latency speedup ratio of 2.7x-3.5x, doubled throughput, while maintaining the distribution of the generated text.

연구 동기 및 목표

• LLM(대형 언어 모델)의 출력을 바꾸지 않으면서 더 빠른 자동회귀 디코딩을 촉진한다.
• 특징 수준의 자동회귀 및 불확실성 처리에 초점을 맞춰 추정적 샘플링을 재검토한다.
• 특징 수준의 불확실성을 해결하기 위해 한 단계 토큰 시프트를 사용하는 경량 초안 프레임워크(EAGLE)를 제안한다.

제안 방법

• EAGLE를 제안한다. 토큰이 아니라 두 번째 위 예층의 특징에서 작동하는 초안 단계를 가진 추정적 샘플링 프레임워크.
• 특징의 샘플링 불확실성을 해결하기 위해 초안 모델의 입력으로 한 시간 스텝 앞당겨진 토큰 시퀀스를 도입한다.
• 임베딩, LM 헤드, 자동회귀 헤드의 세 모듈 초안 모델과 트리 구조 초안이 있는 트리 어텐션을 사용한다.
• 특징 오차 누적 완화를 위해 고정 데이터 세트와 경량 데이터 증강을 사용하여 회귀 및 분류 손실의 결합으로 자동회귀 헤드를 학습한다.
• 출력 분포를 유지하기 위한 트리 어텐션을 갖춘 대상 LLM을 활용하는 검증 단계를 구현한다. 재귀적 추정 샘플링 절차를 따른다.
• 저비용 훈련 체계(70B 모델의 A100에서 1~2일) 제공 및 gpt-fast와 같은 다른 가속 방법과의 호환성을 보인다.

실험 결과

연구 질문

• RQ1 LLM의 추정 샘플링에서 특징 수준의 자동회귀가 토큰 수준 자동회귀보다 더 잘 작동할 수 있는가?
• RQ2 한 단계 앞선 토큰 입력으로 특징 수준의 불확실성을 다루는 것이 정확도와 속도향상에 어떤 영향을 미치는가?
• RQ3 다양한 모델과 작업에서 EAGLE의 지연 이점과 분포 보존 보장은 무엇인가?
• RQ4 EAGLE가 양자화, 컴파일 등 다른 가속 기법 및 배치 크기 규칙과 어떻게 상호작용하는가?
• RQ5 초안 작성 중 훈련 데이터 선택 및 특징 노이즈에 대해 EAGLE가 강건한가?

주요 결과

• EAGLE은 Vicuna 및 LLaMA2-Chat 계열에서 상당한 속도 향상을 달성한다. 예를 들어 LLaMA2-Chat 70B에서 지연 시간 2.7배~3.5배 향상과 처리량 두 배 증가를 보이며 출력 분포를 보존한다.
• 한 단계 앞선 토큰 입력을 가진 특징 수준의 초안 작성은 토큰 전용 초안 작성 및 이전 방법들(Lookahead, Medusa)보다 정확도와 속도를 향상시킨다.
• 초안/검증 단계의 트리 어텐션은 체인 초안 대비 소폭의 추가 이득을 제공하지만 최대 순방향 패스 수를 바꾸지 않고 패스당 토큰 수를 늘린다.
• MT-bench에서 Lookahead(1.7x–2.1x faster) 및 Medusa(1.5x–1.6x faster)보다 더 빠르고, 특정 설정에서 2x–3x 처리량 향상을 보이며 탐욕적/비탐욕적 모드에서도 분포를 보존한다.
• EAGLE은 gpt-fast와 결합하여 LLaMA2-Chat 7B 디코딩을 RTX 3090에서 160.4 토큰/초로 끌어올린다; 훈련 비용은 낮고(자동회귀 헤드의 파라미터 수가 1B 미만; 70B 모델은 A100에서 1–2일),
• 고찰 연구는 특징 및 시프트된 토큰 입력을 통한 무작위성 addressed의 중요성과 고정 데이터 세트 훈련이 대상 LLM이 생성한 데이터와 거의 동등한 성과를 보임을 보여주며, 데이터 민감성이 낮음을 강조한다.