Want to train KANS at scale? Now UKAN!
Not on the map. Why?
Abstract (arXiv)
Kolmogorov-Arnold Networks (KANs) have recently emerged as a powerful alternative to traditional multilayer perceptrons. However, their reliance on predefined, bounded grids restricts their ability to approximate functions on unbounded domains. To address this, we present Unbounded Kolmogorov-Arnold Networks (UKANs), a method that removes the need for bounded grids in traditional Kolmogorov-Arnold Networks (KANs). The key innovation of this method is a coefficient-generator (CG) model that produces, on the fly, only the B-spline coefficients required locally on an unbounded symmetric grid. UKANs couple multilayer perceptrons with KANs by feeding the positional encoding of grid groups into the CG model, enabling function approximation on unbounded domains without requiring data normalization. To reduce the computational cost of both UKANs and KANs, we introduce a GPU-accelerated library that lowers B-spline evaluation complexity by a factor proportional to the grid size, enabling large-scale learning by leveraging efficient memory management, in line with recent software advances such as FlashAttention and FlashFFTConv. Performance benchmarking confirms the superior memory and computational efficiency of our accelerated KAN (warpKAN), and UKANs, showing a 3-30x speed-up and up to 1000x memory reduction compared to vanilla KANs. Experiments on regression, classification, and generative tasks demonstrate the effectiveness of UKANs to match or surpass KAN accuracy. Finally, we use both accelerated KAN and UKAN in a molecular property prediction task, establishing the feasibility of large-scale end-to-end training with our optimized implementation.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 15 of its 38 references that are in the map. First posted on arXiv 20 Aug 2024.
Builds on
15 of 38 references.
References · 15
- 2017Attention Is All You Needcited by 59,932
- 2020Denoising Diffusion Probabilistic Modelscited by 20,851
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
- 2017MoleculeNet: A Benchmark for Molecular Machine Learningcited by 1,037
- 2024KAN: Kolmogorov-Arnold Networkscited by 634
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 507
- 2023Hyena Hierarchy: Towards Larger Convolutional Language Modelscited by 388
- 2017Physics Informed Deep Learning (Part I): Data-driven Solutions of Nonlinear Partial Differential Equationscited by 362
- 2024Wav-KAN: Wavelet Kolmogorov-Arnold Networkscited by 118
- 2024DeepOKAN: Deep Operator Network Based on Kolmogorov Arnold Networks for Mechanics Problemscited by 82
3 more
- 2024TKAN: Temporal Kolmogorov-Arnold Networkscited by 70
- 2017Deep Learning for Computational Chemistrycited by 38
- 2023FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Corescited by 32
Built on by
Citing papers per year: 2025: 1.
1 paper.
- 2025ChemKANs for Combustion Chemistry Modeling and Acceleration17 Apr 2025