code window

2024年5月4日星期六

Hello, KAN!

多層感知器(MLPs),也被稱為全連接前饋神經網絡,是當今深度學習模型的基礎組成部分。



MLPs在機器學習中扮演著至關重要的角色,因為它們是用於近似非線性函數的默認模型,這得益於通用近似定理所保證的表達能力。然而,MLPs真的是我們能構建的最佳非線性回歸器嗎?儘管MLPs被廣泛使用,但它們仍然存在一些顯著缺陷。例如,在Transformer模型中,MLPs幾乎消耗了所有非嵌入參數,並且在沒有使用後續分析工具的情況下,通常不如注意力層那樣易於解釋。

為了解決這些問題,本研究提出了一種MLPs的替代方案,稱為Kolmogorov-Arnold Networks(KANs)。與受通用近似定理啟發的MLPs不同,KANs的靈感來自於Kolmogorov-Arnold定理。KANs與MLPs一樣具有全連接結構,但兩者在激活函數的放置方式上有所不同。MLPs在節點("神經元")上放置固定的激活函數,而KANs則在邊("權重")上放置可學習的激活函數。因此,KANs完全沒有線性權重矩陣,取而代之的是每個權重參數都被一個以樣條線參數化的可學習一維函數所替換。KANs的節點只是簡單地對傳入信號進行求和,而不應用任何非線性變換。有人可能會擔心,由於KANs將MLPs的每個權重參數變成了樣條函數,因此會變得計算成本很高。然而,令人欣喜的是,KANs通常允許使用比MLPs小得多的計算圖。例如,本研究表明,在解決偏微分方程(PDE)問題時,一個2層寬度為10的KAN比4層寬度為100的MLP的均方誤差(MSE)小100倍(10^(-7)對10^(-5)),同時參數效率提高了100倍(10^2對10^4)。

事實上,使用Kolmogorov-Arnold定理來構建神經網絡的可能性已經被研究過。然而,大多數工作都局限於使用原始的深度為2、寬度為(2n+1)的表示,沒有機會利用更現代的技術(例如反向傳播)來訓練網絡。本研究的貢獻在於將原始的Kolmogorov-Arnold表示推廣到任意寬度和深度,並將其置於當前的深度學習應用背景下,通過廣泛的實證實驗來突出其作為AI+科學基礎模型的潛力。

儘管KANs有優雅的數學解釋,但它們實際上只是樣條和MLPs的結合,利用了各自的優點,避免了各自的缺點。樣條在低維度下表現出色,易於局部調整,並且能夠在不同分辨率之間切換。然而,樣條受到嚴重的維度詛咒問題的影響,因為它們無法利用組合結構。另一方面,MLPs由於其特徵學習能力,受維度詛咒的影響較小,但在低維度下的準確性不如樣條,因為它們無法很好地優化一維函數。為了準確學習一個函數,一個模型不僅需要學習組合結構(外部自由度),還需要能夠很好地逼近一維函數(內部自由度)。KANs正是這樣一種模型,因為它們在外部類似於MLPs,在內部類似於樣條。因此,KANs不僅能夠學習特徵(由於其與MLPs在外部的相似性),還能將這些學習到的特徵優化到很高的精度(由於其與樣條在內部的相似性)。例如,給定一個高維函數


\[ f(x_1, ..., x_N) = \exp\left(\frac{1}{N}\sum_{i=1}^N \sin^2(x_i)\right) \]

(1.1)




由於維度詛咒,樣條在N很大時會失效;MLPs可能會學到廣義加性結構,但它們在逼近指數函數和正弦函數方面效率很低,例如使用ReLU激活函數。相比之下,KANs可以很好地學習組合結構和一維函數,因此大大優於MLPs(見圖3.1)。

在本文中,本研究將通過大量的數值實驗來展示KANs在精度和可解釋性方面如何顯著改善MLPs。本文的結構如圖2.1所示。在第二節中,本研究介紹了KAN架構及其數學基礎,引入了網絡簡化技術以使KANs易於解釋,並提出了網格擴展技術以使KANs逐步提高精度。在第三節中,本研究展示了KANs在數據擬合和求解偏微分方程方面優於MLPs的表現:當數據中存在組合結構時,KANs可以克服維度詛咒,實現比MLPs更好的縮放規律。在第四節中,本研究展示了KANs的可解釋性,並將其應用於科學發現。本研究使用來自數學(結理論)和物理學(Anderson局域化)的兩個實例,證明了KANs可以作為幫助科學家(重新)發現數學和物理定律的有用"合作者"。第五節總結了相關工作。在第六節中,本研究通過討論廣泛的影響和未來方向來總結全文。代碼可在https://github.com/KindXiaoming/pykan獲取,也可以通過pip install pykan進行安裝。

Paper: https://arxiv.org/abs/2404.19756

Code: https://github.com/KindXiaoming/pykan

Quick intro: https://kindxiaoming.github.io/pykan/intro.html

Documentation: https://kindxiaoming.github.io/pykan/


2. Kolmogorov-Arnold 網絡(KAN)


多層感知器(MLPs)受到通用近似定理的啟發。相反,該研究關注的是可以通過一種新型神經網絡實現的Kolmogorov-Arnold表示定理,稱為Kolmogorov-Arnold網絡(KAN)。該研究在2.1節回顧Kolmogorov-Arnold定理,以啟發在2.2節中Kolmogorov-Arnold網絡的設計。在2.3節中,該研究提供了KANs的表達能力和神經縮放法則的理論保證。在2.4節中,該研究提出了一種網格擴展技術,以使KANs逐漸更加精確。在2.5節中,該研究提出了簡化技術,使KANs易於解釋。


2.1 Kolmogorov-Arnold表示定理(Kolmogorov–Arnold representation theorem)

弗拉基米爾·阿諾德(Vladimir Arnold)和安德烈·柯爾莫果洛夫(Andrey Kolmogorov)證明,如果f是在有界域上的多變量連續函數,則f可以表示為一個有限組合的連續一維函數和二元加法運算。更具體地說,對於一個平滑的\(f:[0, 1]^n → R\),


\[ f(x) = f(x_1, ..., x_n) = \sum_{q=1}^{2n+1} \Phi_q \left( \sum_{p=1}^n \phi_{q,p}(x_p) \right), \]

(2.1)


其中 \(\phi_{q,p} : [0, 1] → R\) 和 \(\Phi_q : R → R\)。從某種意義上說,他們證明了加法是唯一真正的多變量函數,因為其他所有函數都可以使用一維函數和求和來表示。人們可能會天真地認為這對機器學習是個好消息:學習一個高維函數歸結為學習一定數量的一維函數。然而,這些一維函數可能是不光滑甚至是分形的,因此在實踐中可能無法學習。由於這種病態行為,Kolmogorov-Arnold表示定理基本上在機器學習中被判了死刑,被認為在理論上是健全的但實際上無用。



然而,該研究對Kolmogorov-Arnold定理在機器學習中的用途持更加樂觀的態度。首先,該研究無需堅持原始方程(2.1),該方程僅具有兩層非線性和隱藏層中的少量項目(2n+1):該研究將網絡推廣到任意寬度和深度。其次,科學和日常生活中的大多數功能通常都是平滑的,並且具有稀疏的組合結構,這可能有助於實現平滑的Kolmogorov-Arnold表示。這裡的理念接近物理學家的心態,他們通常更關心典型案例而非最壞情況。畢竟,我們的物理世界和機器學習任務必須具有結構,才能使物理學和機器學習在所有方面都有用或具有普遍性。


2.2 KAN架構

假設該研究有一個監督學習任務,包含輸入輸出對{xi, yi},其中該研究希望找到一個函數f,使得對所有數據點yi ≈ f(xi)。方程(2.1)意味著如果該研究能找到適當的一維函數\(\phi_{q,p}\)和\(\Phi_q\),那麼該研究就完成了。這啟發該研究設計了一個明確參數化方程(2.1)的神經網絡。由於所有需要學習的函數都是一維函數,該研究可以將每個一維函數參數化為一條B-樣條曲線,具有可學習的局部B-樣條基函數係數(見圖2.2右)。現在該研究有了一個KAN的原型,其計算圖由方程(2.1)確切指定,並在圖0.1(b)中顯示(輸入維度n = 2),顯示為具有在邊上放置激活函數而非節點的兩層神經網絡(節點上執行簡單的求和),並在中間層寬度為2n+1。

正如所提及的,這樣的網絡被認為太簡單,無法在實踐中以平滑的樣條任意精確地近似任何函數!因此,該研究將KAN推廣為更寬更深。立即如何使KAN更深並不清楚,因為Kolmogorov-Arnold表示對應於兩層KAN。據該研究所知,尚無與更深KAN對應的“一般化”定理版本。

當該研究注意到MLPs和KANs之間的類比時,突破發生了。在MLPs中,一旦定義了一層(由線性變換和非線性組成),就可以堆疊更多層來使網路更深。為了建立深層KANs,該研究首先需要回答:“什麼是KAN層?”

結果顯示,一個KAN層,其輸入維度為nin,輸出維度為nout,可以被定義為一個1D函數的矩陣:


\[ \Phi = \{\phi_{q,p}\}, p = 1, 2, ..., nin, q = 1, 2, ..., nout, \]

(2.2)


其中函數\(\phi_{q,p}\)具有可訓練的參數。在Kolmogorov-Arnold定理中,內部函數形成了一個KAN層,nin = n 和 nout = 2n+1,外部函數形成了一個KAN層,nin = 2n+1 和 nout = 1。因此,方程(2.1)中的Kolmogorov-Arnold表示只是兩個KAN層的組合。現在變得清楚了,擁有更深的Kolmogorov-Arnold表示意味著什麼:只需堆疊更多的KAN層!

讓該研究引入一些符號。這一段可能有點技術性,但讀者可以參考圖2.2(左)以獲得具體的例子和直觀的理解。一個KAN的形狀由一個整數數組表示:


\[ [n_0, n_1, ..., n_L], \]

(2.3)


其中ni是計算圖中第i層的節點數。該研究用(l, i)表示第l層的第i個神經元,並用xl,i表示(l, i)-神經元的激活值。在第l層和第l+1層之間,有nlnl+1個激活函數:連接(l, i)和(l+1, j)的激活函數用


\[ \phi_{l,j,i}, l = 0, ..., L-1, i = 1, ..., nl, j = 1, ..., nl+1 \]

(2.4)


表示。\( \phi_{l,j,i} \)的預激活簡單地是\(xl,i\);\( \phi_{l,j,i} \)的後激活由\( \tilde{x}_{l,j,i} \equiv \phi_{l,j,i}(x_{l,i}) \)表示。第\((l + 1, j)\)神經元的激活值僅是所有傳入後激活的總和:


\[ x_{l+1,j} = \sum_{i=1}^{nl} \tilde{x}_{l,j,i} = \sum_{i=1}^{nl} \phi_{l,j,i}(x_{l,i}), j = 1, ..., nl+1. \]

(2.5)


以矩陣形式表示,這可以寫為:


\[ x_{l+1} = \begin{pmatrix} \phi_{l,1,1}( \cdot ) & \phi_{l,1,2}( \cdot ) & \cdots & \phi_{l,1,n_l}( \cdot ) \\ \phi_{l,2,1}( \cdot ) & \phi_{l,2,2}( \cdot ) & \cdots & \phi_{l,2,n_l}( \cdot ) \\ \vdots & \vdots & \ddots & \vdots \\ \phi_{l,n_{l+1},1}( \cdot ) & \phi_{l,n_{l+1},2}( \cdot ) & \cdots & \phi_{l,n_{l+1},n_l}( \cdot ) \end{pmatrix} x_l, \]

(2.6)


其中 \(\Phi_l\) 是對應於第 \(l\) 層的函數矩陣。一個通用的KAN網絡是 \(L\) 層的組合:給定一個輸入向量 \(x_0 \in \mathbb{R}^{n_0}\),KAN的輸出為


\[ \text{KAN}(x) = (\Phi_{L-1} \circ \Phi_{L-2} \circ \ldots \circ \Phi_1 \circ \Phi_0)x. \]

(2.7)


該研究還可以重寫上述方程,使其更類似於方程(2.1),假設輸出維度 \(n_L = 1\),並定義 \(f(x) \equiv KAN(x)\):


\[ f(x) = \sum_{i_{L-1}=1}^{n_{L-1}} \phi_{L-1, i_{L}, i_{L-1}}\left(\sum_{i_{L-2}=1}^{n_{L-2}} \cdots \left( \sum_{i_2=1}^{n_2} \phi_{2, i_3, i_2}\left( \sum_{i_1=1}^{n_1} \phi_{1, i_2, i_1}\left( \sum_{i_0=1}^{n_0} \phi_{0, i_1, i_0}(x_{i_0}) \right) \right) \cdots \right) \right), \]

(2.8)


這個表達式雖然複雜,但我們對KAN層的抽象化和視覺化更為清晰直觀。原始的Kolmogorov-Arnold表示方程(2.1)對應於形狀為 \([n, 2n + 1, 1]\) 的2層KAN。值得注意的是,所有操作都是可微的,因此該研究可以使用反向傳播訓練KANs。作為對比,一個MLP可以寫成交錯的仿射變換 \(W\) 和非線性 \(\sigma\):


\[ MLP(x) = (W_{L-1} \circ \sigma \circ W_{L-2} \circ \sigma \circ \cdots \circ W_1 \circ \sigma \circ W_0)(x). \]

(2.9)


很明顯,MLPs將線性變換和非線性分別作為 \(W\) 和 \(\sigma\) 處理,而KANs則將它們結合在 \(\Phi\) 中。在圖0.1的(c)和(d)中,該研究可視化了三層MLP和三層KAN,以澄清它們的差異。


實施細節

儘管KAN層方程(2.5)看起來非常簡單,但使其優化良好並非微不足道。關鍵技巧包括:

1. 殘差激活函數。該研究包括一個基函數 \(b(x)\)(類似於殘差連接),使得激活函數 \(\phi(x)\) 是基函數 \(b(x)\) 和樣條函數的總和:


   \[ \phi(x) = w(b(x) + \text{spline}(x)). \]

(2.10)


該研究設定


\[ b(x) = \text{silu}(x) = \frac{x}{1 + e^{-x}} \]

(2.11)


在大多數情況下。樣條函數 \( \text{spline}(x) \) 則被參數化為B-樣條的線性組合:


\[ \text{spline}(x) = \sum_i c_i B_i(x) \]

(2.12)


其中 \( c_i \) 是可訓練的。原則上 \( w \) 是多餘的,因為它可以被吸收進 \( b(x) \) 和 \( \text{spline}(x) \)。然而,該研究仍然包括這個 \( w \) 因子,以更好地控制激活函數的總體幅度。


2. 初始化規模。每個激活函數被初始化為使 \( \text{spline}(x) \approx 0 \)。\( w \) 根據Xavier初始化法則初始化,這種方法被用來初始化MLPs中的線性層。

3. 樣條網格的更新。該研究根據其輸入激活值即時更新每個網格,以解決樣條在有界區域內定義但激活值在訓練過程中可能超出固定區域的問題。


參數計數

為簡單起見,假設該研究設置一個網絡:

  1. 深度為 \( L \),
  2. 層寬度相等,\( n_0 = n_1 = \ldots = n_L = N \),
  3. 每個樣條的階數通常為 \( k = 3 \),在 \( G \) 個區間(共 \( G+1 \) 個網格點)。

則總共有 \( O(N^2 L (G + k)) \approx O(N^2 LG) \) 參數。相比之下,具有深度 \( L \) 和寬度 \( N \) 的MLP只需要 \( O(N^2 L) \) 參數,這似乎比KAN更有效。幸運的是,KANs通常需要比MLPs小得多的 \( N \),這不僅節省了參數,而且還實現了更好的泛化(參見圖3.1和圖3.3)並促進了解釋性。該研究指出,對於1D問題,可以取 \( N = L = 1 \),而該研究實現中的KAN網絡僅是一個樣條近似。對於更高維度,該研究通過下面的定理描述了KANs的泛化行為。




2.3 KAN的近似能力和縮放定律

回想方程(2.1)中,寬度為(2n+1)的兩層表示可能不是平滑的。然而,更深的表示可能帶來更平滑激活的優勢。例如,對於四變量函數


\[ f(x_1, x_2, x_3, x_4) = \exp\left(\sin(x_1^2 + x_2^2) + \sin(x_3^2 + x_4^2)\right) \]

(2.13)


可以被一個[4, 2, 1, 1] KAN平滑地表示,這是一個三層KAN,但可能不允許兩層KAN具有平滑激活。為了便於進行近似分析,該研究仍然假設激活的平滑性,但允許表示可以任意寬和深,如方程(2.7)。為了強調該研究KAN對有限集合網格點的依賴,該研究在下面使用\(ΦG_l\)和\(ΦG_l,i,j\)替代方程(2.5)和(2.6)中使用的符號\(Φ_l\)和\(Φ_l,i,j\)。


定理2.1(近似理論,KAT):設\(x = (x_1, x_2, ..., x_n)\)。假設一個函數\(f(x)\)允許表示


\[ f = (\Phi_{L-1} \circ \Phi_{L-2} \circ \ldots \circ \Phi_1 \circ \Phi_0)(x), \]

(2.14)


如方程(2.7)所示,其中每一個\(Φ_l,i,j\)都是\((k+1)\)次連續可微的。那麼存在一個依賴於f及其表示的常數C,使得有關於網格大小G的以下近似界限:存在k階B-樣條函數\(ΦG_l,i,j\),對於任何\(0≤m≤k\)有界限


\[ \|f - (\Phi^G_{L-1} \circ \Phi^G_{L-2} \circ \ldots \circ \Phi^G_1 \circ \Phi^G_0)(x)\|_{C^m} \leq CG^{-k-1+m}. \]

(2.15)


這裡採用\(C^m\)範數來測量高達\(m\)階的導數的幅度:


\[ \|g\|_{C^m} = \max_{|\beta| \leq m} \sup_{x \in [0,1]^n} |D^\beta g(x)|. \]


證明:通過經典的1D B-樣條理論和事實,即\(Φ_l,i,j\)作為連續函數可以在有界域上均勻有界,知道存在有限網格B-樣條函數\(ΦG_l,i,j\),使得對於任何\(0≤m≤k\),


\[ \|(\Phi_{l,i,j} \circ \Phi_{l-1} \circ \ldots \circ \Phi_1 \circ \Phi_0)(x) - (\Phi^G_{l,i,j} \circ \Phi_{l-1} \circ \ldots \circ \Phi_1 \circ \Phi_0)(x)\|_{C^m} \leq CG^{-k-1+m}, \]


具有獨立於G的常數C。固定這些B-樣條逼近。因此,通過定義殘差\(R_l\):


\[ R_l = (\Phi^G_{L-1} \circ \ldots \circ \Phi^G_{l+1} \circ \Phi_l \circ \Phi_{l-1} \circ \ldots \circ \Phi_0)x - (\Phi^G_{L-1} \circ \ldots \circ \Phi^G_{l+1} \circ \Phi^G_l \circ \Phi_{l-1} \circ \ldots \circ \Phi_0)x \]


滿足


\[ \|R_l\|_{C^m} \leq CG^{-k-1+m}, \]


具有獨立於G的常數。最後注意到


\[ f - (\Phi^G_{L-1} \circ \Phi^G_{L-2} \circ \ldots \circ \Phi^G_1 \circ \Phi^G_0)(x) = R_{L-1} + R_{L-2} + \ldots + R_1 + R_0, \]


(2.15)成立。


根據定理2.1的假設,KANs在有限網格大小下可以用一個與維度無關的殘差率很好地逼近函數,因此戰勝了維度的詛咒!這自然而然地發生,因為只用樣條來逼近1D函數。特別是對於\(m = 0\),恢復在\(L∞\)範數中的精確度,這反過來提供了在有限域上的均方根誤差(RMSE)的一個界限,給出了\(k+1\)的縮放指數。當然,常數C依賴於表示;因此它將依賴於維度。該研究將在未來的工作中討論常數對維度的依賴。

該研究指出,儘管Kolmogorov-Arnold定理(2.1)對應於形狀為\([d, 2d+1, 1]\)的KAN表示,它的函數並不一定是平滑的。另一方面,如果該研究能夠識別一個平滑的表示(可能需要額外的層或使KAN比理論預測的更寬),那麼定理2.1表明該研究可以戰勝維度的詛咒(COD)。這不應該令人驚訝,因為可以固有地學習函數的結構,並使有限樣本KAN近似可解釋。



神經縮放定律:與其他理論的比較

神經縮放定律是指隨著模型參數的增加,測試損失減少的現象,即\(ℓ ∝ N^(-α)\),其中\(ℓ\)是測試RMSE,\(N\)是參數數量,\(α\)是縮放指數。更大的\(α\)承諾通過簡單地擴大模型來實現更多的改進。不同的理論已被提出來預測\(α\)。Sharma和Kaplan建議,\(α\)來自於在具有內在維度\(d\)的輸入流形上的資料擬合。如果模型函數類是k階分段多項式(對於ReLU,k=1),則標準近似理論暗示\(α = (k+1)/d\)來自近似理論。這個界限受到維度的詛咒的困擾,所以人們尋求其他不依賴於\(d\)的界限,通過利用組合結構。特別是,Michaud等人考慮了只涉及一元(例如,平方、正弦、指數)和二元(+ 和 ×)操作的計算圖,發現\(α = (k+1)/d* = (k+1)/2\),其中\(d* = 2\)是最大元性。Poggio等人利用組合稀疏的想法,證明了給定函數類\(Wm\)(其導數直到第\(m\)階是連續的),需要\(N = O(ε^(-2/m))\)數量的參數來實現誤差\(ε\),這等效於\(α = m/2\)。該研究的方法,假設平滑的Kolmogorov-Arnold表示存在,將高維函數分解為幾個1D函數,給出\(α = k+1\)(其中k是樣條的分段多項式階數)。該研究選擇\(k = 3\)的立方樣條,所以\(α = 4\)是與其他工作相比最大且最好的縮放指數。該研究將在第3.1節展示,這個界限α = 4實際上可以用KANs實現,而以前的工作報告說,即使在滿足較慢的界限(例如,\(α = 1\))和迅速達到平台的情況下,MLPs也存在問題。


KAT和UAT之間的比較

全連接神經網路的威力是由通用逼近定理(UAT)證明的,該定理指出,給定一個函數和誤差容限\(\epsilon > 0\),具有\(k > N(\epsilon)\)個神經元的雙層網路可以將函數逼近至\(\epsilon\)的誤差內。然而,UAT並不保證\(N(\epsilon)\)如何隨\(\epsilon\)縮放。事實上,它遭受維度詛咒(COD)的困擾,並且在某些情況下,\(N\)已被證明隨\(d\)呈指數增長。KAT和UAT之間的差異是KAN利用了函數的內在低維表示,而MLP沒有。將展示KAN與符號函數很好地對齊,而MLP則不然。


2.4 精確性:網格擴展

原則上一個樣條可以被任意精確地調整以符合目標函數,只要網格可以被足夠細化。這一優良特性被KANs繼承。相比之下,MLPs沒有“細化”的概念。確實,增加MLPs的寬度和深度可以提高性能(“神經縮放定律”)。但這些神經縮放定律增長緩慢,且實現起來成本高昂,因為需要獨立訓練不同大小的模型。與此相反,對於KANs,可以首先訓練一個參數較少的KAN,然後通過簡單地細化樣條網格擴展到一個參數更多的KAN,無需從頭開始重新訓練大模型。

接下來,該研究描述如何執行網格擴展(如圖2.2右側所示),基本上是在一個粗糙的樣條上擬合一個新的細糙樣條。假設想在有界區域\([a, b]\)上用B-樣條的\(k\)階來近似一個一維函數\(f\)。一個具有G1區間的粗糙網格在\(\{t_0 = a, t_1, ..., t_{G1} = b\}\)處有網格點,這被擴充為\(\{t_{-k}, ..., t_{-1}, t_0, ..., t_{G1}, t_{G1+1}, ..., t_{G1+k}\}\)。有 \(G_1 + k\) 個B-樣條基函數,其中第 \(i\) 個B-樣條 \(B_i(x)\) 僅在 \([t_{-k+i}, t_{i+1}]\) 上非零。然後,\(f\) 在粗糙網格上用這些B-樣條基函數的線性組合表示:


\[ f_{\text{coarse}}(x) = \sum_{i=0}^{G_1+k-1} c_i B_i(x) \]

(2.16)


給定更細的網格,具有 \(G_2\) 區間,\(f\) 在細網格上相應地表示為:


\[ f_{\text{fine}}(x) = \sum_{j=0}^{G_2+k-1} c'_j B'_j(x) \]


這些 \(c'_j\) 的參數可以從 \(c_i\) 初始化,通過最小化 \(f_{\text{fine}}(x)\) 到 \(f_{\text{coarse}}(x)\) 的距離來獲得(在某些 \(x\) 的分佈上):


\[ \{c'_j\} = \arg \min_{\{c'_j\}} \mathbb{E}_{x \sim p(x)} \left[ \left( \sum_{j=0}^{G_2+k-1} c'_j B'_j(x) - \sum_{i=0}^{G_1+k-1} c_i B_i(x) \right)^2 \right] \]

(2.16)


這可以通過最小平方算法實現。對KAN中的所有樣條進行獨立的網格擴展。

玩具範例:階梯狀的損失曲線。使用一個玩具範例 f(x, y) = exp(sin(πx) + y^2) 來展示網格擴展的效果。在圖 2.3(左上)中,我們展示了一個 [2, 5, 1] KAN 的訓練和測試 RMSE。網格點的數量開始為 3,每 200 個 LBFGS 步驟增加到更高的值,最終達到 1000 個網格點。很明顯,每次細化網格時,訓練損失下降得比之前更快(除了最細的 1000 點網格,優化可能由於不良的損失地形而停止工作)。然而,測試損失先下降然後上升,呈現 U 形,這是由於偏差-方差權衡(欠擬合 vs. 過擬合)。我們推測,當參數數量與數據點數量匹配時,在插值閾值處會達到最佳測試損失。由於我們的訓練樣本為 1000,而 [2, 5, 1] KAN 的總參數為 15G(G 是網格區間的數量),我們預期插值閾值為 G = 1000/15 ≈ 67,這與我們實驗觀察到的值 G ∼ 50 大致一致。

小型 KAN 泛化性能更好。這是我們可以達到的最佳測試性能嗎?請注意,合成任務可以由 [2, 1, 1] KAN 精確表示,所以我們訓練一個 [2, 1, 1] KAN 並在圖 2.3 右上方呈現訓練動態。有趣的是,它可以實現比 [2, 5, 1] KAN 更低的測試損失,具有更清晰的階梯結構,由於參數更少,插值閾值被延遲到更大的網格尺寸。這突顯了選擇 KAN 架構的微妙之處。如果我們不知道問題結構,我們如何確定最小的 KAN 形狀?在第 2.5 節中,我們將提出一種通過正則化和剪枝自動發現這種最小 KAN 架構的方法。

縮放規律:與理論比較。我們也對測試損失如何隨網格參數數量的增加而降低感興趣。在圖 2.3(左下)中,[2,1,1] KAN 的縮放大致為 test RMSE ∝ G^−3。然而,根據定理 2.1,我們預期 test RMSE ∝ G^−4。我們發現樣本之間的誤差並不均勻。這可能歸因於邊界效應 [18]。事實上,有一些樣本的誤差顯著大於其他樣本,使得整體縮放放緩。如果我們繪製平方損失的中位數(非平均值)的平方根,我們得到接近 G^−4 的縮放。儘管存在這種次優性(可能是由於優化),KAN 在數據擬合(圖 3.1)和 PDE 求解(圖 3.3)方面仍然具有比 MLP 更好的縮放規律。此外,訓練時間隨網格點數 G 的增加而呈現良好的縮放,如圖 2.3 右下方所示。

外部與內部自由度。KAN 強調的一個新概念是外部與內部自由度(參數)之間的區別。節點如何連接的計算圖表示外部自由度("dofs"),而激活函數內部的網格點是內部自由度。KAN 受益於同時具有外部 dofs 和內部 dofs。外部 dofs(MLP 也有但樣條沒有)負責學習多個變量的組合結構。內部 dofs(樣條也有但 MLP 沒有)負責學習單變量函數。

 2.5 解釋性:簡化KANs並使它們可互動

上一小節留下的一個問題是,不知道如何選擇最佳匹配數據集結構的KAN形狀。例如,如果知道資料集是通過符號公式 \(f(x, y) = \exp(\sin(\pi x) + y^2)\) 生成的,那麼知道一個 [2, 1, 1] KAN 能夠表達這個函數。但在實踐中不知道這些信息,所以最好有方法能夠自動確定這種形狀。這個想法是從一個足夠大的KAN開始,並用稀疏正則化訓練它,隨後進行剪枝。該研究將展示,這些剪枝後的KAN比未剪枝的更易於解釋。為了使KANs最大限度地可解釋,該研究在2.5.1節提出了一些簡化技術,並在2.5.2節示例了用戶如何與KANs互動以使它們更易於解釋。

2.5.1 簡化技術

1. 稀疏化:對於MLPs,使用L1正則化線性權重以傾向於稀疏化。KANs可以適應這一高級思想,但需要兩個修改:

  •  KANs中沒有線性“權重”。線性權重被可學習的激活函數替代,所以應定義這些激活函數的L1範數。
  •  該研究發現對於KANs的稀疏化,僅使用L1正則化是不夠的;相反,還需要額外的熵正則化。

定義激活函數 \( \phi \) 的L1範數為其在 \( N_p \) 個輸入上的平均幅度,即


\[ |\phi|_1 \equiv \frac{1}{N_p} \sum_{s=1}^{N_p} |\phi(x^{(s)})| \]

(2.17)


對於一個KAN層 \( \Phi \),具有 \( n_{\text{in}} \) 個輸入和 \( n_{\text{out}} \) 個輸出,定義 \( \Phi \) 的L1範數為所有激活函數L1範數的總和,即


\[ |\Phi|_1 \equiv \sum_{i=1}^{n_{\text{in}}} \sum_{j=1}^{n_{\text{out}}} |\phi_{i,j}|_1 \]

(2.18)


此外,定義 \( \Phi \) 的熵為


\[ S(\Phi) \equiv -\sum_{i=1}^{n_{\text{in}}} \sum_{j=1}^{n_{\text{out}}} \frac{|\phi_{i,j}|_1}{|\Phi|_1} \log \left( \frac{|\phi_{i,j}|_1}{|\Phi|_1} \right) \]

(2.19)


總訓練目標 \( \ell_{\text{total}} \) 是預測損失 \( \ell_{\text{pred}} \) 加上所有KAN層的L1和熵正則化:


\[ \ell_{\text{total}} = \ell_{\text{pred}} + \lambda \left( \mu_1 \sum_{l=0}^{L-1} |\Phi_l|_1 + \mu_2 \sum_{l=0}^{L-1} S(\Phi_l) \right) \]

(2.20)


其中,\( \mu_1 \) 和 \( \mu_2 \) 是相對大小,通常設置為 \( \mu_1 = \mu_2 = 1 \),而 \( \lambda \) 控制整體正則化大小。



2. 可視化。當我們將一個KAN可視化時,為了感受量級,我們將一個激活函數\(\phi_{l,i,j}\)的透明度設置為與\(\tanh(\\beta A_{l,i,j})\)成正比,其中\(\beta = 3\)。因此,量級較小的函數會顯得褪色,使我們能夠專注於重要的函數。

3. 修剪。在使用稀疏化懲罰進行訓練之後,我們可能還想將網路修剪為更小的子網路。我們在節點層面(而不是邊層面)對KAN進行稀疏化。對於每個節點(例如第\(l\)層的第\(i\)個神經元),我們定義其輸入和輸出分數為

$$I_{l,i} = \max_k(|\phi_{l-1,k,i}|1),\ O{l,i} = \max_j(|\phi_{l+1,j,i}|_1),$$

(2.21)

如果輸入和輸出分數都大於閾值超參數(\\theta = 10^{-2}\)(預設值),則認為一個節點是重要的。所有不重要的神經元都會被修剪掉。

4. 符號化。在我們懷疑某些激活函數實際上是符號函數(例如\(\cos\)或log)的情況下,我們提供了一個介面來將它們設置為指定的符號形式。fix_symbolic(l,i,j,f)可以將\((l,i,j)\)激活設置為\(f\)。然而,我們不能簡單地將激活函數設置為精確的符號公式,因為其輸入和輸出可能存在偏移和縮放。因此,我們從樣本中獲取預激活\(x\)和後激活\(y\),並擬合仿射參數\((a,b,c,d)\),使得\(y \approx cf(ax + b) + d\)。擬合通過\(a\),\(b\)的迭代網格搜索和線性回歸完成。

除了這些技術,我們還提供了額外的工具,允許用戶對KAN進行更細粒度的控制,列在附錄A中。

2.5.2 人類與KANs的互動實例

該研究提出了一些簡化KANs的技術,可以將這些簡化選項視為可以點擊的按鈕。使用者通過與這些按鈕互動,可以決定哪個按鈕是最有希望點擊的,使KANs更易於解釋。用以下示例展示了使用者如何與KAN互動以獲得最大限度可解釋的結果。

考慮回歸任務


\[ f(x, y) = \exp(\sin(\pi x) + y^2) \]

(2.22)


假設有資料點 \( (x_i, y_i, f_i) \),\( i = 1, 2, \ldots, N_p \),一位假想的使用者Alice有興趣弄清楚符號公式。Alice與KANs互動的步驟如下(圖2.4所示):

步驟1:訓練與稀疏化。從一個完全連接的[2, 5, 1] KAN開始,用稀疏化正則化訓練可以使其變得相當稀疏。隱藏層中的5個神經元中有4個看起來是無用的,因此希望將它们剪除。

步驟2:剪枝。自動剪枝觀察到除了最後一個以外,所有隐藏層神經元都被丢棄,只留下一個 [2, 1, 1] KAN。激活函數是已知的符號函数。

步驟3:設定符号函數。假设用户能從KAN圖中正確猜測这些符號公式,他們可以設置:

```

fix_symbolic(0,0,0,'sin')

fix_symbolic(0,1,0,'x^2')

fix_symbolic(1,0,0,'exp')

```

(2.23)


如果使用者沒有領域知識或不清楚這些激活函數可能對應哪些符號函數,提供了一個 `suggest_symbolic` 函數來建議符號候選者。

步驟4:進一步訓練。在將所有激活函數符號化之後,唯一剩下的參數是仿射參數。繼續訓練這些仿射參數,當看到損失下降到機器精度時,知道找到了正確的符號表達式。

步驟5:輸出符號公式。使用 Sympy 計算輸出節點的符號公式。使用者獲得的公式為 \(1.0e1.0y^2+1.0sin(3.14x)\),這是真正的答案(只顯示了\(π\)的兩位小數)。

為什麼不用符號迴歸(SR)?

在這個例子中使用符號迴歸似乎是合理的。然而,一般來說,符號迴歸法比較脆弱,且難以調試。它們要麼最終成功返回結果,要麼失敗,而不提供可解釋的中間結果。相較之下,KAN通過梯度下降在函數空間中進行連續搜索,因此其結果更加連續且穩定。此外,由於KAN的透明度,使用者可以比使用符號迴歸更好地控制KAN。該研究將KAN的可視化方式比喻為向使用者展示KAN的「大腦」,使用者可以對KAN進行「手術」(調試)。符號迴歸通常不提供這種控制層級。將在第4.4節展示此類實例。更普遍地,當目標函數不是符號表達時,符號迴歸將失敗,但KAN仍然可以提供有意義的輸出。例如,特殊函數(如貝塞爾函數)是不可能被符號迴歸學習的,除非事先提供,但KAN可以使用樣條進行數值近似(見圖4.1(d))。


3 KANs的精確性


在這一節中,該研究將展示KANs在各種任務(如回歸和偏微分方程求解)中表示函數的有效性超過了MLPs。在比較兩類模型時,公平地比較它們的精度(損失)和復雜度(參數數量)是合理的。該研究將展示KANs比MLPs展示出更有利的帕雷托前緣。此外,在3.5節中,該研究將顯示KANs能夠自然地進行持續學習而不會出現災難性遺忘。



3.1 玩具資料集


在第2.3節中,理論建議測試RMSE損失 \( \ell \) 隨模型參數 \( N \) 的增加而按 \( \ell \propto N^{-4} \) 縮放。然而,這依賴於存在柯爾莫哥洛夫-阿諾德表示定理(Kolmogorov-Arnold representation theorem),也被稱為超曲面定理(superposition theorem),是數學分析中的一個重要定理。它由俄羅斯數學家安德烈·柯爾莫哥洛夫(Andrey Kolmogorov)在1956-1957年提出,並由其學生弗拉基米爾·阿諾德(Vladimir Arnold)在1959年完成證明。
該定理的主要內容如下:

任意一個連續函數 \(f(x_1,\ldots,x_n)\),其中自變量 \((x_1,\ldots,x_n)\) 屬於 \(n\) 維單位立方體 \([0,1]^n\),都可以表示為:
\(f(x_1,\ldots,x_n) = \sum_{q=0}^{2n} \Phi_q \left(\sum_{p=1}^n \psi_{q,p}(x_p)\right)\)
其中 \(\Phi_q\) 是連續函數,\(\psi_{q,p}\) 是 \([0,1]\) 上的單變量函數。

這個定理說明,任何多維連續函數都可以由一些單變量函數的疊加來表示,而這些單變量函數與原函數的維度無關。這一結果在當時是非常令人驚訝的,因為它表明高維函數可以用低維函數來表示。

在神經網絡領域,這個定理啟發了新的網絡結構的設計,即論文中提到的KAN(Kolmogorov-Arnold network)。傳統的神經網絡使用固定的激活函數,而KAN使用可學習的、類似於定理中 \(\psi_{q,p}\) 的單變量函數作為激活函數,並將它們放在網絡的邊上而不是節點上。這種結構使得KAN具有更強的表達能力和可解釋性。

然而,柯爾莫哥洛夫-阿諾德表示定理在實踐中也存在一些問題。定理本身是一個存在性結果,並沒有給出如何構造這些單變量函數的方法。此外,定理中的單變量函數可能高度非光滑,這在實際應用中可能會帶來數值不穩定性。儘管如此,該定理仍然是啟發神經網絡設計的重要理論基礎,並為進一步的理論研究提供了方向。。為了作為一個合理性檢驗,構造了五個我們知道具有平滑KA表示的例子:

  1.  \( f(x) = J_0(20x) \),這是貝塞爾函數。由於它是一個一元函數,它可以被樣條表示,這是一個 [1, 1] KAN。
  2.  \( f(x, y) = \exp(\sin(\pi x) + y^2) \)。我們知道它可以被一個 [2, 1, 1] KAN準確表示。
  3.  \( f(x, y) = xy \)。我們從圖4.1知道它可以被一個 [2, 2, 1] KAN準確表示。
  4.  一個高維例子 \( f(x_1, \ldots, x_{100}) = \exp\left(\frac{1}{100} \sum_{i=1}^{100} \sin^2(\pi x_i)\right) \),它可以被一個 [100, 1, 1] KAN表示。
  5.  一個四維例子 \( f(x_1, x_2, x_3, x_4) = \exp\left(\frac{1}{2} (\sin(\pi (x_1^2 + x_2^2)) + \sin(\pi (x_3^2 + x_4^2)))\right) \),它可以被一個 [4, 4, 2, 1] KAN表示。

通過每200步增加網格點數量來訓練這些KANs,總共覆蓋 \( G = \{3, 5, 10, 20, 50, 100, 200, 500, 1000\} \)。該研究還訓練了具有不同深度和寬度的MLPs作為基線。無論是MLPs還是KANs都使用LBFGS訓練了總共1800步。將KANs和MLPs的測試RMSE作為參數數量的函數繪製在圖3.1中,顯示KANs比MLPs有更好的縮放曲線,特別是對於高維例子。為了比較,將從KAN理論預測的紅色虛線(\( \alpha = k+1 = 4 \))和從Sharma & Kaplan 預測的黑色虛線(\( \alpha = (k+1)/d = 4/d \))繪製在圖中。KANs幾乎可以達到更陡峭的紅線,而MLPs即使在達到較慢的黑線時也難以收斂並迅速達到平台期。另外還注意到,對於最後一個例子,2層的KAN [4, 9, 1]的表現比3層的KAN(形狀 [4, 2, 2, 1])差得多。這突顯了更深層KANs的更大表達能力,這對MLPs也是一樣:更深的MLPs比較淺的具有更多的表達能力。


 3.2 特殊函數


上述結果的一個限制是,該研究假設了「真實」KAN形狀的知識。在實踐中,不知道KA表示的存在。即使被告知存在這樣的KA表示,也不知道KAN形狀。多變量特殊函數是這種情況,因為如果多變量特殊函數(例如,貝塞爾函數 \( f(\nu, x) = J_\nu(x) \))能用KA表示表達,只涉及一元函數和加法,這在數學上將是驚人的。以下展示:
  1. 尋找(近似的)緊湊KA表示特殊函數是可能的,這揭示了從柯爾莫哥洛夫-阿諾德表示的角度看特殊函數的新數學屬性。
  2. KANs在表示特殊函數方面比MLPs更有效和精確

該研究收集了15個在數學和物理中常見的特殊函數,彙總在表2中。選擇了固定寬度為5或100的MLPs,並掃描了深度 \(\{2, 3, 4, 5, 6\}\)。運行了帶剪枝和不帶剪枝的KANs。


  • 未剪枝的KANs:固定了KAN的形狀,寬度設置為5,深度也在 \(\{2,3,4,5,6\}\) 中掃描。
  • 帶剪枝的KANs。使用第2.5.1節中的稀疏化(\( \lambda = 10^{-2} \) 或 \( 10^{-3} \))和剪枝技術來獲得一個從固定形狀KAN剪枝得到的更小的KAN。
每個KAN都初始化為有 \( G = 3 \) 的網格點,使用LBFGS訓練,每200步增加網格點數量,覆蓋 \( G = \{3, 5, 10, 20, 50, 100, 200\} \)。對於每個超參數組合,我們運行3個隨機種子。

對於每個資料集和每個模型家族(KANs或MLPs),在(參數數量,RMSE)平面上繪製了帕雷托前緣,顯示在圖3.2中。KANs的表現一致優於MLPs,即KANs在相同參數數量下可以達到比MLPs更低的訓練/測試損失。此外,報告了自動發現的特殊函數KANs的(令人驚訝的緊湊)形狀,列在表2中。一方面,解釋這些緊湊表達在數學上意味著什麼是有趣的(在附錄F的圖F.1和F.2中包括了KAN的插圖)。另一方面,這些緊湊的表達暗示了將高維查找表拆分為幾個一維查找表的可能性,這可以在推理時節省大量內存,並且(幾乎可以忽略不計)進行一些加法運算的開銷。

3.3 費曼數據集


第3.1節的設置是當我們清楚地知道「真實」KAN形狀時。第3.2節的設置是當我們明確不知道「真實」KAN形狀時。這部分調查位於中間的設置:鑑於數據集的結構,我們可能會手工構造KAN,但我們不確定它們是否是最優的。在這個範疇中,比較人工構造的KAN和通過剪枝(在第2.5.1節中的技術)自動發現的KAN是有趣的。





當有資料集的結構資訊時,可以手動構建KAN,但可能不確定這些結構是否完全準確或有效。例如,可能基於資料的某些特性猜測一個特定的KAN結構,然後嘗試用這個結構進行訓練和優化,但隨後的剪枝和調整可能顯示出一個更有效的結構。

這種探索對於理解KAN在實際應用中的彈性和潛力至關重要。透過對比人工構建的模型與經過剪枝精簡後的模型,可以更好地評估KAN設計選擇的有效性,並進一步優化模型的表現和效率。

費曼資料集。費曼資料集從費曼的教科書中收集了許多物理方程式。對於至少有兩個變數的問題感興趣,因為單變數問題對於KAN來說是微不足道的(它們簡化為一維樣條)。費曼資料集的一個樣本方程式是相對論速度加法公式 \( f(u, v) = \frac{u+v}{1+uv} \)。通過隨機抽取 \( u_i \in (-1, 1) \) 和 \( v_i \in (-1, 1) \),並計算 \( f_i = f(u_i, v_i) \),可以建立資料集。給定許多元組 \( (u_i, v_i, f_i) \),訓練一個神經網路,目的是從 \( u \) 和 \( v \) 預測 \( f \)。該研究感興趣的是:
(1) 神經網路在測試樣本上的表現如何;
(2) 可以從神經網路中學到多少關於問題結構的知識。

比較四種神經網路:
  1. 人工構造的KAN。給定一個符號公式,將其改寫為Kolmogorov-Arnold表示。例如,要將兩個數 \( x \) 和 \( y \) 相乘,可以使用恆等式 \( xy = \frac{(x+y)^2}{4} - \frac{(x-y)^2}{4} \),對應於一個 [2, 2, 1] KAN。構造的形狀列在表中的「人工構造KAN形狀」一欄。
  2. 不帶修剪的KAN。我們將KAN形狀固定為寬度5,深度在 {2, 3, 4, 5, 6} 中掃描。
  3. 帶修剪的KAN。使用節中的稀疏化(\( \lambda = 10^{-2} \) 或 \( 10^{-3} \))和修剪技術,從中的固定形狀KAN得到一個更小的KAN。
  4. 固定寬度為20的MLP,深度在 {2, 3, 4, 5, 6} 中掃描,激活函數從 {Tanh, ReLU, SiLU} 中選擇。

每個KAN初始化為 \( G = 3 \),用LBFGS訓練,每200步增加網格點數量以涵蓋 \( G = \{3, 5, 10, 20, 50, 100, 200\} \)。對於每個超參數組合,嘗試3個隨機種子。對於每個資料集(方程式)和每種方法,在表中報告了隨機種子和深度上最佳模型(最小KAN形狀,或最低測試損失)的結果。發現MLP和KAN在平均表現上相當。對於每個資料集和每個模型族(KAN或MLP),在參數數量和RMSE損失構成的平面上繪製帕累托前沿,如附錄D中的圖D.1所示。推測費曼資料集太簡單了,不足以讓KAN做出進一步改進,因為變數依賴關係通常是平滑或單調的,這與經常表現出振盪行為的特殊函數的複雜性形成對比。

自動發現的KAN比人工構造的KAN更小。在表中的兩列中報告了修剪後的KAN形狀;一列是可以達到合理損失(即測試RMSE小於 \(10^{-2}\))的最小修剪KAN形狀;另一列是達到最低測試損失的修剪KAN。為了完整起見,在附錄D(圖D.2和D.3)中可視化了所有54個修剪後的KAN。有趣的是,自動發現的KAN形狀(無論是最小的還是最佳的)通常都比我們的人工構造更小。這意味著Kolmogorov-Arnold表示可能比我們想像的更有效。同時,這可能會使解釋變得微妙,因為資訊被壓縮到一個比我們習慣的更小的空間。

以相對論速度合成公式 \( f(u, v) = \frac{u+v}{1+uv} \) 為例。構造相當深,因為假設 \( u \) 和 \( v \) 的乘法需要兩層,\( 1+uv \) 的倒數需要一層,\( (u+v) \) 和 \( \frac{1}{1+uv} \) 的乘法需要另外兩層,總共需要五層。然而,自動發現的KAN只有兩層深!事後看來,如果回憶相對論中的快度技巧,實際上是意料之中的:定義兩個"快度" \( a \equiv \text{arctanh} u \) 和 \( b \equiv \text{arctanh} v \)。在快度空間中,速度的相對論合成是簡單的加法,即 \( \frac{u+v}{1+uv} = \tanh(\text{arctanh} u + \text{arctanh} v) \),這可以用一個兩層的KAN實現。假裝我們不知道物理中快度的概念,我們可能會直接從KAN中發現這個概念,而無需通過試錯的符號操作。

KAN的可解釋性,有助於科學發現,是第四節的主題。

3.4 求解偏微分方程式

考慮具有零Dirichlet邊界條件的Poisson方程。對於 \(\Omega = [-1, 1]^2\),考慮Partial Differential Equation\(\begin{aligned}u_{xx} + u_{yy} &= f \text{ in } \Omega, \\u &= 0 \text{ on }\partial\Omega.\end{aligned}\)
(3.2)





考慮的資料 \(f\) 為 \(-\pi^2(1 + 4y^2) \sin(\pi x) \sin(\pi y^2) + 2\pi \sin(\pi x) \cos(\pi y^2)\),其中 \(u = \sin(\pi x) \sin(\pi y^2)\) 是真實解。使用物理信息神經網路(PINN)的框架來求解這個PDE,損失函數為\(\text{loss}\text{pde} = \alpha \text{loss}i + \text{loss}b := \alpha \frac{1}{n_i}\sum{i=1}^{n_i} |u{xx}(z_i) + u{yy}(z_i) - f(z_i)|^2 + \frac{1}{n_b}\sum_{i=1}^{n_b} u^2(z_i),\)



其中用 \(\text{loss}_i\) 表示內部損失,通過在區域內均勻採樣 \(n_i\) 個點 \(z_i = (x_i, y_i)\) 來離散化和計算,用 \(\text{loss}_b\) 表示邊界損失,通過在邊界上均勻採樣 \(n_b\) 個點來離散化和計算。 \(\alpha\) 是平衡兩項效果的超參數。

使用相同的超參數 \(n_i = 10000\), \(n_b = 800\) 和 \(\alpha = 0.01\) 來比較KAN架構和MLP架構。

測量 \(L^2\) 範數和能量(\(H^1\))範數中的誤差,發現KAN在使用更小的網路和更少的參數時,實現了更好的縮放規律和更小的誤差;見圖3.3。因此,推測KAN可能有潛力作為PDE模型簡化的良好神經網路表示。

3.5 持續學習

災難性遺忘是當前機器學習中的一個嚴重問題。當人類掌握一項任務並轉向另一項任務時,他們不會忘記如何執行第一項任務。不幸的是,神經網路並非如此。當一個神經網路在任務1上訓練,然後轉向在任務2上訓練時,網路很快就會忘記如何執行任務1。人工神經網路和人腦之間的一個關鍵區別是,人腦在空間上局部放置了功能截然不同的模塊。當學習一個新任務時,結構重組只發生在負責相關技能的局部區域,其他區域保持不變。大多數人工神經網路,包括MLP,都沒有這種局部性的概念,這可能是災難性遺忘的原因。

KAN具有局部可塑性,可以通過利用樣條的局部性來避免災難性遺忘。這個想法很簡單:由於樣條基函數是局部的,一個樣本只會影響少數附近的樣條係數,遠處的係數保持不變(這是理想的,因為遠處的區域可能已經存儲了想要保留的信息)。相比之下,由於MLP通常使用全局激活函數,例如ReLU/Tanh/SiLU等,任何局部變化都可能不受控制地傳播到遠處的區域,破壞存儲在那裡的信息。

用一個玩具示例來驗證這個直覺。1D回歸任務由5個高斯峰組成。圍繞每個峰值的數據按順序(而不是一次全部)呈現給KAN和MLP,如圖3.4頂行所示。每個訓練階段後的KAN和MLP預測分別顯示在中間行和底行。正如預期的那樣,KAN僅對當前階段存在資料的區域進行重塑,保持先前區域不變。相比之下,MLP在看到新的數據樣本後對整個區域進行重塑,導致災難性遺忘。

這裡只是在一個極其簡單的示例上展示了初步結果,演示人們如何可能利用KAN中的局部性(得益於樣條參數化)來減少災難性遺忘。然而,目前尚不清楚方法是否可以推廣到更現實的設置,將此留待未來的工作。還想研究我們的方法如何與持續學習中的最先進方法聯繫和結合。

4 KAN是可解釋的


在本節中,展示KAN是可解釋和可互動的,這要歸功於我們在2.5節中開發的技術。不僅想在合成任務上測試KAN的使用(第4.1和4.2節),而且想在真實的科學研究中測試。證明KAN可以(重新)發現紐結理論中高度非平凡的關係(第4.3節)和凝聚態物理學中的相變邊界(第4.4節)。由於其準確性(上一節)和可解釋性(本節),KAN可能成為AI+科學的基礎模型。

4.1 有監督的玩具資料集

首先檢驗KAN揭示符號公式中組合結構的能力。下面列出了六個示例,KAN在圖4.1中可視化。KAN能夠揭示這些公式中存在的組合結構,以及學習正確的單變量函數。以下是一些示例:

(a) 乘法 \( f(x, y) = xy \)。一個 [2, 5, 1] KAN 被修剪為 [2, 2, 1] KAN。學到的激活函數是線性和二次方的。從計算圖中,看到它計算 \( xy \) 的方式是利用 \( 2xy = (x + y)^2 - (x^2 + y^2) \)。

(b) 正數的除法 \( f(x, y) = x/y \)。一個 [2, 5, 1] KAN 被修剪為 [2, 1, 1] KAN。學到的激活函數是對數函數和指數函數,KAN 通過利用恆等式 \( x/y = \exp(\log x - \log y) \) 來計算 \( x/y \)。

(c) 數值到類別。任務是將 [0, 1] 中的實數轉換為其第一位小數(作為 one-hot),例如,\( 0.0618 \to [1, 0, 0, 0, 0, \cdots] \),\( 0.314 \to [0, 0, 0, 1, 0, \cdots] \)。請注意,激活函數被學習為位於相應十進制數字附近的尖峰。



(d) 特殊函數 \( f(x, y) = \exp(J_0(20x) + y^2) \)。符號回歸的一個局限性是,如果特殊函數不作為先驗知識提供,它將永遠無法找到特殊函數的正確公式。KAN可以學習特殊函數——高度波動的Bessel函數 \( J_0(20x) \) 被KAN(數值地)學習。

(e) 相變 \( f(x_1, x_2, x_3) = \tanh(5(x_1^4 + x_2^4 + x_3^4 - 1)) \)。相變在物理學中備受關注,因此希望KAN能夠檢測相變並識別正確的序參量。使用 tanh 函數來模擬相變行為,序參量是 \( x_1 \), \( x_2 \), \( x_3 \) 的四次方項的組合。KAN訓練後出現了四次方依賴和 tanh 依賴。這是第4.4節中討論的局域化相變的一個示例。

(f) 更深的組合\(f(x_1, x_2, x_3, x_4) = \sqrt{(x_1 - x_2)^2 + (x_3 - x_4)^2}\)。為了計算這個,需要恆等函數、平方函數和平方根,這至少需要一個三層KAN。事實上,該研究發現[4, 3, 3, 1] KAN可以自動修剪為[4, 2, 1, 1] KAN,這正好對應於期望的計算圖。

來自費曼資料集和特殊函數資料集的更多示例在附錄D和F的圖D.2、D.3、F.1、F.2中可視化。

4.2 無監督玩具數據集

通常,科學發現被表述為有監督的學習問題,即給定輸入變量\(x_1, x_2, \cdots, x_d\)和輸出變量\(y\),想找到一個可解釋的函數\(f\)使得\(y \approx f(x_1, x_2, \cdots, x_d)\)。然而,另一種科學發現可以表述為無監督學習,即給定一組變量\((x_1, x_2, \cdots, x_d)\),想發現變量之間的結構關係。具體來說,想找到一個非零的\(f\)使得

$$f(x_1, x_2, \cdots, x_d) \approx 0.$$  
(4.1)



例如,考慮一組特徵\((x_1, x_2, x_3)\)滿足\(x_3 = \exp(\sin(\pi x_1) + x_2^2)\)。那麼一個有效的\(f\)是\(f(x_1, x_2, x_3) = \sin(\pi x_1) + x_2^2 - \log(x_3) = 0\),意味著\((x_1, x_2, x_3)\)的點形成了一個由\(f = 0\)指定的2D子流形,而不是填充整個3D空間。

如果可以設計出求解無監督問題的算法,它相對於有監督問題有相當大的優勢,因為它只需要特徵集\(S = (x_1, x_2, \cdots, x_d)\)。另一方面,有監督問題試圖根據其他特徵預測特徵的子集,即將\(S = S_\text{in} \cup S_\text{out}\)分成要學習的函數的輸入和輸出特徵。如果沒有領域專業知識來指導分割,那麼有\(2^d - 2\)種可能性使得\(|S_\text{in}| > 0\)且\(|S_\text{out}| > 0\)。通過使用無監督方法,可以避免這個指數級大的有監督問題空間。這種無監督學習方法對4.3節中的紐結數據集很有價值。Google Deepmind團隊手動選擇signature作為目標變量,否則他們將面臨上述組合問題。這提出了一個問題,是否可以直接處理無監督學習。在下面介紹我們的方法和一個玩具示例。

通過將無監督學習問題轉化為對所有\(d\)個特徵的有監督學習問題來解決無監督學習問題,而不需要選擇分割。基本思想是學習一個函數\(f(x_1, \ldots, x_d) = 0\)使得\(f\)不是0函數。為此,類似於對比學習,定義正樣本和負樣本:正樣本是真實數據的特徵向量。負樣本通過特徵損壞構造。為了確保每個拓撲不變量的整體特徵分佈保持不變,通過在整個訓練集上隨機排列每個特徵來進行特徵損壞。現在要訓練一個網路\(g\)使得\(g(x_\text{real}) = 1\)且\(g(x_\text{fake}) = 0\),這將問題轉化為有監督問題。但是,請記住,最初想要\(f(x_\text{real}) = 0\)且\(f(x_\text{fake}) \neq 0\)。可以通過令\(g = \sigma \circ f\)來實現這一點,其中\(\sigma(x) = \exp(-\frac{x^2}{2w^2})\)是一個高斯函數,寬度$w$很小,這可以方便地通過一個形狀為\([..., 1, 1]\)的KAN來實現,其最後一個激活設置為高斯函數\(\sigma\),所有前面的層形成\(f\)。除了上述修改外,其他一切與有監督訓練相同。

現在演示無監督範式對合成示例有效。考慮一個6D資料集,其中\((x_1, x_2, x_3)\)是相依變量,使得\(x_3 = \exp(\sin(x_1)+x_2^2)\); \((x_4, x_5)\)是相依變量,\(x_5 = x_4^3\); \(x_6\)與其他變量無關。在圖4.2中,展示了對於seed = 0,KAN揭示了\(x_1\)、\(x_2\)和\(x_3\)之間的函數依賴關係;對於另一個seed = 2024,KAN揭示了\(x_4\)和\(x_5\)之間的函數依賴關係。初步結果依賴於隨機性(不同的種子)來發現不同的關係;在未來,希望研究一種更系統、更可控的方式來發現一組完整的關係。即便如此,目前狀態下的工具可以為科學任務提供見解。在4.3節中介紹了在紐結資料集上的結果。

4.3 數學應用:紐結理論(Knot Theory)

紐結理論是低維拓撲學中的一個主題,它闡明了三維流形和四維流形的拓撲方面,並在生物學和拓撲量子計算等領域有多種應用。從數學上講,紐結\(K\)是\(S^1$到$S^3\)的嵌入。如果兩個紐結\(K\)和\(K'\)可以通過\(S^3\)的環境空間變形相互變形,則它們在拓撲上是等價的,在這種情況下,我們寫成\([K] = [K']\)。

一些紐結在拓撲上是平凡的,這意味著它們可以平滑地變形為標準圓。紐結具有多種變形不變特徵\(f\),稱為拓撲不變量,如果\(f(K) \neq f(K')\),它們可用於表明兩個紐結在拓撲上是不等價的,\([K] \neq [K']\)。在某些情況下,拓撲不變量在本質上是幾何的。例如,雙曲紐結\(K\)的補空間\(S^3 \backslash K\)承認規範的雙曲度量\(g\),使得\(\text{vol}_g(K)\)是一個稱為雙曲體積的拓撲不變量。其他拓撲不變量在本質上是代數的,例如Jones多項式。

鑑於紐結在數學中的基本性質及其應用的重要性,研究ML是否可以帶來新結果是很有趣的。例如,強化學習被用來建立某些紐結的帶性,這排除了許多潛在的反例來反駁光滑4d Poincaré猜想。

有監督學習和人類領域專家被用來得出一個關於代數和幾何紐結不變量的新定理。在這種情況下,梯度顯著性識別出有監督問題的關鍵不變量,這使領域專家提出了一個猜想,隨後被改進和證明。研究KAN是否可以在同一問題上取得良好的可解釋結果,該問題預測紐結的signature。他們從研究紐結理論數據集中得到的主要結果是:

(1) 他們使用網路歸因方法發現signature \(\sigma\)主要依賴於子午距\(\mu\)(實部\(\mu_r\),虛部\(\mu_i\))和經向距\(\lambda\)。 
(2) 人類科學家後來發現\(\sigma\)與斜率\(\equiv \text{Re}(\frac{\lambda}{\mu}) = \frac{\lambda\mu_r}{\mu_r^2+\mu_i^2}\)高度相關,並推導出\(|2\sigma - \text{slope}|\)的界。

下面展示,KAN不僅以更小的網路和更多的自動化方式重新發現了這些結果,而且還提出了一些有趣的新結果和見解。







為了研究(1),將17個紐結不變量視為輸入,將signature視為輸出。signature(偶數)被編碼為 one-hot 向量,網路用交叉熵損失訓練。發現一個極小的[17, 1, 14] KAN能夠達到81.6%的測試準確率(而Deepmind的4層寬度300的MLP達到78%的測試準確率)。[17, 1, 14] KAN(G = 3, k = 3)有\(\approx 200\)個參數,而MLP有\(\approx 3 \times 10^5\)個參數,如表4所示。值得注意的是,KAN可以同時比MLP更準確和更節省參數。

在可解釋性方面,根據每個激活的大小來縮放其透明度,因此無需特徵歸因就可以立即清楚哪些輸入變量很重要(見圖4.3左):signature主要依賴於\(\mu_r\),略微依賴於\(\mu_i\)和\(\lambda\),而對其他變量的依賴很小。然後,在三個重要變量上訓練一個[3, 1, 14] KAN,獲得78.2%的測試準確率。結果與文獻中的結果有一個微妙的區別:他們發現signature主要依賴於\(\mu_i\),而該研究發現signature主要依賴於\(\mu_r\)。這種差異可能是由於微妙的算法選擇,但促使該研究進行以下實驗:
(a)消融研究。表明\(\mu_r\)對準確性的貢獻大於\(\mu_i\)(見圖4.3):例如,僅\(\mu_r\)就可以達到65.0%的準確率,而僅\(\mu_i\)只能達到43.8%的準確率。 
(b)找到一個只涉及\(\mu_r\)和\(\lambda\)的符號公式(在表5中),但可以達到77.8%的測試準確率。

為了研究(2),即獲得\(\sigma\)的符號形式,我們將問題表述為回歸任務。使用2.5.1節中介紹的自動符號回歸,可以將訓練好的KAN轉換為符號公式。訓練形狀為[3, 1]、[3, 1, 1]、[3, 2, 1]的KAN,其對應的符號公式顯示在表5 B-D中。很明顯,通過使用更大的KAN,準確性和複雜性都會增加。因此,KAN提供的不僅僅是一個單一的符號公式,而是一個完整的公式帕累托前沿,在簡單性和準確性之間進行權衡。然而,KAN需要額外的歸納偏置來進一步簡化這些方程,以重新發現中的公式(表5 A)。該研究測試了兩種情況:
(1)在第一種情況下,假設真實公式具有多元Pade表示(兩個多元Taylor級數的商)。首先訓練[3, 2, 1],然後將其擬合到Pade表示。可以獲得表5中的公式E,它與Deepmind的公式有相似之處。 
(2)假設除法對KAN來說不是很直觀,因此訓練兩個KAN(一個用於分子,另一個用於分母)並手動將它們相除。令人驚訝的是,最終得到了公式F(在表5中),它只涉及\(\mu_r\)和\(\lambda\),儘管\(\mu_i\)也被提供但被KAN忽略了。

到目前為止,已經重新發現了主要結果。令人矚目的是,KAN使這一發現變得非常直觀和方便。與其使用特徵歸因方法(這是很好的方法),不如直接盯著KAN的可視化。此外,自動符號回歸也使符號公式的發現變得容易得多。

在下一部分中,我們提出了一種Deepmind論文中沒有包括的新範式"AI for Math",我們旨在使用KAN的無監督學習模式來發現紐結不變量中的更多關係(除了signature)。





無監督學習 正如在4.2節中提到的,無監督學習是更有前景的設置,因為它避免了手動劃分輸入和輸出變量,而手動劃分有組合多種可能性。在無監督學習模式下,將所有18個變量(包括signature)視為輸入,使它們處於同等地位。紐結數據是正樣本,通過隨機洗牌特徵來獲得負樣本。訓練一個[18, 1, 1] KAN來分類給定的特徵向量是否屬於正樣本(1)或負樣本(0)。手動將第二層激活設置為在零處中心為一的高斯函數,因此正樣本在(接近)零處有激活,隱含地給出了紐結不變量之間的關係 

$$\sum_{i=1}^{18} g_i(x_i) = 0$$

其中\(x_i\)表示一個特徵(不變量),\(g_i\)是相應的激活函數,可以從KAN圖中直接讀出。用\(\lambda = {10^{-2}, 10^{-3}}\)訓練KAN以偏好稀疏的輸入組合,種子為\(\text{seed} = {0, 1, \cdots, 99}\)。所有200個網路可以分為三個簇,具有代表性的KAN顯示在圖4.4中。這三組相依變量是:
  1. 第一組相依變量是signature、子午距的實部和經向距(加上另外兩個可以因(3)而移除的變量)。這就是上面研究的signature依賴關係,所以看到這個依賴關係在無監督模式下再次被發現是非常有趣的。
  2. 第二組變量涉及尖點體積\(V\)、子午平移\(\mu_r\)的實部和經向平移\(\lambda\)。它們的激活函數看起來都像對數函數(可以通過2.5.1節中暗示的符號功能驗證)。因此,關係是\(-\log V + \log \mu_r + \log \lambda = 0\),等價於\(V = \mu_r \lambda\),這在定義上是正確的。然而,在沒有任何先驗知識的情況下發現這個關係,這是令人欣慰的。
  3. 第三組變量包括短測地線\(g_r\)的實部和內射半徑。它們的激活看起來定性相同,但相差一個負號,因此推測這兩個變量存在線性相關。我們繪製2D散點圖,發現\(2r\)上限為\(g_r\),這也是一個眾所周知的關係。
有趣的是,KAN的無監督模式可以重新發現幾個已知的數學關係。好消息是KAN發現的結果可能是可靠的;壞消息是還沒有發現任何新東西。值得注意的是,選擇了一個淺層KAN以便於簡單可視化,但如果存在,更深層的KAN可能會發現更多關係。希望在未來的工作中研究如何用更深層的KAN發現更複雜的關係。

4.4 物理應用:Anderson局域化

Anderson局域化是量子系統中無序導致電子波函數局域化的基本現象,導致所有傳輸停止。在一維和二維中,縮放論證表明,對於無限小量的隨機無序,所有電子本征態都是指數局域化的。相比之下,在三維中,臨界能量形成了一個相邊界,將延展態與局域化態分開,稱為遷移率邊緣。對這些遷移率邊緣的理解對於解釋各種基本現象至關重要,如固體中的金屬-絕緣體轉變,以及光在光子器件中的局域化效應。因此,有必要開發展現遷移率邊緣的微觀模型,以便進行詳細研究。在低維中開發此類模型通常更實用,在低維中引入準周期性而不是隨機無序也可能導致分隔局域化相和延展相的遷移率邊緣。此外,分析遷移率邊緣的實驗實現可以幫助解決關於相互作用系統中局域化的爭論。事實上,最近的幾項研究集中在確定此類模型並推導它們的遷移率邊緣的精確解析表達式上。

在這裡,將KAN應用於從準周期緊束縛模型生成的數值資料,以提取它們的遷移率邊緣。我們特地檢查三類模型:馬賽克模型(MM)、廣義Aubry-André模型(GAAM)和修正的Aubry-André模型(MAAM)。對於MM,證明KAN能夠準確提取遷移率邊緣作為能量的1D函數。對於GAAM,發現從KAN獲得的公式與基本事實非常接近。對於更複雜的MAAM,再次展示了該框架符號可解釋性的另一個例子。用戶可以通過"協作"的方式簡化從KAN(和相應的符號公式)獲得的複雜表達式,在協作中,人類生成假設以獲得更好的匹配(例如,對某些激活函數的形式做出假設),然後KAN可以快速進行假設檢驗。

為了量化這些模型中態的局域化,通常使用反參與率(IPR)。第\(k\)個本征態\(\psi^{(k)}\)的IPR由下式給出

$$\text{IPR}_k = \frac{\sum_n |\psi_n^{(k)}|^4}{(\sum_n |\psi_n^{(k)}|^2)^2}$$ 
(4.2)

其中求和是對格點指標進行的。在這裡,使用相關的局域化度量——態的分形維數,由下式給出


$$D_k = -\frac{\log(\text{IPR}_k)}{\log(N)}$$ 
(4.3)


其中\(N\)是系統大小。\(D_k = 0(1)\)表示局域化(延展)態。

馬賽克模型(MM) 我們首先考慮由哈密頓量定義的一類緊束縛模型


$$H = t\sum_n (c_{n+1}^\dagger c_n + \text{H.c.}) + \sum_n V_n(\lambda, \phi)c_n^\dagger c_n,$$ (4.4)


其中$t$是最近鄰耦合,\(c_n(c_n^\dagger)\)是格點\(n\)處的湮滅(產生)算符,勢能\(V_n\)由下式給出


\[ V_n(\lambda, \phi) = \lambda \cos(2\pi nb + \phi) \cdot \mathbf{1}_{\{j = m\kappa\}} \]
(4.5)


為了引入準周期性,將\(b\)設置為無理數(選擇\(b\)為黃金比例\(\frac{1+\sqrt{5}}{2}\))。\(\kappa\)是一個整數,準周期勢以間隔\(\kappa\)出現。對於這個模型,能量(\(E\))譜通常包含由遷移率邊緣分隔的延展和局域化區域。有趣的是,這裡發現的一個獨特特徵是,即使對於任意強的準周期勢(即系統中總是存在延展態,與局域化態共存),遷移率邊緣也存在。

遷移率邊緣可以用\(g(\lambda, E) \equiv \lambda - |f_\kappa(E)| = 0\)來描述。\(g(\lambda, E) > 0\)和\(g(\lambda, E) < 0\)分別對應於局域化相和延展相。因此,學習遷移率邊緣取決於學習"序參量"\(g(\lambda, E)\)。誠然,這個問題可以通過許多其他理論方法來解決這類模型,將在下面演示,KAN框架已經準備好並方便地接受來自人類用戶的假設和歸納偏置。

假設一個假想的用戶Alice,她是凝聚態物理學的一名新博士生,並獲得一個[2, 1] KAN作為該任務的助手。首先,她理解這是一個分類任務,因此明智的做法是使用fix_symbolic功能將第二層中的激活函數設置為sigmoid。其次,她意識到學習整個2D函數\(g(\lambda, E)\)是不必要的,因為最終她只關心由\(g(\lambda, E) = 0\)確定的\(\lambda = \lambda(E)\)。這樣做,假設\(g(\lambda, E) = \lambda - h(E) = 0\)是合理的。Alice只需再次使用fix_symbolic功能將\(\lambda\)的激活函數設置為線性。現在Alice訓練KAN網路,方便地獲得遷移率邊緣,如圖4.5所示。Alice既可以獲得直觀的定性理解(底部),也可以獲得定量結果(中部),與基本事實(頂部)吻合良好。



廣義Andre-Aubry模型(GAAM)接下來,考慮由 Hamiltonian定義的一類緊束縛模型


$$H = t\sum_n (c_{n+1}^\dagger c_n + \text{H.c.}) + \sum_n V_n(\alpha, \lambda, \phi)c_n^\dagger c_n,$$ 
(4.6)

其中\(t\)是最近鄰耦合,\(c_n(c_n^\dagger)\)是格點\(n\)處的湮滅(產生)算符,勢能\(V_n\)由下式給出


$$V_n(\alpha, \lambda, \phi) = 2\lambda\frac{\cos(2\pi nb + \phi)}{1 - \alpha \cos(2\pi nb + \phi)},$$
 (4.7)

對於\(\alpha \in (-1, 1)\),它是光滑的。為了引入準周期性,我們再次將\(b\)設置為無理數(特別地,我們選擇\(b\)為黃金比例)。和之前一樣,我們想得到遷移率邊緣的表達式。對於這些模型,遷移率邊緣由閉合形式表達式給出,


$$\alpha E = 2(t - \lambda).$$ 
(4.8)

隨機採樣模型參數:\(\phi\)、\(\alpha\)和\(\lambda\)(設置能量尺度\(t = 1\)),並計算能量本征值以及相應本征態的分形維數,這構成了訓練資料集。 

在這裡,要學習的"序參量"是\(g(\alpha, E, \lambda, \phi) = \alpha E + 2(\lambda - 1)\),遷移率邊緣對應於\(g = 0\)。再次假設Alice想要確定遷移率邊緣,但只能訪問IPR或分形維數資料,因此她決定使用KAN來幫助她完成任務。Alice希望模型盡可能小,因此她可以從一個大模型開始,使用自動修剪來獲得一個小模型,或者她可以根據對給定問題複雜性的理解來猜測一個合理的小模型。無論哪種方式,假設她得到一個[4, 2, 1, 1] KAN。首先,她將最後一個激活設置為sigmoid,因為這是一個分類問題。



她用一些稀疏正則化訓練她的KAN達到98.7%的準確率,並在圖4.6(a)步驟1中可視化訓練後的KAN。她觀察到KAN完全沒有選擇$\phi$,這使她意識到遷移率邊緣與\(\phi\)無關(與公式(4.8)一致)。此外,她觀察到幾乎所有其他激活函數都是線性或二次方的,因此她打開自動符號捕捉,將函數庫限制為僅線性或二次方。之後,她立即得到一個已經是符號的網路(如圖4.6(a)步驟2所示),具有相當(甚至略好)的準確率98.9%。通過使用symbolic_formula功能,Alice方便地得到\(g\)的符號形式,如表6 GAAM-KAN auto(第三行)所示。也許她想劃掉一些小項並將係數捕捉為小整數,這使她接近真正的答案。

如果Alice使用符號回歸方法,這個假設的故事將完全不同。如果她幸運的話,SR可以返回完全正確的公式。然而,在絕大多數情況下,SR不會返回有用的結果,而且Alice不可能"調試"或與符號回歸的底層過程交互。此外,在運行SR之前,Alice可能會感到不舒服/缺乏經驗,無法將符號項的函數庫作為先驗知識提供給SR。相比之下,在KAN中,Alice不需要將任何先驗信息放入KAN。她可以先盯著訓練後的KAN獲得一些線索,只有在那時她才決定她想做出哪些假設(例如,"所有激活都是線性或二次方的"),並在KAN中實現她的假設。儘管KAN不太可能立即返回正確答案,但KAN總是會返回一些有用的東西,Alice可以與之合作來完善結果。

修正的Andre-Aubry模型(MAAM)考慮的最後一類模型由Hamiltonian定義


$$H = \sum_{n \neq n'} te^{-p|n-n'|} (c_n^\dagger c_{n'} + \text{H.c.}) + \sum_n V_n(\lambda, \phi)c_n^\dagger c_n,$$ 
(4.9)

其中\(t\)是空間中指數衰減耦合的強度,\(c_n(c_n^\dagger)\)是格點\(n\)處的湮滅(產生)算符,勢能\(V_n\)由下式給出


$$V_n(\lambda, \phi) = \lambda \cos(2\pi nb + \phi),$$ 
(4.10)


和之前一樣,為了引入準周期性,將\(b\)設置為無理數(黃金比例)。對於這些模型,遷移率邊緣由閉合形式表達式給出,


$$\lambda \cosh(p) = E + t = E + t_1 \exp(p)$$ 
(4.11)

其中定義\(t_1 \equiv te^{-p}\)為最近鄰跳躍強度,下面設\(t_1 = 1\)。

假設Alice想要確定MAAM的遷移率邊緣。這個任務更加複雜,需要更多人類智慧。和上一個例子一樣,Alice從一個[4, 2, 1, 1] KAN開始訓練,但得到的準確率在75%左右,低於可接受的水平。然後她選擇一個更大的[4, 3, 1, 1] KAN,成功獲得98.4%的準確率,這是可以接受的(圖4.6(b)步驟1)。Alice注意到KAN沒有選擇\(\phi\),這意味著遷移率邊緣與相位因子\(\phi\)無關(與公式(4.11)一致)。如果Alice打開自動符號回歸(使用包含exp、tanh等的大型函數庫),她會得到表6-MAAM-KAN auto中的複雜公式,準確率為97.1%。但是,如果Alice想找到一個更簡單的符號公式,她會想使用手動模式,自己完成符號捕捉。在此之前,她發現訓練後的[4, 3, 1, 1] KAN可以被修剪為[4, 2, 1, 1],同時保持97.7%的準確率(圖4.6(b))。Alice可能認為除了依賴\(p\)的激活函數外,所有激活函數都是線性或二次方的,並通過使用fix_symbolic手動將它們捕捉為線性或二次方。捕捉並重新訓練後,更新的KAN如圖4.6(c)步驟3所示,保持97.7%的準確率。從現在開始,Alice可能會根據她的先驗知識做出兩個不同的選擇。在一種情況下,Alice可能已經猜到\(p\)的依賴關係是\(\cosh\),因此她將\(p\)的激活設置為\(\cosh\)函數。她重新訓練KAN,獲得96.9%的準確率(圖4.6(c)步驟4A)。在另一種情況下,Alice不知道\(\cosh p\)的依賴關係,因此她追求簡單性,再次假設\(p\)的函數是二次方的。她重新訓練KAN,獲得95.4%的準確率(圖4.6(c)步驟4B)。如果她同時嘗試了這兩種情況,她會意識到在準確性方面\(\cosh\)更好,而在簡單性方面二次方更好。這些步驟對應的公式列在表6中。很明顯,Alice進行的手動操作越多,符號公式就越簡單(準確性略有犧牲)。KAN有一個"旋鈕",用戶可以調整以在簡單性和準確性之間進行權衡(有時簡單性甚至可以提高準確性,如GAAM案例中所示)。

5 相關工作

Kolmogorov-Arnold定理與神經網路。Kolmogorov-Arnold定理(KAT)與神經網路之間的聯繫在文獻中並不新鮮,但內函數的病態行為使KAT在實踐中看起來並不樂觀。大多數這些先前的工作都堅持使用原始的2層寬度為\((2n + 1)\)的網路,表達能力有限,其中許多工作甚至早於反向傳播。因此,大多數研究都建立在理論之上,只有相當有限或人工的玩具實驗。該研究的貢獻在於將網路推廣到任意寬度和深度,在當今的深度學習潮流中重新喚起並將其情境化,並強調其作為AI+科學基礎模型的潛在作用。

神經縮放定律(NSL)。NSL是測試損失相對於模型大小、資料、計算等呈冪律行為的現象。NSL的起源仍然是個謎,但有競爭理論,包括內在維度、任務量化、資源理論、隨機特徵、組合稀疏性和最大奇異性。本文通過表明,如果高維函數具有平滑的Kolmogorov-Arnold表示,則可以驚人地縮放為1D函數(這是人們可以希望的最佳界限),為這一空間做出了貢獻。論文為神經縮放定律帶來了新的樂觀情緒,因為它承諾了有史以來最快的縮放指數。在實驗中表明,這種快速神經縮放定律可以在合成數據集上實現,但未來的研究需要解決以下問題:KA表示是否存在於一般任務中?如果存在,訓練是否在實踐中找到這些表示?

機制可解釋性(MI)。MI是一個新興領域,旨在機械地理解神經網路的內部工作原理。MI研究大致可分為被動和主動MI研究。大多數MI研究是被動的,專注於理解用標準方法訓練的現有神經網路。主動MI研究試圖通過設計本質上可解釋的架構或開發訓練方法來明確鼓勵可解釋性來實現可解釋性。該研究的工作屬於第二類,其中模型和訓練方法在設計上是可解釋的。

可學習激活。在機器學習中,可學習激活在神經網路中的想法並不新鮮。可訓練激活函數以可微方式學習或以離散方式搜索。激活函數被參數化為多項式、樣條、sigmoid線性單元或神經網路。KAN使用B樣條來參數化其激活函數。該研究還介紹了關於可學習激活網路(LAN)的初步結果,其性質介於KAN和MLP之間,為了專注於主論文中的KAN,結果被推遲到附錄B。

符號回歸。有許多現成的基於遺傳算法的符號回歸方法(Eureka 、GPLearn 、PySR )、基於神經網路的方法(EQL 、OccamNet )、基於物理啟發的方法(AI Feynman )和基於強化學習的方法。KAN與基於神經網路的方法最相似,但與之前的工作不同,激活函數是在符號捕捉之前連續學習的,而不是手動固定的。

物理信息神經網路(PINN)和物理信息神經算子(PINO)。在3.4小節中,演示了KAN可以取代在求解PDE時使用MLP來施加PDE損失的範式。參考Deep Ritz方法、PINN用於PDE求解,以及Fourier神經算子、PINO 、DeepONet 用於學習解映射的算子學習方法。在上述所有網路中,都有可能用KAN取代MLP。

AI for Mathematics。正如我們在4.3小節中看到的,AI最近被應用於紐結理論中的幾個問題,包括檢測一個紐結是否是未紐結或帶紐結,以及預測紐結不變量並揭示它們之間的關係。

6 討論

在本節中,我們從數學基礎、算法和應用的角度討論KAN的局限性和未來方向。

數學方面:我們提出了KAN的初步數學分析(定理2.1),但對它們的數學理解仍然非常有限。Kolmogorov-Arnold表示定理在數學上已經得到了徹底的研究,但該定理對應於形狀為[n, 2n+1, 1]的KAN,這是KAN的一個非常受限的子類。在更深的KAN上的經驗成功是否意味著數學中的某些基本原理?一個有吸引力的廣義Kolmogorov-Arnold定理可以定義超出2層組合的"更深"Kolmogorov-Arnold表示,並可能將激活函數的平滑性與深度聯繫起來。假設存在一些函數,它們不能在原始(2層)Kolmogorov-Arnold表示中平滑表示,但可能在3層或更深層中平滑表示。我們能否使用這種"Kolmogorov-Arnold深度"的概念來表徵函數類?

算法方面:該研究討論以下內容:
  1. 準確性。架構設計和訓練中的多個選擇沒有得到充分研究,因此替代方案可能進一步提高準確性。例如,樣條激活函數可能被徑向基函數或其他局部核取代。可以使用自適應網格策略。
  2. 效率。KAN運行緩慢的一個主要原因是不同的激活函數無法利用批量計算(大量數據通過相同的函數)。實際上,可以在激活函數全部相同(MLP)和全部不同(KAN)之間進行插值,將激活函數分組為多個組("多頭"),其中一個組內的成員共享相同的激活函數。
  3. KAN和MLP的混合。與MLP相比,KAN有兩個主要區別:
    • 激活函數在邊上而不是在節點上,
    • 激活函數是可學習的而不是固定的。
  4. 哪種變化對解釋KAN的優勢更重要?在附錄B中介紹了初步結果,其中研究了一個模型,它具有(ii),即激活函數是可學習的(如KAN),但不具有(i),即激活函數在節點上(如MLP)。此外,還可以構建另一個模型,其激活函數是固定的(如MLP),但在邊上(如KAN)。
  5. 自適應性。得益於樣條基函數的內在局部性,可以在KAN的設計和訓練中引入自適應性,以提高準確性和效率:參見多層次訓練的想法,如中的多網格方法,或如中的多尺度方法的域相關基函數。
應用方面:提出了一些初步證據,表明KAN在與科學相關的任務中比MLP更有效,例如擬合物理方程和求解PDE。期望KAN在求解Navier-Stokes方程、密度泛函理論或任何其他可以表述為回歸或求解PDE的任務中也很有前景。還希望將KAN應用於與機器學習相關的任務,這需要將KAN集成到當前的架構中,例如將transformer中的MLP替換為KAN,人們可以出"kansformer"。

作為AI+科學的"語言模型"的KAN 大型語言模型之所以如此變革,是因為它們對任何能說自然語言的人都有用。科學的語言是函數。KAN由可解釋的函數組成,因此當人類用戶盯著KAN時,就像用函數的語言與之交流一樣。這一段旨在促進AI-科學家-協作範式,而不是特定工具KAN。就像人們使用不同的語言交流一樣,預計在未來,KAN將只是AI+科學的語言之一,儘管KAN將是最早實現AI和人類交流的語言之一。然而,在KAN的支持下,AI-科學家-協作範式從未如此簡單和方便,使我們重新思考要如何接近AI+科學的範式:想要AI科學家,還是想要幫助科學家的AI?(完全自動化)AI科學家的內在困難在於,很難使人類偏好量化,這將人類偏好編入AI目標。事實上,不同領域的科學家對哪些函數是簡單或可解釋的可能有不同的感受。因此,科學家擁有一個能說科學語言(函數)並能方便地與個別科學家的歸納偏置交互以適應特定科學領域的AI更為可取。

最終要點:我應該使用KAN還是MLP?

目前,KAN的最大瓶頸在於其緩慢的訓練。在給定相同數量的參數的情況下,KAN通常比MLP慢10倍。應該誠實地說,該研究有努力優化KAN的效率,因此KAN的緩慢訓練更像是一個工程問題,未來有待改進,而不是一個根本的限制。如果想要快速訓練模型,應該使用MLP。然而,在其他情況下,KAN應該與MLP相當或更好,這使得它們值得嘗試。圖6.1中的決策樹可以幫助決定何時使用KAN。簡而言之,如果你關心可解釋性和/或準確性,並且緩慢的訓練不是主要問題,建議嘗試KAN。



參考文獻(略)


附錄

A KAN功能

表7包含了用戶可能會發現有用的常見功能。

B 可學習激活網路(LAN)

B.1 架構

除了KAN,該研究還提出了另一種可學習激活網路(LAN),它們幾乎是MLP,但具有參數化為樣條的可學習激活函數。KAN對標準MLP進行了兩項主要更改:
  1. 激活函數變為可學習而不是固定的; 
  2. 激活函數放置在邊上而不是節點上。為了區分這兩個因素,該研究還提出了可學習激活網路(LAN),它只有可學習的激活,但仍然在節點上,如圖B.1所示。

對於寬度為\(N\)、深度為$L$和網格點數為\(G\)的LAN,參數數量為\(N^2L + NLG\),其中\(N^2L\)是權重矩陣的參數數量,\(NLG\)是樣條激活的參數數量,除MLP外,這引起的開銷很小,因為通常\(G \ll N\),所以\(NLG \ll N^2L\)。LAN與MLP相似,因此可以從預訓練的MLP初始化,並通過允許可學習的激活函數進行微調。一個示例是使用LAN來改進SIREN,在B.3節中介紹。

LAN和KAN的比較
    
    LAN的優點:
  1. LAN在概念上比KAN更簡單。它們更接近標準MLP(唯一的變化是激活函數變為可學習)。
  2. LAN比KAN更容易擴展。LAN/KAN在節點/邊上有可學習的激活函數。因此,LAN/KAN中的激活參數分別隨\(N\)/\(N^2\)縮放,其中\(N\)是模型寬度。
    LAN的缺點:
  1. LAN似乎不太可解釋(權重矩陣難以解釋,就像MLP中一樣);
  2. LAN似乎也不如KAN準確,但仍然比MLP更準確。與KAN一樣,如果LAN的激活函數由樣條參數化,LAN也承認網格擴展。




B.2 LAN可解釋性結果

我們在圖B.2中介紹了LAN的初步可解釋性結果。對於圖4.1中KAN完全可解釋的相同示例,由於權重矩陣的存在,LAN似乎可解釋性要差得多。首先,權重矩陣不如可學習的激活函數容易解釋。其次,權重矩陣帶來了太多的自由度,使得可學習的激活函數過於不受約束。我們使用LAN的初步結果似乎暗示,去除線性權重矩陣(通過在邊上具有可學習的激活,如KAN)對於可解釋性是必要的。

B.3 擬合圖像(LAN)

隱式神經表示將圖像視為2D函數\(f(x, y)\),其中像素值\(f\)是像素的兩個坐標\(x\)和\(y\)的函數。為了壓縮圖像,這樣的隱式神經表示(\(f\)是神經網路)可以在保持幾乎原始圖像質量的同時實現令人印象深刻的參數壓縮。SIREN提出使用具有周期性激活函數的MLP來擬合函數\(f\)。在LAN中自然考慮其他激活函數。然而,由於我們將LAN激活初始化為平滑的,而SIREN需要高頻特徵,LAN不能立即工作。請注意,LAN中的每個激活函數都是基函數和樣條函數的總和,即\(\phi(x) = b(x) + \text{spline}(x)\),我們將\(b(x)\)設置為正弦函數,與SIREN中的設置相同,但讓\(\text{spline}(x)\)是可訓練的。對於MLP和LAN,形狀都是[2,128,128,128,128,128,1]。我們使用Adam優化器對它們進行訓練,批量大小為4096,學習率為$10^{-3}$訓練5000步,學習率為$10^{-4}$訓練5000步。如圖B.3所示,LAN(橙色)可以實現比MLP(藍色)更高的PSNR,這是由於LAN微調激活函數的靈活性。我們表明,也可以從MLP初始化LAN,並進一步微調LAN(綠色)以獲得更好的PSNR。在我們的實驗中,我們選擇了\(G = 5\),因此額外的參數增加大約是原始參數的\(G/N = 5/128 \approx 4\%\)。


C 超參數的影響

我們在圖C.1中顯示了超參數對\(f(x, y) = \exp(\sin(\pi x) + y^2)\)案例的影響。為了獲得可解釋的圖,我們希望活躍激活函數的數量盡可能少(理想情況下為3)。

  1.  我們需要熵懲罰來減少活躍激活函數的數量。如果沒有熵懲罰,會有許多重複函數。
  2.  結果可能取決於隨機種子。對於一些不幸的種子,修剪後的網絡可能比需要的更大。
  3.  整體懲罰強度\(\lambda\)有效地控制稀疏性。
  4.  網格數\(G\)也對可解釋性有微妙影響。當\(G\)太小時,由於每個激活函數不是很有表現力,網絡傾向於使用集成策略,使得解釋更加困難。
  5.  分段多項式階數\(k\)對可解釋性只有微妙影響。然而,它的行為有點像隨機種子,在這個玩具示例中沒有顯示任何可見的模式。

D 費曼KAN

我們在費曼資料集(第3.3節)上包含了更多結果。圖D.1顯示了每個費曼數據集的KAN和MLP的帕累托前沿。圖D.3和D.2將每個費曼方程擬合任務的最小KAN(在測試RMSE < 10−2 的約束下)和最佳KAN(具有最低測試RMSE損失)可視化。


E 關於網格大小的備註

對於PDE和回歸任務,當我們在均勻網格上選擇訓練數據時,當網格大小更新到一個大水平時,我們會看到訓練損失突然增加(即性能突然下降),相當於空間一個方向上的不同訓練點。這可能是由於更高維度中B樣條的實現,需要進一步調查。

F 特殊函數的KAN

我們在特殊函數數據集(第3.2節)上包含了更多結果。圖F.2和F.1將每個特殊函數擬合任務的最小KAN(在測試RMSE < 0.01的約束下)和最佳KAN(具有最低測試RMSE損失)可視化。



幾點見解和思考:



1. Kolmogorov-Arnold定理在神經網路領域的應用前景廣闊。論文將KA定理與深度學習結合,提出了KAN這一新穎的神經網路架構。KAN不僅在各種任務上展現出優異的性能,其可解釋性也遠勝於傳統的MLP。這為深度學習模型的可解釋性研究提供了新的思路。

2. KAN在速度上的不足可能限制了其應用。論文坦承,目前KAN的訓練速度比MLP慢10倍左右。作者認為這更多是工程實現的問題,未來有優化的空間。KAN要真正發揮作用,還需要在效率上有所突破。

3. 無監督學習可能是AI輔助科學發現的重要範式。論文在紐結理論的實驗中展示了無監督學習的威力。相比於有監督學習需要專家標註數據,無監督學習可以直接從原始數據出發尋找隱藏的關係,更符合科學發現的實際過程。

4. AI與科學家的協同是大勢所趨。KAN作為一種會"說"科學語言(數學函數)的AI系統,為AI與科學家的溝通提供了渠道。未來,懂科學語言的AI助手或許能夠與科學家進行更加深入的交流,共同推進科學發現的進程。

5. 神經網路的數學基礎仍有待加強。儘管KAN取得了令人矚目的表現,但我們對其背後的數學原理還缺乏深入理解。Kolmogorov-Arnold定理能否推廣到更深層次?網路深度與函數平滑性是否存在內在聯繫?這些問題的解答,需要數學和深度學習領域的專家共同努力。

6. 對科學家而言,擁有正確的工具選擇十分重要。論文最後給出的決策樹,為科學家在KAN和MLP間做出選擇提供了參考。一方面,科學家要善於利用新工具(如KAN)來加速研究;另一方面,也要理性看待新工具的局限性,根據具體問題選擇最合適的方法。

這是一篇將數學理論與深度學習緊密結合,並在多個科學問題上進行實踐的高質量論文。它不僅展示了KAN的強大性能,更重要的是,它體現了一種AI助力科學發現的新範式。儘管KAN目前還不夠完善,但其背後的思想值得我們深入探索。期待未來有更多類似的工作,推動AI與科學的深度融合。












沒有留言:

發佈留言

SambaNova SN40L: 利用Dataflow和專家組合(COE)來克服AI記憶牆的大模型

摘要 GPT-4等整體式大型語言模型(LLM)為現代生成AI應用鋪路。然而,大規模訓練、服務及維護整體式LLM仍然極其昂貴和充滿挑戰。現代AI加速器計算能力與記憶體比例的不成比例增長已經造成了記憶體壁障,需要新的方法來部署AI。最近的研究顯示,許多小型專家模型的組合,每個模型參數...