新智元報道
新智元報道
編輯：KingHZ
【新智元導讀】基於連續概念，Meta團隊新研究提出了超越「下一個token預測」語言建模新正規化。更加重要的是，新方法不僅能增強原有的正規化，而且比起知識蒸餾，資料量減少20%，甚至能從小模型提取概念指導更大的模型！
「下一個token預測」（next token prediction，NTP）是大語言模型（LLMs）不斷取得突破的核心技術。
但這種依賴tokenization的方法，導致LLM「嚴重偏科」。
關注AI的可能也知道GPT-4o不會數字母，不知道Strawberray中有幾個字母「r」。
為了解決此類問題，最近的研究探討了超越token級訊號的方法，利用更豐富的資訊來訓練模型。
比如說，一次性預測多個token；在下一個token預測之前增強輸入，加入豐富的訊號；或用連續的潛在表示替代離散的語言標記，提高推理效率。
這次，受到近期研究發現的啟發，來自Meta的研究人員認為稀疏自編碼器（Sparse Autoencoders，SAEs）可以捕捉高層次的語義概念，在LLM中有效地隔離出有意義的潛在特徵。
由於SAEs是透過稀疏性約束訓練，重構模型的隱狀態，它促使模型集中關注一組緊湊的概念維度。
這可以突出預訓練模型的概念——即支撐模型預測的核心語義方向，同時避免不必要的特徵。
論文連結：https://arxiv.org/abs/2502.08524
新研究在多個語言建模基準和預訓練模型規模（從百萬規模到十億規模的引數模型）上進行廣泛的評估，展示了CoCoMix的有效性。
例如，在應用於1.38B引數模型時，CoCoMix在下一個token預測任務中的表現與傳統方法相當，同時減少了21.5%的訓練資料量。
此外，CoCoMix在弱監督到強監督場景中表現出顯著的提升，其中從小模型中提取的概念甚至可以作為真實標籤，用於監督大模型的訓練。
最後，透過插入壓縮的概念向量，能夠在生成過程中探查預測的概念，從而引導和控制模型。
主要方法：CoCoMix
CoCoMix是一種新的LLM預訓練框架，透過預測概念並將其混入模型的隱狀態中，以提高下一個token預測的準確性。
更高的樣本效率，在下一個token預測、知識蒸餾以及插入暫停token等任務中表現優越，同時提高可解釋性和可引導性，增強模型的可控性。
連續概念混合（CoCoMix）使用基於連續概念的語言建模框架。具體而言，CoCoMix包含三個步驟來學習和使用潛在概念：
1. 從預訓練的SAE中，提取概念並選擇顯著的概念。
2. LLM從其隱藏狀態預測這些概念。
3. 一旦預測出多個概念，就將它們壓縮成一個單一的「連續概念」，並將其「混合」到LLM隱藏狀態中。
圖1：CoCoMix的概覽。
新研究證明了CoCoMix具有更高的樣本效率，並且優於標準的下一個token預測和知識蒸餾基線。
實驗結果
CoCoMix在效能上始終優於下一個token預測和知識蒸餾。
此外，新研究表明CoCoMix可以實現弱監督到強監督的轉換，其中從較小模型中提取的概念可以指導更強（或更大）的學生模型。
由於模型經過訓練可以預測其隱藏狀態中的概念，可以透過檢查概念預測來分析它關注哪些概念。透過放大或抑制預測的概念，我們還可以控制模型的輸出生成。
總而言之，CoCoMix效率更高，並且在不同模型規模下都優於下一個token預測，同時還引入了可解釋性。
具體而言，透過研究以下問題，對CoCoMix進行了實證評估：
-
CoCoMix能否提高LLM預訓練中下一個token預測的效能?（圖2和圖3） -
與其他知識蒸餾方法相比，CoCoMix從弱監督到強監督設定中是否表現出改進？（表1和圖4） -
CoCoMix是否引入了模型的可解釋性和可操縱性？（圖5） -
CoCoMix提出的各個元件對效能貢獻如何？（圖6）
提高NTP效能
圖2展示了CoCoMix與NTP（Next Token Prediction，下一個token預測）在不同訓練檢查點（checkpoint）的效能比較。每個模型包含總共1.38B個引數，都在OpenWebText資料集上進行訓練。對於CoCoMix，概念是從一個1.24億大小的模型（比基礎模型小10倍）中提取的。
顯示了以下方面的改進：(a)驗證困惑度，(b)在LAMBADA、WikiText-103上的平均困惑度，以及(c)在HellaSwag、PIQA、SIQA、Arc-Easy和WinoGrande上的平均準確率。
圖3展示了CoCoMix與NTP在不同模型大小下的效能比較。考慮了各種模型大小，包括69M、386M和1.38B個引數，並在200B個OpenWebText的token上進行訓練。評估了模型在OpenWebText驗證困惑度以及下游資料集LAMBADA、WikiText-103、HellaSwag、PIQA、SIQA、Arc-Easy和WinoGrande上的表現。
與知識蒸餾比較
表1展示了CoCoMix與下一token預測（NTP）與知識蒸餾（KD）的對比。報告了在OpenWebText（OWT）訓練集上的表現，以及在下游任務中的表現。訓練了三種不同規模的模型，其中124M模型作為教師模型。所有模型均在從OpenWebText資料集取樣的20B個token上進行訓練。加粗部分表示最佳結果。
圖4展示了CoCoMix與知識蒸餾(KD)的比較。對於弱監督到強監督設定，訓練一個386M的模型，其中KD的教師（或CoCoMix的概念提取器）是一個124M大小的模型：報告了(a)在OpenWebText、LAMABADA和WikiText上的平均困惑度，以及(b)在HellaSwag、PIQA、SIQA、Arc-Easy和WinoGrande資料集上的平均準確率。對於(c)分佈偏移設定，在OpenWebMath（一個數學特定的預訓練語料庫）上訓練所有方法。
可解釋性和可操縱性
圖5是概念引導效果的定性說明。CoCoMix和GPT2模型分別是350M和124M引數的Transformer，訓練資料集為OpenWebText。對於CoCoMix，透過調整預測的概念logit值z來進行操作，而對於GPT2，透過增加特定概念索引的啟用值來調整SAE概念空間c。這展示了有針對性的概念引導對各自模型輸出的影響。
各元件貢獻
圖6 對CoCoMix的分析： (a) 歸因分數在選擇概念中的有效性。 (b) 概念預測與直接隱藏狀態預測的比較（即，用連續損失預測隱藏狀態，而不是用SAE離散化隱藏狀態）。 (c) 壓縮權重的稀疏性。 (d) 透過分析概念預測和混合的貢獻進行的元件分析。 (e) 透過比較將概念向量新增到原始隱藏狀態和混合（將概念向量與token隱藏表示交替）來選擇概念條件設定的設計。 (f) CoCoMix與暫停token（即新增可學習的tokens）的比較。使用了一個69M的transformer，並且使用來自OpenWebText資料集的20B個tokens進行訓練。
另外，值得一提是，作者中有多位華人，特別是「網紅科學家」田淵棟也參與了本次論文工作。
更為詳細的實驗設定，請參閱原文。
參考資料：
https://arxiv.org/pdf/2502.08524