小模型训练时用渐进式策略提升多词预测效果
Pre-Training Curriculum for Multi-Token Prediction in Language Models
- 从单词预测逐步过渡到多词预测，降低小模型学习难度
- 前向课程使小模型在下游任务中表现更好，生成更流畅
- 适合希望提升小模型生成质量但不追求快速推理的研究者
多词预测（MTP）是语言模型预训练的新目标，不再仅预测下一个词，而是每步预测接下来的 $k$ 个词，使用多个预测头。虽然大模型在该目标下表现出色，提升了下游性能、推理速度和训练效率，但小模型（SLMs）却难以适应。为此，我们提出一种针对MTP的课程学习策略，包含两种变体：前向课程（从NTP逐步增加至MTP）和反向课程（相反过程）。实验表明，前向课程帮助小模型更好地利用MTP，在下游单词预测任务中取得更优表现，并提升生成质量，同时保持自推测解码的优势；而反向课程虽在单词预测和生成上更强，但无法带来推测解码收益。
原文摘要 · Abstract (English)
Multi-token prediction (MTP) is a recently proposed pre-training objective for language models. Rather than predicting only the next token (NTP), MTP predicts the next $k$ tokens at each prediction step, using multiple prediction heads. MTP has shown promise in improving downstream performance, inference speed, and training efficiency, particularly for large models. However, prior work has shown that smaller language models (SLMs) struggle with the MTP objective. To address this, we propose a curriculum learning strategy for MTP training, exploring two variants: a forward curriculum, which gradually increases the complexity of the pre-training objective from NTP to MTP, and a reverse curriculum, which does the opposite. Our experiments show that the forward curriculum enables SLMs to better leverage the MTP objective during pre-training, improving downstream NTP performance and generative output quality, while retaining the benefits of self-speculative decoding. The reverse curriculum achieves stronger NTP performance and output quality, but fails to provide any self-speculative decoding benefits.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。