科尔莫戈洛夫-阿诺德网络的高效实现

本仓库包含科尔莫戈洛夫-阿诺德网络（KAN）的高效实现。 KAN 的原始实现可在此处获取：https://github.com/KindXiaoming/pykan。

原始实现的性能问题主要在于，它需要将所有中间变量展开，以执行不同的激活函数。对于一个具有 in_features 个输入和 out_features 个输出的层，原始实现会将输入扩展为形状为 (batch_size, out_features, in_features) 的张量，以便应用激活函数。然而，所有激活函数实际上都是由一组固定的 B 样条基函数进行线性组合得到的；因此，我们可以重新表述计算过程：先用不同的基函数对输入进行激活，再将其线性组合起来。这种重写方式能够显著降低内存开销，并将计算简化为一次普通的矩阵乘法，同时自然地适用于前向传播和反向传播。

问题在于所谓的“稀疏化”操作，该操作被认为对 KAN 的可解释性至关重要。作者提出了一种基于输入样本定义的 L1 正则化，这种正则化需要对形状为 (batch_size, out_features, in_features) 的张量进行非线性操作，因而与上述重写方式不兼容。为此，我改用权重上的 L1 正则化，这在神经网络中更为常见，且与重写后的计算方式完全兼容。事实上，作者的实现中也同时包含了论文中描述的正则化以及这种权重正则化，因此我认为后者可能更有帮助。不过，仍需更多实验来验证这一点；至少可以确定的是，若追求效率，原始方法是不可行的。

另一个不同之处在于，除了可学习的激活函数（B 样条）之外，原始实现还为每个激活函数引入了一个可学习的缩放参数。我在代码中提供了一个名为 enable_standalone_scale_spline 的选项，默认值为 True，用于启用这一功能；将其禁用虽然能提升模型效率，但可能会对结果产生负面影响。这方面还需要进一步的实验研究。

2024年5月4日更新：@xiaol 指出，在 MNIST 数据集上，base_weight 参数采用常数初始化可能会带来问题。目前，我已将 base_weight 和 spline_scaler 两个矩阵的初始化方式更改为与 nn.Linear 一致的 kaiming_uniform_ 初始化方法。这样做后，模型在 MNIST 上的表现有了显著提升（准确率从约 20% 提升至约 97%），但尚不确定这是否是一种通用的好做法。

efficient-kan 快速上手指南

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.8 及以上
• 核心依赖：
  • PyTorch (推荐最新稳定版)
  • NumPy

请确保已安装 CUDA（如需 GPU 加速）并配置好相应的 PyTorch GPU 版本。

安装步骤

您可以直接通过 pip 从源码或 PyPI 安装（如果已发布）。推荐使用国内镜像源以加快下载速度。

使用清华镜像源安装：

pip install efficient-kan -i https://pypi.tuna.tsinghua.edu.cn/simple

或者从 GitHub 源码安装（获取最新版本）：

pip install git+https://github.com/Blealtan/efficient-kan.git -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

efficient-kan 的 API 设计与 PyTorch 原生层高度兼容，可以轻松替换现有的 nn.Linear 层。

1. 导入与模型定义

最简单的用法是直接实例化 KANLinear 层。以下是一个构建简单分类网络的示例：

import torch
import torch.nn as nn
from efficient_kan import KANLinear

class SimpleKANNet(nn.Module):
    def __init__(self):
        super(SimpleKANNet, self).__init__()
        # 定义一个 KAN 层：输入 784 (如 MNIST), 输出 10
        # grid_size: B-spline 的网格数量
        # spline_order: B-spline 的阶数
        self.kan_layer = KANLinear(
            in_features=784,
            out_features=10,
            grid_size=5,
            spline_order=3,
            enable_standalone_scale_spline=True, # 默认为 True，禁用可提升效率但可能影响精度
        )
        
    def forward(self, x):
        x = x.view(x.size(0), -1) # 展平输入
        return self.kan_layer(x)

# 实例化模型
model = SimpleKANNet()
print(model)

2. 训练循环示例

该层完全支持 PyTorch 的标准前向传播和反向传播流程：

# 模拟输入数据 (batch_size=32, features=784)
x = torch.randn(32, 784)
target = torch.randint(0, 10, (32,))

# 前向传播
output = model(x)

# 计算损失
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)

# 反向传播
loss.backward()

# 优化器步骤 (建议使用 Adam 或 SGD)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.step()
optimizer.zero_grad()

print(f"Loss: {loss.item()}")

关键参数说明

• enable_standalone_scale_spline: 默认为 True，保留原始 KAN 中每个激活函数的可学习缩放因子。若设为 False 可提高计算效率，但需验证对具体任务精度的影响。
• 初始化更新：当前版本已将 base_weight 和 spline_scaler 改为使用 kaiming_uniform_ 初始化，这在 MNIST 等数据集上显著提升了收敛效果（从约 20% 提升至 97%）。