革命性KAN 2.0橫空出世,劍指AI科學(xué)大一統(tǒng)!MIT原班人馬再出神作
4月30日,KAN橫空出世,很多人預(yù)言這會敲響MLP的喪鐘。
沒想到,子彈還沒飛4個月,核心團(tuán)隊(duì)又推出了KAN 2.0,瞄準(zhǔn)AI+Science領(lǐng)域,進(jìn)一步挖掘了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究竟在哪些方面改進(jìn)了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的重要性怎么強(qiáng)調(diào)都不為過,因?yàn)檫@是神經(jīng)網(wǎng)絡(luò)中用于逼近非線性函數(shù)的默認(rèn)模型,其對函數(shù)表達(dá)能力的底層邏輯由「通用逼近定理」保證。
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é)點(diǎn)處
- 給激活函數(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)點(diǎn),但深度學(xué)習(xí)和物理、化學(xué)、生物學(xué)領(lǐng)域依舊有完全不同的「語言」,這構(gòu)成了AI4Science最大的障礙之一。
因此擴(kuò)展后的KAN 2.0的終極目標(biāo)只有一個——使KAN能輕松應(yīng)用于「好奇心驅(qū)動的科學(xué)」。研究人員既能將輔助變量、模塊化結(jié)構(gòu)、符號公式等科學(xué)知識集成到KAN中,也能從KAN的可解釋性分析中得到觀察和見解。
所謂「好奇心驅(qū)動的科學(xué)」,根據(jù)論文的解釋,是過程更具有探索性、提供更基礎(chǔ)層面新發(fā)現(xiàn)和新知識的研究,比如天體運(yùn)動背后的物理原理,而非AlphaFold這類應(yīng)用驅(qū)動的科學(xué)研究。
科學(xué)與KAN的協(xié)同
具體來說,科學(xué)解釋有不同的層次,從最簡單粗略到最精細(xì)、最難發(fā)現(xiàn)、最具因果性,可以有如下幾個分類:
- 重要特征:例如,y完全由x1和x2決定,其他因素并不重要;即存在一個函數(shù)f使得y=f(x1, x2)
- 模塊化結(jié)構(gòu):例如,存在函數(shù)g和h是的y=g(x1)+h(x2)
- 符號公式:例如,y=sin(x1)+exp(x2)
MultKAN
在原始KAN網(wǎng)絡(luò)的基礎(chǔ)上,這篇最新的論文引入了一種稱為MultKAN的新模型,其核心改進(jìn)是引入額外的乘法層進(jìn)行增強(qiáng)。
KAN所依據(jù)的Kolmogorov-Arnold表示定理提出,任何連續(xù)高維函數(shù)都可以分解為單變量連續(xù)函數(shù)和加法的有限組合:
這意味著加法是唯一真正的多元運(yùn)算,,而其他多元運(yùn)算(包括乘法)都可以表示為與單變量函數(shù)組合的加法。因此,原來的KAN中僅包含加法運(yùn)算。
然而,考慮到乘法在科學(xué)和日常生活中的普遍存在,MultKAN中明確包含乘法,能更清楚地揭示數(shù)據(jù)中的乘法結(jié)構(gòu),以期增強(qiáng)可解釋性和表達(dá)能力。
如圖2所示,MultKAN和KAN相似,都包含標(biāo)準(zhǔn)KAN層,但區(qū)別在于插入了乘法節(jié)點(diǎn),對輸入的子節(jié)點(diǎn)進(jìn)行乘法運(yùn)算后再進(jìn)行恒等變換,用Python代碼可表示為:
其中⊙表示逐元素乘法。
根據(jù)上圖,整個MultKAN網(wǎng)絡(luò)進(jìn)行的運(yùn)算就可以寫作:
其中,??L≡??L°??L。
經(jīng)過擴(kuò)展后,論文將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)性)到最精細(xì)(最困難/因果關(guān)系):重要特征、模塊化結(jié)構(gòu)和符號公式。
在KANs中添加重要特征
在回歸問題中,目標(biāo)是找到一個函數(shù)f,使得y=f(x1, x2, ···, xn)。假設(shè)我們希望引入一個輔助輸入變量a=a(x1, x2, ..., xn),將函數(shù)轉(zhuǎn)化為y=f(x1, ···, xn, xa)。
盡管輔助變量a不增加新的信息,但它可以提高神經(jīng)網(wǎng)絡(luò)的表達(dá)能力。這是因?yàn)榫W(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ò)的理解,因?yàn)樗试S我們整體解釋神經(jīng)元群集,而不是單獨(dú)分析每個神經(jīng)元。
結(jié)構(gòu)模塊化的特點(diǎn)是連接群集,其中特征是群集內(nèi)的連接遠(yuǎn)強(qiáng)于群集間的連接。為此,作者引入了module方法:保留群集內(nèi)的連接,同時去除群集間的連接。
模塊由用戶來指定,語法是:
具體而言,模塊化有兩種類型:可分性和對稱性。
可分性:如果說一個函數(shù)是可分的,那么它就可以表示為非重疊變量組的函數(shù)的和或積。
廣義對稱性:如果f(x1, x2, x3, ···)=g(h(x1, x2), x3, ···),則這個函數(shù)在變量(x1, x2)上是對稱的。因?yàn)橹灰猦(x1, x2)保持不變,即使x1和x2發(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é)點(diǎn)表示表達(dá)式,邊表示操作/函數(shù);
2. 然后修改該樹以與KAN圖結(jié)構(gòu)對齊。修改包括通過虛擬邊將所有葉節(jié)點(diǎn)移動到輸入層,并添加虛擬子節(jié)點(diǎn)/節(jié)點(diǎn)以匹配KAN架構(gòu)。這些虛擬邊/節(jié)點(diǎn)/子節(jié)點(diǎn)僅執(zhí)行恒等變換;
3. 在第一層中組合變量,有效地將樹轉(zhuǎn)換為圖。
然而,通過寬度/深度擴(kuò)展來增加表達(dá)能力kanpiler生成的KAN網(wǎng)絡(luò)是緊湊的,沒有冗余邊,這可能限制其表達(dá)能力并阻礙進(jìn)一步的微調(diào)。
為了解決這個問題,作者又提出了expand_width和expand_depth方法來擴(kuò)展網(wǎng)絡(luò),使其變得更寬和更深,如圖5c所示。
KAN to Science
這一節(jié)同樣關(guān)注提取知識的三個層次,從最基本到最復(fù)雜:重要特征,模塊化結(jié)構(gòu)和符號公式。
識別重要特征
給定一個回歸模型f,有y≈f?(x1,x2,…,xn) ,我們的目標(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ò)進(jìn)行剪枝。
識別模塊化結(jié)構(gòu)
歸因分?jǐn)?shù)可以告訴我們哪些邊或節(jié)點(diǎn)更有價值,但它沒有揭示模塊化結(jié)構(gòu),即重要的邊和節(jié)點(diǎn)如何連接。
神經(jīng)網(wǎng)絡(luò)中的模塊化結(jié)構(gòu)可以分為兩種:解剖模塊化(anatomical modularity)和功能模塊化(functional modularity)。
解剖模塊化是指,空間上彼此靠近的神經(jīng)元相比距離較遠(yuǎn)的神經(jīng)元具有更強(qiá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定義了三種類型的功能模塊化:可分性、一般可分性和一般對稱性。
識別符號公式
符號公式信息量最大,因?yàn)榭梢灾苯印⑶宄亟沂竞瘮?shù)中重要的特征和模塊結(jié)構(gòu)。圖9展示了與KAN進(jìn)行交互協(xié)作進(jìn)行符號回歸的3個技巧:
1.發(fā)現(xiàn)并利用模塊化結(jié)構(gòu)
2.稀疏初始化
3.假設(shè)檢驗(yàn)
用KAN助力物理學(xué)研究
除了進(jìn)行原理層面的說明,論文還講解了多個具體案例,如何將KAN融入到現(xiàn)實(shí)的科學(xué)研究中,比如發(fā)現(xiàn)新的物理概念和定律。
論文給出的案例包括守恒量、拉格朗日量、隱藏對稱性和本構(gòu)方程等。這里我們以最簡單的守恒量發(fā)現(xiàn)為例,看看KAN是如何工作的。
守恒量即時間變化過程中保持恒定的物理量,比如能量守恒定律告訴我們,孤立系統(tǒng)的總能量保持不變。
傳統(tǒng)上,科學(xué)家如果不借助計算工具,僅靠紙筆推導(dǎo)守恒量可能非常耗時,并且需要廣泛的領(lǐng)域知識。但機(jī)器學(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, px, py)中具有3個守恒量:x軸方向的能量H1、y軸方向的能量H2和角動量H3。
關(guān)于KAN的其他應(yīng)用,論文也描述了如何從實(shí)驗(yàn)數(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)背后最主要的貢獻(xiàn)者。
他的研究興趣主要集中在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增強(qiáng)物理學(xué)研究,目標(biāo)是讓「讓AI像物理學(xué)家一樣強(qiáng)大」。
為了實(shí)現(xiàn)利用AI和物理學(xué)共建更美好世界的最終目標(biāo),Ziming Liu對包括發(fā)現(xiàn)物理定律、受物理啟發(fā)的生成模型、機(jī)器學(xué)習(xí)理論、機(jī)械解釋性等在內(nèi)的多個主題都有深厚的興趣。
并且,與凝聚態(tài)、高能物理、量子計算等領(lǐng)域的物理學(xué)家以及計算機(jī)科學(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é)位,并曾在微軟亞洲研究院實(shí)習(xí)。
Pingchuan Ma(馬平川)
Pingchuan Ma目前是MIT CSAIL實(shí)驗(yàn)室的博士生,由Wojciech Matusik教授指導(dǎo)。
他的研究方向涵蓋了「基于物理的智能」的整個流程:
1. 重建高效逼真的物理環(huán)境
2. 基于這些環(huán)境生成AI 智能體
3. 在物理世界中實(shí)現(xiàn)這些智能體
此前,他在南開大學(xué)獲得軟件工程專業(yè)學(xué)士學(xué)位,并在麻省理工學(xué)院獲得計算機(jī)科學(xué)碩士學(xué)位。
同時,他還在IBM、字節(jié)、商湯、港大等知名機(jī)構(gòu)從事過研究工作,有著豐富的經(jīng)驗(yàn)。
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é)院計算機(jī)科學(xué)與人工智能實(shí)驗(yàn)室(MIT CSAIL)的電氣工程與計算機(jī)科學(xué)教授,也是計算機(jī)圖形學(xué)小組的成員,負(fù)責(zé)帶領(lǐng)計算設(shè)計與制造團(tuán)隊(duì)。
他的研究興趣包括計算機(jī)圖形學(xué)、計算設(shè)計與制造、計算機(jī)視覺、機(jī)器人學(xué)和人機(jī)交互。
他于2003年在MIT獲得計算機(jī)圖形學(xué)博士學(xué)位,2001年在MIT獲得電氣工程與計算機(jī)科學(xué)碩士學(xué)位,1997年在加州大學(xué)伯克利分校獲得電氣工程與計算機(jī)科學(xué)學(xué)士學(xué)位。
并他曾在三菱電機(jī)研究實(shí)驗(yàn)室、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)實(shí)的終極本質(zhì)。
比如,結(jié)合理論與新的測量技術(shù),精確限定宇宙學(xué)模型及其參數(shù)。在他作為物理學(xué)研究者的前25年里,這種研究方向使他主要關(guān)注宇宙學(xué)和量子信息學(xué)。
雖然他仍與HERA合作研究宇宙學(xué),但目前他的主要研究方向是智能的物理學(xué),即運(yùn)用物理方法深入探索生物智能和AI。
作為麻省理工學(xué)院的物理學(xué)教授,他發(fā)表了超過兩百篇技術(shù)論文,并多次在科學(xué)紀(jì)錄片中出現(xiàn)。他在SDSS項(xiàng)目中關(guān)于星系聚類的研究,贏得了《科學(xué)》雜志「2003年度突破」的第一名。
在此之前,Tegmark于1989年在斯德哥爾摩經(jīng)濟(jì)學(xué)院獲得了經(jīng)濟(jì)學(xué)學(xué)士學(xué)位,1990年在皇家理工學(xué)院獲得物理學(xué)學(xué)士學(xué)位。
畢業(yè)后,他便前往加州大學(xué)伯克利分校繼續(xù)深造,先后獲得物理學(xué)碩士和博士學(xué)位。
在美國西海岸生活四年后,他回到了歐洲,出任馬克斯·普朗克物理研究所的助理研究員。
1996年,他作為Hubble Fellow以及普林斯頓高級研究院的研究員,再次來到美國。
幾年后,他獲得賓夕法尼亞大學(xué)的助理教授職位,并于2003年獲得終身教職。
2004年,他來到MIT并定居在查爾斯河畔的劍橋。