大型語言模型在學習概念時竟然會形成令人驚訝的幾何結構，比如代碼和數(shù)學特征會形成一個「葉（lobe）」，類似于我們在做磁共振功能成像時看到的大腦功能性腦葉。這說明什么呢？

論文通訊作者、MIT 物理學教授 Max Tegmark 的推文。值得注意的是，Max Tegmark 也是著名的 KAN 論文的作者之一，是 KAN 論文一作 ZimingLiu 的導師。

在過去的一年，學術界在理解大型語言模型如何工作方面取得了突破性進展：稀疏自編碼器（SAE）在其激活空間中發(fā)現(xiàn)了大量可解釋為概念的點（「特征」）。最近，此類 SAE 點云已公開發(fā)布，因此研究其在不同尺度上的結構正當其時。

最近，來自 MIT 的一個團隊公布了他們的研究成果。

• 論文標題：The Geometry of Concepts: Sparse Autoencoder Feature Structure
• 論文鏈接：https://arxiv.org/pdf/2410.19750

具體來說，他們發(fā)現(xiàn) SAE 特征的概念宇宙在三個層面上具有有趣的結構：

第一個是「原子」小尺度層面。在這個層面上，作者發(fā)現(xiàn) SAE 特征的概念宇宙包含「晶體」結構，這些晶體的面是平行四邊形或梯形，這泛化了眾所周知的例子，如 (man:woman::king:queen)。他們還發(fā)現(xiàn)，當排除全局干擾方向，如單詞長度時，這類平行四邊形和相關功能向量的質量大大提高，這可以通過線性判別分析有效地完成。

第二個是「大腦」中等尺度層面。在這個層面，作者發(fā)現(xiàn) SAE 特征的概念宇宙具有顯著的空間模塊性。例如，數(shù)學和代碼特征會形成一個「葉（lobe）」，類似于我們在做神經(jīng)磁共振功能成像時看到的大腦功能性葉（如聽覺皮層）。作者用多個度量來量化這些葉的空間局部性，并發(fā)現(xiàn)在足夠粗略的尺度上，共現(xiàn)特征（co-occurring feature）的聚類在空間上也聚集在一起，遠遠超過了特征幾何是隨機的情況下的預期。

第三個是「星系」大尺度層面。在這個層面上，作者發(fā)現(xiàn) SAE 特征點云的結構不是各向同性的，而是呈現(xiàn)出一種特征值的冪律分布，并且在中間層的斜率最陡。此外，他們還量化了聚類熵如何隨層數(shù)的變化而變化。

這項研究吸引了不少研究者的注意。有人評論說，AI 系統(tǒng)在處理信息時自然地發(fā)展出幾何和分形結構，而這些結構與生物大腦中的結構相似。這一現(xiàn)象表明，數(shù)學上的組織模式可能是自然界的基本特性，而不僅僅是人類的認知構造。

也有人提出了一些不同觀點，認為這種結構可能更多是源于 AI 模型從人類數(shù)據(jù)中學習的結果，而不是一種完全獨立的自然特性。反駁者認為，由于人類也是一種生物神經(jīng)網(wǎng)絡，當大規(guī)模 AI 系統(tǒng)基于小規(guī)模神經(jīng)網(wǎng)絡的輸入數(shù)據(jù)進行訓練時，它們自然而然地會接近這種結構模式，因此 AI 模型的結構并非完全出乎意料。反駁者還提出了一個有趣的設想：如果 AI 模型在完全不包含人類數(shù)據(jù)的「外星」數(shù)據(jù)集上進行訓練，那么模型的組織結構可能會有很大的不同 —— 盡管模型仍然可能會產生聚類和分組的結構以有效處理復雜信息，但實際的概念和結構可能會和人類的完全不同。

論文作者表示，他們希望這些發(fā)現(xiàn)有助于大家深入理解 SAE 特征和大型語言模型的工作原理。他們也會在未來繼續(xù)研究，以了解為什么其中一些結構會出現(xiàn)。

以下是論文的詳細信息。

「原子」尺度：晶體結構

在這一部分中，作者尋找他們所說的 SAE 特征點云中的晶體結構。這里的結構指的是反映概念之間語義關系的幾何結構，它泛化了（a, b, c, d）=（man,woman,king,queen）形成近似平行四邊形的經(jīng)典例子，其中 b ? a ≈ d ? c。這可以用兩個功能向量 b ? a 和 c ? a 來解釋，分別將男性實體轉為女性，將普通人轉為皇室成員。他們還尋找只有一對平行邊 b - a ∝ d - c 的梯形（只對應一個功能向量）；圖 1（右）展示了這樣一個例子，其中（a, b, c, d）=（Austria, Vienna, Switzerland, Bern），這里的功能向量可以被解釋為將國家映射到它們的首都。

作者通過計算所有成對差分向量并對其進行聚類來尋找晶體，這應該會產生一個對應于每個功能向量的聚類。一個聚類中的任意一對差分向量應該形成一個梯形或平行四邊形，這取決于差分向量在聚類前是否被歸一化（或者是否通過歐氏距離或余弦相似性來量化兩個差分向量之間的相似性）。

作者最初搜索 SAE 晶體時發(fā)現(xiàn)的大多是噪聲。為了探究原因，他們將注意力集中在第 0 層（token 嵌入）和第 1 層，那里許多 SAE 特征對應于單個詞匯。然后他們研究了 Gemma2-2b 殘差流激活，這些激活是針對之前報告的來自 Todd 等人 (Todd et al., 2023) 數(shù)據(jù)集中的 word->word 功能向量，這澄清了問題。圖 1 說明了候選晶體四元組通常遠非平行四邊形或梯形。這與多篇論文指出的（man, woman, king, queen）也不是一個準確的平行四邊形是一致的。

作者發(fā)現(xiàn)，導致這一問題的原因是存在他們所說的干擾特征。例如，他們發(fā)現(xiàn)圖 1（右）中的水平軸主要對應于單詞長度（圖 10），這在語義上是不相關的，并且對梯形（左）造成了破壞，因為「Switzerland」比其他單詞長得多。

為了消除這些語義上不相關的干擾向量，他們希望將數(shù)據(jù)投影到一個與它們正交的低維子空間上。對于 (Todd et al., 2023) 數(shù)據(jù)集，他們使用線性判別分析（LDA）來實現(xiàn)這一點，它將數(shù)據(jù)投影到信號 - 噪聲的特征模式上，其中「信號」和「噪聲」分別定義為聚類間變化和聚類內變化的協(xié)方差矩陣。圖 1 顯示，這極大地提高了聚類和梯形圖 / 平行四邊形的質量，突出表明干擾特征可能會隱藏現(xiàn)有的晶體。

「大腦」尺度：中尺度模塊性結構

現(xiàn)在放大視野，尋找更大規(guī)模的結構。具體來說，作者研究了功能相似的 SAE 特征組（傾向于一起激活）是否在幾何上也是相似的，從而在激活空間中形成「葉」。

在動物大腦中，這些功能組是眾所周知的神經(jīng)元所在 3D 空間中的聚類。例如，布洛卡區(qū)參與語言產生，聽覺皮層處理聲音，而杏仁體主要與情緒處理相關。作者好奇是否能在 SAE 特征空間中找到類似的功能模塊性。

作者測試了多種自動發(fā)現(xiàn)此類功能「葉」以及量化它們是否具有空間模塊性的方法。他們將葉分區(qū)定義為將點云分割為 k 個子集（「葉」），這些子集的計算不依賴于位置信息。相反，他們識別這些葉的依據(jù)是它們在功能上的相關性，具體來說，就是在一個文檔中趨向于共同激活。

為了自動識別功能葉，作者首先計算 SAE 特征共現(xiàn)的直方圖。他們使用 gemma-2-2b，并將 The Pile Gao et al. (2020) 中的文檔傳遞給它。在這一部分，他們將報告第 12 層殘差流 SAE 的結果，該層有 16k 個特征，平均 L0 為 41。

對于這個 SAE，他們記錄了激活的特征（如果其隱藏激活 > 1，他們認為一個特征被激活）。如果兩個特征在 256 個 token 組成的同一個塊內被激活，則它們就被視為共現(xiàn) —— 這個長度提供了一個粗略的「時間分辨率」，使他們能夠找到在同一文檔中傾向于一起激活的 token。他們使用 1024 的最大上下文長度，并且每個文檔只使用一個這樣的上下文，因此每篇 The Pile 文檔最多只能有 4 個塊（和直方圖更新）。他們計算了 50k 個文檔的直方圖。給定這個直方圖，他們基于它們的共現(xiàn)統(tǒng)計計算每對 SAE 特征之間的親和度（affinity），并對得到的親和度矩陣進行譜聚類。

作者嘗試了以下基于共現(xiàn)的親和概念：簡單匹配系數(shù)、Jaccard 相似性、Dice 系數(shù)、重疊系數(shù)和 Phi 系數(shù)，所有這些都可以僅從共現(xiàn)直方圖計算得出。

作者的 null 假設是，功能相似的點（通常共現(xiàn)的 SAE 特征）在激活空間中均勻分布，沒有空間模塊性。相反，圖 2 顯示了看起來相當空間局部化的葉。為了量化這一點在統(tǒng)計上的重要性，作者使用兩種方法來排除 null 假設。

• 雖然我們可以根據(jù)它們是否共現(xiàn)來聚類特征，但也可以根據(jù) SAE 特征解碼向量之間的余弦相似度執(zhí)行譜聚類。給定一個使用余弦相似度的 SAE 特征聚類和一個使用共現(xiàn)的聚類，計算這兩組標簽之間的互信息。從某種意義上說，這直接衡量了人們從了解功能結構中獲得的幾何結構的信息量。
• 另一個概念上簡單的方法是訓練模型，從其幾何形狀預測一個特征所在的功能葉。為此，作者從基于共現(xiàn)的聚類中獲取一組給定的葉標簽，并訓練一個 logistic 回歸模型，直接從點位置預測這些標簽，采用 80-20 的訓練 - 測試比例，并報告該分類器的平衡測試準確率。

圖 3 顯示，在這兩種測量方法中，Phi 系數(shù)勝出，在功能葉和特征幾何之間實現(xiàn)了最佳對應。為了證明這一點具有統(tǒng)計學意義，作者從基于余弦相似性的聚類中隨機排列聚類標簽，并測量調整后的互信息。他們還從隨機高斯中隨機重新初始化 SAE 特征解碼器方向并歸一化，然后訓練 logistic 回歸模型，從這些特征方向預測功能葉。圖 3（下）顯示，這兩個測試都以高顯著性排除了 null 假設，標準差分別為 954 和 74，這清楚地表明作者看到的葉是真實的，而不是統(tǒng)計上的偶然。

為了評估每個葉專門做什么，作者通過 gemma-2-2b 運行來自 The Pile 的 10k 個文檔，并再次記錄第 12 層的哪些 SAE 特征在 256 個 token 組成的塊內激活。對于每個 token 塊，他們記錄哪個葉有最高比例的特征激活。

作者在圖 4 中展示了三個葉的結果，這些結果是使用 Phi 系數(shù)作為共現(xiàn)度量的，這構成了圖 2 中葉標記的基礎。

圖 5 比較了五種不同的共現(xiàn)度量的效果。盡管作者發(fā)現(xiàn) Phi 是最好的，但所有五種都發(fā)現(xiàn)了「代碼 / 數(shù)學葉」。

「星系」尺度：大規(guī)模點云結構

在本節(jié)中，作者進一步放大視野，研究點云的「星系」尺度結構，主要是其整體形狀和聚類，類似于天文學家研究星系形狀和亞結構的方式。

作者試圖排除的簡單 null 假設是，點云僅僅是從一個各向同性的多元高斯分布中抽取的。圖 6 從視覺上直觀地表明點云的形狀并不僅僅是圓形，即使在其前三個主成分中，一些主軸也比其他的略寬，類似于人腦。

形狀分析

圖 7（左）通過展示點云協(xié)方差矩陣的特征值遞減來量化這一點，揭示它們并不恒定，而是似乎按照冪律下降。為了測試這種令人驚訝的冪律是否顯著，該圖將其與從各向同性高斯分布中抽取的點云的相應特征值譜進行比較，后者看起來更為平坦，與分析預測一致：多元高斯分布的 N 個隨機向量的協(xié)方差矩陣遵循 Wishart 分布，這在隨機矩陣理論中得到了很好的研究。由于最小特征值的突然下降是由數(shù)據(jù)有限引起的，并在 N → ∞的極限中消失，作者將點云降維到其 100 個最大的主成分進行后續(xù)的所有分析。換句話說，點云的形狀像是一個「分形黃瓜」，在連續(xù)的維度中寬度按照冪律下降。作者發(fā)現(xiàn)這種冪律縮放對于激活來說明顯不如對于 SAE 特征那么突出；進一步研究其起源將很有趣。

圖 7（右）顯示了上述冪律斜率如何取決于 LLM 層，計算方法是對 100 個最大特征值進行線性回歸。可以看到一個明顯的模式，即中間層具有最陡峭的冪律斜率：（第 12 層的斜率為 - 0.47，而前面和后面的層（如第 0 層和第 24 層）的斜率較淺（分別為 - 0.24 和 - 0.25）。這可能暗示了中間層起到了瓶頸的作用，將信息壓縮為較少的主成分，或許是為了更有效地表示高層次抽象概念而進行的優(yōu)化。圖 7（右）還顯示了有效云體積（協(xié)方差矩陣的行列式）如何依賴于層（在對數(shù)尺度上。

聚類分析

星系或微觀粒子的聚類通常以功率譜或相關函數(shù)來量化。對于論文中的高維數(shù)據(jù)來說，這種方法比較復雜，因為基本密度隨半徑變化，對于高維高斯分布來說，基本密度主要集中在一個相對較薄的球殼周圍。因此，作者通過估算點云采樣分布的熵來量化聚類。他們使用 k-NN 方法從 SAE 特征點云估計熵 H，計算如下，

對于具有相同協(xié)方差矩陣的高斯分布，熵計算為：

圖 8 顯示了不同層的估計聚類熵。作者發(fā)現(xiàn) SAE 點云特別在中間層有很強的聚類。