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

目標(biāo)受眾

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

動(dòng)機(jī)

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

在 BigScience 和 BigCode 項(xiàng)目中，在數(shù)據(jù)質(zhì)量方面，我們面臨的一個(gè)很大的問題是數(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ù)時(shí)，模型傾向于逐字輸出訓(xùn)練數(shù)據(jù) [1] (這一現(xiàn)象在其他一些領(lǐng)域并不常見 [2])，而且訓(xùn)得的模型也更容易遭受隱私攻擊 [1]。除了能避免上面兩個(gè)問題外，去重還有不少好處:

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

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

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

從 BigScience 到 BigCode

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

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

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

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

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

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

例解 MinHash

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

MinHash 的典型工作流程如下:

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

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

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

詞袋生成

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

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

指紋計(jì)算

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

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

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

在具體實(shí)現(xiàn)時(shí)，我們可以使用 numpy 來對(duì)這些操作進(jìn)行向量化。該操作的時(shí)間復(fù)雜度為 ，其中 是排列數(shù)。以下列出了我們的代碼，它是基于 Datasketch 的實(shí)現(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 的讀者可能會(huì)問，為什么我們要費(fèi)心費(fèi)力剝離 Datasketch 庫提供的所有高級(jí)功能？其主要原因并不是因?yàn)槲覀円獪p少依賴項(xiàng)，而是因?yàn)槲覀兿胍M可能地榨取 CPU 的算力。而將多個(gè)步驟融合到一個(gè)函數(shù)中，是更好利用計(jì)算資源的手段之一。

由于每個(gè)文檔的計(jì)算互相獨(dú)立，因此我們可以充分利用 datasets 庫的 map 函數(shù)來實(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...",
)

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

局部敏感哈希 (LSH)

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

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

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

候選對(duì)生成后 ……

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

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

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

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

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

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

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

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

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

這個(gè)方案引入了一個(gè)很小的迭代開銷，對(duì)中等數(shù)據(jù)集的有不錯(cuò)的效果不錯(cuò)，但在大數(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: 對(duì)大數(shù)據(jù)集，使用 Spark。

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

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

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

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

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

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

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

降低相似度閾值

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

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

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



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

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

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

系統(tǒng)擴(kuò)展性

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



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

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

謹(jǐn)慎行事

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

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

在數(shù)據(jù)泄露和基準(zhǔn)污染方面，還有很多需要探索的地方。由于 HumanEval 也是 GitHub Python 存儲(chǔ)庫之一，我們不得不重新訓(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] 上子串去重是有益的，但尚不清楚是否對(duì)代碼數(shù)據(jù)也有用;

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

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

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

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

審核編輯：劉清