在人工智能領(lǐng)域，尤其是大型語言模型的研究中，我們一直面臨著一個重要挑戰(zhàn)：如何理解這些復(fù)雜模型內(nèi)部的工作機(jī)制？這就像是擁有了一臺神奇的機(jī)器，它能完成各種驚人的任務(wù)，但我們卻不完全明白它內(nèi)部的齒輪是如何運轉(zhuǎn)的。2025年5月28日，來自T-Tech和莫斯科物理技術(shù)學(xué)院的研究團(tuán)隊——Vadim Kurochkin、Yaroslav Aksenov、Daniil Laptev、Daniil Gavrilov和Nikita Balagansky發(fā)表了一篇題為《通過利用特征相關(guān)性高效訓(xùn)練稀疏自編碼器》的研究論文（arXiv:2505.22255v1），為解決這一難題提供了創(chuàng)新方法。

稀疏自編碼器（Sparse Autoencoders，簡稱SAE）是近年來備受關(guān)注的工具，它們能夠幫助我們解釋語言模型的隱藏狀態(tài)，將復(fù)雜的神經(jīng)激活分解成人類可理解的潛在方向。想象一下，這就像是把一束混合的光線通過棱鏡分解成不同顏色的光譜，讓我們能夠看清每種顏色的特性。然而，使用SAE面臨一個主要障礙：當(dāng)我們需要處理現(xiàn)代大型語言模型時，它們的計算成本非常高，尤其是當(dāng)我們使用大型字典（即更多的"棱鏡面"）來獲得更細(xì)致的分解時。

雖然在SAE的解碼器部分已經(jīng)有了一些高效的方法，但編碼器仍然需要進(jìn)行計算密集型的線性操作，這嚴(yán)重限制了SAE在大規(guī)模模型上的應(yīng)用。這就像是我們優(yōu)化了信息輸出的管道，但信息輸入的管道仍然是個瓶頸，大量的數(shù)據(jù)在這里堵塞，影響了整個系統(tǒng)的效率。

針對這一問題，研究團(tuán)隊提出了一個名為KronSAE的創(chuàng)新架構(gòu)。這個架構(gòu)的核心思想是利用克羅內(nèi)克積分解（Kronecker product decomposition）來分解潛在表示，大幅減少內(nèi)存和計算開銷。此外，他們還引入了一個名為mAND的可微分激活函數(shù)，它近似模擬二進(jìn)制AND操作，在他們的分解框架中提高了可解釋性和性能。

一、KronSAE：如何打破編碼器瓶頸？

傳統(tǒng)的稀疏自編碼器面臨著一個明顯的瓶頸：編碼器投影。想象一下，這就像是通過一個巨大的漏斗將信息壓縮——漏斗越大，需要的力氣就越大。在技術(shù)術(shù)語中，這意味著對于輸入維度為d、字典大小為F的SAE，每次前向傳遞都需要O(Fd)的計算復(fù)雜度，這在現(xiàn)代變換器模型上是非常昂貴的。

KronSAE通過一個巧妙的架構(gòu)設(shè)計解決了這個問題。它將潛在空間分解為h個獨立的組件（可以想象為不同的"頭部"），每個頭部k由兩個薄矩陣參數(shù)化：一個是"組合基礎(chǔ)"Pk∈R^(m×d)，另一個是"組合擴(kuò)展"Qk∈R^(n×d)，其中維度m < n << d，總字典大小F = h×m×n。

這種分解方法的工作原理類似于將一個復(fù)雜的拼圖分解成多個小塊，每個小塊都更容易處理。具體來說，系統(tǒng)首先計算兩個預(yù)潛在表示：

pk = ReLU(Pkx) qk = ReLU(Qkx)

這些預(yù)潛在表示通過一個元素級的交互核心（mAND）在每個頭部中獨立組合：

zki,j := mAND(pki, qkj) := { √(pki*qkj), 如果pki > 0且qkj > 0 0, 否則

這里的mAND核心平滑地近似布爾AND門，確保只有當(dāng)兩個輸入都為正時才產(chǎn)生非零激活，同時保持梯度流和激活幅度以實現(xiàn)穩(wěn)定的重建。

最后，我們扁平化并連接所有頭部的結(jié)果，得到后潛在表示f∈R^F，然后應(yīng)用通常的TopK操作（即保留k個最大激活值）。

通過這種方法，每個token的編碼器成本從O(Fd)下降到O(h(m+n)d)，顯著減少了FLOPs和參數(shù)數(shù)量，而不需要像其他方法那樣引入路由開銷。更重要的是，KronSAE與現(xiàn)有的稀疏解碼器核心兼容，因此可以與它們結(jié)合使用，實現(xiàn)端到端的加速。

二、實驗驗證：KronSAE的效果如何？

研究團(tuán)隊在Qwen-2.5-1.5B-Base、Pythia-1.4B-deduped和Gemma-2-2B等語言模型上進(jìn)行了廣泛的實驗。他們使用了FINEWEB-EDU（FineWeb語料庫中經(jīng)過過濾的教育網(wǎng)頁子集）收集激活值，并通過解釋方差（EV）來衡量重建質(zhì)量，其中1.0為最佳。

在等效計算預(yù)算下，研究團(tuán)隊比較了KronSAE和傳統(tǒng)的TopK SAE的性能。實驗結(jié)果令人印象深刻：

在100M token預(yù)算下，所有KronSAE變體在重建質(zhì)量上超過了TopK基線，解釋方差提高了高達(dá)4.3%，同時參數(shù)數(shù)量減少了約54.7%。這就像是用更少的零件構(gòu)建了一個更高效的機(jī)器。

在500M token預(yù)算下，KronSAE在大多數(shù)模型大小上保持了0.8%的優(yōu)勢，參數(shù)減少了43.8%。較小的構(gòu)成基礎(chǔ)維度（m=2）在緊張的計算約束下表現(xiàn)尤為出色。

在1000M token預(yù)算下，盡管TopK SAE縮小了差距，KronSAE仍然能夠匹配基線的重建質(zhì)量，同時參數(shù)數(shù)量減少了約46.1%。

研究團(tuán)隊還進(jìn)行了詳細(xì)的消融實驗，以了解不同設(shè)計選擇對性能的影響：

首先，他們比較了mAND操作與兩個更簡單的交互核心：ReLU(u)·ReLU(v)和原始乘積u·v。結(jié)果顯示，在1B token訓(xùn)練預(yù)算下，mAND變體始終實現(xiàn)最高的解釋方差，明顯優(yōu)于替代方案。

其次，他們系統(tǒng)地改變頭部數(shù)量h和每個頭部的基礎(chǔ)維度m（同時保持n = F/(mh)）。結(jié)果表明，對于大型訓(xùn)練預(yù)算（500M-1B token），較小的m（因此較大的n）產(chǎn)生更高的重建質(zhì)量，因為較小的基礎(chǔ)維度釋放了容量用于更具表現(xiàn)力的擴(kuò)展特征。在更緊張的100M token預(yù)算下，m=4的配置優(yōu)于m=2或m=8，表明每個頭部的表示豐富性與數(shù)據(jù)效率之間存在權(quán)衡。此外，固定m并增加h幾乎線性地提高了解釋方差。

最后，他們評估了KronSAE在不同稀疏性水平和層深度上的穩(wěn)健性。在所有情況下，KronSAE在相同F(xiàn)LOPs預(yù)算下匹配或超過了TopK基線的重建質(zhì)量，證明了其克羅內(nèi)克分解編碼器無論稀疏性水平或深度如何都能保持其優(yōu)勢。

三、特征吸收：KronSAE如何改善特征解釋性？

在可解釋性研究中，一個主要挑戰(zhàn)是"特征吸收"，即一個學(xué)習(xí)特征成為另一個特征的嚴(yán)格子集（例如，"獅子"特征完全被"以L開頭"特征包含），因此無法在滿足更廣泛概念但不滿足其超集表示的實例上激活。

研究團(tuán)隊報告了三個吸收指標(biāo)：平均吸收分?jǐn)?shù)（部分被吸收的特征比例）、平均完全吸收分?jǐn)?shù)（量化完全包含事件）和平均特征分裂數(shù)（單個概念特征分裂成多個激活的頻率）。

實驗結(jié)果表明，在所有稀疏性水平l0∈{16, 32, 64, 128, 256}上，KronSAE變體始終減少了相對于TopK SAE基線的吸收分?jǐn)?shù)和完全吸收分?jǐn)?shù)，同時保持類似的特征分裂率。

研究團(tuán)隊將KronSAE改進(jìn)的解耦歸因于兩個互補(bǔ)的設(shè)計選擇：

1. 平滑mAND激活：通過僅在兩個預(yù)潛在值都為正時才輸出非零值，它引入了一個可微分AND門，防止廣泛的多語義原語完全包含更具體的原語。因此，復(fù)合后潛在主要在其構(gòu)成概念的交集處觸發(fā)，鼓勵每個預(yù)潛在專注于單一語義模式，而不是繼承其"父"激活區(qū)域。

2. 頭部式笛卡爾分解：將潛在空間劃分為h個獨立的子空間（每個子空間都有自己的m×n原語交互網(wǎng)格）確保專門的概念（如"大象"）被限制在單個頭部中，不能完全吸收其他頭部中的更一般概念（如"以E開頭"）。

這些機(jī)制共同產(chǎn)生了更多的單語義特征，簡化了下游因果干預(yù)和目標(biāo)探測。值得注意的是，平均特征分裂數(shù)與TopK基線保持相當(dāng)，因為笛卡爾分解本身不會固有地改變單個原語的分裂。

四、深入分析：KronSAE如何模擬特征關(guān)聯(lián)？

為了評估不同稀疏自編碼器架構(gòu)如何恢復(fù)已知的相關(guān)模式，研究團(tuán)隊構(gòu)建了一個受控實驗，使用合成的、塊結(jié)構(gòu)化的協(xié)方差模型。他們生成帶有不同塊結(jié)構(gòu)的輸入向量，然后訓(xùn)練自編碼器重建它們，檢查所學(xué)習(xí)的特征相關(guān)性。

結(jié)果顯示，KronSAE的解碼器權(quán)重協(xié)方差Cdec = WdecW?dec比TopK SAE更忠實地再現(xiàn)了地面真相分組。特別是，在第三種協(xié)方差模式（其中一些塊非常?。┥?，TopK的學(xué)習(xí)相關(guān)性幾乎消失，而KronSAE仍然發(fā)現(xiàn)了正確的塊結(jié)構(gòu)。

這些觀察通過RV系數(shù)和排列測試得到了量化。即使在最佳匹配TopK原子到密集AE參考后，TopK SAE也只能實現(xiàn)微弱的相關(guān)性對齊（RV≈0.05-0.08），p值不顯著或邊緣顯著。相比之下，KronSAE配置實現(xiàn)了0.11到0.35之間的RV值（所有p<0.001），在相關(guān)性恢復(fù)方面提高了3-6倍。

此外，研究團(tuán)隊分析了在語言數(shù)據(jù)上訓(xùn)練的SAE中的特征相關(guān)性，發(fā)現(xiàn)KronSAE中一個頭部內(nèi)的特征相關(guān)性確實顯著高于不同頭部之間的特征，這表明他們的設(shè)計成功地在SAE潛在空間中施加了所需的相關(guān)結(jié)構(gòu)。

五、學(xué)習(xí)特征的詳細(xì)分析：KronSAE如何提高可解釋性？

研究團(tuán)隊對KronSAE和TopK架構(gòu)進(jìn)行了深入的可解釋性分析，使用自動化管道解釋激活模式，并通過檢測分?jǐn)?shù)和模糊分?jǐn)?shù)評估所獲得的解釋。

KronSAE學(xué)習(xí)的特征更加具體，體現(xiàn)在計算指標(biāo)的較低值和較高的可解釋性分?jǐn)?shù)上。因為后潛在比相應(yīng)的預(yù)潛在更具可解釋性，研究團(tuán)隊推測了編碼和檢索所需語義的隱藏機(jī)制。

通過檢查激活示例和潛在解釋，他們觀察到預(yù)潛在可能攜帶多個不同的、可識別的激活模式。例如，頭部23中的組合基礎(chǔ)元素3展示了與比較描述符、地理區(qū)域和精神概念相關(guān)的子語義。預(yù)潛在的多語義性是減少"工作"編碼器潛在數(shù)量的預(yù)期結(jié)果，因為它們分解了完整的字典大小并減少了編碼器容量。

研究團(tuán)隊假設(shè)特定語義的編碼可能通過幅度來完成，這通過檢查激活示例得到了驗證。例如，上述預(yù)潛在中，"比較"部分編碼在前75%分位數(shù)，而"精神"部分主要在前25%分位數(shù)中找到，"地理"部分主要編碼在四分位范圍內(nèi)。

頭部通常包含語義相關(guān)的預(yù)潛在組，例如，頭部136包含三個基礎(chǔ)元素和一個與數(shù)字和序數(shù)相關(guān)的擴(kuò)展，兩個與地理和空間事物相關(guān)的擴(kuò)展元素，一個與問題相關(guān)的基礎(chǔ)和一個與增長相關(guān)的擴(kuò)展。有趣的是，該頭部的大多數(shù)后潛在具有比其父預(yù)潛在更高的可解釋性分?jǐn)?shù)，這是不尋常的。

檢索主要通過類似邏輯AND電路的機(jī)制發(fā)生，其中一些預(yù)潛在作為多個語義的載體，相應(yīng)的預(yù)潛在（基礎(chǔ)或擴(kuò)展）作為指定器。例如，在基礎(chǔ)包含三個可檢測的子語義的情況下，每個擴(kuò)展然后檢索特定的語義。

其他類型的交互也可能發(fā)生，例如完全新的語義的出現(xiàn)，如頭部23中基礎(chǔ)3和擴(kuò)展1之間的組合，其中出現(xiàn)了醫(yī)學(xué)術(shù)語，不能簡單地解釋為兩個預(yù)潛在語義的交集。

經(jīng)常出現(xiàn)的一種情況是后潛在只繼承一個父語義，或者另一個父語義的影響無法檢測到，這通常發(fā)生在父語義具有非常廣泛的解釋和低分?jǐn)?shù)時。然而，需要更復(fù)雜的技術(shù)來正確識別交互的細(xì)粒度結(jié)構(gòu)。

在后潛在的幾何方面，每個后潛在向量在殘差流中都有一個向量表示，即Wdec中的相應(yīng)列，這是我們在訓(xùn)練SAE時尋找的過完備基向量的近似。研究團(tuán)隊沒有觀察到TopK和KronSAE之間的特征幾何有任何顯著差異，除了KronSAE的架構(gòu)設(shè)計導(dǎo)致聚類，使得由相同頭部、基礎(chǔ)或擴(kuò)展元素產(chǎn)生的后潛在被分組在一個緊密的簇中，幾何結(jié)構(gòu)取決于我們選擇的超參數(shù)h、m、n，這是預(yù)期的，可能對進(jìn)一步應(yīng)用如引導(dǎo)很有用。

六、結(jié)論與未來展望

KronSAE代表了稀疏自編碼器設(shè)計的重要進(jìn)步，通過頭部式克羅內(nèi)克分解和mAND門控直接解決了長期存在的編碼器效率瓶頸。相比標(biāo)準(zhǔn)的TopK SAE，它顯著降低了參數(shù)數(shù)量，同時提高了重建保真度，并通過利用特征相關(guān)性產(chǎn)生了更具解釋性的特征。

研究團(tuán)隊的分析將這些收益歸因于組合潛在結(jié)構(gòu)和邏輯AND風(fēng)格交互的互補(bǔ)效應(yīng)，為稀疏性和分解如何在表示學(xué)習(xí)中協(xié)同作用提供了新的視角。

盡管有這些優(yōu)點，KronSAE也有一些限制。它的收益取決于對(m, n, h)和mAND激活的仔細(xì)調(diào)整——配置不當(dāng)?shù)脑O(shè)置可能會抵消效率和質(zhì)量的改進(jìn)。該評估僅限于中型變換器模型和單一網(wǎng)絡(luò)語料庫，因此它在更大的模型、其他領(lǐng)域或語言上的適用性仍有待測試。

研究團(tuán)隊確定了三個擴(kuò)展這項工作的方向：

1. 轉(zhuǎn)碼：將轉(zhuǎn)碼器視為信息的隱式路由器，研究替代邏輯門控函數(shù)（例如XOR或復(fù)合門）以提高可解釋性和電路分析。

2. 交叉編碼：將KronSAE推廣到交叉編碼器設(shè)置，通過邏輯操作揭示可解釋的、跨層次的組合性。

3. 動態(tài)組合：探索對注意力頭數(shù)量及其維度進(jìn)行可學(xué)習(xí)調(diào)整，實現(xiàn)在不同尺度上對相關(guān)特征組進(jìn)行細(xì)粒度分解。

總之，KronSAE為大型語言模型的解釋提供了一個強(qiáng)大而高效的工具，有望推動我們對這些復(fù)雜系統(tǒng)內(nèi)部工作原理的理解。通過減少計算負(fù)擔(dān)并提高特征質(zhì)量，它使研究人員能夠更深入地探索模型內(nèi)部，最終可能導(dǎo)致更透明、更可控的AI系統(tǒng)。