Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens

发表时间: 2024-01 · arXiv:2401.17377 · COLM 2024 (UW + Allen AI)

文章标题：Infini-gram: 将无界n-gram语言模型扩展至万亿级词元
作者/机构：Jiacheng Liu, Sewon Min, Luke Zettlemoyer, Yejin Choi, Hannaneh Hajishirzi, 来自华盛顿大学保罗·艾伦计算机科学与工程学院 (Paul G. Allen School of Computer Science & Engineering, University of Washington) 和艾伦人工智能研究所 (Allen Institute for Artificial Intelligence)。

A1 主要贡献

本文旨在探究在大型神经网络语言模型（LLM）时代，经典的n-gram语言模型是否仍具价值。研究通过两个方面对n-gram语言模型进行了现代化改造，并证明了其在文本分析和改进神经LLM方面的作用。

  1. 数据规模的扩展：研究将n-gram语言模型的训练数据规模提升至与现代LLM相当的水平，通过整合多个大型开源文本语料库，构建了一个包含5万亿词元（tokens）的n-gram语言模型，这是迄今为止最大的n-gram语言模型。

  2. 无界n值的实现 (∞-gram LM)：传统n-gram模型因n值较小（如n≤5）而性能受限。本文提出了一种新的∞-gram语言模型，它通过一种回退（backoff）机制，允许n值任意大。当较长的n-gram在语料库中计数为零时，模型会回退到使用较短的n-gram。如图1所示，当n值从5增加到16时，模型能够利用更丰富的上下文，从而做出更准确的预测。

  3. 高效引擎 (infini-gram) 的开发：为支持大规模、无界n值的∞-gram模型，研究者开发了一个名为infini-gram的高效引擎。该引擎不依赖于预计算和存储庞大的n-gram计数表，而是基于后缀数组（suffix array）构建索引。这种方法在存储和计算上都非常高效：

    • 存储效率：每个词元的存储开销为7字节（相比原始数据为3.5倍）。
    • 构建效率：在单个128核CPU节点上，为1.4万亿词元的语料库构建索引约需2天，占用10TB磁盘空间。
    • 查询效率：平均查询延迟极低，计数一个n-gram（无论n多大）的延迟低于20毫秒，而执行∞-gram语言建模和解码等其他查询类型的延迟在200毫秒以下。所有索引在推理时均保留在磁盘上。

  4. 新颖的文本分析与发现：利用∞-gram模型，研究者对人类编写和机器生成的文本进行了深入分析，得出以下发现：

    • 高预测准确率：∞-gram在预测人类文档的下一个词元时，准确率高达47%。
    • 对神经LLM的补充：通过与神经LLM（最大至70B参数）进行启发式插值，∞-gram能显著降低模型的困惑度（perplexity），降幅高达73%。
    • 揭示机器生成文本的缺陷：分析发现，机器生成文本与∞-gram的一致性水平随后缀长度的变化呈现不规则波动，这可能揭示了神经LLM预训练过程和Transformer位置编码的不足。

  5. 开源工具与服务：为了促进社区研究，作者发布了相关的源代码、Python包，并提供了公开的Web界面和API端点，支持在Dolma、RedPajama、Pile和C4等流行语料库上进行n-gram/∞-gram查询。

图1：一个例子，其中5-gram LM给出了错误的预测，但∞-gram通过使用提示中最长的、在语料库中计数非零的后缀，给出了正确的预测。∞-gram LM中的计数和分布估计由我们的infini-gram引擎提供支持。
图1：一个例子，其中5-gram LM给出了错误的预测，但∞-gram通过使用提示中最长的、在语料库中计数非零的后缀，给出了正确的预测。∞-gram LM中的计数和分布估计由我们的infini-gram引擎提供支持。

2 ∞-gram LM：将n-gram LM扩展至无界n

n-gram LM背景。n-gram LM是一种经典的统计语言模型，其基础是计算n-gram的出现次数。在其最简单的形式中，给定上下文$w_{i-(n-1):i-1}$，下一个词元$w_i$的概率被估计为$P_n(w_i | w_{i-(n-1):i-1}) = \frac{cnt(w_{i-(n-1):i-1}w_i|D)}{cnt(w_{i-(n-1):i-1}|D)}$，其中$cnt(w | D)$是n-gram $w$在训练数据$D$中出现的次数，n是预定义的超参数。（当n=1时，上下文$w_{i-(n-1):i-1}$定义为空字符串ε，其计数等于$|D|$。）然而，这种朴素的n-gram LM面临稀疏性问题：分子可能为零，导致无限的困惑度。一个常见的解决方案是回退（backoff）【20, Jurafsky & Martin, Speech and language processing, 2000】，即在实例级别上，当分子为零时，将n减一，并重复此过程直到分子变为正数。回退的一个警示是它不会为$P_n(*|w_{i-(n-1):i-1})$产生一个有效的概率分布，因为有效的n依赖于$w_i$。因此，需要进一步的概率折扣来归一化这个分布（例如，Katz backoff【22, Katz, Estimation of probabilities from sparse data for the language model component of a speech recognizer, 1987】）。

传统n-gram LM的局限性。传统上，n-gram LM是通过构建训练数据的n-gram计数表来实现的。该表存储了训练数据中出现的所有唯一n-gram及其计数。这样的n-gram计数表非常庞大，并且其大小几乎随n呈指数增长。例如，一个1.4万亿词元语料库的5-gram计数表将消耗28TB的磁盘空间。因此，以前的n-gram LM仅限于非常小的n值，最常见的是n=5，并且只针对频繁出现的n-gram【11, Franz & Brants, All our n-gram are belong to you, 2006】。正如我们在图1中所示并将在§4中进一步量化的，小n值的问题在于它丢弃了更丰富的上下文，使得这样的n-gram LM对未来词元的预测能力很差。

∞-gram LM的定义。∞-gram LM是n-gram LM的一种泛化，概念上我们从n=∞开始回退。我们使用一种回退的变体：仅当分母为零时才回退。这意味着一旦分母变为正数，我们就会停止回退，此时分子可能仍然为零。在实例级别上，有效的n等于1加上提示（prompt）中在训练数据里出现过的最长后缀的长度。

∞-gram LM的公式。在本文的其余部分，我们将使用“∞-gram”来指代∞-gram LM。∞-gram被正式定义为：

其中，$w_{1:i-1}$是文档中$w_i$之前的所有词元，并且

∞-gram LM的概率分布有效性。与Katz回退不同，$P_∞(*|w_{1:i-1})$通过其构造本身就是一个有效的概率分布，不需要进行折扣。这是因为有效的n完全取决于上下文$w_{1:i-1}$，而不依赖于下一个词元$w_i$，并且$\sum_{w \in V} cnt(w_{i-(n-1):i-1}w | D) = cnt(w_{i-(n-1):i-1} | D)$。

∞-gram估计的稀疏性定义。我们进一步定义了这种∞-gram估计的稀疏性：当对于词汇表V中的某个$w_i$，有$P(w_i|w_{i-(n-1):i-1}) = 1$，而对于所有其他词元该概率为零时，该估计是稀疏的。直观上，这意味着根据训练数据，给定此上下文只有一个可能的下一个词元。我们将在§4中展示，稀疏估计比非稀疏估计更能预测实际的词元。

与神经LM的插值。∞-gram的估计包含零概率，这可能导致无限的困惑度。我们不试图计算∞-gram本身的困惑度，而是将其与神经LM进行插值，并展示相对于单独使用神经LM的困惑度改善（§5）。组合后的模型为：

其中$λ \in [0, 1]$是一个超参数。


图2：左：一个玩具字符串的后缀数组。右：infini-gram索引中后缀数组的图示，数据集中有N=4个词元。

3 Infini-gram：一个用于n-gram/∞-gram查询的高性能引擎

引言：使用后缀数组代替计数表。我们在现代的万亿级词元文本语料库上训练∞-gram。然而，为如此庞大的数据集构建具有无界n的n-gram计数表在实践中是不可行的。本节描述了我们的infini-gram引擎，它能高效处理n-gram/∞-gram查询。Infini-gram由一种名为后缀数组的数据结构驱动。我们将展示如何构建这个后缀数组索引以及如何用它进行n-gram/∞-gram推理。

后缀数组的基本原理。n-gram和∞-gram LM的核心是在训练数据中对给定的n-gram进行计数。因此，我们利用了后缀数组这种数据结构，它最初是为高效地计算一个给定的“针”（长度为L）字符串在一个巨大的“干草堆”（长度为N）字符串中作为子串出现的次数而设计的。当为“干草堆”字符串构建了后缀数组后，计算给定“针”字符串的时间复杂度为$O(L + \log N)$。

infini-gram中的后缀数组结构。后缀数组表示一个数组（或字符串，即字符数组）所有后缀的字典序排序。对于一个长度为N的数组，后缀数组包含N个唯一的整数，其中第i个元素是所有后缀中排名第i的后缀的起始位置。图2（左）展示了一个玩具字符串“aabaca”的后缀数组。如图2（右）所示，我们在分词后数据集的字节数组（即词元数组）上构建后缀数组。文档之间由\xff\xff词元分隔。在词元数组中，每连续两个字节代表一个词元ID（假设词汇表大小$|V| < 2^{16} = 65536$）。给定数据集有N个词元，词元数组有2N字节。后缀数组包含N个元素，每个元素通过存储其字节偏移量指向词元数组中的一个词元。词元数组中的每个词元在后缀数组中都恰好出现一次。每个指针可以用$\lceil\log(2N)/8\rceil$字节存储。对于拥有20亿到5000亿词元的语一料库（这是我们在分片后处理的范围，§A.3），每个指针为5字节，因此后缀数组大小为5N字节。所以，词元数组和后缀数组（即infini-gram索引）的总大小为7N字节。

后缀数组的构建。后缀数组的构建时间相对于词元数组的大小是线性的【21, Kärkkäinen et al., Linear work suffix array construction, 2006】。我们改编自Lee等人【26, Lee et al., Deduplicating training data makes language models better, 2022】的后缀数组实现，并为提高效率进行了进一步优化。我们在一个拥有128个CPU和1TiB RAM的单节点上，花费了约48小时为RedPajama构建了后缀数组。我们已经为Dolma（3T词元）、RedPajama（1.4T词元）、Pile（380B词元）和C4（200B词元）构建了后缀数组。由于infini-gram索引是可加的（§A.2），这些索引可以轻松组合成一个总计5万亿词元的更大索引，其隐含的计数表至少包含2千万亿（quadrillion）个唯一的n-gram（§A.1）。

使用后缀数组进行推理。计算n-gram LM概率涉及计算一个词元字符串$x_1...x_n$的出现次数，即$cnt(x_1...x_n)$。根据构造，以$x_1...x_n$开头的字符串的出现位置位于后缀数组中的一个单一、连续的段中。因此，我们只需要找到第一个和最后一个出现位置，计数就是它们之间的差值。除了计数和n-gram/∞-gram语言建模，infini-gram还可以用于检索包含某个n-gram的文档，或包含多个n-gram的合取范式（CNF）表达式（§A.4）。

推理优化与性能。在推理期间，整个infini-gram索引可以保留在磁盘上，这最大限度地减少了所需的计算资源（无需GPU，且CPU/RAM占用极少）。在§A.4中，我们讨论了应用于推理引擎的几种优化技术：并行化分片处理、提示搜索（hinted search）、内存预取、快速有效n值查找以及摊销查询处理。在RedPajama上，我们最优化的infini-gram引擎可以在平均不到20毫秒的延迟内计算一个给定的n-gram。对于n-gram LM，它可以在40毫秒内计算概率和下一个词元分布；对于∞-gram，则需要200毫秒。支持的查询类型和延迟基准测试的完整列表见§A.5。

A4 实验环境

A4 实验结果

使用∞-gram分析人类与机器文本

本节使用在Pile-train（3.6亿词元）上训练的∞-gram模型，分析其与人类编写文本（来自Pile-val）和机器生成文本的一致性。

使用∞-gram改进神经LM

本节通过将∞-gram模型与多种神经LM进行插值，评估其在降低困惑度（PPL）方面的效果。实验在Pile验证集和测试集上进行。

PROTECTED_IMAGE_15____PROTECTED_IMAGE_16

A5 结论

本文通过将经典的n-gram语言模型扩展至万亿级词元规模和无界n值，成功地对其进行了现代化改造。研究者提出了infini-gram引擎，该引擎能够在此极端设置下进行高效的训练和推理。基于此引擎，本文提出了∞-gram语言模型，并证明了它不仅能为人类编写和机器生成的文本提供新颖的洞见，还能显著改善现有神经语言模型的性能。

A6 附录

A. Infini-gram引擎的附加细节

A.1 infini-gram索引隐含的n-gram计数表有多大？

隐式n-gram数量的估算。n-gram LM的大小通常通过索引的唯一n-gram数量来衡量。对于一个有N个词元的数据集，若考虑所有可能的n值，大约有$\frac{1}{2}N^2$个n-gram。由于我们索引了N=5万亿的词元，我们infini-gram索引所隐含的n-gram计数表大小约等于$1.2 \times 10^{25}$。然而，考虑到文档边界，我们应只计算文档内的n-gram。数据集中有$N=5 \times 10^{12}$个词元和$D=6 \times 10^9$个文档，平均每个文档857个词元。使用柯西-施瓦茨不等式，我们得到n-gram总数为：

因此，infini-gram隐含的计数表中至少有2千万亿（quadrillion）个唯一的n-gram。

A.2 infini-gram索引的附加细节

索引的可加减性。若我们在不相交的数据集上（使用相同分词器）构建了两个或多个索引，可以通过简单地将每个索引的n-gram计数相加来组合它们。此特性在§A.3和§A.4讨论的分片技术中很有用。同样，我们也可以通过计数相减来获得数据集差集的索引。这些操作避免了从头重新索引，虽会增加一些可并行的推理开销，但灵活性很高。

文档偏移与元数据。为实现高效的文档检索，infini-gram索引存储了额外的文档信息。一个文档偏移文件存储每个文档在分词后数据集中的字节偏移，格式类似后缀数组。一个文档元数据文件存储每个文档的元数据（如ID、来源、URL），另有一个元数据偏移文件记录其位置。这些文件相比后缀数组大小可忽略不计。


图6：训练数据上的n-gram/∞-gram查询由一个关联的后缀数组支持。训练数据和后缀数组都作为常规文件存储在磁盘上。白色条带上的内容是文件数据，条带上方的地址是字节偏移量。查询一个特定的n-gram会返回后缀数组的一个连续段，其中每个元素都是一个指向训练数据中该n-gram出现位置的指针。例如，在万亿词元的训练数据中，“Artificial Intelligence, A Modern”出现了42次，并且在所有情况下，接下来的词元都是“Approach”。

A.3 构建后缀数组的附加细节

分片（Sharding）。构建后缀数组需要对字节数组进行大量的随机访问，因此整个字节数组必须保存在内存中以保证合理的构建时间。当字节数组过大无法放入内存时，我们将其分片，并为每个分片构建一个后缀数组。分片会引入额外的推理延迟，我们将在§A.4中讨论并缓解此问题。

A.4 使用infini-gram索引进行推理的附加细节

查找与计数。第一个和最后一个出现位置都可以通过二分查找找到，时间复杂度为$O(n \cdot \log N)$，需要$O(\log N)$次随机数组访问。两个二分查找可以并行化，延迟大约减少2倍。查询长度n的影响可以忽略，因为计算机通常以4K字节的页面获取内存，而字符串比较远快于页面获取。

查找出现位置和文档。使用后缀数组进行n-gram计数的一个副产品是，我们可以免费获得该n-gram在训练数据中出现的所有位置。这些位置信息隐含在计数过程中获得的后缀数组段中。要检索包含该n-gram的原始文档，只需跟随段内的每个指针回到分词数据集中，并通过在文档偏移索引上进行二分查找来确定文档的起止位置。

分片的影响。当后缀数组构建在分片的字节数组上时，我们只需在每个分片上单独执行计数，然后累加所有分片的计数。延迟与分片数量成正比，时间复杂度变为$O(S \cdot \log N)$。不同分片的处理可以并行化，从而将时间复杂度降回$O(\log N)$。

通过重用搜索结果加速n-gram计算。在后缀数组上，$x_1...x_n$的段必须是$x_1...x_{n-1}$段的子段。因此，在计算n-gram概率$P_n(x_n | x_1...x_{n-1})$时，我们可以先计数$x_1...x_{n-1}$，然后在计数$x_1...x_n$时，只需在$x_1...x_n$的段内搜索起止位置，这最多可将延迟减少2倍。

磁盘搜索。字节数组和后缀数组可能太大无法放入内存，因此在实践中，我们将它们保存在磁盘上，并作为内存映射文件读取。然而，这会因二分查找需要随机访问而产生显著延迟。为缓解此问题，我们实现了一种内存预取方法，通知系统我们将来可能要读取的数组偏移量。预取将平均延迟降低了大约5倍。

加速∞-gram计算。为计算∞-gram概率，我们需要对后缀$x_{l-n+1}...x_l$进行计数，直到达到最大n值L，使得后缀仍然满足足够的出现要求。这意味着需要$O(L)$次计数操作。然而，一种简单的“二进制提升+二分查找”算法可以搜索L，将计数操作减少到$O(\log L)$次，因此每个∞-gram计算的时间复杂度变为$O(\log L \cdot \log N)$。

加速密集的∞-gram计算。在评估期间，我们需要计算测试文档中每个词元的∞-gram概率。我们可以通过观察到一个词元的有效n最多比前一个词元长一个词元来节省计算。这使得评估每个词元的摊销时间复杂度降至$O(\log N)$。

A.5 支持的查询类型和延迟基准测试

Infini-gram支持以下类型的n-gram/∞-gram查询：
1. 计数n-gram (COUNT)
2. 从n-gram LM计算词元概率（给定n，无回退）(NGRAMPROB)
3. 从n-gram LM计算完整的下一词元分布 (NGRAMDIST)
4. 从∞-gram LM计算词元概率 (INFGRAMPROB)
5. 从∞-gram LM计算完整的下一词元分布 (INFGRAMDIST)
6. 返回包含n-gram或n-gram项的CNF逻辑表达式的文档 (SEARCHDOC)

我们在不同类型的查询上对infini-gram的延迟进行了基准测试，结果如表3所示。所有类型的查询在万亿级词元训练数据上都表现出亚秒级延迟。用RedPajama计算一个词元的∞-gram概率仅需135毫秒。我们的实现支持计数任意大n值的n-gram，延迟大致恒定在20毫秒。解码需要计算完整的下一词元分布，因此稍慢：n-gram LM每词元39毫秒，∞-gram每词元180毫秒。

表3：infini-gram在不同类型查询上的推理时延。报告了每个查询的平均延迟。基准测试使用C++编写的推理引擎（带并行分片处理），在单个8核CPU节点上运行。时间复杂度符号：N=参考数据中的词元数；S=后缀数组的分片数；L=查询文档中的词元数；V=词汇表大小。
表3：infini-gram在不同类型查询上的推理时延。报告了每个查询的平均延迟。基准测试使用C++编写的推理引擎（带并行分片处理），在单个8核CPU节点上运行。时间复杂度符号：N=参考数据中的词元数；S=后缀数组的分片数；L=查询文档中的词元数；V=词汇表大小。

B. 参考数据的净化

净化过程。为了在Pile的评估集上正确评估∞-gram LM的效果，我们在使用Pile的训练集和RedPajama作为∞-gram LM的参考数据之前，对它们进行了数据净化。我们运行了Big Friendly Filter (BFF)【14, Groeneveld, The big friendly filter, 2023】工具，过滤掉与Pile评估集有过多n-gram重叠的文档。表4报告了净化的统计数据。

净化设置。在使用BFF时，我们总是移除整个文档，而不是按段落移除。遵循默认设置，我们考虑n=13的n-gram，如果一个文档中至少80%的n-gram出现在评估集中，则丢弃该文档。对于Pile的训练集，我们将所有文档转为小写以捕获更多潜在的污染。

表4：RedPajama（左）和Pile训练集（右）的去污统计。
表4：RedPajama（左）和Pile训练集（右）的去污统计。

C. 附加分析

图7是对图3中间图的扩展，更细致地分析了人类编写文本与∞-gram之间逐词元的一致性。图8扩展了图4，包含了Llama-2-13b/7b的结果。图9扩展了图5，包含了GPT-Neo模型的结果。

PROTECTED_IMAGE_17__图8：图4的延续，包含Llama-2-13b/7b的结果。__PROTECTED_IMAGE_19

D. 使用∞-gram改进神经LM的附加实验

D.1 实验设置的附加细节

评估数据处理。我们将评估数据中的每个文档分割成最大序列长度为1024、滑动窗口为512的批次，这是先前语言建模文献中的标准设置【5, Baevski & Auli, Adaptive input representations for neural language modeling, 2019】、【24, Khandelwal et al., Generalization through memorization: Nearest neighbor language models, 2020】。

神经LM。以下是我们使用的神经LM家族的附加细节：
- GPT-2: 早期自回归语言模型，大小从117M到1.6B不等。
- GPT-Neo/J: 在Pile上训练的语言模型，大小从125M到6.7B不等。
- Llama-2: 在2万亿词元上训练，是当时性能最强的开源权重模型之一。
- SILO: 仅在宽松许可数据上训练的1.3B语言模型，面临极端领域泛化挑战。我们使用了PD、PDSW和PDSWBY三个变体。

表5：在时间漂移数据上的评估。评估数据来自2023年4月以来新增的维基百科文章，这在Pile和RedPajama创建之后。神经模型是Llama-2（13B），∞-gram参考数据是Pile + RPJ。
表5：在时间漂移数据上的评估。评估数据来自2023年4月以来新增的维基百科文章，这在Pile和RedPajama创建之后。神经模型是Llama-2（13B），∞-gram参考数据是Pile + RPJ。

D.2 在时间漂移数据上进行评估

泛化能力验证。为进一步展示∞-gram的有效性并消除性能增益可能源于净化不充分的疑虑，我们在时间漂移数据上进行评估：这些文档创建于∞-gram参考数据的截止时间之后。我们使用2023年4月和8月期间创建的新维基百科文章。

结果。表5报告了神经LM以及组合模型的困惑度。在五个月中的四个月的文档上，与∞-gram插值改善了神经LM的困惑度。我们发现，通过应用随机森林来决定实例级的插值超参数，可以进一步提升这种改善，其中随机森林的特征是后缀长度（从1到有效n）以及每个后缀在参考数据中的频率。应用随机森林后，困惑度改善范围为3% – 20%。

D.3 消融实验

参考数据大小的影响。图10报告了组合模型性能随参考数据大小的变化。我们看到，随着参考数据规模的增长，∞-gram带来的改善也随之扩大，并且这种关系大致是对数线性的。

参考数据领域的影响。图10还比较了∞-gram使用完整参考数据或仅使用领域内参考数据时的性能。仅使用领域内参考数据与使用完整参考数据的效果大致相当，这表明我们取得的几乎所有改进都归功于领域内数据（这些数据已经过净化）。这意味着使用完整的参考数据不会有坏处，特别是当测试领域未知或没有领域内参考数据时；然而，拥有领域内参考数据是最有帮助的。

图10：扩展∞-gram数据存储的影响，全部使用Llama-2模型（7B、13B和70B）作为神经LM，并以Pile作为参考数据。---：仅神经LM（基线）。•：∞-gram使用完整的Pile；◦：∞-gram仅使用Pile的领域内部分。随着数据存储规模的扩大，增益持续增加。
图10：扩展∞-gram数据存储的影响，全部使用Llama-2模型（7B、13B和70B）作为神经LM，并以Pile作为参考数据。---：仅神经LM（基线）。•：∞-gram使用完整的Pile；◦：∞-gram仅使用Pile的领域内部分。随着数据存储规模的扩大，增益持续增加。

G. 查询示例和Web界面

图11至图16展示了infini-gram支持的六种查询类型各一个示例。我们在这些示例中查询Dolma语料库。截图来自我们的Web界面。

PROTECTED_IMAGE_20__图12：查询类型2示例（NGRAMPROB）：从n-gram LM计算词元概率（给定n，无回退）。图13：查询类型3示例（NGRAMDIST）：从n-gram LM计算完整的下一词元分布。由于空间限制，仅显示前10个词元。图14：查询类型4示例（INFGRAMPROB）：从∞-gram LM计算词元概率。图15：查询类型5示例（INFGRAMDIST）：从∞-gram LM计算完整的下一词元分布。由于空间限制，仅显示前10个词元。__PROTECTED_IMAGE_25

A7 补充细节

6. 相关工作

n-gram语言模型。n-gram是自然语言处理诞生以来最经典的语言建模方法之一【20, Jurafsky & Martin, Speech and language processing, 2000】。人们一直通过扩大训练数据来推动n-gram LM的极限。迄今为止，最大的n-gram表【7, Brants et al., Large language models in machine translation, 2007】在一个2万亿词元的语料库中统计了5-gram。虽然n-gram LM目前在很大程度上被神经LM超越，但最近有工作重新审视了n-gram。Mikolov和Zweig【30, Mikolov & Zweig, Context dependent recurrent neural network language model, 2012】发现与5-gram Kneser-Ney模型插值可以改善RNN模型的困惑度，而Khandelwal等人【24, Khandelwal et al., Generalization through memorization: Nearest neighbor language models, 2020】则发现与Transformer插值n-gram模型并不能显著改善困惑度。Li等人【27, Li et al., Residual learning of neural text generation with n-gram language model, 2022】发现n-gram模型与小型神经LM相当有竞争力。然而，这些工作都使用了有限的参考数据和较小的神经LM。我们的工作将n-gram LM的训练数据扩展到万亿级词元，并通过扩展n值，显著改善了高达70B参数的SOTA神经模型。

无界n-gram、后缀数组、后缀树。之前的工作曾探索使用基于后缀的数据结构来实现无界n的n-gram查询，但数据规模有限。Stehouwer和van Zaanen【43, Stehouwer & van Zaanen, Using suffix arrays as language models, 2010】提出使用后缀数组实现∞-gram，但其公式不能产生有效的概率分布。Kennington等人【23, Kennington et al., Suffix trees as language models, 2012】提出使用后缀树，但其存储开销巨大。在这些工作中，只有Shareghi等人【40, Shareghi et al., Compact, efficient and unlimited capacity: Language modeling with compressed suffix trees, 2015】在通用语言建模任务上进行了评估，但困惑度数值过高。我们的训练数据比这些先前工作中最大的还要大500倍。

非参数语言模型。非参数LM指其复杂性不受先验限制，可根据参考数据变化的LM【24, Khandelwal et al., 2020; 6, Borgeaud et al., 2022; 3, Asai et al., 2023】。∞-gram LM是其中一个实例，其简单性使得用适度资源即可显著扩展参考数据。据我们所知，我们的∞-gram LM在参考数据大小（5万亿词元）和基础神经LM大小（70B）方面都是最大的。

E. 扩展讨论

infini-gram的潜在应用。我们相信infini-gram可以支持更广泛的研究和应用，包括但不限于：
- 理解海量文本语料库：使用COUNT查询快速了解语料库内容。
- 数据策管：使用SEARCHDOC查询检索并移除语料库中的有害内容（如毒性、PII）。
- 文档检索：作为检索增强LM的数据存储，利用其可扩展性提升性能。
- 减少事实性知识的幻觉：通过从训练数据中逐字读取来缓解幻觉。
- 检测数据污染、记忆和抄袭：检查评估查询是否泄露到训练数据中。
- 防止版权侵权：当模型即将生成受版权保护的长n-gram时，引导其走向其他路径。
- 衡量实体流行度：通过计算实体名称在海量语料中的出现次数来作为流行度的代理指标。
- 衡量文本的新颖性和创造性：通过生成文本与预训练语料的n-gram重叠度来定义新颖性。
- 归因：通过n-gram查找追溯影响模型决策的训练数据。
- 非参数推测性解码：利用infini-gram的低延迟，将∞-gram作为推测性解码中的快速解码器。

表6：与其他非参数语言建模方法的比较。# tokens：推理时参考数据中的词元数。# entries：索引中的表示（计数）数。max n：考虑的最大上下文词元数。对于infini-gram，我们将Pile-train和RedPajama的组合视为参考数据。
表6：与其他非参数语言建模方法的比较。# tokens：推理时参考数据中的词元数。# entries：索引中的表示（计数）数。max n：考虑的最大上下文词元数。对于infini-gram，我们将Pile-train和RedPajama的组合视为参考数据。

F. 扩展相关工作

超越n=5的n-gram。Mochihashi和Sumita【33, Mochihashi & Sumita, The infinite markov model, 2007】以及Wood等人【50, Wood et al., A stochastic memoizer for sequence data, 2009】考虑了无限阶马尔可夫模型，∞-gram LM是其一个特例。其他工作也发现了扩展n值的价值，例如KiloGrams【38, Raff et al., Kilograms: Very large n-grams for malware classification, 2019】使用1024-gram作为恶意软件检测的特征。

其他文本索引数据结构。除后缀数组和后缀树外，Koala【48, Vu et al., Koala: An index for quantifying overlaps with pre-training corpora, 2023】使用压缩后缀数组分析语料重叠。ROOTS搜索工具【36, Piktus et al., The roots search tool: Data transparency for llms, 2023】构建BM25索引。Data Portraits【29, Marone & Durme, Data portraits: Recording foundation model training data, 2023】提出基于布隆过滤器的轻量级索引。ElasticSearch也被用于在C4等语料库中搜索文档和计数n-gram【9, Dodge et al., 2021; 10, Elazar et al., 2023】。

非参数语言模型。先前的工作主要分为两种：一种是词元检索方法，另一种是块检索方法。在非参数LM中扩展参考数据非常昂贵。据我们所知，先前工作中拥有最大参考数据的是RETRO【6, Borgeaud et al., Improving language models by retrieving from trillions of tokens, 2022】，它使用了1.8万亿词元的参考数据，存储和搜索280亿个向量，估计消耗432TB的磁盘空间。（详细比较见表6。）

引用文献汇总