BigCode 背后的通用近似去重技术

目标受众

对大规模文档级近似去重感兴趣，并对哈希、图和文本处理有一定了解的人员。

动机

在将数据输入模型（至少在我们的例子中是大型语言模型）之前，妥善处理数据非常重要，正如老话所说，垃圾进，垃圾出。尽管现在越来越难以做到这一点，因为那些引人注目的模型（或者我们应该说是 API）正在制造一种数据质量不那么重要的错觉。

我们在 BigScience 和 BigCode 中都面临的一个数据质量问题是重复，包括可能的基准污染。研究表明，当存在大量重复数据时，模型倾向于逐字输出训练数据[1]（尽管在其他一些领域不那么明显[2]），这也使模型容易受到隐私攻击[1]。此外，去重的一些典型优点还包括：

1. 高效训练：您可以通过更少的训练步骤实现相同甚至更好的性能[3] [4]。
2. 防止可能的数据泄露和基准污染：非零重复数据会降低您的评估可信度，并可能使所谓的改进成为虚假声明。
3. 可访问性。我们大多数人都无法反复下载或传输数千 GB 的文本，更不用说用它来训练模型了。对于固定大小的数据集，去重使其更容易研究、传输和协作。

从 BigScience 到 BigCode

请允许我先分享一个故事，讲述我是如何踏上这次近似去重之旅的，结果是如何进展的，以及我一路学到了什么经验。

这一切都始于 LinkedIn 上的一次对话，当时 BigScience 已经开始了几个月。Huu Nguyen 注意到我在 GitHub 上的个人项目，于是联系我，问我是否有兴趣为 BigScience 做去重工作。当然，我的回答是肯定的，完全没有意识到仅仅是数据量的庞大就需要付出多少努力。

这既有趣又充满挑战。说它挑战，是因为我以前没有处理如此大规模数据的研究经验，而且每个人都非常欢迎我，并信任我使用数千美元的云计算预算。是的，我好几次在睡梦中醒来，只为再次检查是否关掉了那些机器。因此，我不得不通过反复试验在工作中学习，最终这为我打开了一个全新的视角，如果没有 BigScience，我想我永远不会有这样的机会。

快进一年后，我将所学到的知识应用到 BigCode 中，处理更大的数据集。除了用于英语的 LLM[3]，我们还确认去重也能改进代码模型[4]，而且使用的数据集要小得多。现在，我将我所学到的知识分享给亲爱的读者，希望您也能通过去重的视角，了解 BigCode 幕后正在发生的事情。

如果您感兴趣，这里是我们在 BigScience 中开始的去重比较的更新版本

这是我们为 BigCode 创建的代码数据集。当数据集名称不可用时，使用模型名称。

MinHash + LSH 参数 $(P, T, K, B, R)$

1. $P$ 排列/哈希的数量
2. $T$ Jaccard 相似度阈值
3. $K$ n-gram/shingle 大小
4. $B$ band 数量
5. $R$ row 数量

为了了解这些参数可能如何影响您的结果，这里有一个简单的演示，以数学方式说明计算：MinHash 数学演示。

MinHash 详解

在本节中，我们将介绍 MinHash 的每个步骤（BigCode 中使用的方法），以及潜在的扩展问题和解决方案。我们将通过一个包含三篇英文文档的示例来演示工作流程

MinHash 的典型工作流程如下：

1. 切分（分词）和指纹（MinHashing），我们将每篇文档映射到一组哈希值。
2. 局部敏感哈希（LSH）。此步骤通过将具有相似 band 的文档分组在一起，来减少比较次数。
3. 重复删除。此步骤是我们决定保留或删除哪些重复文档的地方。

字符串片段（Shingles）

与大多数涉及文本的应用程序一样，我们需要从分词开始。N-gram，又名 shingle，常被使用。在我们的例子中，我们将使用词级别的三元组，不带任何标点符号。我们将在后面章节中回顾 n-gram 大小如何影响性能。

此操作的时间复杂度为 $\mathcal{O}(NM)$，其中 $N$ 是文档数量，$M$ 是文档长度。换句话说，它线性依赖于数据集的大小。此步骤可以通过多进程或分布式计算轻松并行化。

指纹计算

在 MinHash 中，每个 shingle 通常会经历以下两种操作之一：1) 使用不同的哈希函数多次哈希；或 2) 使用一个哈希函数多次排列。在这里，我们选择对每个哈希进行 5 次排列。MinHash 的更多变体可以在 MinHash - Wikipedia 中找到。

取每个文档中每列的最小值——“MinHash”中的“Min”部分，我们得到该文档的最终 MinHash

严格来说，我们不必使用每列的最小值，但最小值是最常见的选择。也可以使用其他顺序统计量，例如最大值、第 k 小值或第 k 大值[21]。

在实现中，您可以使用 `numpy` 轻松地向量化这些步骤，并期望时间复杂度为 $\mathcal{O}(NMK)$，其中 $K$ 是您的排列次数。代码修改自 Datasketch。

def embed_func(
    content: str,
    idx: int,
    *,
    num_perm: int,
    ngram_size: int,
    hashranges: List[Tuple[int, int]],
    permutations: np.ndarray,
) -> Dict[str, Any]:
    a, b = permutations
    masks: np.ndarray = np.full(shape=num_perm, dtype=np.uint64, fill_value=MAX_HASH)
    tokens: Set[str] = {" ".join(t) for t in ngrams(NON_ALPHA.split(content), ngram_size)}
    hashvalues: np.ndarray = np.array([sha1_hash(token.encode("utf-8")) for token in tokens], dtype=np.uint64)
    permuted_hashvalues = np.bitwise_and(
        ((hashvalues * np.tile(a, (len(hashvalues), 1)).T).T + b) % MERSENNE_PRIME, MAX_HASH
    )
    hashvalues = np.vstack([permuted_hashvalues, masks]).min(axis=0)
    Hs = [bytes(hashvalues[start:end].byteswap().data) for start, end in hashranges]
    return {"__signatures__": Hs, "__id__": idx}

如果您熟悉 Datasketch，您可能会问，为什么我们还要费心剥离该库提供的所有高级功能呢？这并不是因为我们想避免添加依赖项，而是因为我们打算在并行化过程中尽可能地压榨 CPU 计算。将几个步骤融合到一个函数调用中，可以让我们更好地利用计算资源。

由于一个文档的计算不依赖于其他任何东西。一个好的并行化选择是使用 `datasets` 库的 `map` 函数

embedded = ds.map(
    function=embed_func,
    fn_kwargs={
        "num_perm": args.num_perm,
        "hashranges": HASH_RANGES,
        "ngram_size": args.ngram,
        "permutations": PERMUTATIONS,
    },
    input_columns=[args.column],
    remove_columns=ds.column_names,
    num_proc=os.cpu_count(),
    with_indices=True,
    desc="Fingerprinting...",
)

指纹计算完成后，一个特定的文档被映射到一个整数值数组。为了找出哪些文档彼此相似，我们需要根据这些指纹对它们进行分组。这就是 **局部敏感哈希 (LSH)** 登场的时候了。

局部敏感哈希（Locality Sensitive Hashing）

LSH 将指纹数组分成多个 band，每个 band 包含相同数量的行。如果剩下任何哈希值，它们将被忽略。让我们使用 $b=2$ 个 band 和 $r=2$ 行来对这些文档进行分组

如果两个文档在特定位置（band 索引）的 band 中共享相同的哈希值，它们将被聚类到同一个 bucket 中，并被视为候选对象。

对于 `doc_ids` 列中的每一行，我们可以通过将每两个文档配对来生成候选对。从上表中，我们可以生成一个候选对：`(0, 1)`。

超越重复对

这是许多论文或教程中去重描述止步的地方。我们仍然面临如何处理它们的问题。通常，我们可以采取两种选择：

1. 由于 MinHash 的估计性质，通过计算它们的 shingle 重叠来双重检查它们实际的 Jaccard 相似度。两个集合的 Jaccard 相似度定义为交集的大小除以并集的大小。现在这比计算所有对的相似度更容易实现，因为我们可以只关注集群内的文档。这也是我们最初为 BigCode 所做的工作，效果相当不错。
2. 将它们视为真阳性。您可能已经注意到了这里的问题：Jaccard 相似度不具有传递性，这意味着 $A$ 与 $B$ 相似，$B$ 与 $C$ 相似，但 $A$ 和 $C$ 不一定共享相似性。然而，我们从 The Stack 实验中发现，将所有这些都视为重复会最好地提高下游模型的性能。现在我们逐渐转向这种方法，它也节省了时间。但是，要将其应用于您自己的数据集，我们仍然建议您仔细检查数据集并查看重复项，然后做出数据驱动的决策。

无论这些对是否经过验证，我们现在都可以用这些对作为边构建一个图，重复项将被聚类到社区或连接组件中。在实现方面，不幸的是，`datasets` 在这里帮不上什么忙，因为我们现在需要像 `groupby` 这样的功能，可以根据文档的*band 偏移*和*band 值*进行聚类。以下是我们尝试过的一些选项：

选项 1：以老式方式遍历数据集并收集边。然后使用图库进行社区检测或连通分量检测。

这在我们的测试中扩展性不佳，原因有很多。首先，大规模遍历整个数据集速度慢且占用内存。其次，像 `graphtool` 或 `networkx` 这样的流行图库在图创建方面有很大的开销。

选项 2：使用流行的 Python 框架（如 `dask`）以实现更高效的 `groupby` 操作.

但您仍然面临迭代速度慢和图创建速度慢的问题。

选项 3：迭代数据集，但使用并查集数据结构对文档进行聚类。

这为迭代增加了可忽略不计的开销，并且对于中型数据集来说效果相对较好。

for table in tqdm(HASH_TABLES, dynamic_ncols=True, desc="Clustering..."):
    for cluster in table.values():
        if len(cluster) <= 1:
            continue
        idx = min(cluster)
        for x in cluster:
            uf.union(x, idx)

选项 4：对于大型数据集，使用 Spark。

我们已经知道，MinHash 之前的步骤可以并行化，这在 Spark 中也可以实现。除此之外，Spark 本身就支持分布式 `groupBy`，并且也很容易实现像 [18] 这样的连通分量检测算法。如果您想知道为什么我们没有使用 Spark 的 MinHash 实现，答案是我们目前所有的实验都源于 Datasketch，它的实现与 Spark 完全不同，我们希望确保我们能从中吸取经验和见解，而不是陷入另一个消融实验的无底洞。

edges = (
    records.flatMap(
        lambda x: generate_hash_values(
            content=x[1],
            idx=x[0],
            num_perm=args.num_perm,
            ngram_size=args.ngram_size,
            hashranges=HASH_RANGES,
            permutations=PERMUTATIONS,
        )
    )
    .groupBy(lambda x: (x[0], x[1]))
    .flatMap(lambda x: generate_edges([i[2] for i in x[1]]))
    .distinct()
    .cache()
)

一个基于 [18] 在 Spark 中实现的简单连通分量算法。

a = edges
while True:
    b = a.flatMap(large_star_map).groupByKey().flatMap(large_star_reduce).distinct().cache()
    a = b.map(small_star_map).groupByKey().flatMap(small_star_reduce).distinct().cache()
    changes = a.subtract(b).union(b.subtract(a)).collect()
    if len(changes) == 0:
        break

results = a.collect()

此外，感谢云服务提供商，我们可以通过 GCP DataProc 等服务轻松搭建 Spark 集群。**最终，我们能够以每小时 15 美元的预算，在不到 4 小时内舒适地运行程序，对 1.4 TB 的数据进行去重。**

质量至上

爬梯子并不能带我们上月球。这就是为什么我们需要确保方向正确，并以正确的方式使用它。

早期，我们的参数主要继承自 CodeParrot 实验，我们的消融实验表明，这些设置确实改善了模型的下游性能[16]。我们随后进一步探索这条路径，并可以确认[4]

1. 近似去重用小得多的数据集（6 TB 对比 3 TB）提高了模型的下游性能
2. 我们还没有找出极限，但更积极的去重（6 TB 对比 2.4 TB）可以进一步提高性能
   1. 降低相似度阈值
   2. 增加 shingle 大小（unigram → 5-gram）
   3. 放弃误报检查，因为我们可以承受少量误报的损失

图片：两张图展示了相似度阈值和 shingle 大小的影响，第一张图使用 unigram，第二张图使用 5-gram。红色虚线表示相似度截止线：低于该线的任何文档都将被视为误报——它们与集群中其他文档的相似度低于阈值。

这些图可以帮助我们理解为什么 CodeParrot 和早期版本的 Stack 需要仔细检查误报：使用 unigram 会产生许多误报；它们还表明，通过将 shingle 大小增加到 5-gram，误报的百分比显著降低。如果我们要保持去重的激进性，则需要更小的阈值。

额外实验也表明，降低阈值会移除更多具有高相似度对的文档，这意味着在我们最希望移除的部分中召回率有所提高。

扩展性

图片：去重时间与原始数据集大小的关系。这是在 GCP 上使用 15 台 c2d-standard-16 工作机器实现的，每台机器每小时成本约 0.7 美元。

图片：处理 JSON 数据集期间集群的 CPU 使用情况截图。

这并不是最严格的扩展性证明，但在固定计算预算下，去重时间看起来与数据集的物理大小呈线性关系。如果您仔细观察处理 JSON 数据集（Stack 中最大的子集）时集群的资源使用情况，您会发现 MinHash + LSH（阶段 2）占据了总实际计算时间（阶段 2 + 3）的主导地位，根据我们之前的分析，其时间复杂度为 $\mathcal{O}(NM)$，即与数据集的物理体积呈线性关系。

谨慎行事

去重并不能免除您进行彻底的数据探索和分析。此外，这些去重发现对于 Stack 来说是成立的，但这并不意味着它能轻易适用于其他数据集或语言。它是构建更好数据集的第一步，但仍需进一步调查数据质量过滤（例如，脆弱性、毒性、偏见、生成模板、PII）等问题。

我们仍然鼓励您在训练之前对数据集进行类似分析。例如，如果您时间紧迫且计算预算有限，进行去重可能作用不大：@geiping_2022 提到子串去重并未改善其模型的下游性能。现有数据集在使用前也可能需要彻底检查，例如，@gao_2020 指出他们只确保了 Pile 本身及其拆分是去重的，他们不会主动对任何下游基准进行去重，并将该决定留给读者。

在数据泄露和基准污染方面，仍有许多值得探索的地方。我们不得不重新训练我们的代码模型，因为 HumanEval 发布在一个 GitHub Python 仓库中。早期的近似去重结果也表明，MBPP[19]（最受欢迎的编码基准之一）与许多 Leetcode 问题（例如，MBPP 中的任务 601 基本上是 Leetcode 646，任务 604 ≈ Leetcode 151）存在大量相似性。我们都知道 GitHub 上不乏这些编码挑战和解决方案。如果有人恶意上传所有基准，以 Python 脚本或其他不那么明显的方式，并污染您的所有训练数据，那将更加困难。

未来方向

1. 子串去重。尽管它对英语显示出了一些好处[3]，但目前尚不清楚这是否也应应用于代码数据；
2. 重复：一个文档中多次重复的段落。@rae_2021 分享了一些检测和删除它们的有趣启发式方法。
3. 使用模型嵌入进行语义去重。这是一个全新的研究问题，涉及扩展、成本、消融实验以及与近似去重的权衡。对此有一些有趣的看法[20]，但我们仍需要更多具体证据才能得出结论（例如，@abbas_2023 唯一的文本去重参考文献是 @lee_2022a，其主要主张是去重有帮助，而不是试图达到 SOTA）。
4. 优化。总有优化的空间：更好的质量评估、扩展、下游性能影响分析等。
5. 然后还有另一个看待问题的方向：近似去重在何种程度上开始损害性能？在何种程度上需要相似性来增加多样性而不是被视为冗余？

致谢

标题图片包含来自 Noto Emoji 的表情符号（拥抱脸、圣诞老人、文档、巫师和魔杖）（Apache 2.0）。这篇博文完全由人工撰写，未使用任何生成式 API。

非常感谢 Huu Nguyen @Huu 和 Hugo Laurençon @HugoLaurencon 在 BigScience 中的合作，以及 BigCode 的每位成员在整个过程中的帮助！如果您发现任何错误，请随时联系我：mouchenghao at gmail dot com。

支持资源

• Datasketch (MIT)
• simhash-py 和 simhash-cpp (MIT)
• 去重训练数据让语言模型更好 (Apache 2.0)
• Gaoya (MIT)
• BigScience (Apache 2.0)
• BigCode (Apache 2.0)

参考文献

• [1]：Nikhil Kandpal、Eric Wallace、Colin Raffel，去重训练数据可降低语言模型中的隐私风险，2022
• [2]：Gowthami Somepalli 等人，扩散艺术还是数字伪造？调查扩散模型中的数据复制，2022
• [3]：Katherine Lee, Daphne Ippolito, et al., Deduplicating Training Data Makes Language Models Better, 2022
• [4]：Loubna Ben Allal, Raymond Li, et al., SantaCoder: Don't reach for the stars!, 2023
• [5]：Leo Gao、Stella Biderman 等人，The Pile：一个 800GB 的多样化文本数据集，用于语言建模，2020
• [6]：Asier Gutiérrez-Fandiño、Jordi Armengol-Estapé 等人，MarIA：西班牙语语言模型，2022
• [7]：Jack W. Rae, Sebastian Borgeaud, et al., Scaling Language Models: Methods, Analysis & Insights from Training Gopher, 2021
• [8]：Xi Victoria Lin、Todor Mihaylov 等人，使用多语言语言模型进行少样本学习，2021
• [9]：Hugo Laurençon、Lucile Saulnier 等人，The BigScience ROOTS Corpus：一个 1.6TB 的复合多语言数据集，2022
• [10]：Daniel Fried、Armen Aghajanyan 等人，InCoder：用于代码填充和合成的生成模型，2022
• [11]：Erik Nijkamp、Bo Pang 等人，CodeGen：用于多轮程序合成的开源大型代码语言模型，2023
• [12]：Yujia Li, David Choi, et al., Competition-Level Code Generation with AlphaCode, 2022
• [13]：Frank F. Xu, Uri Alon, et al., A Systematic Evaluation of Large Language Models of Code, 2022
• [14]：Aakanksha Chowdhery、Sharan Narang 等人，PaLM：通过 Pathways 扩展语言建模，2022
• [15]：Lewis Tunstall, Leandro von Werra, Thomas Wolf, Natural Language Processing with Transformers, Revised Edition, 2022
• [16]：Denis Kocetkov、Raymond Li 等人，The Stack：3 TB 的宽松许可源代码，2022
• [17]：洛奇 | 祝福号维基 | Fandom
• [18]：Raimondas Kiveris、Silvio Lattanzi 等人，MapReduce 及更远领域的连通分量，2014
• [19]：Jacob Austin、Augustus Odena 等人，使用大型语言模型进行程序合成，2021
• [20]：Amro Abbas, Kushal Tirumala, et al., SemDeDup: Data-efficient learning at web-scale through semantic deduplication, 2023
• [21]：Edith Cohen，MinHash Sketches：简要调查，2016