比知識蒸餾好用，田淵棟等提出連續概念混合，再度革新Transformer預訓練框架
近年來，大型語言模型（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 進行學習和推理，被許多聲音認為是大模型范式變革的新起點?，F在看來，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 所獲得的性能增益隨著訓練步驟的增加而增加，顯示出很強的泛化性能。
與 KD 基線的比較
研究者還將 CoCoMix 與 KD 基線在多種情況下進行了比較，包括：
- 較強的教師模型教授較小的學生模型；
- 弱到強的監督，即較弱的教師教授較大的學生模型；
- 分布轉移，即學生在不同于教師預訓練分布的語料庫上接受訓練。
如表 1 所示，在這些模型配置中，CoCoMix 都比 KD 有所改進。特別是，CoCoMix 在從弱到強的監督設置中表現出了顯著的性能提升，例如在 386M 中平均復雜度提高了 2.8，而 KD 并沒有表現出很大的改進。這是因為較弱的教師可能會引入噪聲或次優知識，尤其是當學生的能力超過教師時。
在圖 4 中也可以觀察到這一趨勢，當學生的能力超過教師時（尤其是在分布轉移場景中），使用 KD 訓練的模型在訓練中途會落后于標準訓練。相比之下，CoCoMix 選擇性地利用了有用的概念，從而實現了持續的性能提升。
CoCoMix 的可解釋性和可操縱性
CoCoMix 的另一個核心優勢在于其可解釋性和模型操縱。具體來說，由于模型經過訓練后可以預測其隱藏狀態中的概念，因此研究者可以根據概念預測結果來分析其重點關注的概念。此外，通過放大預測概念 z_t 的大小，可以控制模型的輸出生成。
為了驗證可操縱性是否達到預期效果，研究者在預訓練模型的 SAE 潛在空間 c 中對同一概念的激活進行操縱，并確認輸出是否表現出相應的概念。這里使用的是用 CoCoMix 訓練的 386M 參數模型，其中預訓練模型為 GPT-2。
如圖 5 所示，當與「網站地址」相關的概念被放大時，兩個模型都開始生成實際的網站地址。這表明本文模型成功地學習了 GPT-2 對齊概念。
CoCoMix 有效性分析
隨后，研究者對 CoCoMix 進行了詳細分析，以驗證每個提出的組件的效果。角度如下：
- 歸因分數對選擇概念的有效性；
- 概念預測與直接隱藏狀態預測（即用連續損失預測隱藏狀態，而不是用 SAE 將隱藏狀態離散化）之間的比較；
- 壓縮權重的稀疏性；
- 通過分析概念預測和混合的貢獻進行成分分析；
- 通過比較將概念向量添加到原始隱藏狀態和混合（將概念向量與 token 隱藏表示交錯），對概念調節進行設計選擇；
- CoCoMix 與 Pause token（即添加可學習 token）之間的比較。
此處，研究者使用了 69M Transformer，并在 OpenWebText 數據集中的 20B token 上進行訓練。
更多研究細節，可參考原論文。
