Transformer要變Kansformer?用了幾十年的MLP迎來挑戰者KAN

机器之心發表於2024-05-03
MLP(多層感知器)用了幾十年了,真的沒有別的選擇了嗎?

多層感知器(MLP),也被稱為全連線前饋神經網路,是當今深度學習模型的基礎構建塊。

MLP 的重要性無論怎樣強調都不為過,因為它們是機器學習中用於逼近非線性函式的預設方法。

然而,MLP 是否就是我們能夠構建的最佳非線性迴歸器呢?儘管 MLP 被廣泛使用,但它們存在明顯的缺陷。例如,在 Transformer 模型中,MLP 幾乎消耗了所有非嵌入式引數,並且通常在沒有後處理分析工具的情況下,相對於注意力層來說,它們的可解釋性較差。

所以,是否有一種 MLP 的替代選擇?

今天,KAN 出現了。

圖片

圖片

圖片

這是一個靈感來源於 Kolmogorov-Arnold 表示定理的網路。

  • 連結:https://arxiv.org/pdf/2404.19756

  • Github:https://github.com/KindXiaoming/pykan

該研究一經發布,就在國外社交平臺引起了廣泛的關注與討論。

有網友稱,Kolmogorov 早在 1957 年就發現了多層神經網路,比 Rumerhart、Hinton 和 William 的 1986 年論文發表的時間要早得多,但他卻被西方忽視了。

圖片

也有網友表示,這篇論文釋出意味著深度學習的喪鐘已經敲響。

圖片

有網友思考,該研究是否會像 Transformer 的論文一樣具有顛覆性。

圖片

但也有作者表示,他們在 2018-19 年就基於改進的 Kolmogrov-Gabor 技術做了同樣的事。

圖片

接下來,讓我們看下這篇論文講了什麼?

論文概覽

本文提出了一種有前景的多層感知器(MLP)的替代方案,稱為 Kolmogorov-Arnold Networks(KAN)。MLP 的設計靈感來源於通用近似定理,而 KAN 的設計靈感則來源於 Kolmogorov-Arnold 表示定理。與 MLP 類似,KAN 擁有全連線的結構。然而,MLP 在節點(神經元)上放置固定啟用函式,KAN 則在邊(權重)上放置可學習的啟用函式,如圖 0.1 所示。因此,KAN 完全沒有線性權重矩陣:每個權重引數都被替換為一個可學習的一維函式,引數化為樣條(spline)。KAN 的節點僅對傳入訊號進行求和,而不應用任何非線性變換。

圖片

有人可能擔心 KAN 的成本過高,因為每個 MLP 的權重引數都變成了 KAN 的樣條函式。不過,KAN 允許的計算圖比 MLP 要小得多。例如,研究者展示了 PED 求解:一個兩層寬度為 10 的 KAN 比一個四層寬度為 100 的 MLP 精確 100 倍(MSE 分別為 10^-7 和 10^-5 ),並且在引數效率上也提高了 100 倍(引數量分別為 10^2 和 10^4 )。

使用 Kolmogorov-Arnold 表示定理來構建神經網路的可能性已經被研究過。不過大多數工作都停留在原始的深度為 2、寬度為 (2n + 1) 的表示上,並且沒有機會利用更現代的技術(例如,反向傳播)來訓練網路。本文的貢獻在於將原始的 Kolmogorov-Arnold 表示泛化到任意寬度和深度,使其在當今的深度學習領域煥發新生,同時利用大量的實證實驗來突出其作為「AI + 科學」基礎模型的潛在作用,這得益於 KAN 的準確性和可解釋性。

儘管 KAN 數學解釋能力不錯,但實際上它們只是樣條和 MLP 的組合,利用了二者的優點,避免了缺點的出現。樣條在低維函式上準確度高,易於區域性調整,並且能夠在不同解析度之間切換。然而,由於樣條無法利用組合結構,因此它們存在嚴重 COD 問題。另一方面,MLP 由於其特徵學習能力,較少受到 COD 的影響,但在低維空間中卻不如樣條準確,因為它們無法最佳化單變數函式。

為了準確學習一個函式,模型不僅應該學習組合結構(外部自由度),還應該很好地近似單變數函式(內部自由度)。KAN 就是這樣的模型,因為它們在外部類似於 MLP,在內部類似於樣條。結果,KAN 不僅可以學習特徵(得益於它們與 MLP 的外部相似性),還可以將這些學習到的特徵最佳化到很高的精度(得益於它們與樣條的內部相似性)。

例如,對於一個高維函式:

圖片

當 N 很大時,樣條會因為 COD 而失敗;MLP 雖然有可能學習到廣義的加性結構,但使用例如 ReLU 啟用函式來近似指數和正弦函式卻非常低效。相比之下,KAN 能夠很好地學習組合結構和單變數函式,因此以很大的優勢超越了 MLP(見圖 3.1)。

圖片

在本篇論文中,研究者展示了大量的實驗數值,體現了 KAN 在準確性和可解釋性方面對 MLP 的顯著改進。論文的結構如下圖 2.1 所示。程式碼可在 https://github.com/KindXiaoming/pykan 獲取,也可以透過 pip install pykan 安裝。

圖片

Kolmogorov-Arnold 網路 (KAN)

Kolmogorov-Arnold 表示定理

Vladimir Arnold 和 Andrey Kolmogorov 證明了,如果 f 是一個在有界域上的多變數連續函式,那麼 f 可以寫成一個單變數連續函式和二元加法運算的有限組合。更具體地說,對於一個平滑的函式 f : [0, 1]^n → R,它可以表示為:

圖片其中圖片以及圖片

在某種意義上,他們展示了唯一的真正多變數函式是加法,因為所有其他函式都可以透過單變數函式和求和來表示。有人可能會認為這對機器學習是個好訊息:學習一個高維函式可以歸結為學習多項式數量的一維函式。然而,這些一維函式可能是非平滑的,甚至是分形的,因此在實踐中可能無法學習。因此,Kolmogorov-Arnold 表示定理機器學習中基本上被判處了死刑,被認為是理論上正確但實踐中無用的。

然而,研究者對 Kolmogorov-Arnold 定理在機器學習中的實用性持更樂觀的態度。首先,不必堅持原始的方程,它只有兩層非線性和一個隱藏層中的少量項(2n + 1):研究者將將網路泛化為任意寬度和深度。其次,科學和日常生活中的大多數函式通常是平滑的,並且具有稀疏的組合結構,這可能有助於平滑的 Kolmogorov-Arnold 表示。

KAN 架構

假設有一個監督學習任務,由輸入輸出對 {x_i , y_i} 組成,研究者希望找到一個函式 f,使得對於所有資料點 y_i ≈ f (x_i) 。方程(2.1)意味著,如果能找到適當的單變數函式圖片圖片,那麼任務就完成了。這啟發研究者設計一個顯式引數化方程(2.1)的神經網路。由於所有要學習的函式都是單變數函式,研究者將每個一維函式引數化為 B 樣條曲線,具有可學習的區域性 B 樣條基函式的係數(見圖 2.2 右側)。現在就有了一個 KAN 的原型,其計算圖完全由方程(2.1)指定,並在圖 0.1(b)中說明(輸入維度 n = 2),它看起來像是一個兩層的神經網路啟用函式放置在邊而不是節點上(節點上執行簡單的求和),中間層寬度為 2n + 1。

圖片

如前所述,在實踐中,這樣的網路被認為過於簡單,無法用平滑樣條來任意精確地逼近任何函式。因此,研究者將 KAN 泛化為更寬和更深的網路。由於 Kolmogorov-Arnold 表示對應於兩層 KAN,因此如何使 KAN 更深尚不清楚。

突破點在於研究者注意到了 MLP 和 KAN 之間的類比。在 MLP 中,一旦定義了一個層(由線性變換和非線性組成),就可以堆疊更多的層來使網路更深。要構建深度 KAN,首先應該回答:「什麼是一個 KAN 層?」研究者發現,一個具有 n_in 維度輸入和 n_out 維度輸出的 KAN 層可以被定義為一個一維函式矩陣。

圖片

其中函式圖片具有可訓練引數,如下所述。在 Kolmogorov-Arnold 定理中,內層函式形成一個 KAN 層,其中 n_in = n 和 n_out = 2n+1,外層函式形成一個 KAN 層,其中 n_in = 2n + 1 和 n_out = 1。因此,方程(2.1)中的 Kolmogorov-Arnold 表示僅僅是兩個 KAN 層的組合。現在,擁有更深的 Kolmogorov-Arnold 表示意味著:只需堆疊更多的 KAN 層!

更進一步理解需要引入一些符號,你可以參考圖 2.2(左)來獲得具體示例和直觀理解。KAN 的形狀由一個整數陣列表示:

圖片

其中,n_i 是計算圖第 i 層的節點數。這裡用 (l, i) 表示第 l 層的第 i 個神經元,用 x_l,i 表示 (l, i) 神經元的啟用值。在第 l 層和第 l + 1 層之間,有 n_l*n_l+1 個啟用函式:連線 (l, j) 和 (l + 1, i) 的啟用函式表示為

圖片

函式 ϕ_l,i,j 的預啟用值簡單表示為 x_l,i;ϕ_l,i,j 的後啟用值為圖片≡ ϕ_l,i,j (x_l,i)。第 (l + 1, j) 神經元的啟用值是所有傳入後啟用值的總和:

圖片

用矩陣形式表示如下:

圖片

其中,Φ_l 是對應於第 l 層 KAN 層的函式矩陣。一個通用的 KAN 網路是 L 層的組合:給定一個輸入向量 x_0 ∈ R^n0,KAN 的輸出是

圖片

上述方程也可以寫成類似於方程(2.1)的形勢,假設輸出維度 n_L = 1,並定義 f (x) ≡ KAN (x):

圖片

這樣寫起來相當繁瑣。相比之下,研究者對 KAN 層的抽象及其視覺化更加簡潔直觀。原始的 Kolmogorov-Arnold 表示公式(2.1)對應於形狀為 [n, 2n + 1, 1] 的 2 層 KAN。請注意,所有操作都是可微分的,因此可以用反向傳播來訓練 KAN。作為比較,MLP 可以寫成仿射變換 W 和非線性 σ 的交織:

圖片

很明顯,MLP 將線性變換和非線性分別處理為 W 和 σ,而 KAN 則將它們一併處理為 Φ。在圖 0.1 (c) 和 (d) 中,研究者展示了三層 MLP 和三層 KAN,以說明它們之間的區別。

KAN 的準確性

在論文中,作者還證明了在各種任務(迴歸和偏微分方程求解)中,KAN 在表示函式方面比 MLP 更有效。而且他們還表明 KAN 可以自然地在持續學習中發揮作用,而不會出現災難性遺忘。

toy 資料集

我們在圖 3.1 中繪製了 KAN 和 MLP 的測試 RMSE 作為引數數量的函式,展示了 KAN 比 MLP 有更好的縮放曲線,特別是在高維示例中。為了比較,作者繪製了根據他們的 KAN 理論預測的線條,為紅色虛線(α = k + 1 = 4),以及根據 Sharma & Kaplan [17] 預測的線條,為黑色虛線(α = (k + 1)/d = 4/d)。KAN 幾乎可以填滿更陡峭的紅色線條,而 MLP 甚至難以以更慢的黑色線條的速度收斂,並迅速達到平臺期。作者還注意到,對於最後一個示例,2 層 KAN 的表現遠不如 3 層 KAN(形狀為 [4, 2, 2, 1])。這突出了更深的 KAN 有更強的表達能力,對於 MLP 也是如此:更深的 MLP 比更淺的 MLP 具有更強的表達能力。

圖片

特殊函式

我們在這部分展示了以下兩點:

(1) 找到特殊函式的(近似)緊湊的 KA 表示是可能的,這從 Kolmogorov-Arnold 表示的角度揭示了特殊函式的新數學屬性。

(2) 在表示特殊函式方面,KAN 比 MLP 更有效、更準確。

對於每個資料集和每個模型族(KAN 或 MLP),作者在引數數量和 RMSE 平面上繪製了帕累託邊界,如圖 3.2 所示。

圖片

KAN 的效能一致優於 MLP,即在相同數量的引數下,KAN 能夠實現比 MLP 更低的訓練 / 測試損失。此外,作者在表 2 中報告了他們自動發現的特殊函式的 KAN(出人意料地緊湊)的形狀。一方面,從數學上解釋這些緊湊表示的意義是有趣的。另一方面,這些緊湊表示意味著有可能將一個高維查詢表分解為幾個一維查詢表,這可以潛在地節省大量記憶體,而在推理時執行一些加法運算的開銷(幾乎可以忽略不計)。

圖片

Feynman 資料集

上上節的設定是我們清楚地知道「真實」的 KAN 形狀。上節的設定是我們顯然不知道「真實」的 KAN 形狀。這一部分研究了一箇中間的設定:給定資料集的結構,我們可能手工構建 KAN,但我們不確定它們是否最優。

對於每種超引數組合,作者嘗試了 3 個隨機種子。對於每個資料集(方程)和每種方法,他們在表 3 中報告了最佳模型(最小 KAN 形狀或最低測試損失)在隨機種子和深度上的結果。

圖片

他們發現 MLP 和 KAN 平均表現相當。對於每個資料集和每個模型族(KAN 或 MLP),作者在引數數量和 RMSE 損失構成的平面上繪製了帕累託邊界,如圖 D.1 所示。他們推測費曼資料集太簡單,無法讓 KAN 做出進一步改進,在這個意義上,變數依賴通常是平滑的或單調的,這與特殊函式的複雜性形成對比,特殊函式經常表現出振盪行為。

圖片

解決偏微分方程

作者使用相同的超引數,比較了 KAN 與 MLP 架構。他們測量了 L^2 norm 和能量(H^1)norm 的誤差,並觀察到 KAN 在使用更小的網路和更少的引數的情況下,實現了更好的 scaling law 和更小的誤差,見圖 3.3。因此,他們推測 KAN 可能有潛力作為偏微分方程(PDE)模型約簡的良好神經網路表示。

圖片

持續學習

作者展示了 KAN 具有區域性可塑性,並且可以透過利用 spline 的區域性性來避免災難性遺忘。這個想法很簡單:由於 spline 基是區域性的,一個樣本只會影響少數附近的 spline 係數,同時保持遠處的係數不變(這是我們希望的,因為遠處的區域可能已經儲存了我們想要保留的資訊)。相比之下,由於 MLP 通常使用全域性啟用函式,例如 ReLU/Tanh/SiLU 等,任何區域性變化都可能無法控制地傳播到遠處的區域,破壞那裡儲存的資訊。

作者使用了一個簡單的示例來驗證這種直覺。一個一維迴歸任務由 5 個高斯峰組成。每個峰周圍的資料是順序呈現的(而不是一次性全部呈現),如圖 3.4 頂部行所示,這些資料分別呈現給 KAN 和 MLP。KAN 和 MLP 在每個訓練階段後的預測結果分別顯示在中間和底部行。正如預期的那樣,KAN 只重構當前階段存在資料的區域,而保持之前的區域不變。相比之下,MLP 在看到新的資料樣本後,會重構整個區域,導致災難性遺忘。

圖片

KAN 是可解釋的

在文章第 4 章,作者展示了 KAN 由於在第 2.5 節中開發的技術而具有可解釋性和互動性。他們想要測試 KAN 的應用,不僅在合成任務(第 4.1 和 4.2 節)上,而且也在現實生活科學研究中。他們展示了 KANs 能夠(重新)發現結理論中的複雜關係(第 4.3 節)和凝聚態物理學中的相變邊界(第 4.4 節)。由於其準確性和可解釋性,KAN 有潛力成為 AI + Science 的基礎模型。

討論

在論文中,作者從數學基礎、演算法和應用的角度討論了 KAN 的侷限性和未來的發展方向。

數學方面:儘管作者已經對 KAN 進行了初步的數學分析(定理 2.1),但對它們的數學理解仍然非常有限。Kolmogorov-Arnold 表示定理在數學上已經被徹底研究,但該定理對應的 KAN 形狀為 [n, 2n + 1, 1],這是 KAN 的一個非常受限的子類。在更深的 KAN 上的實證成功是否意味著數學上的某些基本原理?一個吸引人的廣義 Kolmogorov-Arnold 定理可以定義超出兩層組合的「更深」的 Kolmogorov-Arnold 表示,並可能將啟用函式的平滑度與深度相關聯。假設存在一些函式,它們不能在原始的(深度為 2)Kolmogorov-Arnold 表示中平滑表示,但可能在深度為 3 或更深時平滑表示。我們能否使用這種「Kolmogorov-Arnold 深度」的概念來表徵函式類?

演算法方面,他們討論了以下幾點:

準確性。在架構設計和訓練中存在多種選擇,尚未完全研究,因此可能存在進一步提高準確性的替代方案。例如,spline 啟用函式可能被徑向基函式或其他區域性核函式所取代。可以使用自適應網格策略。

效率。KAN 執行緩慢的主要原因之一是因為不同的啟用函式不能利用批計算(大量資料透過同一個函式)。實際上,我們可以透過將啟用函式分組為多組(「多頭」),在 MLP(所有啟用函式都相同)和 KAN(所有啟用函式都不同)之間進行插值,其中組內成員共享相同的啟用函式

KAN 和 MLP 的混合。與 MLP 相比,KAN 有兩個主要區別:

  • (i) 啟用函式位於邊而不是節點上;

  • (ii) 啟用函式是可學習的而不是固定的。

哪種改變更能解釋 KAN 的優勢?作者在附錄 B 中展示了他們的初步結果,他們研究了一個模型,該模型具有 (ii),即啟用函式是可學習的(像 KAN 一樣),但沒有 (i),即啟用函式位於節點上(像 MLP 一樣)。此外,人們還可以構建另一個模型,其啟用函式是固定的(像 MLP 一樣),但位於邊上(像 KAN 一樣)。

自適應性。由於 spline 基函式的固有區域性性,我們可以在 KAN 的設計和訓練中引入自適應性,以提高準確性和效率:參見 [93, 94] 中的多級訓練思想,如多重網格方法,或 [95] 中的領域依賴基函式,如多尺度方法。

應用方面:作者已經提出了一些初步證據,表明 KAN 在科學相關任務中比 MLP 更有效,例如擬合物理方程和解決 PDE。他們預計 KAN 在解決 Navier-Stokes 方程、密度泛函理論或任何可以表述為迴歸或 PDE 解決的其他任務方面也可能很有前景。他們還希望將 KAN 應用於與機器學習相關的任務,這將需要將 KAN 整合到當前的架構中,例如 transformer—— 人們可以提出「kansformers」,在 transformer 中用 KAN 替換 MLP。

KAN 作為 AI + Science 的語言模型:大型語言模型之所以具有變革性,是因為它們對任何能夠使用自然語言的人來說都是有用的。科學的語言是函式。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。

更多細節,請閱讀原論文。

相關文章