
引言
深度学习的基本构成单元 Multi-Layer Perceptron(MLP)自 1958 年 Frank Rosenblatt 提出感知机以来,60 多年来一直是神经网络架构的核心。MLP 的结构很清晰:每个神经元(节点)应用固定的激活函数(ReLU、SiLU 等),边(权重)则执行可学习的线性变换。这一简单却强大的结构,凭借 Universal Approximation Theorem 获得了理论上的正当性,又因为矩阵乘法针对 GPU 运算做了优化,在实践中也一直占据主导地位。
然而在 2024 年 4 月,MIT 的 Ziming Liu、Max Tegmark 等人发表的 "KAN: Kolmogorov-Arnold Networks" 论文,向这个持续了 60 年的范式提出了根本性的质疑:激活函数一定要固定在节点上吗?边(权重)一定要是单纯的标量乘法吗?KAN 把这两个假设都推翻了——把激活函数放在边而不是节点上,并将每个激活函数用 B 样条参数化,使其可学习。这一想法的数学依据,来自 1957 年 Andrey Kolmogorov 及其学生 Vladimir Arnold 证明的 Kolmogorov-Arnold 表示定理。
KAN 已被 ICLR 2025 接收,在函数拟合、求解偏微分方程(PDE)、符号回归(symbolic regression)等科学探索任务上,用 比 MLP 更少的参数达到了更高的精度。此外,由于学习到的激活函数可以被可视化并进行符号化解释,KAN 在 可解释性(interpretability) 方面也大幅领先于 MLP。
本文将全面剖析 KAN 论文——从数学基础、架构设计、B 样条激活函数的工作原理,到与 MLP 的系统性比较、PyTorch 实现代码、可视化方法、实际应用案例,以及当前存在的局限性。
Kolmogorov-Arnold 表示定理
定理的数学定义
KAN 的理论基础是 1957 年被证明的 Kolmogorov-Arnold 表示定理(Kolmogorov-Arnold Representation Theorem)。这一定理是对希尔伯特第 13 问题(Hilbert's 13th Problem)的解答,证明了任意多元连续函数都可以表示为单变量连续函数的复合。
用数学语言描述该定理如下。对有界域 上的任意连续函数 ,下式成立:
这里 是 内部函数(inner function), 是 外部函数(outer function)。关键在于,内部函数和外部函数都是 单变量(univariate)连续函数。也就是说,无论多元函数多么复杂,都可以仅用单变量函数的复合与加法来精确表示。
与 MLP 的 Universal Approximation Theorem 的区别
MLP 的理论基础 Universal Approximation Theorem 保证的是 近似(approximation),而 Kolmogorov-Arnold 表示定理保证的是 精确表示(exact representation)。两个定理的比较如下:
| 项目 | Universal Approximation(MLP) | Kolmogorov-Arnold(KAN) |
|---|---|---|
| 保证程度 | 任意精度的近似 | 精确表示 |
| 结构 | 固定激活函数 + 可学习权重 | 可学习的单变量函数复合 |
| 神经元数量 | 可以无限多 | 个内部求和即可 |
| 激活函数位置 | 节点(神经元) | 边(权重) |
| 可解释性 | 黑箱 | 可对单变量函数进行可视化 |
从神经网络视角的再诠释
从神经网络的视角重新诠释 Kolmogorov-Arnold 表示定理,可以把它看作一个拥有 2 个隐藏层(hidden layer)的网络。第一个隐藏层对每个输入维度应用 个单变量函数 ,第二个隐藏层则对内部求和应用 个外部函数 。
然而,原始定理中的 与 可能并不光滑,甚至具有分形结构,因此难以直接用于实际训练。KAN 论文的核心贡献,是把这一定理 推广为任意深度的网络,并通过将每个单变量函数 用 B 样条参数化,使其真正具备了实用性。
KAN 架构详细分析
MLP 与 KAN 的结构性差异
MLP 与 KAN 的根本差异在于,激活函数与权重的角色发生了对调。
MLP:在边上执行线性变换(),在节点上应用非线性激活
KAN:在边上应用可学习的非线性函数 ,节点上只做简单的求和
这里 是放置在「从第 个输入通往第 个输出」这条边上的可学习单变量函数。
KAN 层架构图
下面这张图展示了 KAN 层的结构。关键在于,每条边都不是单纯的标量权重,而是一个可学习的 B 样条函数。
KAN Layer: [2, 3] (2 inputs -> 3 outputs)
Input Layer Learnable Activations (Edges) Output Layer
(B-spline functions)
+------ phi_1,1(x) --------+
| |
x_1 -------+------ phi_2,1(x) --------+-----> y_1 = phi_1,1(x_1) + phi_1,2(x_2)
| |
+------ phi_3,1(x) --------+
| |
| +-----> y_2 = phi_2,1(x_1) + phi_2,2(x_2)
| |
x_2 -------+------ phi_1,2(x) --------+
| |
+------ phi_2,2(x) --------+-----> y_3 = phi_3,1(x_1) + phi_3,2(x_2)
| |
+------ phi_3,2(x) --------+
[Nodes: SUM only] [Edges: Learnable] [Nodes: SUM only]
Total learnable functions: n_in x n_out = 2 x 3 = 6
Each phi_{j,i} is a B-spline with (G + k) parameters
where G = grid intervals, k = spline order
多层 KAN 网络
Kolmogorov-Arnold 表示定理的原始结构被限制在 2 层,而 KAN 论文将其扩展到 任意深度。在深度为 的 KAN 网络中,每一层 拥有 个输入和 个输出,整个网络可以表示为:
这里每个 都是一个 KAN 层,层 中共有 个可学习的单变量函数。整个网络的可学习函数总数为:
KAN 的参数量计算
如果每个单变量函数用 个网格区间和 阶 B 样条来参数化,那么每个函数需要 个系数。因此,整个网络的参数量为:
举例来说,对于宽度(width)为 、、 的 KAN:
- 第 0 层: 个参数
- 第 1 层: 个参数
- 总参数量: 个
相同结构的 MLP 有 个参数。虽然 KAN 的参数更多,但论文主张,就 参数与精度之比 而言,KAN 全面压制了 MLP。
B 样条激活函数
B 样条的定义
B 样条(Basis Spline)是由分段多项式(piecewise polynomial)线性组合而定义出的光滑函数。 阶 B 样条基函数由 De Boor 递归公式定义:
这里 是 节点(knot) 向量的元素。使用 个网格区间和 阶样条,会生成 个 B 样条基函数,样条函数则表示为这些基函数的线性组合:
这里 是可学习的系数。
KAN 中激活函数的结构
放置在 KAN 每条边上的激活函数 由两部分组成:
这里 是 残差函数(residual function),默认取 , 与 分别是对应的缩放系数。残差函数在训练初期提供稳定性,帮助样条部分逐步学习更复杂的模式。
B 样条基函数的实现
下面是用 PyTorch 实现 B 样条基函数的代码:
import torch
import torch.nn as nn
import numpy as np
def compute_bspline_basis(x: torch.Tensor, grid: torch.Tensor, k: int) -> torch.Tensor:
"""用 De Boor 递归公式计算 B 样条基函数。
Args:
x: 输入张量 [batch_size, in_features]
grid: 节点(knot)向量 [in_features, grid_size + 2k + 1]
k: 样条阶数(通常为 3,即三次样条)
Returns:
基函数值 [batch_size, in_features, grid_size + k]
"""
# 扩展 x 的维度,以便与 grid 比较
# x: [batch, in_features] -> [batch, in_features, 1]
x = x.unsqueeze(-1)
# 0 阶基函数:落在区间内为 1,否则为 0
bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).float()
# De Boor 递归:重复 k 次以计算 k 阶基函数
for degree in range(1, k + 1):
# 左侧项: (x - t_i) / (t_{i+degree} - t_i)
left_num = x - grid[:, :-(degree + 1)]
left_den = grid[:, degree:-1] - grid[:, :-(degree + 1)]
left = left_num / left_den.clamp(min=1e-8)
# 右侧项: (t_{i+degree+1} - x) / (t_{i+degree+1} - t_{i+1})
right_num = grid[:, (degree + 1):] - x
right_den = grid[:, (degree + 1):] - grid[:, 1:(-degree)]
right = right_num / right_den.clamp(min=1e-8)
# 递归组合
bases = left * bases[:, :, :-1] + right * bases[:, :, 1:]
return bases # [batch, in_features, grid_size + k]
# 使用示例
batch_size, in_features = 32, 4
grid_size, k = 5, 3 # 5 个网格区间,三次样条
# 生成均匀网格:在 [-1, 1] 范围内两侧各扩展 k 个
grid_points = torch.linspace(-1, 1, grid_size + 1)
extended_grid = torch.cat([
grid_points[0:1] - torch.arange(k, 0, -1) * (grid_points[1] - grid_points[0]),
grid_points,
grid_points[-1:] + torch.arange(1, k + 1) * (grid_points[1] - grid_points[0])
])
# 复制到 in_features 维度
grid = extended_grid.unsqueeze(0).repeat(in_features, 1)
x = torch.randn(batch_size, in_features)
basis_values = compute_bspline_basis(x, grid, k)
print(f"基函数 shape: {basis_values.shape}")
# 输出: 基函数 shape: torch.Size([32, 4, 8]) # G + k = 5 + 3 = 8
网格更新机制
KAN 的一个重要特点是,训练过程中可以 动态更新网格。初期使用均匀网格,但如果按训练数据的分布重新排布网格,就能大幅提升近似精度。grid_eps 超参数负责控制这一点:
grid_eps = 1.0:完全均匀的网格grid_eps = 0.0:按数据样本的分位数(quantile)排布的网格0 < grid_eps < 1:两个极端之间的插值
逐步增加网格区间数 (网格扩展,grid extension)可以提高分辨率,从而表示更复杂的函数。这与在 MLP 中增加神经元数量类似,但优点在于,能够在保留已有训练结果的同时,只提升分辨率。
KAN 与 MLP 的比较分析
系统性比较表
| 比较项 | MLP | KAN |
|---|---|---|
| 激活函数位置 | 节点(固定:ReLU、SiLU 等) | 边(可学习:B 样条) |
| 权重的角色 | 可学习的线性变换 | 无(由激活函数取代) |
| 参数效率 | 低(需要加宽才能提升精度) | 高(靠网格分辨率提升精度) |
| 缩放规律 | , 较小 | , 更大 |
| 可解释性 | 黑箱 | 可对激活函数进行可视化 |
| 符号回归 | 不可 | 训练后可提取符号函数 |
| 训练速度 | 快(GPU 矩阵乘法已优化) | 慢(样条计算带来额外开销) |
| GPU 并行化 | 出色(批量矩阵乘法) | 受限(每条边的函数各不相同) |
| 高维输入 | 出色(图像、NLP 等) | 受限(受维度灾难影响) |
| 主要应用领域 | 通用深度学习 | 科学发现、符号回归 |
精度缩放比较
KAN 论文中最令人印象深刻的结果之一,是在 神经缩放规律(Neural Scaling Law)上表现出的差异。将测试 RMSE 表示为参数量 的函数时:
对 MLP 而言, 的上限由 Universal Approximation Theorem 决定;而 KAN 基于 Kolmogorov-Arnold 表示定理,展现出更快的缩放速度(更大的 )。在论文针对 5 个特殊函数的实验中,KAN [2,5,1] 用 80 个参数,比 MLP [2,100,1] 的 201 个参数少 2-3 倍,却达到了相当甚至更好的精度。
训练速度与现实考量
KAN 的训练速度比 MLP 慢得多,主要原因如下:
- 样条基函数计算:每条边都需要递归计算 个基函数。
- GPU 并行化的局限:MLP 可以通过权重矩阵的一次矩阵乘法完成计算,而 KAN 在每条边上应用不同的函数,因此难以有效利用批量矩阵乘法。
- 网格更新开销:训练过程中的网格重新排布会带来额外的计算成本。
根据 "KAN or MLP: A Fairer Comparison"(Yu et al., 2024)论文,在大多数标准机器学习、计算机视觉、NLP、音频处理任务上,MLP 都表现出了更高的平均精度。KAN 展现出优势的,仅限于符号公式表示(symbolic formula representation)等特定科学任务。
实现代码
KAN 层的 PyTorch 实现
下面是用 PyTorch 从零(from scratch)实现 KAN 层的代码:
import torch
import torch.nn as nn
import torch.nn.functional as F
class KANLayer(nn.Module):
"""Kolmogorov-Arnold Network 层。
在每条边上放置可学习的 B 样条激活函数。
节点上只执行简单的求和。
"""
def __init__(
self,
in_features: int,
out_features: int,
grid_size: int = 5,
spline_order: int = 3,
scale_base: float = 1.0,
scale_spline: float = 1.0,
grid_range: tuple = (-1.0, 1.0),
):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.grid_size = grid_size
self.spline_order = spline_order
self.scale_base = scale_base
self.scale_spline = scale_spline
# 生成并扩展均匀网格(两侧各扩展 spline_order 个)
h = (grid_range[1] - grid_range[0]) / grid_size
grid = torch.arange(-spline_order, grid_size + spline_order + 1) * h + grid_range[0]
# grid shape: [grid_size + 2 * spline_order + 1]
# 扩展维度,使每个 (in, out) 组合都可以拥有独立的网格
self.register_buffer(
"grid",
grid.unsqueeze(0).unsqueeze(0).expand(out_features, in_features, -1)
)
# 样条系数:每条边 (grid_size + spline_order) 个
self.spline_weight = nn.Parameter(
torch.randn(out_features, in_features, grid_size + spline_order)
* 0.1
)
# 残差(base)权重
self.base_weight = nn.Parameter(
torch.randn(out_features, in_features) * 0.1
)
def b_splines(self, x: torch.Tensor) -> torch.Tensor:
"""计算 B 样条基函数。
Args:
x: [batch_size, in_features]
Returns:
基函数值 [batch_size, in_features, grid_size + spline_order]
"""
# 将 x 扩展为 [batch, 1, in_features, 1],以便与 grid 比较
x = x.unsqueeze(1).unsqueeze(-1) # [B, 1, in, 1]
# grid: [out, in, G+2k+1] -> 便于广播
grid = self.grid # [out, in, G+2k+1]
# 0 阶基函数:落在区间内为 1
# 使用 grid[:, :, :-1] 与 grid[:, :, 1:]
# 只使用第一个 out 维度的网格(因为都相同)
g = grid[0] # [in, G+2k+1]
bases = ((x[:, 0] >= g[:, :-1]) & (x[:, 0] < g[:, 1:])).float()
# bases: [B, in, G+2k]
for degree in range(1, self.spline_order + 1):
left_num = x[:, 0] - g[:, :-(degree + 1)]
left_den = g[:, degree:-1] - g[:, :-(degree + 1)]
left = left_num / left_den.clamp(min=1e-8)
right_num = g[:, (degree + 1):] - x[:, 0]
right_den = g[:, (degree + 1):] - g[:, 1:(-degree)]
right = right_num / right_den.clamp(min=1e-8)
bases = left * bases[:, :, :-1] + right * bases[:, :, 1:]
return bases # [B, in, G+k]
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""前向传播:残差函数 + 样条函数。
Args:
x: [batch_size, in_features]
Returns:
输出 [batch_size, out_features]
"""
batch_size = x.size(0)
# 1) 残差部分:SiLU 激活后做线性组合
base_output = F.silu(x) # [B, in]
base_output = torch.einsum("bi,oi->bo", base_output, self.base_weight)
# base_output: [B, out]
# 2) 样条部分:B 样条基函数 * 系数
spline_basis = self.b_splines(x) # [B, in, G+k]
spline_output = torch.einsum(
"big,oig->bo", spline_basis, self.spline_weight
)
# spline_output: [B, out]
# 3) 合并
return self.scale_base * base_output + self.scale_spline * spline_output
完整的 KAN 网络实现
class KAN(nn.Module):
"""多层 Kolmogorov-Arnold Network。
Args:
width: 各层神经元数量的列表。例如:[2, 5, 1]
grid_size: B 样条网格区间数
spline_order: B 样条阶数(默认值:3,即三次样条)
"""
def __init__(
self,
width: list,
grid_size: int = 5,
spline_order: int = 3,
):
super().__init__()
self.width = width
self.layers = nn.ModuleList()
for i in range(len(width) - 1):
self.layers.append(
KANLayer(
in_features=width[i],
out_features=width[i + 1],
grid_size=grid_size,
spline_order=spline_order,
)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
for layer in self.layers:
x = layer(x)
return x
def parameter_count(self) -> dict:
"""返回各层的参数量。"""
info = {}
total = 0
for i, layer in enumerate(self.layers):
layer_params = sum(p.numel() for p in layer.parameters())
info[f"layer_{i}"] = layer_params
total += layer_params
info["total"] = total
return info
# 使用示例
model = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
print(f"模型结构: {model.width}")
print(f"参数量: {model.parameter_count()}")
# 前向传播测试
x = torch.randn(16, 2)
y = model(x)
print(f"输入: {x.shape}, 输出: {y.shape}")
# 输出示例:
# 模型结构: [2, 5, 1]
# 参数量: {'layer_0': 90, 'layer_1': 45, 'total': 135}
# 输入: torch.Size([16, 2]), 输出: torch.Size([16, 1])
用于对比的 MLP 实现
实现一个相同结构的 MLP,以进行公平的比较:
class MLP(nn.Module):
"""用于对比的 Multi-Layer Perceptron。
使用与 KAN 相同的 width 结构,
但采用固定激活函数(SiLU)和线性权重。
"""
def __init__(self, width: list):
super().__init__()
layers = []
for i in range(len(width) - 1):
layers.append(nn.Linear(width[i], width[i + 1]))
if i < len(width) - 2: # 排除最后一层
layers.append(nn.SiLU())
self.network = nn.Sequential(*layers)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.network(x)
# MLP vs KAN 参数量比较
mlp = MLP(width=[2, 5, 1])
kan = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
mlp_params = sum(p.numel() for p in mlp.parameters())
kan_params = sum(p.numel() for p in kan.parameters())
print(f"MLP 参数量: {mlp_params}")
print(f"KAN 参数量: {kan_params}")
print(f"KAN/MLP 参数量比率: {kan_params / mlp_params:.1f}x")
# 输出示例:
# MLP 参数量: 21
# KAN 参数量: 135
# KAN/MLP 参数量比率: 6.4x
训练与可视化
训练循环的实现
KAN 论文推荐使用 LBFGS 优化器,但在大规模数据集上也可以使用 Adam。下面是用于函数近似任务的完整训练循环:
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
def create_dataset(f, n_train=1000, n_test=200, input_dim=2):
"""生成训练/测试数据集。
Args:
f: 待近似的目标函数
n_train: 训练样本数
n_test: 测试样本数
input_dim: 输入维度
"""
x_train = torch.rand(n_train, input_dim) * 2 - 1 # [-1, 1]
y_train = f(x_train)
x_test = torch.rand(n_test, input_dim) * 2 - 1
y_test = f(x_test)
return x_train, y_train, x_test, y_test
def train_model(model, x_train, y_train, x_test, y_test,
epochs=500, lr=1e-2, optimizer_type="adam"):
"""训练 KAN 或 MLP 模型。
Args:
model: KAN 或 MLP 实例
optimizer_type: 'adam' 或 'lbfgs'
Returns:
训练历史字典
"""
if optimizer_type == "lbfgs":
optimizer = torch.optim.LBFGS(
model.parameters(), lr=lr, max_iter=20,
history_size=10, line_search_fn="strong_wolfe"
)
else:
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
history = {"train_loss": [], "test_loss": []}
for epoch in range(epochs):
model.train()
if optimizer_type == "lbfgs":
def closure():
optimizer.zero_grad()
pred = model(x_train)
loss = criterion(pred, y_train)
loss.backward()
return loss
loss = optimizer.step(closure)
else:
optimizer.zero_grad()
pred = model(x_train)
loss = criterion(pred, y_train)
loss.backward()
optimizer.step()
# 评估
model.eval()
with torch.no_grad():
train_loss = criterion(model(x_train), y_train).item()
test_loss = criterion(model(x_test), y_test).item()
history["train_loss"].append(train_loss)
history["test_loss"].append(test_loss)
if (epoch + 1) % 100 == 0:
print(f"Epoch {epoch+1}/{epochs} | "
f"Train MSE: {train_loss:.6f} | "
f"Test MSE: {test_loss:.6f}")
return history
# 目标函数: f(x1, x2) = exp(sin(pi * x1) + x2^2)
target_fn = lambda x: torch.exp(
torch.sin(torch.pi * x[:, 0:1]) + x[:, 1:2] ** 2
)
x_train, y_train, x_test, y_test = create_dataset(target_fn)
# 训练 KAN
kan_model = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
kan_history = train_model(kan_model, x_train, y_train, x_test, y_test,
epochs=500, lr=1e-2, optimizer_type="adam")
# 训练 MLP(为了对齐参数量而加宽)
mlp_model = MLP(width=[2, 50, 50, 1])
mlp_history = train_model(mlp_model, x_train, y_train, x_test, y_test,
epochs=500, lr=1e-3, optimizer_type="adam")
# 可视化损失曲线对比
fig, ax = plt.subplots(1, 1, figsize=(10, 6))
ax.semilogy(kan_history["test_loss"], label="KAN [2,5,1]", linewidth=2)
ax.semilogy(mlp_history["test_loss"], label="MLP [2,50,50,1]", linewidth=2)
ax.set_xlabel("Epoch")
ax.set_ylabel("Test MSE (log scale)")
ax.set_title("KAN vs MLP: Function Approximation")
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("kan_vs_mlp_loss.png", dpi=150)
plt.show()
使用 pykan 库进行训练与可视化
在实际研究或实验中,使用官方 pykan 库会更加方便:
# pip install pykan
from kan import KAN as PyKAN
import torch
# 定义目标函数
target_fn = lambda x: torch.exp(
torch.sin(torch.pi * x[:, [0]]) + x[:, [1]] ** 2
)
# 生成数据集(使用 pykan 内置工具)
dataset = {
"train_input": torch.rand(1000, 2) * 2 - 1,
"test_input": torch.rand(200, 2) * 2 - 1,
}
dataset["train_label"] = target_fn(dataset["train_input"])
dataset["test_label"] = target_fn(dataset["test_input"])
# 创建并训练 KAN 模型
model = PyKAN(width=[2, 5, 1], grid=5, k=3, seed=42)
# 执行训练(使用 LBFGS 优化器)
results = model.fit(dataset, opt="LBFGS", steps=50, lamb=0.01)
# 可视化学习到的激活函数
# 以图形方式展示每条边学习到的 B 样条函数
model.plot()
# 符号回归:尝试从学习到的函数中提取数学公式
# 通过 model.auto_symbolic() 映射为 sin、exp 等符号函数
model.auto_symbolic()
formula = model.symbolic_formula()
print(f"提取出的公式: {formula}")
# 预期输出: exp(sin(pi*x1) + x2^2)
实际应用案例
1. 求解偏微分方程(PDE)
在物理学与工程学领域,KAN 正作为求解偏微分方程的神经网络(Physics-Informed Neural Networks, PINNs)的替代方案受到关注。论文中针对泊松方程 的实验表明,KAN 用比基于 MLP 的 PINN 少 100 倍的参数,达到了相同的精度。
应用领域:
- 流体力学模拟
- 热传导方程
- 波动方程
2. 符号回归(Symbolic Regression)
KAN 最强大的应用,是 从数据中发现数学公式 的符号回归。将学习到的激活函数可视化并与已知的符号函数(sin、cos、exp 等)进行匹配,就能提取出隐藏在数据中的数学关系。
KAN 2.0 论文(Liu et al., 2024)在此基础上进一步扩展,重新发现了以下物理定律:
- 守恒量(Conserved Quantities):力学系统的能量守恒定律
- 拉格朗日量(Lagrangians):由变分原理导出的运动方程
- 对称性(Symmetries):基于诺特定理的对称结构
- 本构定律(Constitutive Laws):材料的应力-应变关系
3. 纽结理论(Knot Theory)
在数学领域,KAN 同样给出了有意义的结果——成功重新发现了纽结不变量(knot invariants)之间的非平凡关系,这也表明 AI 在纯数学研究中的角色有望进一步扩大。
4. 凝聚态物理学
在识别相变(phase transition)边界的任务中,KAN 学习到的边界比 MLP 更精确,并且通过学习到的激活函数,还能够解释相变的物理意义。
局限性与未来方向
当前的局限性
1. 训练速度
KAN 的训练速度明显慢于 MLP。由于需要计算样条基函数,并且每条边都要应用不同的函数,这种结构特性使其无法充分利用 GPU 的批量矩阵乘法加速。在相同任务上,也有报告显示 KAN 的训练时间是 MLP 的 10 倍以上。
2. 高维输入的可扩展性
随着输入维度升高(例如图像的数万像素),KAN 的性能会下降。在 CIFAR-10 这类标准计算机视觉基准上,KAN 相较于 MLP 并未展现出明显优势,样条参数的调优也变得困难。
3. 对噪声的敏感性
KAN 在干净的函数型数据上表现出色,但在噪声较大的真实数据上,可能比 MLP 更敏感。这是因为样条函数存在对噪声过拟合(overfit)的倾向。
4. 内存占用
由于每条边都需要存储 个样条系数以及网格信息,与相同结构的 MLP 相比,内存占用相当高。在大规模网络中,这可能成为实际的瓶颈。
5. 标准任务上的表现
在机器学习、计算机视觉、NLP、音频处理等常规任务上,MLP 的平均精度普遍高于 KAN。KAN 真正展现优势的领域,仅限于符号公式表示、科学发现等特定方向。
未来研究方向
1. 高效的 KAN 实现
EfficientKAN、FastKAN、FourierKAN 等多种旨在提升 KAN 计算效率的变体已被提出。研究者们正在探索用切比雪夫多项式或傅里叶基替代 B 样条,能否在提升训练速度的同时,保留 KAN 的优势。
2. 混合架构
结合 MLP 与 KAN 各自优点的混合架构,是一个有前景的方向。例如,有研究尝试在 Transformer 的 FFN 层中用 KAN 替代 MLP,或者只在特定层中选择性地使用 KAN。
3. KAN 2.0 与科学发现
KAN 2.0 论文通过 MultKAN(引入乘法节点)、网格自动优化,以及在多个科学领域的应用,大幅扩展了 KAN 的实用性。这项发表于 Physical Review X 的后续研究表明,KAN 有潜力超越单纯的函数近似工具,成为一种 科学发现工具。
故障排查与优化技巧
1. 训练不收敛时
KAN 训练中最常见的问题是不收敛。可以检查以下几点:
# 检查网格范围是否覆盖了输入数据的范围
x_min, x_max = x_train.min().item(), x_train.max().item()
print(f"数据范围: [{x_min:.2f}, {x_max:.2f}]")
# 把网格范围设置得比数据范围略宽一些
# 可调整为 grid_range=(-1.5, 1.5) 等
# 使用 LBFGS 优化器时,必须使用 closure 模式
optimizer = torch.optim.LBFGS(
model.parameters(),
lr=0.1, # LBFGS 通常使用更大的 lr
max_iter=20, # 内部迭代次数
tolerance_grad=1e-7,
tolerance_change=1e-9,
line_search_fn="strong_wolfe" # 稳定的线搜索
)
# 通过调整 grid_eps 来适配网格
# 建议初期设为 0.5 左右,训练后期再降到 0.0
2. 防止过拟合
KAN 的表达能力非常强,因此容易过拟合:
- 正则化:pykan 中可以用
lamb参数施加 L1 正则化,也可以使用熵正则化(lamb_entropy) - 调整网格大小:初期从较小的 (例如 3)开始,等训练稳定后再逐步增大
- 剪枝(Pruning):训练结束后,移除不必要的边(激活函数接近 0 的那些)
3. 内存优化
# 使用梯度检查点(gradient checkpointing)以节省内存
from torch.utils.checkpoint import checkpoint
class MemoryEfficientKAN(nn.Module):
def __init__(self, width, grid_size=5, spline_order=3):
super().__init__()
self.layers = nn.ModuleList([
KANLayer(width[i], width[i+1], grid_size, spline_order)
for i in range(len(width) - 1)
])
def forward(self, x):
for layer in self.layers:
# 梯度检查点: 节省内存(反向传播时重新计算)
x = checkpoint(layer, x, use_reentrant=False)
return x
4. 通过可视化进行调试
KAN 的一大优势,是可以直接可视化学习到的激活函数。要确认训练是否正常进行,定期绘制每条边的激活函数会很有帮助:
import matplotlib.pyplot as plt
import torch
import numpy as np
def visualize_kan_activations(model, layer_idx=0, input_range=(-1, 1)):
"""可视化 KAN 层学习到的激活函数。
Args:
model: KAN 模型实例
layer_idx: 要可视化的层索引
input_range: 输入范围
"""
layer = model.layers[layer_idx]
n_in = layer.in_features
n_out = layer.out_features
# 生成用于可视化的输入
x_vis = torch.linspace(input_range[0], input_range[1], 200)
fig, axes = plt.subplots(n_out, n_in, figsize=(4 * n_in, 3 * n_out))
if n_out == 1:
axes = axes.reshape(1, -1)
if n_in == 1:
axes = axes.reshape(-1, 1)
for j in range(n_out):
for i in range(n_in):
ax = axes[j][i]
# 计算单个输入维度对应的激活函数
x_input = torch.zeros(200, n_in)
x_input[:, i] = x_vis
with torch.no_grad():
# 仅提取样条部分
basis = layer.b_splines(x_input) # [200, in, G+k]
spline_out = torch.einsum(
"big,g->bi", basis[:, i:i+1, :],
layer.spline_weight[j, i, :]
)
ax.plot(x_vis.numpy(), spline_out[:, 0].numpy(),
linewidth=2, color="blue")
ax.set_title(f"phi_{j+1},{i+1}(x)", fontsize=10)
ax.grid(True, alpha=0.3)
ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)
plt.suptitle(f"Layer {layer_idx} Activation Functions", fontsize=14)
plt.tight_layout()
plt.savefig(f"kan_layer{layer_idx}_activations.png", dpi=150)
plt.show()
# 训练后进行可视化
# visualize_kan_activations(kan_model, layer_idx=0)
5. 性能比较时的注意事项
要公平地比较 KAN 与 MLP,需要注意以下几点:
- 控制参数量:应该以相同的参数量而非相同的网络结构来比较
- 控制训练时间:以总训练时间(wall-clock time)而非 epoch 数来比较更符合实际
- 任务适配性:KAN 在低维函数近似上更有优势,MLP 在高维模式识别上更强,比较时应涵盖多种任务
- 优化器选择:KAN 往往在 LBFGS 下表现最佳,MLP 则在 Adam 下表现最佳,因此为每个模型使用各自最优的优化器才算公平
参考资料
-
KAN: Kolmogorov-Arnold Networks - Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljacic, Thomas Y. Hou, Max Tegmark (ICLR 2025) https://arxiv.org/abs/2404.19756
-
KAN 2.0: Kolmogorov-Arnold Networks Meet Science - Ziming Liu et al. (Physical Review X, 2024) https://arxiv.org/abs/2408.10205
-
pykan: Official Python Library for KAN - KindXiaoming (GitHub) https://github.com/KindXiaoming/pykan
-
KAN or MLP: A Fairer Comparison - Runpeng Yu, Weihao Yu, Xinchao Wang (2024) https://arxiv.org/abs/2407.16674
-
Kolmogorov-Arnold Representation Theorem - Wikipedia https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Arnold_representation_theorem
-
pykan Documentation - Kolmogorov Arnold Network Documentation https://kindxiaoming.github.io/pykan/
-
Kolmogorov-Arnold Networks: The Latest Advance in Neural Networks, Simply Explained - Towards Data Science https://towardsdatascience.com/kolmogorov-arnold-networks-the-latest-advance-in-neural-networks-simply-explained-f083cf994a85/
-
A Comprehensive and FAIR Comparison Between MLP and KAN - ScienceDirect (2024) https://www.sciencedirect.com/science/article/abs/pii/S0045782524005462
현재 단락 (1/495)
深度学习的基本构成单元 **Multi-Layer Perceptron**(MLP)自 1958 年 Frank Rosenblatt 提出感知机以来,60 多年来一直是神经网络架构的核心。MLP ...