Review

TL; DR

• 저자

  

  • 메타 슈퍼 인텔리전스 랩 (소속 5명), 취리히 연방 공과대학교, 엔트로픽, GenAI, MAI

• cited: 28

Preliminary: Speculative Decoding

• paper : https://arxiv.org/abs/2211.17192

왜 LLM은 느린가

• LLM의 구조적인 한계: Auto-regressive 한 디코딩 방식
  • 단어 하나를 출력할 때마다 엄청난 양의 메모리와 계산이 필요

Speculative Decoding으로 해결해보자!

• 목표: LLM의 Inference Speed 향상

• 핵심 아이디어: 생성은 느리지만, 검증은 병렬로 빠르게 할 수 있다

• 구성요소
  • Draft Model (DM): 아주 빠르지만 약간 덜 똑똑한 학생 역할. 대략적인 초안을 빠르게 작성
  • Target Model (TM): 느리지만 아주 똑똑하고 정확한 선생님 역할. Draft 결과를 검토

• 작동원리
  1. Drafting: DM이 먼저 문장의 뒷부분을 추측해서 미래에 올 토큰 $K$개(예: 4개)를 빠르게 생성
     • e.g., "The cat is [sitting on the mat]" (괄호 안이 DM이 추측한 부분)
  2. Verification: TM이 DM이 생성한 $K$개의 토큰을 입력 받고 한 번의(Parallel) 연산으로 $K$개의 토큰 확률을 한번에 (Foward Pass)로 계산
     • 생성(Generation)은 순차적이어야 하지만, 검증(Verification)은 병렬로 할 수 있음
     • Teacher forcing 방식으로 전체 문맥을 한 번에 확인
       • teacher forcing: target word(Ground Truth)를 디코더의 다음 입력으로 넣어주는 기법
         • https://blog.naver.com/sooftware/221790750668
  3. Accept/Reject: 만약 TM이 DM의 출력이 옳다고 판단하면, TM은 승인(Accept)만 하면 됨 (시간 절약!)
     • 중간에 틀린 부분이 있다면(e.g., DM은 "mat"라고 썼는데 TM은 "sofa"라고 생각함), 그 이후의 토큰은 모두 버리고(Reject), 해당 지점부터 다시 생성함

Introduction

• Scaling Law의 현실적인 문제
  • Meta는 최근 4,050억 개의 파라미터를 가진 사상 최대 규모, 최고 성능의 모델인 Llama-3.1-405B를 공개함
  • 이런 대형 모델들은 배포에 막대한 자원을 요구하며, 추론 효율성이 중요한 문제로 떠오름
  • 이런 문제를 해결하기 위해 Speculative Decoding (SD)이 제안됨

• Speculative Decoding의 문제점
  • 기존 Standard Speculative Decoding의 검증 방식:
    • 이 토큰이 문맥적으로 좋은가?가 아니라 이 토큰이 TM이 고를 토큰과 얼마나 일치(alignment)하는가?로 판단함!
      • 문제점 DM이 충분히 맞는 답을 생성해도 TM과 Align이 안되면 초반에 reject가 자주 발생
      • 그래서 실제 환경에서 M (DM이 미리 뽑는 후보 토큰 수) 를 크게 못 키움 (보통 5~7 정도의 작은 값으로 설정)

        

  • RQ 검증 과정을 TM과의 정렬(Alignment)이 아닌 토큰 자체의 품질을 평가하도록 바꿀 순 없을까?
    • 핵심 아이디어: LLM-as-a-judge
      • LLM judge는 유연한 방식으로 답변을 평가함
      • target과 완전히 정렬(Align)되지 않았더라도 올바른(Correct) 응답을 긍정적으로 평가

      ⇒ LLM judge를 통해 alignment과 correctness를 구분해보자!

Contribution

1. 기존 SD가 고품질 토큰을 많이 Reject한다는 한계를 실험적으로 입증

2. LLM-as-a-judge 개념을 SD 검증에 적용

3. Llama-8B/70B-Judge로 최대 9배의 속도 향상, Llama-405B 수준 품질 유지

Judge Decoding

• 기존 Speculative Decoding의 검증 방식

  용어 정리

  • $LLM_{targ}$ : target model

  • $LLM_{draft}$ : draft model

  • $V = \{１,…,V\}$: 어휘 집합 V

  • $M ∈ ℕ$: 후보 토큰의 개수

  • $m^*$: 실제 Accept된 토큰의 개수

  • $s ∈ V^L$: 현재 문맥(Context)

  • $(t₁, p₁), …, (t_m, p_m) = \text{LLM}^{(m)}(s)$
    • 문맥 s가 주어졌을 때, LLM으로부터 m개의 토큰을 auto-regressive 하게 샘플링한 결과
      • $tᵢ$: $i$ 번째로 생성된 토큰
      • $pᵢ$: 해당 시점의 softmax 분포

  • $p₁, …, p_{m+1} = \text{LLM}(t₁,…,t_m; s)$
    • 문맥 $s$ 와 토큰 $t_1, \dots, t_m$을 한 번에 입력
    • Target 모델을 병렬 forward pass로 실행
    • 이 때의 각 위치에서의 확률 분포를 $p₁,…,p_{m+1}$로 정의

  Draft Model

  • DM은 현재 문맥 s를 기준으로 보통 greedy decoding 방식을 통해 M개의 후보 토큰을 생성

  • $(c₁, q₁), …, (c_M, q_M) = \text{LLM}^{(M)}_{draft}(s)$: draft 모델이 생성한 M개의 후보 토큰과 각 토큰 확률
    • $cᵢ$: 후보 토큰
    • $qᵢ$: DM의 Softmax 분포
    • $qᵢ[cᵢ]$: DM이 토큰 $cᵢ$에 부여한 확률

  Target model의 검증 방식(Acceptance Rule)

  1. target 모델은 이 후보 토큰들을 병렬로 처리하여 $p₁,…,p_{M+1}$ 확률 벡터를 생성

  2. 각 후보 토큰 $c_i$에 대하여 아래 조건 검사(이전 토큰들이 모두 Accept된 경우에만 현재 토큰을 검사)

     

  • Acceptance Rule: 이전 토큰들이 모두 Accept되었고, 균등분포에서 샘플한 $εᵢ$가 $pᵢ[cᵢ]/qᵢ[cᵢ]$ (TM의 $C_i$ 생성 확률 나누기 draft 모델에서 $c_i$ 생성 확률) 보다 작으면 토큰 $cᵢ$를 Accept
    • alignment 기반 검증
      • $p_i[c_i]≥q_i[c_i]$ (TM 확률 ≥ DM 확률) ⇒ $\varepsilon_i \in [0,1]$ ⇒ 무조건 accept
      • $p_i[c_i]<q_i[c_i]$(TM 확률 < DM 확률) ⇒ 확률적으로 accept
        • e.g., 비율이 0.3이면 → 30% 확률로 Accept/70%는 Reject

    ⇒ TM이 해당 토큰을 DM보다 더 높은 확률로 평가하면 accept

  ⇒ Standard Speculative Decoding에서는 검증 방식(alignment) 자체의 한계 때문에 draft 토큰 수를 늘려도 acceptance가 포화되어 M을 크게 쓰는 것이 오히려 비효율적임

기존 검증 방식(Alignment)의 한계

RQ 어떤 종류의 토큰들이 거절되는가?

→ GSM8K, MT-Bench, HumanEval 등의 여러 벤치마크에서 SD의 동작을 분석

• draft 모델로는 Llama-8B를, target 모델로는 Llama-405B 사용

• draft 모델도 성능 나쁘지 않기 때문에 acceptance rate를 높여도 품질이 반드시 저하되지는 않음
  • 특히 비교적 단순한 질문의 경우 많은 draft 답변은 그대로 Accept되어도 갠춘
  • GSM8K 같은 문제에서 특히 강함

• 문제: draft 모델이 완전히 정확한 해답을 생성한 경우도, 기존의 검증 방식의 한계로 인해 target 모델이 많은 토큰을 자주 거절함!
  • 이러한 거절의 이유는 DM의 응답의 correctness와 무관하게 발생
  • TM이 문맥적 정확성보다 자기 자신의 응답과의 정렬만 보기 때문

    

• Standard SD의 문제점
  • Standard SD는 토큰을 Accept/Reject 할 때 정렬(alignment)만 봄
  • draft가 만든 토큰이 정답이고 문맥적으로 좋은데도, TM이 선호하는 표현과 다르면 Reject 때려버림

  ⇒ 목표: 후보 토큰이 문맥적으로 올바른(Correct) 경우 Accept하도록 TM을 학습하자!

새로운 검증 방식인 Judge Decoding 제안

• Judge Decoding 목표: TM과의 alignment 말고, 토큰이 틀렸는지/맞는지(correctness)를 판단해서 Accept를 늘리자!!
  • LLM-as-a-Judge에서 착안
  • But, LLM-as-a-Judge의 문제점
    1. 긴 시스템 프롬프트와 CoT 추론이 필요한데 이것들이 추론 속도를 저하시킴
    2. LLM judge는 전체 답변을 평가하는 방식을 사용하는데 SD는 짧고 부분적인 연속 토큰을 평가해야 한다는 점에서 실용적이지 않음

    ⇒ 기존 방식을 장점은 살리되, LLM-as-a-Judge 의 느낌은 살리도록 설계해보자!

임베딩은 이미 오류를 알고 있다..!

• TM은 잘못된 토큰을 처리하면, 마지막 hidden layer embedding에서 이상 감지 신호를 발생시킴!
  • 잘못된 토큰의 마지막 hidden layer embedding 에 오류가 표시(flag) 돼있음

    

  • 모델은 이후 토큰에서 해당 오류를 수정하려는 방향으로 출력을 생성하게 됨
    (위 그림의 왼쪽 어시스턴트 참고)

  모델은 이미 이 토큰이 틀렸다는 걸 내부적으로 알고 있다..!

Judge Head

• TM의 embedding 위에 붙는 작은 이진(binary) 분류기
  • 목적: target 임베딩에 담긴 신호를 이용해, 각 후보 토큰 $c_i$에 대해 이 토큰을 문맥상(correctness) 통과시켜도 되는가?를 빠르게 판정
  • 입력: 토큰 임베딩 $e_i$
  • 출력: 해당 토큰이 통과(accept) 가능할 확률(score) ⇒ $σ(f_{judge}(e_i))$
  • 결정: 임계값 $\delta$를 넘으면 Accept ⇒ $\sigma(f_{judge}(e_i)) > \delta$

• Linear head (logistic regression)로 구현
  • 이미 target embedding에 “오류 신호”가 존재 → 복잡한 모델이 필요 없음
    • 복잡한 MLP/Transformer는 오히려 과적합 위험
  • 장점
    • 파라미터가 매우 작고(약 16.4k)
    • 학습 빠름 (~1.5h)
    • TM의 파라미터는 동결(frozen)

Judge head 학습을 위한 데이터셋 구축

• 총 500개의 고품질 사용자 질문과 그에 대한 정답/오답 답변 쌍으로 구성된 데이터셋 구축
  • 새롭게 생성한 질문과 Alpaca와 ARC 데이터셋에서 필터링한 질문들 사용
  • 입력 질문만 사용하고 그에 대한 정답은 사용하지 않았음!

• Mistral-Large-2, Llama-8B, Llama-405B을 활용해 정답과 오답을 다양하게 생성
  • 인간이 실제로 오류 토큰에 주석도 달음

• 학습 과정: 정답 답변에 포함된 모든 토큰을 positive 로 라벨링, 오답 답변에서는 오류가 발생하기 전까지의 모든 토큰을 positive로, 오류가 발생한 토큰들을 negative으로 라벨링하였음
  • positive가 negative 보다 20배 많음

모델 설계 및 학습

• 구축한 데이터셋을 바탕으로 target 모델의 임베딩 위에 linear head인 $f_{judge}$를 학습시킴
  • 데이터 불균형을 보정하기 위해 가중 cross-entropy loss 사용
  • 잘못된 토큰을 잘 잡아내도록 negative 샘플에 더 큰 가중치를 둠

• 학습 파라미터: 16.4k

• 학습 데이터: 30k 토큰

• 학습 시간: 1.5시간 이내

• target 모델 파라미터: 고정(frozen)

실제 Inference 과정(Judge + Standard SD 결합)

• TM이 후보 토큰을 보고 두 가지 질문을 동시에 함
  1. alignment: 내가 원래 뽑을 토큰이랑 비슷해?
  2. correctness: 문맥상 맞는 토큰인지

  ⇒ 둘 중 하나라도 Yes면 Accept

==== 추론 과정 ====

1. DM이 후보 토큰 M개 생성 (기존 방식과 동일)

   $c_1, c_2, \dots, c_M$

2. TM이 후보 토큰을 한 번에 검토 (2개의 판정 동시 수행)
   • Standard SD 검증 마스크 $z_{stand}$
     • 기준: alignment (확률 비율)
       • $z_{stand}[i] = 1$ → Accept
       • $z_{stand}[i] = 0$→ Reject
   • Judge 마스크 $z_{judge}$
     • 기준: correctness (embedding 기반)
       • 토큰 임베딩 $e_i$ 를 보고 judge head가 점수 계산
     • 계산: ${Accept\ if\ } \sigma(f_{judge}(e_i)) > \delta$
     • 결과:
       • 점수 > δ $\delta$$z_{judge}[i] = 1$→ Accept
       • 점수≤ δ → $z_{judge}[i] = 0$ → Reject

3. 최종 Accept/Reject는 OR로 결합
   • $z = z_{stand} \lor z_{judge}$
     • Standard SD가 Accept면 → 무조건 Accept
     • Standard SD가 Reject여도, Judge가 Accept면 → Accept

Experiment

Draft 품질이 아주 좋은 경우 (GPT-4o)

• 실험 목적: DM의 올바른 응답(correct)조차 높은 Reject을 겪는다는 점을 증명하기 ㅜ이해

• 실험 셋업:
  • DM: GPT-4o
  • TM: Llama-405B
  • 데이터셋: MT-Bench, GSM8K, HumanEval

• 실험 방식: 데이터셋의 질문에 대한 전체 답변을 생성한 뒤, greedy 검증에서 첫 거절이 발생하기 전까지 TM이 몇 개의 토큰을 Accept하는지 측정

• 실험 결과:

  

  • Standard SD acceptance: 약 2개 토큰 Accept
  • Judge SD acceptance: 20~27 토큰 accept

• insight
  • Draft의 품질이 좋아진다고 acceptance가 좋아지지 않는다!
  • Judge Decoding 방식을 쓰자..!

실험 세팅을 반대(draft ↔ target 모델 바꾸기)로 해도 결과는 동일!

• draft 모델: Llama-405B

• target 모델: GPT-4o (위 실험과 반대 세팅)

• 실험 결과

  

  • 8B/405B 일 때의 acceptance ≈ 6.6 토큰
  • 405B/8B 일 때의 acceptance ≈ 6.3 토큰

    → 거의 차이 없음

Human expert drafting

• 인간 전문가가 생성한 Draft 토큰(성능 최상)을 검증해보자!
  • 고품질의 커뮤니티 검증 요약문을 포함한 wikipedia-summary 데이터셋 일부를 사용해, greedy SD 검증에서의 토큰 Accept률을 평가

• 실험 결과
  • Human text (draft) → Llama-405B

    

  • Standard SD 일 때의 acceptance ≈ 3.1 토큰
  • Judge SD 일 때의 acceptance ≈ 12.3 토큰
  • → human이 직접 draft 한 토큰들도 reject 때리는거면 TM의 기존 검증 방식 문제 확실히 있음!

⇒ target model과의 alignment를 통해 검증하는 기존의 검증 방식이 문제 있음!!

Judge Decoding Benchmark Results

• 실험 세팅
  • 모델: DM (Llama-8B) + TM (Llama-70B/405B)
  • 비교 방법(Decoding / Verification)
    • Draft only (Llama-8B)
    • Target only (Llama-70B, 405B)
    • Top-K verification (표준 SD 완화 휴리스틱): Target 모델이 보기에 확률이 높은 K개 안에 들면 그냥 통과시키는 검증 방식
    • Judge Decoding
  • 하이퍼파라미터
    • Top-K: M=10
    • Judge Decoding: M=25
  • 벤치마크: GSM8K, HumanEval, ARC-Challenge, MMLU, MT-Bench

• 실험 목적
  1. Top-K (휴리스틱 완화): 검증 기준을 느슨하게 했을 때 정확도가 얼마나 무너지는지 확인
  2. Judge Decoding이 더 긴 후보 시퀀스(M=25) 를 accept하면서도 Target 수준 정확도를 유지하는지 검증

• 실험 결과
  • Top-K는 정확도 크게 하락
    • 일부 벤치마크에서 두드러진 성능 저하 발생
    • 상위 K 안에만 들면 accept하는 Top-K 방식은 틀린 토큰도 통과시키기 쉬움
    • 품질-속도 trade-off가 심하게 발생(속도는 빠르지만 품질은 나쁘다는 뜻)
  • Judge Decoding은 정확도 거의 보존
    • 모든 벤치마크에서 Target-only 대비 거의 차이 없이 유지
    • 즉, 약 20개 수준 토큰을 한 번에 더 많이 accept해도 품질이 깨지지 않음을 보여줌
  • 70B/405B 모두에서 일관된 경향
    • Target이 커져도(70B→405B) 결과 패턴이 유지됨

      → 방법이 특정 모델 크기에만 맞는 트릭이 아니라는 근거

• 인사이트
  • Judge Decoding은 top-K보다 2.5배 많은 샘플을 생성하면서도 더 높은 정확도 달성
    → 샘플 개수보다 품질과 검증 메커니즘이 중요
    • Top-K: 단순 규칙이라 correctness 판단 실패 → 정확도 하락
    • Judge Decoding: 임베딩 기반 correctness 판정 → 더 길게 accept해도 정확도 유지
  • Target 모델이 직접 최종 토큰을 선택하여 모델 정렬(alignment) 문제를 근본적으로 해결

분포 외 일반화 실험

• 학습되지 않은 상황에서도 judge decoding이 얼마나 일반화되는지를 평가
  • 코딩 예제를 제거한 데이터로 judge를 학습한 뒤, HumanEval에서 평가

• 성능이 86.6%에서 80.4%로 하락하긴 했지만, 여전히 DM(71.3%)보다는 훨씬 높음