預(yù)訓(xùn)練數(shù)據(jù)集通常從網(wǎng)絡(luò)內(nèi)容收集，缺乏固有的領(lǐng)域劃分。例如，廣泛使用的Common Crawl等數(shù)據(jù)集不包含明確的領(lǐng)域標(biāo)簽，而手動策劃像The Pile這樣的標(biāo)注數(shù)據(jù)集則非常耗費(fèi)人力。因此，盡管識別最優(yōu)預(yù)訓(xùn)練數(shù)據(jù)混合對預(yù)訓(xùn)練性能有顯著益處，但這仍然是一個具有挑戰(zhàn)性的問題。

NVIDIA和多所大學(xué)的研究團(tuán)隊(duì)推出了CLIMB（CLustering-based Iterative Data Mixture Bootstrapping），這是一個自動化框架，能夠在預(yù)訓(xùn)練環(huán)境中發(fā)現(xiàn)、評估和優(yōu)化數(shù)據(jù)混合配方。

圖片

預(yù)訓(xùn)練數(shù)據(jù)混合的挑戰(zhàn)

預(yù)訓(xùn)練數(shù)據(jù)集已擴(kuò)展到數(shù)萬億token，通常結(jié)合大規(guī)模網(wǎng)絡(luò)爬取和較小的高質(zhì)量領(lǐng)域特定數(shù)據(jù)集。這些語料庫使得能夠開發(fā)處理多樣化任務(wù)的通用模型。然而，其巨大的規(guī)模和異質(zhì)性在平衡通用知識與領(lǐng)域?qū)I(yè)知識方面帶來了挑戰(zhàn)。

現(xiàn)有問題：

缺乏領(lǐng)域標(biāo)簽：像Common Crawl這樣的大規(guī)模數(shù)據(jù)集提供了無與倫比的多樣性和規(guī)模，但缺乏明確的領(lǐng)域標(biāo)簽，難以提取領(lǐng)域相關(guān)內(nèi)容。數(shù)據(jù)過濾通常依賴于困惑度或教育價值等通用啟發(fā)式方法，這些方法不一定能捕獲特定領(lǐng)域最有信息量或最高質(zhì)量的內(nèi)容。

混合優(yōu)化困難：即使有像The Pile這樣帶有領(lǐng)域注釋的精選數(shù)據(jù)集，選擇最優(yōu)數(shù)據(jù)混合也并非易事，因?yàn)閿?shù)據(jù)集組成與模型性能之間存在復(fù)雜的非線性關(guān)系。例如，為編碼任務(wù)優(yōu)化模型不僅需要編程相關(guān)內(nèi)容，還需要來自數(shù)學(xué)、推理和安全等領(lǐng)域的互補(bǔ)知識。

計(jì)算成本高昂：傳統(tǒng)的網(wǎng)格搜索或隨機(jī)搜索方法需要訓(xùn)練大量模型來評估不同的數(shù)據(jù)混合，計(jì)算成本極高。

CLIMB框架

CLIMB提出了一個新穎的框架，用于在預(yù)訓(xùn)練期間自動搜索最優(yōu)數(shù)據(jù)混合。框架包含三個關(guān)鍵步驟：

圖片

1. 嵌入和聚類

數(shù)據(jù)嵌入：使用預(yù)訓(xùn)練的語言模型（如BERT）將大規(guī)模數(shù)據(jù)集中的文檔嵌入到語義空間中。每個文檔被表示為一個高維向量，捕獲其語義內(nèi)容。

聚類：在嵌入空間中對文檔進(jìn)行聚類，將語義相似的文檔分組。CLIMB使用K-means聚類算法，將1.2萬億token的數(shù)據(jù)集聚類為20個簇。

每個簇代表一個語義主題或領(lǐng)域。例如：

?C0：科學(xué)和技術(shù)

?C1：社會科學(xué)

?C2：編程和代碼

?C3：數(shù)學(xué)

?C4：醫(yī)學(xué)和健康

?...

這種聚類方法的優(yōu)勢在于：

?自動化：無需人工標(biāo)注領(lǐng)域標(biāo)簽

?語義一致性：同一簇內(nèi)的文檔語義相關(guān)

?可擴(kuò)展性：可以處理數(shù)萬億token的數(shù)據(jù)

2. 混合-性能對構(gòu)建

采樣數(shù)據(jù)混合：從20個簇中采樣不同的數(shù)據(jù)混合配置。每個配置指定從每個簇中采樣的token比例。

訓(xùn)練代理模型：對每個數(shù)據(jù)混合配置，訓(xùn)練一個小型代理模型（proxy model）。CLIMB使用350M參數(shù)的模型作為代理，而不是直接訓(xùn)練大型目標(biāo)模型（如1B參數(shù)）。

評估性能：在下游任務(wù)上評估代理模型的性能，如MMLU、ARC、HellaSwag等。

構(gòu)建數(shù)據(jù)集：將數(shù)據(jù)混合配置作為輸入特征，性能指標(biāo)作為目標(biāo)標(biāo)簽，構(gòu)建混合-性能對數(shù)據(jù)集。

3. 預(yù)測器訓(xùn)練

回歸模型：訓(xùn)練一個回歸模型作為預(yù)測器，學(xué)習(xí)數(shù)據(jù)混合配置與性能之間的關(guān)系。

迭代優(yōu)化：使用預(yù)測器指導(dǎo)下一輪的數(shù)據(jù)混合采樣，逐步優(yōu)化搜索空間。

CLIMB采用自舉策略（bootstrapping）：在每次迭代中，候選混合被提出、修剪和優(yōu)化，以優(yōu)化多樣性和領(lǐng)域相關(guān)性。與靜態(tài)混合策略不同，該方法使用弱預(yù)測器方法在整個訓(xùn)練過程中動態(tài)調(diào)整數(shù)據(jù)混合，迭代集成多個預(yù)測器以發(fā)現(xiàn)領(lǐng)域適應(yīng)的有效配置。

圖片

核心創(chuàng)新

極致的計(jì)算效率：CLIMB優(yōu)先考慮計(jì)算效率，證明了在固定訓(xùn)練預(yù)算內(nèi)，迭代數(shù)據(jù)混合搜索能夠取得優(yōu)異結(jié)果。

代理模型策略：使用350M參數(shù)的小型代理模型評估混合質(zhì)量，而不是直接訓(xùn)練1B參數(shù)的目標(biāo)模型。這將每次評估的計(jì)算成本降低了約3倍。

漸進(jìn)式修剪：在每次迭代中，基于預(yù)測器的輸出修剪搜索空間，只保留最有希望的候選配置。這顯著減少了需要評估的配置數(shù)量。

迭代優(yōu)化：通過3-5次迭代逐步優(yōu)化數(shù)據(jù)混合，而不是一次性搜索整個空間。每次迭代都基于前一次的結(jié)果，使搜索更加高效。

實(shí)驗(yàn)表明，CLIMB的總計(jì)算成本約為傳統(tǒng)網(wǎng)格搜索的10-20%，同時取得了更好的結(jié)果。

自適應(yīng)學(xué)習(xí)：CLIMB主動學(xué)習(xí)根據(jù)環(huán)境驗(yàn)證的真實(shí)反饋來優(yōu)化和優(yōu)化數(shù)據(jù)混合，而不是被動依賴預(yù)定義的啟發(fā)式或人工標(biāo)注的領(lǐng)域標(biāo)簽。這種迭代自我改進(jìn)的能力使CLIMB更加靈活，能夠適應(yīng)新的數(shù)據(jù)分布和領(lǐng)域特定需求。

可解釋性：通過分析最終的數(shù)據(jù)混合配置，CLIMB能夠揭示最優(yōu)數(shù)據(jù)混合的特征，為理解不同領(lǐng)域的數(shù)據(jù)需求提供洞察。

圖片

實(shí)驗(yàn)結(jié)果

通用推理任務(wù)

任務(wù)：在12個通用推理基準(zhǔn)上評估模型性能，包括PIQA、ARC-Challenge、ARC-Easy、HellaSwag、Winogrande、SIQA等。

訓(xùn)練設(shè)置：使用CLIMB發(fā)現(xiàn)的最優(yōu)數(shù)據(jù)混合，在400B token上連續(xù)訓(xùn)練1B參數(shù)模型。

結(jié)果：

?超越Llama-3.2-1B：CLIMB訓(xùn)練的模型在平均性能上超越了最先進(jìn)的Llama-3.2-1B模型2.0%

?更好的擴(kuò)展效果：如圖1所示，CLIMB訓(xùn)練的模型在相同token預(yù)算下展現(xiàn)出更好的擴(kuò)展效果

對比基線：

?隨機(jī)采樣：從所有簇中均勻隨機(jī)采樣

?Best@N：在N個隨機(jī)配置中選擇最佳的

?CLIMB-iter1/2/3：CLIMB的第1/2/3次迭代結(jié)果

CLIMB-iter3在所有基線上都取得了最佳性能，證明了迭代優(yōu)化的有效性。

圖片

領(lǐng)域特定任務(wù)

社會科學(xué)（MMLU-Social-Sciences）：

?隨機(jī)采樣：36.69%

?CLIMB優(yōu)化：41.72%

?提升：5.03%（相對提升13.7%）

STEM（MMLU-STEM）：

?隨機(jī)采樣：32.45%

?CLIMB優(yōu)化：35.87%

?提升：3.42%（相對提升10.5%）

人文（MMLU-Humanities）：

?隨機(jī)采樣：35.21%

?提升：約4%

這些結(jié)果表明，針對特定領(lǐng)域優(yōu)化數(shù)據(jù)混合能夠顯著提升該領(lǐng)域的性能。

代理模型的有效性

跨規(guī)模遷移：使用350M代理模型發(fā)現(xiàn)的最優(yōu)混合，在1B目標(biāo)模型上仍然有效。這證明了代理模型方法的可行性。

更小的代理模型：實(shí)驗(yàn)還測試了62M參數(shù)的代理模型，發(fā)現(xiàn)即使將代理模型規(guī)模縮小5倍，性能仍然保持強(qiáng)勁。這進(jìn)一步降低了計(jì)算成本。

Spearman相關(guān)性：預(yù)測器的預(yù)測準(zhǔn)確率與真實(shí)性能之間的Spearman秩相關(guān)系數(shù)達(dá)到94%，表明預(yù)測器能夠準(zhǔn)確捕捉數(shù)據(jù)混合與性能的關(guān)系。

ClimbLab和ClimbMix數(shù)據(jù)集

ClimbLab

規(guī)模：1.2萬億token的過濾語料庫

結(jié)構(gòu)：20個語義簇，每個簇代表一個主題或領(lǐng)域

用途：作為研究平臺，供研究人員探索數(shù)據(jù)混合優(yōu)化

特點(diǎn)：

?高質(zhì)量過濾：移除低質(zhì)量、重復(fù)和廣告內(nèi)容

?語義組織：通過聚類實(shí)現(xiàn)語義一致性

?開放訪問：在Hugging Face上公開發(fā)布

ClimbMix

規(guī)模：400億token的緊湊數(shù)據(jù)集

設(shè)計(jì)目標(biāo)：在相同token預(yù)算下提供卓越性能的高效預(yù)訓(xùn)練數(shù)據(jù)集

優(yōu)勢：

?性能優(yōu)異：在多個基準(zhǔn)上超越使用更大數(shù)據(jù)集訓(xùn)練的模型

?效率高：僅需400B token即可達(dá)到或超越使用數(shù)萬億token訓(xùn)練的模型

?精心策劃：基于CLIMB發(fā)現(xiàn)的最優(yōu)數(shù)據(jù)混合

應(yīng)用場景：

?資源受限的研究團(tuán)隊(duì)

?快速原型開發(fā)

?領(lǐng)域特定模型訓(xùn)練

最優(yōu)數(shù)據(jù)混合的特征分析

圖片

通過分析CLIMB發(fā)現(xiàn)的最優(yōu)數(shù)據(jù)混合，研究團(tuán)隊(duì)揭示了幾個關(guān)鍵特征：

簇的重要性差異

不同的簇對性能的貢獻(xiàn)差異很大。例如：

通用推理任務(wù)：

?C8（科學(xué)內(nèi)容）：權(quán)重18.5%

?C9（技術(shù)文檔）：權(quán)重15.2%

?C18（問答內(nèi)容）：權(quán)重12.8%

?C19（教育材料）：權(quán)重11.3%

社會科學(xué)任務(wù)：

?C1（社會科學(xué)）：權(quán)重22.1%

?C5（歷史和文化）：權(quán)重16.7%

?C18（問答內(nèi)容）：權(quán)重14.3%

STEM任務(wù)：

?C0（科學(xué)和技術(shù)）：權(quán)重25.4%

?C3（數(shù)學(xué)）：權(quán)重19.8%

?C8（科學(xué)內(nèi)容）：權(quán)重17.6%

稀疏性

最優(yōu)數(shù)據(jù)混合往往是稀疏的，即只有少數(shù)幾個簇占據(jù)主要權(quán)重，大多數(shù)簇的權(quán)重接近零。這表明：

?并非所有數(shù)據(jù)都同等重要

?選擇性采樣比均勻采樣更有效

?領(lǐng)域特定優(yōu)化需要針對性的數(shù)據(jù)選擇

互補(bǔ)性

某些簇之間存在互補(bǔ)關(guān)系。例如，編碼任務(wù)不僅需要編程內(nèi)容（C2），還需要數(shù)學(xué)（C3）、算法（C8）和安全（C12）相關(guān)的內(nèi)容。

領(lǐng)域特異性

不同領(lǐng)域的最優(yōu)混合差異顯著。社會科學(xué)任務(wù)需要更多的人文和社會內(nèi)容，而STEM任務(wù)需要更多的科學(xué)和數(shù)學(xué)內(nèi)容。這強(qiáng)調(diào)了針對特定領(lǐng)域優(yōu)化數(shù)據(jù)混合的重要性。

技術(shù)細(xì)節(jié)

聚類方法

嵌入模型：使用預(yù)訓(xùn)練的BERT模型生成文檔嵌入

聚類算法：K-means，簇?cái)?shù)K=20

距離度量：余弦相似度

簇質(zhì)量評估：使用輪廓系數(shù)（Silhouette Coefficient）評估聚類質(zhì)量

預(yù)測器架構(gòu)

模型類型：梯度提升樹（Gradient Boosting Trees）

輸入特征：20維向量，每個維度表示對應(yīng)簇的采樣比例

輸出：預(yù)測的性能指標(biāo)（如準(zhǔn)確率）

訓(xùn)練數(shù)據(jù)：數(shù)百個混合-性能對

采樣策略

初始采樣：使用Dirichlet分布生成多樣化的初始配置

迭代采樣：基于預(yù)測器的輸出，使用貝葉斯優(yōu)化或進(jìn)化算法生成新配置

修剪策略：保留預(yù)測性能最高的top-K配置進(jìn)行下一輪評估

CLIMB的成功表明，通過智能的數(shù)據(jù)混合優(yōu)化，我們可以在相同或更少的計(jì)算預(yù)算下訓(xùn)練出更強(qiáng)大的模型。這為預(yù)訓(xùn)練數(shù)據(jù)策劃提供了新的范式。

論文標(biāo)題：CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training

本文轉(zhuǎn)載自???AI帝國???，作者：無影寺