arXiv:2505.10518cs.CLcs.AI2025-05NeurIPS被引 12

通过插入可学习的寄存器令牌,提升语言模型多标记预测能力。

Multi-Token Prediction Needs Registers

  • 在输入序列中插入可学习的寄存器令牌,分别预测未来目标。
  • 仅增加极少参数,兼容现有预训练模型,支持长程预测。
  • 适用于微调、高效微调及预训练,在文本与视觉任务中表现优异。

多标记预测已成为提升语言模型预训练的有前景目标,但其优势未能一致推广到微调等场景。本文提出MuToR,一种简单有效的多标记预测方法:将可学习的寄存器令牌交错插入输入序列,每个寄存器负责预测未来目标。相比现有方法,MuToR仅引入极少量额外参数,无需架构改动,确保与现成预训练模型兼容,且保持与下一标记预测目标对齐,特别适合监督微调。此外,其天然支持可扩展的预测范围。我们在多种应用场景中验证了MuToR的有效性与通用性,涵盖监督微调、参数高效微调(PEFT)及预训练,在语言与视觉领域的挑战性生成任务上均表现良好。代码将开源于:https://github.com/nasosger/MuToR。

原文摘要 · Abstract (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.

多标记预测语言模型高效微调可扩展

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。