[논문 리뷰] DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Yongchao Zhou, Kaifeng Lyu외 6명|arXiv (Cornell University)|2023. 10. 12.

Topic Modeling인용 수 4

한 줄 요약

DistillSpec는 지식 정련을 통해 타겟 모델과 일치시키는 보조 모델의 성능을 햖थ하여 사전 추측 추론을 향상시킨다. 이는 정책 기반으로 생성된 데이터와 임무에 특화된 산란 함수를 사용한다. 이는 표준 사전 추측 추론 대비 10–45%의 속도 향상을 가져오며, 타겟 모델을 먼저 정련한 후 적용할 경우 지연 시간을 최대 6–10배까지 줄일 수 있다.

ABSTRACT

Speculative decoding (SD) accelerates large language model inference by employing a faster draft model for generating multiple tokens, which are then verified in parallel by the larger target model, resulting in the text generated according to the target model distribution. However, identifying a compact draft model that is well-aligned with the target model is challenging. To tackle this issue, we propose DistillSpec that uses knowledge distillation to better align the draft model with the target model, before applying SD. DistillSpec makes two key design choices, which we demonstrate via systematic study to be crucial to improving the draft and target alignment: utilizing on-policy data generation from the draft model, and tailoring the divergence function to the task and decoding strategy. Notably, DistillSpec yields impressive 10 - 45% speedups over standard SD on a range of standard benchmarks, using both greedy and non-greedy sampling. Furthermore, we combine DistillSpec with lossy SD to achieve fine-grained control over the latency vs. task performance trade-off. Finally, in practical scenarios with models of varying sizes, first using distillation to boost the performance of the target model and then applying DistillSpec to train a well-aligned draft model can reduce decoding latency by 6-10x with minimal performance drop, compared to standard decoding without distillation.

연구 동기 및 목표

• 사전 추측 추론에서 높은 수락률을 달성하기 위해 보조 모델과 타겟 언어 모델 간의 일치도를 향상시키는 데 도전한다.
• 표준 지식 정련의 한계를 극복하여, 사전 추측 추론을 위한 분포 일치를 최적화하고, 하류 작업 성능을 최적화하지는 않는다.
• 보조 모델 품질과 타겟 모델과의 호환성에 영향을 주는 설계 선택 사항—데이터 생성 전략과 산란 함수 선택—을 체계적으로 조사한다.
• 손실이 있는 사전 추측 추론과 결합함으로써 지연 시간 대 성능의 트레이드오���을 세밀하게 제어할 수 있도록 한다.
• 먼저 큰 모델을 더 작은 타겟 모델로 정련한 후, 높은 호환성을 갖는 보조 모델을 훈련하기 위해 DistillSpec를 적용하는 실용적인 추론 파이프라인을 개발한다.

제안 방법

• 학생 모델이 자신의 타당한 완성 부분에서 학습할 수 있도록, 보조 모델의 정책 기반으로 생성된 데이터를 사용하여 정련 목표를 훈련한다.
• 해당 추론 전략(예: 근사적 vs. 비근사적)과 임무에 따라 산란 함수(예: KL 산란, 교차 엔트로피)를 조정하여 타겟 모델의 출력 분포와의 일치도를 향상시킨다.
• 지식 정련을 적용하여 타겟 모델의 토큰 수준 및 시퀀스 수준 분포를 정밀하게 근사하는 작고 컴act한 보조 모델을 훈련한다. 이는 독립적인 작업 성능가 감소를 감수할 수 있다.
• 손실이 있는 사전 추측 추론과 DistillSpec를 통합하여 타겟 모델 분포에서의 허용 가능한 이탈을 허용함으로써 지연 시간-성능 트레이드오프를 더 잘 제어할 수 있다.
• 두 단계 파이프라인을 제안한다: 먼저 큰 모델을 더 작고 고성능의 타겟 모델로 정련한 후, 사전 추측 추론을 위한 잘 일치하는 보조 모델을 훈련하기 위해 DistillSpec를 적용한다.

실험 결과

연구 질문

• RQ1 보조 모델의 정책 기반으로 생성된 데이터를 사용할 경우, 사전 추측 추론에서 보조 모델과 타겟 모델 간의 일치도에 어떤 영향을 미치는가?
• RQ2 사전 추측 추론에서 지식 정련에 가장 적합한 산란 함수는 무엇이며, 이는 추론 전략(근사적 대비 비근사적)에 따라 어떻게 달라지는가?
• RQ3 DistillSpec는 다양한 NLP 작업과 모델 크기에서 사전 추측 추론의 수락률과 블록 효율성을 향상시킬 수 있는가?
• RQ4 DistillSpec를 손실이 있는 사전 추측 추론과 결합했을 경우, 표준 사전 추측 추론 대비 지연 시간-성능 트레이드오프에 어떤 영향을 미치는가?
• RQ5 여러 크기의 모델이 존재할 경우 가장 효과적인 추론 파이프라인은 무엇인가? 사전 추측 추론 이전에 정련을 적용할 것인가, 이후에 적용할 것인가?

주요 결과

• DistillSpec는 다양한 벤치마크, 특히 GSM8K, XSum, CNNDM에서 근사적 및 비근사적 추론 모두에서 표준 사전 추측 추론 대비 10–45%의 속도 향상을 달성한다.
• 정련 과정에서 정책 기반으로 생성된 데이터를 사용할 경우 보조-타겟 모델 간의 일치도가 크게 향상되어 수락률과 블록 효율성이 높아진다.
• 정련에 가장 적합한 산란 함수는 추론 전략에 크게 의존한다: 근사적 추론에서는 KL 산란이 가장 우수하고, 비근사적 샘플링에서는 교차 엔트로피가 더 효과적이다.
• DistillSpec를 손실이 있는 사전 추측 추론과 결합하면 더 평탄하고 고성능인 지연 시간-성능 트레이드오프 곡선을 얻을 수 있어, 성능 저하를 최소화하면서도 더 큰 지연 시간 감소 효과를 얻을 수 있다.
• 타겟 모델을 먼저 정련하고, 그 다음에 DistillSpec를 적용하는 순서로 적용할 경우, 표준 추론 대비 디코딩 지연 시간을 6–10배까지 줄일 수 있으며, 성능 저하도 최소한이다.
• GSM8K에서 비근사적 샘플링과 최적의 초모수를 사용할 경우, DistillSpec는 표준 사전 추측 추론 대비 Rouge2 점수를 15% 향상시키고, 블록 효율성을 25% 향상시킨다.