和人工標記數據說拜拜，利用預訓練語言模型中的注意力機制就能選擇可激發推理能力的訓練數據！

字節Seed團隊最新宣布了一個重要成果——AttentionInfluence。

無需訓練，無需標簽，只需用1.3B模型給7B模型選擇數據，就能提升模型推理能力，甚至也能提升代碼生成能力。

以往，篩選數據的方法通常依賴于監督分類器，需要人工或大語言模型進行標注，難免引入領域特定偏見。

字節Seed團隊注意到：

預訓練模型中的檢索頭與檢索和上下文推理緊密相關。

檢索頭在訓練早期就會出現，逐漸增強，并最終在訓練的中后期階段牢固建立，對模型性能起到至關重要的作用。

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

但如果直接關閉它們會怎樣？

他們用小型預訓練語言模型通過簡單的注意力頭屏蔽操作，充當強大的模型的數據選擇器。

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

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

將AttentionInfluence方法應用于1.3B參數預訓練語言模型，對SmolLM語料庫進行數據選擇，篩選出73.1B tokens與完整的SmolLM語料庫組合，使用WSD方法預訓練7B模型。

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

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

這項研究發布后引來不少網友關注，谷歌DeepMind研究科學家都轉發為其點贊：

有網友看后表示：

多么簡單而巧妙的數據選擇思路！

關于這項研究的更多細節，我們接著往下看。

讓預訓練模型識別重要數據

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

實現方法包含兩個主要步驟：檢測特定重要頭部、計算AttentionInfluence分數。

檢測特定重要頭部

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

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

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

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

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

包含三個上下文中的示例（即3-shot）旨在模擬小樣本學習場景，并幫助模型理解任務。考慮到現有預訓練模型的上下文長度限制，團隊將每個測試樣本的總長度（包括輸入提示和答案）限制為接近但不超4096token。

接下來，計算每個注意力頭在測試樣本上的檢索分數。

在這項工作中，團隊使用一個基于類Llama 2架構的1.3B參數模型作為小型預訓練語言模型，使用平均分數作為頭的最終檢索分數，并按該分數對其進行排序，選擇排名前5%的頭作為特別重要的頭。

計算AttentionInfluence分數

獲得重要頭部后，接下來計算每個樣本的AttentionInfluence分數。

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

然后，使用基礎模型計算語料庫中每個樣本的平均token級交叉熵損失（L_base），使用參考模型計算相應的損失（L_ref）。

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

由于語言模型對來自不同領域（如通用領域、數學領域、代碼領域）的數據的損失，因分布差異顯著而無法直接比較，團隊將注意力影響分數的比較限制在相同領域內。

團隊認為，注意力影響分數越高，表明樣本具有更高的推理強度。

效果全面分析

實驗設置上，團隊使用Llama2類似的1.3B參數模型作為小型預訓練語言模型，對SmolLM語料庫進行數據選擇.

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

作為對比，基線模型則是僅使用SmolLM語料庫訓練的相同架構和大小的模型。

然后，在小樣本學習設置下，團隊采用一套涵蓋四大類別的綜合基準評估，對模型與基線模型進行全面比較：

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

主要研究結果顯示，使用AttentionInfluence選擇的數據訓練的模型在多個關鍵基準上顯著優于基線：

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

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

將AttentionInfluence應用于更大的7B參數模型進行數據選擇時，能在多個知識密集型和推理密集型基準上獲得更好的性能，表明增加模型規模有助于選擇更高推理強度的樣本。

為驗證AttentionInfluence的有效性，團隊還設計了兩個指標來量化所選數據的質量：

• Education Score：評估內容的教育價值
• Reasoning Score：評估內容的推理強度

實驗結果顯示，AttentionInfluence和FineWeb-Edu分類器在教育相關內容上獲得相當的分數；在推理方面，AttentionInfluence獲得顯著更高的分數；在Python-Edu和OpenWebMath領域，AttentionInfluence選擇的樣本平均長度幾乎是FineWeb-Edu分類器選擇樣本的兩倍。

團隊也對所選數據進行了多樣性分析，感興趣的童鞋可以查看原論文。

總之，團隊表示這些結果驗證了AttentionInfluence法能有效地識別高質量的預訓練數據，從而增強大語言模型的知識和推理能力，尤其在需要綜合知識和復雜推理的基準測試中取得了顯著提升。

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

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