用查表法实现高效高维映射,推理速度提升6倍以上。
Lookup multivariate Kolmogorov-Arnold Networks
- 用可训练的低维分段函数替代传统线性层,通过查表加速计算。
- 在高维函数逼近中性能媲美MLP,推理浮点运算量降低6.0倍。
- 适合追求推理效率的模型部署,尤其适用于图像与分子数据。
高维线性映射(即线性层)主导了现代深度学习模型的参数量和计算成本。本文提出通用的即插即用替代方案——查找型多变量柯尔莫戈洛夫-阿诺德网络(lmKAN),显著优化了模型容量与推理开销的权衡。其核心思想是将高维映射表示为可训练的低维多变量函数,每个函数可包含数十至数百个可训练参数,但通过样条查表实现,仅需少量乘法即可计算。实验表明,lmKAN在一般高维函数逼近任务中性能与MLP相当,推理FLOPs降低最多达6.0倍;在随机位移甲烷构型的表格类数据集上,lmKAN在相同精度下使H100吞吐量提升超10倍。在卷积神经网络框架中,基于lmKAN的CNN在保持精度的前提下,推理FLOPs分别在CIFAR-10和ImageNet-1k数据集上减少1.6-2.1倍和1.7倍。相关代码(含专用CUDA核)已开源:https://github.com/schwallergroup/lmkan。
原文摘要 · Abstract (English)
High-dimensional linear mappings, or linear layers, dominate both the parameter count and the computational cost of most modern deep-learning models. We introduce a general-purpose drop-in replacement, lookup multivariate Kolmogorov-Arnold Networks (lmKANs), which deliver a substantially better trade-off between capacity and inference cost. Our construction expresses a general high-dimensional mapping through trainable low-dimensional multivariate functions. These functions can carry dozens or hundreds of trainable parameters each, and yet it takes only a few multiplications to compute them because they are implemented as spline lookup tables. Empirically, lmKANs reduce inference FLOPs by up to 6.0x while matching the flexibility of MLPs in general high-dimensional function approximation. In another feedforward fully connected benchmark, on the tabular-like dataset of randomly displaced methane configurations, lmKANs enable more than 10x higher H100 throughput at equal accuracy. Within frameworks of Convolutional Neural Networks, lmKAN-based CNNs cut inference FLOPs at matched accuracy by 1.6-2.1x and by 1.7x on the CIFAR-10 and ImageNet-1k datasets, respectively. Our code, including dedicated CUDA kernels, is available online at https://github.com/schwallergroup/lmkan.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。