An Efficient Implementation of Kolmogorov-Arnold Network

This repository contains an efficient implementation of Kolmogorov-Arnold Network (KAN). The original implementation of KAN is available here.

It borrows a lot from efficient-kan, but include also plotting, locks and pruning. Radial basis function are also implemented and allow for extra speedup.

Contrary to efficient-kan, it use the original L1 regularization of pykan without sacrifying too much too performance.

Note on RBF

There is some fast implementations of KAN that include RBF, but there are not equivalent to the original paper when there is more than one input. Indeed, most RBF implementation use for node $q$, inputs $x_p$ and centers $c_i$

$$ \sumi c{q,i} \phi\left(\sqrt{\sum_p (x_p-c_i)^2 }\right) $$

whereas the correct formule should be

$$ \sum_p \sumi c{q,i} \phi \left( |x_p-c_i|\right) .$$

The last formula is the one implemented here.

Current benchmark (19/07/2024)

The implementation is not as fast as efficient-kan when using slow version but still allow for significant speedup with respect to pykan. Based on this benchmarck

python3 benchmark.py --batch-size 1000 --inp-size 100 --hid-size 1000 --reps 50

However, when swicthing to fast version, speed becomes comparable to faster KAN implementation. But this fast version does not allow to use the original regularization neither the plotting and pruning utilities. Comparaison with pykan use model.speed()

python3 benchmark.py --batch-size 1000 --inp-size 100 --hid-size 1000 --reps 50 --fast

Comparing to pykan v1 (with smaller number of parameters)

python3 benchmark.py --batch-size 1000 --inp-size 100 --hid-size 100 --reps 10 --method pykan

Comparing to pykan v2 (with smaller number of parameters)

python3 benchmark.py --batch-size 1000 --inp-size 100 --hid-size 100 --reps 10 --method pykan