InKAN:用截断幂形式替代Cox-de Boor递归,提速KAN前向传播

新论文提出InKAN,将均匀三次B样条激活改为截断幂形式并融合为单一GPU内核,消除递归和散聚操作,同时提供开源pip包inkan,可作为现有KAN层的即插即用替代。

AI解读:KAN网络把可学习的B样条激活放在边而非节点上,但标准Cox-de Boor递归在每次前向传播里串行执行k次,竟能消耗超过 90% 的时间,这让KAN的实用性大打折扣。InKAN的突破口是用逼近论里的截断幂形式来表达均匀三次B样条——每条样条只需五个移位的 (x)_+^3 项,然后把torch.compile融合进单一GPU内核,彻底移除递归、跨度查找和散聚操作。为避免经典数值问题,它还加入有界坐标稳定化,把归一化输入钳制到 [0, k+1],防止灾难性抵消。对机器学习研究者或工程师来说,这意味着可以执行pip install inkan直接替换现有KAN层,无需重写模型代码,前向速度有望显著提升,尤其适合批量处理或迭代训练场景。不过论文仅 7 页且正在评审,尚未提供具体性能对比数据,实际加速幅度仍有待验证。

arXiv论文(编号 2609.01956,正在ICLR 2027评审)提出InKAN,用截断幂形式替代Cox-de Boor递归来评估B样条激活,声称可消除递归和散聚操作,并发布开源pip包inkan,作为现有KAN层的即插即用替代。

技术方案与实现

Kolmogorov-Arnold Networks(KAN)将可学习B样条激活放置在网络边上,而非节点上。标准Cox-de Boor递归对k次样条激活需要k次串行遍历,前向传播时间中超过 90% 消耗在此。

InKAN用截断幂形式(truncated power form)替换该递归。这是逼近论中的经典结果,将每个均匀三次B样条表示成五个在移位节点位置的 (x)_+^3 项。

论文的三个贡献为:(1) 使用torch.compile融合实现,将所有操作折叠为单一GPU内核,消除所有递归、跨度查找(span lookup)和散聚(scatter-gather)操作;(2) 有界坐标稳定化,将归一化输入钳制到 [0, k+1],防止历史上促使Cox-de Boor递归出现的灾难性抵消;(3) 提供生产就绪的开源包,可通过pip install inkan安装。

信息来源