NeurIPS 2024 | 大模型的词表大小，同样适用于Scaling Law

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

• 论文：https://arxiv.org/abs/2407.13623
• 代码：https://github.com/sail-sg/scaling-with-vocab/
• Demo （快速预计合适的词表大小）：https://huggingface.co/spaces/sail/scaling-with-vocab-demo

本文发现，模型中的非词表参数与相应的最优词表参数之间的关系遵循幂律，其中的增长速度应慢于，即。实证结果与我们所提出的 3 种预测最优词表大小的方法的结果基本一致。其中较大的圆圈表示较高的损失值。这里指的是词表大小。

扩展法则 (scaling laws) 考虑了一个计算预算（以 FLOPs 衡量），目标是在模型参数 N 和训数据量 D 之间最优地分配这个算力的预算：

通常情况，我们使用语言模型损失 来度量语言模型：

其中 是在给定上下文 和词表大小为 的分词器的情况下单词 的输出概率。

因为对于同一个训练语料和给定的分词算法，不同的词表大小会得到不同的词元量 (D)， 因此我们以训练字符量（H）来衡量训练数据量。为了将我们的发现与现有的关于扩展规律的研究联系起来，我们需要能够从 H 映射到 D。这种映射是分词器的压缩比，可以通过 计算。分词器需要表示 H 的标记越多，D 就越大，因此压缩得越少。我们设计了一个简单的函数 来仅从选择的词表大小 V 估计这个比例：

在基于 FLOPs 和损失函数的估计方法中，我们可以直接统计出训练的词元数量。在基于导数的估计方法，我们是通过 FLOPs 关于的解析式来导出最优的，而不是使用具体的大量实验数据来拟合。这时候，我们需要使用来估计 D 和 H 的关系。附录中我们可视化了拟合结果，并展示了我们的近似方法适用于不同的分词器，并且对不同的词表大小具有鲁棒性。

语言模型损失或者常用的困惑度(perplexity)是随词表大小变化而变的，因此在模型词表大小是非固定的情况，我们不能直接在扩展法则中直接采样这个损失来度量模型。为了公平地评估不同词表大小的模型，我们采样了一种归一化损失函数，来消减由词表大小对语言建模损失带来的影响：

其中，是给定上下文和词表大小V的条件下，词的输出概率。是词在分词之后的语料库中的出现频率。

备注：BPC （平均每字的比特数）也是一个值得尝试的、对词表大小非敏感的指标，本文认为和 BPC 分别是从词元（token）和字符 (character) 的角度对语言建模损失进行归一化，达到对词表大小非敏感的特点。

我们分析了随着词表大小的增长，模型性能先是提高然后降低的原因。词表大小 对语言模型的性能的影响：

在第 4 章中，论文描述了三种不同的方法来估计大型语言模型（LLMs）的最优词表大小。这三种方法包括：通过 IsoFLOPs 估计幂律、基于导数的快速估计和损失公式的参数拟合。每一种方法都旨在预测在给定计算预算下，最优的词表大小应该是多少。不失一般性的，我们建立了词表大小和词表参数的关系，， 其中 d 是词表征的维度。

这一方法的核心思想是通过保持 FLOPs 不变，变化词表配置，来探索词表大小对模型性能的影响。本文定义了 6 组模型，每组的非词表参数（模型总参数 - 词表参数）从 33M 到 1.13B 不等。在每组中，只改变词表大小 V，从 4K 到 96K 选择了 10 种大小的词表，并在相同的 FLOPs 预算下评估不同的模型。模型架构遵循 Llama 模型的设计，训练数据集为 SlimPajama，一个大规模的文本清理和去重数据集，采用了 bfloat16 混合精度训练。本文选择了每个 FLOPs 预算下归一化损失最小的数据点，曲线如图所示：

我们发现了非词表参数、词表参数和训练字符数与 FLOPs 预算之间的关系可以用幂律表示。基于以前的研究 [1]，数据量和模型参数在最优的算力分配下应该同比例的放缩，我们在拟合过程中，加入了假设：对于词表参数和训练字符数， 都和 FLOPs 保持同样的放缩比例。通过拟合，本文得到了以下幂律关系：

• LLMs 对数据需求量大。与非词表参数 相比，从业者应分配更多计算资源用于训练数据 
• 词表参数与 FLOPs 呈幂律关系（）。随着模型计算量的增加，更大的词表大小增强了模型理解更多样化文本的能力，因此词表大小对模型扩展至关重要。
• 词表参数应比非词表参数增长得更慢。这种差异可以从它们的幂律指数中看出，即。

通过设置导数等于 0，我们可以求解最优的，也就是。这个方程的解将给出在给定的 FLOPs 预算下，能够使 FLOPs 最小化的词表大小。这种方法的优势在于它不需要大量的实验数据，而是通过数学推导和数值方法来快速估计最优词表大小。这对于初步模型设计和快速迭代非常有用。

在具体的使用过程，我们是通过导数得到的最优的词表参数 和非词表参数 拟合出了一个符合幂律参数，也就是 中的， 从通过一组轻量化小模型的实验数据点找出一组满足 和最优的 作为初始点，结合，去预测任意 情况下的最优词表参数，详见原文。

其中，是可学习的参数。

通过收集不同非词表参数、词表大小和训练数据量的实验点，并使用这些点来拟合上述损失函数的参数，我们拟合出这个有有关的损失函数，从而可以通过对关于求导的方式，找到最优的。这个方法的好处在于，它可以给出任意的非词表参数和训练数据量 的组合情况下的局部最优的词表大小，而不仅仅是 和训练数据量 等比例放缩情况下的最优词表大小。

我们报告了在给定 的情况下，通过提出的三种方法预测的最优词表参数 和词表大小。我们假设训练 FLOPs 被最优分配，即非词表参数和训练数据按相同比例扩展。结果分析：

• 预测结果显示，随着非词表参数的增加，最优的词表参数和词表大小也随之增加。
• 这表明对于更大的模型，更大的词表是必要的，以充分利用模型的容量。

我们之前的实验主要集中在训练计算预算为主要约束条件的情况下，我们寻求将其最优分配给参数和训练数据。这是扩展规律研究中的典型设置。然而，在实践中，我们经常面临数据稀缺或者数据相对于模型大小过量的情况，迫使我们进行算力次优分配时候的训练。为了验证我们的方法 3 能够处理这些实际场景中由于训练数据量变化对最优词表大小的影响，我们将词表大小为 的模型与方法 3 预测的最优词表大小 的模型进行了比较。如表所示，我们的预测可以根据不同的训练数据量，有效调整词表大小，实现了更好的模型。