MLP的代替：KAN

受柯尔莫哥洛夫-阿诺德表示定理的启发，作者提出柯尔莫哥洛夫-阿诺德网络（KAN）作为多层感知器（MLP）有前途的替代品。MLP 在节点（“神经元”）上具有固定的激活函数，而 KAN 在边（“权重”）上具有可学习的激活函数。KAN 没有线性权重----每个权重参数都被参数化为spline的单变量函数所取代。作者证明，这种看似简单的改变使得 KAN 在准确性和可解释性方面优于 MLP。就准确性而言，在数据拟合和 PDE 求解中，较小的 KAN 可以比较大的 MLP 获得可比或更好的准确性。从理论上和经验上来说，KAN 比 MLP 拥有更快的神经尺度法则（neural scaling laws：随着模型参数的增加，测试损失减小）。对于可解释性，KAN 可以直观地可视化，并且可以轻松地与人类交互。通过数学和物理领域的两个例子，KAN 被证明是有用的“合作者”，帮助科学家（重新）发现数学和物理定律。总之，KAN 是 MLP 的有前途的替代品，为进一步改进当今严重依赖 MLP 的深度学习模型提供了机会。

来自：KAN: Kolmogorov–Arnold Networks

引言

多层感知器 MLP，也称为全连接前馈神经网络，是当今深度学习模型的基础构建块。MLP 的重要性怎么强调都不为过，因为它们是机器学习中用于逼近非线性函数的默认模型，MLP的表达能力由通用逼近定理保证。然而，MLP 是我们可以构建的最好的非线性回归器吗？尽管 MLP 被广泛使用，但它们也有明显的缺点。 例如，在 Transformer 中，MLP 占用几乎所有参数，并且在没有后期分析工具的情况下难以解释。

作者提出了一种有希望的MLP替代方案，称为Kolmogorov-Arnold网络 KANs。MLP受到通用近似定理的启发，而KAN则受到Kolmogorov-Arnold表示定理的启发。与MLP一样，KAN具有全连接的结构。然而，MLP将固定的激活函数放在节点（“神经元”）上，而KAN将可学习的激活函数放在边（“权重”）上，如图0.1所示。因此，KAN根本没有线性权重矩阵：取而代之的是，每个权重参数都被一个可学习的一维函数参数化为spline所取代。KAN的节点只是简单地对输入信号求和，而不应用任何非线性。

• 图0.1：MLP与KAN的比较。

有人可能会担心，由于每个MLP的权重参数都变成了KAN的spline，因此KAN的成本非常昂贵。幸运的是，KAN拥有比MLP更小的规模。比如：对于PDE solving，2层，每层宽度10的KAN比4层，每层宽度100的MLP精确100倍（$10^{-7}$ vs $10^{-5}$ MSE），参数效率高100倍（$10^2$ vs $10^4$参数）。

其实之前已经有很多研究使用Kolmogorov-Arnold表示定理构建神经网络。然而，大多数工作都坚持使用原始的depth-2 width-(2$n$+1)，并且没有利用更现代的技术，例如反向传播来训练网络。KAN的贡献在于将原始的Kolmogorov-Arnold表示推广到任意宽度和深度，并在今天的深度学习世界中重新激活它，以及使用广泛的实验来突出其作为AI+Science基础模型的优势，因为它的准确性和可解释性。

KAN是spline和MLP的组合，利用各自的优势并避免各自的弱点。spline对于低维函数是精确的，易于局部调整，并且能够在不同的分辨率之间切换。然而，spline曲线由于无法利用组合结构而存在严重的COD（维数诅咒）。另一方面，MLP由于其特征学习而较少受到COD的影响，但由于其无法优化单变量函数，因此在低维情况下不如spline曲线准确。

为了准确地学习一个函数，一个模型不仅要学习组成结构（外部自由度），而且要很好地近似单变量函数（内部自由度）。KAN是这样的模型，因为它们在外部有MLP，在内部有spline。因此，KAN不仅可以学习特征（由于它们与MLP的外部相似性），而且还可以以很高的精度优化这些学习到的特征（由于它们与spline的内部相似性）。比如，给定一个高维函数：$$f(x_1,...,x_N)=exp(\frac{1}{N}\sum _{i=1}^{N}si{n}^2(x_i))$$当$N$较大时，由于COD的影响，spline失效；MLP可以潜在地学习广义加性结构，但它们对于用ReLU激活来近似指数函数和正弦函数是非常低效的。相比之下，KANs可以很好地学习组合结构和单变量函数，因此在很大程度上优于MLP（见图3.1）。

• 图3.1：用5个toy例子来比较KAN和MLP。KAN 几乎可以实现理论预测的最快缩放定律，而 MLP 的缩放定律很缓慢。

KAN

MLP的灵感来自于通用近似定理。本节组织形式为：

• 回顾Kolmogorov-Arnold定理
• Kolmogorov-Arnold网络的设计
• 提出简化技术以使KANs可解释

Kolmogorov-Arnold表示理论

Vladimir Arnold和Andrey Kolmogorov证明了如果$f$是一个有界域上的多元连续函数，那么$f$可以写成单变量连续函数的有限复合，更具体的，对于函数$f:[0,1{]}^n\to \mathbb{R}$：$$f(\text{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}(2.1)}$$其中，${\varphi }_{q,p}:[0,1]\to \mathbb{R}$，${\Phi }_q:\mathbb{R}\to \mathbb{R}$。从某种意义上说，他们证明了唯一真正的多元函数是加法，因为所有其他函数都可以用一元函数与求和来表示。有人可能会天真地认为这对机器学习来说是个好消息：学习一个高维函数可以归结为学习一个多项式数量的一维函数。然而，这些一维函数可能是非平滑的，甚至是分形的，因此在实践中可能无法学习。由于这种病态的行为，Kolmogorov-Arnold表示定理在机器学习中基本上被判了死刑，被认为理论上是正确的，但实际上毫无用处。

然而，作者对Kolmogorov-Arnold定理在机器学习中的有用性更为乐观：

• 首先，不需要拘谨于原始的Eq.(2.1)，它只有两层非线性和隐藏层中少量的项（2$n$+1）：作者将网络推广到任意的宽度和深度。
• 其次，科学和日常生活中的大多数函数通常是光滑的，并且具有稀疏的组成结构，这有助于光滑的Kolmogorov-Arnold表示。这里的哲学接近于物理学家的心态，他们通常更关心典型的情况，而不是最坏的情况。

KAN架构

假设我们有一个由输入输出对 { x i , y i } \left\{\textbf{x}_{i},y_{i}\right\} {xi​,yi​}组成的监督学习任务，我们想要找到一个$f$对所有data point有$y_i\approx f({\text{x}}_i)$。Eq.(2.1)暗示，如果我们能找到合适的单变量函数${\varphi }_{q,p}$和${\Phi }_q$，我们就完成了。这启发作者设计一个神经网络，显式参数化Eq.(2.1)。由于所有需要学习的函数都是单变量函数，我们可以将每个一维函数参数化为一条B-spline，以及具有局部B-spline基函数的可学习系数（见图2.2右)。

• 图2.2：左----流经网络的激活符号（从下往上前向计算）。右----激活函数被参数化为B-spline，它允许在粗粒度和细粒度网格之间切换。

现在我们有了一个KAN的原型，其计算图由Eq.(2.1)精确指定，如图0.1 b所示（输入维数$n$=2），表现为一个两层神经网络，激活函数放置在边缘而不是节点上（对节点进行简单求和），中间层宽度为 2$n$ + 1。

如前所述，这样的网络被认为太简单，无法在实践中任意地用光滑spline近似任何函数！因此，作者希望将KAN扩展到更宽和更深。目前还不清楚如何使KAN更深，因为Kolmogorov-Arnold表示对应于两层的KAN。然而，目前还没有一个“一般化”版本的定理对应于更深层次的KAN。

作者注意到MLP和KAN之间的类比时，突破出现了。在MLP中，一旦定义了一个层（它由线性变换和非线性组成），我们就可以堆叠更多的层来使网络更深。为了建立更深的KAN，我们应该首先回答：什么是KAN layer？事实证明，具有$n_{in}$-dim输入和$n_{out}$-dim输出的KAN层可以定义为一维函数的矩阵： Φ = { ϕ q , p } ,    p = 1 , 2 , . . . , n i n ,    q = 1 , 2 , . . . , n o u t (2.2) \Phi=\left\{\phi_{q,p}\right\},\thinspace\thinspace p=1,2,...,n_{in},\thinspace\thinspace q=1,2,...,n_{out}\tag{2.2} Φ={ϕq,p​},p=1,2,...,nin​,q=1,2,...,nout​(2.2)其中${\varphi }_{q,p}$有可学习参数，在Kolmogov-Arnold定理中，内部函数形成一个$n_{in}=n,n_{out}=2n+1$的KAN层，外部函数形成一个$n_{in}=2n+1,n_{out}=n$的KAN层。因此，Eq.(2.1)中的Kolmogov-Arnold表示只是两个KAN层的简单组合。现在很清楚拥有更深的Kolmogorov-Arnold表示意味着什么：简单地堆叠更多的KAN层！

作者引入一些符号。这段话有点技术性，但我们可以参考图2.2-左来获得具体的例子和直观的理解。KAN的形状由整数数组表示：$$[n_0,n_1,...,n_L]$$其中，$n_i$是计算图第$i$层的节点数量。用$(l,i)$表示第$l$层的第$i$个神经元，用$x_{l,i}$表示$(l,i)$神经元的激活值。在第$l$层和第$l+1$层之间，有$n_l{n}_{l+1}$个激活函数：连接$(l,j)$和$(l+1,i)$的激活函数表示为$${\varphi }_{l,i,j},l=0,...,L-1,i=1,...,n_{l+1},j=1,...,n_l$$其中，${\varphi }_{l,i,j}$的激活之前是$x_{l,i}$，${\varphi }_{l,i,j}$的激活之后被记为${\tilde{x}}_{l,i,j}={\varphi }_{l,i,j}(x_{l,i})$。$(l+1,j)$神经元的激活值就是所有传入的后激活值之和：$$x_{l+1,j}=\sum _{i=1}^{n_l}{\tilde{x}}_{l,i,j}=\sum _{i=1}^{n_l}{\varphi }_{l,i,j}(x_{l,i}),j=1,...,n_{l+1}\text{\hspace{1em}(2.5)}$$在矩阵形式中，有：

其中，${\Phi }_l$是第$l$层KAN对应的函数矩阵。一般的KAN网络是$L$层的组合：给定一个输入向量${\text{x}}_0\in {\mathbb{R}}^{n_0}$，KAN的输出为：$$KAN(\text{x})=({\Phi }_{L-1}\circ {\Phi }_{L-2}\circ \cdot \cdot \cdot \circ {\Phi }_1\circ {\Phi }_0)\text{x}$$也可以重写上述方程，假设输出维数$n_L=1$，则有$f(\text{x})=KAN(\text{x})$：

原始Kolmogorov-Arnold表示Eq.(2.1)对应于形状为$[n,2n+1,1]$的2层KAN。所有的运算都是可微的，所以我们可以用反向传播来训练KAN。为了比较，一个MLP可以写成仿射变换$\text{W}$和非线性$\sigma$的交织：$$MLP(\text{x})=({\text{W}}_{L-1}\circ \sigma \circ {\text{W}}_{L-2}\circ \sigma \circ \cdot \cdot \cdot \circ {\text{W}}_1\circ \sigma \circ {\text{W}}_0)\text{x}$$很明显，MLP将线性变换和非线性分别处理为$\text{W}$和$\sigma$，而KAN在$\Phi$中将它们一起处理。在图0.1 ©和(d)中，可视化了三层MLP和三层KAN，以澄清它们的区别。

实现细节

1.残差激活函数
我们包含一个基函数$b(x)$（类似于残差连接），使得激活函数$\varphi (x)$是基函数$b(x)$和spline函数的和：$$\varphi (x)=w(b(x)+spline(x))$$设置$$b(x)=silu(x)=\frac{x}{1+e^{-x}}$$在大部分情况下，spline函数被参数化为B-spline的线性组合：$$spline(x)=\sum _i{c}_i{B}_i(x)$$其中，$c_i$是可学习的。原则上$w$是多余的，因为它可以被吸收到$b(x)$和$spline(x)$中。然而，我们仍然包括这个$w$因子，以更好地控制激活函数的总体大小。

2.初始化scales
每个激活函数初始化为$spline(x)\approx 0$

3.更新spline grids
作者根据输入激活动态更新每个网格。

化简KAN使其可解释

这个想法是从一个足够大的KAN开始，用稀疏性正则化训练它，然后进行修剪。作者将证明这些经过修剪的KAN比未经过修剪的KAN更易于解释。

1.稀疏化
对于MLP，线性权重的L1正则化用于支持稀疏性。KAN可以适应这个高层次的想法，但需要两个修改：

• 在KAN中没有线性的“权重”。线性权重被可学习的激活函数所取代，因此我们应该定义这些激活函数的L1范数。
• 实验发现L1不足以使KAN稀疏化；相反，一个额外的熵正则化是必要的。

定义激活函数$\varphi$的L1范数为在$N_p$输入上的平均幅度：$$|\varphi {|}_1=\frac{1}{N_p}\sum _{s=1}^{N_p}|\varphi (x^{(s)})|$$然后对于输入$n_{in}$和输出$n_{out}$的KAN layer $\Phi$，其L1范数为：$$|\Phi {|}_1=\sum _{i=1}^{n_{in}}\sum _{j=1}^{n_{out}}|{\varphi }_{i,j}{|}_1$$因此，训练目标为：$$l_{total}=l_{pred}+\lambda \sum _{l=0}^{L-1}|{\Phi }_l{|}_1$$

2.可视化
当可视化KAN时，将激活函数的透明度设置为与$tanh(\beta A_{l,i,j})$成比例，其中$\beta =3$。因此，贡献较小的函数会逐渐消失，让我们专注于重要的函数。

3.修剪
在使用稀疏化惩罚进行训练后，我们可能还希望将网络修剪成更小的子网。作者在节点级别（而不是在边缘级别）对KANs进行稀疏化。对于每个节点（假设是第1层的第$i$个神经元），作者将其传入和传出的分数定义为：$$I_{l,i}=ma{x}_k(|{\varphi }_{l-1,k,i}{|}_1),O_{l,i}=ma{x}_j(|{\varphi }_{l+1,j,i}{|}_1)$$如果传入和传出的分数都大于阈值超参数$\theta =10^{-2}$，则认为节点是重要的。不重要的节点被修剪。

4.符号化
在我们怀疑某些激活函数实际上是符号的情况下，例如，cos或log，作者提供了一个接口来将它们设置为指定的符号形式。

思考

KAN的设计看起来更像是为了求解物理公式，尤其是可解释性符号化部分的内容。在修剪的时候，我们可以发现保留的节点是随着输入动态改变的。尽管我一开始总在想KAN相比MLP有什么缺陷，但是现在我真的觉得KAN是代替MLP的一个未来方案。