AI自己長出了類似大腦的腦葉?新研究揭示LLM特徵的驚人幾何結構
機器之心報道
編輯:張倩、蛋醬
論文通訊作者、MIT 物理學教授 Max Tegmark 的推文。值得注意的是,Max Tegmark 也是著名的 KAN 論文的作者之一,是 KAN 論文一作 ZimingLiu 的導師。
在過去的一年,學術界在理解大型語言模型如何工作方面取得了突破性進展:稀疏自編碼器(SAE)在其激活空間中發現了大量可解釋爲概念的點(「特徵」)。最近,此類 SAE 點雲已公開發布,因此研究其在不同尺度上的結構正當其時。
最近,來自 MIT 的一個團隊公佈了他們的研究成果。
具體來說,他們發現 SAE 特徵的概念宇宙在三個層面上具有有趣的結構:
第一個是「原子」小尺度層面。在這個層面上,作者發現 SAE 特徵的概念宇宙包含「晶體」結構,這些晶體的面是平行四邊形或梯形,這泛化了衆所周知的例子,如 (man:woman::king:queen)。他們還發現,當排除全局干擾方向,如單詞長度時,這類平行四邊形和相關功能向量的質量大大提高,這可以通過線性判別分析有效地完成。
第二個是「大腦」中等尺度層面。在這個層面,作者發現 SAE 特徵的概念宇宙具有顯著的空間模塊性。例如,數學和代碼特徵會形成一個「葉(lobe)」,類似於我們在做神經磁共振功能成像時看到的大腦功能性葉(如聽覺皮層)。作者用多個度量來量化這些葉的空間局部性,並發現在足夠粗略的尺度上,共現特徵(co-occurring feature)的聚類在空間上也聚集在一起,遠遠超過了特徵幾何是隨機的情況下的預期。
第三個是「星系」大尺度層面。在這個層面上,作者發現 SAE 特徵點雲的結構不是各向同性的,而是呈現出一種特徵值的冪律分佈,並且在中間層的斜率最陡。此外,他們還量化了聚類熵如何隨層數的變化而變化。
這項研究吸引了不少研究者的注意。有人評論說,AI 系統在處理信息時自然地發展出幾何和分形結構,而這些結構與生物大腦中的結構相似。這一現象表明,數學上的組織模式可能是自然界的基本特性,而不僅僅是人類的認知構造。
也有人提出了一些不同觀點,認爲這種結構可能更多是源於 AI 模型從人類數據中學習的結果,而不是一種完全獨立的自然特性。反駁者認爲,由於人類也是一種生物神經網絡,當大規模 AI 系統基於小規模神經網絡的輸入數據進行訓練時,它們自然而然地會接近這種結構模式,因此 AI 模型的結構並非完全出乎意料。反駁者還提出了一個有趣的設想:如果 AI 模型在完全不包含人類數據的「外星」數據集上進行訓練,那麼模型的組織結構可能會有很大的不同 —— 儘管模型仍然可能會產生聚類和分組的結構以有效處理複雜信息,但實際的概念和結構可能會和人類的完全不同。
論文作者表示,他們希望這些發現有助於大家深入理解 SAE 特徵和大型語言模型的工作原理。他們也會在未來繼續研究,以瞭解爲什麼其中一些結構會出現。
以下是論文的詳細信息。
「原子」尺度:晶體結構
在這一部分中,作者尋找他們所說的 SAE 特徵點雲中的晶體結構。這裡的結構指的是反映概念之間語義關係的幾何結構,它泛化了(a, b, c, d)=(man,woman,king,queen)形成近似平行四邊形的經典例子,其中 b − a ≈ d − c。這可以用兩個功能向量 b − a 和 c − a 來解釋,分別將男性實體轉爲女性,將普通人轉爲皇室成員。他們還尋找只有一對平行邊 b - a ∝ d - c 的梯形(只對應一個功能向量);圖 1(右)展示了這樣一個例子,其中(a, b, c, d)=(Austria, Vienna, Switzerland, Bern),這裡的功能向量可以被解釋爲將國家映射到它們的首都。
作者通過計算所有成對差分向量並對其進行聚類來尋找晶體,這應該會產生一個對應於每個功能向量的聚類。一個聚類中的任意一對差分向量應該形成一個梯形或平行四邊形,這取決於差分向量在聚類前是否被歸一化(或者是否通過歐氏距離或餘弦相似性來量化兩個差分向量之間的相似性)。
作者最初搜索 SAE 晶體時發現的大多是噪聲。爲了探究原因,他們將注意力集中在第 0 層(token 嵌入)和第 1 層,那裡許多 SAE 特徵對應於單個詞彙。然後他們研究了 Gemma2-2b 殘差流激活,這些激活是針對之前報告的來自 Todd 等人 (Todd et al., 2023) 數據集中的 word->word 功能向量,這澄清了問題。圖 1 說明了候選晶體四元組通常遠非平行四邊形或梯形。這與多篇論文指出的(man, woman, king, queen)也不是一個準確的平行四邊形是一致的。
作者發現,導致這一問題的原因是存在他們所說的干擾特徵。例如,他們發現圖 1(右)中的水平軸主要對應於單詞長度(圖 10),這在語義上是不相關的,並且對梯形(左)造成了破壞,因爲「Switzerland」比其他單詞長得多。
爲了消除這些語義上不相關的干擾向量,他們希望將數據投影到一個與它們正交的低維子空間上。對於 (Todd et al., 2023) 數據集,他們使用線性判別分析(LDA)來實現這一點,它將數據投影到信號 - 噪聲的特徵模式上,其中「信號」和「噪聲」分別定義爲聚類間變化和聚類內變化的協方差矩陣。圖 1 顯示,這極大地提高了聚類和梯形圖 / 平行四邊形的質量,突出表明干擾特徵可能會隱藏現有的晶體。
「大腦」尺度:中尺度模塊性結構
現在放大視野,尋找更大規模的結構。具體來說,作者研究了功能相似的 SAE 特徵組(傾向於一起激活)是否在幾何上也是相似的,從而在激活空間中形成「葉」。
在動物大腦中,這些功能組是衆所周知的神經元所在 3D 空間中的聚類。例如,布洛卡區參與語言產生,聽覺皮層處理聲音,而杏仁體主要與情緒處理相關。作者好奇是否能在 SAE 特徵空間中找到類似的功能模塊性。
作者測試了多種自動發現此類功能「葉」以及量化它們是否具有空間模塊性的方法。他們將葉分區定義爲將點雲分割爲 k 個子集(「葉」),這些子集的計算不依賴於位置信息。相反,他們識別這些葉的依據是它們在功能上的相關性,具體來說,就是在一個文檔中趨向於共同激活。
爲了自動識別功能葉,作者首先計算 SAE 特徵共現的直方圖。他們使用 gemma-2-2b,並將 The Pile Gao et al. (2020) 中的文檔傳遞給它。在這一部分,他們將報告第 12 層殘差流 SAE 的結果,該層有 16k 個特徵,平均 L0 爲 41。
對於這個 SAE,他們記錄了激活的特徵(如果其隱藏激活 > 1,他們認爲一個特徵被激活)。如果兩個特徵在 256 個 token 組成的同一個塊內被激活,則它們就被視爲共現 —— 這個長度提供了一個粗略的「時間分辨率」,使他們能夠找到在同一文檔中傾向於一起激活的 token。他們使用 1024 的最大上下文長度,並且每個文檔只使用一個這樣的上下文,因此每篇 The Pile 文檔最多隻能有 4 個塊(和直方圖更新)。他們計算了 50k 個文檔的直方圖。給定這個直方圖,他們基於它們的共現統計計算每對 SAE 特徵之間的親和度(affinity),並對得到的親和度矩陣進行譜聚類。
作者嘗試了以下基於共現的親和概念:簡單匹配係數、Jaccard 相似性、Dice 係數、重疊係數和 Phi 係數,所有這些都可以僅從共現直方圖計算得出。
作者的 null 假設是,功能相似的點(通常共現的 SAE 特徵)在激活空間中均勻分佈,沒有空間模塊性。相反,圖 2 顯示了看起來相當空間局部化的葉。爲了量化這一點在統計上的重要性,作者使用兩種方法來排除 null 假設。
1、雖然我們可以根據它們是否共現來聚類特徵,但也可以根據 SAE 特徵解碼向量之間的餘弦相似度執行譜聚類。給定一個使用餘弦相似度的 SAE 特徵聚類和一個使用共現的聚類,計算這兩組標籤之間的互信息。從某種意義上說,這直接衡量了人們從瞭解功能結構中獲得的幾何結構的信息量。
2、另一個概念上簡單的方法是訓練模型,從其幾何形狀預測一個特徵所在的功能葉。爲此,作者從基於共現的聚類中獲取一組給定的葉標籤,並訓練一個 logistic 迴歸模型,直接從點位置預測這些標籤,採用 80-20 的訓練 - 測試比例,並報告該分類器的平衡測試準確率。
圖 3 顯示,在這兩種測量方法中,Phi 係數勝出,在功能葉和特徵幾何之間實現了最佳對應。爲了證明這一點具有統計學意義,作者從基於餘弦相似性的聚類中隨機排列聚類標籤,並測量調整後的互信息。他們還從隨機高斯中隨機重新初始化 SAE 特徵解碼器方向並歸一化,然後訓練 logistic 迴歸模型,從這些特徵方向預測功能葉。圖 3(下)顯示,這兩個測試都以高顯著性排除了 null 假設,標準差分別爲 954 和 74,這清楚地表明作者看到的葉是真實的,而不是統計上的偶然。
爲了評估每個葉專門做什麼,作者通過 gemma-2-2b 運行來自 The Pile 的 10k 個文檔,並再次記錄第 12 層的哪些 SAE 特徵在 256 個 token 組成的塊內激活。對於每個 token 塊,他們記錄哪個葉有最高比例的特徵激活。
作者在圖 4 中展示了三個葉的結果,這些結果是使用 Phi 係數作爲共現度量的,這構成了圖 2 中葉標記的基礎。
圖 5 比較了五種不同的共現度量的效果。儘管作者發現 Phi 是最好的,但所有五種都發現了「代碼 / 數學葉」。
「星系」尺度:大規模點雲結構
在本節中,作者進一步放大視野,研究點雲的「星系」尺度結構,主要是其整體形狀和聚類,類似於天文學家研究星系形狀和亞結構的方式。
作者試圖排除的簡單 null 假設是,點雲僅僅是從一個各向同性的多元高斯分佈中抽取的。圖 6 從視覺上直觀地表明點雲的形狀並不僅僅是圓形,即使在其前三個主成分中,一些主軸也比其他的略寬,類似於人腦。
形狀分析
圖 7(左)通過展示點雲協方差矩陣的特徵值遞減來量化這一點,揭示它們並不恆定,而是似乎按照冪律下降。爲了測試這種令人驚訝的冪律是否顯著,該圖將其與從各向同性高斯分佈中抽取的點雲的相應特徵值譜進行比較,後者看起來更爲平坦,與分析預測一致:多元高斯分佈的 N 個隨機向量的協方差矩陣遵循 Wishart 分佈,這在隨機矩陣理論中得到了很好的研究。由於最小特徵值的突然下降是由數據有限引起的,並在 N → ∞的極限中消失,作者將點雲降維到其 100 個最大的主成分進行後續的所有分析。換句話說,點雲的形狀像是一個「分形黃瓜」,在連續的維度中寬度按照冪律下降。作者發現這種冪律縮放對於激活來說明顯不如對於 SAE 特徵那麼突出;進一步研究其起源將很有趣。
圖 7(右)顯示了上述冪律斜率如何取決於 LLM 層,計算方法是對 100 個最大特徵值進行線性迴歸。可以看到一個明顯的模式,即中間層具有最陡峭的冪律斜率:(第 12 層的斜率爲 - 0.47,而前面和後面的層(如第 0 層和第 24 層)的斜率較淺(分別爲 - 0.24 和 - 0.25)。這可能暗示了中間層起到了瓶頸的作用,將信息壓縮爲較少的主成分,或許是爲了更有效地表示高層次抽象概念而進行的優化。圖 7(右)還顯示了有效雲體積(協方差矩陣的行列式)如何依賴於層(在對數尺度上。
聚類分析
星系或微觀粒子的聚類通常以功率譜或相關函數來量化。對於論文中的高維數據來說,這種方法比較複雜,因爲基本密度隨半徑變化,對於高維高斯分佈來說,基本密度主要集中在一個相對較薄的球殼周圍。因此,作者通過估算點雲採樣分佈的熵來量化聚類。他們使用 k-NN 方法從 SAE 特徵點雲估計熵 H,計算如下,
對於具有相同協方差矩陣的高斯分佈,熵計算爲:
圖 8 顯示了不同層的估計聚類熵。作者發現 SAE 點雲特別在中間層有很強的聚類。
這項研究的結果非常有趣,如果你有關於這篇論文的見解,歡迎在評論區留言。