In the paper they use some spline interpolation to represent 1d function that they sum. Their code seemed aimed at smaller sizes. Instead I chose a different representation, aka fourier coefficients that are used to interpolate the functions of individual coordinates.
It should give an idea of Kolmogorov-Arnold networks representation power, it should probably converge easier than their spline version but spline version have less operations.
Of course, if my code doesn't work, it doesn't mean theirs doesn't.
Feel free to experiment and publish paper if you want.
When I played around with implementing this last night I found using a radial basis function instead of Fourier coefficients (I tried the same, nice and parallel and easy to write) to be more well behaved in training networks of depth greater than 2.
Hi Noesis, I just noticed that your implementation, combined with the efficientKAN by Blealtan (https://github.com/Blealtan/efficient-kan), results in a structure very similar to Siren(MLP with Sin activations). efficientKAN first computes the common basis functions for all the edge activations and the output can be calculated with a linear combination of the basis. If the basis functions are fourier, then a KAN layer can be viewed as a linear layer with fixed weights + Sin activation + a linear layer with learnable weights, which is a special form of Siren.
I think this may show some connection between KAN and MLP.
How could this help us understand the difference between the learned parameters and their gradients? Can the gradients become one with the parameters a la exponential function?
The symbolic library (type of activations) requires a branching at the very core of the kernel. GPU will need to serialized on these operations warp-wise.
To optimize, you might want to do a scan operation beforehand and dispatch to activation funcs in a warp specialized way, this, however, makes the global memory read/write non-coalesced.
You then may sort the input based on type of activations and store it in that order, this makes the gmem IO coalesced but requires gather and scatter as pre and post processing.
Conceptually it's simpler to think about and optimize. But you can also write it use einsum to do the sum product reductions (I've updated some comment to show how) to use less memory, but it's more intimidating.
You can probably use KeOps library to fuse it further (einsum would get in the way).
But the best is probably a custom kernel. Once you have written it as sums and product, it's just iterating. Like the core is 5 lines, but you have to add roughly 500 lines of low-level wrapping code to do cuda parallelisation, c++ to python, various types, manual derivatives. And then you have to add various checks so that there are no buffer overflows. And then you can optimize for special hardware operations like tensor cores. Making sure along the way that no numerical errors where introduced.
So there are a lot more efforts involved, and it's usually only worth it if the layer is promising, but hopefully AI should be able to autocomplete these soon.
https://github.com/GistNoesis/FourierKAN/
The core is really just a few lines.
In the paper they use some spline interpolation to represent 1d function that they sum. Their code seemed aimed at smaller sizes. Instead I chose a different representation, aka fourier coefficients that are used to interpolate the functions of individual coordinates.
It should give an idea of Kolmogorov-Arnold networks representation power, it should probably converge easier than their spline version but spline version have less operations.
Of course, if my code doesn't work, it doesn't mean theirs doesn't.
Feel free to experiment and publish paper if you want.