Autoregressive 특성에 대하여 (LLM이 답변을 하나씩 하는 이유)

일반적인 LLM은 Autoregressive model이다. 지금까지 생성한 Token을 조건으로 다음 Token의 확률 분포를 계산하고, 그중 하나를 선택한 뒤 다시 같은 과정을 반복한다.

질문에 대한 답변이 다음과 같다고 해보자.

Speculative Decoding은 LLM의 출력을 빠르게 만든다.

모델은 이 문장을 통째로 생성하지 않는다.

1. Speculative을 생성한다.
2. Speculative Decoding을 보고 은을 생성한다.
3. 지금까지 생성된 문장을 보고 LLM을 생성한다.
4. 다시 전체 Context를 보고 다음 Token을 생성한다.

수식으로 표현하면 길이 $T$인 답변의 확률은 다음과 같이 분해된다.

$$p(x_{1:T}) = \prod_{t=1}^{T} p(x_t \mid x_{<t}).$$

문제는 $x_t$를 알아야 $x_{t+1}$을 계산할 수 있다는 점이다. 답변에 100개의 Token이 필요하다면 Target model을 적어도 100번 순차적으로 호출해야 한다.

물론 KV Cache를 사용하면 이전 Token에 대한 Key와 Value를 매번 다시 계산하지 않아도 된다. 그러나 다음 Token을 생성하기 위해 수십억 개의 Parameter를 가진 모델의 각 Layer를 계속 통과해야 한다는 사실은 변하지 않는다.

특히 Batch size가 작을 때의 LLM Decoding은 연산량보다 GPU Memory Bandwidth에 묶이는 경우가 많다. Matrix multiplication을 굉장히 많이 계산해서 느린 것이 아니라, 다음 Token 하나를 계산할 때마다 거대한 Weight를 GPU Memory에서 계속 읽어와야 해서 느린 것이다.

그렇다면 여기서 다음과 같은 의문이 생긴다.

어차피 Weight를 읽어와야 한다면, 다음 Token 하나만 계산하지 말고 여러 Token을 한 번에 검사할 수는 없을까?

이 발상에서 시작된 방법이 Speculative Decoding이다.

Speculative Decoding

Speculative라는 단어는 ‘추측에 근거한’ 정도의 의미를 가진다. 이름 그대로 다음에 나올 Token 여러 개를 미리 추측하고, Target model은 그 추측이 맞는지만 한꺼번에 검사한다.

여기에는 두 모델이 등장한다.

• Target model: 원래 답변을 생성하던 크고 무거운 LLM
• Draft model: 다음에 나올 법한 Token을 빠르게 제안하는 작은 모델

Draft model이 $q$, Target model이 $p$라고 하자. Prefix $x_{1:t}$가 주어졌을 때 전체 과정은 다음과 같다.

1. Draft

작은 Draft model이 다음 $k$개의 Token을 Autoregressive하게 먼저 생성한다.

$$\hat{x}_{t+1},\hat{x}_{t+2},\cdots,\hat{x}_{t+k} \sim q.$$

Draft model도 순차적으로 Token을 생성하지만 Target model보다 훨씬 작기 때문에 이 과정은 상대적으로 저렴하다.

2. Verification

Target model에 기존 Prefix와 Draft Token을 함께 입력한다.

$$[x_{1:t},\hat{x}_{t+1:t+k}]$$

Transformer는 학습할 때 Causal mask를 사용하여 여러 위치의 Next-token distribution을 병렬적으로 계산한다. 같은 원리로 Target model은 Draft가 만든 여러 위치의 확률 분포를 한 번의 Forward pass로 함께 계산할 수 있다.

3. Accept or Reject

앞쪽 Token부터 Draft의 제안을 검사한다.

• Target model의 분포와 잘 맞는 Draft Token은 Accept한다.
• 처음으로 Reject가 발생하면 이후 Draft Token은 버린다.
• Reject된 위치는 Target distribution에 맞게 보정하여 다시 Sampling한다.
• 모든 Draft가 Accept되면 Target model이 계산해둔 다음 Token까지 하나 더 얻을 수 있다.

쉽게 말하면 Draft model은 문장의 초안을 빠르게 작성하고, Target model은 빨간 펜을 들고 초안을 한 번에 검수한다.

Draft가 꽤 정확하다면 Target model을 한 번 호출하고 여러 Token을 확정할 수 있다.

$$\text{Vanilla AR} : 1\text{ Target Forward} \rightarrow 1\text{ Token}$$ $$\text{Speculative Decoding} : 1\text{ Drafting} + 1\text{ Target Verification} \rightarrow \text{여러 Token}$$

한 번에 여러 Token을 검사하면 왜 빨라질까?

Target model이 $k$개의 Draft Token을 검사한다고 해서 연산이 공짜가 되는 것은 아니다. 하나의 Token을 검사하는 것보다는 당연히 더 많은 연산이 필요하다.

그럼에도 불구하고 $k$번의 Target Forward를 순차적으로 수행하는 것보다는, 여러 위치를 한 차례의 더 큰 Matrix multiplication으로 묶는 편이 GPU에서 훨씬 효율적일 수 있다.

일반적인 Token-by-token Decoding에서는 Matrix의 한쪽 축이 매우 작다. GPU의 수많은 연산 장치를 충분히 사용하지 못한 채 매번 Model Weight를 읽는 비용을 치르게 된다.

반면 Verification에서는 여러 Token 위치를 같이 계산한다.

• 한 번 읽어온 Weight로 더 많은 연산을 수행한다.
• 작은 Matrix multiplication 여러 번을 더 큰 연산 한 번으로 묶는다.
• Target model의 순차 호출 횟수를 줄인다.
• GPU의 남아있는 연산 여유를 활용한다.

Speculative Decoding은 Target model의 전체 연산량을 무조건 없애는 방법이 아니다.

Memory-bound인 Decoding을 조금 더 Compute-bound한 Verification으로 바꾸고, 비싼 Target model의 순차 호출 횟수를 줄이는 것이다.

여기서 실제 속도는 다음 조건에 의해 결정된다.

$$\text{실제 이득} \approx \frac{ \text{한 Cycle에서 확정한 Token 수} }{ \text{Draft 비용} + \text{Verification 비용} }.$$

Draft가 너무 무거우면 후보를 만드는 데 시간이 오래 걸린다. Draft가 너무 부정확하면 대부분 Reject되어 Target model을 호출하고도 한두 Token밖에 얻지 못한다.

좋은 Speculative Decoding의 조건은 단순하다.

1. Draft model이 충분히 작고 빨라야 한다.
2. Target model이 선택할 Token을 Draft model도 높은 확률로 제안해야 한다.
3. 여러 Candidate를 GPU에서 효율적으로 검증할 수 있어야 한다.

EAGLE 계열은 이 중 두 번째 문제, 즉 작은 Draft가 어떻게 큰 Target의 다음 Token을 더 잘 맞힐 수 있을까에 집중한 방법이다.

작은 모델의 답을 사용해도 성능이 떨어지지 않을까?

Speculative Decoding을 처음 보면 가장 먼저 드는 의문이다.

작은 Draft model은 Target model보다 당연히 성능이 낮을 텐데, 작은 모델이 만든 Token을 사용하면 원래 LLM의 답변 품질도 같이 떨어지는 것이 아닐까?

결론부터 말하면 Draft model은 최종 결정권을 가지지 않는다.

Draft model은 Candidate를 제안할 뿐이고, 실제 Token의 채택과 보정은 Target model의 분포를 기준으로 진행한다.

Greedy decoding이라면 이해하기 쉽다.

• Draft가 제안한 Token과 Target model의 Argmax Token이 같으면 Accept
• 다르면 Target model의 Token을 사용하고 이후 Draft를 폐기

Draft가 아무리 이상한 Token을 제안해도 결과는 원래 Target model이 Greedy decoding으로 생성했을 Token과 같다. Draft의 품질은 정답이 아니라 속도에만 영향을 준다.

Sampling에서는 조금 더 정교한 Rejection sampling을 사용한다.

Draft가 제안한 Token $x$에 대해 Draft probability를 $q(x)$, Target probability를 $p(x)$라고 하면 Acceptance probability는 다음과 같다.

$$A(x) = \min \left( 1, \frac{p(x)}{q(x)} \right).$$

Draft가 어떤 Token을 제안하고 그 Token이 Accept될 전체 확률 질량은 다음과 같다.

$$q(x)A(x) = q(x) \min \left( 1, \frac{p(x)}{q(x)} \right) = \min(p(x),q(x)).$$

Draft distribution과 Target distribution이 겹치는 부분만 Draft의 제안으로 가져온다.

Reject가 발생하면 Target에는 있지만 Draft가 충분히 제안하지 못한 나머지 확률 질량에서 Token을 다시 Sampling한다.

$$p_{\mathrm{residual}}(x) = \operatorname{norm} \left( \max(0,p(x)-q(x)) \right).$$

최종적으로 Draft에서 Accept된 질량과 Reject 이후 보정된 질량을 합치면 정확히 Target distribution $p$가 된다.

$$\min(p(x),q(x)) + \max(0,p(x)-q(x)) = p(x).$$

엄밀하게는 Residual distribution의 Normalization과 전체 Reject probability까지 함께 계산해야 하지만 핵심 논리는 위와 같다.

Draft가 Target과 비슷하면 겹치는 확률 질량이 커져 Acceptance rate가 높아진다. Draft가 Target과 다르면 Reject가 자주 발생하지만, 보정 과정을 거친 최종 Sampling distribution은 여전히 Target과 같다.

그래서 Speculative Decoding에서 말하는 Lossless는 다음 의미를 가진다.

더 작은 모델의 답변으로 원래 모델을 대체하는 것이 아니라, Target model과 동일한 출력 분포를 더 적은 순차 호출로 Sampling한다.

Target model의 Weight를 바꾸지도 않는다. 따라서 EAGLE-3가 더 좋은 답을 만드는 것은 아니며, 원래 Target model이 만들 답을 더 빠르게 만드는 것이다.

Vanilla Speculative Decoding의 한계

가장 단순한 방식은 Target model과 같은 Family의 작은 LLM을 Draft model로 사용하는 것이다.

예를 들어 70B Target model을 빠르게 만들기 위해 1B 또는 3B 모델이 Draft를 작성하게 할 수 있다.

하지만 작은 모델은 Target model과 독립적으로 동작한다. 같은 Tokenizer와 비슷한 학습 데이터를 사용했더라도 두 모델이 계산한 다음 Token distribution은 다를 수 있다.

특히 다음 Token이 명확하지 않은 문장, 복잡한 추론, 전문 지식이 필요한 문제에서는 작은 Draft가 큰 Target의 선택을 맞히기 어렵다.

Draft model을 키우면 Acceptance rate는 올라갈 수 있지만 Drafting 비용도 증가한다. Draft model을 줄이면 빠르지만 Reject가 많아진다.

$$\text{Draft 크기 증가} \Rightarrow \begin{cases} \text{Acceptance 증가} \\ \text{Draft 비용 증가} \end{cases}$$

Vanilla Speculative Decoding은 결국 이 Trade-off에 묶인다.

EAGLE의 질문은 여기서 출발한다.

Target model이 이미 계산한 내부 Feature를 Draft model에게 보여주면, 아주 작은 Draft도 Target의 다음 선택을 더 잘 따라갈 수 있지 않을까?

EAGLE: Token이 아니라 Feature를 예측하자

EAGLE은 Target model과 완전히 분리된 작은 LLM을 하나 더 사용하는 대신, Target model이 직전 Forward에서 계산한 Hidden feature를 재사용한다.

일반적인 LLM의 마지막 부분을 단순화하면 다음과 같다.

$$h_t \xrightarrow{\text{LM Head}} z_t \xrightarrow{\text{Softmax}} p(x_{t+1} \mid x_{\le t}).$$

$h_t$는 Target model의 Top-layer feature이고, LM Head를 통과하면 Vocabulary 크기의 Logit $z_t$가 된다.

Vanilla Draft model은 Token sequence만 보고 다음 Token을 예측한다. 반면 EAGLE의 Draft model은 Target이 계산한 Feature를 받아 다음 시점의 Feature를 예측한다.

$$\hat{h}_{t+1} = D_\theta(h_{\le t}, e_{\le t+1}).$$

여기서 $D_\theta$는 작은 Draft model이고, $e$는 실제로 Sampling된 Token의 Embedding이다. EAGLE은 Feature autoregression에 존재하는 불확실성을 줄이기 위해 한 Step 앞선 Token sequence도 같이 입력한다.

예측한 Feature를 Target model의 LM Head에 통과시키면 Draft Token distribution을 얻을 수 있다.

$$q(x_{t+2}) = \operatorname{softmax} (W_{\mathrm{LM}}\hat{h}_{t+1}).$$

이 방식의 장점은 명확하다.

• Draft가 Target model의 풍부한 내부 정보를 재사용한다.
• Target model 전체보다 훨씬 작은 하나의 Decoder layer로 Draft를 만들 수 있다.
• Target과 독립적인 작은 LLM보다 Target distribution을 따라가기 쉽다.
• Target model의 LM Head와 Token Embedding도 재사용할 수 있다.

EAGLE은 Feature prediction loss와 Token prediction loss를 함께 사용한다.

$$\mathcal{L}_{\mathrm{EAGLE}} = \lambda_{\mathrm{fea}} \mathcal{L}_{\mathrm{fea}} + \mathcal{L}_{\mathrm{token}}.$$

Feature loss는 Draft output이 실제 Target feature와 가까워지도록 만들고, Token loss는 LM Head를 통과한 Token distribution이 Target의 Token을 잘 예측하도록 만든다.

덕분에 한 Step의 Feature prediction을 학습하고도 자기 예측을 다시 입력으로 사용하는 Multi-step Drafting을 어느 정도 수행할 수 있었다.

EAGLE-2: 모든 문장에 같은 Draft Tree가 필요할까?

Speculative Decoding은 하나의 Chain만 Draft할 수도 있지만 여러 후보를 가진 Tree를 만들 수도 있다.

다음 Token으로 나는과 오늘의 확률이 모두 높다면 하나만 고집하지 않고 두 Branch를 만들고, 각 Branch 뒤에 다시 Candidate를 붙이는 것이다.

Target model은 Tree attention을 사용하여 여러 Branch를 한 번에 검증한다. 각 Node는 공통 Prefix와 자기 Branch의 Ancestor만 볼 수 있도록 Attention mask를 구성한다.

EAGLE-1의 Draft tree는 미리 정한 고정된 형태였다. 하지만 모든 Context의 난이도가 같지는 않다.

• 다음 Token이 뻔한 Code나 상투적인 문장은 깊게 Draft해도 잘 맞는다.
• 여러 표현이 가능한 대화는 한 Branch를 깊게 파는 것보다 초반 Candidate를 넓게 잡는 것이 나을 수 있다.

EAGLE-2는 Draft model의 Confidence가 실제 Acceptance rate를 꽤 잘 근사한다는 점을 이용한다.

Confidence가 높은 Branch는 더 확장하고, 가능성이 낮은 Branch는 Pruning하여 Context마다 다른 Dynamic draft tree를 만든다.

EAGLE-2는 무엇을 예측할 것인가를 바꾼 모델이라기보다, 제한된 Draft budget을 어떤 Tree에 배분할 것인가를 개선한 방법이다.

EAGLE-3도 이 Dynamic draft tree를 그대로 사용한다.

그런데 Feature를 정확히 맞힐 필요가 있을까?

EAGLE-3의 출발점은 의외로 EAGLE의 장점처럼 보였던 Feature prediction을 다시 의심하는 것이다.

우리가 최종적으로 원하는 것은 Target feature와 똑같은 Vector가 아니다.

우리가 원하는 것은 Target model이 Accept할 Token Candidate다.

예를 들어 서로 꽤 멀리 떨어진 두 Hidden vector $a$와 $b$가 있다고 해보자.

$$\Vert a-b \Vert_2 \gg 0$$

그럼에도 두 Vector가 LM Head를 통과한 뒤 같은 Token에 높은 확률을 부여할 수 있다.

$$\arg\max W_{\mathrm{LM}}a = \arg\max W_{\mathrm{LM}}b.$$

Speculative Decoding의 목적에서는 이것으로 충분하다. Target이 선택할 Token만 잘 맞히면 된다.

그런데 Feature regression loss는 Draft output이 Target feature의 좌표까지 닮도록 강제한다.

$$\mathcal{L}_{\mathrm{fea}} = \Vert \hat{h}_{t+1}-h_{t+1} \Vert.$$

Token을 맞히는 데에는 여러 가능한 표현이 존재할 수 있는데, Feature loss는 그중 Target feature 하나로 Draft의 표현 공간을 제한한다.

논문에서는 이것을 Feature prediction constraint라고 부른다.

실제로 EAGLE의 Training data를 늘려도 성능 향상이 금방 둔화되었다. Draft model이 더 많은 Token pattern을 학습할 여지가 있어도 Feature regression이라는 추가 Constraint에 묶여 충분히 이득을 얻지 못한 것이다.

그렇다면 Feature loss를 없애고 Token만 직접 예측하면 되지 않을까?

실제로 Feature constraint를 제거하면 첫 번째 Draft Token의 Acceptance rate는 크게 올라간다.

하지만 바로 다음 문제가 생긴다.

Feature constraint를 없애자 Training과 Inference가 달라졌다

EAGLE의 Feature loss가 귀찮은 Constraint였던 것은 맞지만, 한 가지 중요한 역할도 하고 있었다.

Draft가 출력한 $\hat{h}_{t+1}$을 실제 Target feature $h_{t+1}$과 비슷하게 만들었기 때문에, 다음 Draft Step에 자기 출력을 다시 넣어도 학습 때 보던 Feature와 어느 정도 비슷한 입력이 들어왔다.

Feature loss를 제거하면 Draft output $a_{t+1}$은 더 이상 Target feature $h_{t+1}$을 닮을 필요가 없다.

첫 Step에서는 Target model이 계산한 정확한 Feature를 입력받는다.

$$g_1,g_2,\cdots,g_t \rightarrow a_{t+1}$$

하지만 두 번째 Step부터는 아직 Target이 검증하지 않은 위치의 정확한 Feature를 구할 수 없다. 그것을 구하려고 Target model을 호출하면 Speculative Decoding을 하는 의미가 사라진다.

그래서 이전 Step에서 Draft가 만든 $a_{t+1}$을 대신 입력해야 한다.

$$g_1,\cdots,g_t,a_{t+1} \rightarrow a_{t+2}.$$

문제는 일반적인 Teacher forcing 방식으로 한 Step만 학습하면 모델이 Training 중에는 주로 정확한 Target feature만 본다는 점이다.

실제 Inference에서는 자신의 불완전한 출력을 다시 입력으로 받는다.

$$\begin{aligned} \text{Training input} &: g_1,g_2,\cdots,g_t \\ \text{Inference input} &: g_1,\cdots,g_t,a_{t+1},a_{t+2},\cdots \end{aligned}$$

첫 번째 예측의 작은 오차가 두 번째 입력의 분포를 바꾸고, 두 번째 오차가 다시 세 번째 입력으로 들어간다. 뒤로 갈수록 Training distribution에서 멀어지면서 Acceptance rate가 급격히 떨어진다.

이는 Sequence model에서 흔히 이야기하는 Exposure bias와 비슷하다.

EAGLE-3는 이 문제를 Training-time Test로 해결한다.

Training-time Test

이름 때문에 Test 중에 모델의 Weight를 학습시키는 Test-time Training과 헷갈릴 수 있다.

EAGLE-3의 Training-time Test는 추론 도중 Gradient update를 수행하는 방법이 아니다.

Training 중에 실제 Test-time의 Multi-step Drafting 과정을 미리 펼쳐서 학습하는 방식이다.

학습 과정에서 Draft model을 한 번 호출하고 끝내지 않는다.

1. Target feature를 입력하여 첫 번째 Draft output $a_{t+1}$을 만든다.
2. $a_{t+1}$을 다시 Draft model의 입력에 넣어 $a_{t+2}$를 만든다.
3. 다시 자기 출력을 입력하여 $a_{t+3}$를 만든다.
4. 각 Step에서 Token prediction loss를 계산한다.

개념적인 Objective는 다음과 같다.

$$\mathcal{L}_{\mathrm{TTT}} = \sum_{s=1}^{S} \mathcal{L}_{\mathrm{token}}^{(s)}.$$

모델은 학습할 때부터 자기 출력이 섞인 Context를 보게 된다.

$$[g_{\le t}] \rightarrow [g_{\le t},a_{t+1}] \rightarrow [g_{\le t},a_{t+1},a_{t+2}]$$

Draft depth가 깊어져도 처음 보는 입력 분포에 갑자기 던져지지 않는다.

논문의 표현을 빌리면 Test 과정을 Training 안으로 가져온 것이다.

이 과정은 단일 Chain만 흉내 내는 것이 아니라 여러 위치에서 생성된 Prediction을 다음 Round의 입력으로 사용한다. 각 Prediction은 서로 다른 Parent를 가진 Tree 구조가 되므로 Attention mask도 일반적인 Lower-triangular mask와 달라진다.

• 원본 Training sequence는 일반적인 Causal attention을 사용한다.
• Draft가 만든 Node는 공통 Prefix를 볼 수 있다.
• 각 Node는 자신의 Parent와 Ancestor만 볼 수 있다.
• 다른 Branch의 미래 정보는 볼 수 없다.

이렇게 Test-time의 Draft tree와 같은 Context dependency를 Training 중에 재현한다.

EAGLE-3가 Feature loss를 제거하고도 Multi-step Drafting을 안정적으로 수행할 수 있는 이유가 바로 이 부분이다.

Top-layer 하나만 보지 않고 여러 Layer를 섞는다

Feature constraint를 제거하면 Draft input도 꼭 Target의 Top-layer feature와 같은 공간에 머물 필요가 없다.

EAGLE-3는 이 자유도를 이용해 Target model의 Low-, Middle-, High-level feature를 함께 사용한다.

각 Layer에서 얻은 Hidden state를 $l_t,m_t,h_t \in \mathbb{R}^{d}$라고 하자.

먼저 세 Feature를 Concatenate한다.

$$[l_t;m_t;h_t] \in \mathbb{R}^{3d}.$$

이후 Fully-connected layer로 다시 Hidden size $d$에 맞춰 Fusion한다.

$$g_t = W_g[l_t;m_t;h_t].$$

그리고 현재 Sampling된 Token의 Embedding $e_t$와 결합하여 작은 Draft decoder의 입력으로 사용한다.

$$u_t = W_u[g_{t-1};e_t].$$

정확한 Index 표현은 구현을 단순화한 것이지만, 중요한 구조는 다음과 같다.

$$\text{Target의 여러 Layer Feature} + \text{현재 Token Embedding} \rightarrow \text{1-layer Draft Decoder} \rightarrow \text{LM Head} \rightarrow \text{Draft Token}.$$

Top-layer feature는 LM Head 바로 앞에 있으므로 다음 Token Logit에 강하게 맞춰진 표현이다.

논문에서는 LM Head의 Weight가 Full-rank라면 Logit에 대응되는 Top-layer feature가 유일하다는 관점에서, 이 Feature가 사실상 바로 다음 Token을 위한 정보에 정렬되어 있다고 설명한다.

하지만 Draft model은 다음 Token 하나만이 아니라 그다음, 그다음 Token까지 미리 내다봐야 한다.

일반적으로 Transformer의 서로 다른 깊이에는 Local pattern부터 더 추상적인 Context까지 서로 다른 정보가 존재할 수 있다. EAGLE-3는 Low/Middle/High feature를 Fusion하여 Top-layer 하나에만 의존할 때보다 Drafting에 유용한 정보를 넓게 가져온다.

Fusion feature $g$는 다음 Step마다 Target model에서 새로 구하지 않는다.

• Target이 이미 검증한 Prefix에는 Target의 Fusion feature $g$를 사용한다.
• 아직 검증하지 않은 Draft 영역에는 이전 Draft output $a$를 대신 사용한다.
• Training-time Test가 $g$와 $a$가 섞인 입력을 미리 학습한다.

EAGLE-3의 전체 Inference 과정

지금까지의 내용을 하나로 합쳐보자.

Prefix가 주어졌을 때 EAGLE-3는 다음 과정을 반복한다.

1. Target Feature 재사용

Prefill 또는 직전 Verification에서 Target model이 계산한 Low/Middle/High hidden state를 저장하고 Fusion feature $g$를 만든다.

2. Draft Tree 생성

작은 1-layer Draft decoder가 Fusion feature, Token embedding 그리고 이전 Draft output을 이용해 여러 미래 Token을 생성한다.

EAGLE-2의 Dynamic draft tree를 이용하여 Confidence가 높은 Branch에 더 많은 Candidate budget을 배분한다.

3. Tree Attention으로 병렬 검증

생성된 Tree를 Target model에 한 번 입력한다.

Attention mask를 이용하여 각 Candidate가 공통 Prefix와 자기 Ancestor만 보게 만들고, 여러 Branch의 Target probability를 병렬적으로 계산한다.

4. Strict Acceptance

Greedy decoding에서는 Target의 선택과 일치하는 가장 긴 Prefix를 Accept한다.

Sampling에서는 $\min(1,p/q)$ Acceptance와 Residual correction을 적용하여 Target distribution을 보존한다.

5. KV Cache 갱신

Accept된 경로에 해당하는 KV Cache만 남기고 Reject된 Branch는 버린다. 이후 확정된 Prefix에서 다시 Drafting을 시작한다.

전체 Flow를 줄이면 다음과 같다.

$$\text{Target Feature} \rightarrow \text{Cheap Dynamic Draft Tree} \rightarrow \text{One Target Verification} \rightarrow \text{Accepted Token Prefix} \rightarrow \text{Repeat}.$$

EAGLE-3에서 각 요소가 해결하는 문제

EAGLE-3는 하나의 Trick으로 빨라진 것이 아니다. 각 요소가 서로 다른 Bottleneck을 해결한다.

여기서 Feature constraint 제거만 단독으로 적용하면 첫 Draft는 잘 맞지만 뒤쪽 Draft의 Acceptance가 무너질 수 있다.

반대로 Training-time Test만 적용하면서 여전히 Feature를 정확히 회귀하도록 강제하면 Draft의 표현력과 Input feature 선택이 제한된다.

EAGLE-3의 논리는 두 변화를 같이 적용하는 데 있다.

Feature를 닮게 만드는 Constraint를 버려 Token 예측의 자유도를 얻고, 그 결과 생긴 자기 출력의 Distribution shift는 Training-time Test로 직접 학습한다.

데이터가 많아졌을 때 비로소 좋아진 이유

EAGLE-3 논문의 또 다른 핵심은 Draft model에서도 Data scaling의 효과를 확인했다는 점이다.

기존 EAGLE은 Training data를 늘려도 Acceptance rate와 Speedup의 개선이 제한적이었다.

Data가 부족해서라기보다 Feature regression constraint가 모델이 배울 수 있는 함수의 형태를 제한했기 때문이다.

EAGLE-3에서는 다음 변화가 일어난다.

1. 최종 목적과 직접 관련된 Token prediction에 집중한다.
2. Multi-layer feature를 사용하여 Input information을 풍부하게 만든다.
3. Training-time Test로 실제 Multi-step input distribution을 학습한다.

Data가 늘어나면 더 다양한 Context에서 Target model의 Token distribution과 자기 예측 이후의 복구 방법을 배울 수 있다.

논문에서는 기존 ShareGPT 약 68K에 UltraChat-200K 약 464K를 추가하고, Target model이 생성한 Response를 이용해 Draft를 학습하였다. EAGLE보다 약 8배 많은 Data를 사용했을 때 EAGLE-3의 Scaling curve가 계속 증가하는 것을 확인했다.

여기서 말하는 Scaling law는 모든 환경에서 무조건 같은 비율로 빨라진다는 이론적 보장이라기보다, 논문의 실험 범위에서 Data scale 증가에 따라 Acceptance와 Speedup이 계속 개선된 경험적인 관찰로 이해하는 것이 적절하다.

실제로 얼마나 빨라졌을까?

논문은 Chat model과 Reasoning model을 대상으로 MT-Bench, HumanEval, GSM8K, Alpaca, CNN/DailyMail에서 실험하였다.

Temperature 0에서 일부 결과만 정리하면 다음과 같다.

Vicuna 13B의 HumanEval에서는 최대 6.47x Speedup과 7.54의 평균 확정 길이를 기록했다.

Code generation은 반복되는 Syntax와 고정된 Pattern이 많아 Draft가 Candidate를 맞히기 쉬운 편이다. 그래서 HumanEval에서 특히 높은 Speedup이 나타났다.

Ablation 결과도 EAGLE-3의 설계를 이해하는 데 중요하다.

Feature constraint 제거와 Multi-layer fusion이 각각 실제 Acceptance length와 Speedup에 기여한 것을 볼 수 있다.

Batch가 커져도 항상 빨라질까?

Speculative Decoding은 보통 작은 Batch에서 가장 큰 효과를 낸다.

Batch가 작을 때 Target Decoding은 GPU의 연산 장치를 충분히 사용하지 못하고 Memory bandwidth에 묶이기 쉽다. Draft Verification은 그 남는 연산 여유를 활용한다.

반대로 Batch가 커지면 Vanilla Decoding 자체가 이미 큰 Matrix multiplication을 수행한다. GPU utilization이 높아진 상태에서는 Drafting과 Tree 관리가 순수한 Overhead가 될 수 있다.

EAGLE-3 논문이 이 점에서 흥미로운 이유는 Production serving framework에서도 비교적 큰 Batch까지 Throughput 개선을 보였기 때문이다.

SGLang 0.4.4, H100, LLaMA 3.1 8B, MT-Bench 환경의 결과는 다음과 같다.

기존 EAGLE은 Batch 24부터 Baseline보다 느려졌지만, EAGLE-3는 Batch 64에서도 1.38x의 Throughput을 기록했다.

다만 이 숫자를 모든 GPU와 모든 Serving 환경에 그대로 적용할 수는 없다.

같은 논문의 vLLM 실험에서는 Batch 56에서 EAGLE-3가 1.01x까지 내려갔다. Draft length, Tree width, Kernel implementation, GPU, Quantization, Context length, Request scheduling에 따라 손익분기점이 달라진다.

EAGLE-3는 Lossless이지만 항상 Speedup을 보장하지는 않는다.

• Lossless: 올바른 Verification을 사용하면 Target의 출력 분포를 유지한다.
• Speedup: 실제 Hardware와 Workload에서 Draft 비용보다 절약한 Target 비용이 커야 한다.

두 개념은 구분해야 한다.

성능 저하가 없다는 말을 정확히 해석하기

EAGLE-3 논문은 Target model의 Generation quality를 별도로 비교하지 않는다.

그 이유는 EAGLE-3가 Target model의 Weight를 수정하지 않고, Strict speculative sampling 조건으로 Target distribution을 보존하기 때문이다.

그렇다고 “무엇을 해도 결과가 완전히 같은 문자열이다”라는 의미는 아니다.

Greedy Decoding

같은 Numerical environment와 Deterministic kernel을 가정하면 원래 Target의 Argmax 경로와 같은 Token sequence를 얻는다.

Stochastic Sampling

매번 같은 문자열을 보장하는 것이 아니라, Vanilla Target sampling과 같은 확률 분포를 보장한다.

Random seed와 Sampling 순서가 달라지면 개별 실행의 문자열은 달라질 수 있지만, 알고리즘이 Sampling하는 분포는 Target distribution과 같다.

Draft가 틀렸을 때

답변 품질이 떨어지는 대신 Reject가 늘고 속도가 느려진다.

최악의 경우 Draft가 계속 틀리면 Speculative Decoding의 이득이 사라지고 추가 Overhead만 생길 수 있다. 그래도 Strict verifier가 유지되는 한 작은 Draft의 잘못된 분포가 최종 답변 분포를 대신하지는 않는다.

EAGLE-3에서 Draft model의 성능은 Quality knob보다 Speed knob에 가깝다.

실전에서 고려할 한계

Target model마다 Draft를 학습해야 한다

EAGLE-3는 Target의 Hidden state, Token embedding, LM Head에 강하게 의존한다.

하나의 EAGLE-3 Draft를 전혀 다른 Target model에 그대로 붙이기는 어렵다. Target model별로 호환되는 Draft checkpoint가 필요하다.

Model Weight를 Fine-tuning하거나 Vocabulary, Chat template, Architecture가 달라지면 Acceptance rate가 영향을 받을 수 있다.

Training 비용이 사라진 것은 아니다

Vanilla Speculative Decoding은 이미 존재하는 작은 모델을 Draft로 쓸 수 있지만, EAGLE-3는 Target response와 Hidden feature를 수집하고 전용 Draft head를 학습해야 한다.

Inference에서 절약할 비용이 충분히 크고 Request가 반복적으로 들어오는 서비스일수록 이 초기 비용을 회수하기 쉽다.

구현이 단순한 Two-model 방식보다 복잡하다

여러 Layer의 Hidden state를 가져와 Fusion하고, Dynamic tree와 Tree attention mask를 구성하고, Accept된 Branch에 맞춰 KV Cache를 정리해야 한다.

Serving framework가 이 과정을 효율적인 Kernel과 Scheduler로 지원하지 않으면 논문에서 보고한 Speedup을 그대로 얻기 어렵다.

Domain에 따라 Acceptance가 달라진다

Code처럼 Pattern이 반복되는 Domain은 Draft가 미래를 맞히기 쉽다.

반면 표현의 자유도가 높은 Creative writing, Target이 불확실한 Open-ended generation 또는 Training data와 크게 다른 Domain에서는 Acceptance가 낮아질 수 있다.

Reasoning model도 무조건 Draft하기 어려운 것은 아니다. 논문에서는 DeepSeek-R1 Draft를 Mathematical data로 추가 학습했을 때 GSM8K에서 높은 Speedup을 얻었다. 결국 Target이 자주 생성하는 Output distribution을 Draft training data가 얼마나 잘 포함하는지가 중요하다.

Lossless와 Numerical identity는 다르다

논리적으로는 Target distribution을 보존하지만 실제 System에서는 Floating-point precision, Kernel, Quantization, Random number consumption 순서에 따라 Numerical difference가 생길 수 있다.

구현체가 성능을 위해 Strict acceptance가 아닌 근사 Threshold를 사용한다면 더 이상 같은 의미의 Lossless를 주장할 수 없다.

EAGLE-3에서 흥미로웠던 부분

처음 EAGLE-3를 접하면 “작은 모델이 다음 Token을 여러 개 예측하는 방법” 정도로 이해하기 쉽다.

논문의 핵심은 Draft model을 하나 더 붙인 구조보다 학습 방식에 있다.

첫 번째로, Speculative Decoding은 제안과 결정의 역할을 분리한다.

• Draft model은 빠르게 제안한다.
• Target model은 최종 분포를 결정한다.

이 구조 덕분에 Draft가 가벼워도 답변 품질을 Draft 수준으로 낮출 필요가 없다.

두 번째로, EAGLE은 Target의 내부 Feature를 재사용하여 작은 Draft가 큰 Target의 의도를 더 잘 따라가게 만들었다.

세 번째로, EAGLE-3는 Feature를 정확히 복제하는 일이 최종 목적이 아니라는 점을 짚었다.

• Feature regression을 제거해 Token prediction의 자유도를 높인다.
• 그로 인해 생긴 자기 출력의 Distribution shift는 Training-time Test로 학습한다.
• Top-layer 하나 대신 여러 Layer의 정보를 Fusion한다.
• EAGLE-2의 Dynamic tree로 가능성 높은 Candidate에 계산을 집중한다.

마지막으로 Strict verification을 사용하여 최종 출력은 원래 Target distribution에 남겨둔다.

결국 EAGLE-3가 답변을 효과적으로 내뱉는 논리는 다음 한 줄로 정리할 수 있다.

가벼운 Draft가 Target의 미래를 여러 개 미리 제안하고, 무거운 Target은 그 제안을 병렬로 검증하며, 틀린 제안은 Target distribution으로 보정한다.

빠른 모델이 초안을 쓰고 좋은 모델이 최종 결재를 한다. 초안을 잘 쓸수록 한 번에 통과되는 Token이 많아져 빨라지고, 초안이 틀려도 결재권자가 바뀌지 않으므로 답변의 분포는 그대로 유지된다.

개인적으로 EAGLE-3에서 가장 재미있는 부분은 Feature prediction을 더 잘하는 방향이 아니라, 애초에 Feature를 맞혀야 한다는 문제 정의를 버린 것이라고 생각한다.

기존 Constraint를 제거하면 당장은 Multi-step 예측이 무너진다. 하지만 실제 Test 과정을 Training에 넣어 그 문제를 정면으로 학습했고, 덕분에 더 다양한 Feature와 더 많은 Data를 사용할 수 있게 되었다.

EAGLE-3에서는 다음과 같은 모델링 원칙을 읽을 수 있다.

최종 목표가 Token acceptance라면 중간 표현의 모양까지 강제하지 말고, 실제 배포 환경에서 모델이 만나게 될 자기 출력을 학습 과정에 포함시켜라.

이것이 EAGLE-3가 Draft를 더 정확하게 만들고, 더 긴 Candidate를 통과시키며, 원래 모델의 성능 저하 없이 Autoregressive Decoding의 순차적인 Bottleneck을 줄이는 논리다.