（本文閱讀時間：9分鐘）
在自然語言處理領域，預訓練語言模型常因大規模噪聲資料而面臨挑戰。對此，微軟亞洲研究院的研究員們提出了一種新型的基於選擇 token 建模的預訓練方法。該方法透過選擇性語言建模（Selective Language Modeling, SLM）策略，精準篩選出對模型訓練有價值的 token，有效提升了資料效率和模型效能。這一突破不僅優化了模型訓練過程，也為自然語言處理技術的進一步發展提供了新思路。本篇論文在 NeurIPS 2024 上榮獲最佳論文 Runner-Up 獎。
論文共同作者之一、微軟 GenAI 副總裁陳偉柱在微軟研究院播客中，分享了本篇論文的技術細節，歡迎收聽。
論文連結：
https://arxiv.org/abs/2404.07965
專案連結：
https://github.com/microsoft/rho
現有大模型基於大批次文字語料進行預訓練，在各類文字生成、文字理解和文字邏輯推理等任務上表現突出。然而，預訓練過程中從各種來源獲取的原始語料存在大量噪聲，因此科研人員經常採用一些質量過濾方法對原始語料進行過濾，使其可以用於模型預訓練。例如，文件級（document-level）過濾可以去除一些干擾文件，進一步地還可以在行級（line-level）過濾單個文件中的噪聲，從而得到高質量的語料，用於預訓練。
圖1：語料清洗示意圖
在以往的方法中，過濾出來的高質量語料輸入到以因果語言建模（Causal Language Modeling）方式的模型當中，計算每個 token 的損失並平均後求梯度，然後更新模型的引數。然而，當使用這種 next-token prediction 的形式對完整的句子序列進行建模時，可能忽略一些 token 級（token-level）的內容。
比如看到圖2具體的案例：“The farm has 35 hens <Apr12 1:24> and 12 pigs. ##davidjl123 says totaling 47 animals.”，其中包含像“<Apr12 1:24>”這樣的時間資訊，以及“##davidjl123”這樣的使用者 id。這種 token 級的噪聲在語料中較為常見，這些細粒度的噪聲很難透過以前採用的文件級和行級過濾察覺。即使語料質量較高，中間仍然可能存在一些高度不確定性的 token，例如在圖2的案例中，透過“12 pigs”和“35 hens”可以推測農場裡一共有“47 animals”，但是在沒有先驗條件和上下文的前提下，讓語言模型學習準確預測農場裡有幾隻豬是困難的。
圖2：Token 級噪聲樣例圖
然而這些高度不確定的 token 和噪聲 token 在因果語言建模中沒有區分，以平均的權重參與到最終的模型更新中，會使語言模型感到困惑。
為了進一步研究 token 級對於模型訓練的影響，研究員們對語言模型預訓練過程做了 token 損失（token loss）的動態分析。研究員們使用了 15B 的 OpenWebMath 語料來訓練 Tinyllama-1B 模型，而且在每訓練 1B 的 token 後於驗證集上評估所有 token 的損失。透過獲取所有檢查點在驗證集上的 token 損失資料，研究員們為驗證集中的每個 token 擬合了損失的變化趨勢，並重點關注訓練初期和末期的 token 損失，以及訓練前後的損失差值。基於訓練前後損失差值和整體token平均的損失，研究員們將驗證集中的 token 分為四類：H→H、L→H、H→L、L→L，分別代表 token 損失在訓練過程中動態變化的趨勢。（具體的劃分依據如圖3所示。）
圖3：Token 類別劃分依據
H→H 代表一直保持較高損失的 token，L→L 則一直保持較低的損失。L→H 代表損失上升的 token，而 H→L 是最常見的變化趨勢——token 損失下降。從圖4(a)中可以看到，僅26%的 token 屬於 H→L 類別，大多數 token 的損失變化不大，甚至有12%的 token 的損失呈現上升趨勢。當研究員們隨機取樣 H→H 和 L→L 類別的 token 並單獨觀察其損失曲線（圖4(b)和圖4(c)）時，可以發現這些 token 在訓練過程中處於反覆波動的狀態，可能影響模型的收斂速度。因此，研究員們認為如果有一種方法可以合理選擇適合學習且更有用的 token，讓其參與訓練，將可以減少噪聲並提升模型的資料效率。
圖4：Token 級損失的動態示意圖
基於此，研究員們提出了選擇性語言建模（Selective Language Modeling, SLM），在保證原有輸入序列的情況下，透過在損失端裁剪模型所需的 token 損失，來選擇有用的 token，如圖5所示。
圖5：選擇性語言建模示意圖
首先，需要有一個高質量的語料庫。第一步，使用傳統的因果語言建模損失在高質量的語料庫上訓練一個參考模型（reference model）來建模高質量 token 的分佈。第二步，用訓練好的參考模型在離線階段對預訓練語料中的每個 token 打分，最終得分由 Token Scoring 公式計算得到。這樣的打分方式不會在實際的訓練過程中引入額外的時間開銷。第三步，用打好分的預訓練語料訓練模型，對每個 token 的分數排序後，取 topk% 作為選擇的 token，對應 Token Selection 的公式。最後，透過 SLM 方式訓練的損失公式迭代更新模型。
圖6：選擇性語言建模流程圖
圖7顯示，在相同的資料集中，透過 SLM 方式訓練的模型比直接訓練的模型有效提高了資料效率，加快了收斂速度。
圖7：選擇性語言建模的資料效率
研究員們在數學領域進行了實驗。實驗中，研究員們使用 14B OpenWebMath 語料，分別在 Tinyllama-1B 和 Mistral-7 上繼續預訓練該模型，並採用 SLM 訓練方式，選擇比例分別為60%和70%。從圖8中的結果可以看到，使用 SLM 訓練的 Rho-1 Math 1B 和 7B 模型相較於直接繼續預訓練，效能分別提升了16%和10%。
圖8：數學領域中的 Rho-1 預訓練結果
為了進一步驗證預訓練的結果，研究員們基於上述訓練的基礎模型進行了推理微調對比實驗。圖9顯示，該模型在數學領域上取得了與 DeepSeekmath7B 相當的成績，在數學基準上的準確率均高於50%。值得注意的是，在預訓練過程中，該模型僅使用了 14B 的 OpenWebMath，遠少於 DeepSeekMath 使用的 120B 數學相關語料，進一步證明了使用 SLM 訓練的資料效率。
圖9：Rho-1 SFT 在數學領域的結果
研究員們也在通用領域（general domain）上進行了類似的實驗，採用了包含總計 80B token 的預訓練語料對 Tinyllama1B 進行繼續預訓練。如圖10的評估結果所示，Rho-1 在各項通用基準測試中平均提升了約6%。
圖10：Rho-1 在通用領域上的結果
同時，本篇論文還探討了在缺乏高質量語料作為參考的情況下，SLM 是否能夠正常運作。如圖11所示，可以直接將預訓練好的基礎模型作為參考模型進行自我參考（self-reference）迭代。在 Tinyllama 上繼續預訓練以驗證自我參考的可行性。在圖12中可以看到，僅透過一輪迭代，SLM 就可以顯著提升模型效能，這在模型的自我提升方面具有重要意義。
圖11：自我參考流程示意圖
圖12：自我參考結果
總而言之，該研究表明並非所有 token 在語言模型預訓練過程中都是同等重要的。透過基於 token 級資料篩選的 SLM 建模方式能夠極大提高模型的資料效率。這種 token 級的思路不僅適用於預訓練，還可以應用於微調、強化學習、多模態等領域。此外，選擇 token 的形式可以多種多樣，要根據具體場景及需求確定不同的 token 選擇方法。研究員們希望未來能夠出現更多有效的 token 選擇策略和重新加權策略。
你也許還想看：