캥거루: 이중 조기 종료를 통한 무손실 자기 추론 디코딩

초록

스펙큘레이티브 디코딩(speculative decoding)은 대형 언어 모델의 추론 속도를 가속화하면서도 일관된 샘플링 분포를 유지하는 데 효과적임이 입증되었습니다. 그러나 기존의 별도 드래프트 모델을 훈련시켜 만족스러운 토큰 수용률을 달성하는 방식은 비용이 많이 들 수 있습니다. 이를 해결하기 위해, 우리는 초기 종료(early exiting)에서 영감을 받아 고정된 얕은 서브 네트워크를 자체 드래프트 모델로 사용하고, 나머지 레이어를 더 큰 타겟 모델로 활용하는 새로운 자기 스펙큘레이티브 디코딩 프레임워크인 Kangaroo를 제안합니다. 우리는 서브 네트워크와 전체 모델의 표현 능력 간의 격차를 줄이기 위해 서브 네트워크 위에 가볍고 효율적인 어댑터 모듈을 훈련시킵니다. 특히, 자체 드래프트 모델의 추론 지연 시간이 더 이상 대형 모델에 비해 무시할 수 없을 정도로 커질 수 있으므로, 작은 모델의 드래프팅 단계를 최소화하면서 토큰 수용률을 높이는 전략이 필요합니다. 이 문제를 해결하기 위해, 우리는 드래프트 토큰을 생성하기 위한 추가적인 초기 종료 메커니즘을 도입합니다. 구체적으로, 드래프팅 단계에서 현재 토큰에 대한 신뢰도가 특정 임계값 아래로 떨어지면 작은 모델의 후속 예측을 중단합니다. Spec-Bench에서의 광범위한 실험을 통해 Kangaroo의 효과를 입증했습니다. 단일 시퀀스 검증(single-sequence verification) 하에서 Kangaroo는 Spec-Bench에서 최대 1.68배의 속도 향상을 달성하며, Medusa-1 대비 88.7% 더 적은 추가 파라미터(67M 대 591M)를 사용합니다. Kangaroo의 코드는 https://github.com/Equationliu/Kangaroo에서 확인할 수 있습니다.

English

Speculative decoding has demonstrated its effectiveness in accelerating the inference of large language models while maintaining a consistent sampling distribution. However, the conventional approach of training a separate draft model to achieve a satisfactory token acceptance rate can be costly. Drawing inspiration from early exiting, we propose a novel self-speculative decoding framework Kangaroo, which uses a fixed shallow sub-network as a self-draft model, with the remaining layers serving as the larger target model. We train a lightweight and efficient adapter module on top of the sub-network to bridge the gap between the sub-network and the full model's representation ability. It is noteworthy that the inference latency of the self-draft model may no longer be negligible compared to the large model, necessitating strategies to increase the token acceptance rate while minimizing the drafting steps of the small model. To address this challenge, we introduce an additional early exiting mechanism for generating draft tokens. Specifically, we halt the small model's subsequent prediction during the drafting phase once the confidence level for the current token falls below a certain threshold. Extensive experiments on the Spec-Bench demonstrate the effectiveness of Kangaroo. Under single-sequence verification, Kangaroo achieves speedups up to 1.68times on Spec-Bench, outperforming Medusa-1 with 88.7\% fewer additional parameters (67M compared to 591M). The code for Kangaroo is available at https://github.com/Equationliu/Kangaroo.

캥거루: 이중 조기 종료를 통한 무손실 자기 추론 디코딩

Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

초록

Support