Predição Eficiente de Múltiplos Tokens sem Treinamento por meio de Sondagem do Espaço de Embeddings

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

March 18, 2026

Autores: Raghavv Goel, Mukul Gagrani, Mingu Lee, Chris Lott

cs.AI

Resumo

Os grandes modelos de linguagem (LLMs) exibem capacidades latentes de previsão multi-token (MTP), apesar de serem treinados apenas para geração do próximo token. Propomos uma abordagem MTP simples e livre de treinamento que investiga um LLM usando tokens de máscara gerados dinamicamente a partir de seu espaço de incorporação, permitindo a previsão paralela de tokens futuros sem modificar os pesos do modelo ou depender de modelos auxiliares de rascunho. Nosso método constrói uma árvore especulativa de tokens através da amostragem dos principais K candidatos dos logits dos tokens de máscara e aplica uma estratégia leve de poda para reter continuações de alta probabilidade. Durante a decodificação, as previsões candidatas são verificadas em paralelo, resultando em geração sem perdas enquanto reduz substancialmente o número de chamadas do modelo e melhora a taxa de transferência de tokens. Em benchmarks, nosso MTP baseado em investigação supera consistentemente as linhas de base livres de treinamento existentes, aumentando o comprimento de aceitação em aproximadamente 12% no LLaMA3 e 8–12% no Qwen3, e alcançando ganhos de taxa de transferência de até 15–19%. Finalmente, fornecemos insights teóricos e evidências empíricas mostrando que as camadas do decodificador alinham naturalmente as representações dos tokens de máscara com os estados do próximo token, permitindo previsão multi-etapa precisa sem retreinamento ou modelos auxiliares.

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.