4月30日，KAN橫空出世，很多人預(yù)言這會敲響MLP的喪鐘。

沒想到，子彈還沒飛4個月，核心團隊又推出了KAN 2.0，瞄準(zhǔn)AI+Science領(lǐng)域，進一步挖掘了KAN的潛力。

這篇論文更雄心勃勃的地方在于，作者希望通過一種框架來彌合AI世界的連接主義（connectionism）和科學(xué)世界的符號主義（symbolism）之間的不相容性。

通過提出pykan等工具，作者還展現(xiàn)了KAN發(fā)現(xiàn)各種物理定律的能力，包括守恒量、拉格朗日量、隱藏對稱性和本構(gòu)方程等等。

論文地址：https://arxiv.org/abs/2408.10205

這次KAN 2.0依舊出自初代架構(gòu)原班人馬之手。

深度學(xué)習(xí)變天了，MLP成過去式？

我們先簡要回顧一下，今年4月首次提出的KAN究竟在哪些方面改進了MLP。

MLP（multi-layer perceptron）又被稱為全連接神經(jīng)網(wǎng)絡(luò)，是當(dāng)今幾乎所有深度學(xué)習(xí)模型的基礎(chǔ)構(gòu)建塊，它的出世甚至可以追溯到第一波人工智能浪潮方興未艾的1958年。

論文地址：https://www.ling.upenn.edu/courses/cogs501/Rosenblatt1958.pdf

KAN的論文中都表示，MLP的重要性怎么強調(diào)都不為過，因為這是神經(jīng)網(wǎng)絡(luò)中用于逼近非線性函數(shù)的默認(rèn)模型，其對函數(shù)表達能力的底層邏輯由「通用逼近定理」保證。

Transformer和其他架構(gòu)中常見的FFN本質(zhì)上就是一個MLP。但由于網(wǎng)絡(luò)稠密、參數(shù)量大，MLP往往占據(jù)了模型中幾乎所有的非編碼層參數(shù)。

而且相比注意力層，在沒有后期分析工具時，MLP中的大量參數(shù)也缺乏可解釋性。

受到Kolmogorov-Arnold表示定理的啟發(fā)，KAN打破了對通用逼近定理的遵循。

雖然底層邏輯變了，但是他們做出的修改相當(dāng)簡潔且直觀：

- 將激活函數(shù)放在網(wǎng)絡(luò)邊緣而非節(jié)點處

- 給激活函數(shù)賦予可學(xué)習(xí)參數(shù)，而非固定的函數(shù)

KAN中沒有任何線性權(quán)重，網(wǎng)絡(luò)中的每個權(quán)重都變成了B-spline型單變量函數(shù)的可學(xué)習(xí)參數(shù)。

這種看似簡單的改變讓KAN獲得了擬合準(zhǔn)確性和可解釋性方面的優(yōu)勢。今年4月的第一篇論文中，作者們就發(fā)現(xiàn)KAN在數(shù)學(xué)和物理定律方面的潛力。

下面這個動圖展示了簡單的3層KAN網(wǎng)絡(luò)擬合一個復(fù)雜函數(shù)的訓(xùn)練過程，相當(dāng)簡潔清楚。

此外，KAN也能從根本上很好地解決MLP中普遍存在的「災(zāi)難性遺忘」問題。

以上這些優(yōu)勢，都奠定了KAN作為「科學(xué)家合作助手」的基本能力。

KAN2.0問世，一統(tǒng)AI+科學(xué)

雖然第一版的KAN網(wǎng)絡(luò)本身有很多適合科學(xué)研究的優(yōu)點，但深度學(xué)習(xí)和物理、化學(xué)、生物學(xué)領(lǐng)域依舊有完全不同的「語言」，這構(gòu)成了AI4Science最大的障礙之一。

因此擴展后的KAN 2.0的終極目標(biāo)只有一個——使KAN能輕松應(yīng)用于「好奇心驅(qū)動的科學(xué)」。研究人員既能將輔助變量、模塊化結(jié)構(gòu)、符號公式等科學(xué)知識集成到KAN中，也能從KAN的可解釋性分析中得到觀察和見解。

所謂「好奇心驅(qū)動的科學(xué)」，根據(jù)論文的解釋，是過程更具有探索性、提供更基礎(chǔ)層面新發(fā)現(xiàn)和新知識的研究，比如天體運動背后的物理原理，而非AlphaFold這類應(yīng)用驅(qū)動的科學(xué)研究。

科學(xué)與KAN的協(xié)同

具體來說，科學(xué)解釋有不同的層次，從最簡單粗略到最精細、最難發(fā)現(xiàn)、最具因果性，可以有如下幾個分類：

- 重要特征：例如，y完全由x₁和x₂決定，其他因素并不重要；即存在一個函數(shù)f使得y=f(x₁, x₂)

- 模塊化結(jié)構(gòu)：例如，存在函數(shù)g和h是的y=g(x₁)+h(x₂)

- 符號公式：例如，y=sin(x₁)+exp(x₂)

MultKAN

在原始KAN網(wǎng)絡(luò)的基礎(chǔ)上，這篇最新的論文引入了一種稱為MultKAN的新模型，其核心改進是引入額外的乘法層進行增強。

KAN所依據(jù)的Kolmogorov-Arnold表示定理提出，任何連續(xù)高維函數(shù)都可以分解為單變量連續(xù)函數(shù)和加法的有限組合：

這意味著加法是唯一真正的多元運算，，而其他多元運算（包括乘法）都可以表示為與單變量函數(shù)組合的加法。因此，原來的KAN中僅包含加法運算。

然而，考慮到乘法在科學(xué)和日常生活中的普遍存在，MultKAN中明確包含乘法，能更清楚地揭示數(shù)據(jù)中的乘法結(jié)構(gòu)，以期增強可解釋性和表達能力。

如圖2所示，MultKAN和KAN相似，都包含標(biāo)準(zhǔn)KAN層，但區(qū)別在于插入了乘法節(jié)點，對輸入的子節(jié)點進行乘法運算后再進行恒等變換，用Python代碼可表示為：

其中⊙表示逐元素乘法。

根據(jù)上圖，整個MultKAN網(wǎng)絡(luò)進行的運算就可以寫作：

其中，??_L≡??_L°??_L。

經(jīng)過擴展后，論文將KAN和MultKAN視為同義詞，即默認(rèn)情況下的KAN都將允許乘法層的存在，除非有特殊說明。

GitHub倉庫中的KAN代碼已經(jīng)更新，可以通過pip快捷命令直接安裝使用。

倉庫地址：https://github.com/KindXiaoming/pykan

Science to KAN

在科學(xué)領(lǐng)域，領(lǐng)域知識至關(guān)重要，讓我們可以在數(shù)據(jù)稀少或不存在的情況下，也能有效工作。

因此，對KAN采用基于物理的方法會很有幫助：將可用的歸納偏置整合到KAN中，同時保持其從數(shù)據(jù)中發(fā)現(xiàn)新物理規(guī)律的靈活性。

文中作者探討了三種可以整合到KAN中的歸納偏置，從最粗略（最簡單/相關(guān)性）到最精細（最困難/因果關(guān)系）：重要特征、模塊化結(jié)構(gòu)和符號公式。

在KANs中添加重要特征

在回歸問題中，目標(biāo)是找到一個函數(shù)f，使得y=f(x₁, x₂, ···, x_n)。假設(shè)我們希望引入一個輔助輸入變量a=a(x₁, x₂, ..., x_n)，將函數(shù)轉(zhuǎn)化為y=f(x₁, ···, x_n, x_a)。

盡管輔助變量a不增加新的信息，但它可以提高神經(jīng)網(wǎng)絡(luò)的表達能力。這是因為網(wǎng)絡(luò)無需消耗資源來計算輔助變量。此外，計算可能變得更簡單，從而提升可解釋性。

這里，用戶可以使用augment_input方法向輸入添加輔助特征：

圖3顯示了包含輔助變量和不包含這些輔助變量的KAN：（a）由符號公式編譯而成的KAN，需要5條連接邊；（b）（c）包含輔助變量的KAN，僅需2或3條連接邊，損失分別為10??和10??。

為KAN構(gòu)建模塊化結(jié)構(gòu)

模塊化在自然界中非常普遍：比如，人類大腦皮層被劃分為幾個功能不同的模塊，每個模塊負(fù)責(zé)特定任務(wù)，如感知或決策。模塊化簡化了對神經(jīng)網(wǎng)絡(luò)的理解，因為它允許我們整體解釋神經(jīng)元群集，而不是單獨分析每個神經(jīng)元。

結(jié)構(gòu)模塊化的特點是連接群集，其中特征是群集內(nèi)的連接遠強于群集間的連接。為此，作者引入了module方法：保留群集內(nèi)的連接，同時去除群集間的連接。

模塊由用戶來指定，語法是：

具體而言，模塊化有兩種類型：可分性和對稱性。

可分性：如果說一個函數(shù)是可分的，那么它就可以表示為非重疊變量組的函數(shù)的和或積。

廣義對稱性：如果f(x₁, x₂, x₃, ···)=g(h(x₁, x₂), x₃, ···)，則這個函數(shù)在變量(x₁, x₂)上是對稱的。因為只要h(x₁, x₂)保持不變，即使x₁和x₂發(fā)生變化，f的值仍然保持不變。

將符號公式編譯成KAN

為了結(jié)合「符號方程」和「神經(jīng)網(wǎng)絡(luò)」這兩種方法的優(yōu)勢，作者提出了一個兩步程序：（1）將符號方程編譯成KAN，（2）使用數(shù)據(jù)微調(diào)這些KAN。

其中，第一步可以將已知的領(lǐng)域知識嵌入到KAN中，而第二步則專注于從數(shù)據(jù)中學(xué)習(xí)新的「物理」知識。

具體來說，作者首先提出了用于將符號公式編譯成KAN的kanpiler（KAN編譯器）。過程如圖5a所示：

1. 將符號公式解析為樹結(jié)構(gòu)，其中節(jié)點表示表達式，邊表示操作/函數(shù)；

2. 然后修改該樹以與KAN圖結(jié)構(gòu)對齊。修改包括通過虛擬邊將所有葉節(jié)點移動到輸入層，并添加虛擬子節(jié)點/節(jié)點以匹配KAN架構(gòu)。這些虛擬邊/節(jié)點/子節(jié)點僅執(zhí)行恒等變換；

3. 在第一層中組合變量，有效地將樹轉(zhuǎn)換為圖。

然而，通過寬度/深度擴展來增加表達能力kanpiler生成的KAN網(wǎng)絡(luò)是緊湊的，沒有冗余邊，這可能限制其表達能力并阻礙進一步的微調(diào)。

為了解決這個問題，作者又提出了expand_width和expand_depth方法來擴展網(wǎng)絡(luò)，使其變得更寬和更深，如圖5c所示。

KAN to Science

這一節(jié)同樣關(guān)注提取知識的三個層次，從最基本到最復(fù)雜：重要特征，模塊化結(jié)構(gòu)和符號公式。

識別重要特征

給定一個回歸模型f，有y≈f?(x₁,x₂,…,x_n) ，我們的目標(biāo)是為輸入變量分配重要性分?jǐn)?shù)。

論文提出，之前所使用的L1范數(shù)（圖6a）只考慮到了局部信息，因此得出的結(jié)果可能存在問題。

依據(jù)KAN網(wǎng)絡(luò)，作者提出了一種更有效的歸因分?jǐn)?shù)，能比L1范數(shù)更好反映變量的重要性，還可以根據(jù)這種歸因分?jǐn)?shù)對網(wǎng)絡(luò)進行剪枝。

 識別模塊化結(jié)構(gòu)

歸因分?jǐn)?shù)可以告訴我們哪些邊或節(jié)點更有價值，但它沒有揭示模塊化結(jié)構(gòu)，即重要的邊和節(jié)點如何連接。

神經(jīng)網(wǎng)絡(luò)中的模塊化結(jié)構(gòu)可以分為兩種：解剖模塊化（anatomical modularity）和功能模塊化（functional modularity）。

解剖模塊化是指，空間上彼此靠近的神經(jīng)元相比距離較遠的神經(jīng)元具有更強的連接趨勢。論文采用了之前研究提出的「神經(jīng)元交換」方法，在代碼中被稱為auto_swap，可以在保留網(wǎng)絡(luò)功能的同時縮短連接，有助于識別模塊。

圖7展示了兩個成功識別模塊的auto_swap任務(wù)：多任務(wù)匹配和分層多數(shù)投票。其中，KAN的模塊結(jié)構(gòu)相比MLP更加簡單且富有組織性。

但無論auto_swap結(jié)構(gòu)如何，網(wǎng)絡(luò)全局的模塊化結(jié)構(gòu)仍和整體功能仍不清楚，這就需要用到功能模塊化分析，通過輸入和輸出的前向和后向傳遞來收集有關(guān)信息。

圖8定義了三種類型的功能模塊化：可分性、一般可分性和一般對稱性。

 識別符號公式

符號公式信息量最大，因為可以直接、清楚地揭示函數(shù)中重要的特征和模塊結(jié)構(gòu)。圖9展示了與KAN進行交互協(xié)作進行符號回歸的3個技巧：

1.發(fā)現(xiàn)并利用模塊化結(jié)構(gòu)

2.稀疏初始化

3.假設(shè)檢驗

用KAN助力物理學(xué)研究

除了進行原理層面的說明，論文還講解了多個具體案例，如何將KAN融入到現(xiàn)實的科學(xué)研究中，比如發(fā)現(xiàn)新的物理概念和定律。

論文給出的案例包括守恒量、拉格朗日量、隱藏對稱性和本構(gòu)方程等。這里我們以最簡單的守恒量發(fā)現(xiàn)為例，看看KAN是如何工作的。

守恒量即時間變化過程中保持恒定的物理量，比如能量守恒定律告訴我們，孤立系統(tǒng)的總能量保持不變。

傳統(tǒng)上，科學(xué)家如果不借助計算工具，僅靠紙筆推導(dǎo)守恒量可能非常耗時，并且需要廣泛的領(lǐng)域知識。但機器學(xué)習(xí)方法可以將守恒量參數(shù)化，轉(zhuǎn)化為求解微分方程的問題。

此處所用的方法基本類似于作者Ziming Liu等人2022年發(fā)表的論文，但將其中的MLP網(wǎng)絡(luò)換成了KAN。

論文地址：https://pubmed.ncbi.nlm.nih.gov/36397460/

比如使用KAN可以發(fā)現(xiàn)二維諧振子(x, y, p_x, p_y)中具有3個守恒量：x軸方向的能量H₁、y軸方向的能量H₂和角動量H₃。

關(guān)于KAN的其他應(yīng)用，論文也描述了如何從實驗數(shù)據(jù)中推斷出拉格朗日量（圖11）。

或者，發(fā)現(xiàn)Schwarzschild黑洞中的隱藏對稱性（圖12）。

還有數(shù)據(jù)驅(qū)動的本構(gòu)定律發(fā)現(xiàn)（圖13）。本構(gòu)定律通過模擬材料對外力或變形的響應(yīng)，定義材料的行為和屬性，比如描述彈簧的胡克定律。

作者介紹

Ziming Liu（劉子鳴）

Ziming Liu目前是MIT和IAIFI的三年級博士生，由Max Tegmark教授指導(dǎo)。他是兩篇KAN論文的第一作者，可以說是這個架構(gòu)背后最主要的貢獻者。

他的研究興趣主要集中在AI與物理學(xué)（以及其他科學(xué)領(lǐng)域）的交叉區(qū)域：

1. Physics of AI：從物理學(xué)原理來理解AI，目標(biāo)是讓「AI像物理學(xué)一樣簡單」；

2. Physics for AI：受物理學(xué)啟發(fā)的AI，目標(biāo)是讓「AI像物理學(xué)一樣自然」；

3. AI for physics：利用AI增強物理學(xué)研究，目標(biāo)是讓「讓AI像物理學(xué)家一樣強大」。

為了實現(xiàn)利用AI和物理學(xué)共建更美好世界的最終目標(biāo)，Ziming Liu對包括發(fā)現(xiàn)物理定律、受物理啟發(fā)的生成模型、機器學(xué)習(xí)理論、機械解釋性等在內(nèi)的多個主題都有深厚的興趣。

并且，與凝聚態(tài)、高能物理、量子計算等領(lǐng)域的物理學(xué)家以及計算機科學(xué)家、生物學(xué)家、神經(jīng)科學(xué)家和氣候科學(xué)家等建立了緊密合作關(guān)系。

他多次在頂尖的物理期刊和AI會議上發(fā)表論文，并擔(dān)任IEEE、Physical Review、NeurIPS、ICLR等的審稿人。同時，還共同組織了NeurIPS 2021和ICML 2022的AI4Science workshop。

在攻讀博士學(xué)位之前，他在北京大學(xué)獲得了物理學(xué)學(xué)士學(xué)位，并曾在微軟亞洲研究院實習(xí)。

Pingchuan Ma（馬平川）

Pingchuan Ma目前是MIT CSAIL實驗室的博士生，由Wojciech Matusik教授指導(dǎo)。

他的研究方向涵蓋了「基于物理的智能」的整個流程：

1. 重建高效逼真的物理環(huán)境

2. 基于這些環(huán)境生成AI 智能體

3. 在物理世界中實現(xiàn)這些智能體

此前，他在南開大學(xué)獲得軟件工程專業(yè)學(xué)士學(xué)位，并在麻省理工學(xué)院獲得計算機科學(xué)碩士學(xué)位。

同時，他還在IBM、字節(jié)、商湯、港大等知名機構(gòu)從事過研究工作，有著豐富的經(jīng)驗。

Yixuan Wang

Yixuan Wang目前是加州理工學(xué)院，應(yīng)用及計算數(shù)學(xué)專業(yè)的博士生。

他的研究方向十分廣泛，包括數(shù)值分析、偏微分方程、應(yīng)用概率，以及AI for Science。

此前，他在北京大學(xué)獲得數(shù)學(xué)學(xué)士學(xué)位。

Wojciech Matusik

Wojciech Matusik是麻省理工學(xué)院計算機科學(xué)與人工智能實驗室（MIT CSAIL）的電氣工程與計算機科學(xué)教授，也是計算機圖形學(xué)小組的成員，負(fù)責(zé)帶領(lǐng)計算設(shè)計與制造團隊。

他的研究興趣包括計算機圖形學(xué)、計算設(shè)計與制造、計算機視覺、機器人學(xué)和人機交互。

他于2003年在MIT獲得計算機圖形學(xué)博士學(xué)位，2001年在MIT獲得電氣工程與計算機科學(xué)碩士學(xué)位，1997年在加州大學(xué)伯克利分校獲得電氣工程與計算機科學(xué)學(xué)士學(xué)位。

并他曾在三菱電機研究實驗室、Adobe和迪士尼蘇黎世研究所工作。

2004年，他被「麻省理工科技評論」評為全球100位頂尖青年創(chuàng)新者之一。2009年，獲得了ACM Siggraph的杰出新研究者獎。2012年，獲得了DARPA青年教師獎，并被評為斯隆研究學(xué)者。2014年，獲得了Ruth和Joel Spira卓越教學(xué)獎。

Max Tegmark

Max Tegmark被大家親切地稱為「瘋狂的麥克斯」（Mad Max）。

憑借著自己創(chuàng)新的思維和對冒險的熱情，他的科研興趣涵蓋從精確宇宙學(xué)到探索現(xiàn)實的終極本質(zhì)。

比如，結(jié)合理論與新的測量技術(shù)，精確限定宇宙學(xué)模型及其參數(shù)。在他作為物理學(xué)研究者的前25年里，這種研究方向使他主要關(guān)注宇宙學(xué)和量子信息學(xué)。

雖然他仍與HERA合作研究宇宙學(xué)，但目前他的主要研究方向是智能的物理學(xué)，即運用物理方法深入探索生物智能和AI。

作為麻省理工學(xué)院的物理學(xué)教授，他發(fā)表了超過兩百篇技術(shù)論文，并多次在科學(xué)紀(jì)錄片中出現(xiàn)。他在SDSS項目中關(guān)于星系聚類的研究，贏得了《科學(xué)》雜志「2003年度突破」的第一名。

在此之前，Tegmark于1989年在斯德哥爾摩經(jīng)濟學(xué)院獲得了經(jīng)濟學(xué)學(xué)士學(xué)位，1990年在皇家理工學(xué)院獲得物理學(xué)學(xué)士學(xué)位。

畢業(yè)后，他便前往加州大學(xué)伯克利分校繼續(xù)深造，先后獲得物理學(xué)碩士和博士學(xué)位。

在美國西海岸生活四年后，他回到了歐洲，出任馬克斯·普朗克物理研究所的助理研究員。

1996年，他作為Hubble Fellow以及普林斯頓高級研究院的研究員，再次來到美國。

幾年后，他獲得賓夕法尼亞大學(xué)的助理教授職位，并于2003年獲得終身教職。

2004年，他來到MIT并定居在查爾斯河畔的劍橋。