A Previsão de Múltiplos Tokens Requer Registros

Multi-Token Prediction Needs Registers

May 15, 2025

Autores: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Resumo

A previsão de múltiplos tokens surgiu como um objetivo promissor para melhorar o pré-treinamento de modelos de linguagem, mas seus benefícios não se generalizaram de forma consistente para outros cenários, como o ajuste fino. Neste artigo, propomos o MuToR, uma abordagem simples e eficaz para a previsão de múltiplos tokens que intercala tokens de registro aprendíveis na sequência de entrada, cada um responsável por prever alvos futuros. Em comparação com métodos existentes, o MuToR oferece várias vantagens-chave: introduz apenas um número insignificante de parâmetros adicionais, não requer alterações arquitetônicas—garantindo compatibilidade com modelos de linguagem pré-treinados prontos para uso—e permanece alinhado com o objetivo de pré-treinamento de próximo token, tornando-o especialmente adequado para ajuste fino supervisionado. Além disso, ele suporta naturalmente horizontes de previsão escaláveis. Demonstramos a eficácia e versatilidade do MuToR em uma variedade de casos de uso, incluindo ajuste fino supervisionado, ajuste fino com eficiência de parâmetros (PEFT) e pré-treinamento, em tarefas generativas desafiadoras tanto no domínio da linguagem quanto da visão. Nosso código estará disponível em: https://github.com/nasosger/MuToR.

English

Multi-token prediction has emerged as a promising objective for improving language model pretraining, but its benefits have not consistently generalized to other settings such as fine-tuning. In this paper, we propose MuToR, a simple and effective approach to multi-token prediction that interleaves learnable register tokens into the input sequence, each tasked with predicting future targets. Compared to existing methods, MuToR offers several key advantages: it introduces only a negligible number of additional parameters, requires no architectural changes--ensuring compatibility with off-the-shelf pretrained language models--and remains aligned with the next-token pretraining objective, making it especially well-suited for supervised fine-tuning. Moreover, it naturally supports scalable prediction horizons. We demonstrate the effectiveness and versatility of MuToR across a range of use cases, including supervised fine-tuning, parameter-efficient fine-tuning (PEFT), and pretraining, on challenging generative tasks in both language and vision domains. Our code will be available at: https://github.com/nasosger/MuToR.