FlashKAN: B-Spline KANs via Truncated Power Form

2026-09-01

AI Summary

FlashKAN通过使用截断幂形式替代Cox-de Boor递归，优化B样条激活函数的计算，减少前向传递时间，并提供了一个生产就绪的开源包。

Abstract

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. FlashKAN 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. This paper makes three contributions: (1) a torch.compile-fused implementation that collapses these operations into a single GPU kernel, eliminating all recursion, span lookup, and scatter-gather operations; (2) a bounded-coordinate stabilization that clamps the normalized input to [0, k+1], preventing the catastrophic cancellation that historically motivated the Cox-de Boor recursion; and (3) a production-ready, open-source package (pip install flashkan) that serves as a drop-in replacement for existing KAN layers.

Problem

• Kolmogorov-Arnold Networks
• B-spline activations
• Cox-de Boor recursion
• truncated power form

Innovation

• truncated power form
• torch.compile-fused implementation
• bounded-coordinate stabilization