La prédiction multi-jetons nécessite des registres.

Multi-Token Prediction Needs Registers

May 15, 2025

Auteurs: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Résumé

La prédiction multi-tokens s'est imposée comme un objectif prometteur pour améliorer le pré-entraînement des modèles de langage, mais ses avantages ne se sont pas systématiquement généralisés à d'autres contextes, tels que le réglage fin. Dans cet article, nous proposons MuToR, une approche simple et efficace pour la prédiction multi-tokens qui entrelace des tokens de registre apprenables dans la séquence d'entrée, chacun étant chargé de prédire des cibles futures. Par rapport aux méthodes existantes, MuToR offre plusieurs avantages clés : il introduit un nombre négligeable de paramètres supplémentaires, ne nécessite aucune modification architecturale—garantissant ainsi la compatibilité avec les modèles de langage pré-entraînés disponibles—et reste aligné avec l'objectif de prédiction du token suivant, ce qui le rend particulièrement adapté au réglage fin supervisé. De plus, il supporte naturellement des horizons de prédiction évolutifs. Nous démontrons l'efficacité et la polyvalence de MuToR à travers une gamme de cas d'utilisation, incluant le réglage fin supervisé, le réglage fin efficace en paramètres (PEFT), et le pré-entraînement, sur des tâches génératives complexes dans les domaines du langage et de la vision. Notre code sera disponible à l'adresse : 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.