BigCode背后的大規(guī)模數(shù)據(jù)去重方法有哪些？

目標(biāo)受眾

本文面向?qū)Υ笠?guī)模文檔去重感興趣，且對散列 (hashing) 、圖 (graph) 及文本處理有一定了解的讀者。

動機

老話說得好: 垃圾進，垃圾出 (garbage in, garbage out)，把數(shù)據(jù)處理干凈再輸入給模型至關(guān)重要，至少對大語言模型如此。雖然現(xiàn)在一些明星大模型 (嚴(yán)格來講，它們很多是 API) 的存在讓大家恍惚產(chǎn)生了數(shù)據(jù)質(zhì)量好像不那么重要了的錯覺，但事實絕非如此。

在 BigScience 和 BigCode 項目中，在數(shù)據(jù)質(zhì)量方面，我們面臨的一個很大的問題是數(shù)據(jù)重復(fù)，這不僅包括訓(xùn)練集內(nèi)的數(shù)據(jù)重復(fù)，還包括訓(xùn)練集中包含測試基準(zhǔn)中的數(shù)據(jù)從而造成了基準(zhǔn)污染 (benchmark contamination)。已經(jīng)有研究表明，當(dāng)訓(xùn)練集中存在較多重復(fù)數(shù)據(jù)時，模型傾向于逐字輸出訓(xùn)練數(shù)據(jù) [1] (這一現(xiàn)象在其他一些領(lǐng)域并不常見 [2])，而且訓(xùn)得的模型也更容易遭受隱私攻擊 [1]。除了能避免上面兩個問題外，去重還有不少好處:

讓訓(xùn)練更高效: 你可以用更少的訓(xùn)練步驟獲得相同的，甚至是更好的性能 [3] [4]。

防止可能的數(shù)據(jù)泄漏和基準(zhǔn)污染: 數(shù)據(jù)重復(fù)會損害你的模型性能報告的公信力，并可能讓所謂的改進淪為泡影。

提高數(shù)據(jù)可得性。我們大多數(shù)人都負擔(dān)不起重復(fù)下載或傳輸數(shù)千 GB 文本的成本，更不用說由此帶來的額外訓(xùn)練成本了。對數(shù)據(jù)集進行去重，能使其更易于學(xué)習(xí)、傳輸及協(xié)作。

從 BigScience 到 BigCode

我想先分享一個故事，故事主要講述我如何接受數(shù)據(jù)去重這一任務(wù)，過程如何，以及在此過程中我學(xué)到了什么。

一切開始于 LinkedIn 上的一次對話，當(dāng)時 BigScience 已經(jīng)開始幾個月了。Huu Nguyen 注意到我在 GitHub 上的一個小項目并找到了我，問我是否有興趣為 BigScience 做數(shù)據(jù)去重工作。我當(dāng)然愿意了，盡管當(dāng)時我完全沒意識到由于數(shù)據(jù)量巨大，這項工作比想象中麻煩很多。

這項工作既有趣又充滿挑戰(zhàn)。挑戰(zhàn)在于，我對處理如此大規(guī)模的數(shù)據(jù)并沒有太多經(jīng)驗。但項目組的每個人仍然歡迎我、信任我，還給了我數(shù)千美元的云計算預(yù)算。有多少回，我不得不從睡夢中醒來，反復(fù)確認我是否關(guān)閉了那些云實例。我不停地在試驗和錯誤中學(xué)習(xí)，在此過程中，新的視角被打開了。如果沒有 BigScience，可能我永遠不會有這種視角。

一年后的今天，我正在把從 BigScience 學(xué)到的東西應(yīng)用到 BigCode 項目中去，去處理更大的數(shù)據(jù)集。除了英語 [3] LLM 之外，我們已經(jīng)再次證明數(shù)據(jù)去重也能改進代碼模型 [4] 的性能。有了數(shù)據(jù)去重，我們可以用更小的數(shù)據(jù)集達到更優(yōu)的性能?，F(xiàn)在，親愛的讀者，我想與你分享我學(xué)到的知識，希望你能透過數(shù)據(jù)去重的鏡頭一瞥 BigCode 項目的幕后故事。

下表列出了 BigScience 項目中各數(shù)據(jù)集使用的去重方法，以供參考:

下表是我們在創(chuàng)建 BigCode 的訓(xùn)練數(shù)據(jù)集 (訓(xùn)練數(shù)據(jù)皆為代碼) 時所用的方法。這里，如果當(dāng)遇到?jīng)]有名字的數(shù)據(jù)集時，我們就用模型名稱來代替。

例解 MinHash

在本節(jié)中，我們將詳細介紹在 BigCode 中使用的 MinHash 方法的每個步驟，并討論該方法的系統(tǒng)擴展性問題及其解決方案。我們以一個含有三個英文文檔為例來演示整個工作流程:

MinHash 的典型工作流程如下:

詞袋生成 (生成 n- 元組) 及指紋生成 (生成 MinHash): 將每個文檔映射成一組哈希值。

局部敏感哈希 (LSH): 逐條帶 (band) 的比較文檔的相似性，并將相似的文檔聚類以減少后續(xù)比較的次數(shù)。

去重: 決定保留或刪除哪些重復(fù)文檔。

詞袋生成

與大多數(shù)文本應(yīng)用一樣，我們需要先把文本表示成詞袋，這里我們通常使用 N- 元組詞袋。在本例中，我們使用以單詞為基本單元的 3- 元組 (即每 3 個連續(xù)單詞組成一個元組)，且不考慮標(biāo)點符號。我們后面會回過頭來討論元組大小對性能的影響。

這個操作的時間復(fù)雜度為 ，其中 表示文檔數(shù)，而 表示文檔長度。也就是說，時間復(fù)雜度與數(shù)據(jù)集大小呈線性關(guān)系。我們可以用多進程或分布式計算來并行化詞袋生成過程。

指紋計算

使用 MinHash 方法時，每個 N- 元組需要生成多個哈希值，此時我們通常要么 1) 使用不同的哈希函數(shù)進行多次哈希，要么 2) 使用一個哈希函數(shù)進行哈希后再進行多次重排。本例中，我們選擇第二種方法，重排生成 5 個哈希值。更多 MinHash 的變體可以參考 MinHash - 維基百科。

對以上文檔哈希矩陣中的每一列取最小值 —— 即 “MinHash” 中的 “Min” 的題中之義，我們就能得到該文檔最終的 MinHash 值:

從技術(shù)上講，雖然我們通常取最小值，但這并不代表我們一定要取每列的最小值。其他順序統(tǒng)計量也是可以的，例如最大值、第 k 個最小值或第 k 個最大值 [21]。

在具體實現(xiàn)時，我們可以使用 numpy 來對這些操作進行向量化。該操作的時間復(fù)雜度為 ，其中 是排列數(shù)。以下列出了我們的代碼，它是基于 Datasketch 的實現(xiàn)修改而得的。

defembed_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)fortinngrams(NON_ALPHA.split(content),ngram_size)}
hashvalues:np.ndarray=np.array([sha1_hash(token.encode("utf-8"))fortokenintokens],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)forstart,endinhashranges]
return{"__signatures__":Hs,"__id__":idx}

熟悉 Datasketch 的讀者可能會問，為什么我們要費心費力剝離 Datasketch 庫提供的所有高級功能？其主要原因并不是因為我們要減少依賴項，而是因為我們想要盡可能地榨取 CPU 的算力。而將多個步驟融合到一個函數(shù)中，是更好利用計算資源的手段之一。

由于每個文檔的計算互相獨立，因此我們可以充分利用 datasets 庫的 map 函數(shù)來實現(xiàn)并行化:

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...",
)

指紋計算完畢之后，每個文檔都被映射成了一個整數(shù)數(shù)組。為了弄清楚哪些文檔彼此相似，我們需要根據(jù)這些指紋對它們進行聚類。輪到 局部敏感哈希 (Locality Sensitive Hashing，LSH) 閃亮登場了。

局部敏感哈希 (LSH)

LSH 將指紋數(shù)組按行分成若干個條帶 (band)，每個條帶的行數(shù)相同，如果遇到最后一個條帶行數(shù)不足，我們就直接忽略它。以條帶數(shù) 為例，每個條帶有 行，具體組織如下:

若兩個文檔在某條帶上 MinHash 值相同，這兩個文檔就會被聚到同一個桶中備選。

遍歷 doc_ids 列的每一行，將其中的文檔兩兩配對就生成了候選對。上表中，我們能生成一個候選對: (0, 1) 。

候選對生成后 ……

很多數(shù)據(jù)去重的論文或教程講完上一節(jié)就結(jié)束了，但在實際項目中我們還涉及如何處理這些候選對的問題。通常，候選對生成后，我們有兩個選擇:

由于 MinHash 只是一個近似，所以仍需計算兩個文檔的 N- 元組集合的交并比來算得準(zhǔn)確的 Jaccard 相似性。此時，因為 LSH 已經(jīng)幫我們過濾了不少，所以最終參與計算的候選對的量會大大減少。在 BigCode 項目中，我們起初就采用了這種做法，效果相當(dāng)不錯。

我們還可以直接認可 LSH 選出來的相似對。這里面可能會有個問題: Jaccard 相似性不具傳遞性，也就是說 相似于 且 相似于 ，并不意味著 相似于 。所以這里可能會有不少假陽性。通過在 The Stack 數(shù)據(jù)集上的實驗，我們發(fā)現(xiàn)，直接認可 LSH 選出來的相似對在很大程度上能提高下游模型的性能，同時還節(jié)省了處理時間和訓(xùn)練時間。因此目前我們正慢慢開始轉(zhuǎn)向這種方法。但是，這個經(jīng)驗并不是放之四海而皆準(zhǔn)的，如果你準(zhǔn)備在自己的數(shù)據(jù)集上仿效我們的做法，我們建議你在此之前好好檢查你的數(shù)據(jù)集及其特點，然后作出數(shù)據(jù)驅(qū)動的決策。

最后，我們可以用生成的相似文本對構(gòu)建一個圖，在這個圖中，重復(fù)的文檔會被聚至同一個社區(qū)或同一個連通子圖中。不幸的是， datasets 在這方面幫不上什么忙，因為現(xiàn)在我們需要類似 groupby 的功能，以根據(jù) 條帶 ID 及 文檔在該條帶上的取值 對文檔進行聚類。下面列出了我們嘗試過的一些方案:

方案 1: 老辦法，迭代數(shù)據(jù)集以創(chuàng)建圖，然后用一個圖處理庫對其做社區(qū)檢測或者連通分量檢測。

我們測試下來，該方案的擴展性不怎么好，其原因是多方面的: 首先，整個數(shù)據(jù)集迭代起來很慢，而且內(nèi)存消耗很大; 其次，諸如 graphtool 或 networkx 的市面上流行的圖處理庫創(chuàng)建圖的開銷較大。

方案 2: 使用流行的 Python 框架 (如 dask ) 及其高效的 groupby 操作。

但迭代慢和創(chuàng)建圖慢的問題仍然存在。

方案 3: 迭代數(shù)據(jù)集并使用并查集 (union find data structure) 對文檔進行聚類。

這個方案引入了一個很小的迭代開銷，對中等數(shù)據(jù)集的有不錯的效果不錯，但在大數(shù)據(jù)集上還是慢。

fortableintqdm(HASH_TABLES,dynamic_ncols=True,desc="Clustering..."):
forclusterintable.values():
iflen(cluster)<=?1:
???continue
??idx?=?min(cluster)
??for?x?in?cluster:
???uf.union(x,?idx)

方案 4: 對大數(shù)據(jù)集，使用 Spark。

我們已經(jīng)知道到 LSH 的有些步驟是可以并行化的，我們可以用 Spark 來實現(xiàn)它們。Spark 的好處是，它開箱即支持分布式 groupBy ，而且也能很輕松地實現(xiàn)像 [18] 這樣的連通分量檢測算法。注意，這里我們并沒有使用 Spark 的原生 MinHash 實現(xiàn)，其原因是迄今為止我們所有的實驗都源于 Datasketch，而 Datasketch 的 MinHash 實現(xiàn)與 Spark 的原生實現(xiàn)完全不同。我們希望之前的經(jīng)驗和教訓(xùn)能幫助到后面的工作，而不是另起爐灶，進入另一個消融實驗的輪回，因此我們選擇在 Spark 中自己實現(xiàn) Datasketch 的 MinHash 算法。

edges=(
records.flatMap(
lambdax: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(lambdax:(x[0],x[1]))
.flatMap(lambdax:generate_edges([i[2]foriinx[1]]))
.distinct()
.cache()
)

以下是基于 [18] 的簡單連通分量檢測算法的 Spark 實現(xiàn)。

a=edges
whileTrue:
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()
iflen(changes)==0:
break

results=a.collect()

多虧了云計算提供商，我們可以使用 GCP DataProc 等服務(wù)輕松地搭建 一個 Spark 集群。最終，我們把程序運行起來，只用了不到 4 小時就完成了 1.4 TB 數(shù)據(jù)的去重工作，每小時僅需 15 美元。

數(shù)據(jù)質(zhì)量很重要

我們不可能爬著梯子登上月球。因此我們不僅要確保方向正確，還要確保方法正確。

早期，我們使用的參數(shù)主要來自 CodeParrot 的實驗，消融實驗表明這些參數(shù)確實提高了模型的下游性能 [16]。后來，我們開始沿著這條路進一步探索，由此進一步確認了以下結(jié)論 [4]:

數(shù)據(jù)去重可以在縮小數(shù)據(jù)集 (6 TB VS. 3 TB) 規(guī)模的同時提高模型的下游性能

雖然我們還沒有完全搞清楚其能力邊界及限制條件，但我們確實發(fā)現(xiàn)更激進的數(shù)據(jù)去重 (6 TB VS. 2.4 TB) 可以進一步提高性能，方法有:

降低相似度閾值

使用更長的元組 (如: 一元組 → 五元組)

放棄誤報檢查，承受一小部分誤報帶來的數(shù)據(jù)損失

1- 元組時不同設(shè)置影響的小提琴圖



5- 元組時不同設(shè)置影響的小提琴圖 圖例: 上述兩幅圖展示了相似性閾值和元組大小帶來的影響，第一幅圖使用 1- 元組，第二幅圖使用 5- 元組。紅色虛線表示相似性閾值: 低于該值的文檔與同一簇中其他文檔的相似性低于閾值，我們將其視為誤報。

上面兩幅圖可以幫助我們理解為什么有必要仔細檢查 CodeParrot 以及早期版本的 The Stack 訓(xùn)練數(shù)據(jù)上的誤報: 這是使用 1- 元組的誤報比例會很大; 上圖還表明，將元組大小增加到 5，誤報比例會顯著降低。如果想激進點去重的話，閾值可以設(shè)低點。

還有實驗表明，降低閾值會刪除更多包含部分相似內(nèi)容的文檔，因此意味著提高了我們最想刪除的那部分文檔的查全率。

系統(tǒng)擴展性

Scaling results for dataset size and deduplication time 圖例: 數(shù)據(jù)去重時間與原始數(shù)據(jù)集規(guī)模的關(guān)系。測試基于 GCP 上的 15 個 c2d-standard-16 實例，每個實例每小時的成本約為 0.7 美元。



CPU usage screenshot for the cluster during processing JSON dataset 圖例: 集群在處理 JSON 數(shù)據(jù)集時的 CPU 使用率。

上述擴展性數(shù)據(jù)未必非常嚴(yán)格，但也足夠說明，在給定預(yù)算的情況下，數(shù)據(jù)去重耗時與數(shù)據(jù)集規(guī)模的關(guān)系應(yīng)該是線性的。如果你仔細看一下處理 JSON 數(shù)據(jù)集 (The Stack 數(shù)據(jù)集的最大子集) 的集群資源使用情況，你會發(fā)現(xiàn)實際總計算時間 (圖中第 2 和第 3 階段) 主要都花在了 MinHash + LSH (圖中第 2 階段) 上，這與我們先前的分析一致，即第 2 階段 d 的時間復(fù)雜度為 — 與數(shù)據(jù)體量成線性關(guān)系。

謹(jǐn)慎行事

數(shù)據(jù)去完重并不意味著萬事大吉了，你仍然需要對數(shù)據(jù)進行徹底的探索和分析。此外，上文這些有關(guān)數(shù)據(jù)去重的發(fā)現(xiàn)來自于 The Stack 數(shù)據(jù)集，并不意味著它能無腦適用于其他數(shù)據(jù)集或語言。要構(gòu)建一個好的訓(xùn)練數(shù)據(jù)集，我們僅僅邁出了萬里長征的第一步，后面還有很多工作要做，例如數(shù)據(jù)質(zhì)量過濾 (如過濾漏洞數(shù)據(jù)、毒性數(shù)據(jù)、偏見數(shù)據(jù)、模板生成的數(shù)據(jù)、個人身份數(shù)據(jù)等)。

我們還鼓勵你在訓(xùn)練前像我們一樣對數(shù)據(jù)集進行徹底的分析，因為大家的情況可能各不相同。例如，如果你的時間和計算預(yù)算都很緊張，那么數(shù)據(jù)去重可能不是很有幫助: @geiping_2022 提到基于子字符串的數(shù)據(jù)去重并沒有提高他們模型的下游性能。在使用前，可能還需要對現(xiàn)存數(shù)據(jù)集進行徹底檢查，例如，@gao_2020 聲明他們只確保 Pile 本身及其子集都已去重，但不保證其與任何下游基準(zhǔn)數(shù)據(jù)集沒有重復(fù)，要不要對 Pile 與下游基準(zhǔn)數(shù)據(jù)集進行去重取決于使用者自己。

在數(shù)據(jù)泄露和基準(zhǔn)污染方面，還有很多需要探索的地方。由于 HumanEval 也是 GitHub Python 存儲庫之一，我們不得不重新訓(xùn)練了我們的代碼模型。早期的工作還發(fā)現(xiàn)，最流行的編碼基準(zhǔn)之一的 MBPP[19] 與許多 Leetcode 問題有很多相似之處 (例如，MBPP 中的任務(wù) 601 基本上是 Leetcode 646，任務(wù) 604 ? Leetcode 151)。我們都知道 GitHub 中不乏很多編程挑戰(zhàn)賽題及其答案代碼。如果居心叵測的人把所有基準(zhǔn)測試的 Python 代碼以不易察覺的方式上傳到 Github，污染你所有的訓(xùn)練數(shù)據(jù)，這事兒就更難了。

后續(xù)方向

子串去重。盡管在英語 [3] 上子串去重是有益的，但尚不清楚是否對代碼數(shù)據(jù)也有用;

重復(fù)段落: 在一篇文檔中重復(fù)多次的段落。@rae_2021 分享了一些關(guān)于如何檢測和刪除它們的有趣的啟發(fā)式方法。

使用模型嵌入進行語義級的去重。這是另外一套思路了，需要一整套去重、擴展性、成本、銷蝕等各方面的實驗和權(quán)衡。對此 [20] 提出了一些有趣的看法，但我們?nèi)匀恍枰鄬嶋H證據(jù)才能得出結(jié)論 (其文本去重工作僅參考了 @lee_2022a 的工作，而 @lee_2022a 的主張主要是去重有作用而并未證明其效果達到了 SOTA)。

優(yōu)化。還有不少優(yōu)化空間: 更好的質(zhì)量評估標(biāo)準(zhǔn)、擴展性、對下游性能影響的分析等。

換個角度: 對相似數(shù)據(jù)，去重到什么程度就會開始損害性能？需要保留多少相似數(shù)據(jù)以保留數(shù)據(jù)的多樣性又不至冗余？

審核編輯：劉清