직접 기대 디코딩 라운드를 최소화하여 병렬 추측 초안 모델 훈련하기

추측 디코딩(Speculative decoding)은 저비용의 초안 모델(draft model)을 사용하여 전체 크기의 타겟 모델(target model)이 병렬로 검증하는 토큰을 제안함으로써 대규모 언어 모델(LLM) 추론 속도를 가속화합니다. 병렬 및 준-자동회귀식(semi-AR) 초안기(drafter)는 단일 순방향 패스(forward pass)에서 전체 블록을 제안하여 초안 효율성을 개선하지만, 이를 훈련하는 과정에서 새로운 어려움이 발생합니다. 즉, 주어진 위치에 대한 초안 분포가 디코딩 라운드가 어디서 시작하느냐에 따라 달라지며, 라운드 시작 지점은 앞선 라운드들이 몇 개의 토큰을 수락했는지에 따라 달라집니다. 기존의 훈련 목표는 이러한 교차 라운드 결합(cross-round coupling)을 무시하는 블록 로컬 서로게이트(block-local surrogates)에 의존하며, 따라서 전역 디코딩 효율성을 직접적으로 최적화하지 못합니다. 본 연구에서는 추측 디코딩을 마르코프 보상 과정(Markov reward process)으로 표현하여 이러한 초안기들을 훈련하고 평가하기 위한 이론적 프레임워크를 개발합니다. 이 공식화는 기대 디코딩 라운드(Expected Decoding Rounds, EDR) 목표를 도출하며, 이는 상태 점유율에 따라 로컬 거부 비용을 가중치 부여하고 정확히 기대 디코딩 라운드 수와 일치합니다. 이전의 서로게이트 목표들과 달리, EDR은 보조 하이퍼파라미터를 도입하지 않습니다. 이후 우리는 타겟 모델 롤아웃(target-model rollouts)으로부터 편향되지 않은 확률적 최적화(unbiased stochastic optimization)를 지원하는 정확한 시간 차분 그래디언트(temporal-difference gradient)를 유도합니다. 동일한 프레임워크는 라운드 수에 대한 정확한 오프라인 평가기(offline evaluator)를 제공하여, 추측 디코딩을 실행하지 않고 공유된 타겟 롤아웃 상에서 쌍을 이루는 초안기 비교를 가능하게 합니다. EDR로 두 가지 최첨단 초안기인 DSpark와 DFly를 미세 조정(Finetuning)한 결과, 수학적 추론, 코드 생성 및 채팅에 걸친 아홉 개의 벤치마크 전반에서 평균 수락 길이가 지속적으로 개선되었으며 기존의 훈련 목표들을 능가하는 성능을 보였습니다.