EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

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.

한 줄 요약

EAGLE는 LLaMA2-Chat 70B 모델에서 2.7x-3.5x의 지연 감소를 달성한 효율적인 추측 샘플링 프레임워크이다.

핵심 기여도

**특징 수준(두 번째 상위 레이어)에서의 자동 회귀가 토큰 수준보다 더 직관적이라는 통찰**을 제시.
**특징 수준 자동 회귀의 불확실성**이 성능을 제한한다는 문제를 식별하고, 이를 해결하기 위한 **EAGLE 알고리즘** 제안.
**LLaMA2-Chat 70B 모델에서 2.7x-3.5x의 지연 감소**와 **2배의 처리량 증가**를 달성.
**EAGLE는 생성된 텍스트 분포를 유지하면서도 속도 향상**을 보장.

핵심 아이디어

EAGLE는 기존 추측 샘플링 방식에서 토큰 수준 자동 회귀를 사용하는 대신, **두 번째 상위 레이어**(second-to-top-layer)의 **특징 수준**(feature level)에서 자동 회귀를 수행함으로써 더 정확하고 효율적인 샘플링이 가능하다는 점에 주목한다. 이는 토큰 수열보다 더 규칙적인 패턴을 보이는 특징 수열을 이용하기 때문에, **LM Head를 통해 토큰을 유도하는 방식**이 토큰을 직접 예측하는 방식보다 더 효과적이라는 점에서 차별화된다.

또한, **특징 수준 자동 회귀의 불확실성**은 샘플링 과정에서 발생하는 **다양한 토큰 선택**(예: "am" vs. "always")에 따라 서로 다른 특징 수열을 생성하게 되어, 예측 정확도를 낮춘다. 이를 해결하기 위해 EAGLE는 **하나의 시간 단계 앞선 토큰 수열**을 드래프트 모델에 입력하여 불확실성을 해소한다. 예를 들어, 입력 `f_I`와 `t_always`를 기반으로 `f_always`를 예측하거나, `t_am`을 기반으로 `f_am`을 예측함으로써 **정확도 0.8 수준**의 드래프트를 생성한다.

기술적 접근법

**EAGLE는 두 번째 상위 레이어 특징**(second-to-top-layer feature)을 기반으로 자동 회귀를 수행.
드래프트 모델은 **하나의 시간 단계 앞선 토큰 수열**을 입력으로 받아, 불확실성을 해소.
**LLaMA2-Chat 70B 모델**을 기반으로 실험 수행.
**EAGLE는 ShareGPT 데이터셋**을 사용하여 훈련, 학습률은 3e-5, AdamW 최적화기 사용.
**70B 모델의 경우 1-2일 내에 A100 40G 서버에서 학습 가능**.
**EAGLE의 학습 가능한 파라미터 수는 0.99B**.

주요 결과

**LLaMA2-Chat 70B 모델에서 EAGLE는 2.7x-3.5x의 지연 감소**를 달성.
**처리량은 2배 증가**.
**Medusa 대비 1.5x-1.6x, Lookahead 대비 1.7x-2.1x의 속도 향상**.
**MT-bench, HumanEval, GSM8K, Alpaca 데이터셋**에서 평가.
**EAGLE는 생성된 텍스트 분포를 유지**하면서도 속도 향상.

의의 및 한계

EAGLE는 기존 추측 샘플링 방식에서 **드래프트 모델의 정확도와 불확실성 문제를 해결**함으로써, **LLM 추론 속도를 극적으로 향상**시키는 기술적 기여를 한다. 특히, **두 번째 상위 레이어 특징을 활용한 자동 회귀**는 기존 토큰 수준 자동 회귀보다 더 정확하고 구조화된 예측을 가능하게 하며, **하나의 시간 단계 앞선 토큰을 입력으로 사용하는 방식**은 불확실성을 효과적으로 해소한다.

하지만, EAGLE는 **드래프트 모델의 정확도**(0.8 수준)가 완전히 완벽하지 않기 때문에, **일부 토큰이 거부될 수 있으며**, 이는 **드래프트 모델의 정확도 향상**이 필요하다는 한계를 내포한다. 또한, **트리 어텐션**(tree attention) 방식과는 호환되지 않아, **체인 드래프트**(chain draft)만 사용 가능하다는 점도 한계로 작용할 수 있다.

실용적 활용

EAGLE는 대규모 언어 모델의 **대화, 코드 생성, 수학적 추론, 명령어 수행** 등 다양한 작업에서 활용 가능하다. 특히, **LLaMA2-Chat 70B와 같은 대형 모델의 추론 속도를 3배 이상 향상**시키므로, **실시간 대화 시스템, 코드 생성 도구, 고객 지원 챗봇** 등에서 실용적 가치가 크다. 또한, **정량화**(quantization), **컴파일**(compilation)과 같은 다른 가속 기술과 병행 사용이 가능하여, **LLM 시스템의 운영 비용을 추가로 절감**할 수 있다.