第一作者陶超凡（Chaofan Tao）是香港大學(xué)（HKU）的四年級(jí)博士生，導(dǎo)師是黃毅教授和羅平教授。他本科畢業(yè)于電子科技大學(xué)的英才榮譽(yù)學(xué)院。他的研究論文發(fā)表在ACL、EMNLP、ECCV、NeurIPS、ICML、T-NNLS等期刊和會(huì)議上。他獲得了 ACL 2022 年的杰出論文獎(jiǎng)。陶超凡的研究興趣包括：1) 高效機(jī)器學(xué)習(xí)與模型加速：以低成本對(duì)模型進(jìn)行高效調(diào)優(yōu)和部署。2) 支持多種任務(wù)的通用大型模型，涵蓋不同模態(tài)。

本文是一篇發(fā)表在 NeurIPS 2024 上的論文，單位是香港大學(xué)、Sea AI Lab、Contextual AI 和俄亥俄州立大學(xué)。論文主要探討了大型語言模型（LLMs）的詞表大小對(duì)模型性能的影響。

• 論文：https://arxiv.org/abs/2407.13623
• 代碼：https://github.com/sail-sg/scaling-with-vocab/
• Demo （快速預(yù)計(jì)合適的詞表大小）：https://huggingface.co/spaces/sail/scaling-with-vocab-demo

摘要

研究大型語言模型（LLMs）的擴(kuò)展法則（scaling laws）時(shí)，以前的工作主要關(guān)注模型參數(shù)和訓(xùn)練數(shù)據(jù)的大小，而忽略了詞表大小的作用。本研究通過訓(xùn)練包含不同詞表配置的模型（參數(shù)范圍從 33M 到 3B，字符數(shù)最多 500B），提出了三種方法來預(yù)測(cè)計(jì)算最優(yōu)的詞表大小：基于 FLOPs 的、基于導(dǎo)數(shù)的和基于損失函數(shù)參數(shù)擬合的估計(jì)方法。研究結(jié)果表明，更大的模型應(yīng)該配備更大的詞表，且在給定算力的情況下，最優(yōu)的詞表大小是有上限的。例如，預(yù)測(cè) Llama2-70B 的最優(yōu)詞表大小應(yīng)該是至少 216K，遠(yuǎn)大于其實(shí)際的 32K。通過在不同 FLOPs 預(yù)算下訓(xùn)練 3B 參數(shù)的模型驗(yàn)證了這些預(yù)測(cè)，發(fā)現(xiàn)僅僅把原始詞表的大小替換成預(yù)測(cè)的最優(yōu)詞表大小，就可以提高模型在多個(gè)下游任務(wù)的性能。

本文發(fā)現(xiàn)，模型中的非詞表參數(shù)與相應(yīng)的最優(yōu)詞表參數(shù)之間的關(guān)系遵循冪律，其中的增長速度應(yīng)慢于，即。實(shí)證結(jié)果與我們所提出的 3 種預(yù)測(cè)最優(yōu)詞表大小的方法的結(jié)果基本一致。其中較大的圓圈表示較高的損失值。這里指的是詞表大小。

第 1 章 引言

LLMs 通過在大量文本語料庫上進(jìn)行預(yù)訓(xùn)練，利用巨大的計(jì)算資源，已經(jīng)取得了顯著的性能。以往的研究主要集中在模型參數(shù)、訓(xùn)練數(shù)據(jù)量和計(jì)算資源（如 FLOPs）的變化對(duì)模型性能的影響，而忽略了詞表大小這一重要因素。事實(shí)上，詞表大小對(duì)語言模型的性能有著不小的影響。所以，本研究旨在填補(bǔ)這一空白，探討詞表大小對(duì) LLMs 性能的影響，并提出預(yù)測(cè)最優(yōu)詞表大小的方法。

如圖，我們提出 3 種預(yù)測(cè)最優(yōu)詞表大小的方法 (基于 FLOPs 的、基于導(dǎo)數(shù)的和基于損失函數(shù)參數(shù)擬合的估計(jì)方法)，并且列出了當(dāng)前主流的大型語言模型（LLMs）的詞表參數(shù)和預(yù)測(cè)最優(yōu)詞表參數(shù)的關(guān)系。當(dāng)前大多數(shù) LLMs 的詞表參數(shù)由于詞表大小小于預(yù)測(cè)的最優(yōu)值而處于次優(yōu)狀態(tài)。

第 2 章 預(yù)備知識(shí)

2.1 擴(kuò)展法則

擴(kuò)展法則 (scaling laws) 考慮了一個(gè)計(jì)算預(yù)算（以 FLOPs 衡量），目標(biāo)是在模型參數(shù) N 和訓(xùn)數(shù)據(jù)量 D 之間最優(yōu)地分配這個(gè)算力的預(yù)算：

通常情況，我們使用語言模型損失 來度量語言模型：

其中 是在給定上下文 和詞表大小為 的分詞器的情況下單詞 的輸出概率。

2.2 考慮詞表的擴(kuò)展法則

對(duì)訓(xùn)練數(shù)據(jù)量的統(tǒng)計(jì)方式

因?yàn)閷?duì)于同一個(gè)訓(xùn)練語料和給定的分詞算法，不同的詞表大小會(huì)得到不同的詞元量 (D)， 因此我們以訓(xùn)練字符量（H）來衡量訓(xùn)練數(shù)據(jù)量。為了將我們的發(fā)現(xiàn)與現(xiàn)有的關(guān)于擴(kuò)展規(guī)律的研究聯(lián)系起來，我們需要能夠從 H 映射到 D。這種映射是分詞器的壓縮比，可以通過 計(jì)算。分詞器需要表示 H 的標(biāo)記越多，D 就越大，因此壓縮得越少。我們?cè)O(shè)計(jì)了一個(gè)簡(jiǎn)單的函數(shù) 來僅從選擇的詞表大小 V 估計(jì)這個(gè)比例：

在基于 FLOPs 和損失函數(shù)的估計(jì)方法中，我們可以直接統(tǒng)計(jì)出訓(xùn)練的詞元數(shù)量。在基于導(dǎo)數(shù)的估計(jì)方法，我們是通過 FLOPs 關(guān)于的解析式來導(dǎo)出最優(yōu)的，而不是使用具體的大量實(shí)驗(yàn)數(shù)據(jù)來擬合。這時(shí)候，我們需要使用來估計(jì) D 和 H 的關(guān)系。附錄中我們可視化了擬合結(jié)果，并展示了我們的近似方法適用于不同的分詞器，并且對(duì)不同的詞表大小具有魯棒性。

對(duì)詞表大小非敏感的損失

語言模型損失或者常用的困惑度(perplexity)是隨詞表大小變化而變的，因此在模型詞表大小是非固定的情況，我們不能直接在擴(kuò)展法則中直接采樣這個(gè)損失來度量模型。為了公平地評(píng)估不同詞表大小的模型，我們采樣了一種歸一化損失函數(shù)，來消減由詞表大小對(duì)語言建模損失帶來的影響：

其中，是給定上下文和詞表大小V的條件下，詞的輸出概率。是詞在分詞之后的語料庫中的出現(xiàn)頻率。

備注：BPC （平均每字的比特?cái)?shù)）也是一個(gè)值得嘗試的、對(duì)詞表大小非敏感的指標(biāo)，本文認(rèn)為和 BPC 分別是從詞元（token）和字符 (character) 的角度對(duì)語言建模損失進(jìn)行歸一化，達(dá)到對(duì)詞表大小非敏感的特點(diǎn)。

第 3 章 分析：為什么最優(yōu)詞表大小受計(jì)算限制

我們分析了隨著詞表大小的增長，模型性能先是提高然后降低的原因。詞表大小 對(duì)語言模型的性能的影響：

較小的 V：增加詞表大小可以提高標(biāo)記化分詞的效率，也就是用更短的詞元去表示文本，從而提高模型性能。

較大的 V：逐漸增加詞表大小的時(shí)候，分詞效率提高的收益會(huì)逐漸減少，且可能導(dǎo)致詞表有關(guān)參數(shù)的欠擬合，特別是針對(duì)低頻詞的詞表征。

進(jìn)一步地，我們研究了在固定 FLOP 預(yù)算下，詞表如何影響損失，并發(fā)現(xiàn)對(duì)于每個(gè) FLOPs 預(yù)算，存在一個(gè)使損失最小化的最優(yōu)詞表大小。

圖上是在不同 FLOP 預(yù)算下，不同詞表大小的損失曲線。對(duì)于每個(gè)預(yù)算，都存在一個(gè)最小化損失的最優(yōu)詞表大小。并且隨著 FLOPs 預(yù)算的增加，這個(gè)最優(yōu)詞表大小也會(huì)增加（向右移動(dòng)）。

第 4 章：估計(jì)模型的最優(yōu)詞表大小

在第 4 章中，論文描述了三種不同的方法來估計(jì)大型語言模型（LLMs）的最優(yōu)詞表大小。這三種方法包括：通過 IsoFLOPs 估計(jì)冪律、基于導(dǎo)數(shù)的快速估計(jì)和損失公式的參數(shù)擬合。每一種方法都旨在預(yù)測(cè)在給定計(jì)算預(yù)算下，最優(yōu)的詞表大小應(yīng)該是多少。不失一般性的，我們建立了詞表大小和詞表參數(shù)的關(guān)系，， 其中 d 是詞表征的維度。

4.1 方法 1：通過 IsoFLOPs 估計(jì)

這一方法的核心思想是通過保持 FLOPs 不變，變化詞表配置，來探索詞表大小對(duì)模型性能的影響。本文定義了 6 組模型，每組的非詞表參數(shù)（模型總參數(shù) - 詞表參數(shù)）從 33M 到 1.13B 不等。在每組中，只改變?cè)~表大小 V，從 4K 到 96K 選擇了 10 種大小的詞表，并在相同的 FLOPs 預(yù)算下評(píng)估不同的模型。模型架構(gòu)遵循 Llama 模型的設(shè)計(jì)，訓(xùn)練數(shù)據(jù)集為 SlimPajama，一個(gè)大規(guī)模的文本清理和去重?cái)?shù)據(jù)集，采用了 bfloat16 混合精度訓(xùn)練。本文選擇了每個(gè) FLOPs 預(yù)算下歸一化損失最小的數(shù)據(jù)點(diǎn)，曲線如圖所示：

我們發(fā)現(xiàn)了非詞表參數(shù)、詞表參數(shù)和訓(xùn)練字符數(shù)與 FLOPs 預(yù)算之間的關(guān)系可以用冪律表示。基于以前的研究 [1]，數(shù)據(jù)量和模型參數(shù)在最優(yōu)的算力分配下應(yīng)該同比例的放縮，我們?cè)跀M合過程中，加入了假設(shè)：對(duì)于詞表參數(shù)和訓(xùn)練字符數(shù)， 都和 FLOPs 保持同樣的放縮比例。通過擬合，本文得到了以下冪律關(guān)系：

從擬合結(jié)果我們可以看出：

• LLMs 對(duì)數(shù)據(jù)需求量大。與非詞表參數(shù) 相比，從業(yè)者應(yīng)分配更多計(jì)算資源用于訓(xùn)練數(shù)據(jù) 
• 詞表參數(shù)與 FLOPs 呈冪律關(guān)系（）。隨著模型計(jì)算量的增加，更大的詞表大小增強(qiáng)了模型理解更多樣化文本的能力，因此詞表大小對(duì)模型擴(kuò)展至關(guān)重要。
• 詞表參數(shù)應(yīng)比非詞表參數(shù)增長得更慢。這種差異可以從它們的冪律指數(shù)中看出，即。

4.2 方法 2：基于導(dǎo)數(shù)的快速估計(jì)

這一方法的核心思想是通過計(jì)算 FLOPs 相對(duì)于詞表大小的導(dǎo)數(shù)，并找到零點(diǎn)解，來估計(jì)最優(yōu)詞表大小。根據(jù)前人的研究，transformer 架構(gòu)的訓(xùn)練 FLOPs 可以近似表示為：

通過對(duì) V 求導(dǎo)，我們可以得到：

通過設(shè)置導(dǎo)數(shù)等于 0，我們可以求解最優(yōu)的，也就是。這個(gè)方程的解將給出在給定的 FLOPs 預(yù)算下，能夠使 FLOPs 最小化的詞表大小。這種方法的優(yōu)勢(shì)在于它不需要大量的實(shí)驗(yàn)數(shù)據(jù)，而是通過數(shù)學(xué)推導(dǎo)和數(shù)值方法來快速估計(jì)最優(yōu)詞表大小。這對(duì)于初步模型設(shè)計(jì)和快速迭代非常有用。

在具體的使用過程，我們是通過導(dǎo)數(shù)得到的最優(yōu)的詞表參數(shù) 和非詞表參數(shù) 擬合出了一個(gè)符合冪律參數(shù)，也就是 中的， 從通過一組輕量化小模型的實(shí)驗(yàn)數(shù)據(jù)點(diǎn)找出一組滿足 和最優(yōu)的 作為初始點(diǎn)，結(jié)合，去預(yù)測(cè)任意 情況下的最優(yōu)詞表參數(shù)，詳見原文。

4.3 方法 3：損失公式的參數(shù)擬合

這一方法的核心思想是直接預(yù)測(cè)給定非詞表參數(shù)、詞表參數(shù)和訓(xùn)練字符數(shù)量的損失，然后通過找到損失相對(duì)于詞表的最小點(diǎn)來預(yù)測(cè)最優(yōu)詞表配置。本文設(shè)計(jì)了一個(gè)依賴于詞表大小的損失函數(shù)：

其中，是可學(xué)習(xí)的參數(shù)。

通過收集不同非詞表參數(shù)、詞表大小和訓(xùn)練數(shù)據(jù)量的實(shí)驗(yàn)點(diǎn)，并使用這些點(diǎn)來擬合上述損失函數(shù)的參數(shù)，我們擬合出這個(gè)有有關(guān)的損失函數(shù)，從而可以通過對(duì)關(guān)于求導(dǎo)的方式，找到最優(yōu)的。這個(gè)方法的好處在于，它可以給出任意的非詞表參數(shù)和訓(xùn)練數(shù)據(jù)量 的組合情況下的局部最優(yōu)的詞表大小，而不僅僅是 和訓(xùn)練數(shù)據(jù)量 等比例放縮情況下的最優(yōu)詞表大小。

第 5 章 進(jìn)一步討論

5.1 預(yù)測(cè)語言模型的理想詞表大小

在這一節(jié)中，本文報(bào)告了基于三種方法預(yù)測(cè)的最優(yōu)詞表參數(shù)和大小，遵從以前的關(guān)于 scaling laws 中數(shù)據(jù)量和模型參數(shù)的算力分配有關(guān)工作 [1]，訓(xùn)練數(shù)據(jù)的量與非詞表參數(shù)等比例地隨 FLOPs 預(yù)算縮放。  

我們報(bào)告了在給定 的情況下，通過提出的三種方法預(yù)測(cè)的最優(yōu)詞表參數(shù) 和詞表大小。我們假設(shè)訓(xùn)練 FLOPs 被最優(yōu)分配，即非詞表參數(shù)和訓(xùn)練數(shù)據(jù)按相同比例擴(kuò)展。結(jié)果分析：

• 預(yù)測(cè)結(jié)果顯示，隨著非詞表參數(shù)的增加，最優(yōu)的詞表參數(shù)和詞表大小也隨之增加。
• 這表明對(duì)于更大的模型，更大的詞表是必要的，以充分利用模型的容量。

值得注意的是，主流 LLMs 通常分配給詞表參數(shù)偏少。然而，學(xué)界和工業(yè)界已經(jīng)開始轉(zhuǎn)向更大的詞表大小。例如 Llama3 的詞表大小從 Llama2 的 32K 增加到 128K。然而，擴(kuò)展數(shù)據(jù)仍然是最關(guān)鍵的部分，解決數(shù)據(jù)稀缺問題應(yīng)成為未來工作的重點(diǎn)。

為了驗(yàn)證這些預(yù)測(cè)，本文在 3B 參數(shù)的模型上進(jìn)行了實(shí)驗(yàn)，這些模型在不同的 FLOPs 預(yù)算下進(jìn)行了訓(xùn)練。實(shí)驗(yàn)中，本文比較了使用常規(guī)詞表大小（例如 32K）與使用預(yù)測(cè)的最優(yōu)詞表大小的模型性能。性能通過多個(gè)下游任務(wù)進(jìn)行評(píng)估，包括 ARC-Challenge、Hellaswag 等。實(shí)驗(yàn)結(jié)果表明，使用預(yù)測(cè)的最優(yōu)詞表大小的模型在多個(gè)任務(wù)上一致地優(yōu)于使用常規(guī)詞表大小的模型。

5.2 訓(xùn)練數(shù)據(jù)量對(duì)最優(yōu)詞表大小的影響

我們之前的實(shí)驗(yàn)主要集中在訓(xùn)練計(jì)算預(yù)算為主要約束條件的情況下，我們尋求將其最優(yōu)分配給參數(shù)和訓(xùn)練數(shù)據(jù)。這是擴(kuò)展規(guī)律研究中的典型設(shè)置。然而，在實(shí)踐中，我們經(jīng)常面臨數(shù)據(jù)稀缺或者數(shù)據(jù)相對(duì)于模型大小過量的情況，迫使我們進(jìn)行算力次優(yōu)分配時(shí)候的訓(xùn)練。為了驗(yàn)證我們的方法 3 能夠處理這些實(shí)際場(chǎng)景中由于訓(xùn)練數(shù)據(jù)量變化對(duì)最優(yōu)詞表大小的影響，我們將詞表大小為 的模型與方法 3 預(yù)測(cè)的最優(yōu)詞表大小 的模型進(jìn)行了比較。如表所示，我們的預(yù)測(cè)可以根據(jù)不同的訓(xùn)練數(shù)據(jù)量，有效調(diào)整詞表大小，實(shí)現(xiàn)了更好的模型。

第 6 章 結(jié)論

本文通過實(shí)驗(yàn)驗(yàn)證了詞表大小對(duì)語言模型性能有顯著影響。他們發(fā)現(xiàn)，對(duì)于給定的計(jì)算預(yù)算，存在一個(gè)最優(yōu)的詞表大小，能夠最大化模型性能。詞表大小是影響模型性能的關(guān)鍵因素之一。更大的模型受益于更大的詞表，因?yàn)樗鼈冃枰S富的詞表來表達(dá)更復(fù)雜的語言模式。另一方面，詞表參數(shù)應(yīng)該比非詞表參數(shù)增長得慢，但仍然對(duì)性能至關(guān)重要。論文提出了三種方法來預(yù)測(cè)最優(yōu)詞表大小，這些方法都基于計(jì)算預(yù)算和模型性能之間的關(guān)系，論文強(qiáng)調(diào)了在設(shè)計(jì)和訓(xùn)練 LLMs 時(shí)，需要綜合考慮模型參數(shù)、訓(xùn)練數(shù)據(jù)和詞表大小。本文建議在分配計(jì)算資源時(shí)，應(yīng)該考慮到詞表大小的影響。