기술 요약: Draft-OPD

문제 정의

추측적 디코딩 (Speculative Decoding, SD) 은 경량 드래프트 모델을 사용하여 토큰 블록을 제안하고, 이를 더 큰 타겟 모델이 병렬로 검증함으로써 대규모 언어 모델 (LLM) 추론을 가속화합니다. 현재 주류 방법들 (예: EAGLE-3, DFlash) 은 타겟 모델이 생성한 고정된 궤적에 대한 지도 미세 조정 (Supervised Fine-Tuning, SFT) 을 통해 이러한 드래프트 모델을 훈련시키지만, 이 접근법은 성능의 한계에 직면합니다.

저자들은 근본적인 오프라인 - 추론 불일치 (offline-to-inference mismatch) 를 식별했습니다:

1. 훈련 단계: 드래프트 모델은 타겟 모델이 생성한 모든 접두사를 포함하는 정적 궤적에서 학습합니다.
2. 추론 단계: 드래프트 모델은 자신이 제안한 블록에 대해 평가됩니다. 수용 길이 (검증 단계당 수용된 토큰 수) 는 SFT 동안 관찰된 정적 상태가 아닌, 드래프트 모델 자체의 정책 유도 상태에 의해 결정됩니다.

결과적으로 지속적인 SFT 는 빠르게 수렴하며, 드래프트 모델은 테스트 데이터에서 수용 길이를 개선하지 못합니다. 또한, 표준 온 - 정책 증류 (On-Policy Distillation, OPD) 는 드래프트 모델에 직접 적용하기 어렵습니다.

• 드래프트 전용 롤아웃 (Draft-only rollouts) (드래프트 모델이 독립적으로 시퀀스를 생성하는 경우) 은 불안정하며 종종 반복적인 출력으로 퇴화합니다.
• 타겟 지원 롤아웃 (Target-assisted rollouts) (타겟 모델이 드래프트를 수정하는 경우) 은 고품질 시퀀스를 생성하지만 "온 - 정책" 신호를 제거합니다. 타겟 모델이 드래프트의 오류를 폐기하기 때문에, 드래프트 모델은 타겟의 분포로부터만 학습하여 추측적 수용을 제한하는 특정 오류를 해결하지 못합니다.

방법론: Draft-OPD

이 격차를 해소하기 위해 저자들은 훈련 기반 드래프트 모델을 위해 설계된 Draft-OPD, 즉 온 - 정책 증류 프레임워크를 제안합니다. 이 방법은 드래프트 모델이 성공적인 제안뿐만 아니라 검증 실패를 초래하는 특정 오류를 포함한 추론 시 상태를 학습할 수 있도록 합니다.

이 프레임워크는 세 가지 조율된 구성 요소로 이루어집니다:

1. 오류 위치 수집을 통한 롤아웃

드래프트 모델이 전체 시퀀스를 단독으로 생성하도록 강요하는 대신, Draft-OPD 는 안정적이고 고품질의 연속성을 보장하기 위해 타겟 지원 롤아웃을 사용합니다.

• 추측적 디코딩 동안 드래프트 모델은 토큰 블록을 제안하고 타겟 모델이 이를 검증합니다.
• 시스템은 각 드래프트된 블록의 시작 위치를 앵커 (anchor) 로 기록합니다.
• 이 과정은 타겟 분포를 따르는 유효한 롤아웃을 생성하면서도 드래프트 모델이 행동을 취한 특정 상태를 보존합니다.

2. 로그 확률 계산을 위한 재생

타겟 지원 롤아웃 동안 손실된 "온 - 정책" 신호를 복원하기 위해 시스템은 기록된 앵커에서 드래프트 과정을 재생 (replay) 합니다.

• 앵커 상태 (타겟에 의해 검증된 접두사) 에서 시작하여 시스템은 후속 블록에 대한 드래프트 모델의 제안을 다시 시뮬레이션합니다.
• 이 드래프트 생성 접두사에 대해 드래프트 모델 (학생) 과 타겟 모델 (교사) 의 로그 확률을 계산합니다.
• 핵심적으로, 이는 초기 검증 동안 거부된 토큰에 대한 타겟 모델의 분포를 평가할 수 있게 하여, 드래프트 모델이 저지른 특정 오류를 포착합니다.

3. 수용 인지 증류 목적 함수

훈련 목적 함수는 추측적 수용을 최적화하기 위해 수용된 토큰과 거부된 토큰을 다르게 처리합니다:

• 수용된 토큰: 검증을 통과한 토큰의 경우, 이 방법은 Forward KL 발산 ($D_{KL}(P_{target} || P_{draft})$) 을 사용합니다. 이는 이미 일치하는 영역에서 타겟 분포를 드래프트 분포가 커버하도록 장려하여 신뢰할 수 있는 행동을 강화합니다.
• 거부된 토큰: 검증을 실패한 토큰의 경우, 이 방법은 Reverse KL 발산 ($D_{KL}(P_{draft} || P_{target})$) 을 사용합니다. 이는 타겟 모델이 거부하는 토큰에 높은 확률을 할당한 드래프트 모델을 처벌하여, 드래프트 모델이 자신의 고신뢰 오류를 피하도록 효과적으로 가르칩니다.
• 가중치: 거부된 토큰은 블록 내 위치 ($w_k = \gamma^{k-1}$) 에 따라 지수적으로 가중치가 부여됩니다. 블록 초기에 발생하는 오류는 전체 후속 접미사를 무효화하므로 전체 수용 길이에 더 큰 영향을 미치기 때문에 더 높은 가중치를 받습니다.

주요 기여

1. 한계 식별: 이 논문은 훈련 상태 (타겟 생성) 와 추론 상태 (드래프트 생성) 간의 불일치로 인해 드래프트 모델에 대한 오프라인 SFT 가 빠르게 수렴함을 입증합니다.
2. OPD 도전 과제 분석: 표준 OPD 가 드래프트 모델에 실패하는 이유를 설명합니다. 독립적인 롤아웃은 불안정하며, 단순한 타겟 지원 롤아웃은 온 - 정책 오류 신호를 지워버립니다.
3. Draft-OPD 프레임워크: 타겟 지원 롤아웃과 오류 위치 재생을 결합한 새로운 프레임워크를 도입하여 검증 시 오류에 대한 효율적인 사후 훈련을 가능하게 합니다.
4. 수용 인지 목적 함수: 추측적 가속화를 제한하는 오류를 구체적으로 대상으로 하기 위해 수용된 토큰과 거부된 토큰을 구분하는 하이브리드 KL 발산 목적 함수를 제안합니다.

실험 결과

저자들은 수학 추론 (GSM8K, MATH-500, AIME), 코드 생성 (MBPP, HumanEval, SWE-bench Lite), 일반 채팅 (MT-Bench) 전반에 걸쳐 Qwen3 계열 (4B, 8B, 30B-A3B-Thinking) 에서 Draft-OPD 를 평가했습니다.

• 성능 향상: 일치된 훈련 FLOPs 예산 하에서 Draft-OPD 는 EAGLE-3 및 DFlash 보다 일관되게 우수한 성능을 보였습니다.
  • 사고 모델 (thinking models) 에 대해 손실 없는 가속화 (lossless acceleration) 5 배 이상을 달성했습니다.
  • 평균 수용 길이와 가속 속도를 EAGLE-3 대비 23%, DFlash 대비 13% 개선했습니다.
• 처리량: SGLang 배포 벤치마크에서 Draft-OPD 는 동시성 수준 (1~32) 전반에 걸쳐 일관된 처리량 향상을 보였으며, DFlash 대비 평균 수용 길이는 약 11.2% 개선되었습니다.
• 애블레이션 연구:
  • 재생 메커니즘 없이 OPD 생성 데이터에 단순히 SFT 를 적용하는 것은 성능 향상을 이루지 못해 온 - 정책 오류 신호의 필요성을 확인했습니다.
  • 단일 KL 방향 (모든 Forward 또는 모든 Reverse) 은 혼합된 수용 인지 목적 함수보다 열등했습니다.
  • 오류 위치에 기반하여 앵커를 선택하는 것보다 무작위로 재생 앵커를 선택하는 것이 덜 효과적이었으며, 이는 실패 상태에 초점을 맞춘 훈련의 중요성을 강조합니다.

중요성과 주장

이 논문은 Draft-OPD가 훈련 기반 드래프트 모델 최적화에서 중요한 진전을 제공한다고 주장합니다. 정적 오프라인 감독에서 검증 시 오류를 명시적으로 대상으로 하는 동적 온 - 정책 증류로 전환함으로써, 이 방법은 전통적 SFT 의 성능 수렴을 돌파합니다.

저자들은 이 접근법이 타겟 모델의 출력 분포를 유지하는 추측적 디코딩의 손실 없는 (lossless) 성격을 보존하면서 효율성을 크게 향상시킨다고 강조합니다. 그들은 이 작업을 표준 드래프트 모델이 긴 추론 경로의 높은 불확실성으로 인해 종종 어려움을 겪는 "사고" 모델 및 복잡한 추론 작업을 가속화하기 위한 실용적인 해결책으로 위치시킵니다. 논문은 추론 시 자신의 오류에 대해 드래프트 모델을 사후 훈련하는 것이 추측적 디코딩을 개선하기 위한 중요한 방향이며, 생성 품질을 훼손하지 않고 더 높은 수용 길이와 더 빠른 추론을 위한 길을 제시한다고 결론지었습니다.