임베딩 공간 탐색을 통한 효율적인 학습 없이 다중 토큰 예측

초록

대규모 언어 모델(LLM)은 다음 토큰 생성만을 위해 훈련되었음에도 잠재적인 다중 토큰 예측(MTP) 능력을 보여준다. 우리는 모델 가중치 수정이나 보조 드래프트 모델에 의존하지 않고 미래 토큰을 병렬 예측할 수 있도록, 임베딩 공간에서 추출한 마스크 토큰을 활용하여 LLM을 탐색하는 간단한 훈련 불필요(training-free) MTP 접근법을 제안한다. 우리의 방법은 마스크 토큰 로짓(logits)에서 상위 K개 후보를 샘플링하여 추론적 토큰 트리(speculative token tree)를 구성하고, 가벼운 가지치기 전략을 적용하여 높은 확률의 후속 토큰을 보존한다. 디코딩 과정에서 후보 예측들은 병렬로 검증되어, 모델 호출 횟수를 상당히 줄이고 토큰 처리량을 향상시키면서 무손실(lossless) 생성을 달성한다. 다양한 벤치마크에서 우리의 탐색(probing) 기반 MTP는 기존의 훈련 불필요 베이스라인을 지속적으로 능가하며, LLaMA3에서 약 12%, Qwen3에서 8-12%의 수용 길이(acceptance length) 증가와 최대 15-19%의 처리량 향상을 보였다. 마지막으로, 디코더 레이어가 마스크 토큰 표현을 다음 토큰 상태와 자연스럽게 정렬시켜 재훈련이나 보조 모델 없이도 정확한 다단계 예측을 가능하게 한다는 이론적 통찰과 실증적 증거를 제시한다.

English

Large language models (LLMs) exhibit latent multi-token prediction (MTP) capabilities despite being trained solely for next-token generation. We propose a simple, training-free MTP approach that probes an LLM using on-the-fly mask tokens drawn from its embedding space, enabling parallel prediction of future tokens without modifying model weights or relying on auxiliary draft models. Our method constructs a speculative token tree by sampling top-K candidates from mask-token logits and applies a lightweight pruning strategy to retain high-probability continuations. During decoding, candidate predictions are verified in parallel, resulting in lossless generation while substantially reducing the number of model calls and improving token throughput. Across benchmarks, our probing-based MTP consistently outperforms existing training-free baselines, increasing acceptance length by approximately 12\% on LLaMA3 and 8--12\% on Qwen3, and achieving throughput gains of up to 15--19\%. Finally, we provide theoretical insights and empirical evidence showing that decoder layers naturally align mask-token representations with next-token states, enabling accurate multi-step prediction without retraining or auxiliary models.

임베딩 공간 탐색을 통한 효율적인 학습 없이 다중 토큰 예측

Efficient Training-Free Multi-Token Prediction via Embedding-Space Probing

초록

Support