AI 中文总结
FlashKAN用截断幂形式替代KAN的Cox-de Boor递归,通过GPU内核优化、坐标稳定化实现加速,提供开源可直接替换的KAN层软件包。
AI 中文摘要
Kolmogorov-Arnold Networks(KANs)在网络边放置可学习的B样条激活函数,而非在节点上使用固定激活函数。标准Cox-de Boor递归法对k次样条需进行k次顺序运算来评估这些激活函数,占前向传播时间的90%以上。FlashKAN用截断幂形式(近似理论中的经典结果,将每个均匀三次B样条表示为偏移节点位置处的5个(x)_+^3项)取代该递归法。本文有三项贡献:(1)融合了http URL的实现,将这些运算合并为单个GPU内核,消除了所有递归、跨度查找和散度-聚集运算;(2)有界坐标稳定化,将归一化输入钳位到[0, k+1]区间,防止了历史上推动Cox-de Boor递归的灾难性抵消;(3)提供了可直接替换现有KAN层的生产级开源软件包(pip install flashkan)。
英文摘要
Kolmogorov-Arnold Networks (KANs) place learnable B-spline activations on network edges rather than fixed activations on nodes. The standard Cox-de Boor recursion evaluates these activations through $k$ sequential passes for degree-$k$ splines, consuming over 90% of forward-pass time. InKAN replaces this recursion with the truncated power form, a classical result from approximation theory that expresses each uniform cubic B-spline as five $(x)_+^3$ terms at shifted knot positions. The resulting expression computes exact B-spline basis values: the same mathematical function as the Cox-de Boor recursion, evaluated without sequential passes. This paper documents three contributions: (1) an implementation structured for torch . compile fusion, eliminating recursion, span lookup, and scatter-gather operations; (2) a bounded-coordinate evaluation that clamps the normalized input to $[0, k{+}1]$, preventing the growth of cancellation error at large off-support coordinates; and (3) an open-source package (pip install inkan). In the tested configurations, InKAN has 2.8--3.5$\times$ lower forward-pass latency than the Cox-de Boor recursion. Partition-of-unity errors remain below $10^{-5}$ for grid sizes up to 200.
CommentsOpen-source implementation: pip install inkan, experiments (speed, classification, regression), 8 Pages