Long post ahead - sorry about this, but might be worth it. You never know.
........
When we started this project, we were motivated by the following question: If we would have way more cheap compute and memory, in some time in the future, would that enable something better than training weights with gradient descent to optimize AI models? And to cut some long story short: there might be :) !
But let's start from the beginning:
In our weekly seminar at our Paradigms of Intelligence Team, we invited Kazuki Irie who presented results from his intriguing paper: The dual form of neural networks revisited: Connecting test time predictions to training patterns via spotlights of attention, 2022.
In that paper, Kazuki puts forward the simple duality of gradient descent on weight matrices within a linear layer y=Wx. What is that? Let's look at how weight matrices are optimized:
You might already see connections here to things well known in the sequence modeling literature, but let’s first look quickly on how trained weights are used during inference:
Although I thought I knew quite a bit about how NNs are optimized, I was quite astonished that I didn't know this fact about trained weight matrices: During inference, one can regard the computation of a trained linear layer y=Wx as the incoming activation retrieving and recombining backproped gradients w.r.t. y saved during training, i.e. the values, weighted by a simple dot-product attention on the respective inputs x, i.e. the keys! Although functionally this doesn't change anything, we can of course equivalently collapse the database into the trained W, it requires you to store all inputs as well as gradients in a database D containing all inputs and backproped gradients that you saw during training for that single linear layer - so maybe billions of vectors, for every weight matrix in your model! This is absolutely crazy, but remember - we don't care about compute and memory overheads in this project, for now ;)!
Well, given this duality, you might wonder if we can first and foremost get a better intuition of how neural networks and optimization works - I am not entirely sure about that, but it's at least cool to visualize what, e.g., learning rate scheduling, weight decay or momentum do, given this dual lense:
The red shades here tell you how much of that entry in the database is weighted differently, when compared to plaing GD, when doing inference at the end of training, by the linear attention mechanism! We are not sure what analyses could be done, but it seems there might be quite some untouched potential here!
Now, what we did in this project is to ask what happens if we replace the simple dot-product, i.e. linear attention with something more powerful, e.g., the softmax self-attention which, as everybody knows, revolutionized sequence modeling and empowers all of large-scale modern AI!
So from that perspective, we wanted to replace all weights with a powerful attention function - closing an existing gap on how to do "weights learning", of course with the hope of improving it! As everyone knows, global softmax self-attention is costly, compared to linear attention, but introduces a step function in performance and model behaviour - so we just pay the price in sequence modeling!
So how do we go about this? In reality, we went down a huge rabbit hole by trying to directly use the softmax self-attention as our kernel of choice - which, as we will see, is actually highly non-trivial (and therefore interesting)! So let's start with the simple RBF kernel which we only investigated after weeks of struggling ...
Now you could ask yourself if this algorithm actually makes sense? Can we simply drop in the RBF as our kernel and hope we are doing something sensible? We didn't derive this algorithm from any principle but simply replaced the linear kernel with some well-known non-linear kernel.
Luckily, it turns out that this algorithm is indeed a principled learning algorithm - and actually we could use any positive-definite kernel (PDK)!
Please check our paper for the details, but it's quite a simple proof - and as we then discovered a not widely known but published result in a paper from Grubb & Bagnell (2010)!
It's kind of cool and surprising, I think, what we ended up with: A growing neural network layer that, similar to trained weight matrices, retrieves and recombines backproped gradients - but now weighted with some non-linear attention function: Algorithm 5 of our paper :)
And what about our original quest: the softmax self-attention? Well, problems arise :) and we learned the hard way. Basically, Max tried and well things kind of went downhill (or uphill on the loss landscape).
But, it's always amazing to have people on the team (not me :)) that are actually real mathematical geniuses: after some time wondering what the hell was going on, we started to take a step back and look at the optimization of the softmax function closely. This ended up, with a little help of AI ;), to be a 68 page thesis entitled "Foundations of Deep Growing Softmax Attention Networks: Theory and (Possibly) Practical Algorithms"
Ok, so let's dive in!
The first insight, which is quite straight forward to see is that, naively plugging in, the softmax self-attention kernel in the learning rule derived for linear attention does not do anything good!
What we can see here, is that adding a new entry in the database is definitely not a continuous operation. The update, even if we reduce the learning rate to 0, induces a nonzero change in the function - as the additional E term doesn't vanish!
So, basically what we went out to do, is to create a new softmax database structure which allows for adding new entries in a continuous way. One elegant way to do this, is to introduce a weighted softmax layer:
With this new softmax function, we can actually add entries to the database in a smooth way, but how and what triples should we add, so that this leads to principled optimization, is not clear at all a priori.
Due to space constraints, I refer to the paper to study the mathematics, but simply put, we can setup the space of infinite databases, weighted by w, in which doing functional gradient descent on these weights, therefore choosing smoothly some databases, corresponds actually when discretizing to adding a certain new triple to our database! I am brushing over important details here, but the resulting learning rule is quite similar to the RBFs one.
This results in our "Growing Weighted Softmax Attention layer", i.e., Algorithm 3 of our paper which allows replacing any weight matrix, i.e., linear attention over inputs and gradients, with a weighted softmax kernel - and optimizing it in a principled way. In short: We could now build and optimize AI models without any weight matrices.
We call this paradigm Retrieval-Centric Deep Learning (RCDL)!
Well, we are almost done with equations and gradient descent in infinite dimensional spaces and the like, but let's ask another time - are there more kernels to consider, on top of RBF and softmax?
Sequence layer enthusiasts like myself, who worked on linear attention variants, would wonder if we can use for example DeltaNet or MesaNet (we were a bit biased towards trying those two :)) as learning rules to optimize weights?
It turns out that yes, you can port these modern sequence layers from in-context learnign to weight learning in a principled way - which, generally, might provide us with an angle to come up with new powerful (weight) optimizers.
The cool part here, I think, is that you don't motivate optimizing weights by GD - but how to find weight matrices that map well from keys to values, i.e., backpropagated gradients! So we are still in database attention land.
Note that in the image above, I show only the definition of the kernels that correspond to Mesa and Delta, but there exist, of course, incremental ways to update, as in their respective sequence layer formulations, the resulting weights incrementally batch by batch.
Ok, ok, now, let’s look at all of these new layers finally in practice! Some sanity checks first, on for our novel weighted softmax layer:
In a classic student-teacher setting, a single of our novel weighted softmax layer optiumizes well. While the overparametrized ReLU student is fitting the teacher very well, at the end of training, we do so without matching the non-linearity of the teacher. Also, we are beating significantly a parametrized softmax baseline, i.e., FFN with softmax non-linearity, even if we increase the width so it matches the fully grown weighted softmax layers. Note, that the student ReLU MLP shows some form of saddle point behaviour here while we do not, mhhhm :)
On to good old CIFAR.
On a classic residual MLP backbone, where we replace FFNs with our RCDL layers, our Mesa optimizer (really bad name, nothing mesa here at all & actually a reinvention of a not widely known but great FOOF optimizer of our friend Frederik Benzing) shows interesting competitive behaviour with Muon. Also, our RCDL layers, which are significantly more costly than the parametric baselines, match but do not outperform these strong baselines. Note that we did some quite intense hyperparameter and model architecture scanning here, while showing the best setting by overfitting on the test set with every method. If Muon wouldn't exist, this would be an impressive first result. Note that Muon, on this architecture and dataset, introduces a significant step function improvement on e.g. Adam and RMSprop.
One thing to note here: RCDL layers use, for now, plain gradient descent which arguably is a big disadvantage compared to weights trained with Adam and Muon - an Adam-like optimization of growing RCDL would somewhat need to operate in the functional space, which we haven't fully developed yet. Some first results, including hybrids of growing RCDL and parametric learning rules, are already in the paper.
We also provide some analyses on the temperature influence in the softmax kernels, as well as observe very different generalization behaviour then parametric baseline.
Despite these let's call them mixed empirical results, we got nice and interested reviewers & were accepted at NeurIPS - we are happy to present the paper in December in Sydney and Paris!
Now, let’s get to hot off the press, unpublished & spicy empirical results on language! Well, synthetic language 🙂.
Here, we started with a beautiful result from our friend Nicolas Zucchet, who looked carefully at the training dynamics of an adaptation of a biography dataset from the great Physics of Language modeling saga from Zeyuan Allen-Zhu & Yuanzhi Li. Here a short description:
What Nicolas observed is that transformers trained on the biography dataset acquire knowledge in three phases:
- Model learns generic attribute value statistics.
- Performance plateaus at a level achievable by an ideal model lacking individual-specific knowledge.
- Model learns associations between subjects and attributes.
Intriguingly, the plateau’s duration is nearly proportional to the number of individuals. Here a reproduction of the result with modern optimization and architecture choices, e.g., a Gemma 3 backbone.
Furthermore, Nicolas found that fine-tuning on new individuals leads to strong “forgetting”, i.e., he observes a rapid drop in performance on individuals learned during pre-training, with new knowledge acquisition being a slower process, when fine-tuning on a new set of individuals. Also, incorporating classic replay of pre-training data partially mitigates the final performance drop, but not the initial decline
Ok, now let’s see what happens if we use RCDL layers on this task.
But first, note that for any reasonable training duration and model architecture used for language experiments, if we replace any weight matrix with RCDL layers, this leads to a whole datacenter explosion.
So, we asked ourselves if we can develop more reasonable RCDL layers which are at least somewhat manageable w.r.t. to compute and memory requirements. So let’s go back to the theory drawing board.
It is well known, also tightly related to the fact we can do principled GD on any PDK kernel, that we can approximate any PDK with a certain finite feature vectors, depending on which kernel we want to approximate. After this approximation, doing GD in this (finite dimensional) space is equivalent to approximating our growing RCDL layers rule for PDKs.
Note, this results simply in a FFN with a non-standard cosine non-linearity, in the case of growing RBF layers, with some special parametrization - kinda cool, as this approximates, in the limit of large number of features, our growing RBF layer.
There are many deep learning papers that use these kinds of insights to design non-linearities and FFN layers, in particular, the connection of random features with non-parametric kernel method is well known, but we believe the connection to growing RBF layers could be novel, especially in context of the following empirical results.
Now, here is a quick overview of how we integrated 1) GeGlu layers 2) the Cos Bochner approximation and 3) the full blown growing RBF layers into a Gemma 3 transformer backbone, namely by replacing FFN layers with either of these.
To compare fairly baselines with the growing RBF layers, we again scanned the width of the parametric GeGlu as well as parametric Cos Bochner layer such that the final inference costs in terms of flops are (kinda) comparable to a fully grown RCDL RBF layer.
Here, first some results on the 2L model, trained on 16k population size:
As we can see, the plateau shrinks drastically for the RBF as well as wide Cos Bochner layers. Intriguingly, Cos Bochner improves with increased width matching growing RBF layers, pointing strongly towards the fact that they indeed approximate them better with more width, while standard GeGLU FFNs reach quickly a ceiling of performance.
For now, we do not know exactly why this happens and we think it's mysterious. Note however, that the growing RBF layer as well as the parametrized Cos Bochner layers have huge width. Although a cool result, these layers are still not ready to be incorporated into models at scale!
These findings beautifully translate to different population sizes and also the hybrid architectures, see the architecture image above for more info, where we somewhere in the model replace standard GeGLU FFNs with a single or two growing RBF or Cos Bochner layer:
Finally, in the fine-tuning setting, the RBF, and also the Cos Bochner, seem to do better than the GeGLU. We are still exploring the interesting phenomenon that the RCDL layers can “unforget” without replay.
Alright, this was a long one. Quick summary:
- We introduce the retrieval-centric deep learning paradigm which allows to replace weight matrices with growing databases consisting of inputs and backpropagated gradients - which are recalled and interpolated at inference time.
- The RCDL framework not only brings non-linear kernels to “weights” learning, which scale for now to immense size, but also allow you to motivate new optimizers from the lens of database compression.
- On top of showing intriguing behaviour on simple regression and image classification tasks, growing RCDL layers and its approximations improve the speed of fact acquisition in language significantly.
- Nevertheless, to observe significant advantages, we need to scale the width to the millions yielding very expensive layers.
We continue to work on the language modeling results which we will write up soon into a paper. Furthermore, we are of course trying to find ways to make the RBF layers grow sublinearly, use, e.g. KV-cache compression techniques and develop more powerful optimizers for our RBF layers as well as understand why and when they show advantages to common GD optimized weight matrices. Lot's of fun ahead!
-------
This was an amazing collaboration among many great friends: the jax tunneler Maximilian Schlegel, the math wizzAIrd Rajai Nasser, my good old mesa friend Seijin Kobayashi (@SeijinKobayashi), our dropmaxer Yanick Schimpf (@yaschimpf), weightdeltadecay fanatic Oliver Sieberling (@osieberling), the amazing Robert Obryk, freshly baked Prof. Kazuki (@kzkirie) and my long term partner in crime João Sacramento. The for-now-unpublished language modeling results were produced by the amazing Kaitlin Maile (@kaitlinmaile), Robert, Rajai, João and also myself.
There are tons of related work, including, of course, kernel machines and gradient boosting techniques, which I didn’t mention in this article, please check the paper for a discussion, and probably we also missed citing relevant works - please reach out to us if that is the case.
Code here, Arxiv link here - Happy softmaxing everyone :)
