MatrixKAN: Parallelized Kolmogorov-Arnold Network

TL;DR

MatrixKAN introduces a parallelized approach to accelerate Kolmogorov-Arnold Networks by leveraging matrix operations, significantly reducing training and inference times especially for high-degree B-splines.

Contribution

The paper presents MatrixKAN, a novel method that parallelizes B-spline calculations in KANs using matrix operations, enhancing computational efficiency and scalability.

Findings

Achieves approximately 40x speedup over traditional KAN.

Demonstrates superior scaling with B-spline degree.

Potential for further speedup with larger datasets or higher degrees.

Abstract

Kolmogorov-Arnold Networks (KAN) are a new class of neural network architecture representing a promising alternative to the Multilayer Perceptron (MLP), demonstrating improved expressiveness and interpretability. However, KANs suffer from slow training and inference speeds relative to MLPs due in part to the recursive nature of the underlying B-spline calculations. This issue is particularly apparent with respect to KANs utilizing high-degree B-splines, as the number of required non-parallelizable recursions is proportional to B-spline degree. We solve this issue by proposing MatrixKAN, a novel optimization that parallelizes B-spline calculations with matrix representation and operations, thus significantly improving effective computation time for models utilizing high-degree B-splines. In this paper, we demonstrate the superior scaling of MatrixKAN's computation time relative to B-spline degree. Further, our experiments demonstrate speedups of approximately 40x relative to KAN, with significant additional speedup potential for larger datasets or higher spline degrees.