PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashKAN: B-Spline KANs via Truncated Power Form

Naveen Mysore

arXiv:2609.01956v1cs.LGmath.NA

TL;DR

KANs face a forward-pass bottleneck because Cox-de Boor recursion evaluates B-spline bases through sequential passes. FlashKAN replaces that recursion with a fused truncated-power formulation and bounded-coordinate stabilization, packaged as a drop-in KAN layer; its supported scope is uniform knot vectors.

  • Problem

    Cox-de Boor B-spline evaluation requires sequential passes and accounts for 91% of a KAN layer’s forward-pass time.

  • Method

    FlashKAN evaluates uniform cubic B-splines with five truncated-power terms, fuses operations into one GPU kernel, and clamps normalized coordinates to [0, k+1].

  • Results

    FlashKAN provides a single broadcastable evaluation expression while preserving compact-support and partition-of-unity structure.

  • Takeaways & Limitations

    The open-source package offers an API-compatible replacement for standard KAN layers, requiring one import change.

  • Takeaways & Limitations

    The closed form applies only to uniform knot vectors; extension to non-uniform grids is left for future work.

Abstract

from arXiv · show

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.

1 INTRODUCTION

KANs place learnable B-spline activations on network edges, but Cox-de Boor evaluation makes basis computation the dominant forward-pass cost. FlashKAN uses the truncated power form with compiler fusion and bounded coordinates to target this deployment bottleneck.

  • KANs place learnable B-spline activations on network edges, whose weighted basis functions can be inspected and sometimes symbolically recovered.
  • 91% of a KAN layer’s forward-pass time is consumed by B-spline basis computation through k sequential Cox-de Boor passes for degree-k splines.For cubic splines, three passes are required, each depending on the preceding pass.
  • The truncated power form expresses uniform cubic B-splines as five max(0, ·)^3 terms with fixed coefficients.Its mathematical elegance was historically offset by numerical cancellation on early hardware.
  • FlashKAN fuses truncated-power evaluation into one GPU kernel, eliminating recursion, span lookup, and data-dependent memory access.This targets both sequential dependency and fragmented GPU memory access identified in prior approaches.
  • FlashKAN adds bounded-coordinate stabilization and packages the method as an open-source, API-compatible replacement for standard KAN layers.The normalized coordinate is clamped to [0, k+1] before evaluation, and installation requires changing one import.

2 FROM RECURSION TO CLOSED FORM

The section transforms recursive spline evaluation into a closed-form truncated-power expression whose elementwise structure supports fused computation, then stabilizes it by bounding normalized coordinates. This preserves B-spline properties while addressing the dominant computational and numerical costs.

  • 2.1 STEP 1: DE CASTELJAU TO BERNSTEIN: De Casteljau’s recursive interpolations expand into the Bernstein basis, separating curve evaluation into input powers, fixed coefficients, and control points.For cubic curves, the input-dependent power vector is [1, t, t^2, t^3].
  • 2.3 STEP 3: BÉZIER SEGMENTS TO B-SPLINES: B-splines provide compact support and C^k−1 continuity by construction, so local control and smoothness do not require separately imposed constraints.Uniform B-spline segments retain the same T(t) M P factorization while changing only the basis coefficients.
  • 2.3 STEP 3: BÉZIER SEGMENTS TO B-SPLINES: 91% of a KAN layer’s forward-pass time is consumed by cubic B-spline basis computation through three sequential Cox-de Boor passes.Each pass depends on the previous one.
  • 2.4 STEP 4: COX-DE BOOR TO TRUNCATED POWER FORM: For uniform cubic splines, expanding Cox-de Boor yields five truncated-power terms with coefficients {1, −4, 6, −4, 1}/6.The positive-part function is (z)+ = max(0, z).
  • 2.4 STEP 4: COX-DE BOOR TO TRUNCATED POWER FORM: The resulting expression replaces sequential interpolation with one subtraction, one multiply, and elementwise operations that torch.compile can fuse into a single kernel.This removes sequential dependency, span lookup, and data-dependent memory access.
  • 2.5 BOUNDED-COORDINATE STABILIZATION: Out-of-support evaluation causes catastrophic cancellation because alternating terms of size Θ(u^3) cancel to an exact zero, while finite precision leaves residual error O(ϵmach · u^3).In KANs, hidden activations can leave the nominal grid domain, and float16 overflows when u exceeds approximately 40.
  • 2.5 BOUNDED-COORDINATE STABILIZATION: Clamping the normalized coordinate to [0, k+1] preserves all in-support values while bounding cubic intermediate quantities by 64.The clamp is folded into the same fused kernel at no measurable cost.
  • 2.6 INTEGRATION WITH KAN LAYERS: Replacing Cox-de Boor with bounded truncated-power evaluation changes only basis computation; learnable weights, residual connections, and gradient flow remain identical.The modified basis therefore integrates directly into the existing KAN layer computation.

3 RELATED WORK

Related work reduces spline cost through recursion restructuring, matrix precomputation, or replacing B-splines with Gaussian bases. FlashKAN instead combines the truncated-power equivalence with compiler fusion and bounded-coordinate stabilization.

  • Existing KAN implementations: Efficient-KAN reduces constant factors but preserves recursive dependency, while FastKAN uses Gaussian bases that lose exact compact support and automatic partition of unity.The Gaussian basis remains non-negative and C∞ smooth.
  • Truncated power basis and change of basis: Southworth et al. establish a Toeplitz change-of-basis equivalence between spline KAN layers and multichannel MLPs with power-ReLU activations.Their primary goal is multilevel training, including orders-of-magnitude PINN accuracy improvements.
  • Truncated power basis and change of basis: FlashKAN pursues compiler-level speed, bounded-coordinate numerical stabilization, and production deployment, complementing Southworth et al.’s multilevel training framework.The paper proposes that FlashKAN could serve as the fast evaluation engine within a multilevel training loop.
  • Numerical stability of the truncated power form: The truncated-power basis is poorly conditioned far outside support because terms of size Θ(u^k) cancel to produce zero.Bounding the coordinate removes this far-outside-support regime.

4 DISCUSSION

FlashKAN’s truncated power form preserves compact support and partition of unity while providing a single broadcastable expression. Bounded-coordinate evaluation extends reliable operation beyond inputs near the grid.

  • Comparison with Gaussian RBF: Gaussian RBFs have global influence and require explicit normalization because their basis sum is not guaranteed to equal one.Their non-negativity and infinite differentiability are not the properties sacrificed in the comparison.
  • Comparison with Gaussian RBF: The truncated power form preserves exact compact support and automatic partition of unity while simplifying evaluation to one broadcastable expression.Unlike Gaussian RBFs, it retains these structural B-spline properties without explicit normalization for partition of unity.

5 LIMITATIONS

FlashKAN’s closed-form Eq. 6 is restricted to uniform knot vectors. Supporting non-uniform grids would reintroduce data-dependent indexing, so that extension remains future work.

  • Scope boundary: FlashKAN’s closed form applies only to uniform knot vectors, leaving non-uniform-grid evaluation for future work.Non-uniform grids require per-span polynomial coefficients and reintroduce data-dependent indexing.

6 CONCLUSION

FlashKAN turns the truncated power representation into a deployable KAN implementation through stabilization, kernel fusion, and packaging. It is presented as the first KAN implementation using this evaluation form and includes broad correctness tests.

  • Conclusion: FlashKAN is presented as the first KAN implementation to evaluate B-spline basis functions through the truncated power form.The implementation converts a classical identity and recent algebraic equivalence into a deployable system.
  • Conclusion: Bounded-coordinate stabilization removes the numerical cancellation that historically discouraged truncated power evaluation.The method clamps normalized coordinates before evaluation to bound intermediate terms.
  • Conclusion: torch.compile fuses evaluation into a single GPU kernel, while the package provides a drop-in replacement available through pip under the MIT license.The package includes 26 unit tests covering basis correctness, partition of unity, compact support, C2 smoothness, and gradient flow.
Loading 2609.01956v1…