KAN（Kolmogorov-Arnold Networks）：用可学习激活函数替代MLP层的探索
各位同学，大家好。今天我们来聊聊一个最近比较火的神经网络架构——Kolmogorov-Arnold Networks (KANs)。它通过用可学习激活函数替代传统MLP（Multilayer Perceptron）层中的固定激活函数，在某些任务上展现出了令人惊喜的效果。我们将深入探讨KANs的原理、优势、局限性，以及如何在实际中应用它们。
1. KANs的理论基础：Kolmogorov-Arnold表示定理
KANs的设计灵感来源于Kolmogorov-Arnold表示定理。这个定理指出，任何多元连续函数都可以被表示成单变量连续函数的有限次叠加和复合。具体来说，对于一个函数 f(x₁, x₂, …, xₙ)，可以找到单变量函数 φᵢ 和 ψᵢⱼ，使得：
f(x₁, x₂, ..., xₙ) = Σᵢ ψᵢ ( Σⱼ φᵢⱼ(xⱼ) )这个定理表明，我们可以将一个复杂的多元函数分解成更简单的单变量函数的组合。KANs正是基于这个思想，尝试将MLP中的权重矩阵和固定激活函数替换为可学习的单变量激活函数。
2. KANs的结构：可学习激活函数的神经网络
与MLP不同，KANs不再使用固定的激活函数（如ReLU、Sigmoid等）。取而代之的是，每个连接都关联着一个可学习的单变量激活函数。
- 传统MLP层: 输入向量
x经过权重矩阵W，然后通过一个固定的激活函数σ，得到输出σ(Wx + b)。 - KAN层: 输入向量
x的每个分量xᵢ通过一个可学习的激活函数σᵢⱼ，得到σᵢⱼ(xᵢ)。然后将这些结果加权求和，得到输出。
更具体地说，一个KAN层可以表示为：
yᵢ = Σⱼ wᵢⱼ σᵢⱼ(xⱼ)其中：
xⱼ是输入向量x的第j个分量。σᵢⱼ(xⱼ)是连接输入节点j和输出节点i的可学习激活函数。wᵢⱼ是连接输入节点j和输出节点i的权重。yᵢ是输出向量y的第i个分量。
这种结构使得KANs能够更好地拟合复杂函数，因为激活函数不再是固定的，而是可以根据数据进行学习。
3. KANs的实现：B样条基函数
在实际应用中，可学习激活函数 σᵢⱼ(xⱼ) 通常使用B样条基函数来表示。B样条是一种分段多项式函数，具有良好的平滑性和逼近能力。
具体来说，我们可以将激活函数 σᵢⱼ(xⱼ) 表示为一组B样条基函数的线性组合：
σᵢⱼ(xⱼ) = Σₖ cᵢⱼₖ Bₖ(xⱼ)其中：
Bₖ(xⱼ)是第k个B样条基函数。cᵢⱼₖ是第k个B样条基函数的系数，这些系数是可学习的参数。
通过学习这些系数 cᵢⱼₖ，KANs可以学习到任意形状的激活函数，从而更好地拟合数据。
Python代码示例：使用PyTorch实现KAN层
import torch
import torch.nn as nn
import numpy as np
class BSpline(nn.Module):
def __init__(self, n_basis, degree=3, min_val=-1.0, max_val=1.0):
super().__init__()
self.n_basis = n_basis
self.degree = degree
self.min_val = min_val
self.max_val = max_val
self.delta = (max_val - min_val) / (n_basis - degree)
self.register_buffer("knots", torch.linspace(min_val - degree * self.delta, max_val + degree * self.delta, n_basis + degree + 1))
def forward(self, x):
"""
x: (batch_size, )
return: (batch_size, n_basis)
"""
x = x.unsqueeze(-1) # (batch_size, 1)
basis_funcs = torch.zeros(x.shape[0], self.n_basis, device=x.device)
for i in range(self.n_basis):
basis_funcs[:, i] = self.basis_func(x, i)
return basis_funcs
def basis_func(self, x, i):
"""
x: (batch_size, 1)
i: int, the index of the basis function
return: (batch_size, )
"""
if self.degree == 0:
return ((self.knots[i] <= x.squeeze(-1)) & (x.squeeze(-1) < self.knots[i+1])).float()
else:
coeff1 = (x - self.knots[i]) / (self.knots[i+self.degree] - self.knots[i] + 1e-8)
coeff2 = (self.knots[i+self.degree+1] - x) / (self.knots[i+self.degree+1] - self.knots[i+1] + 1e-8)
return coeff1.squeeze(-1) * self.basis_func(x, i, self.degree-1) + coeff2.squeeze(-1) * self.basis_func(x, i+1, self.degree-1)
class KANLayer(nn.Module):
def __init__(self, in_features, out_features, n_basis=20, bspline_degree=3):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.n_basis = n_basis
self.bspline_degree = bspline_degree
self.bsplines = nn.ModuleList([BSpline(n_basis=self.n_basis, degree=self.bspline_degree) for _ in range(in_features * out_features)])
self.weights = nn.Parameter(torch.randn(out_features, in_features))
self.coeffs = nn.Parameter(torch.randn(out_features, in_features, n_basis)) # Learnable coefficients for B-splines
def forward(self, x):
"""
x: (batch_size, in_features)
return: (batch_size, out_features)
"""
batch_size = x.size(0)
outputs = torch.zeros(batch_size, self.out_features, device=x.device)
for i in range(self.out_features):
for j in range(self.in_features):
bspline_index = i * self.in_features + j
basis_funcs = self.bsplines[bspline_index](x[:, j]) # (batch_size, n_basis)
activation = torch.sum(basis_funcs * self.coeffs[i, j], dim=1) # (batch_size, )
outputs[:, i] = outputs[:, i] + self.weights[i, j] * activation
return outputs
class KAN(nn.Module):
def __init__(self, in_features, hidden_features, out_features, n_basis=20, bspline_degree=3):
super().__init__()
self.kan1 = KANLayer(in_features, hidden_features, n_basis=n_basis, bspline_degree=bspline_degree)
self.kan2 = KANLayer(hidden_features, out_features, n_basis=n_basis, bspline_degree=bspline_degree)
def forward(self, x):
x = self.kan1(x)
x = self.kan2(x)
return x
# 示例用法
if __name__ == '__main__':
# 示例用法
input_size = 2
hidden_size = 5
output_size = 1
batch_size = 10
# 创建KAN模型
kan_model = KAN(input_size, hidden_size, output_size)
# 生成随机输入数据
input_data = torch.randn(batch_size, input_size)
# 前向传播
output_data = kan_model(input_data)
# 打印输出
print("Input shape:", input_data.shape)
print("Output shape:", output_data.shape)这段代码定义了一个BSpline类，用于计算B样条基函数的值。KANLayer类实现了KAN层，它使用BSpline类来表示可学习的激活函数。每个连接都有自己的B样条基函数和权重，这些参数都是可学习的。 KAN类将两个KANLayer连接在一起构成一个简单的KAN网络。
4. KANs的优势
- 更好的函数逼近能力: 由于使用了可学习的激活函数，KANs能够更好地逼近复杂的函数，特别是在高维空间中。
- 更高的计算效率（理论上）： KANs的作者认为，KANs可以用更少的参数达到与MLP相当的性能。这主要是因为KANs将参数集中在激活函数上，而不是权重矩阵上。更少的参数意味着更少的计算量，可以加速训练和推理。
- 更好的可解释性： 由于每个连接都有自己的激活函数，我们可以分析这些激活函数的形状，从而了解KANs是如何学习数据的。
- 潜在的泛化能力提升： KANs的学习过程可以看作是在学习输入特征的不同表达方式。这种学习方式可能提高模型的泛化能力，使其在未见过的数据上表现更好。
5. KANs的局限性
- 训练难度： 训练KANs可能比训练MLP更困难。由于激活函数也是可学习的，训练过程可能会更加不稳定。
- 计算复杂度： 尽管KANs的参数量可能比MLP少，但计算B样条基函数的值可能会增加计算复杂度。特别是在激活函数数量很多的情况下。
- 内存占用： 存储大量的激活函数可能会增加内存占用。
- 缺乏成熟的优化技巧： 相对MLP而言，KANs还比较新，缺乏成熟的优化技巧，这使得训练KANs更具挑战。
- 实际应用验证不足： 目前KANs主要是在一些特定的数学函数拟合任务上表现出色，在更广泛的实际应用中的效果还有待验证。
6. KANs的应用场景
- 函数逼近： KANs非常适合用于逼近复杂的数学函数。
- 科学计算： KANs可以用于解决科学计算中的一些问题，例如求解微分方程。
- 数据降维： KANs可以用于学习数据的低维表示。
- 强化学习： KANs可以用于学习策略函数和价值函数。
7. KANs vs. MLPs：对比分析
| 特性 | KANs | MLPs |
|---|---|---|
| 激活函数 | 可学习的单变量函数 (通常使用B样条) | 固定的非线性函数 (ReLU, Sigmoid, 等) |
| 连接 | 每个连接都有一个激活函数 | 连接之间共享权重矩阵 |
| 参数量 | 理论上更少 (但实际情况取决于网络结构和基函数数量) | 通常更多 |
| 计算复杂度 | 计算B样条基函数可能增加计算复杂度 | 矩阵乘法和简单的激活函数计算 |
| 可解释性 | 更高 (可以分析激活函数的形状) | 较低 |
| 训练难度 | 可能更高 (需要同时学习权重和激活函数) | 相对较低 |
| 成熟度 | 较低 (仍在发展中) | 较高 (有大量的研究和实践经验) |
| 适用场景 | 函数逼近、科学计算等 | 图像识别、自然语言处理等 |
8. 训练KANs的技巧
- 初始化： 合适的初始化对于KANs的训练至关重要。可以尝试使用一些特殊的初始化方法，例如将B样条系数初始化为零，或者使用预训练的MLP来初始化KANs。
- 正则化： 为了防止过拟合，可以使用一些正则化技术，例如L1正则化或L2正则化。
- 学习率调整： 可以使用一些自适应学习率调整算法，例如Adam或RMSprop。
- 激活函数平滑化： 为了保证训练的稳定性，可以对激活函数进行平滑化处理。
- Grid Search: 通过Grid Search找到最佳的B样条基函数的数量和度数。
9. 未来研究方向
- 更有效的激活函数表示方法： 除了B样条，还可以尝试使用其他函数来表示可学习的激活函数，例如傅里叶级数或小波变换。
- 更高效的训练算法： 开发更高效的训练算法，以降低KANs的训练难度。
- KANs的理论分析： 对KANs的理论性质进行更深入的研究，例如泛化能力和收敛性。
- KANs在实际应用中的探索： 将KANs应用于更多的实际问题，例如图像识别、自然语言处理等。
- 硬件加速： 针对KANs的特点，设计专门的硬件加速器，以提高KANs的计算效率。
10. 总结KANs的特点和潜在价值
KANs是一种很有前景的神经网络架构，它通过用可学习激活函数替代MLP层中的固定激活函数，在函数逼近等方面表现出了优异的性能。虽然KANs还存在一些局限性，例如训练难度和计算复杂度，但随着研究的深入和技术的进步，相信KANs会在未来得到更广泛的应用。KANs的出现为神经网络的设计提供了一种新的思路，也为解决一些复杂的机器学习问题带来了新的希望。
希望今天的讲座能够帮助大家更好地理解KANs，并激发大家对神经网络架构的思考。谢谢大家！