ICML 2024 || 多token预测的训练显著提高大模型性能和推理速度

大家好,今天为大家介绍一篇有意思的研究论文。这篇论文提出了一种名为 多token预测的训练方法,能显著提高大模型的性能和推理速度 。这项工作为改进当前语言模型训练范式提供了一种简单有效的方法,值得我们关注。

1. 基本信息

• 论文题目: Better & Faster Large Language Models via Multi-token Prediction
• 作者: Fabian Gloeckle*, Badr Youbi Idrissi*, Baptiste Rozière, David Lopez-Paz+, Gabriel Synnaeve+
  (* 共同第一作者, + 共同通讯作者)
• 作者单位:
• FAIR at Meta
• CERMICS Ecole des Ponts ParisTech
• LISN Université Paris-Saclay

2. 研究背景

近年来,大模型(LLM)取得了显著成就,但仍存在一些局限性。传统的next-token prediction训练方法虽然简单有效,但在获取语言、世界知识和推理能力方面效率不高。具体来说,teacher forcing与next-token prediction的结合导致模型过于关注局部模式,忽视了"困难"的决策。这使得当前最先进的next-token predictor 需要比人类儿童多几个数量级的数据才能达到相同的语言水平。

为解决这个问题,本文提出了multi-token prediction训练方法。该方法指导LLM同时并行预测未来的n个token, 而不仅仅预测下一个token 。作者认为,这种方法将提高模型的样本效率。

3. 方法

multi-token prediction方法的核心思想是:在训练语料的每个位置,让模型一次性并行预测 未来的n个token 。具体来说,模型架构包含一个共享的transformer主干和n个独立的输出头,这些输出头都建立在共享主干之上。

形式化描述如下:

标准语言建模通过实现next-token prediction任务来学习大型文本语料 。学习目标是最小化交叉熵损失:

其中 是正在训练的大模型,目标是最大化给定过去token历史 的情况下,下一个未来token 的概率。

本文将上述方法推广为multi-token prediction任务,在训练语料的每个位置,模型被指示一次性预测n个未来token。这转化为以下交叉熵损失:

为了使计算可行,作者假设大模型 使用共享主干生成观察到的上下文 的潜在表示 ,然后将 其输入到n个独立的头部,以并行预测每个未来token 。这导致multi-token prediction交叉熵损失的以下分解:

在实际应用中,模型架构包括:

1. 共享的transformer主干 ,从观察到的上下文 生成隐藏表示
2. n个独立的输出头,由transformer层 实现
3. 共享的unembedding矩阵

因此,要预测n个未来token,模型计算:

其中 ,特别地,