La Predicción Multi-Token Requiere Registros

Multi-Token Prediction Needs Registers

May 15, 2025

Autores: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Resumen

La predicción de múltiples tokens ha surgido como un objetivo prometedor para mejorar el preentrenamiento de modelos de lenguaje, pero sus beneficios no se han generalizado de manera consistente a otros escenarios, como el ajuste fino. En este artículo, proponemos MuToR, un enfoque simple y efectivo para la predicción de múltiples tokens que intercala tokens de registro aprendibles en la secuencia de entrada, cada uno encargado de predecir objetivos futuros. En comparación con los métodos existentes, MuToR ofrece varias ventajas clave: introduce solo un número insignificante de parámetros adicionales, no requiere cambios arquitectónicos—asegurando compatibilidad con modelos de lenguaje preentrenados disponibles—y se mantiene alineado con el objetivo de preentrenamiento de siguiente token, lo que lo hace especialmente adecuado para el ajuste fino supervisado. Además, admite de manera natural horizontes de predicción escalables. Demostramos la efectividad y versatilidad de MuToR en una variedad de casos de uso, incluyendo ajuste fino supervisado, ajuste fino eficiente en parámetros (PEFT) y preentrenamiento, en tareas generativas desafiantes tanto en dominios de lenguaje como de visión. Nuestro código estará disponible en: 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.