KAN 2.0: Kolmogorov-Arnold Networks Meet Science
https://github.com/KindXiaoming/pykan
Abstract
A major challenge of AI + Science lies in their inherent incompatibility: today’s AI is primarily based on connectionism, while science depends on symbolism.
To bridge the two worlds, we propose a framework to seamlessly synergize Kolmogorov-Arnold Networks (KANs) and science.
The framework highlights KANs’ usage for three aspects of scientific discovery: identifying relevant features, revealing modular structures, and discovering symbolic formulas.
The synergy is bidirectional: science to KAN (incorporating scientific knowledge into KANs), and KAN to science (extracting scientific insights from KANs).
We highlight major new functionalities in pykan:
(1) MultKAN: KANs with multiplication nodes.
(2) kanpiler: a KAN compiler that compiles symbolic formulas into KANs.
(3) tree converter: convert KANs (or any neural networks) to tree graphs.
Based on these tools, we demonstrate KANs’ capability to discover various types of physical laws, including conserved quantities, Lagrangians, symmetries, and constitutive laws.
“AI+科学”面临的一大挑战在于它们之间固有的不兼容性:当前的AI主要基于连接主义,而科学则依赖于符号主义。
为了弥合这两个领域的鸿沟,提出了一个框架,旨在将Kolmogorov-Arnold网络(KANs)与科学无缝融合。
该框架着重展示了KANs在科学发现三个方面的应用:识别相关特征、揭示模块化结构以及发现符号公式。
这种融合是双向的:科学融入KAN(将科学知识整合到KANs中),KAN反哺科学(从KANs中提取科学洞见)。
在pykan中强调了以下几项主要新功能:
(1)MultKAN:具有乘法节点的 KANs。
(2)kanpiler:一个KAN编译器,可将符号公式编译为KANs。
(3)tree converter:将KANs(或任何神经网络)转换为树形图。
基于这些工具,我们展示了KANs在发现多种物理定律方面的能力,包括守恒量、拉格朗日量、对称性以及本构关系等。
Introduction
In recent years, AI + Science has emerged as a promising new field, leading to significant scientific advancements including protein folding prediction [37], automated theorem proving [95, 83], weather forecast [41], among others. A common thread among these tasks is that they can all be well formulated into problems with clear objectives, optimizable by black-box AI systems. While this paradigm works exceptionally well for application-driven science, a different kind of science exists: curiosity-driven science. In curiosity-driven research, the procedure is more exploratory, often lacking clear goals beyond “gaining more understanding”. To clarify, curiosity-driven science is far from useless; quite the opposite. The scientific knowledge and understanding gained through curiosity often lay a solid foundation for tomorrow’s technology and foster a wide range of applications.
Although both application-driven and curiosity-driven science are invaluable and irreplaceable, they ask different questions. When astronomers observe the motion of celestial bodies, application-driven researchers focus on predicting their future states, while curiosity-driven researchers explore the physics behind the motion. Another example is AlphaFold, which, despite its tremendous success in predicting protein structures, remains in the realm of application-driven science because it does not provide new knowledge at a more fundamental level (e.g., atomic forces). Hypothetically, AlphaFold must have uncovered important unknown physics to achieve its highly accurate predictions. However, this information remains hidden from us, leaving AlphaFold largely a black box. Therefore, we advocate for new AI paradigms to support curiosity-driven science. This new paradigm of AI + Science demands a higher degree of interpretability and interactivity in AI tools so that they can be seamlessly integrated into scientific research.
Recently, a new type of neural network called Kolmogorov-Arnold Network (KAN) [57], has shown promise for science-related tasks. Unlike multi-layer perceptrons (MLPs), which have fixed activation functions on nodes, KANs feature learnable activation functions on edges. Because KANs can decompose high-dimensional functions into one-dimensional functions, interpretability can be gained by symbolically regressing these 1D functions. However, their definition of interpretability is somewhat narrow, equating it almost exclusively with the ability to extract symbolic formulas. This limited definition restricts their scope, as symbolic formulas are not always necessary or feasible in science. For example, while symbolic equations are powerful and prevalent and physics, systems in chemistry and biology the systems are often too complex to be represented by such equations. In these fields, modular structures and key features may be sufficient to characterize interesting aspects of these systems. Another overlooked aspect is the reverse task of embedding knowledge into KANs: How can we incorporate prior knowledge into KANs, in the spirit of physics-informed learning?
We enhance and extend KANs to make them easily used for curiosity-driven science. The goal of this paper can be summarized as follows:
Goal: Synergize Kolmogorov-Arnold Networks ⇔ Science. ⇐: Build in scientific knowledge to KANs (Section 3). ⇒: Extract out scientific knowledge from KANs (Section 4).
To be more concrete, scientific explanations may have different levels, ranging from the coarsest/easiest/correlational to the finest/hardest/causal:
• Important features: For example, “y is fully determined by x1 and x2, while other factors do no matter.” In other words, there exists a function f such that y = f(x1, x2).
• Modular structures: For instance, “x1 and x2 contributes to y independently in an additive way.“ This means there exists functions g and h such that y = g(x1) + h(x2).
• Symbolic formulas: For example, “y depends on x1 as a sine function and on x2 as an exponential function”. In other words, y = sin(x1) + exp(x2).
The paper reports on how to incorporate and extract these properties from KANs. The structure of the paper is as follows (illustrated in Figure 1): In Section 2, we augment the original KAN with multiplication nodes, introducing a new model called MultKAN. In Section 3, we explore ways to embed scientific inductive biases into KANs, focusing on important features (Section 3.1), modular structures (Section 3.2), and symbolic formulas (Section 3.3). In Section 4, we propose methods to extract scientific knowledge from KANs, again covering important features (Section 4.1), modular structures (Section 4.2), and symbolic formulas (Section 4.3). In Section 5, we apply KANs to various scientific discovery tasks using the tools developed in the previous sections. These tasks include discovering conserved quantities, symmetries, Lagrangians, and constitutive laws. Codes are available at https://github.com/KindXiaoming/pykan and can also be installed via pip install pykan. Although the title of the paper is “KAN 2.0”, the release version of pykan is 0.2.x.
近年来,“AI+科学”作为一个新兴且有前途的领域崭露头角,推动了包括蛋白质折叠预测[37]、自动定理证明[95, 83]、天气预报[41]在内的众多重大科学进步。这些任务的一个共同点是,它们都可以被很好地构建成具有明确目标的问题,这些问题可以通过黑盒AI系统进行优化。虽然这种范式在应用驱动的科学研究中表现出色,但科学界还存在另一种类型的研究:好奇心驱动的科学。在好奇心驱动的研究中,过程更具探索性,往往除了“获取更多理解”之外,没有明确的目标。需要明确的是,好奇心驱动的科学绝非无用,恰恰相反,它常常为未来的技术奠定坚实的科学知识和理解基础,并催生出广泛的应用。
尽管应用驱动和好奇心驱动的科学都是无价且不可替代的,但它们提出的问题却截然不同。当天文学家观测天体的运动时,应用驱动的研究者专注于预测它们未来的状态,而好奇心驱动的研究者则探索运动背后的物理学原理。另一个例子是AlphaFold,尽管它在预测蛋白质结构方面取得了巨大成功,但它仍然属于应用驱动科学的范畴,因为它并没有在更基本的层面上(如原子力)提供新知识。理论上,AlphaFold为了实现其高度准确的预测,必然已经揭示了重要的未知物理学原理。然而,这些信息对我们来说是隐藏的,使得AlphaFold在很大程度上仍然是一个黑盒。因此,本文倡导新的AI范式来支持好奇心驱动的科学。这种“AI+科学”的新范式要求AI工具具有更高的可解释性和交互性,以便能够无缝集成到科学研究中。
最近,一种名为Kolmogorov-Arnold网络(KAN)[57]的新型神经网络在与科学相关的任务中显示出了潜力。与具有固定节点激活函数的多层感知机(MLPs)不同,KANs在边上具有可学习的激活函数。由于KANs可以将高维函数分解为一维函数,因此可以通过对这些一维函数进行符号回归来获得可解释性。然而,它们对可解释性的定义相对狭隘,几乎等同于提取符号公式的能力。这种有限的定义限制了它们的应用范围,因为在科学中符号公式并不总是必要或可行的。例如,虽然符号方程在物理学中强大且普遍,但化学和生物学中的系统往往过于复杂,无法用此类方程来表示。在这些领域中,模块化结构和关键特征可能就足以描述这些系统的有趣方面。另一个被忽视的方面是将知识嵌入到KANs中的反向任务:如何以物理信息学习的精神将先验知识融入KANs?
本文增强并扩展了KANs,使其更易于用于好奇心驱动的科学。本文的目标可以概括为:
目标:协同Kolmogorov-Arnold网络科学。
:将科学知识构建到KANs中(第3节)。
:从KANs中提取科学知识(第4节)。
更具体地说,科学解释可能有不同的层次,从最粗略/最简单/相关性最强到最精细/最难/因果性最强:
- 重要特征:例如,“y 完全由 x1 和 x2 决定,而其他因素无关紧要。” 换句话说,存在一个函数 f,使得 y = f(x1, x2)。
- 模块化结构:例如,“x1和x2以相加的方式独立地对 y 做出贡献。”这意味着存在函数 g 和 h,使得 y = g(x1) + h(x2)。
- 符号公式:例如,“y 与 x1 的关系是正弦函数,与 x2 的关系是指数函数 ”。换句话说,y = sin(x1) + exp(x2)。
本文报告了如何将这些属性从KANs中融入和提取出来。论文的结构如下(如图1所示):
在第2节中,通过添加乘法节点来增强原始的KAN,引入了一种名为MultKAN的新模型。
在第3节中,探索了将科学归纳偏差嵌入到KANs中的方法,重点关注重要特征(第3.1节)、模块化结构(第3.2节)和符号公式(第3.3节)。
在第4节中,提出了从KANs中提取科学知识的方法,同样涵盖了重要特征(第4.1节)、模块化结构(第4.2节)和符号公式(第4.3节)。
在第5节中,使用前几节开发的工具将KANs应用于各种科学发现任务,包括发现守恒量、对称性、拉格朗日量和本构关系。代码可在 https://github.com/KindXiaoming/pykan 上获得,也可以通过pip install pykan进行安装。虽然论文的标题是“KAN 2.0”,但pykan的发布版本是0.2.x。
2. MultKAN: Augmenting KANs with multiplications
Kolmogorov-Arnold表示定理(KART)指出,任何连续的高维函数都可以被分解为有限数量的单变量连续函数和加法的组合:
这意味着加法是唯一的真正的多元操作,而其他多元操作(包括乘法)可以表示为加法和单变量函数的组合。例如,为了计算两个正数x和y的乘积,我们可以将其表示为xy=exp(logx+logy),其右侧仅包含加法和单变量函数(log和exp)。
然而,鉴于乘法在科学和日常生活中无处不在,因此可望在KANs中明确包含乘法操作,这可能会增强其解释性和容量。
Kolmogorov-Arnold Network (KAN)
虽然KART等式(1)对应于一个两层网络,但Liu等人[57]通过认识到看似不同的外部函数Φq和内部函数 ϕ_{q,p} 可以通过他们提出的KAN层统一起来,从而将其扩展到任意深度。一个深度为L的KAN可以通过简单地堆叠L个KAN层来构造。深度为L的KAN的形状由一个整数数组[n0,n1,…,nL]表示,其中nl表示第l层神经元的数量。第l个KAN层,具有nl个输入维度和nl+1个输出维度,将输入向量:
整个网络是L个KAN层的组合,即
在图表中,KANs可以直观地表示为由节点(求和)和边(可学习的激活)组成的网络,如图2左上角所示。当在由f(x,y)=xy生成的数据集上进行训练时,KAN(图2左下角)使用两个加法节点,这使得网络的具体作用不清晰。然而,经过一些思考,我们意识到它利用了等式,但这远非显而易见。
Multiplicative Kolmogorov-Arnold Networks (MultKAN)
为了显式地引入乘法操作,提出了MultKAN,它可以在数据中更清晰地揭示乘法结构。MultKAN(如图2顶部右侧所示)与KAN类似,都包含标准的KAN层。我们将KAN层的输入节点称为节点,将KAN层的输出节点称为子节点。KAN和MultKAN之间的区别在于从当前层的子节点到下一层节点的转换方式。在KAN中,节点直接从前一层的子节点复制而来。在MultKAN中,一些节点(加法节点)从前一层的对应子节点复制而来,而其他节点(乘法节点)则对前一层的k个子节点执行乘法操作。为了简化,设k=2(但pykan包允许k为任何整数k≥2)。
基于MultKAN的图示(图2顶部右侧),可以直观地理解为MultKAN是一个带有可选乘法插入的标准KAN。为了数学上的精确性,定义以下符号:第l层的加法(乘法)操作数分别表示为n_al(n_ml),这些被收集到数组中:加法宽度n_a ≡ [n_a0, n_a1, ..., n_aL]和乘法宽度n_m ≡ [n_m0, n_m1, ..., n_mL]。当n_m0 = n_m1 = ... = n_mL = 0时,MultKAN简化为KAN。例如,图2(顶部右侧)展示了一个MultKAN,其n_a = [2, 2, 1]且n_m = [0, 2, 0]。
一个MultKAN层由一个标准的KAN层Φ_l和一个乘法层M_l组成。Φ_l接收一个输入向量x_l ∈ R(n_al + n_ml)并输出z_l = Φ_l(x) ∈ R(n_a(l+1) + 2n_m(l+1))。乘法层由两部分组成:乘法部分对子节点对执行乘法操作,而另一部分执行恒等变换。用Python编写,M_l将z_l转换如下:
其中⊙表示逐元素乘法。MultKAN层可以简洁地表示为Ψ_l ≡ M_l ◦ Φ_l。整个MultKAN因此是:
由于乘法层中没有可训练参数,所有用于KAN的稀疏正则化技术(如ℓ1和熵正则化)都可以直接应用于MultKAN。
公式与结构
-
Kolmogorov-Arnold表示定理(KART):该定理指出任何连续的高维函数都可以分解为有限个一元连续函数和加法的组合。然而,这并没有直接包含乘法操作。
-
MultKAN的扩展:为了引入乘法操作,MultKAN在KAN的基础上增加了乘法层M_l。乘法层在标准KAN层Φ_l的输出上执行乘法操作,同时保留其他加法节点。这种设计使得MultKAN能够更直接地学习和表示乘法结构。
-
数学表示:MultKAN层的输出z_l通过乘法层M_l进一步处理,其中乘法部分对选定的子节点对执行逐元素乘法,然后将结果与未参与乘法的子节点拼接。这种结构允许MultKAN在保持KAN灵活性的同时,增加了对乘法操作的显式建模能力。
优势与应用
-
增强解释性:通过显式引入乘法节点,MultKAN能够更清晰地揭示数据中的乘法结构,从而提高模型的解释性。这对于科学发现任务尤为重要,因为乘法结构往往与物理定律和机制紧密相关。
-
提升容量:虽然KART表明任何函数都可以仅通过加法和一元函数表示,但在实际应用中,乘法操作的引入可能使模型更容易学习和表示某些复杂函数,从而提升模型的容量和性能。
-
科学发现:MultKAN能够应用于各种科学发现任务,如发现守恒量、对称性、拉格朗日量和本构定律等。
3 Science to KANs
在科学中,领域知识是至关重要的,它使我们即使在数据较少或没有数据的情况下也能有效工作。因此,对于KANs(Kolmogorov-Arnold Networks)而言,采用一种基于物理信息的方法是有益的:我们应该将可用的归纳偏差整合到KANs中,同时保持它们从数据中发现新物理现象的灵活性。
3.1 Adding important features to KANs
3.1 Adding important features to KANs
在回归问题中,目标是找到一个函数f,使得y=f(x1,x2,…,xn)。假设我们想要引入一个辅助输入变量a=a(x1,x2,…,xn),从而将函数变换为y=f(x1,…,xn,a)。尽管辅助变量a没有增加新的信息,但它可以增加神经网络的表达能力。这是因为网络不需要花费资源来计算辅助变量,而且计算可能会变得更简单,从而提高可解释性。用户可以使用augment_input
方法将辅助特征添加到输入中:
model.augment_input(original_variables, auxiliary_variables, dataset) |
以相对论质量公式为例,,其中m0是静止质量,v是质点的速度,c是光速。由于物理学家经常使用无量纲数 β≡v/c 和
,他们可能会将β和γ与v和c一起作为输入。
图3展示了有无这些辅助变量的KANs:(a)展示了从符号公式编译而来的KAN(见第3.3节),它需要5条边;(b)(c)展示了具有辅助变量的KANs,分别只需要2条或3条边,并达到了10−6和10−4的损失。注意,(b)和(c)仅随机种子不同。种子1代表一个次优解,因为它也将β=v/c识别为一个关键特征。这并不奇怪,因为在经典极限v≪c下,γ≡。
--------
这一节讨论了如何将科学领域的知识融入KANs中,以提高其在数据不足或零数据情况下的性能。主要方法包括:
添加重要特征(Adding Important Features):
在回归问题中,通过添加辅助输入变量来增加神经网络的表达能力。
使用augment_input
方法将辅助特征添加到输入中,这样网络无需自行计算这些辅助变量,从而提高计算效率和可解释性。
以相对论质量公式为例,展示了添加无量纲数β和γ作为辅助变量后,KANs的性能和可解释性显著提高。
通过这些方法,KANs能够更好地利用科学领域的先验知识,提高其在科学发现任务中的准确性和可解释性。
3.2 Building modular structures to KANs
在科学中,模块化结构是一种常见的现象,它极大地简化了复杂系统的理解和分析。例如,人类的大脑皮层被划分为多个功能各异的模块,每个模块负责特定的任务,如感知或决策制定。这种模块化也适用于神经网络,允许我们集体地解释神经元簇,而不是单独分析每个神经元。结构模块化由连接集群内部远强于集群之间的连接所表征。为了增强KANs的模块化特性,引入了模块方法,该方法在保留集群内部连接的同时,移除集群之间的连接。模块的指定由用户完成。
模块方法的语法如下:
model.module(start_layer_id, '[nodes_id]->[subnodes_id]->[nodes_id]...') |
其中,start_layer_id
指定了开始定义模块的层,'[nodes_id]->[subnodes_id]->[nodes_id]...'
描述了模块内节点和子节点的连接关系。例如,如果用户想要将特定的节点/子节点分配给一个模块——比如第1层的第0个节点,第1层的第1个和第3个子节点,以及第2层的第1个和第3个节点——他们可能会使用如下命令:
model.module(1, '[0]->[1,3]->[1,3]') |
具体地,模块化有两种类型:可分离性和对称性。
可分离性(Separability)
如果一个函数可以表示为非重叠变量组的函数之和或积,那么我们说这个函数是可分离的。例如,一个四变量函数 f(x1,x2,x3,x4) 如果可以表示为 f(x1,x2)+f(x3,x4) 或 f(x1)⋅f(x2)⋅f(x3)⋅f(x4),则它是可分离的。对于乘法可分离性,我们可以通过构建特定的模块结构来增强KANs识别这种类型结构的能力。通过将相关的节点分组到不同的模块中,可以使KANs更容易地学习到变量之间的独立贡献。
对称性(Symmetries)
在物理学和其他科学领域,对称性是一个重要的概念。对称性指的是系统在某种变换下保持不变的性质。例如,在某些物理系统中,改变粒子的位置或速度的方向可能不会对系统的行为产生影响。在KANs中,可以通过在模块设计中引入对称性约束来模拟这种物理特性。通过指定某些节点或模块在输入变换下(如交换变量顺序)的行为不变,可以训练KANs学习并尊重这些对称性。
模块化结构不仅有助于KANs更好地理解数据中的复杂关系,还可以提高模型的可解释性和泛化能力。通过将大型神经网络分解为可解释的小模块,科学家可以更容易地理解模型的决策过程,并发现隐藏在数据中的潜在科学规律。
3.3 Compiling symbolic formulas to KANs
在科学研究中,符号公式是表达物理定律和数学关系的重要工具。为了将科学知识无缝地融入KANs中,开发了一个名为kanpiler
的工具,它能够将符号公式编译成KANs。通过这种方式,可以利用已有的科学理论来指导KANs的学习过程,同时保持KANs从数据中发现新物理的灵活性。
编译过程
kanpiler
的工作流程可以概括为以下几个步骤:
-
解析符号公式:首先,
kanpiler
解析输入的符号公式,识别出其中的变量、函数和运算符。 -
构建KAN结构:基于解析的结果,
kanpiler
自动构建出一个与符号公式相对应的KAN结构。这个结构包括必要的KAN层和边缘(即学习激活函数)。 -
初始化参数:KAN结构中的参数(即边缘的激活函数)会被初始化为能够近似符号公式中相应函数的形式。这通常是通过预训练或设置合理的初始值来实现的。
-
微调与优化:最后,如果需要,可以使用数据对KAN进行微调,以进一步优化其性能。然而,由于KAN的初始结构已经基于符号公式构建,因此微调通常不是必需的,除非要处理与原始公式略有不同的变体或噪声数据。
示例
假设我们有一个简单的符号公式 y=sin(x1)+exp(x2),想要将其编译成一个KAN。使用kanpiler
,可以得到如下的KAN结构:
- 输入层:包含两个节点,分别对应 x1 和 x2。
- 隐藏层:包含两个KAN层,每个层有两个节点(或子节点)。第一个层中的节点分别计算 sin(x1) 和 exp(x2),第二个层中的节点将这两个结果进行加法运算。
- 输出层:包含一个节点,输出最终的结果 y。
在这个过程中,kanpiler
会自动设置KAN层中的激活函数,使得它们能够近似 sin 和 exp 函数。然后,通过加法节点将这两个函数的结果相加,得到最终的输出 y。
4 KANs to Science
当今的黑盒深度神经网络功能强大,但解释这些模型仍然具有挑战性。科学家不仅追求高性能的模型,还希望从这些模型中提取出有意义的知识。在本节中,我们专注于提高KANs(知识增强网络)在科学应用中的可解释性。我们将从KANs中提取知识的三个层次进行探讨,从最基础到最复杂:重要特征(第4.1节)、模块化结构(第4.2节)和符号公式(第4.3节)。
4.1 Identifying important features from KANs
识别重要变量对于许多任务至关重要。给定一个回归模型f,其中y ≈ f(x1, x2, ..., xn),我们的目标是为输入变量分配分数以衡量它们的重要性。Liu等人[57]使用L1范数来表示边的重要性,但这种度量可能存在问题,因为它仅考虑局部信息。
为了解决这个问题,我们引入了一个更有效的归因分数,它比L1范数更能反映变量的重要性。为了简化,假设存在乘法节点,因此不需要区分节点和子节点。假设有一个L层的KAN,其宽度为[n0, n1, ..., nL]。将E_{l,i,j} 定义为 (l, i, j) 边上激活的标准差,N_{l,i} 定义为 (l, i) 节点上激活的标准差。然后,定义节点(归因)分数 A_{l,i} 和边(归因)分数B_{l,i,j}。在[57]中,我们简单地定义了B_{l,i,j} = E_{l,i,j} 和 A_{l,i} = N_{l,i}。然而,这种定义没有考虑到网络的后续部分;即使一个节点或边本身具有较大的范数,如果网络的其余部分实际上是零函数,它也可能不会对输出做出贡献。因此,我们现在从输出层到输入层迭代地计算节点和边的分数。我们将所有输出维度的分数设置为1,即A_{L,i} = 1, i = 0, 1, ..., n_L-1,并按以下方式计算分数:
比较 E_{l,i,j} 和 B_{l,i,j}
我们发现 B_{l,i,j} 更准确地反映了边的重要性。在图6中,我们比较了在两个方程 y = exp(sin(πx_1) + x^2_2)和y = (x^2_1 + x22)2 + (x^2_3 + x^2_4)^2上训练的KANs,并可视化了使用E(L1范数)或B(归因分数)作为重要性分数的KANs。对于第一个方程,归因分数揭示了一个比L1范数更清晰的图,因为第一层中的许多活跃边由于后续的不活跃边而不贡献于最终输出。归因分数考虑了这一点,从而产生了更有意义的图。对于第二个方程y = (x^2_1 + x^2_2)^2 + (x^2_3 + x^2_4)^2,我们可以从符号方程中看出四个变量同等重要。归因分数正确地反映了这四个变量同等重要,而L1范数则错误地表明x3和x4比x1和x2更重要。
基于归因分数的输入剪枝
在真实数据集中,输入维度可能很大,但只有少数变量可能是相关的。为了解决这个问题,我们提出基于归因分数剪枝不相关的特征,以便我们能够关注最相关的特征。用户可以使用prune_input来仅保留最相关的变量。例如,如果在函数中,有100个输入特征按相关性递减排序,且训练后只有前五个特征的归因分数显著较高,则prune_input方法将仅保留这五个特征。剪枝后的网络变得紧凑且可解释,而原始的具有100个输入的KAN则过于密集,难以直接解释。
4.2 Identifying modular structures from KANs
尽管归因分数为哪些边或节点是重要的提供了有价值的见解,但它并未揭示模块化结构,即重要的边和节点是如何相互连接的。在本部分中,我们旨在通过检查两种类型的模块化——解剖模块化和功能模块化,来揭示训练后的KANs(知识感知网络)和MLPs(多层感知机)中的模块化结构。
4.2.1 解剖模块化
解剖模块化指的是在空间上彼此接近的神经元具有比距离较远的神经元更强的连接倾向。尽管人工神经网络缺乏物理空间坐标,但引入物理空间的概念已被证明能够增强可解释性[51, 52]。我们采用了[51, 52]中的神经元交换方法,该方法在保持网络功能的同时缩短了连接。我们称这种方法为auto_swap。通过神经元交换揭示的解剖模块化结构便于甚至通过视觉方式轻松识别模块,如图7所示的两个任务:
(1)多任务稀疏奇偶校验;(2)层次化多数投票。对于多任务稀疏奇偶校验,我们有10个输入位xi ∈ {0, 1},i = 1, 2, ..., 10,以及输出yj = x2j-1 ⊕ x2j, j = 1, ..., 5,其中⊕表示模2加法。该任务展现出模块化特性,因为每个输出仅依赖于输入的一个子集。auto_swap成功地为KANs和MLPs识别了模块,KAN发现的模块更为简单。对于层次化多数投票,有9个输入位xi ∈ {0, 1},i = 1, ..., 9,以及输出,其中maj表示多数投票(如果两个或三个输入为1,则输出1,否则为0)。KAN在auto_swap之前就已经揭示了模块化结构,并且在auto_swap之后图表变得更加有序。MLP从第一层权重的模式中显示出一些模块化结构,表明变量之间的相互作用,但无论是否进行auto_swap,全局模块化结构仍然不清晰。
4.2.2 功能模块化
功能模块化与神经网络所表示的整体功能有关。给定一个Oracle网络(其内部细节如权重和隐藏层激活状态无法访问或过于复杂而无法分析),我们仍然可以通过在输入和输出上进行前向和后向传递来收集关于功能模块化的信息。我们主要基于[84]定义了三种类型的功能模块化(见图8(a))。
可分性:如果一个函数f是加性可分的,则
注意,当1 ≤ i ≤ k, k + 1 ≤ j ≤ n时,。为了检测可分性,可以计算Hessian矩阵
并检查其块结构。如果对于所有1 ≤ i ≤ k和k + 1 ≤ j ≤ n,Hij = 0,则我们知道f是加性可分的。对于乘性可分性,我们可以通过取对数将其转换为加性可分性:
为了检测乘性可分性,我们定义,并检查其块结构。用户可以调用test_separability来测试一般可分性。
广义可分性:如果函数f满足以下形式,则称其具有广义可分性
为了检测广义可分性,我们计算
这里我们使用了∂F/∂g = ∂F/∂h的假设。并且注意到在乘法上是可分的,因此可以通过上述提出的可分性测试来检测。用户可以调用test_general_separability来检查加法可分性或乘法可分性。
广义对称性:如果函数在前k个变量上具有广义对称性,则满足
我们记y = (x1, · · · , xk)和z = (xk+1, · · · , xn)。这一属性被称为广义对称性,因为只要h保持不变,无论x1, · · · , xk的单个值如何变化,f都保持相同的值。我们计算f关于y的梯度:。由于∂g/∂h是一个标量函数,它不会改变∇yh的方向。因此,
的方向与z无关,即
这是对称性的条件。用户可以调用test_symmetry方法来检查对称性。
树转换器:三种功能模块化类型形成了一个层次结构:对称性是最一般的,广义可分性位于中间,可分性是最具体的。数学上,
Separability ⊂ Generalized Separability ⊂ Generalized Symmetry (16)
为了获得模块化结构的最大层次,我们递归地应用广义对称性检测,形成小到k=2个变量、大到k=n个变量的组。例如,考虑一个8变量的函数
它具有四个k=2的广义对称性,涉及组(x1, x2)、(x3, x4)、(x5, x6)、(x7, x8);以及两个k=4的广义对称性,涉及组(x1, x2, x3, x4)和(x5, x6, x7, x8)。因此,每个k=4的组包含两个k=2的组,展示了层次结构。对于每个广义对称性,我们还可以测试它是否进一步具有广义可分性或可分性。用户可以使用plot_tree方法来获取函数的树状图(该函数可以是任何Python表达式、神经网络等)。对于神经网络模型,用户可以直接调用model.tree()。树状图可以是默认的‘tree’样式或‘box’样式。
示例
图8(b)提供了两个示例。当将精确的符号函数输入到plot_tree中时,会获得真实的树图。我们特别关注树转换器是否适用于神经网络。对于这些简单情况,只要训练充分,KANs和MLPs都能找到正确的图。图8(b)(底部)展示了KAN和MLP训练过程中树图的演变。特别有趣的是看到神经网络如何逐渐学习正确的模块化结构。在第一个示例中,KAN和MLP都逐渐吸收了更多的归纳偏置(它们的中间状态不同),直到达到正确的结构。在第二个示例
中,两个模型最初都检测到了所有三个变量的乘法可分性,显示出比正确结构更高的对称性。随着训练的进行,两个模型都“意识到”:为了更好地拟合数据(损失变得更低),这种高对称结构无法再满足,应该放松为更不严格的结构。另一个观察结果是,KAN具有MLP中没有的中间结构。我们想要提到两个注意事项:
(1)结果可能依赖于种子和/或阈值。
(2)所有测试都依赖于二阶导数,由于模型仅使用零阶信息进行训练,因此可能不够鲁棒。
对抗性构造如fϵ(x) = f(x) + ϵsin(xϵ)可能导致问题,因为虽然随着,但
。尽管在实践中这种极端情况不太可能发生,但平滑性对于确保我们方法的成功是必要的。
4.3 Identifying symbolic formulas from KANs
符号公式是最具信息量的,因为它们一旦已知,就能清晰地揭示出重要的特征和模块化结构。在Liu等人的研究中[57],作者展示了一系列示例,从中可以提取符号公式,并在需要时利用一些先验知识。借助上述提出的新工具(特征重要性、模块化结构和符号公式),用户可以轻松地利用这些新工具与KANs进行交互和协作,从而使符号回归变得更加容易。下面将介绍三种技巧,并在图9中进行说明。
技巧A:发现和利用模块化结构
我们可以首先训练一个通用网络并探测其模块化特性。一旦识别出模块化结构,我们就以这种模块化结构作为归纳偏置来初始化一个新模型。例如,考虑函数f(q, v, B, m) = qvB/m。我们首先初始化一个大型KAN(假定其表达能力足够强)以将数据集拟合到合理的准确度。训练后,从训练好的KAN中提取树图(参考第4.2节),该图显示了乘法可分性。然后,我们可以将模块化结构构建到第二个KAN中(参考第3.2节),对其进行训练,并将所有一维函数符号化以推导出公式。
技巧B:稀疏初始化
符号公式通常对应于具有稀疏连接的KANs(见图5(b)),因此稀疏初始化KANs可以使其更好地与符号公式的归纳偏置对齐。否则,密集初始化的KANs需要仔细的正则化以促进稀疏性。稀疏初始化可以通过向KAN初始化器传递参数“sparse_init=True”来实现。例如,对于函数 f(q, E, v, B, θ) = q(E + vBsinθ),稀疏初始化的KAN与最终训练好的KAN非常相似,训练时只需进行微调。相比之下,密集初始化则需要大量的训练来移除不必要的边。
技巧C:假设检验
当面临多个合理的假设时,我们可以尝试所有假设(分支到“平行宇宙”)以测试哪个假设最准确和/或最简单。为了促进假设检验,我们构建了一个检查点系统,该系统在发生更改(如训练、剪枝)时自动保存模型版本。例如,考虑函数f。我们从一个随机初始化的KAN开始,其版本为0.0。训练后,它演变为版本0.1,此时它在β = v/c和
上均被激活。假设可能只需要β或γ中的一个。我们首先将γ上的边设置为零,并训练模型,得到6.5×10-4的测试RMSE(版本0.2)。为了测试另一个假设,我们需要返回到分支点(版本0.1)——我们调用model.rewind(‘0.1’),它将模型回滚到版本0.1。为了表示调用了回滚,版本0.1被重命名为版本1.1。现在我们将β上的边设置为零,训练模型,得到2.0×10-6的测试RMSE(版本变为1.2)。比较版本0.2和1.2表明,第二个假设更好,因为在相同复杂度下损失更低(两个假设都有两条非零边)。