第一屆“歐洲基礎物理人工智能會議”(EuCAIFCon)于2024年4月30日至5月3日在阿姆斯特丹舉行。會議上許多討論都集中在基礎模型上，探討是否可以利用人工智能發(fā)現(xiàn)潛在的新物理定律。令人驚訝的是，最近一篇名為KAN的論文在arXiv上發(fā)布，探討了利用神經網(wǎng)絡發(fā)現(xiàn)或重新發(fā)現(xiàn)物理和數(shù)學模型的可能性。

接下來我們將深入探討構建KANs的數(shù)學公式和概念。以了解為什么這個網(wǎng)絡會引起如此大的轟動！

KA表示定理

KAN或Kolmogorov-Arnold網(wǎng)絡基于著名數(shù)學家Kolmogorov和Arnold的表示定理。首先，讓我們先退幾步來理解這個定理。

KA定理是由這些數(shù)學家提出的，用于解決希爾伯特的第13個問題：是否所有七次方程的解都可以用兩個變量的代數(shù)函數(shù)來表達？這是什么意思呢？

假設有如下的七次多項式：

• 方程1

希爾伯特問，其解x，被認為是三個變量a，b，c的函數(shù)，是否可以表示為有限數(shù)目的二元函數(shù)的組合：

KA表示定理指出，對于任何連續(xù)函數(shù)

存在一元連續(xù)函數(shù)g_q, ψ_{p,q}，使得：

• 方程2：KA表示定理

這意味著（2d+1)(d+1)個一元函數(shù)g_q, ψ_{p,q}足以精確表示一個d變量函數(shù)。這里的關鍵點是，唯一的真正多變量函數(shù)是加法，因為其他所有函數(shù)都可以使用一元函數(shù)和求和來表示。

我們需要記住，g_q, ψ_{p,q}都是一元函數(shù)。因此，任何多變量的連續(xù)函數(shù)都可以表示為一元函數(shù)的組合。哇！我之前不知道。太酷了！

思考KA定理和多層感知器（MLPs）

論文的作者指出，這些一維函數(shù)（之前定義的）可能是不光滑甚至是分形的，因此在實踐中可能無法學習。為了構建KANs，作者們超越了KA表示定理的定義，但首先我們來思考MLPs。

多層感知器（MLPs）基于通用逼近定理，該定理指出任何連續(xù)函數(shù)f:[0, 1]? → [0, 1]都可以通過至少包含一個隱藏層的神經網(wǎng)絡（權重、偏差和非線性激活函數(shù)）來任意精確地逼近。在反向傳播過程中，網(wǎng)絡學習優(yōu)化權重和偏差以充當函數(shù)逼近器，而激活函數(shù)保持不變。

現(xiàn)在，我們能否根據(jù)上述KA表示定理構建一個基于神經網(wǎng)絡的架構？

如果考慮一個監(jiān)督學習問題，給定{x_i, y_i}對，我們想找到一個函數(shù)使得y_i ≈ f(x_i)，那么KA表示定理告訴我們，需要找到方程2中的一元函數(shù)（g_q, ψ_{p,q）。

• 圖1：思考和構建KAN，摘自論文

在這里，作者們認為，由于我們只需要學習一元函數(shù)，我們可以將每個一維函數(shù)參數(shù)化為B樣條曲線（見下文），其局部B樣條基函數(shù)的系數(shù)是可學習的。這導致了KAN的原型，并在圖1（b）中進行了說明，輸入維度n=2表現(xiàn)為一個兩層神經網(wǎng)絡，激活函數(shù)放置在邊上而不是節(jié)點上（在節(jié)點上進行簡單的求和），中間層的寬度為2n+1。網(wǎng)絡的構造，即激活數(shù)、節(jié)點數(shù)等，將很快清晰。鑒于方程2中的定義，原始KA表示定理可以被視為深度為2，每層包含（2d+1）項。

B樣條：我們可以將B樣條函數(shù)理解為由多個控制點控制的靈活帶組成，用于創(chuàng)建平滑曲線。從數(shù)學角度更嚴格的定義是，p+1階的B樣條是由變量t中p階的分段多項式函數(shù)B_{i, p}組成的集合。分段多項式相接處的t值被稱為結點。

再次說明，B樣條由分段多項式（基函數(shù)）構建，其階數(shù)比其基多項式的度數(shù)多一。例如，二次B樣條的多項式度數(shù)為2，階數(shù)為3。這正是KAN論文中展示和使用的內容。

構建KAN層

已經提到，代表原始KA表示定理的兩層網(wǎng)絡太簡單，無法任意精確地逼近任何函數(shù)。我們如何讓KAN更寬、更深呢？

這里，作者提出了KAN和MLPs之間的絕佳類比以便更深入探討。首先，我們需要了解什么是KAN層以及如何將它們堆疊起來構建深度神經網(wǎng)絡。

首先，可以將KA表示以矩陣形式表達：

• 方程3：以矩陣形式思考方程2

具有n_{in}維輸入和n_{out}維輸出的KAN層可以定義為一維函數(shù)的矩陣：

• 方程4：設置矩陣的維度

在Kolmogov-Arnold定理（方程2）中，內部函數(shù)構成一個n_{in}=n和n_{out}=2n+1的KAN層，外部函數(shù)構成一個n_{in}=2n+1和n_{out}=1的KAN層。此時，我們可以將KA表示視為兩個KAN層的組合。讓我們嘗試習慣于堆疊更多KAN層時的符號。

我們可以使用作者提供的示例圖來討論網(wǎng)絡維度等內容：

• 圖2：思考KAN層：摘自論文

作者將n_i表示為KAN中第i層的節(jié)點數(shù)，第l層的第i個神經元由(l, i)表示，其中該神經元的激活由x_{l, i}給出。我們認為激活函數(shù)是位于網(wǎng)絡圖邊緣的可學習函數(shù)，節(jié)點表示求和運算。因此，在第1層（0層）和第2層（1層）之間，我們看到有10個激活函數(shù)，分別由?_{0,1,1}, ?_{0,1,2}等表示。激活函數(shù)的數(shù)量由0層和1層的節(jié)點數(shù)決定。

這里我們可以清楚地看到MLPs與KANs的區(qū)別。KANs的激活函數(shù)位于邊緣，而MLPs的激活函數(shù)位于節(jié)點上。

在第0層，我們有兩個節(jié)點x_{0,1}, x_{0,2}，在第一層，有5個，所以激活函數(shù)的數(shù)量將是n_l × n_{l+1}。

n_l和n_{l+1}是根據(jù)方程4中定義的內部函數(shù)的輸入和輸出維度確定的。因此我們從兩個輸入n_{in}=2開始，所以n_{out}必須為2n+1=5。這反過來決定了隱藏層中激活函數(shù)的數(shù)量。

如果我們繼續(xù)以節(jié)點數(shù)n_1=5和n_2=1(n_{out})，在該層有5個激活函數(shù)是合理的。這將是外部函數(shù)。重申一遍，KA表示由兩個KAN層組成。

KAN層的矩陣形式

現(xiàn)在我們可以開始編寫激活函數(shù)。讓我們看看：連接第l層和l+1層兩個節(jié)點的激活函數(shù)由?_{l, j, i}表示，其中{j, i}分別代表那兩層中的第j和第i個神經元。

因此，在第l層和l+1層之間的可學習激活函數(shù)：

• 方程5：KAN邊緣的可學習激活函數(shù)

我們可以再次檢查圖2，

• 通過與圖2比較來理解節(jié)點數(shù)量。

我們將?_{l, j, i}的輸入前激活表示為x_{l, i}；然后在激活后有：

• 激活前后的符號

第(l+1, j)神經元的激活值簡單地是所有傳入激活后的求和。

利用這些，我們可以定義激活的可學習變換矩陣：

• 方程6：KAN不同層的完整可學習激活函數(shù)集

利用這個我們也可以編寫變換規(guī)則：

• 方程7：完整變換規(guī)則：給定激活矩陣的預激活和激活后。

再次檢查我們的理解，因此與圖2比較：

• 方程8：鑒于圖2，檢查我們是否能理解變換矩陣的維度！

確實有5個輸出x_{1,1}, x_{1,2}, x_{1,3}, x_{1,4}, x_{1,5}。

一旦我們準備好了變換矩陣，我們可以簡單地將它們組合（堆疊層）以便更深入地探討，如下所示：

• 方程9：通過堆疊若干KAN層來組成一個KAN

此時我們也可以認識到，所有運算都是可微分的（假設1D函數(shù)也是），梯度可以通過網(wǎng)絡流動，即我們可以進行反向傳播！

我們還可以將KAN與MLP層進行比較，在MLP層中有權重矩陣（線性變換）和激活函數(shù)（非線性）分開：

• 方程10：將KAN與MLP比較；權重（線性）和激活（非線性）。

權重矩陣中的值會更新，但一旦定義，MLP中的激活函數(shù)就是固定的。這是KAN與MLP層之間的關鍵區(qū)別，我們的激活函數(shù)是可學習的。

由于對于KAN來說，現(xiàn)在一切都歸結為激活函數(shù)，作者定義了如何構建這些函數(shù)。

可學習的激活函數(shù)

為了構建激活函數(shù)?(x)，作者提議使用基函數(shù)(b(?))和樣條函數(shù)，并將它們組合如下：

• 方程11：將基函數(shù)和樣條函數(shù)作為線性組合的可學習激活函數(shù)。

作者選擇的基函數(shù)為SiLU：

• 方程12：基函數(shù)b(?)的選擇

對于樣條函數(shù)，它是B樣條的線性疊加：

• 方程13：樣條函數(shù)作為B樣條的線性組合。

如果回顧第二張圖，我們看到它是k=3的B樣條的線性組合，即階數(shù)為3，所以B樣條中的多項式的度數(shù)為2。像這樣定義樣條的一個優(yōu)勢是，通過增加曲線的數(shù)量可以使其任意平滑。這也在圖2中顯示，作者增加了我們連接不同多項式的區(qū)間數(shù)量，從7增加到12。

B樣條的權重，即c_i，是可訓練的，作者認為方程11中的因子w的唯一用途是更好地控制激活函數(shù)的總體大小。

MLP與KAN的參數(shù)數(shù)量

作者還討論了通常情況下，KAN比MLP運行較慢。為了理解這一點，我們可以簡單地通過假設網(wǎng)絡深度為L，每層都有相同數(shù)量的節(jié)點n_i=N，每個樣條的階數(shù)為k（通常為3）在G個區(qū)間上，來計算參數(shù)數(shù)量：

• KAN與MLP的參數(shù)數(shù)量對比

然而，KAN所需的寬度即N比MLP中的小，且KAN是可解釋的，我們將看到作者提出的一個例子。作者強調訓練KAN比MLP慢的另一個原因是，由于激活函數(shù)是可學習的，不可能利用“批處理計算”，即大量數(shù)據(jù)通過相同的函數(shù)。這在MLP中不是問題，因為在訓練和測試時間內激活是固定的。

結語

這篇論文中還有很多復雜的細節(jié)，但對我個人而言最突出的是KAN的可解釋性。作者展示了KAN可以“發(fā)現(xiàn)”從簡單的除法法則到結理論中的非平凡關系。這可能會進一步推動KAN在AI與科學的基礎模型中的應用。

作者建議，KAN可能比符號回歸更具“吸引力”；作者給出了一個例子，通過KAN學習非常曲折的20階貝塞爾函數(shù)(J_{20}(x))，這通過符號回歸在沒有任何關于那個特殊函數(shù)（本例中為貝塞爾函數(shù)）的先驗知識的情況下是不可能的。

通過KAN的“發(fā)現(xiàn)”示例

在作者提出的許多示例中，我喜歡一個相對簡單但引人入勝的“自動可發(fā)現(xiàn)”特性的KAN。我們總是喜歡這類物理示例；比如，我們從相對論速度加法公式開始：

• 兩個相對論速度的加法

我們可以考慮KAN的深度，將每一層KAN視為發(fā)現(xiàn)一個數(shù)學運算；因此，看看上面的公式，首先我們考慮乘法；作者顯示，學習的激活函數(shù)將是線性和二次的，所以：

• 使用KAN學習乘法，將其視為線性和二次函數(shù)的組合

對(1+v_1 * v_2)的求逆將使用一層，而(v_1 + v_2)與(1/(1+v_1 * v_2))的乘法將需要另外兩層；總共5層。

但研究人員發(fā)現(xiàn)，“自動發(fā)現(xiàn)”的KAN只有2層深，這可以通過速率技巧來解釋。

在相對論中，我們可以通過速率技巧簡化變換規(guī)則；可以定義速率為：

我們可以使用雙曲正切加法公式：

使用這個，可以看到：

現(xiàn)在只有兩層完全有意義。如果我們不知道速率技巧，試圖理解這個2層的自動發(fā)現(xiàn)KAN可能會引導我們發(fā)現(xiàn)這個技巧。我們可以像這個例子中那樣使用KAN來發(fā)現(xiàn)/重新發(fā)現(xiàn)一些基本的物理定律嗎？

加載和運行KAN

讓我們仔細看看上面的例子，作者在論文中也提到了這個例子，并且可以在GitHub上找到。我們將在Colab上運行它。

我們將開始本地安裝并加載必要的庫。

!pip install pykan
from kan import KAN, create_dataset

import torch

我們創(chuàng)建了包含兩組速度的數(shù)據(jù)集，這些速度在訓練和測試集中分開，并且有相應的回歸值，即標簽(f(v_1, v_2))。為了保險起見，我們可以使用matplotlib檢查輸入速度及其相應加成相對速度的分布。

f = lambda x: (x[:,[0]]+x[:,[1]])/(1+x[:,[0]]*x[:,[1]]) # dataset creation where x[:, [0]] represents v1, x[:, [1]]: v2dataset = create_dataset(f, n_var=2, ranges=[-0.9,0.9])#plot the distribution import matplotlib.pyplot as plt
### check train and test input distribution
fig = plt.figure(figsize=(10, 5))fig.add_subplot(131)plt.hist(dataset['train_input'][:, 0], bins=20, alpha=0.7,          label=r'$v_1$-train', color='orange')plt.hist(dataset['train_input'][:, 1], bins=20, alpha=0.7,          label=r'$v_2$-train', histtype='step')plt.legend(fontsize=12)fig.add_subplot(132)plt.hist(dataset['test_input'][:, 0], bins=20, alpha=0.7,          label=r'$v_1$-test', color='orange')plt.hist(dataset['test_input'][:, 1], bins=20, alpha=0.7,          label=r'$v_2$-test', histtype='step')plt.legend(fontsize=12)fig.add_subplot(133)plt.hist(dataset['train_label'].numpy(), bins=20, alpha=0.7,          label=r'$\frac{v_1+v_2}{1+v_1\, v_2}$-train', color='orange')plt.hist(dataset['test_label'].numpy(), bins=20, alpha=0.7,          label=r'$\frac{v_1+v_2}{1+v_1\, v_2}$-test', histtype='step')plt.legend(fontsize=12)plt.tight_layout()plt.show()

我們得到了數(shù)據(jù)和標簽的這些直方圖。

• 圖3：在[-1, 1]之間隨機分布的速度及其相應的（相對論性）加和值。

• 圖4：具有學習到的激活函數(shù)的兩層KAN

第一層的激活函數(shù)已經看起來像arctanh，第二層的激活函數(shù)看起來像tanh。這真的很酷！

嘗試從模型中獲取第一層激活函數(shù)的符號函數(shù)表示的建議，揭示了我們實際看到的內容：

model.suggest_symbolic(0, 1, 0)>>> function , r2arctanh , 0.9986623525619507tan , 0.9961022138595581arcsin , 0.968244731426239

第二層也一樣：

model.suggest_symbolic(1, 0, 0)
>>> function , r2tanh , 0.9995558857917786arctan , 0.995667040348053gaussian , 0.9793974757194519

我們確實得到tanh作為符號函數(shù)的最佳建議。

我非常興奮地看到基礎AI、物理學和數(shù)學的研究者將如何合并KAN和MLP，或者修改KAN使其更快、更好、可能更具解釋性（如果這是可能的）；此外，發(fā)現(xiàn)/重新發(fā)現(xiàn)物理定律的可能性，可能在天體物理學、宇宙學領域，應該是使用KAN需要探索的另一個方面。