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

但這種依賴tokenization的方法，導(dǎo)致LLM「嚴(yán)重偏科」。

比如，Karpathy發(fā)現(xiàn)一個(gè)表情包相當(dāng)于53個(gè)token!

關(guān)注AI的可能也知道GPT-4o不會(huì)數(shù)字母，不知道Strawberray中有幾個(gè)字母「r」。

為了解決此類問題，最近的研究探討了超越token級(jí)信號(hào)的方法，利用更豐富的信息來訓(xùn)練模型。

比如說，一次性預(yù)測多個(gè)token；在下一個(gè)token預(yù)測之前增強(qiáng)輸入，加入豐富的信號(hào)；或用連續(xù)的潛在表示替代離散的語言標(biāo)記，提高推理效率。

Meta的下一代系統(tǒng)「大概念模型」，徹底超越token級(jí)別語言建模，直接在語句級(jí)別上語言建模，擺脫人類語言類型對模型性能的制約。

這次，受到近期研究發(fā)現(xiàn)的啟發(fā)，來自Meta的研究人員認(rèn)為稀疏自編碼器（Sparse Autoencoders，SAEs）可以捕捉高層次的語義概念，在LLM中有效地隔離出有意義的潛在特征。

由于SAEs是通過稀疏性約束訓(xùn)練，重構(gòu)模型的隱狀態(tài)，它促使模型集中關(guān)注一組緊湊的概念維度。

這可以突出預(yù)訓(xùn)練模型的概念——即支撐模型預(yù)測的核心語義方向，同時(shí)避免不必要的特征。

論文鏈接：https://arxiv.org/abs/2502.08524

新研究在多個(gè)語言建模基準(zhǔn)和預(yù)訓(xùn)練模型規(guī)模（從百萬規(guī)模到十億規(guī)模的參數(shù)模型）上進(jìn)行廣泛的評(píng)估，展示了CoCoMix的有效性。

例如，在應(yīng)用于1.38B參數(shù)模型時(shí)，CoCoMix在下一個(gè)token預(yù)測任務(wù)中的表現(xiàn)與傳統(tǒng)方法相當(dāng)，同時(shí)減少了21.5%的訓(xùn)練數(shù)據(jù)量。

此外，CoCoMix在弱監(jiān)督到強(qiáng)監(jiān)督場景中表現(xiàn)出顯著的提升，其中從小模型中提取的概念甚至可以作為真實(shí)標(biāo)簽，用于監(jiān)督大模型的訓(xùn)練。

最后，通過插入壓縮的概念向量，能夠在生成過程中探查預(yù)測的概念，從而引導(dǎo)和控制模型。

主要方法：CoCoMix

CoCoMix是一種新的LLM預(yù)訓(xùn)練框架，通過預(yù)測概念并將其混入模型的隱狀態(tài)中，以提高下一個(gè)token預(yù)測的準(zhǔn)確性。

更高的樣本效率，在下一個(gè)token預(yù)測、知識(shí)蒸餾以及插入暫停token等任務(wù)中表現(xiàn)優(yōu)越，同時(shí)提高可解釋性和可引導(dǎo)性，增強(qiáng)模型的可控性。

連續(xù)概念混合（CoCoMix）使用基于連續(xù)概念的語言建模框架。具體而言，CoCoMix包含三個(gè)步驟來學(xué)習(xí)和使用潛在概念：

1. 從預(yù)訓(xùn)練的SAE中，提取概念并選擇顯著的概念。

2. LLM從其隱藏狀態(tài)預(yù)測這些概念。

3. 一旦預(yù)測出多個(gè)概念，就將它們壓縮成一個(gè)單一的「連續(xù)概念」，并將其「混合」到LLM隱藏狀態(tài)中。

圖1：CoCoMix的概覽。

新研究證明了CoCoMix具有更高的樣本效率，并且優(yōu)于標(biāo)準(zhǔn)的下一個(gè)token預(yù)測和知識(shí)蒸餾基線。

實(shí)驗(yàn)結(jié)果

CoCoMix在性能上始終優(yōu)于下一個(gè)token預(yù)測和知識(shí)蒸餾。

此外，新研究表明CoCoMix可以實(shí)現(xiàn)弱監(jiān)督到強(qiáng)監(jiān)督的轉(zhuǎn)換，其中從較小模型中提取的概念可以指導(dǎo)更強(qiáng)（或更大）的學(xué)生模型。

由于模型經(jīng)過訓(xùn)練可以預(yù)測其隱藏狀態(tài)中的概念，可以通過檢查概念預(yù)測來分析它關(guān)注哪些概念。通過放大或抑制預(yù)測的概念，我們還可以控制模型的輸出生成。

總而言之，CoCoMix效率更高，并且在不同模型規(guī)模下都優(yōu)于下一個(gè)token預(yù)測，同時(shí)還引入了可解釋性。

具體而言，通過研究以下問題，對CoCoMix進(jìn)行了實(shí)證評(píng)估：

• CoCoMix能否提高LLM預(yù)訓(xùn)練中下一個(gè)token預(yù)測的性能?（圖2和圖3）
• 與其他知識(shí)蒸餾方法相比，CoCoMix從弱監(jiān)督到強(qiáng)監(jiān)督設(shè)置中是否表現(xiàn)出改進(jìn)？（表1和圖4）
• CoCoMix是否引入了模型的可解釋性和可操縱性？（圖5）
• CoCoMix提出的各個(gè)組件對性能貢獻(xiàn)如何？（圖6）

提高NTP性能

圖2展示了CoCoMix與NTP（Next Token Prediction，下一個(gè)token預(yù)測）在不同訓(xùn)練檢查點(diǎn)（checkpoint）的性能比較。每個(gè)模型包含總共1.38B個(gè)參數(shù)，都在OpenWebText數(shù)據(jù)集上進(jìn)行訓(xùn)練。對于CoCoMix，概念是從一個(gè)1.24億大小的模型（比基礎(chǔ)模型小10倍）中提取的。

顯示了以下方面的改進(jìn)：(a)驗(yàn)證困惑度，(b)在LAMBADA、WikiText-103上的平均困惑度，以及(c)在HellaSwag、PIQA、SIQA、Arc-Easy和WinoGrande上的平均準(zhǔn)確率。

圖3展示了CoCoMix與NTP在不同模型大小下的性能比較。考慮了各種模型大小，包括69M、386M和1.38B個(gè)參數(shù)，并在200B個(gè)OpenWebText的token上進(jìn)行訓(xùn)練。評(píng)估了模型在OpenWebText驗(yàn)證困惑度以及下游數(shù)據(jù)集LAMBADA、WikiText-103、HellaSwag、PIQA、SIQA、Arc-Easy和WinoGrande上的表現(xiàn)。

與知識(shí)蒸餾比較

表1展示了CoCoMix與下一token預(yù)測（NTP）與知識(shí)蒸餾（KD）的對比。報(bào)告了在OpenWebText（OWT）訓(xùn)練集上的表現(xiàn)，以及在下游任務(wù)中的表現(xiàn)。訓(xùn)練了三種不同規(guī)模的模型，其中124M模型作為教師模型。所有模型均在從OpenWebText數(shù)據(jù)集采樣的20B個(gè)token上進(jìn)行訓(xùn)練。加粗部分表示最佳結(jié)果。

圖4展示了CoCoMix與知識(shí)蒸餾(KD)的比較。對于弱監(jiān)督到強(qiáng)監(jiān)督設(shè)置，訓(xùn)練一個(gè)386M的模型，其中KD的教師（或CoCoMix的概念提取器）是一個(gè)124M大小的模型：報(bào)告了(a)在OpenWebText、LAMABADA和WikiText上的平均困惑度，以及(b)在HellaSwag、PIQA、SIQA、Arc-Easy和WinoGrande數(shù)據(jù)集上的平均準(zhǔn)確率。對于(c)分布偏移設(shè)置，在OpenWebMath（一個(gè)數(shù)學(xué)特定的預(yù)訓(xùn)練語料庫）上訓(xùn)練所有方法。

可解釋性和可操縱性

圖5是概念引導(dǎo)效果的定性說明。CoCoMix和GPT2模型分別是350M和124M參數(shù)的Transformer，訓(xùn)練數(shù)據(jù)集為OpenWebText。對于CoCoMix，通過調(diào)整預(yù)測的概念logit值z來進(jìn)行操作，而對于GPT2，通過增加特定概念索引的激活值來調(diào)整SAE概念空間c。這展示了有針對性的概念引導(dǎo)對各自模型輸出的影響。

各組件貢獻(xiàn)

圖6 對CoCoMix的分析：

(a) 歸因分?jǐn)?shù)在選擇概念中的有效性。

(b) 概念預(yù)測與直接隱藏狀態(tài)預(yù)測的比較（即，用連續(xù)損失預(yù)測隱藏狀態(tài)，而不是用SAE離散化隱藏狀態(tài)）。

(c) 壓縮權(quán)重的稀疏性。

(d) 通過分析概念預(yù)測和混合的貢獻(xiàn)進(jìn)行的組件分析。

(e) 通過比較將概念向量添加到原始隱藏狀態(tài)和混合（將概念向量與token隱藏表示交替）來選擇概念條件設(shè)定的設(shè)計(jì)。

(f) CoCoMix與暫停token（即添加可學(xué)習(xí)的tokens）的比較。使用了一個(gè)69M的transformer，并且使用來自O(shè)penWebText數(shù)據(jù)集的20B個(gè)tokens進(jìn)行訓(xùn)練。

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

更為詳細(xì)的實(shí)驗(yàn)設(shè)置，請參閱原文。