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

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

• 論文：https://arxiv.org/abs/2407.13623
• 代碼：https://github.com/sail-sg/scaling-with-vocab/
• Demo （快速預計合適的詞表大?。篽ttps://huggingface.co/spaces/sail/scaling-with-vocab-demo

摘要

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

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

第 1 章 引言

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

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

第 2 章 預備知識

2.1 擴展法則

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

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

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

2.2 考慮詞表的擴展法則

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

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

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

對詞表大小非敏感的損失

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

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

備注：BPC （平均每字的比特數(shù)）也是一個值得嘗試的、對詞表大小非敏感的指標，本文認為和 BPC 分別是從詞元（token）和字符 (character) 的角度對語言建模損失進行歸一化，達到對詞表大小非敏感的特點。

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

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

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

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

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

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

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

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

4.1 方法 1：通過 IsoFLOPs 估計

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

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

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

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

4.2 方法 2：基于導數(shù)的快速估計

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

通過對 V 求導，我們可以得到：

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

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

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

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

其中，是可學習的參數(shù)。

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

第 5 章 進一步討論

5.1 預測語言模型的理想詞表大小

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

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

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

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

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

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

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

第 6 章 結(jié)論

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