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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

研究團隊還進行了詳細的消融實驗，以了解不同設計選擇對性能的影響：

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

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

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

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

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

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

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

研究團隊將KronSAE改進的解耦歸因于兩個互補的設計選擇：

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

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

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

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

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

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

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

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

五、學習特征的詳細分析：KronSAE如何提高可解釋性？

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

KronSAE學習的特征更加具體，體現(xiàn)在計算指標的較低值和較高的可解釋性分數(shù)上。因為后潛在比相應的預潛在更具可解釋性，研究團隊推測了編碼和檢索所需語義的隱藏機制。

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

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

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

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

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

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

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

六、結論與未來展望

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

研究團隊的分析將這些收益歸因于組合潛在結構和邏輯AND風格交互的互補效應，為稀疏性和分解如何在表示學習中協(xié)同作用提供了新的視角。

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

研究團隊確定了三個擴展這項工作的方向：

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

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

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

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