La Predizione Multi-Token Richiede Registri

Multi-Token Prediction Needs Registers

May 15, 2025

Autori: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Abstract

La previsione multi-token è emersa come un obiettivo promettente per migliorare il pre-addestramento dei modelli linguistici, ma i suoi benefici non si sono generalizzati in modo consistente ad altri contesti, come il fine-tuning. In questo articolo, proponiamo MuToR, un approccio semplice ed efficace alla previsione multi-token che intercala token di registro apprendibili nella sequenza di input, ciascuno incaricato di prevedere target futuri. Rispetto ai metodi esistenti, MuToR offre diversi vantaggi chiave: introduce solo un numero trascurabile di parametri aggiuntivi, non richiede modifiche architetturali—garantendo la compatibilità con modelli linguistici pre-addestrati pronti all'uso—e rimane allineato con l'obiettivo di pre-addestramento del token successivo, rendendolo particolarmente adatto per il fine-tuning supervisionato. Inoltre, supporta naturalmente orizzonti di previsione scalabili. Dimostriamo l'efficacia e la versatilità di MuToR in una gamma di casi d'uso, tra cui fine-tuning supervisionato, fine-tuning efficiente in termini di parametri (PEFT) e pre-addestramento, su compiti generativi impegnativi sia nel dominio linguistico che in quello visivo. Il nostro codice sarà disponibile all'indirizzo: 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.