KAN: Kolmogorov–Arnold Networks

KAN: Kolmogorov–Arnold Networks

基本信息

博客贡献人

燕青

作者

Ziming Liu, Yixuan Wang, Sachin Vaidya, et al.

标签

Deep Learning

摘要

受Kolmogorov–Arnold表示定理的启发，本文提出了Kolmogorov–Arnold Networks ( KAN )，作为多层感知机( MLP )的一种可行的替代方案。MLP在每一个节点（神经元）上具有固定的激活函数，KAN在边（权重）上具有可学习的激活函数。KAN中完全没有线性权重——每个权重参数都被替换为一个单变量样条函数（spline function）。本文表明，这种看似简单的变化使得KAN在准确性和可解释性方面优于MLP。

样条函数：分段定义的多项式参数曲线

在精度方面，较小的KAN在数据拟合和偏微分方程（PDE）求解方面可以达到与较大的MLP相当或更好的精度。对于可解释性，KAN可以直观地可视化，并且可以方便地与人类用户进行交互。通过数学和物理学中的两个例子，KAN被证明是帮助科学家（重新）发现数学和物理规律的有用的"合作者"。

总之，KAN是MLP的有前途的替代品，为进一步改善当今严重依赖MLP的深度学习模型提供了机会。

图1. MLP（左） vs. KAN（右）

问题定义

多层感知器( Multi-layer Perceptrons，MLP )，也称为全连接前馈神经网络，是当今深度学习模型的基础构建模块。MLP是机器学习中用于逼近非线性函数的默认模型，因为它们的表达能力由万能近似定理保证。然而，MLP是否是我们所能建立的最好的非线性回归器？尽管MLP的使用非常普遍，但它们都有很大的缺点。例如，在TransFormer中，MLP几乎消耗了所有的非嵌入层参数，并且在没有后期分析工具的情况下，通常是可解释性较差的(相对于注意力层)。

本文提出了一种有前途的MLP的替代方案，称为Kolmogorov–Arnold Networks ( KAN )。与MLP一样，KAN也具有全连接结构。然而，MLP在节点( “神经元”)上具有固定的激活函数，KAN在边( “权重”)上具有可学习的激活函数，如图1所示。因此，KAN没有线性权重矩阵：相反，每个权重参数被一个可学习的1D函数替代，该函数被参数化为样条函数（spline）。

KAN的节点只需对输入信号进行简单的求和，而不需要施加任何非线性。人们可能会担心KAN的成本太高，因为每个MLP的权重参数都变成了KAN的样条函数。幸运的是，KAN通常仅需比MLP小得多的计算图。例如，对于PDE求解，深2层宽为10的KAN，比深4层宽为100的MLP在精确度上优于100倍，参数效率上优于100倍。

利用Kolmogorov–Arnold表示定理来构建神经网络的可能性已经被研究过了。然而，大多数工作都停留在原始的深度为2，宽度为( 2n + 1)的表示上，没有利用更多的现代技术（如反向传播）来训练网络。本文的贡献在于将原始的Kolmogorov–Arnold表示推广到任意的宽度和深度，在当今的深度学习世界中重新激活，并使用广泛的实证实验来突出其作为AI + Science的基础模型的潜在作用，因为它具有准确性和可解释性。

方法

总体框架

图2. KAN在数学上是可靠的、准确的和可解释的。

Kolmogorov-Arnold表示定理

Arnold和Kolmogorov证明了如果 $f$ 是有界区域上的多元连续函数，那么 $f$ 可以写成单变量连续函数和加法二元运算的有限组合。更具体地，对于光滑的 $f:[0,1{]}^n\to \mathbb{R}$，其中 ${\varphi }_{q,p}:[0,1]\to \mathbb{R}$，${\Phi }_q:\mathbb{R}\to \mathbb{R}$：
$$f(x)=f(x_1,...,x_n)=\sum _{q=1}^{2n+1}{\Phi }_q(\sum _{p=1}^n{\varphi }_{q,p}(x_p))\text{\hspace{1em}(1)}$$
从上述公式似乎可以看出：学习一个高维函数可以归结为学习一个多项式个数的一维函数。然而，这些一维函数可能是非光滑的，甚至是分形的，因此它们在实际中可能是不可学习的。由于这种病态行为，Kolmogorov-Arnold表示定理在机器学习中基本被判为死亡，被认为是理论上可靠但实际上无用的。

然而，本文更乐观地看待Kolmogorov-Arnold定理对机器学习的有用性。首先，不必拘泥于原方程，方程(1)只有两层非线性项，且隐含层中含有少量项(2n+1)：本文则将网络推广到任意宽度和深度。其次，科学和日常生活中的大多数函数通常是光滑的，并且具有稀疏的组成结构，这可能有利于光滑的Kolmogorov-Arnold表示。这里的理念接近于物理学家的思维定式，物理学家往往更关心典型案例而非最坏情况。

KAN架构

假设有一个由输入-输出对 { x i , y i } \{x_i,y_i\} {xi​,yi​} 组成的监督学习任务，其中我们想要找到 $f$ 使得对于所有的数据点，$y_i\approx f(x_i)$。方程(1)意味着，如果能够找到合适的单变量函数 ${\varphi }_{q,p}$ 和 ${\Phi }_q$，就可以找到 $f$。这启发我们设计一个显式参数化方程的神经网络。方程(1)中由于所有要学习的函数都是一元函数，可以将每个1D函数参数化为B样条曲线（B-spline），其中局部B样条基函数（图3右）的系数是可学习的。现在就有了一个KAN的原型，它的计算图正好由方程(1)所表示，表现为一个两层神经网络，激活函数放置在边上而不是节点，隐层宽度为2n + 1。

图3. 左：流经网络的激活标记。右：激活函数被参数化为B样条，允许在粗粒度和细粒度网格之间切换。

如前文所述，这样的网络过于简单，在实践中无法用光滑的样条很好地逼近任意函数。因此，本文将KAN推广到更广和更深的范围。由于Kolmogorov-Arnold表示对应两层KAN，如何使KAN更深并不清楚，目前还没有一个"广义"版本的定理对应于更深层次的KAN。

当注意到MLP和KAN之间的类似性时，这种突破就发生了。在MLP中，一旦定义了一层（由线性变换和非线性项组成），就可以堆叠更多的层来使网络更深。要构建深层KAN，首先要回答：“什么是KAN层? ”结果表明，一个具有 $n_{in}$ 维输入和 $n_{out}$ 维输出的KAN层可以被定义为一个一维函数的矩阵：
Φ = { ϕ q , p } , p = 1 , 2 , . . . , n i n , q = 1 , 2 , . . . , n o u t (2) \Phi=\{\phi_{q,p}\},\quad p=1,2,...,n_{in},\quad q=1,2,...,n_{out}\tag{2} Φ={ϕq,p​},p=1,2,...,nin​,q=1,2,...,nout​(2)
其中，函数 ${\varphi }_{q,p}$ 具有可训练参数，具体如下。在Kolmogov - Arnold定理中，内层函数形成 $n_{in}=n$ 和 $n_{out}=2n+1$ 的KAN层，外层函数形成 $n_{in}=2n+1$，$n_{out}=1$ 的KAN层。因此，方程(1)中的表示是两个KAN层的简单叠加。现在可以清楚地知道，拥有更深的Kolmogov - Arnold表示意味着：简单地堆叠更多的KAN层。

一个完整KAN的形状由下列数组定义：
$$[n_0,n_1,...,n_L]\text{\hspace{1em}(3)}$$
其中 $n_i$ 是计算图的第 $i$ 层的节点数。用 $(l,i)$ 表示第 $l$ 层的第 $i$ 个神经元，用 $x_{l,i}$ 表示 $(l,i)$ 神经元的激活值。在层 $l$ 和层 $l+1$ 之间，有 $n_l{n}_{l+1}$ 个激活函数：连接 $(l,i)$ 和 $(l+1,j)$ 的激活函数记为：
$${\varphi }_{l,j,i},l=0,...,L-1,i=1,...,n_l,j=1,...,n_{l+1}\text{\hspace{1em}(4)}$$
${\varphi }_{l,j,i}$ 激活前为 $x_{l,i}$ ；${\varphi }_{l,j,i}$激活后为 ${\tilde{x}}_{l,j,i}\equiv {\varphi }_{l,j,i}(x_{l,i})$ 。$(l+1,j)$ 神经元的激活值就是所有激活后值的总和：
$$x_{l+1,j}=\sum _{i=1}^{n_l}{\tilde{x}}_{l,j,i}=\sum _{i=1}^{n_l}{\varphi }_{l,j,i}(x_{l,i}),j=1,...,n_{l+1}\text{\hspace{1em}(5)}$$
一般的KAN网络是由L层组成：给定一个输入向量 $x_0\in {\mathbb{R}}^{n_0}$ ，KAN的输出为:
$$KAN(x)=({\Phi }_{L-1}\circ {\Phi }_{L-2}\circ ...{\Phi }_1\circ {\Phi }_0)x\text{\hspace{1em}(6)}$$

实现细节

• 残差激活函数：包含一个基函数 $b(x)$ (类似于残差连接)，使得激活函数 $\varphi (x)$ 是基函数 $b(x)$ 和样条函数的和：
  $$\varphi (x)=w(b(x)+spline(x))\text{\hspace{1em}(7)}$$
  设：
  $$b(x)=silu(x)=x/(1+e^{-x})\text{\hspace{1em}(8)}$$
  在大多数情况下，$spline(x)$ 被参数化为B样条的线性组合：
  $$spline(x)=\sum _i{c}_i{B}_i(x)\text{\hspace{1em}(9)}$$
  其中 $c_i$ 是可训练的。原则上 $w$ 是多余的，因为它可以被归入 $b(x)$ 和 $spline(x)$ 中。但是，本文仍然加入了这个 $w$ 因子，以更好地控制激活函数的整体大小。

• 初始化尺度：每个激活函数初始化为 $spline(x)\approx 0$ ， $w$ 的初始化使用MLP线性层采用的Xavier初始化来进行。

• 样条网格的更新：根据输入的激活值动态更新每个网格，以解决样条函数定义在有界区域上，但激活值在训练过程中会超出固定区域的问题。

网格拓展

原则上，样条可以精确到任意目标函数，正如网格可以划分为任意细粒度一样。这种良好的特性被KAN继承。相比之下，MLP没有"细粒度"的概念。诚然，增加MLP的宽度和深度可以导致性能的提高。然而，这种做法是缓慢且昂贵的，因为不同大小的模型都是独立训练的。相比之下，对于KAN，可以先训练一个参数较少的KAN，然后通过简单地将其样条网格变得更细，将其扩展为一个参数更多的KAN，而不需要从头开始重新训练更大的模型。

KAN缩放/标度定律（Neural scaling laws）

神经缩放定律是测试损失随模型参数增加而减小的现象。$\ell \propto N^{-\alpha }$其中，$\ell$ 为检验均方根误差RMSE，$N$ 为参数个数，$\alpha$ 为缩放指数。

通过简单地放大模型，较大的 $\alpha $ 可以带来更多的改进。

本文的方法假设存在光滑的Kolmogorov-Arnold表示，将高维函数分解成若干个一维函数，给出 $\alpha =k+1$ (其中 $k$ 是样条的分段多项式阶数)。

本文选择 $k=3$ 的三次样条，因此 $\alpha =4$ ，这是与其他工作相比最大和最好的缩放指数。

可解释性

剪枝

• 稀疏化

  对于MLP，使用线性权重参数的L1正则化来进行模型的稀疏化。KAN可以适应这种高层次的想法，但需要两个修改：

  • KAN中不存在线性的"权重"。线性权重由可学习的激活函数代替，因此需要定义这些激活函数的L1范数。
  • 本文发现L1不足以对KAN进行稀疏化；为此，需要一个额外的熵正则化。

  总的训练目标 ${\ell }_{total}$ 是所有KAN层的预测损失 ${\ell }_{pred}$ 加上L1和熵正则化,其中 $S$ 为熵正则化函数：
  $${\ell }_{\text{total }}={\ell }_{\text{pred }}+\lambda \left({\mu }_1\sum _{l=0}^{L-1}{\left|{\mathbf{\Phi }}_l\right|}_1+{\mu }_2\sum _{l=0}^{L-1}S\left({\mathbf{\Phi }}_l\right)\right)\text{\hspace{1em}(10)}$$

• 剪枝

  在使用稀疏化惩罚进行训练后，若还希望将网络剪枝到一个更小的子网络中，可以在节点级别上对KAN进行稀疏化。对于每个节点，定义其输入和输出分数为：
  $$I_{l,i}=\underset{k}{\max }\left({\left|{\varphi }_{l-1,k,i}\right|}_1\right),O_{l,i}=\underset{j}{\max }\left({\left|{\varphi }_{l+1,j,i}\right|}_1\right)\text{\hspace{1em}(11)}$$
  默认情况下，如果输入和输出分数都大于阈值超参数 $\theta =10^{-2}$，则认为节点是重要的，所有不重要的神经元都被修剪掉。

• 符号化

  在怀疑某些激活函数实际上是某些符号(例如cos或log)的情况下，本文提供了一个接口将它们设置为指定的符号形式，fixed _ symbol( l , i , j , f) 函数可以将 $(l,i,j)$ 的激活函数设置为 $f$ 。但是，不能简单地将激活函数设定为精确的符号公式，因为它的输入和输出可能存在偏移和缩放。因此，从样本中获得激活前的 $x$ 和激活后的 $y$ ，并拟合仿射参数$(a、b、c、d)$，使得 $y\approx cf(ax+b)+d$ 。拟合是通过$a、b$ 和线性回归的迭代网格搜索来完成的。

人类如何与KAN交互

给定一个回归任务：
$$f(x,y)=\exp \left(\sin (\pi x)+y^2\right)\text{\hspace{1em}(12)}$$

• Step1：稀疏化训练

  从一个全连接的 [2, 5, 1] KAN开始，用稀疏正则化进行训练可以使模型变得稀疏。

• Step2：剪枝

  可以看到自动剪枝舍弃了除最后一个隐层神经元之外的所有隐层神经元，留下了一个 [ 2, 1, 1] KAN。激活函数看起来是已知的符号函数。

• Step3：设置符号函数

  假设用户从KAN图中可以直接猜出这些符号公式，则可以对其进行设置。
  $$\begin{array}{l}fix\_symbolic(0,0,0,sin)\\ fix\_symbolic(0,1,0,x^2)\\ fix\_symbolic(1,0,0,exp)\end{array}\text{\hspace{1em}(13)}$$
  在用户没有领域知识或者不知道这些激活函数可能是哪些符号函数的情况下，本文提供了一个函数suggest _ symbol来推荐符号化的候选者。

• Step4：进一步训练

  将网络中所有的激活函数符号化后，剩下的参数只有仿射参数（平移和缩放）。继续训练这些仿射参数，当看到损失下降到机器精度时，就找到了正确的符号表达式。

• Step5：输出符号公式

  Sympy用于计算输出节点的符号公式。用户获得 $1.0e^{1.0y+1.0sin(3.14x)}$，即真实答案（对于 $\pi$，只显示了两个小数）。

图4. KAN进行符号回归的举例

实验

准确性

在这一部分中，本文证明了在各种任务(回归和PDE求解)中，KAN比MLP更有效地表示函数。

小规模数据集(Toy Dataset)

我们清楚地知道“真实”KAN的形状。

• $f(x)=J_0(20x)$ 为贝塞尔函数。由于它是一元函数，可以用样条来表示，也可以被表示为[1 , 1] 的KAN。
• $f(x)=exp(sin(\pi x)+y^2)$ 可以被表示为[2 , 1, 1] 的KAN。
• $f(x)=exp(\frac{1}{100}{\sum }_{i=1}^{100}{\sin }^2(\frac{\pi x_i}{2}))$ 可以被表示为 [100, 1, 1] 的KAN。

在图4中绘制了KAN和MLP的测试RMSE随参数个数的变化曲线，表明KAN比MLP具有更好的缩放曲线，特别是对于高维样本。

图5. 在五个小例子上比较KAN和MLP。 KAN几乎可以饱和本文理论预测的最快缩放定律（scaling law）( α = 4)，而MLP缩放缓慢且很快达到瓶颈。

特殊函数

我们清楚地不知道"真实" KAN形状。

• 寻找特殊函数的近似KA表示是可能的，从KA表示的角度揭示了特殊函数新的数学性质。
• 与MLP相比，KAN在表示特殊函数时更加高效和准确。

本文收集了数学和物理中常见的15个特殊函数，汇总在图6中。选取固定宽度为5或100，深度为{ 2，3，4，5，6 }的MLP。分别运行有剪枝和无剪枝的KAN。无剪枝的KAN：固定KAN的形状，其宽度设置为5，深度在{ 2，3，4，5，6 }中扫描。剪枝Kan：利用KAN稀疏化和剪枝技术，从一个固定形状的KAN中剪枝得到一个更小的KAN。每个KAN初始化为G=3，用LBFGS进行训练，每200步增加网格点数，覆盖G = { 3，5，10，20，50，100，200 }。对于每个超参数组合，运行3个随机种子。

图6. KAN拟合特殊函数表现

图7. KAN拟合特殊函数对比MLP

Feynman数据集

图8. Feynman数据集函数拟合表现对比

可解释性

在这一部分中，本文证明了KAN具有可解释性和交互性。本节不仅测试了KAN在人造任务上的使用情况，还测试了KAN在真实科学研究中的使用情况。由于其准确性和可解释性，KAN可能成为AI+Science的基础模型。

• 有监督小规模数据集

  本文首先考察了KAN揭示符号公式组成结构的能力。图9中列举了六个例子，并可视化了它们的KAN。KAN能够揭示这些公式中存在的组成结构，并学习正确的单变量函数。

  图9. KAN对于简单的符号任务具有可解释性
  

• 无监督小规模数据集

  除有监督学习外，另一种类型的科学发现可以表述为无监督学习，即给定一组输入变量 $x_1,x_2,...,x_d$ ,希望发现变量之间的结构关系。具体来说，期望得到一个非零的 $f$ 使得：
  $$f(x_1,x_2,...,x_d)\approx 0\text{\hspace{1em}(14)}$$
  例如，考虑一组满足 $x_3=exp(sin(\pi x_1)+x_2^2)$ 的特征 $(x_1,x_2,x_3)$。则一个有效的 $f$ 是 $f(x_1,x_2,x_3)=sin(\pi x_1)+x_2^2-log(x_3)=0$，这意味着 $(x1,x2,x3)$ 的点构成由 $f=0$ 指定的2维子流形，而不是填充整个3维空间。

  本文通过将无监督学习问题转化为对所有 $d$ 个特征的有监督学习问题来解决，而不需要划分。其基本思想是学习一个函数 $f(x_1,...,x_d)=0$，使得 $f$ 不是0-函数。为此，类似于对比学习，本文定义了正样本和负样本：正样本是真实数据的特征向量。负样本通过特征分解（feature corruption）构造。为了确保每个拓扑不变量的整体特征分布保持不变，本文通过在整个训练集中对每个特征进行随机排列来执行特征分解。

  现在训练一个网络 $g$，使得$g(x_{real})=1$，$g(x_{fake})=0$，从而将问题转化为一个有监督的问题。然而，最初的期望是使 $f(x_{real})=0$ 和$f(x_{fake})\ne 0$ 。可以通过令 $g=\sigma \circ f$ 来实现，其中$σ ( x ) = exp(-\frac{x^2}{2w2}) $是一个具有较小宽度 $w$ 的高斯函数，这可以通过一个形状为 $[...,1,1]$ 的KAN方便地实现，其最后一个激活设置为高斯函数 $\sigma$，所有的前几层都形成 $f$ 。除了上面提到的修改之外，其他的对于有监督训练来说都是一样的。

  本文证明了无监督范式对于一个合成的例子是有效的。考虑一个6维数据集，当 $(x1,x2,x3)$ 是因变量时， $x_3=exp(sin(x_1)+x_2^2)$；当$(x_4,x_5)$ 为因变量时，$x_5=x_4^3$；$x_6$独立于其他变量。在图10中，本文证明了对于seed = 0，KAN揭示了x1，x2和x3之间的函数依赖关系。对于另一个种子= 2024，KAN揭示了x4和x5之间的函数依赖关系。目前，本文的初步结果依赖于随机性(不同的种子)来发现不同的关系；在未来，本文希望研究一种更系统、更可控的方法来发现一组完整的关系。

  图10. 小规模任务的无监督学习。KAN可以识别因变量组( x1, x2, x3)和(x4, x5)。
  

• 数学应用：纽结理论（Knot Theory）

  图11. KAN在纽结数据集上，监督模式下的表现。
  

相关知识链接

B样条函数

数值计算 — B样条函数（B-spline）

论文原文

KAN: Kolmogorov–Arnold Networks

总结

• 受Kolmogorov–Arnold表示定理的启发，本文提出了KAN作为MLP的一种可行的替代方案
• 相比于MLP在每一个节点（神经元）上具有固定的激活函数，KAN在边（权重）上具有可学习的激活函数。
• KAN在准确性方面优于MLP，且具有一定的可解释性

[局限]

• KAN的理论基础为学习一个高维函数可以归结为学习一个多项式个数的一维函数。然而，这些一维函数可能是非光滑甚至是分形的，因此在现实中可能是不可学习的。尽管本文认为典型情况下拟合的绝大部分1维函数都是光滑的，但仍无法解决最坏的情况。

[启发]

• 将经典数学定理结合当下技术，分析两者的结合点并提出创新（KA理论结合神经网络层，提出KAN层）
• 舍弃最坏情况进行一般化假设
• 增强深度神经网络模型可解释性的思路

BibTex

@misc{liu2024kan,
      title={KAN: Kolmogorov-Arnold Networks}, 
      author={Ziming Liu and Yixuan Wang and Sachin Vaidya and Fabian Ruehle and James Halverson and Marin Soljačić and Thomas Y. Hou and Max Tegmark},
      year={2024},
      eprint={2404.19756},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}