Better & Faster Large Language Models via Multi-token Prediction

📄 arXiv: 2404.19737v1 📥 PDF

作者: Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

分类: cs.CL

发布日期: 2024-04-30


💡 一句话要点

通过多标记预测提升大型语言模型的效率与性能

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 多标记预测 样本效率 自然语言处理 代码生成 推理速度 生成任务

📋 核心要点

  1. 现有大型语言模型主要依赖下一个标记预测,导致样本效率低下,尤其在处理复杂任务时表现不佳。
  2. 本文提出通过多标记预测来训练语言模型,利用多个独立输出头同时预测多个标记,从而提高训练效率和下游任务能力。
  3. 实验结果显示,使用该方法的模型在多个基准测试中显著超越传统模型,尤其在生成任务中表现出色。

📝 摘要(中文)

大型语言模型如GPT和Llama通常使用下一个标记预测损失进行训练。本文提出同时预测多个未来标记的训练方法,以提高样本效率。具体而言,在训练语料的每个位置,模型被要求使用n个独立的输出头预测接下来的n个标记,这些输出头在共享的模型主干上操作。将多标记预测视为辅助训练任务,我们在代码和自然语言模型上观察到下游能力的提升,且训练时间没有增加。该方法在模型规模增大时效果更为显著,并在多轮训练中保持吸引力。在生成基准测试中,模型的表现优于多个强基线,13B参数的模型在HumanEval和MBPP上分别解决了12%和17%更多的问题。实验还表明,多标记预测有助于归纳头和算法推理能力的发展,且使用4标记预测的模型在推理时速度提高了最多3倍。

🔬 方法详解

问题定义:现有的大型语言模型主要依赖下一个标记预测,这种方法在处理复杂任务时效率较低,且在样本利用上存在不足。

核心思路:本文提出通过多标记预测来训练语言模型,即在每个训练位置同时预测多个未来标记,利用多个独立的输出头来实现。这种设计旨在提高样本效率和下游任务的表现。

技术框架:整体架构包括一个共享的模型主干和多个独立的输出头。每个输出头负责预测一个未来标记,模型在训练过程中同时优化这些输出。

关键创新:最重要的创新在于将多标记预测作为辅助训练任务引入,显著提升了模型在生成任务中的能力,并在训练时间上没有增加额外负担。

关键设计:在参数设置上,模型使用n个独立的输出头进行多标记预测,损失函数设计为综合考虑所有输出头的预测结果,确保模型能够有效学习多个标记的关系。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,使用多标记预测的13B参数模型在HumanEval和MBPP基准上分别解决了12%和17%更多的问题,且在推理速度上提高了最多3倍,显示出显著的性能提升和效率优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、代码生成和其他需要高效文本生成的任务。通过提高模型的样本效率和推理速度,该方法可以在实际应用中显著提升用户体验和系统性能,尤其是在需要快速响应的场景中。

📄 摘要(原文)

Large language models such as GPT and Llama are trained with a next-token prediction loss. In this work, we suggest that training language models to predict multiple future tokens at once results in higher sample efficiency. More specifically, at each position in the training corpus, we ask the model to predict the following n tokens using n independent output heads, operating on top of a shared model trunk. Considering multi-token prediction as an auxiliary training task, we measure improved downstream capabilities with no overhead in training time for both code and natural language models. The method is increasingly useful for larger model sizes, and keeps its appeal when training for multiple epochs. Gains are especially pronounced on generative benchmarks like coding, where our models consistently outperform strong baselines by several percentage points. Our 13B parameter models solves 12 % more problems on HumanEval and 17 % more on MBPP than comparable next-token models. Experiments on small algorithmic tasks demonstrate that multi-token prediction is favorable for the development of induction heads and algorithmic reasoning capabilities. As an additional benefit, models trained with 4-token prediction are up to 3 times faster at inference, even with large batch sizes.