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

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

KA表示定理

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

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

假設(shè)有如下的七次多項式：

• 方程1

希爾伯特問，其解x，被認(rèn)為是三個變量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ù)。這里的關(guān)鍵點是，唯一的真正多變量函數(shù)是加法，因為其他所有函數(shù)都可以使用一元函數(shù)和求和來表示。

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

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

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

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

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

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

• 圖1：思考和構(gòu)建KAN，摘自論文

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

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

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

構(gòu)建KAN層

已經(jīng)提到，代表原始KA表示定理的兩層網(wǎng)絡(luò)太簡單，無法任意精確地逼近任何函數(shù)。我們?nèi)绾巫孠AN更寬、更深呢？

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

首先，可以將KA表示以矩陣形式表達(dá)：

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

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

• 方程4：設(shè)置矩陣的維度

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

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

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

作者將n_i表示為KAN中第i層的節(jié)點數(shù)，第l層的第i個神經(jīng)元由(l, i)表示，其中該神經(jīng)元的激活由x_{l, i}給出。我們認(rèn)為激活函數(shù)是位于網(wǎng)絡(luò)圖邊緣的可學(xué)習(xí)函數(shù)，節(jié)點表示求和運(yùn)算。因此，在第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中定義的內(nèi)部函數(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個神經(jīng)元。

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

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

我們可以再次檢查圖2，

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

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

• 激活前后的符號

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

利用這些，我們可以定義激活的可學(xué)習(xí)變換矩陣：

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

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

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

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

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

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

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

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

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

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

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

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

由于對于KAN來說，現(xiàn)在一切都?xì)w結(jié)為激活函數(shù)，作者定義了如何構(gòu)建這些函數(shù)。

可學(xué)習(xí)的激活函數(shù)

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

• 方程11：將基函數(shù)和樣條函數(shù)作為線性組合的可學(xué)習(xí)激活函數(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樣條的權(quán)重，即c_i，是可訓(xùn)練的，作者認(rèn)為方程11中的因子w的唯一用途是更好地控制激活函數(shù)的總體大小。

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

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

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

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

結(jié)語

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

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

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

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

• 兩個相對論速度的加法

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

• 使用KAN學(xué)習(xí)乘法，將其視為線性和二次函數(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可能會引導(dǎo)我們發(fā)現(xiàn)這個技巧。我們可以像這個例子中那樣使用KAN來發(fā)現(xiàn)/重新發(fā)現(xiàn)一些基本的物理定律嗎？

加載和運(yùn)行KAN

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

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

!pip install pykan
from kan import KAN, create_dataset

import torch

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

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ù)和標(biāo)簽的這些直方圖。

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

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

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

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

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ù)的最佳建議。

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