LUT-KAN:面向快速KAN推理的分段LUT量化
机器学习
2026-01-08 v1
摘要
Kolmogorov-Arnold网络(KAN)用可学习的单变量函数取代标量权重,通常用B样条实现。这种设计既准确又具可解释性,但会使CPU上的推理成本高昂,因为每层需要大量样条求值。标准量化工具链也难以应用,因为主要计算不是矩阵乘法,而是重复的样条基求值。本文引入LUT-KAN,一种针对PyKAN风格KAN层的分段查找表(LUT)编译和量化方法。LUT-KAN将每个边缘函数转换为每个分段的LUT,采用仿射int8/uint8量化和线性插值。该方法提供了明确且可重复的推理契约,包括边界约定和超出边界(OOB)策略。我们提出了一种“诚实基线”方法,用于速度评估:在相同的后端优化下(NumPy vs NumPy和Numba vs Numba),将B样条求值与LUT求值进行比较,这将表示收益与向量化和JIT效果分离。实验包括对LUT分辨率L在16、32、64、128上的受控扫描,以及两种量化方案(对称int8和非对称uint8)。我们报告准确率、速度和内存指标,包括多个随机种子下的均值和标准差。一个两-by-two OOB鲁棒性矩阵评估了不同边界模式和OOB策略下的行为。在一个案例研究中,我们将训练好的KAN模型用于DoS攻击检测(CICIDS2017管道)编译为LUT制品。编译后的模型在保持分类质量(F1下降不到0.0002)的同时,在NumPy后端下将稳态CPU推理延迟降低12倍,在Numba后端下降低10倍(诚实基线)。内存开销在L=64时约为10倍。所有代码和制品均公开可用,附有固定发布标签以确保可重复性。
引用
@article{arxiv.2601.03332,
title = {LUT-KAN: Segment-wise LUT Quantization for Fast KAN Inference},
author = {Oleksandr Kuznetsov},
journal= {arXiv preprint arXiv:2601.03332},
year = {2026}
}