LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

이 논문은 KL 발산 대신 직접적인 수용률 최적화를 목표로 하는 'LK 손실'을 제안하여, 다양한 크기의 모델과 도메인에서 스펙큘레이티브 디코딩의 수용 길이를 8~10%까지 향상시키고 구현의 용이성과 추가 비용 부재를 입증했습니다.

핵심

🚀 1. 배경: 왜 '예측'이 필요할까요?

대형 인공지능이 글을 쓸 때는 한 글자씩 순서대로 써야 합니다. 마치 한 번에 한 칸만 움직이는 말처럼요. 이 방식은 메모리 대역폭이라는 병목 현상 때문에 속도가 느립니다.

**'Speculative Decoding(예측적 디코딩)'**은 이 문제를 해결하기 위해 두 명의 작가를 고용합니다.

1. 주작가 (Target Model): 거대하고 똑똑하지만 느립니다. (예: 685B 파라미터 모델)
2. 보조 작가 (Draft Model): 작고 빠르지만, 실수는 좀 합니다. (예: 8B~70B 모델)

작동 원리:
보조 작가가 "주작가, 다음 글자는 A, B, C, D, E, F, G 순서로 올 거예요!"라고 7 개를 한 번에 예측합니다.
그리고 주작가는 이 7 개를 한 번에 확인합니다.

• "아, A 는 맞네! B 는 맞네! C 는 틀렸네!"
• 틀린 글자 (C) 이후의 글자는 모두 버리고, A 와 B 만 받아서 다음 단계를 진행합니다.

핵심: 보조 작가가 얼마나 많이 맞히느냐 (Acceptance Rate) 에 따라 속도가 결정됩니다. 7 개를 다 맞히면 속도가 7 배 빨라지지만, 1 개만 맞혀도 속도는 barely 느려집니다.

---

🎯 2. 문제점: "완벽한 모방"은 속도를 보장하지 않습니다.

기존에는 보조 작가를 훈련시킬 때 **"주작가의 말투를 완벽하게 따라해라"**라고 가르쳤습니다. 수학적으로는 **KL 발산 (KL Divergence)**이라는 지표를 최소화하는 방식입니다.

• 비유: 주작가가 "오늘 날씨가 좋네요"라고 했을 때, 보조 작가도 "오늘 날씨가 좋네요"라고 완벽하게 똑같이 말하도록 훈련하는 것입니다.

하지만 여기서 함정이 있습니다.
보조 작가는 주작가보다 훨씬 작고 능력치가 낮습니다. (주작가가 100 점짜리라면 보조 작가는 10 점짜리입니다.)
작은 능력으로 거대한 주작가의 말투를 완벽하게 따라가는 것은 불가능합니다. 그래서 보조 작가는 "완벽한 모방"을 하려다 보니, 실제로는 주작가가 가장 많이 고를 만한 글자를 찍는 데 실패하고, 엉뚱한 곳에 에너지를 쏟게 됩니다.

• 결과: "말투는 비슷해졌는데 (KL 이 줄어듦), 정작 주작가가 고를 글자를 맞추는 확률 (Acceptance Rate) 은 여전히 낮다."

---

💡 3. 해결책: LK Loss (직접적인 목표 달성)

이 논문은 **"말투를 비슷하게 만드는 게 아니라, 주작가가 고르는 글자를 맞추는 것 자체를 훈련 목표"**로 바꾸자고 제안합니다. 이를 LK Loss라고 부릅니다.

🏋️‍♂️ 비유: 마라톤 선수의 훈련

• 기존 방식 (KL Loss): "주작가 (마라토너) 의 발걸음 리듬을 100% 똑같이 따라해."
  • 작은 보조 작가는 리듬을 완벽히 따라 하려다 지쳐버리고, 실제 달리기 (예측) 는 느려집니다.
• 새로운 방식 (LK Loss): "주작가가 어디로 갈지 정확히 맞춰봐. 리듬은 나중에 맞춰도 돼."
  • 보조 작가는 주작가가 가장 확실히 갈 길을 먼저 예측하는 데 집중합니다.

🧩 두 가지 새로운 훈련법

저자들은 이 목표를 달성하기 위해 두 가지 방법을 썼습니다.

1. 직접적인 확률 최적화 (Likelihood-based):

   • "네가 예측한 글자가 주작가에 의해 받아들여질 확률을 최대화해."
   • 마치 도박에서 "어떤 숫자가 나올 확률이 가장 높은지"만 계산하는 것과 같습니다.

2. 혼합 훈련 (Hybrid Objective - Adaptive Blending):

   • 초반: 보조 작가가 아직 서툴 때 (주작가와 많이 다름) 는 "리듬을 비슷하게 잡아라 (KL)"라고 가르칩니다. 그래야 방향을 잃지 않고 훈련이 시작됩니다.
   • 후반: 어느 정도 비슷해지면, "리듬은 그만, 정답을 맞춰라 (TV Distance)"로 목표를 바꿉니다.
   • 비유: 처음에는 "주작가처럼 걷는 법"을 배우다가, 실력이 늘면 "주작가가 가는 길을 정확히 예측하는 법"으로 훈련을 전환하는 것입니다.

---

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

이 논문은 다양한 크기의 인공지능 (8 억 개 파라미터부터 6850 억 개까지) 과 다양한 작업 (일반 대화, 코딩, 수학) 에서 실험했습니다.

• 성공: 기존 방식보다 평균적으로 8~10% 더 많은 글자를 한 번에 맞췄습니다.
• 특징: 보조 작가가 작을수록 (능력이 부족할수록) 이新方法의 효과가 더 컸습니다. 작은 보조 작가일수록 "완벽한 모방"보다는 "핵심 예측"이 중요하기 때문입니다.
• 장점: 훈련 속도를 늦추지 않고, 기존 시스템에 바로 적용할 수 있습니다.

📝 한 줄 요약

"작은 보조 AI 가 거대한 주 AI 를 완벽하게 흉내 내려고 애쓰지 말고, 주 AI 가 다음에 무엇을 고를지 맞추는 데만 집중하게 훈련시키면, AI 가 글을 쓰는 속도가 훨씬 빨라진다!"

이 논문은 AI 의 속도를 높이기 위해, **무엇을 훈련시킬지 (목표 함수)**를 단순히 '모방'에서 '성공'으로 바꾼 획기적인 아이디어를 제시했습니다.

기술 요약

1. 문제 정의 (Problem)

• 배경: 대규모 언어 모델 (LLM) 추론은 계산 처리량보다는 메모리 대역폭에 의해 제한받으며, 자동회귀 (autoregressive) 방식은 이를 더욱 악화시킵니다. 이를 해결하기 위해 Speculative Decoding(가상 모델이 토큰을 제안하고 타겟 모델이 병렬로 검증하는 방식) 이 도입되었습니다.
• 핵심 병목: Speculative Decoding 의 가속화 정도는 **수용률 (Acceptance Rate)**에 의해 결정됩니다. 즉, 제안된 토큰이 타겟 모델에 의해 얼마나 자주 받아들여지는지가 중요합니다.
• 기존 방법의 한계: 현재 표준적인 훈련 방법은 타겟 모델과 가상 (Draft) 모델 간의 분포 차이를 최소화하는 **KL 발산 (Kullback-Leibler Divergence)**을 손실 함수로 사용합니다.
  • 이론적으로 KL 발산이 0 이 되면 수용률도 1 이 되지만, 이는 **전역 최적점 (Global Optimum)**에서만 성립합니다.
  • 실제 가상 모델은 타겟 모델보다 파라미터 수가 훨씬 적어 (1~5% 수준) 용량 제약 (Capacity Constraints) 으로 인해 전역 최적점에 도달하지 못하고 **국소 최적점 (Suboptimal Solutions)**에 머무릅니다.
  • 이 국소 최적점에서는 KL 발산을 최소화하는 것이 수용률을 최대화하는 것을 보장하지 않습니다. 즉, KL 은 수용률 최적화를 위한 불완전한 대리 목표 (Proxy Objective) 로 작용합니다.

2. 방법론 (Methodology)

저자들은 수용률을 직접적으로 최적화하는 새로운 손실 함수인 LK Losses를 제안합니다.

가. TV 거리 (Total Variation Distance) 의 한계와 통찰

• 수용률 ($\alpha$) 은 TV 거리와 $\alpha = 1 - TV(p, q)$ 관계가 있어, 수용률 최대화는 TV 거리 최소화와 수학적으로 동일합니다.
• 그러나 무작위 초기화된 가상 모델을 훈련할 때 TV 거리를 직접 최소화하면 **기울기 소실 (Vanishing Gradients)**과 비연속적인 손실 지형 (Non-smooth Landscape) 문제로 인해 훈련이 불안정해집니다.

나. 제안된 LK Losses 두 가지 변형

1. 하이브리드 목적 함수 (Hybrid Objective, $L^\lambda_{LK}$):
   • 개념: 훈련 초기에는 KL 발산을, 후기에는 TV 거리를 점진적으로 강조하는 **적응형 스케줄링 (Adaptive Blending)**을 사용합니다.
   • 수식: $L^\lambda_{LK} = \lambda \cdot KL(p\|q) + (1-\lambda) \cdot TV(p, q)$
   • 동작: 수용률 ($\alpha$) 이 낮을 때 (훈련 초기) 는 $\lambda \approx 1$로 KL 을 통해 안정적인 기울기를 제공하고, 수용률이 높아질수록 $\lambda$를 감소시켜 TV 를 통해 수용률을 직접 최적화합니다. 이는 Trust-Region 접근법과 유사합니다.
2. 가능도 기반 접근법 (Likelihood-based Approach, $L^\alpha_{LK}$):
   • 개념: 수용 확률 ($\alpha$) 의 음의 로그 가능도 (Negative Log-Marginal Likelihood) 를 직접 최소화합니다.
   • 수식: $L^\alpha_{LK} = -\log \sum \min(p(x), q(x))$
   • 특징: 기울기 분석 결과, 이 방법은 TV 최적화를 수행하되 수용률이 낮을 때 기울기를 자동 증폭 ($1/\alpha$) 하여 기울기 소실 문제를 해결합니다.

다. 어휘 축소 (Vocabulary Truncation) 처리

• EAGLE-3 등 일부 아키텍처는 계산 효율을 위해 어휘를 축소합니다. 기존 KL 기반 훈련은 축소된 어휘 밖의 토큰에 대해 KL 발산을 무한대로 만들 수 있는 문제가 있지만, LK Losses 는 수용률 계산 시 축소된 어휘 밖의 토큰 기여도가 0 이므로 자연스럽게 처리됩니다.

3. 주요 기여 (Key Contributions)

1. 직접 최적화 목표: KL 발산을 우회하여 수용률을 직접 타겟팅하는 두 가지 새로운 손실 함수 (LK Losses) 를 제안했습니다.
2. 범용성 입증: 다양한 타겟 모델 (8B~685B 파라미터) 과 4 가지 가상 모델 아키텍처 (EAGLE-3, MEDUSA, MLP, MTP) 에 걸쳐 일관된 성능 향상을 입증했습니다.
3. 실용성: 계산 오버헤드가 없으며, 기존 Speculator 훈련 프레임워크에 즉시 통합 가능합니다.
4. 오픈소스: 훈련 데이터셋과 가중치를 공개하여 재현성을 보장했습니다.

4. 실험 결과 (Results)

• 실험 설정: MT-bench (대화), HumanEval (코딩), GSM8K (수학) 등 6 가지 타겟 모델과 4 가지 아키텍처에서 평가.
• 성능 향상:
  • **평균 수용 길이 (Average Acceptance Length, $\tau$)**에서 KL 기반 훈련 대비 **8~10%**까지 향상되었습니다.
  • 저용량 모델 효과: 용량이 제한된 아키텍처 (예: MEDUSA, MLP) 일수록 LK Losses 의 개선 효과가 더 컸습니다 (약 7.8~8.3% 향상).
  • 대규모 모델: 685B 파라미터의 DeepSeek-V3 와 같은 거대 모델에서도 5.6% 의 추가 향상을 기록했습니다.
  • 적응형 스케줄링의 중요성: 고정된 가중치 ($\lambda=0.5$) 를 사용한 하이브리드 모델보다, 수용률에 따라 $\lambda$를 동적으로 조절하는 적응형 스케줄링이 훨씬 우수한 성능을 보였습니다.
• 비교: 순수 TV 손실은 훈련 초기 기울기 문제로 인해 성능이 낮았으며, LK Losses 는 이를 해결하여 KL 보다 우월한 결과를 도출했습니다.

5. 의의 및 결론 (Significance)

• 패러다임 전환: Speculative Decoding 훈련에서 "분포 정렬 (Distribution Alignment)"을 위한 KL 발산이라는 간접적인 목표에서, "수용률 (Acceptance Rate)"이라는 직접적인 목표를 최적화하는 방식으로 전환했습니다.
• 실무 적용 가능성: 구현이 간단하고 추가 계산 비용이 없으며, 기존 파이프라인에 바로 적용 가능하여 LLM 추론 속도를 높이는 데 즉각적인 효과를 기대할 수 있습니다.
• 이론적 기여: KL 과 TV 거리 간의 최적화 역학 (Gradient Dynamics) 을 분석하고, 용량 제약이 있는 모델에서 왜 KL 이 수용률 최적화를 보장하지 않는지 명확히 규명했습니다.

이 논문은 Speculative Decoding 의 핵심 병목인 수용률을 직접적으로 개선함으로써, 다양한 크기와 아키텍처의 LLM 에 대해 추론 속도를 획기적으로 높일 수 있는 실용적이고 강력한 솔루션을 제시합니다.