「下一個token預測」（next token prediction，NTP）是大語言模型（LLMs）不斷取得突破的核心技術。

但這種依賴tokenization的方法，導致LLM「嚴重偏科」。

比如，Karpathy發現一個表情包相當于53個token!

關注AI的可能也知道GPT-4o不會數字母，不知道Strawberray中有幾個字母「r」。

為了解決此類問題，最近的研究探討了超越token級信號的方法，利用更豐富的信息來訓練模型。

比如說，一次性預測多個token；在下一個token預測之前增強輸入，加入豐富的信號；或用連續的潛在表示替代離散的語言標記，提高推理效率。

Meta的下一代系統「大概念模型」，徹底超越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進行訓練。

另外，值得一提是，作者中有多位華人，特別是「網紅科學家」田淵棟也參與了本次論文工作。

更為詳細的實驗設置，請參閱原文。