當前位置: 華文世界 > 科學

MLP一夜被幹掉!MIT加州理工等發現數學定理碾壓DeepMind

2024-05-05科學

編輯:桃子 LRS

【新智元導讀】 無需懷念MLP,新網路KAN基於柯爾莫哥洛夫-艾諾特定理,帶著更少的參數、更強的效能、更好的可解釋性來了,深度學習架構革新進入新時代!

一夜之間,機器學習範式要變天了!

當今,統治深度學習領域的基礎架構便是,多層感知器(MLP)——將啟用函式放置在神經元上。

那麽,除此之外,我們是否還有新的路線可走?

就在今天,來自MIT、加州理工、東北大學等機構的團隊重磅釋出了,全新的神經網路結構——Kolmogorov–Arnold Networks(KAN)。

研究人員對MLP做了一個簡單的改變,即將可學習的啟用函式從節點(神經元)移到邊(權重)上!

論文地址:https://arxiv.org/pdf/2404.19756

這個改變乍一聽似乎毫無根據,但它與數學中的「逼近理論」(approximation theories)有著相當深刻的聯系。

事實證明,Kolmogorov-Arnold表示對應兩層網路,在邊上,而非節點上,有可學習的啟用函式。

正是從表示定理得到啟發,研究人員用神經網路顯式地,將Kolmogorov-Arnold表示參數化。

值得一提的是,KAN名字的由來,是為了紀念兩位偉大的已故數學家Andrey Kolmogorov和Vladimir Arnold。

實驗結果顯示,KAN比傳統的MLP有更加優越的效能,提升了神經網路的準確性和可解釋性。

而最令人意想不到的是,KAN的視覺化和互動性,讓其在科學研究中具有潛在的套用價值,能夠幫助科學家發現新的數學和物理規律。

研究中,作者用KAN重新發現了紐結理論(knot theory)中的數學定律!

而且,KAN以更小的網路和自動化方式,復現了DeepMind在2021年的結果。

在物理方面,KAN可以幫助物理學家研究Anderson局域化(這是凝聚態物理中的一種相變)。

對了,順便提一句,研究中KAN的所有範例(除了參數掃描),在單個CPU上不到10分鐘就可以復現。

KAN的橫空出世,直接挑戰了一直以來統治機器學習領域的MLP架構,在全網掀起軒然大波。

機器學習新紀元開啟

有人直呼,機器學習的新紀元開始了!

谷歌DeepMind研究科學家稱,「Kolmogorov-Arnold再次出擊!一個鮮為人知的事實是:這個定理出現在一篇關於置換不變神經網路(深度集)的開創性論文中,展示了這種表示與集合/GNN聚合器構建方式(作為特例)之間的復雜聯系」。

一個全新的神經網路架構誕生了!KAN將極大地改變人工智慧的訓練和微調方式。

難道是AI進入了2.0時代?

還有網友用通俗的語言,將KAN和MLP的區別,做了一個形象的比喻:

Kolmogorov-Arnold網路(KAN)就像一個可以烤任何蛋糕的三層蛋糕配方,而多層感知器(MLP)是一個有不同層數的客製蛋糕。MLP更復雜但更通用,而KAN是靜態的,但針對一項任務更簡單、更快速。

論文作者,MIT教授Max Tegmark表示,最新論文表明,一種與標準神經網路完全不同的架構,在處理有趣的物理和數學問題時,以更少的參數實作了更高的精度。

接下來,一起來看看代表深度學習未來的KAN,是如何實作的?

重回牌桌上的KAN

KAN的理論基礎

柯爾莫哥洛夫-艾諾特定理(Kolmogorov–Arnold representation theorem)指出,如果f是一個定義在有界域上的多變量連續函式,那麽該函式就可以表示為多個單變量、加法連續函式的有限組合。

對於機器學習來說,該問題可以描述為:學習高維函式的過程可以簡化成學習多項式數量的一維函式。

但這些一維函式可能是非光滑的,甚至是碎形的(fractal),在實踐中可能無法學習,也正是由於這種「病態行為」,柯爾莫哥洛夫-艾諾特表示定理在機器學習領域基本上被判了「死刑」,即理論正確,但實際無用。

在這篇文章中,研究人員仍然對該定理在機器學習領域的套用持樂觀態度,並提出了兩點改進:

1、原始方程式中,只有兩層非線性和一個隱藏層(2n+1),可以將網路泛化到任意寬度和深度;

2、科學和日常生活中的大多數函式大多是光滑的,並且具有稀疏的組合結構,可能有助於形成平滑的柯爾莫哥洛夫-艾諾特表示。類似於物理學家和數學家的區別,物理學家更關註典型場景,而數學家更關心最壞情況。

KAN架構

柯爾莫哥洛夫-艾諾特網路(KAN)設計的核心思想是將多變量函式的逼近問題轉化為學習一組單變量函式的問題。在這個框架下,每個單變量函式可以用B樣條曲線來參數化,其中B樣條是一種局部的、分段的多項式曲線,其系數是可學習的。

為了把原始定理中的兩層網路擴充套件到更深、更寬,研究人員提出了一個更「泛化」的定理版本來支持設計KAN:

受MLPs層疊結構來提升網路深度的啟發,文中同樣引入了一個類似的概念,KAN層,由一個一維函式矩陣組成,每個函式都有可訓練的參數。

根據柯爾莫哥洛夫-艾諾特定理,原始的KAN層由內部函式和外部函陣列成,分別對應於不同的輸入和輸出維度,這種堆疊KAN層的設計方法不僅擴充套件了KANs的深度,而且保持了網路的可解釋性和表達能力,其中每個層都是由單變量函陣列成的,可以對函式進行單獨學習和理解。

下式中的f就等價於KAN

實作細節

雖然KAN的設計理念看起來簡單,純靠堆疊,但最佳化起來也並不容易,研究人員在訓練過程中也摸索到了一些技巧。

1、殘留誤差啟用函式:透過引入基函式b(x)和樣條函式的組合,使用殘留誤差連線的概念來構建啟用函式ϕ(x),有助於訓練過程的穩定性。

2、初始化尺度(scales):啟用函式的初始化設定為接近零的樣條函式,權重w使用Xavier初始化方法,有助於在訓練初期保持梯度的穩定。

3、更新樣條網格:由於樣條函式定義在有界區間內,而神經網路訓練過程中啟用值可能會超出這個區間,因此動態更新樣條網格可以確保樣條函式始終在合適的區間內執行。

參數量

1、網路深度:L

2、每層的寬度:N

3、每個樣條函式是基於G個區間(G+1個網格點)定義的,k階(通常k=3)

所以KANs的參數量約為

作為對比,MLP的參數量為O(L*N^2),看起來比KAN效率更高,但KANs可以使用更小的層寬度(N),不僅可以提升泛化效能,還能提升可解釋性。

KAN比MLP,勝在了哪?

效能更強

作為合理性檢驗,研究人員構造了五個已知具有平滑KA(柯爾莫哥洛夫-艾諾特)表示的例子作為驗證數據集,透過每200步增加網格點的方式對KANs進行訓練,覆蓋G的範圍為{3,5,10,20,50,100,200,500,1000}

使用不同深度和寬度的MLPs作為基線模型,並且KANs和MLPs都使用LBFGS演算法總共訓練1800步,再用RMSE作為指標進行對比。

從結果中可以看到,KAN的曲線更抖,能夠快速收斂,達到平穩狀態;並且比MLP的縮放曲線更好,尤其是在高維的情況下。

還可以看到,三層KAN的效能要遠遠強於兩層,表明更深的KANs具有更強的表達能力,符合預期。

互動解釋KAN

研究人員設計了一個簡單的回歸實驗,以展現使用者可以在與KAN的互動過程中,獲得可解釋性最強的結果。

假設使用者對於找出符號公式感興趣,總共需要經過5個互動步驟。

步驟 1:帶有稀疏化的訓練。

從全連線的KAN開始,透過帶有稀疏化正則化的訓練可以使網路變得更稀疏,從而可以發現隱藏層中,5個神經元中的4個都看起來沒什麽作用。

步驟 2:剪枝

自動剪枝後,丟棄掉所有無用的隱藏神經元,只留下一個KAN,把啟用函式匹配到已知的符號函式上。

步驟 3:設定符號函式

假設使用者可以正確地從盯著KAN圖表猜測出這些符號公式,就可以直接設定

如果使用者沒有領域知識或不知道這些啟用函式可能是哪些符號函式,研究人員提供了一個函式suggest_symbolic來建議符號候選項。

步驟 4:進一步訓練

在網路中所有的啟用函式都符號化之後,唯一剩下的參數就是仿射參數;繼續訓練仿射參數,當看到損失降到機器精度(machine precision)時,就能意識到模型已經找到了正確的符號運算式。

步驟 5:輸出符號公式

使用Sympy計算輸出節點的符號公式,驗證正確答案。

可解釋性驗證

研究人員首先在一個有監督的玩具數據集中,設計了六個樣本,展現KAN網路在符號公式下的組合結構能力。

可以看到,KAN成功學習到了正確的單變量函式,並透過視覺化的方式,可解釋地展現出KAN的思考過程。

在無監督的設定下,數據集中只包含輸入特征x,透過設計某些變量(x1, x2, x3)之間的聯系,可以測試出KAN模型尋找變量之間依賴關系的能力。

從結果來看,KAN模型成功找到了變量之間的函式依賴性,但作者也指出,目前仍然只是在合成數據上進行實驗,還需要一種更系統、更可控的方法來發現完整的關系。

帕累托最優

透過擬合特殊函式,作者展示了KAN和MLP在由模型參數數量和RMSE損失跨越的平面中的帕累托前沿(Pareto Frontier)。

在所有特殊函式中,KAN始終比MLP具有更好的帕累托前沿。

求解偏微方程式

在求解偏微方程式任務中, 研究人員繪制了預測解和真實解之間的L2平方和p平方損失。

下圖中,前兩個是損失的訓練動態,第三和第四是損失函式數量的擴充套件定律(Sacling Law)。

如下結果所示,與MLP相比,KAN的收斂速度更快,損失更低,並且具有更陡峭的擴充套件定律。

持續學習,不會發生災難性遺忘

我們都知道,災難性遺忘是機器學習中,一個嚴重的問題。

人工神經網路和大腦之間的區別在於,大腦具有放置在空間局部功能的不同模組。當學習新任務時,結構重組僅發生在負責相關技能的局部區域,而其他區域保持不變。

然而,大多數人工神經網路,包括MLP,卻沒有這種局部性概念,這可能是災難性遺忘的原因。

而研究證明了,KAN具有局部可塑性,並且可以利用樣條(splines)局部性,來避免災難性遺忘。

這個想法非常簡單,由於樣條是局部的,樣本只會影響一些附近的樣條系數,而遠處的系數保持不變。

相比之下,由於MLP通常使用全域啟用(如ReLU/Tanh/SiLU),因此,任何局部變化都可能不受控制地傳播到遠處的區域,從而破壞儲存在那裏的資訊。

研究人員采用了一維回歸任務(由5個高斯峰組成)。每個峰值周圍的數據按順序(而不是一次全部)呈現給KAN和MLP。

結果如下圖所示,KAN僅重構當前階段存在數據的區域,而使之前的區域保持不變。

而MLP在看到新的數據樣本後會重塑整個區域,從而導致災難性的遺忘。

發現紐結理論,結果超越DeepMind

KAN的誕生對於機器學習未來套用,意味著什麽?

紐結理論(Knot theory)是低維拓撲學中的一門學科,它揭示了三流形和四流形的拓撲學問題,並在生物學和拓撲量子計算等領域有著廣泛的套用。

2021年,DeepMind團隊曾首次用AI證明了紐結理論(knot theory)登上了Nature。

論文地址:https://www.nature.com/articles/s41586-021-04086-x

這項研究中,透過監督學習和人類領域專家,得出了一個與代數和幾何結不變量相關的新定理。

即梯度顯著性辨識出了監督問題的關鍵不變量,這使得領域專家提出了一個猜想,該猜想隨後得到了完善和證明。

對此,作者研究KAN是否可以在同一問題上取得良好的可解釋結果,從而預測紐結的簽名。

在DeepMind實驗中,他們研究紐結理論數據集的主要結果是:

1 利用網路歸因法發現,簽名

主要取決於中間距離

和縱向距離λ。

2 人類領域專家後來發現

與斜率有很高的相關性

並得出

為了研究問題(1),作者將17個紐結不變量視為輸入,將簽名視為輸出。

與DeepMind中的設定類似,簽名(偶數)被編碼為一熱向量,並且網路透過交叉熵損失進行訓練。

結果發現,一個極小的KAN能夠達到81.6%的測試精度,而DeepMind的4層寬度300MLP,僅達到78%的測試精度。

如下表所示,KAN (G = 3, k = 3) 有約200參數,而MLP約有300000參數量。

值得註意的是,KAN不僅更準確,而且更準確。同時比MLP的參數效率更高。

在可解釋性方面,研究人員根據每個啟用的大小來縮放其透明度,因此無需特征歸因即可立即清楚,哪些輸入變量是重要的。

然後,在三個重要變量上訓練KAN,獲得78.2%的測試準確率。

如下是,透過KAN,作者重新發現了紐結數據集中的三個數學關系。

物理Anderson局域化有解了

而在物理套用中,KAN也發揮了巨大的價值。

Anderson是一種基本現象,其中量子系統中的無序會導致電子波函式的局域化,從而使所有傳輸停止。

在一維和二維中,尺度論證表明,對於任何微小的隨機無序,所有的電子本征態都呈指數級局域化。

相比之下,在三維中,一個臨界能量形成了一個相分界,將擴充套件態和局域態分開,這被稱為移動性邊緣。

理解這些移動性邊緣對於解釋固體中的金屬-絕緣體轉變等各種基本現象至關重要,以及在光子裝置中光的局域化效應。

作者透過研究發現,KANs使得提取移動性邊緣變得非常容易,無論是數值上的,還是符號上的。

顯然,KAN已然成為科學家的得力助手、重要的合作者。

總而言之,得益於準確性、參數效率和可解釋性的優勢,KAN將是AI+Science一個有用的模型/工具。

未來,KAN的進一步在科學領域中的套用,還待挖掘。