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.
>Of course, if my code doesn't work, it doesn't mean theirs doesn't.
But, _does_ it work?
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.
Sums and products can get you surprisingly far.
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.
It works as advertised with the parameters selected by the authors, but if we modified the network shape in the second half of the tutorial (Classification formulation) from (2, 2) to (2, 2, 2), it fails to generalize. The training loss gets down to 1e-9, while test loss stays around 3e-1. Getting to larger network sizes does not help either.
I would really like to see a bigger example with many more parameters and more data complexity and if it could be trained at all. MNIST would be a good start.
Update: I increased the training dataset size 100x, and that helps with the overfitting, but now I can't get training loss below 1e-2. Still iterating on it; a GPU acceleration would really help - right now, my progress is limited by the speed of my CPU.
1. https://github.com/KindXiaoming/pykan/blob/master/tutorials/...
Changes:
1. Increased the training set from 1000 to 100k samples. This solved overfitting.
2. In the dataset generation, slightly reduced noise (0.1 -> 0.07) so that classes don't overlap. With an overlap, naturally, it's impossible to hit 100%.
3. Most important & specific to KANs: train for 30 steps with grid=5 (5 segments for each activation function), then 30 steps with grid=10 (and initializing from the previous model), and then 30 steps with grid=20. This is idiomatic to KANs and covered in the Example_1_function_fitting.ipynb: https://github.com/KindXiaoming/pykan/blob/master/tutorials/...
Overall, my impressions are:
- it works!
- the reference implementation is very slow. A GPU implementation is dearly needed.
- it feels like it's a bit too non-linear and training is not as stable as it's with MLP + ReLU.
- Scaling is not guaranteed to work well. Really need to see if MNIST is possible to solve with this approach.
I will definitely keep an eye on this development.
Solved over fitting or created more? Even if your sets are completely disjoint with something like two moons the more data you have the lower the variance.
This. I don't think toy examples are useful for modern ML techniques. If you tested big ideas in ML (transformers, LSTM's, ADAM) on a training dataset of 50 numbers trying to fit a y=sin(x) curve, I think you'd wrongly throw these ideas out.
Unfortunately, I had to modify KAN.py and KANLayer.py to make it work as not all relevant tensor are put on the correct device. In some places the formatting even suggests that there was previously a device argument.
GLMs in turn generalize logistic-, linear and other popular regression models.
Neural GAMs with learned basis functions have already been proposed, so I'm a bit surprised that the prior art is not mentioned in this new paper. Previous applications focused more on interpretability.
It's not clear from the paper how well this algorithm will scale, both in terms of the algorithm itself (does it still train well with more layers?), and ability to make use of hardware acceleration, (e.g. it's not clear to me that the structure, with its per-weight activation functions, can make use of fast matmul acceleration).
It's an interesting idea, that seems to work well and have nice properties on a smaller scale; but whether it's a good architecture for imagenet, LLMs, etc. is not clear at this stage.
Sounds like something which could be approximated by a DCT (discrete cosine transform). JPEG compression does this, and there are hardware accelerations for it.
> can make use of fast matmul acceleration
Maybe not, but matmul acceleration was done in hardware because it's useful for some problems (graphics initially).
So if these per weight activations functions really work, people will be quick to figure out how to run them in hardware.
The best thing about this new work is that it's not an either/or proposition. The proposed "learnable spline interpolations as activation functions" can be used in conventional DNNs, to improve their expressivity. Now we just have to test the stuff to see if it really works better.
Very nice. Thank you for sharing this work here!
---
References:
[1] A grad.-based way to optimize axis-parallel and oblique decision trees: the Tree Alternating Optimization (TAO) algorithm https://proceedings.neurips.cc/paper_files/paper/2018/file/1.... An extension was the softmax tree https://aclanthology.org/2021.emnlp-main.838/.
[2] XAI explains models, but can you recommend corrective actions? FACE: feasible and Actionable Counterfactual Explanations https://arxiv.org/pdf/1909.09369, Algorithmic Recourse: from Counterfactual Explanations to Interventions https://arxiv.org/pdf/2002.06278
[3] OBOE: Collaborative Filtering for AutoML Model Selection https://arxiv.org/abs/1808.03233
Yes, I agree. The two most common patterns I've noticed in research that does show up on HN are: 1) It outright improves, or has the potential to improve, applications currently used in production by many HN readers. In other words, it's not just navel-gazing. 2) The authors and/or their organizations are well-known, as you suggest.
Everything described was laughably basic by modern standards, but the motivation given in that book was the Kolmogorov representation theorem: a modest 3 layer networks with the right activation function can represent any continuous m-to-n function.
Most research back then focused on 3 layer networks, possibly for that reason. Sigmoid activation was king, and vanishing gradients the main issue. It took 2 decades until AlexNet brought NN research back from the AI winter of the 1990’s
This is science as is :)
95% percent will produce mediocre-to-nice improvements to what we already have so there were reserachers that eventually grow up and do something really exciting
> A practical construction of g in cases with discontinuous bounded and un- bounded functions is not yet known. For such cases Theorem 2.1 gives only a theoretical understanding of the representation problem. This is because for the representation of discontinuous bounded functions we have derived (2.1) from the fact that the range of the operator Z∗ is the whole space of bounded functions B(Id). This fact directly gives us a formula (2.1) but does not tell how the bounded one-variable function g is attained. For the representation of unbounded functions we have used a linear extension of the functional F , existence of which is based on Zorn’s lemma (see, e.g., [19, Ch. 3]). Application of Zorn’s lemma provides no mechanism for practical construction of such an extension. Zorn’s lemma helps to assert only its existence.
If you look at the OP post arxiv link, you will see they are using splines .
https://arxiv.org/abs/2404.19756
Still interesting and potentially useful, but not useful for discontinuous functions without further discoveries.
If I am wrong please provide a link, it is of great interest to me.
1957: https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Arnold_repr...
1958: https://en.wikipedia.org/wiki/Multilayer_perceptron
2. Another advantage of this approach is that it has only one class of parameters (the coefficients of the local activation functions) as opposed to MLP which has three classes of parameters (weights, biases, and the globally uniform activation function).
3. Everybody is talking transformers. I want to see diffusion models with this approach.
There isn't much difference between weights of a linear sum and coefficients of a spline.
Granted, however this approach does not require that constant-one input either.
> There isn't much difference between weights of a linear sum and coefficients of a function.
Yes, the trained function coefficients of this approach are the equivalent to the trained weights of MLP. Still this approach does not require the globally uniform activation function of MLP.
One might argue this via parsimony (Occam’s razor). Is this your thinking? / Anything else?
I'm not seeing decision trees, though. Am I missing something?
> "KANs’ nodes simply sum incoming signals without applying any non-linearities." (page 2 of the PDF)
- PyTorch Module of the KAN GPT
- Deployed to PyPi
- MIT Licence
- Test Cases to ensure forward-backward passes work as expected
- Training script
I am currently working on training it on the WebText dataset to compare it to the original gpt2. Facing a few out-of-memory issues at the moment. Perhaps the vocab size (50257) is too large?
I'm open to contributions and would love to hear your thoughts!
Seminar 2021: https://warwick.ac.uk/fac/sci/maths/research/events/seminars...
Article in archive 2023: https://arxiv.org/abs/2305.08194
Video 2021: https://www.youtube.com/watch?v=eS_k6L638k0
Extension to stochastic models where KAN builds the distribution 2023: https://www.youtube.com/watch?v=0hhJIpzxPR0
At the end of this example, they recover the symbolic formula that generated their training set: exp(x₂² + sin(3.14x₁)).
It's like a computation graph with a library of "activation functions" that is optimised, and then pruned. You can recover good symbolic formulas from the pruned graph.
Maybe not meaningful for MNIST.
It will easily recover this formula, because it is separable under the log transformation (which ACE recovers as well).
But ACE doesn’t work well on unseparable problems - not sure how well KAN will.
Given its sparse, Will this be just replacement for MoE.
> Unsurprisingly, the possibility of using Kolmogorov-Arnold representation theorem to build neuralnetworks has been studied [8, 9, 10, 11, 12, 13]. However, most work has stuck with the original depth-2 width-(2n + 1) representation, and did not have the chance to leverage more modern techniques (e.g., back propagation) to train the networks. Our contribution lies in generalizing the original Kolmogorov-Arnold representation to arbitrary widths and depths, revitalizing and contextualizing it in today’s deep learning world, as well as using extensive empirical experiments to highlight its potential role as a foundation model for AI + Science due to its accuracy and interpretability.
Are you assuming that this particular "progress" would be relatively innocent?
If and when something radically better comes along, say an alternative to back-propagation that is more like the way our brains learn, it will need a lot of scaling and refinement to catch up with the then-current LLM.
I would worry if I'd own Nvidia shares.
1. A new architecture would make all/most of these upcoming Transformer accelerators obsolete => back to GPUs.
2. Higher performance LLMs on GPUs => we can speed up LLMs with 1T+ parameters. So, LLMs become more useful, so more of GPUs would be purchased.
In the article "Distilling free-form natural laws from experimental data", Schmidt and Lipson introduced the idea that free-form natural laws can be learned from experimental measurements in a physical system using symbolic (genetic) regression algorithms. An important claim in this work is that the algorithm finds laws in data without having incorporated any prior knowledge of physics. Upon close inspection, however, we show that their method implicitly incorporates Hamilton's equations of motions and Newton's second law, demystifying how they are able to find Hamiltonians and special classes of Lagrangians from data.
I cannot get a PDF of your article and instead I will read a commentary on it which appears to be very interesting.
Would this approach (with non-linear learning) still be able to utilize GPUs to speed up training?
https://en.wikipedia.org/wiki/Bayesian_network#Graphical_mod...
> Each node is associated with a probability function that takes, as input, a particular set of values for the node's parent variables, and gives (as output) the probability (or probability distribution, if applicable) of the variable represented by the node. For example, if m parent nodes represent m Boolean variables, then the probability function could be represented by a table of 2^m entries, one entry for each of the 2^m possible parent combinations.
I mean it's great but at the current state it seems better suited for tasks where an explicit formula exists (though not known) and the goal is to predict it on unknown points (and be able to interpret the formula as a side effect). Deep learning tasks are more of a statistical nature (think models with a cross entropy loss - it's statistically predicting the frequency of different choices of the class/next token), it requires a specialized training procedure and it is designed to fit 100% rather than somewhat close (think linear algebra - it won't be good at it). It would very likely take a radically different idea to apply it to deep learning tasks. The recently updated "Author's note" also mentions this: "KANs are designed for applications where one cares about high accuracy and/or interpretability."
It's great but let's be patient before we see this improve LLM accuracy or be used elsewhere.
I wonder how many more new architectures are going to be found in the next few years
> doesn't KA representation require continuous univariate functions?
All multivariate continuous functions (on a bounded domain) can be represented as compositions of addition and univariate continuous functions. Much like an MLP, you can also approximate discontinuous functions well on most of the domain (learning a nearby continuous function instead).
> do B-splines actually cover the space of all continuous functions
Much like an MLP, you can hit your favorite accuracy bound with more control points.
> wouldn't... MLPs be better for the learnable activation functions
Perhaps. A B-spline is comparatively very fast to compute. Also, local training examples have global impacts on an MLP's weights. That's good and bad. One property you would expect while training a KAN in limited data regimes is that some control points are never updated, leading to poor generalization due to something like a phase shift as you cross over control points (I think the entropy-based regularizer they have in the paper probably solves that, but YMMV). The positive side of that coin is that you neatly side-step catastrophic forgetting.
It’s weird to just ignore MLPs when approximating a continuous univariate function. But if the paper did use MLPs theyd have ended up with something that looks a lot more like conventional neural networks, so maybe thats why?