除了提升資料效率之外，本文方法 MeCo 保證了計算開銷與複雜性也幾乎不會增加。
普林斯頓大學電腦科學系助理教授陳丹琦團隊又有了新論文，這次將重點放在了「使用後設資料來加速預訓練」上來。
我們知道，語言模型透過在大量網路語料庫上進行訓練來實現卓越的通用能力。多樣性訓練資料凸顯了一個根本性挑戰：人們自然地根據資料來源來調整他們的理解，與之不同，語言模型將所有內容作為等效樣本來處理。
這種以相同方式處理異構源資料的做法會帶來兩個問題：一是忽略了有助於理解的重要上下文訊號，二是在專門的下游任務中阻礙模型可靠地展示適當的行為，比如幽默或事實。
面對以上這些挑戰，併為了提供每個文件來源的更多資訊，陳丹琦團隊在本文中提出透過在每個文件之前新增廣泛可用的源 URL，從而在預訓練期間使用文件相應的後設資料進行調節。並且為了確保模型在推理過程中無論有無後設資料都能高效地執行，在最後 10% 的訓練中實施了冷卻（cooldown）。他們將這種預訓練方法稱為 Metadata Conditioning then Cooldown（MeCo）。
先前的工作中已經有人使用後設資料條件來引導模型生成並提高模型對惡意提示的穩健性，但研究者透過關鍵的兩點確認了所提方法的通用實用性。首先，他們證明這一正規化可以直接加速語言模型的預訓練並提高下游任務效能。其次，MeCo 的冷卻階段確保模型在沒有後設資料的情況下可以執行推理，這點與以往的方法不同。
本文的主要貢獻包括如下：
一、MeCo 大大加速了預訓練過程。研究者證明，MeCo 使得 1.6B 的模型在少用 33％訓練資料的情況下，實現與標準預訓練模型相同的平均下游效能。MeCo 在模型規模（600M、1.6B、3B 和 8B）和資料來源（C4、RefinedWeb 和 DCLM）表現出了一致的增益。
二、MeCo 開闢了一種引導模型的新方法。在推理過程中，在提示之前新增合適的真實或合成 URL 可以誘導期望的模型行為。舉個例子，使用「factquizmaster.com」（並非真實 URL）可以增強常識知識任務的效能，比如零樣本常識問題絕對效能可以提升 6%。相反，使用「wikipedia.org」（真實 URL）可以將有毒生成的可能性比標準無條件推理降低數倍。
三、MeCo 設計選擇的消融實驗表明，它能與不同型別的後設資料相容。使用雜湊 URL 和模型生成主題的消融實驗表明，後設資料的主要作用是按照來源對文件進行分組。因此，即使沒有 URL，MeCo 也可以有效地合併不同型別的後設資料，包括更細粒度的選項。
研究結果表明，MeCo 可以顯著提高語言模型的資料效率，同時幾乎不會增加預訓練過程的計算開銷和複雜性。此外，MeCo 提供了增強可控性，有望建立更可控的語言模型，並且它與更細粒度和創造性的後設資料的普遍相容性值得進一步探索。
總之，作為一種簡單、靈活、有效的訓練正規化，MeCo 可以同時提高語言模型的實用性和可控性。
論文標題：Metadata Conditioning Accelerates Language Model Pre-training
論文地址：https://arxiv.org/pdf/2501.01956v1
程式碼地址：https://github.com/princeton-pli/MeCo
論文一作高天宇（Tianyu Gao）還在評論區與讀者展開了互動，並回答了一個問題「MeCo 是否需要平衡過擬合和欠擬合」。他表示，本文的一個假設是 MeCo 進行隱式資料混合最佳化（DoReMi、ADO）並上取樣欠擬合和更多有用域。
OpenAI 一位研究人員 Lucas Beyer 表示，他很久之前就對視覺語言模型（VLM）做過類似的研究，很有趣，但最終用處不大。
方法概覽
本文方法包括以下兩個訓練階段，如下圖 1 所示。
使用後設資料條件進行預訓練（前 90%）：模型在串接的後設資料和文件上進行訓練，並遵循以下模板「URL: en.wikipedia.org\n\n [document]」。使用其他型別的後設資料時，URL 替換為相應的後設資料名稱。研究者僅計算文件 token 的交叉熵損失，而忽略出自模板或後設資料的 token。他們在初步實驗中發現：使用這些 token 訓練會損害下游任務效能。
使用標準資料進行冷卻（後 10%）：對於僅使用後設資料增強的資料進行訓練的模型，在沒有後設資料的情況下效能會下降（具體可見下表 4）。為了確保通用性，研究者在冷卻階段，使用了沒有任何後設資料的標準預訓練文件來訓練模型，該階段涵蓋了預訓練過程最後 10% 的步驟。
冷卻階段繼承了來自後設資料條件階段的學習率計劃和最佳化器狀態，即它從上一個階段的最後一個檢查點初始化學習率、模型引數和最佳化器狀態，並繼續根據計劃來調整學習率。
研究者還在所有實驗中採用了以下兩項技術，並且初步實驗表明它們提高了基線預訓練模型的效能：
停用了跨文件注意力，此舉既加快了訓練速度（1.6B 模型的速度提升了 25％），又提高了下游任務的效能； 將多個文件打包成一個序列時，確保每個序列都從一個新文件開始，而不是從一個文件的中間開始，這可能會導致在將文件打包為一個固定長度時丟棄一些資料，但被證明有利於提高下游任務效能。
實驗結果
研究者在所有實驗中使用了 Llama 系列模型使用的 Transformer 架構和 Llama-3tokenizer，使用了四種規模的模型大小，分別是 600M、1.6B、3B 和 8B。他們對語言模型採用了標準最佳化設定，即 AdamW 最佳化器和餘弦學習率計劃。
少用 33% 資料，MeCo 效能與標準預訓練方法相當
下表 1 顯示了研究者在 DCLM 上的 160B token 上，對 1.6B 語言模型進行預訓練的主要結果。他們首先觀察到，在大多數任務中，MeCo 的效能顯著優於標準預訓練方法。MeCo 還超越了資料挑選基線。並且與資料挑選方法不同的是，MeCo 不會產生任何計算開銷，它利用了預訓練資料中隨時可用的 URL 資訊。
更重要的是，MeCo 實現了與標準預訓練方法相當的效能，同時使用的資料和計算量減少了 33%，代表了資料效率的顯著提高。
下表 1 為困惑度指標，表明了驗證困惑度與下游效能無關。值得注意的是，當將 240B 基線模型與 160B MeCo 模型比較時，由於資料量較大，基線模型表現出的困惑度要低得多，但這兩個模型實現了類似的平均效能。
研究者在下圖 2 中展示了整個預訓練過程中下游任務的效能變化。對於 MeCo，圖中的每個檢查點都包含使用 16B token（佔總訓練 token 的 10％）的冷卻階段。例如，80B 檢查點包含了 64B token 的條件訓練和 16B token 的冷卻。他們觀察到，MeCo 始終超越了基線模型，尤其是在訓練後期。
MeCo 在所有模型規模下均提升了效能
下圖 3 顯示了不同模型規模（600 M、1.6B、3B 和 8B）的結果。研究者使用相同的最佳化超引數和相同的資料量（DCLM 上的 160B）來訓練所有模型，其中 8B 模型是個個例，它使用 80B token 進行訓練，由於資源限制和訓練不穩定而導致學習率較低。
研究者觀察到，MeCo 在所有規模下均提升了模型效能。並且 MeCo 看起來可以為更大的模型帶來更多的改進，十億級引數的模型與 600M 相比顯示出更顯著的收益。不過需要注意，這是一個定性觀察，與預訓練損失相比，下游任務效能的擴充套件不太平穩。
MeCo 提升了不同訓練語料庫的效能
研究者基於三個不同的資料來源（C4、RefinedWeb 和 DCLM），在 160B token 上訓練了 1.6B 模型，結果如下圖 4 所示。如果將平均下游效能作為資料質量指標，三個資料來源的排序為 DCLM > RefinedWeb > C4。他們觀察到，MeCo 在不同資料來源上實現了一致且顯著的增益，平均準確率和單個任務均是如此。
更多技術細節請參閱原論文。