効率的な訓練不要の埋め込み空間プロービングによるマルチトークン予測

要旨

大規模言語モデル（LLM）は、次トークン生成のみで学習されているにもかかわらず、潜在的なマルチトークン予測（MTP）能力を示す。本論文では、モデルの重みを変更したり補助的なドラフトモデルに依存したりすることなく、将来のトークンを並列予測可能にする、訓練不要の簡易なMTP手法を提案する。この手法では、埋め込み空間から動的にマスクトークンを抽出してLLMをプローブし、マスクトークンのロジットから上位K個の候補をサンプリングして投機的トークンツリーを構築する。さらに、軽量な枝刈り戦略を適用し、高確率の継続トークンを保持する。デコーディング時には候補予測を並列検証することで、モデル呼び出し回数を大幅に削減しつつロスレス生成を実現し、トークン処理量を向上させる。ベンチマーク評価では、提案するプローブベースMTPが既存の訓練不要ベースラインを一貫して上回り、LLaMA3では受入長を約12%、Qwen3では8～12%向上させ、最大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