Многотокенное предсказание требует использования регистров.

Multi-Token Prediction Needs Registers

May 15, 2025

Авторы: Anastasios Gerontopoulos, Spyros Gidaris, Nikos Komodakis

cs.AI

Аннотация

Прогнозирование нескольких токенов стало перспективным подходом для улучшения предварительного обучения языковых моделей, однако его преимущества не всегда распространяются на другие задачи, такие как тонкая настройка. В данной статье мы предлагаем MuToR — простой и эффективный метод прогнозирования нескольких токенов, который встраивает обучаемые регистровые токены во входную последовательность, каждый из которых отвечает за предсказание будущих целей. По сравнению с существующими методами, MuToR обладает несколькими ключевыми преимуществами: он вводит лишь незначительное количество дополнительных параметров, не требует изменений архитектуры — обеспечивая совместимость с готовыми предобученными языковыми моделями — и остается согласованным с целью предсказания следующего токена, что делает его особенно подходящим для контролируемой тонкой настройки. Кроме того, он естественным образом поддерживает масштабируемые горизонты прогнозирования. Мы демонстрируем эффективность и универсальность MuToR в различных сценариях использования, включая контролируемую тонкую настройку, параметрически эффективную тонкую настройку (PEFT) и предварительное обучение, на сложных генеративных задачах в областях языка и зрения. Наш код будет доступен по адресу: 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.