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