Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs

1. 문제 정의 (Problem)

Speculative Decoding 은 작은 '드래프트 모델 (Draft Model)'이 토큰 시퀀스를 예측하고, 이를 큰 '타겟 모델 (Target Model)'이 병렬로 검증하여 추론 속도를 높이는 기술입니다.

• 현재의 한계: 단일 드래프트 모델은 특정 작업 (예: 수학, 코딩) 에서는 뛰어나지만 다른 작업에서는 성능이 급격히 떨어지는 경우가 많습니다.
• 질문: 여러 개의 후보 드래프트 모델 (예: 수학 전용, 코딩 전용, 일반용 등) 이 주어졌을 때, **각 입력 쿼리 **(Query)
• 기존 접근법: MetaSD 나 BanditSpec 같은 기존 연구들은 이를 **다중 암 밴디트 **(Multi-armed Bandit) 문제로 모델링했습니다. 이는 '탐색 (Exploration)'과 '이용 (Exploitation)'을 균형 있게 조절해야 하므로, 최적의 모델을 찾기까지 많은 시간과 리소스가 소모되며 수렴 속도가 느립니다.

2. 방법론 (Methodology: HedgeSpec)

저자들은 Speculative Decoding 의 구조적 특성을 활용하여 **탐색 **(Exploration)을 증명했습니다. 이를 통해 문제를 밴디트 문제에서 **풀 정보 **(Full-Information)로 전환했습니다.

핵심 아이디어: 풀 정보 피드백 (Full-Information Feedback)

• 기존 밴디트 방식: 선택된 드래프트 모델의 결과만 관찰 (부분 정보).
• HedgeSpec 방식: 타겟 모델이 검증한 **단 하나의 토큰 시퀀스 **(Verified Trajectory)를 사용하여, **선택되지 않은 모든 다른 드래프트 모델들이 어떻게 반응했을지 **(Counterfactual)를 추정합니다.
  • 타겟 모델에 대한 추가 호출 없이도, 모든 후보 모델의 '수용 확률 (Acceptance Probability)'과 '기대 수용 길이 (Expected Acceptance Length)'를 계산할 수 있습니다.
  • 이는 Theorem 3에서 증명된 바와 같이, 편향되지 않은 추정량 (Unbiased Estimator) 을 제공합니다.

알고리즘 및 학습 과정

1. 피드백 수집: 타겟 모델이 토큰을 검증한 후, 해당 시퀀스를 모든 다른 드래프트 모델에 대해 'prefill'하여 각 모델이 해당 토큰을 수용할 확률 ($\gamma_t[i]$) 을 계산합니다.
2. 손실 함수 (Loss Function):
   • 수용 길이 최적화: $ft[i] = 1 - \frac{1}{K+1} \sum k(1-\gamma)\dots$ 형태로, 최대 가능한 길이 대비 기대 수용 길이의 차이를 손실로 정의합니다.
   • 수용 확률 최적화: $ft[i] = 1 - \gamma_t[i]$를 사용합니다.
3. 지연된 피드백 처리 (Delayed Feedback):
   • 블록 (Chunk) 단위로 검증이 완료되므로, 피드백이 즉시 제공되지 않습니다. 이는 **지연된 피드백 **(Delayed Feedback) 문제로 모델링됩니다.
   • Joulani et al. (2013) 의 알고리즘을 적용하여 지연을 처리하며, NormalHedge 또는 Hedge 알고리즘을 기반으로 드래프트 모델의 가중치를 업데이트합니다.
4. 시스템 효율성: 모든 드래프트 모델의 평가를 병렬로 수행하거나, 업데이트 주기를 조절하여 지연 시간을 최소화합니다.

3. 주요 기여 (Key Contributions)

1. 풀 정보 온라인 학습 프레임워크: 드래프트 모델 선택 문제에서 탐색이 불필요하며, 모든 모델에 대한 정확한 피드백을 얻을 수 있음을 최초로 증명했습니다.
2. **이론적 보장 **(No-Regret Guarantee):
   • 수용 확률과 기대 수용 길이 두 가지 목적 함수 모두에 대해 **Regret **(후회)을 증명했습니다.
   • 드래프트 모델의 수 ($N$) 가 증가함에 따라 밴디트 기반 방법 ($O(\sqrt{N})$) 보다 지수적으로 빠르게 최적 모델에 수렴함을 보였습니다.
3. 실제 시스템 효율성: 타겟 모델 호출 수를 늘리지 않으면서도, 평가 오버헤드를 최소화하는 시스템 설계 (예: 지연된 업데이트, 배치 처리) 를 제안했습니다.
4. 광범위한 실험 검증: 오픈소스 LLM (Llama-3.1, Qwen-3) 과 다양한 도메인 (수학, 코딩, 의학, 화학 등) 에서 21 개의 전문 드래프트 모델을 사용하여 검증했습니다.

4. 실험 결과 (Results)

저자들은 Llama-3.1-8B, Qwen-3-8B, Qwen-3-32B 를 타겟 모델로 하여 7 가지 도메인 (Python, Math, Biology 등) 에서 실험을 수행했습니다.

• 성능 향상:
  • EAGLE3(최고의 단일 일반 모델): 평균적으로 **46.1%**의 토큰 생성 속도 (Token/s) 향상. 특정 도메인 (예: SQL) 에서는 최대 **83.7%**까지 향상되었습니다.
  • BanditSpec(UCB, EXP3 기반): 평균적으로 **49%**의 평균 수용 토큰 수 (MAT) 향상 및 **41%**의 속도 향상을 기록했습니다.
• **확장성 **(Scalability):
  • 드래프트 모델 풀 (Pool) 의 크기가 커질수록 밴디트 기반 방법은 성능이 급격히 떨어지지만, HedgeSpec 은 드래프트 모델 수가 증가해도 성능이 유지되거나 향상되었습니다. 이는 풀 정보 피드백이 더 많은 전문가 모델을 효과적으로 활용하게 하기 때문입니다.
• **분포 외 **(OOD):
  • 오프라인 라우터 (Static Router) 는 훈련 분포와 다른 입력 (OOD) 에 대해 성능이 급격히 저하되지만, HedgeSpec 은 런타임 피드백을 통해 실시간으로 적응하여 2.34 배까지 더 나은 성능을 보였습니다.

5. 의의 및 결론 (Significance)

• 패러다임 전환: Speculative Decoding 의 드래프트 선택 문제를 '탐색이 필요한 밴디트 문제'가 아닌 '풀 정보가 가능한 온라인 학습 문제'로 재정의했습니다.
• 실용성: 훈련 데이터에 접근할 수 없거나, 다양한 도메인 전문가 모델들이 따로 존재하는 실제 환경 (Cloud Serving) 에서 매우 효과적입니다.
• 일반성: EAGLE, REST, Medusa 등 어떤 Speculative Decoding 프레임워크와도 호환되어 적용 가능합니다.

결론적으로, HedgeSpec은 여러 전문가 드래프트 모델들을 동적으로 조율하여, 단일 모델이나 기존 밴디트 기반 방법보다 훨씬 빠르고 효율적인 LLM 추론을 가능하게 하는 No-Regret 솔루션입니다.