Yuezhou Hu Jiaxin Guo Xinyu Feng Tuo Zhao

초록

사전 추론(Speculative Decoding, SD)은 소규모 초안 모델( draft model)을 활용해 예측을 생성하고, 이를 더 큰 타겟 모델이 검증함으로써 대규모 언어 모델의 추론 속도를 가속화한다. SD의 성능은 두 모델 간의 일치도에 크게 의존하며, 이는 일반적으로 지식 전이(Knowledge Distillation, KD)를 통해 향상된다. 그러나 기존의 KD 방법은 모든 토큰에 대해 초안 모델과 타겟 모델 간의 KL 발산을 최소화하는 것을 목표로 하며, 이는 SD의 실제 목적인 토큰 수용률을 최대화하는 것과 부합하지 않는다. 결과적으로 초안 모델은 용량 제약으로 인해 타겟 모델의 지식을 충분히 흡수하지 못하는 경우가 많아 성능이 최적화되지 못한다. 이 문제를 해결하기 위해, 본 연구는 KD 과정에 선택적 토큰 필터링을 도입하는 새로운 방법인 AdaSPEC을 제안한다. AdaSPEC은 참조 모델을 활용해 학습이 어려운(-fit이 어려운) 토큰을 식별하고 제거함으로써, 보다 간단한 토큰에 대해 타겟 모델과 더 잘 일치하는 초안 모델을 전이할 수 있도록 한다. 이 방식은 생성 품질을 훼손하지 않으면서도 전반적인 토큰 수용률을 향상시킨다. 우리는 31M/1.4B 및 350M/2.7B 파라미터 규모의 모델 구성에서 산술 추론, 지시어 따르기, 코딩, 요약 등 다양한 작업에 대해 AdaSPEC을 평가하였다. 실험 결과, AdaSPEC은 최신 기술인 DistillSpec을 일관되게 상회하며, 모든 작업에서 더 높은 토큰 수용률을 달성하였고(최대 15% 향상), 특히 복잡한 작업에서도 뛰어난 성능을 보였다. 코드는 공개적으로 제공되며, https://github.com/yuezhouhu/adaspec 에서 확인할 수 있다.