科技

比知識蒸餾好用，田淵棟等提出連續概念混合，再度革新Transformer預訓練框架

02月16日 10:16 新浪網 tech-auto-hilite

機器之心報導

編輯：蛋醬、陳陳

近年來，大型語言模型（LLMs）的進展徹底改變了自然語言處理領域，並因此成為各種現實應用中的核心技術，例如代碼助手、搜索引擎和個人 AI 助手。

這些突破的核心在於對「下一個 token 預測」的範式。

然而，自然語言 token 代表的意思通常是表層的（例如 the 或 a 這樣的功能性詞彙），需要模型進行大量訓練才能獲得高級推理和對概念的理解能力，同時也限制了它們處理長期任務（如規劃）的能力。

為瞭解決這一問題，最近的研究探索了超越 token 層面信號的方法。例如有研究表明稀疏自編碼器（Sparse Autoencoders, SAEs）能夠通過捕捉高級語義概念，有效地分離出大型語言模型（LLMs）中有意義的潛在特徵。

本文，來自 Meta 等機構的研究者提出了一種新穎且高效的預訓練框架：連續概念混合（Continuous Concept Mixing, CoCoMix），其將離散的下一個 token 預測與連續概念相結合。

CoCoMix 用來預測從預訓練的稀疏自編碼器中學習到的連續概念，並通過與 token 隱藏表示交錯的方式將其混合到模型的隱藏狀態中。

具體來說，本文使用經過預訓練的 SAE 提取語義概念，並根據歸因（attribution）分數選擇最具影響力的概念，這些分數量化了每個概念對模型輸出的影響。然後，模型通過交叉熵損失訓練，從其隱藏狀態中預測這些選定的概念。一旦預測出多個概念，就將它們壓縮為單個連續概念，並通過與 token 嵌入交錯的方式混合（或插入）到隱藏狀態中，從而直接貢獻於下一個 token 的預測。

本文通過在多個語言建模基準和不同規模的預訓練模型（從百萬級到十億級參數規模）上進行了評估，從而證明了 CoCoMix 的有效性。

結果表明，CoCoMix 采樣效率更高，優於標準的下一個 token 預測、知識蒸餾以及插入停頓 token。本文發現，在端到端的框架中結合概念學習和交錯技術對於性能提升至關重要。

CoCoMix 性能有多好？舉例來說，當將 CoCoMix 應用於一個 1.38B 規模的模型時，CoCoMix 在減少 21.5% 訓練 token 的情況下，實現了與下一個 token 預測相當的性能。

此外，CoCoMix 在弱監督到強監督的場景中表現出顯著改進，其中從小模型中提取的概念甚至可以用作監督更大模型訓練的標籤。

論文標題：LLM Pretraining with Continuous Concepts
論文地址：https://arxiv.org/pdf/2502.08524
項目地址：https://github.com/facebookresearch/RAM/tree/main/projects/cocomix

此前，Meta 提出了大型概念模型（LCM），同樣也是通過概念而非 token 進行學習和推理，被許多聲音認為是大模型範式變革的新起點。現在看來，CoCoMix 出現，讓 Meta 在創新算法逐步取代「連續預測下一個 token」這條道路上又邁出了重要的一步。

CoCoMix 介紹

CoCoMix 是一個使用連續概念擴展下一個 token 預測的框架。

其核心訓練流程包括：一個概念選擇框架（參見圖 1 左），以及兩個用於學習和利用連續概念的訓練步驟（step）（參見圖 1 右）。

首先，本文使用歸因分數選擇重要概念，該分數衡量了每個概念對輸出的影響。
然後，本文提出通過交叉熵損失從模型的隱藏狀態中預測選定的概念，使模型能夠隱式地學習哪些概念應被編碼為隱藏表示。
最後，本文利用預測的概念創建一個連續概念，並將其交錯插入到隱藏狀態中，使模型能夠顯式地學習如何使用連續概念以及 token 隱藏狀態。直觀上，模型選擇性地學習哪些概念對下一個 token 預測有用，以及如何將這些概念與 token 表示混合。

實驗

實驗部分，研究者主要通過以下幾個問題對 CoCoMix 進行了實證評估：

CoCoMix 能否提高 LLM 預訓練中下一個 token 預測的性能？(圖 2 和圖 3）
與其他知識提煉方法相比，CoCoMix 在弱到強監督設置中是否有所改進？(表 1 和圖 4）
CoCoMix 是否引入了模型的可解釋性和可操縱性？(圖 5）
CoCoMix 的每個建議組件對性能有何貢獻？(圖 6）

首先是兩個核心結果：

在相對大規模的預訓練設置中與 NTP 的比較；
與 KD 基線的比較，尤其是在從小模型中提取的概念用於指導大模型的弱到強監督場景中。

大規模使用 CoCoMix 改進 NTP

如圖 3 所示，CoCoMix 在各種規模的模型上都能持續顯著提高下遊任務的整體性能。結果還表明，較大的模型（如 386M 和 1.38B）可以從使用從較小的 124M 模型中提取的概念中獲益，顯示了有效的弱到強監督。

如圖 2 所示，在十億級規模的模型上，CoCoMix 與 NTP 相比持續提高了性能。例如，CoCoMix 的性能與 NTP 相近，但使用的 token 卻減少了 21.5%，顯示了很高的采樣效率。最後，值得注意的是，使用 CoCoMix 所獲得的性能增益隨著訓練步驟的增加而增加，顯示出很強的泛化性能。