KAN是否公平

KAN和MLP之间更公平和全面的比较

[论文]

摘要 (1) - 对于符号公式表示，使用KAN；对于计算机视觉、机器学习、文本处理和音频处理任务，使用MLP。

摘要 (2) - 通过对KAN和MLP的网络架构进行消融实验，发现KAN的主要优势在于使用B样条函数。将MLP中的激活函数替换为B样条函数后，MLP在KAN之前具有优势的数据集上也能超越KAN的表现。

目录

• 安装
• 运行
• 参数数量和浮点运算次数

安装

我们的代码在CUDA 12.2下进行测试。实验环境可以通过yml文件构建。

cd KANbeFair
conda env create -f environment.yml

运行

使用以下命令运行代码。

cd src

python train.py \
    --model KAN \ # 或MLP，用于指定模型架构
    --layers_width 10 10 \ # 隐藏层宽度。例如，对于MNIST上的MLP，这将导致Linear(28*28,10) -> Linear(10,10) -> Linear(10,10)的模型；对于MNIST上的KAN，这将导致kan.KAN(width = (28*28,10,10,10))
    --dataset MNIST \ 
    --batch-size 128 \
    --epochs 20 \
    --lr 0.001 \
    --seed 1314 \
    --activation_name gelu \ # 仅适用于MLP，激活函数
    --kan_bspline_grid 20 \ # 仅适用于KAN，样条网格数量
    --kan_bspline_order 5 \ # 仅适用于KAN，样条阶数
    --kan_grid_range -4 4 \ # 仅适用于KAN，样条定义域

参数数量和浮点运算次数

我们为KAN、MLP和带B样条激活的MLP分别提供了两个实例方法来计算其参数数量和浮点运算次数。

import torch
from argparse import Namespace

from models.mlp import *
from models.kanbefair import *
from models.bspline_mlp import *

# 实例化KANbeFair网络
args_KAN = Namespace(**{
    "input_size":28*28,
    "layers_width":[10,10],
    "output_size":10,
    "kan_bspline_grid":20,
    "kan_bspline_order":5,
    "kan_shortcut_function":torch.nn.SiLU(),
    "kan_grid_range":[-1,1]
    })
model_KAN = KANbeFair(args_KAN)
# 打印KAN的参数数量和浮点运算次数
num_parameters_KAN = model_KAN.total_parameters()
flops_KAN = model_KAN.total_flops()
print(f"KAN: 参数数量: {num_parameters_KAN:,}; 浮点运算次数: {flops_KAN:,}")

# 实例化MLP网络
args_MLP = Namespace(**{
    "input_size":28*28,
    "layers_width":[10,10],
    "output_size":10,
    "batch_norm": False,
    "activation":torch.nn.ReLU
    })
model_MLP = MLP(args_MLP)
# 打印MLP的参数数量和浮点运算次数
num_parameters_MLP = model_MLP.total_parameters()
flops_MLP = model_MLP.total_flops()
print(f"MLP: 参数数量: {num_parameters_MLP:,}; 浮点运算次数: {flops_MLP:,}")

引用

如果您觉得KANbeFair有用，请使用以下BibTeX进行引用：

@misc{yu2024kanbefair,
      title={KAN or MLP: A Fairer Comparison}, 
      author={Runpeng Yu and Weihao Yu and Xinchao Wang},
      year={2024},
      eprint={2407.16674},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2407.16674}, 
}

致谢

• pykan：我们要感谢作者的出色工作和代码库。