VOCABTRIM: 대규모 언어 모델의 효율적 스펙츌레이티브 디코딩을 위한 어휘 가지치기

초록

본 논문에서는 드래프트 기반 스펙큘레이티브 디코딩(SpD) 방법의 성능을 향상시키기 위한 간단한 학습 불필요 기법을 소개한다. 이 기법은 드래프팅 과정에서 언어 모델링 헤드(LM 헤드)를 통합한다. 드래프트 기반 스펙큘레이티브 디코딩은 하나 이상의 더 작은 언어 모델, 즉 드래프터 또는 드래프트 모델을 활용하여 여러 토큰으로 구성된 드래프트 시퀀스 또는 트리를 샘플링한 후, 기본 LLM인 타겟 모델이 이를 검증하여 일부를 유효한 생성으로 수락한다. 일반적으로 스펙큘레이티브 디코딩은 타겟 모델과 드래프트 모델의 어휘 간 일대일 매핑이 필요하다고 여겨지기 때문에, 이 둘 간에 어휘를 공유하거나 EAGLE이나 Medusa에서와 같이 LM 헤드를 공유하는 것이 자연스러운 접근이었다. 우리는 먼저 이 드래프트 토큰 샘플링 방식이 특히 매우 큰 어휘를 가진 일부 타겟 LLM의 경우, 드래프팅 과정에서 불필요한 추론 오버헤드를 포함하고 있음을 확인했다. 그런 다음, 메모리 제약 환경에서 생성 속도를 향상시키기 위해 드래프팅 오버헤드를 완화하는 간단한 기법인 VocabTrim을 제안한다. VocabTrim은 드래프터 LM 헤드를 재구성하여 타겟 모델의 어휘에서 가장 빈번하게 샘플링된 토큰들로만 제한된 집합을 포함하도록 한다. 드래프팅에서 어휘를 제한하는 것은 수락률을 약간 저하시키지만, 에지 디바이스에서 흔히 발생하는 메모리 제약 프로세스에서 드래프팅 지연을 크게 줄여 더 높은 메모리 제약 속도 향상(MBSU)을 가져온다. 우리는 이 방법이 Spec-Bench에서 Llama-3 모델의 메모리 제약 속도 향상을 촉진할 수 있음을 보여주며, 특히 Llama-3.2-3B-Instruct의 경우 16%의 속도 향상을 달성했다.

English

In this paper, we introduce a simple training-free technique to improve the performance of drafter-based speculative decoding (SpD) methods that incorporates language modeling head (LM head) during drafting process. A drafter-based speculative decoding leverages one or more smaller language models, a.k.a. drafters or draft models, to sample a draft sequence or tree consisting of multiple tokens, followed by verification by a base LLM, a target model, accepting a subset as its valid generation. As it is usually considered that the speculative decoding requires one-to-one mapping between vocabularies of the target model and the draft model, it has been natural to share the vocabulary between them, or even share the LM head as in EAGLE or Medusa. We first identify that this draft token sampling scheme inherently contains an unnecessary inference overhead in drafting, especially for some target LLMs with very large vocabularies. Then, we propose a simple technique, VocabTrim, to mitigate the drafting overhead to improve the generation speed in memory-bound environment. VocabTrim reconstructs the drafter LM head to contain only a limited set of tokens, selected by the most frequently sampled from the vocabulary of the target model. While limiting the vocabulary in drafting slightly degrades the acceptance rate, it significantly reduces the drafting latency in memory-bound process which is often the case on edge devices, resulting in higher memory-bound speed up (MBSU). We show that our method can boost the memory-bound speed-up for Llama-3 models on Spec-Bench, specifically by 16% for Llama-3.2-3B-Instruct.

VOCABTRIM: 대규모 언어 모델의 효율적 스펙츌레이티브 디코딩을 위한 어휘 가지치기

VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs

초록

Support