和人工標(biāo)記數(shù)據(jù)說拜拜，利用預(yù)訓(xùn)練語言模型中的注意力機(jī)制就能選擇可激發(fā)推理能力的訓(xùn)練數(shù)據(jù)！

字節(jié)Seed團(tuán)隊最新宣布了一個重要成果——AttentionInfluence。

無需訓(xùn)練，無需標(biāo)簽，只需用1.3B模型給7B模型選擇數(shù)據(jù)，就能提升模型推理能力，甚至也能提升代碼生成能力。

以往，篩選數(shù)據(jù)的方法通常依賴于監(jiān)督分類器，需要人工或大語言模型進(jìn)行標(biāo)注，難免引入領(lǐng)域特定偏見。

字節(jié)Seed團(tuán)隊注意到：

預(yù)訓(xùn)練模型中的檢索頭與檢索和上下文推理緊密相關(guān)。

檢索頭在訓(xùn)練早期就會出現(xiàn)，逐漸增強(qiáng)，并最終在訓(xùn)練的中后期階段牢固建立，對模型性能起到至關(guān)重要的作用。

1.3B參數(shù)稠密模型中檢索頭的演化過程，be like：

但如果直接關(guān)閉它們會怎樣？

他們用小型預(yù)訓(xùn)練語言模型通過簡單的注意力頭屏蔽操作，充當(dāng)強(qiáng)大的模型的數(shù)據(jù)選擇器。

具體操作是，識別重要檢索頭，屏蔽這些頭以創(chuàng)建性能下降的“弱”模型，計算“弱”模型與原始“強(qiáng)”模型之間的損失差異，根據(jù)損失增加幅度對數(shù)據(jù)進(jìn)行排名，形成影響分?jǐn)?shù)。

沒想到，實驗后他們得到了一個驚人結(jié)果。

將AttentionInfluence方法應(yīng)用于1.3B參數(shù)預(yù)訓(xùn)練語言模型，對SmolLM語料庫進(jìn)行數(shù)據(jù)選擇，篩選出73.1B tokens與完整的SmolLM語料庫組合，使用WSD方法預(yù)訓(xùn)練7B模型。

在知識密集型和推理密集型基準(zhǔn)測試中模型性能均有提升，具體來說：

MMLU+1.4個百分點、MMLU-Pro+2.7個百分點、AGIEval-en+1.8個百分點、GSM8K+2.7個百分點、HumanEval+3.5個百分點。

這項研究發(fā)布后引來不少網(wǎng)友關(guān)注，谷歌DeepMind研究科學(xué)家都轉(zhuǎn)發(fā)為其點贊：

有網(wǎng)友看后表示：

多么簡單而巧妙的數(shù)據(jù)選擇思路！

關(guān)于這項研究的更多細(xì)節(jié)，我們接著往下看。

讓預(yù)訓(xùn)練模型識別重要數(shù)據(jù)

如前所述，AttentionInfluence方法的核心思想是通過比較基礎(chǔ)模型和弱化參考模型之間的損失差異來評估訓(xùn)練數(shù)據(jù)的對推理的影響程度。

實現(xiàn)方法包含兩個主要步驟：檢測特定重要頭部、計算AttentionInfluence分?jǐn)?shù)。

檢測特定重要頭部

在本研究中，作者主要關(guān)注檢索頭，此前已有研究表明檢索頭與大語言模型的檢索和推理能力高度相關(guān)。

受CLongEval中提出的關(guān)鍵段落檢索評估任務(wù)啟發(fā)，團(tuán)隊采用了一種類似的簡單proxy task，在可控環(huán)境下評估大語言模型的檢索能力，并識別與檢索和推理強(qiáng)相關(guān)的注意力頭。

為此，他們構(gòu)建了一個包含800個樣本的合成測試數(shù)據(jù)集。每個樣本被格式化為一個3-shot自然語言檢索任務(wù)，由上下文、三個上下文中的示例和一個查詢hash_key組成。

每個上下文是一個包含k個鍵值對（key-value pairs）的JSON對象，其中每個鍵是隨機(jī)生成的32字符字母數(shù)字字符串（hash_key），每個值（text_val）是從網(wǎng)絡(luò)文檔語料庫中采樣的自然語言句子。

該任務(wù)要求模型從上下文中檢索text_val，并輸出與給定查詢hash_key對應(yīng)的text_val。

包含三個上下文中的示例（即3-shot）旨在模擬小樣本學(xué)習(xí)場景，并幫助模型理解任務(wù)?？紤]到現(xiàn)有預(yù)訓(xùn)練模型的上下文長度限制，團(tuán)隊將每個測試樣本的總長度（包括輸入提示和答案）限制為接近但不超4096token。

接下來，計算每個注意力頭在測試樣本上的檢索分?jǐn)?shù)。

在這項工作中，團(tuán)隊使用一個基于類Llama 2架構(gòu)的1.3B參數(shù)模型作為小型預(yù)訓(xùn)練語言模型，使用平均分?jǐn)?shù)作為頭的最終檢索分?jǐn)?shù)，并按該分?jǐn)?shù)對其進(jìn)行排序，選擇排名前5%的頭作為特別重要的頭。

計算AttentionInfluence分?jǐn)?shù)

獲得重要頭部后，接下來計算每個樣本的AttentionInfluence分?jǐn)?shù)。

首先是創(chuàng)建參考模型，通過屏蔽在第一階段檢測到的基礎(chǔ)模型的重要頭部，獲得參考模型。

然后，使用基礎(chǔ)模型計算語料庫中每個樣本的平均token級交叉熵?fù)p失（L_base），使用參考模型計算相應(yīng)的損失（L_ref）。

最后將L_base和L_ref之間的相對差值作為注意力影響分?jǐn)?shù)，以量化每個樣本的推理影響程度，其計算公式如下：

由于語言模型對來自不同領(lǐng)域（如通用領(lǐng)域、數(shù)學(xué)領(lǐng)域、代碼領(lǐng)域）的數(shù)據(jù)的損失，因分布差異顯著而無法直接比較，團(tuán)隊將注意力影響分?jǐn)?shù)的比較限制在相同領(lǐng)域內(nèi)。

團(tuán)隊認(rèn)為，注意力影響分?jǐn)?shù)越高，表明樣本具有更高的推理強(qiáng)度。

效果全面分析

實驗設(shè)置上，團(tuán)隊使用Llama2類似的1.3B參數(shù)模型作為小型預(yù)訓(xùn)練語言模型，對SmolLM語料庫進(jìn)行數(shù)據(jù)選擇.

根據(jù)AttentionInfluence分?jǐn)?shù)選擇排名前20%的樣本，約73.1B tokens，使用選定的73.1B tokens與完整的SmolLM語料庫組合，預(yù)訓(xùn)練7B參數(shù)模型。

作為對比，基線模型則是僅使用SmolLM語料庫訓(xùn)練的相同架構(gòu)和大小的模型。

然后，在小樣本學(xué)習(xí)設(shè)置下，團(tuán)隊采用一套涵蓋四大類別的綜合基準(zhǔn)評估，對模型與基線模型進(jìn)行全面比較：

• 綜合基準(zhǔn)，包括AGIEval-en、MMLU、MMLU-Pro、GPQA、C-Eval；
• 數(shù)學(xué)、代碼和推理，包括GSM8K、MATH、HumanEval、ARC Challenge、DROP、BBH；
• 常識推理與理解，包括HellaSwag、ARC-Easy、WinoGrande、CommonSenseQA、PiQA、OpenBookQA、TriviaQA；
• 閱讀理解，以RACE為代表。

主要研究結(jié)果顯示，使用AttentionInfluence選擇的數(shù)據(jù)訓(xùn)練的模型在多個關(guān)鍵基準(zhǔn)上顯著優(yōu)于基線：

研究還跟蹤了預(yù)訓(xùn)練過程中的性能演變，AttentionInfluence模型在整個預(yù)訓(xùn)練過程中始終優(yōu)于基線，性能差距在訓(xùn)練早期（約100B tokens之前）就已顯現(xiàn)，并在整個訓(xùn)練過程中保持穩(wěn)定，即使在學(xué)習(xí)率衰減（LRD）階段，性能優(yōu)勢仍然存在。

另外，當(dāng)特定重要頭部被屏蔽時，1.3B模型在某些任務(wù)上的性能顯著下降，而AttentionInfluence方法選擇的數(shù)據(jù)往往能改善7B模型在這些任務(wù)上的性能，表明該方法具有預(yù)測能力。

將AttentionInfluence應(yīng)用于更大的7B參數(shù)模型進(jìn)行數(shù)據(jù)選擇時，能在多個知識密集型和推理密集型基準(zhǔn)上獲得更好的性能，表明增加模型規(guī)模有助于選擇更高推理強(qiáng)度的樣本。

為驗證AttentionInfluence的有效性，團(tuán)隊還設(shè)計了兩個指標(biāo)來量化所選數(shù)據(jù)的質(zhì)量：

• Education Score：評估內(nèi)容的教育價值
• Reasoning Score：評估內(nèi)容的推理強(qiáng)度

實驗結(jié)果顯示，AttentionInfluence和FineWeb-Edu分類器在教育相關(guān)內(nèi)容上獲得相當(dāng)?shù)姆謹(jǐn)?shù)；在推理方面，AttentionInfluence獲得顯著更高的分?jǐn)?shù)；在Python-Edu和OpenWebMath領(lǐng)域，AttentionInfluence選擇的樣本平均長度幾乎是FineWeb-Edu分類器選擇樣本的兩倍。

團(tuán)隊也對所選數(shù)據(jù)進(jìn)行了多樣性分析，感興趣的童鞋可以查看原論文。

總之，團(tuán)隊表示這些結(jié)果驗證了AttentionInfluence法能有效地識別高質(zhì)量的預(yù)訓(xùn)練數(shù)據(jù)，從而增強(qiáng)大語言模型的知識和推理能力，尤其在需要綜合知識和復(fù)雜推理的基準(zhǔn)測試中取得了顯著提升。

此外，AttentionInfluence可與 FineWeb-Edu分類器結(jié)合使用，以在需要簡單事實性知識、高級推理或兩者兼具的任務(wù)中全面提升大語言模型的性能。

論文鏈接：https://arxiv.org/pdf/2505.07293