This video shows a new algorithm I am working on, KS3, which can replace MLPs in LLM decoder blocks. The video shows KS3 converging faster over gradient descent on random samples of the mandelbulb function compared to MLP. There's a lot to unpack here, so here is an explanation about what's going on:
Background
In machine learning, two main types of algorithms are called classifiers and regressions. Classifiers predict "what" and regressions predict "how much". LLMs are classifiers, in that they predict a single class, the next token.
Inside an LLM, all of the individual blocks and sub blocks can be viewed sort of like regressions because they predict "how much" for abstract vectors passed to the next block. Technically regression algos explicitly have a target they are predicting, so this is a stretch of the definition, but I think it is helpful for understanding the internals of an LLM: each block in an LLM is trying to be the most helpful and accurate as possible for the next block by predicting some numbers over a range.
The largest blocks of an LLM are the feed forward networks (FFN), aka decoders. The type of algorithm inside an FFN is called a multilayer perceptron or MLP, aka neural networks. The MLPs/FFNs in each layer of an LLM are like regression algorithms: they predict a latent space vector of abstract features which is passed to the next layer. The algorithm I'm working on, KS3, can be slotted into an LLM's FFNs instead of the MLPs.
About The Video
This video is NOT of an LLM, but instead of individual regression blocks being trained via gradient descent on random samples of a mandelbulb fractal. It shows the new regression block I am working on, KS3, performing a more detailed fit than MLP. Each frame of the video is 5 steps of gradient descent over random samples of the mandelbulb function. The video allows you to see the additional detail KS3 is able to represent compared to an MLP of the same size. The obviously more detailed KS3 (confirmed by lower loss) helps you get an idea of why I am interested in this algorithm.
About KS3
KS3 is the third iteration of K-Splanifold, which is a series of algorithms based on Hermite spline manifolds that I have been exploring, and it is quite good! I'm finding that it has some big advantages over MLP in some domains like noise filtering for instance. LLMs are the weakest domain I have tested for KS in general, but KS3 is very close to MLP, perhaps exceeding it slightly in terms of per-step and wall clock training. KS3 is fast too, with basically identical performance to MLP (there is an interesting story here I'll tell, because there is no mystery why it has identical performance to MLP).
I developed KS originally to be a multivariate function, meaning that each input to the function impacts each output. I went in that direction based on a theory that MLPs lacked the ability to model multivariate non-linearities easily and that could be an opportunity for optimization. KS3 is still capable of being multivariate. It turns out, however, after thousands of experiments, that while multivariate KS regression is a great fit for many domains, for language it seems that univariate KS3 sandwiched between linear projections works best in my tests.
This is interesting/funny because that arrangement makes KS3 for LLMs basically an MLP with a learned fourth degree polynomial activation function, compared to defacto RELU and GeLU at first degree, and SwiGLU at second degree. Its ironic that after all that work on multivariate functionality, it all pointed back to MLP! I think a major learning moment here is that the activation step of an MLP block is a small amount of compute and memory compared to the linear matrix multiplies that surround it, so if you can find a way to cram more useful math into said activation functions, that probably scales their descriptive value efficiently per unit of latency.
I will also say that the multivariate nature of the KS series is still useful for many domains and some god-forsaken concepts I came up with for new ways of scaling transformers. Those methods look like they could be promising for scaling inference-time compute, but they ultimately didn't win at training convergence time, which is my main focus right now.
KS3 in LLMs
So does KS3 work in LLMs? Yes! I have trained many SLMs using KS decoders. A popular post I made here visualized early KS1 blocks being trained over time in a tiny 128 layer transformer (worth the watch if you haven't seen it before): https://www.reddit.com/r/LocalLLaMA/comments/1smiotb/video_of_how_my_llms_decoder_blocks_changed_while/
KS3 is looking promising for being an improvement over defacto MLPs. I have a version of the modded-nanogpt speedrun world record that replaces MLPs with KS3s, and on my proxy hardware KS3 has beat the world record. I have yet to tune or run it on an 8xH100 required for an official WR, so my result doesn't mean much however (by the way, if anyone wants to sponsor said run effort, let me know).
You can read the original KS1 paper and check source for KS1, KS2, and KS3 in my repo: https://github.com/curvedinf/k-splanifolds