提升科学计算中KAN网络的GPU效率，速度最高快12倍。
PolyKAN: Efficient Fused GPU Operators for Polynomial Kolmogorov-Arnold Network Variants
- 用融合核函数优化KAN层前向与反向传播。
- 在多种任务上实现1.2到12倍加速，精度不变。
- 适合需要高效训练科学模型的研究者使用。
Kolmogorov-Arnold Networks（KANs）在科学人工智能领域展现出比多层感知机更强的表达能力与可解释性，但现有并行实现的GPU利用率低，限制了实际应用。为此，我们提出首个通用开源的KAN及其变体加速库PolyKAN。该库将多项式KAN层的前向与反向传播融合为一组优化的CUDA内核，采用四项关键技术：(i) 查表线性插值替代高开销数学函数；(ii) 二维分块以暴露线程级并行性并保持内存局部性；(iii) 两阶段归约将散乱原子更新转为单一可控合并步骤；(iv) 系数布局重排，在分块调度下实现单位步长读取。以切比雪夫KAN为例，PolyKAN在高端与消费级GPU上对语音、音频增强和表格回归任务均达到与Triton + cuBLAS基线相同的精度，推理速度提升1.2–10倍，训练速度提升1.4–12倍。
原文摘要 · Abstract (English)
Kolmogorov-Arnold Networks (KANs) promise higher expressive capability and stronger interpretability than Multi-Layer Perceptron, particularly in the domain of AI for Science. However, practical adoption has been hindered by low GPU utilization of existing parallel implementations. To address this challenge, we present a GPU-accelerated operator library, named PolyKAN which is the first general open-source implementation of KAN and its variants. PolyKAN fuses the forward and backward passes of polynomial KAN layers into a concise set of optimized CUDA kernels. Four orthogonal techniques underpin the design: (i) \emph{lookup-table} with linear interpolation that replaces runtime expensive math-library functions; (ii) \emph{2D tiling} to expose thread-level parallelism with preserving memory locality; (iii) a \emph{two-stage reduction} scheme converting scattered atomic updates into a single controllable merge step; and (iv) \emph{coefficient-layout reordering} yielding unit-stride reads under the tiled schedule. Using a KAN variant, Chebyshev KAN, as a case-study, PolyKAN delivers $1.2$--$10\times$ faster inference and $1.4$--$12\times$ faster training than a Triton + cuBLAS baseline, with identical accuracy on speech, audio-enhancement, and tabular-regression workloads on both highend GPU and consumer-grade GPU.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。