Multi-Token Voorspelling Vereist Registers

Multi-Token Prediction Needs Registers

May 15, 2025

Auteurs: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Samenvatting

Multi-token predictie is naar voren gekomen als een veelbelovend doel voor het verbeteren van het vooraf trainen van taalmodelen, maar de voordelen ervan zijn niet consistent gegeneraliseerd naar andere instellingen, zoals fine-tuning. In dit artikel stellen we MuToR voor, een eenvoudige en effectieve benadering van multi-token predictie die leerbare registertokens in de invoerreeks verweeft, elk belast met het voorspellen van toekomstige doelen. In vergelijking met bestaande methoden biedt MuToR verschillende belangrijke voordelen: het introduceert slechts een verwaarloosbaar aantal extra parameters, vereist geen architectuurwijzigingen—wat compatibiliteit met kant-en-klare vooraf getrainde taalmodelen garandeert—en blijft afgestemd op het next-token pretrainingsdoel, waardoor het bijzonder geschikt is voor supervised fine-tuning. Bovendien ondersteunt het op natuurlijke wijze schaalbare voorspellingshorizons. We demonstreren de effectiviteit en veelzijdigheid van MuToR in een reeks use cases, waaronder supervised fine-tuning, parameter-efficiënt fine-tuning (PEFT), en pretraining, op uitdagende generatieve taken in zowel taal- als visiedomeinen. Onze code zal beschikbaar zijn op: 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.