比知識蒸餾好用！田淵棟等提出CoCoMix：革新Transformer預訓練框架

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達

點選進入—>【Mamba/多模態/擴散】交流群

新增微訊號：CVer2233，小助手會拉你進群！

掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：機器之心｜編輯：蛋醬、陳陳

近年來，大型語言模型（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 所獲得的效能增益隨著訓練步驟的增加而增加，顯示出很強的泛化效能。

與 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 上進行訓練。

更多研究細節，可參考原論文。

何愷明在MIT授課的課件PPT下載

在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集

CVPR 2024 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2024，即可下載CVPR 2024論文和程式碼開源的論文合集

Mamba、多模態和擴散模型交流群成立

掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-Mamba、多模態學習或者擴散模型微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。

一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者擴散模型+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群