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

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

第一作者陶超凡（Chaofan Tao）是香港大学（HKU）的四年级博士生，导师是黄毅教授和罗平教授。他本科毕业于电子科技大学的英才荣誉学院。他的研究论文发表在ACL、EMNLP、ECCV、NeurIPS、ICML、T-NNLS等期刊和会议上。他获得了 ACL 2022 年的杰出论文奖。陶超凡的研究兴趣包括：1) 高效机器学习与模型加速：以低成本对模型进行高效调优和部署。2) 支持多种任务的通用大型模型，涵盖不同模态。

本文是一篇发表在 NeurIPS 2024 上的论文，单位是香港大学、Sea AI Lab、Contextual AI 和俄亥俄州立大学。论文主要探讨了大型语言模型（LLMs）的词表大小对模型性能的影响。

• 论文：https://arxiv.org/abs/2407.13623

• 代码：https://github.com/sail-sg/scaling-with-vocab/

• Demo （快速预计合适的词表大小）：https://huggingface.co/spaces/sail/scaling-with-vocab-demo

摘要

研究大型语言模型（LLMs）的扩展法则（scaling laws）时，以前的工作主要关注模型参数和训练数据的大小，而忽略了词表大小的作用。本研究通过训练包含不同词表配置的模型（参数范围从 33M 到 3B，字符数最多 500B），提出了三种方法来预测计算最优的词表大小：基于 FLOPs 的、基于导数的和基于损失函数参数拟合的估计方法。研究结果表明，更大的模型应该配备更大的词表，且在给定算力的情况下，最优的词表大小是有上限的。例如，预测 Llama2-70B 的最优词表大小应该是至少 216K，远大于其实际的 32K。通过在不同 FLOPs 预算下训练 3B 参数的模型验证了这些预测，发现仅仅把原始词表的大小替换成预测的最优词表大小，就可以提高模型在多个下游任务的性能。

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

第 1 章 引言

LLMs 通过在大量文本语料库上进行预训练，利用巨大的计算资源，已经取得了显著的性能。以往的研究主要集中在模型参数、训练数据量和计算资源（如 FLOPs）的变化对模型性能的影响，而忽略了词表大小这一重要因素。事实上，词表大小对语言模型的性能有着不小的影响。所以，本研究旨在填补这一空白，探讨词表大小对 LLMs 性能的影响，并提出预测最优词表大小的方法。

如图，我们提出 3 种预测最优词表大小的方法 (基于 FLOPs 的、基于导数的和基于损失函数参数拟合的估计方法)，并且列出了当前主流的大型语言模型（LLMs）的词表参数和预测最优词表参数的关系。当前大多数 LLMs 的词表参数由于词表大小小于预测的最优值而处于次优状态。

第 2 章 预备知识

2.1 扩展法则

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

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

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

2.2 考虑词表的扩展法则

对训练数据量的统计方式

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

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

对词表大小非敏感的损失

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

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

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

第 3 章 分析：为什么最优词表大小受计算限制

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

较小的 V：增加词表大小可以提高标记化分词的效率，也就是用更短的词元去表示文本，从而提高模型性能。

较大的 V：逐渐增加词表大小的时候，分词效率提高的收益会逐渐减少，且可能导致词表有关参数的欠拟合，特别是针对低频词的词表征。

进一步地，我们研究了在固定 FLOP 预算下，词表如何影响损失，并发现对于每个 FLOPs 预算，存在一个使损失最小化的最优词表大小。

图上是在不同 FLOP 预算下，不同词表大小的损失曲线。对于每个预算，都存在一个最小化损失的最优词表大小。并且随着 FLOPs 预算的增加，这个最优词表大小也会增加（向右移动）。

第 4 章：估计模型的最优词表大小

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

4.1 方法 1：通过 IsoFLOPs 估计

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

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

从拟合结果我们可以看出：

• LLMs 对数据需求量大。与非词表参数 相比，从业者应分配更多计算资源用于训练数据 

• 词表参数与 FLOPs 呈幂律关系（）。随着模型计算量的增加，更大的词表大小增强了模型理解更多样化文本的能力，因此词表大小对模型扩展至关重要。

• 词表参数应比非词表参数增长得更慢。这种差异可以从它们的幂律指数中看出，即。

4.2 方法 2：基于导数的快速估计

这一方法的核心思想是通过计算 FLOPs 相对于词表大小的导数，并找到零点解，来估计最优词表大小。根据前人的研究，transformer 架构的训练 FLOPs 可以近似表示为：

通过对 V 求导，我们可以得到：

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

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

4.3 方法 3：损失公式的参数拟合

这一方法的核心思想是直接预测给定非词表参数、词表参数和训练字符数量的损失，然后通过找到损失相对于词表的最小点来预测最优词表配置。本文设计了一个依赖于词表大小的损失函数：

其中，是可学习的参数。

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

第 5 章 进一步讨论

5.1 预测语言模型的理想词表大小

在这一节中，本文报告了基于三种方法预测的最优词表参数和大小，遵从以前的关于 scaling laws 中数据量和模型参数的算力分配有关工作 [1]，训练数据的量与非词表参数等比例地随 FLOPs 预算缩放。  

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

• 预测结果显示，随着非词表参数的增加，最优的词表参数和词表大小也随之增加。

• 这表明对于更大的模型，更大的词表是必要的，以充分利用模型的容量。

值得注意的是，主流 LLMs 通常分配给词表参数偏少。然而，学界和工业界已经开始转向更大的词表大小。例如 Llama3 的词表大小从 Llama2 的 32K 增加到 128K。然而，扩展数据仍然是最关键的部分，解决数据稀缺问题应成为未来工作的重点。

为了验证这些预测，本文在 3B 参数的模型上进行了实验，这些模型在不同的 FLOPs 预算下进行了训练。实验中，本文比较了使用常规词表大小（例如 32K）与使用预测的最优词表大小的模型性能。性能通过多个下游任务进行评估，包括 ARC-Challenge、Hellaswag 等。实验结果表明，使用预测的最优词表大小的模型在多个任务上一致地优于使用常规词表大小的模型。

5.2 训练数据量对最优词表大小的影响

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

第 6 章 结论

本文通过实验验证了词表大小对语言模型性能有显著影响。他们发现，对于给定的计算预算，存在一个最优的词表大小，能够最大化模型性能。词表大小是影响模型性能的关键因素之一。更大的模型受益于更大的词表，因为它们需要更丰富的词表来表达更复杂的语言模式。另一方面，词表参数应该比非词表参数增长得慢，但仍然对性能至关重要。论文提出了三种方法来预测最优词表大小，这些方法都基于计算预算和模型性能之间的关系，论文强调了在设计和训练 LLMs 时，需要综合考虑模型参数、训练数据和词表大小。本文建议在分配计算资源时，应该考虑到词表大小的影响。

引用

[1] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. 2022. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556.