Turning Back Time — Frequency Dimension Modelling with SSMs

“Written as a proposal for the experiments that followed”

There are few fields in Mathematics as routinely tested as Numerical Analysis, in some sense my very ability to write this proposal is only possible because of numerical integration techniques born from the field. The process of taking a continuous signal, discretizing it to map it onto the discrete computer world of 1’s and 0’s to remarkable accuracy is not just a solved problem but it’s perhaps one of the most fundamental and consequential solved problems for the development of our modern world. Along this line of thinking, researchers inevitably began thinking of the sequence modelling task as a similar problem, mapping a continuous/discrete signal found in nature to discrete computation friendly analog with remarkable accuracy. This is the problem of Quadrature and it is fundamentally what SSMs are designed to achieve.

In the interest of transparency and good writing, I will provide a small caveat early. Saying quadrature is a solved problem is a bit like saying baking croissants is a solved problem, while with the right technique and equipment it most certainly is, in the wrong hands with the wrong equipment and or wrong environment it can seem an impossible task. Nonetheless, the evolution of SSMs over the years has been mostly geared towards changing this, making SSMs more brownie-like than croissant-like, fool-proof. In doing so many of the strengths of SSMs have silently fallen away, yet its primary benefit seemed to persist until much later on.

The primary benefit of SSMs seemed to be their ability to handle long sequences with a fixed dimension hidden state, as opposed to transformers that scaled quadratically with sequence length. This requires a lot less of the HiPPO theory that birthed SSMs and the structural machinery was consistently stripped away for easier implementations. Now we’re at a place where SSMs can’t even match the performance of Transformers on most NLP tasks, putting the extremal length test beyond value.

Individually, a lot of my interest in ML stemmed from my interest in SSMs, with primarily a physics background, seeing reasoning posed as a quadrature problem signaled to me the start of something revolutionary. However, as noted above, the plateau of the rise of SSMs, short of expectations, researchers have begun to move toward variants of SSMs and combination architectures, yet Quadrature remains a solved problem. So what’s wrong?

Over the last few months, and for the foreseeable future, that question has been (will be) my primary research focus. I have implemented several SSM based architectures on a bevy of problems, from the stock market, to RUL (remaining useful life) prediction, and the sum total of my engineering experience is that SSMs hardly if ever blow transformers out the water. An equivalent statement for those not as mathematically inclined is two plus two equals five.

The brownification of SSMs has come at a fatal cost, they no longer guarantee a near perfect approximation of input signals. They no longer attempt to solve the solved problem. A lot of this was deliberately done, the mathematical rigor of accurate quadrature came with significant initialization and optimization burdens that couldn’t be justified in the all too important NLP task. Language could suffice with handicapped HiPPO theory, up to a point at least.

The question I’ve come to ask is whether that optimization and initialization burden can be eased by other means. My aim is to re-introduce some of the rigor stripped away over the years — and to keep training tractable, I have to confront the exploding and vanishing gradients that arise when eigenvalues clump together and become intolerant to small perturbations. My hypothesis is that holding the modes on a continuous spectrum is the heart of the problem: small perturbations can shift clustered eigenvalues unpredictably, and the gradients pay for it. A discrete spectrum would largely defuse this — perturbations smaller than the spacing between allowed values simply round away, so the modes stop drifting under them. My plan is to realize that discrete spectrum through quantization-aware training.