Bridging Draft Policy Misalignment: Group Tree Optimization for Speculative Decoding

이 논문은 추론 속도를 높이는 스펙큘레이티브 디코딩에서 드래프트 모델의 학습 목표와 디코딩 시의 트리 정책 간 불일치를 해결하기 위해, 드래프트 트리의 기대 수용 길이를 직접 최적화하는 '그룹 트리 최적화 (GTO)'를 제안하여 기존 최첨단 방법보다 더 높은 수용 길이와 가속도를 달성함을 보여줍니다.

핵심

🚀 GTO: AI 가 더 빨리, 더 똑똑하게 말하게 만드는 새로운 방법

이 논문은 거대한 인공지능 (LLM) 이 글을 쓸 때, 속도를 높이면서도 실수를 줄이는 새로운 방법을 소개합니다. 기존 방법들의 약점을 찾아내고, 이를 해결하는 'GTO'라는 기술을 제안했죠.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 문제: "예측하는 사람"과 "검증하는 사람"의 불일치

인공지능이 글을 쓸 때, 보통 한 글자씩 천천히 씁니다. (예: "안녕" -> "하세요" -> "반갑습니다") 이렇게 하면 너무 느리죠.

그래서 **'스펙큘레이티브 디코딩 (Speculative Decoding)'**이라는 기술을 썼습니다.

• 드래프트 모델 (Draft Model): 작고 빠른 AI 가 먼저 "다음 글자는 뭐가 될까?"라고 여러 가지 후보를 미리 내밉니다. (예: "안녕하세요", "안녕하십니까", "안녕하세요?" 등)
• 타겟 모델 (Target Model): 크고 똑똑한 AI 가 그 후보들을 한 번에 확인하고, 맞는 것만 채택합니다.

🚨 여기서 문제가 생겼습니다.
기존에는 훈련 (학습) 할 때와 **실제 쓸 때 (추론)**의 방식이 달랐습니다.

• 훈련할 때: 작은 AI 는 "가장 확률이 높은 하나의 길만 골라" 정답을 맞추는 연습을 했습니다. (예: "안녕하세요"만 골라 맞췄다)
• 실제 쓸 때: 작은 AI 는 **여러 가지 길 (나무 가지)**을 모두 만들어 내고, 큰 AI 가 그중 가장 좋은 것을 골라냅니다.

🌲 나무 비유:
훈련할 때는 작은 AI 가 "가장 높은 나뭇가지 하나만 올라가면 돼"라고 배웠습니다. 하지만 실제로는 큰 AI 가 "모든 나뭇가지를 다 보고, 그중 가장 튼튼한 가지를 골라라"라고 지시합니다.
결과적으로, 훈련할 때 열심히 배운 "하나의 가지"가 실제 상황에서는 버려지거나, 더 좋은 가지가 있는데도 불구하고 쓸모없는 가지를 고집하게 되어 속도 향상 효과가 반감됩니다.

---

2. 해결책: GTO (Group Tree Optimization)

이 논문은 이 불일치를 해결하기 위해 GTO를 제안합니다. 핵심은 **"훈련할 때부터 실제 상황 (나무 전체) 을 그대로 시뮬레이션하자"**는 것입니다.

🌟 핵심 아이디어 1: "나무 전체의 점수"를 가르친다 (Draft Tree Reward)

기존에는 "다음 글자가 맞았나?" (단일 글자 점수) 를 보았지만, GTO 는 **"내가 만든 나무 전체가 얼마나 잘 받아들여졌나?"**를 점수로 매깁니다.

• 비유: 시험을 볼 때, 단순히 "단어 하나를 맞췄나?"를 보는 게 아니라, "내가 쓴 전체 문장이 얼마나 자연스럽게 이어졌나?"를 평가하는 것과 같습니다.
• 이렇게 하면 작은 AI 는 "하나의 길"만 고집하지 않고, 여러 가지 가능성을 만들어내어 큰 AI 가 선택할 수 있도록 훈련받게 됩니다.

🌟 핵심 아이디어 2: "비교 그룹"을 만들어 안정적으로 학습 (Group-Based Optimization)

새로운 방식을 가르치려면 학습이 불안정해질 수 있습니다. (예: "아, 내가 만든 나무가 왜 안 좋은 거지?"라고 혼란스러워함)

• GTO 는 비교 그룹을 만듭니다.
  • 현재 AI: 지금 만든 나무.
  • 참고 AI (Frozen): 예전에 잘 훈련된 AI 가 만든 나무.
• 이 두 나무를 비교해서, "어떤 부분이 더 좋았는지"를 명확하게 알려줍니다.
• 비유: 운동 선수 (현재 AI) 가 코치 (참고 AI) 와 함께 훈련할 때, "네가 오늘 한 동작은 코치보다 10% 더 좋았어, 하지만 이 부분은 5% 나빠"라고 구체적인 피드백을 주는 것과 같습니다. 이렇게 하면 학습이 흔들리지 않고 빠르게 발전합니다.

---

3. 결과: 얼마나 빨라졌나요?

이 새로운 방법 (GTO) 을 적용한 결과, 기존에 가장 좋았던 방법 (EAGLE-3) 보다 약 7.7% 더 빨라졌습니다.

• 숫자로 보면: AI 가 한 번에 받아들이는 글자 수가 평균 56 개에서 **67 개**로 늘어났습니다.
• 실생활 비유:
  • 기존 방법: 택시 (AI) 가 1 분에 500m 를 갔다면,
  • GTO 방법: 같은 시간 동안 535m를 갔습니다.
  • 겉보기엔 작아 보이지만, 수천 번의 대화를 이어갈 때 엄청난 시간 절약이 됩니다.

이 효과는 대화 (MT-Bench), 코딩 (HumanEval), 수학 문제 (GSM8K) 등 모든 분야에서 똑같이 나타났습니다. 특히 복잡한 코딩이나 수학 문제처럼 논리가 중요한 작업에서 더 큰 향상을 보였습니다.

---

4. 요약: 왜 이것이 중요한가요?

• 기존의 문제: AI 가 훈련할 때와 실제 쓸 때 전략이 달라서, 속도를 높이는 데 한계가 있었습니다. (훈련은 '하나의 길', 실전은 '여러 가지 길')
• GTO 의 해결: 훈련할 때부터 **여러 가지 길 (나무)**을 만들고, 그중 가장 좋은 것을 선택하는 방식을 배웁니다.
• 결과: AI 가 더 빠르게, 더 정확하게 글을 쓸 수 있게 되었습니다.

한 줄 요약:

"AI 가 글을 쓸 때, 하나의 길만 고집하는 훈련을 버리고, 여러 가지 가능성을 열어두고 최선의 길을 찾는 훈련을 시켜주니, 속도가 훨씬 빨라졌습니다!"

이 기술은 AI 가 더 효율적으로 작동하게 만들어, 우리가 더 빠르게 답변을 받을 수 있게 해주는 실용적인 혁신입니다.

기술 요약

1. 문제 정의: 드래프트 정책 불일치 (Draft Policy Misalignment)

대형 언어 모델 (LLM) 의 추론 속도를 높이기 위해 Speculative Decoding(추측적 디코딩) 이 널리 사용되지만, 기존 방법론에는 근본적인 한계가 존재합니다.

• 훈련과 추론의 불일치:
  • 훈련 단계: 기존 드래프트 모델 (Draft Model) 은 단일 경로 (Greedy Path) 에서 다음 토큰의 확률을 최대화하도록 훈련됩니다. 즉, 매 단계에서 가장 확률이 높은 토큰 하나만 선택하는 'Greedy Drafting'을 최적화합니다.
  • 추론 (디코딩) 단계: 실제 추론 과정에서는 단일 경로가 아닌 트리 구조 (Tree Policy) 를 사용합니다. 드래프트 모델이 여러 가지 후보 토큰을 생성하고, 이를 타겟 모델이 병렬로 검증하며, 최종적으로 가장 신뢰도가 높은 경로를 선택합니다.
• 결과적 비효율: 훈련 시 최적화된 'Greedy Path'가 실제 디코딩 시 트리 재순위화 (Re-ranking) 과정에서 잘려나가거나 (Pruning), 타겟 모델이 다른 경로를 선택하는 경우, 훈련에 투자된 노력이 낭비됩니다.
• 실증적 증거: 실험 결과, 훈련 시 Greedy Path 의 약 1934% 가 디코딩 과정에서 잘려나갔으며, 최종적으로 수용된 경로가 Greedy Path 와 일치하는 비율은 3649% 에 불과했습니다. 이는 추론 효율성을 제한하는 주요 병목 현상임을 보여줍니다.

2. 제안 방법: 그룹 트리 최적화 (Group Tree Optimization, GTO)

저자들은 훈련과 추론 간의 정책 불일치를 해결하기 위해 GTO를 제안했습니다. 이는 단일 경로가 아닌 드래프트 트리 전체를 최적화 대상으로 삼는 새로운 학습 프레임워크입니다.

핵심 구성 요소

1. 드래프트 트리 보상 (Draft Tree Reward):

   • 기존 방식 (토큰 단위 정확도) 이 아닌, 디코딩 시의 실제 성능을 직접 측정하는 보상을 정의합니다.
   • 드래프트 모델이 생성한 트리 (여러 가지 시퀀스) 를 타겟 모델이 검증했을 때, 기대 수용 길이 (Expected Acceptance Length) 를 보상으로 사용합니다.
   • 모든 시퀀스의 기대 수용 길이를 평균내는 것이 아니라, Log-Sum-Exp (LSE) 를 사용하여 강건한 분기 (Strong Branches) 에 더 큰 가중치를 두면서도 미분 가능한 형태로 보상을 계산합니다.
   • 이론적 보장: 이 보상을 최대화하면 타겟 모델의 샘플링 온도와 무관하게 기대 수용 길이가 증가함이 수학적으로 증명되었습니다.

2. 그룹 기반 드래프트 정책 학습 (Group-based Draft Policy Training):

   • 희소하고 분산이 큰 트리 보상을 안정적으로 학습하기 위해 그룹 기반 강화학습 (Group-based RL) 전략을 도입했습니다.
   • 그룹화 (Grouping): 동일한 시퀀스 내 인접한 위치 (Prefix) 들을 그룹으로 묶습니다. 이는 문맥이 유사한 경우끼리 비교하여 위치별 난이도 편향을 제거하고 분산을 줄입니다.
   • 편향 제거 (Debiasing): 현재 모델과 고정된 참조 모델 (Frozen Reference Model, 예: EAGLE-3 로 훈련된 모델) 이 생성한 트리의 보상을 비교하여 차이를 계산합니다. 이를 통해 문맥의 난이도에 따른 시스템적 편향을 제거하고, 모델의 실제 개선분을 포착합니다.
   • PPO 스타일 최적화: 그룹 내에서 표준화된 이점 (Advantage) 을 계산하고, 가장 긴 수용 시퀀스 (Longest Accepted Sequence) 에 대해 PPO 의 Clipped Surrogate Objective 를 적용하여 안정적인 업데이트를 수행합니다.

3. 주요 기여 (Contributions)

1. 정책 정렬 (Policy Alignment): 훈련 목표 (단일 Greedy Path) 와 추론 목표 (Tree Policy) 간의 불일치를 해결하는 최초의 체계적인 접근법을 제시했습니다.
2. 새로운 보상 함수: 추론 효율성을 직접적으로 반영하는 'Draft Tree Reward'를 도입하여, 훈련된 모델이 실제 디코딩 환경에서 더 긴 토큰 시퀀스를 수용하도록 유도합니다.
3. 안정적인 학습 알고리즘: 희소 보상 문제를 해결하기 위한 그룹 기반 편향 제거 및 표준화 기법을 개발하여, 학습의 안정성과 수렴 속도를 보장합니다.
4. 이론적 증명: 드래프트 트리 보상 최대화가 기대 수용 길이 증가로 이어진다는 것을 수학적으로 증명했습니다.

4. 실험 결과

다양한 LLM (LLaMA-3.1-8B, 70B, Vicuna-13B, DeepSeek-R1, Qwen3 등) 과 벤치마크 (MT-Bench, HumanEval, GSM8K) 에서 실험을 수행했습니다.

• 성능 향상:
  • 기존 State-of-the-Art 인 EAGLE-3 대비 수용 길이 (Acceptance Length) 는 평균 7.4% 증가했습니다.
  • 속도 향상 (Speedup Ratio) 은 평균 7.7% 추가 개선되었습니다.
  • 특히 HumanEval(코드 생성) 과 GSM8K(수학 추론) 과 같은 구조화된 생성 작업에서 두드러진 성능 향상을 보였습니다.
• 범용성:
  • EAGLE-3 뿐만 아니라 GRIFFIN, HASS 등 다른 방법으로 훈련된 드래프트 모델을 GTO 로 미세 조정 (Fine-tuning) 했을 때도 일관된 성능 향상을 보였습니다.
  • 온도 (Temperature) 가 0 인 Greedy 디코딩 환경과 1 인 확률적 디코딩 환경 모두에서 우수한 성능을 발휘했습니다.
• Ablation Study:
  • 보상 집계: 단순 평균 (Sum) 이나 최대값 (Max) 보다 LSE 기반의 집계 방식이 더 우수함을 확인했습니다.
  • 그룹 크기: 그룹 크기 $m=8$이 최적의 균형을 제공하며, 너무 작거나 크면 성능이 저하됨을 확인했습니다.
  • 편향 제거: 참조 모델을 이용한 편향 제거 (Debiasing) 가 학습 안정성과 최종 성능에 결정적인 역할을 함을 입증했습니다.
  • 계산 비용: 단순히 학습 시간을 늘리는 것 (Continual Training) 이 아니라, GTO 알고리즘 자체의 개선이 성능 향상의 원인임을 확인했습니다.

5. 의의 및 결론

이 논문은 Speculative Decoding 의 효율성을 극대화하기 위한 핵심 병목인 '훈련 - 추론 정책 불일치' 문제를 명확히 규명하고, 이를 해결하는 GTO를 제안했습니다.

• 실용성: 드래프트 모델의 아키텍처를 변경하거나 검증 스택을 수정할 필요 없이, 기존 훈련된 모델 위에 GTO 를 적용하여 추가적인 추론 속도 향상을 얻을 수 있습니다.
• 일반성: 다양한 모델과 작업에 적용 가능한 범용 솔루션으로, LLM 추론의 지연 시간 (Latency) 을 줄이는 데 중요한 기여를 할 것으로 기대됩니다.
• 비용 효율성: 학습 단계에서의 추가 계산 비용은 추론 단계에서 얻어지는 큰 속도 향상 (7% 이상) 으로 상쇄될 수 있어, 실제 배포 환경에서 매우 효율적인 접근법입니다.

결론적으로, GTO 는 추측적 디코딩의 잠재력을 최대한 끌어올리기 위해 훈련과 추론을 정렬시키는 획기적인 방법론으로 평가됩니다.