Hacker News Reader

Comment by smusamashah | Hacker News Reader

> One downside of KANs is that they take longer per parameter to train—in part because they can’t take advantage of GPUs. But they need fewer parameters. Liu notes that even if KANs don’t replace giant CNNs and transformers for processing images and language, training time won’t be an issue at the smaller scale of many physics problems.

They don't even say that it might be possible to take advantage of GPUs in future. Reads like a fundamental problem with these.

View on HN

From the KAN repo itself, it appears they already have GPU support https://github.com/KindXiaoming/pykan/blob/master/tutorials/...

This looks interesting for sure:

"ReLU-KAN: New Kolmogorov-Arnold Networks that Only Need Matrix Addition, Dot Multiplication, and ReLU" https://arxiv.org/abs/2406.02075#

I’ve seen neural nets combined with decision trees. There’s a few ways to do such hybrids. One style essentially uses the accurate, GPU-trained networks to push the interpretable networks to higher accuracy.

Do any of you think that can be done cost-effectively with KAN’s? Especially using pre-trained, language models like LlaMa-3 to train the interpretable models?

I wonder what's the issue ... GPUs can do very complex stuff

From my limited understanding: No one has written GPU code for it yet.

There's not a lot of details there, but GPUs tend to not like code with a lot of branching. I'm guessing that's probably the issue.

I suspect it is because they have different activation functions on each edge, rather than using the same one over lots of data.

Are the activation functions truly different, or just different parameter values to one underlying function?

A usual problem is that GPUs don't branch on instructions efficiently. A next most likely problem is that they don't branch on data efficiently. Ideas fundamentally requiring the former or the latter are hard to port efficiently.

A simple example of something hard to port to a GPU is a deep (24 lvls) binary tree with large leaf sizes (4kb). Particular trees can be optimized further, particular operations on trees might have further optimizations, and trees with nicer dimensionality might have tricks available, but solving that problem in the abstract is 32x slower on a GPU that "good" GPU problems. That's not a death knell, but it cuts down substantially the constraints which would make a GPU a better fit than a CPU.

Instruction branching is much worse, when required. Runtime is exponential.

As far as KANs are concerned, the problem is more with data branching. Each spline computation requires its own set of data and is only used once. The math being done on the aggregate computations is non-negligible, but fast relative to the memory loads. You quickly enter a regime where (1) you're bottlenecked on RAM bandwidth, and (2) for a given RAM load you can't efficiently use the warp allocated to it.

You can tweak the parameters a bit to alleviate that problem (smaller splines allow you to load and parallelize a few at once, larger ones allow you to do more work at once), but it's a big engineering challenge to fully utilize a GPU for that architecture. Your best bets are (1) observing something clever allowing you to represent the same result with different computations, and (2) a related idea, construct a different KAN-inspired algorithm with similar expressivity and more amenable to acceleration. My gut says (2) is more likely, but we'll see.

More succinctly: The algorithm as written is not a good fit for the GPU primitives we have. It might be possible to bridge that gap, but that isn't guaranteed.

What about cards with higher memory bandwidth, like Groq’s LPUs? Would that help with data branching?

Data branching in general, no (pulling from 32 places is still 32x as expensive in that architecture, but you might be able to load bigger chunks in each place). For a KAN, a bit (it shifts the constants involved when I was talking about smaller vs bigger splines above -- sparsity and dropout will tend to make the GPU tend toward that worst-case though). You still have the problem that you're heavily underutilizing the GPU's compute.

What if instead of splines there were Fourier series or something like that? Would that be easier to infer and learn on GPU if it was somehow teachable?

EDIT: FourierKAN exists https://arxiv.org/html/2406.01034v1

Fourier subcomponents are definitely teachable in general. I'd expect that FourierKAN to have similar runtime to a normal KAN, only really benefitting from a GPU on datasets where you get better predictive performance than a normal KAN.

I'd expect Chebyshev polynomials to be much faster and easier to work with than splines, certainly, and probably Fourier series as well. (Especially if there aren't trig instructions in hardware, because then each sine or cosine is itself a chebyshev polynomial to evaluate).

Reply on news.ycombinator.com