First Contact with ML Research - Ablating Regularizers on Spectral Koopman Attention

Introduction

In this project, fellow Stanford frosh Max Yang and I were under the supervison of Alexander Johansen, a Computer Science PhD student at Stanford. We ablated several regularizers for training the Spectral Koopman Attention module, which is a drop-in replacement for attention layers. Its main selling point is that it stores information about previous tokens through three sufficient statistics, which each new token additively updates. This means that unlike attention, it does not have to store representations of the entire preceding context and maintain a linearly growing KV cache, and unlike SSMs, the context is not held by a hidden state that decays multiplicatively with every new token.

SKA Architecture and Motivation for Regularization

For brevity, I omit the concrete formulation for SKA, which can be found in the Echo paper (Sridhar & Johansen, 2026). The two terms of interest for us in its computations were $G^{-1}=(W_k x_t)^{-1}$ and the Koopman power filter $A_W^K$.

Our ablation with regularizations was motivated by two concerns. First, if the keys of inputs close to being linearly dependent, $G$, which is the sum of their outer products, becomes ill-conditioned, which make $G^{-1}$ sensitive to small perturbations in the keys. we would like the keys $z_t=W_kx_t$ to not become linearly dependent and have similar magnitudes, we also want $W_k$ to be full rank and have singular near 1. Thus, we experimented with several forms of orthogonal regularizers from Bansal et al. 2018.

Second, if the Koopman operator $A_W$ has a high spectral radius, this would cause $A_W^K$ to explode, resulting in numerical instability of outputs and a loss of expressivity (as it would become a projection onto its largest eigenvalue). Thus, we also experimented with spectral regularization (Yoshida & Miyato 2017) and normalization (Miyato et al. 2018). Ablating these spectral control methods was my main focus.

Experiments

I made my initial comparisons by training single layers of the SKA module fitted with different regularization/normalization configurations on the MQAR task on the same configuration Sridhar & Johansen 2026 used to evaluate SKA.

I monitored two main trends:

1. Optimization of Training: whether the spectral normalization or regularization techniques would improve the final loss or the speed of training.
2. Health metrics: whether the condition number and spectral radius would become large over the course of training, as we feared, and whether regularization can help with this.

Note: We also monitor the learned gate value, which controls how much the module contributes to the residual stream. This value should steadily climb if the model is training as we expect. This is less important for the single-layer case, but in the next experiment we run, which involves inserting these layers in a frozen Qwen model, we do actually need to ensure that the performance is improving as we’d expect.

I make the following observations:

1. Results are homogenous: All of the regularization techniques that we tested had generally similar performance, except for weight decay, which made things much worse. The control case with no regularization (in purple) is one of the best-performing configurations.
2. Stability of health metrics without regularization: When loss decreases significantly (e.g. steps 0-1000, 3000-4500), the condition number tends to increase. However, it drops gradually once the loss plateaus. This seems to show that gradient descent naturally maintains a low condition number, since computations with ill-conditioned matrices likely increase the final loss. Spectral radius remained stable throughout

Given these observations, I suspected that the regularization does not yield meaningful improvements for SKA and run two follow-up experiments to confirm this:

1. Average the performance three of the best-performing configurations over five seeds and graph error bounds to see if regularizers actually improve results.
2. Vary the difficulty of MQAR up to its hardest setting to ensure we haven’t missed any behaviors that may emerge if the task changed in difficulty in various aspects (volcab size, distractor gap, and number of queries)

I waited for the results with baited breath, as I was about to see whether all of the analysis that we’d done with the linear algebra we’d learned in the first two quarters of freshman year and Claude, as well as all of the literature search + the nights we spent with Claude learning how to read the literature we found, was all about to come to nought.

Unfortunately, I found that the best technique (normalization) performed pretty much exactly the same as using no regularization or normalization whatsover. Indeed, spectral normalization overlapped completely with the control because the spectral radius barely fluctuated around 1.

At this point, it was quite clear that the null hypothesis was probably correct. Analysis beyond this point was mainly Max and I wanting to train stuff and graph things.

Messing Around

Indeed, undeterred and with several weeks left in the quarter (as we had sped through the initial analysis), I proceeded to sweep our regularizers/normalizers on every configuration of MQAR to check if perhaps there was anything we missed:

There were none.

Still undeterred, and having newly obtained access to a high-performance computing cluster at Stanford, Max and I decided that we would also do a sweep of the configurations on the model our PI was training for NeurIPS. We hoped that perhaps, language modeling, which was noisier than MQAR, may reveal some interesting behaviors worth analysis. Also, having access to remote GPUs for the first time in our lives (the toy models were small enough to train on our laptops), we wanted to learn how to work with those.

We pull Qwen-3.5-9b, replace 8 layers with SKA modules, freeze the original Qwen weights, and train for a preliminary 2000 steps on FineWeb-Edu:

While there appeared to be some differences in loss, this is mostly due to the spikiness of the graph and the average values across the (more optimal) configurations were not significantly far apart. Furthermore, as spectral radius of the Koopman operator barely exceeded 2.5, we decided that the original motivation for spectral regularization–preventing the Koopman operator from exploding–was an unnecessary concern.

Conclusion

This was a humbling experience. Having spent the first two weeks of this project trying to understand the math and figure out sources of instability from theoretical motivations, and then doing another week of literature search to find adequate regularizers, it was pretty awe-inpsirng to see that, in the end, the magic of gradient descent was just able to account for all of the potential structural issues we identified with the architecture.

On the technical side, I gained confidence working with GPUs, training models, collaboratively working in ML directories, and building test harnesses. One takeaway I’ll bring to future projects is that when runs seem close, the next best step is almost certainly to run a sweep through different seeds to ensure that whatever difference I’m observing is not noise.