[논문리뷰]Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding(ACL, 2024)

Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Gang Chen, and Sharad Mehrotra. 2024. Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Lun-Wei Ku, Andre Martins, and Vivek Srikumar (Eds.). Association for Computational Linguistics, Bangkok, Thailand, 11263–11282. https://doi.org/10.18653/v1/2024.acl-long.607

Problem Statement

Draft & Verify는 별도의 draft 모델 없이 하나의 LLM으로 초안(draft) 생성과 검증(verify)을 모두 수행하는 self-speculative decoding을 제안한다. 기존 speculative decoding은 몇 가지 문제점이 존재한다.

[Auxiliary Draft Model] 기존의 연구는 드래프팅을 위한 모듈을 별도의 모듈을 사용한다. 그럼에도 불구하고 LLaMA-2-Chat과 같은 이미 fine-tuniung이 된 모델과의 일치율을 높이기 어렵다.

[Memory Overhead] 기존과 같이 별도의 Draft 모듈과 Verify 모듈을 둘 경우 두 개 모델을 항상 GPU에 올려두어야 하므로 메모리 사용량이 증가한다.

[Static Hyperparameter] 고정된 draft길이나 고정 임계치는 인스턴스 난이도 변동을 제대로 반영하지 못해 acceptance rate의 최적점을 결정하는데 어렵다.

Draft & Verify는 출력 품질을 유지하면서 토큰 당 추론 시간을 단축하는 것을 목표로 한다. 핵심은 드래프팅 단계에서 일부 중간 레이어를 선택적으로 건너뛰고, 이후 검증 단계에서 원래 LLM으로 한 번의 시퀀스 전방 패스로 초안 토큰을 일괄 검증하는 것이다 (최종 출력은 원래 LLM과 동일하게 유지됨).

Methodology

Overview

Draft & Verify의 핵심은 하나의 LLM을 가지고 드래프팅과 검증을 모두 하는 것이다. 드래프팅 단계에서는 중간 레이어의 일부를 건너뛰어 빠르게 K개 후보를 만든 뒤( \(p(\cdot \mid z^*,\cdot)\) ), 원본 LLM \(p(\cdot \mid\vec 0, \cdot)\)으로 한 번의 forward pass로 검증한다.

Self-Speculative Decoding

• 입력: 프롬프트 \(x_1, \cdots, x_t\), 최대 Draft 길이 \(K\), 건너뛸 레이어를 나타내는 벡터 \(z^* \in \{0, 1\}^L\)
• 출력: 최종 생성 \(x_1, \cdots, x_T\) (원본 LLM과 완전히 동일)

Draft → Verify로 구성되며, 보조 모델 없이 동일한 LLM 이 동작한다. 먼저 드래프팅 단계에서는 건너뛴 레이어 구성이 적용된 \(p(x \mid z^*, x_1, \cdots, x_j)\)로 \(x_{j+1} = \arg \max p(\cdot)\)를 최대 \(K\)개 생성한다.

다음으로 검증 단계에서는 LLM \(p(x \mid \vec 0, x_1, \cdots, x_i)\)로 순서대로 비교하며 불일치가 난 첫 지점 \(x_{i+1}\)을 원본 예측으로 대체하고 해당 지점부터 다시 드래프팅 단계를 시작한다 (수용 시 다음 토큰으로 진행). 이 절차는 한 번의 forward pass로 여러 draft 토큰을 검증하므로 토큰별 디코딩보다 효율적이며, 검증 기반이 원본 분포이므로 lossless(출력 동일)를 보장한다.

Selection of Skipped Layers (Bayesian Optimization)

문제 정의: 최종 목표는 검증된 토큰당 평균 추론 시간을 최소화하는 레이어 스킵 조합 \(z\)를 찾는 것이다.

• 목적함수: \(f(z)\) = 입력에 대해서 검증된 토큰당 평균 지연(Draft + Verify)
• 최적화

$$z^* = \arg \min_z f(z), \quad \text{s.t.} z \in \{ 0, 1\}^L$$

• 입력: 탐색 후보 \(z \in \{ 0, 1\}^L\)
• 출력: \(z^*\) (Draft구성; 어떤 attention/MLP 서브 레이어를 스킵할지)

Bayesian Optimization으로 가우시안 프로세스 기반 대리모델과 획득함수를 사용해 \(z\) 탐색을 반복하고, 모델 단위로 고정된 \(z^*\)를 사용한다.

Adaptive Draft-Exiting Mechanism

• 초기 단계의 종료 규칙: \(p(x_{t+1} \mid x_{1:t}) < \gamma\)이면 드래프팅 생성을 즉시 중단한다

핵심은 드래프트 길이를 자동으로 조절해 목표 수용률(acceptance rate) \(\alpha\)에 맞추는 feedback 제어 규칙이다. \(\gamma\)는 “드래프트 생성 중 조기 종료 임계값”으로, 초기 단계에는 확률 \(p(x_{t+1} \mid x_{1:t}) < \gamma\)을 만족하면 즉지 드래프팅 과정을 중단한다. 임계값이 커질수록 더 일찍 멈추므로 드래프트가 짧아져 수용률 (AR)이 올라가기 쉽고, 반대로 작아지면 더 길게 드래프트를 뽑아 수용률이 내려가기 쉬움을 이용해 세 단계로 \(AR\)을 \(\alpha\)로 수렴시키는 매커니즘이다. \(e\)번째 드래프팅 단계의 수용률을 \(AR_e\)라고 할 때,

$$AR \leftarrow \beta_1 AR + (1-\beta_1)AR_e$$ $$\tilde{\gamma} = \begin{cases} \gamma + \varepsilon, & \text{if } AR \le \alpha \\ \gamma - \varepsilon, & \text{otherwise} \end{cases}$$ $$\gamma \leftarrow \beta_{2} \gamma + (1-\beta_{2}) \tilde{\gamma}$$

Sampling Variant for Self-Speculative Decoding

하나의 LLM을 두 역할로 쓰되, 먼저 “가벼운 드래프팅 모드”로 몇 개 토큰을 미리 뽑고, 이어서 “정식 검증 모드”로 그 초안을 빠르게 확인·수용하는 방식이다. 보조 모델이 전혀 필요 없고, 최종 출력은 원래 LLM의 그리디 출력과 완전히 동일하다. 원래 그리디 디코딩을 쓰면서도 지연을 줄이고 싶을 때, 그리고 보조 드래프트 모델을 두기 어렵거나 메모리를 아끼고 싶을 때 효과적이다. 하지만, 가장 높은 점수를 가지는 샘플만 선택하면 다양성 부족을 야기할 수 있기 때문에 더 발전된 알고리즘을 제안한다.

전체 구조는 동일하지만 가장 높은 점수만 고르는 대신, 여러 샘플링 규칙을 적용해 드래프팅과 검증을 진행할 수 있다. 정답이 하나로 고정되지 않고 다양성이 가치 있는 과업에서, 즉 원래도 샘플링을 쓰는 환경에서 지연을 줄이되 생성 품질의 성격은 유지하고 싶을 때 적합하다.

알고리즘 4에서는 일반적인 LLM의 샘플링 (temperature top-p) 규칙을 그대로 사용한 변형 방법을 보여준다. 구체적으로, 초안 단계에서는 스킵된 레이어 구성으로 빠르게 샘플을 뽑고, 검증 단계에서 원본 LLM의 샘플링 정책(동일한 temperature/top-p)으로 해당 초안을 확인·수용·교체하는 절차를 따른다.

NR 카테고리 내 다른 글 보러가기