Want to train KANS at scale? Now UKAN!

Alireza Moradzadeh, Srimukh Prasad Veccham, Lukasz Wawrzyniak, Miles Macklin, Saee Gopal Paliwal

Keywords: GPU, Kolmogorov–Arnold, KAN, Unbounded-KAN, B-spline

Abstract: While Kolmogorov–Arnold Networks (KANs) are powerful MLP alternatives, their reliance on bounded grids limits their scalability and domain range. We propose Unbounded Kolmogorov–Arnold Networks (UKANs), which utilize a coefficient-generator (CG) model to dynamically produce B-spline coefficients on an unbounded symmetric grid. By integrating positional encodings, UKANs enable function approximation without data normalization. To address computational overhead, we introduce warpKAN, a GPU-accelerated library that optimizes memory management and reduces B-spline complexity proportional to grid size. Benchmarks show a $3\text{--}30\times$ speed-up and up to $1000\times$ memory reduction over vanilla KANs. Experiments across regression, classification, and generative tasks—including molecular property prediction—confirm that UKANs match or exceed KAN accuracy while enabling large-scale, end-to-end training.

Submission Number: 100