Multi-Token-Vorhersage benötigt Register

Multi-Token Prediction Needs Registers

May 15, 2025

Autoren: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Zusammenfassung

Die Vorhersage mehrerer Tokens hat sich als vielversprechendes Ziel zur Verbesserung des Vortrainings von Sprachmodellen erwiesen, doch ihre Vorteile haben sich nicht konsistent auf andere Anwendungen wie das Feinabstimmen übertragen lassen. In diesem Artikel schlagen wir MuToR vor, einen einfachen und effektiven Ansatz zur Vorhersage mehrerer Tokens, der lernbare Register-Tokens in die Eingabesequenz einfügt, von denen jeder für die Vorhersage zukünftiger Ziele verantwortlich ist. Im Vergleich zu bestehenden Methoden bietet MuToR mehrere entscheidende Vorteile: Es führt nur eine vernachlässigbare Anzahl zusätzlicher Parameter ein, erfordert keine architektonischen Änderungen – was die Kompatibilität mit vorgefertigten vortrainierten Sprachmodellen sicherstellt – und bleibt mit dem Next-Token-Vortrainingsziel abgestimmt, was es besonders gut für das überwachte Feinabstimmen geeignet macht. Darüber hinaus unterstützt es auf natürliche Weise skalierbare Vorhersagehorizonte. Wir demonstrieren die Wirksamkeit und Vielseitigkeit von MuToR in einer Reihe von Anwendungsfällen, einschließlich überwachtem Feinabstimmen, parameter-effizientem Feinabstimmen (PEFT) und Vortraining, bei anspruchsvollen generativen Aufgaben in den Bereichen Sprache und Bildverarbeitung. Unser Code wird unter folgender Adresse verfügbar sein: 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.