Focusing on the moments that matter with WaveSSM

Transformers, the architecture that powers large language models (LLMs), are part of a larger group of deep learning technologies known as sequence models, which can track how pieces of information early in a sequence relate to later ones. This capability makes transformers good for processing text. But because the computational cost of running a transformer increases quadratically with sequence length, it’s inefficient to process very large chunks of sequential data, such as long audio recordings or continuous sensor data, with transformers.

Researchers have developed an alternative architecture known as state space models (SSMs) that can capture long-range dependencies in sequential data at a cost that only grows linearly with sequence length. SSMs do this by generating an overall representation of a signal that is continuously updated over time, with the representation acting like a running summary of the data.

But the ability of SSMs to compress a long signal into a holistic representation is also a limitation, explains Ruben Solozabal, previously a postdoctoral associate in Machine Learning at MBZUAI. “A traditional SSM generates a continuous projection that has a global temporal span, but it can’t focus on localized events in the data,” Solozabal says. And localized events are often where meaningful information in a signal can be found.

Consider the case of data gathered by a wearable heart monitor. Because traditional SSMs blend the entire recording into one global summary, brief spikes in heart rate get mixed together with everything else, leaving the model no way to isolate the anomalies a physician might care about.

Solozabal and researchers from MBZUAI are authors of a recent study that proposes a way to provide SSMs with both global and localized representations of sequential data. Called WaveSSM, the model uses what are known as wavelets to briefly store information from different temporal regions of the signal as different subsets of the overall state. “We use wavelet functions that have limited temporal support,” Solozabal says. “They appear and disappear very fast, so if there is an event, our model can catch it in particular dimensions.”

The work was presented at the 43rd International Conference on Machine Learning (ICML), in Seoul. Velibor Bojkovic, Hilal Alquabeh, Klea Ziu, Kentaro Inui, and Martin Takáč are co-authors of the study.

How WaveSSM works

WaveSSM builds on an SSM architecture known as S4. Both are what are known as linear time-invariant models. At each step, the input is wrapped into the model’s hidden state, and this process happens consistently across the sequence.

But in a traditional SSM like S4, two events from different moments end up layered on top of each other in the same compressed state, making it difficult to recover one without the other. WaveSSM stores these events separately while still generating an overall representation. “Other models look at events from a superposition state where the whole signal is imprinted,” Solozabal says. “You can’t pick one small window or another and do a simple operation because they appear in this superposition state.”

WaveSSM avoids this problem by giving different events their own coordinates, making the stored information “addressable,” as the researchers describe it in their study. The model can go back to a specific window and act on it without disturbing the rest.

What WaveSSM can do

The researchers conducted several experiments to compare the performance of different configurations of WaveSSM to S4 and a transformer. The WaveSSM configurations varied according to which kind of wavelet they used.

For one experiment, they created a long noisy signal with several flagged, non-overlapping windows in the sequence. The task was to retrieve the windows and add them together. Compared to S4, WaveSSM reduced the error on the task by roughly an order of magnitude.

The researchers also compared the performance of the models on real-world data.

On an electrocardiogram benchmark called PTB-XL, which is composed of nearly 22,000 recordings, the WaveSSM variants topped almost every diagnostic category, with a version of WaveSSM taking the best overall score. WaveSSM’s gains were largest where the clinical signal was transient, but it didn’t lead on a rhythm task, where a model called MS-SSM took the top score.

On raw-audio keyword classification, WaveSSM surpassed the S4 baseline at full resolution. But under a zero-shot setting where the audio was down-sampled to half its training frequency, S4 performed better.

The researchers also tested the models on a benchmark called Long Range Arena, designed to evaluate a model’s ability to capture long-range dependencies across different modalities, including symbolic reasoning and serialized images. WaveSSM beat S4 on four out of six tasks but struggled on two very-long-range settings, failing to learn the longest one at all.

The authors note that the mixed performance is due to a trade-off: the locality that helps on transient signals is in tension with the global, ultra-long-range dependencies those tasks demand.

Getting back to fundamentals

Solozabal says that much of the recent interest in SSMs has been due to a model called Mamba. Unlike WaveSSM and S4, Mamba is a linear time-varying SSM that bases its dynamics on the input itself. This approach, and scaling Mamba up to the size of frontier models, has allowed it to achieve a kind of attention to important parts of a sequence, but in a way that’s different from linear time-invariant SSMs or transformers, he explains.

With WaveSSM, however, Solozabal and his co-authors deliberately decided to stick with the time-invariant approach. And their findings show that they were able to gain a greater level of precision not by making the model’s dynamics adaptive to the input — and dealing with the computational costs that come with doing so — but by implementing the concept of wavelets.

Solozabal says that the results show that “it’s not always necessary to try to improve performance by scaling alone, but to return to theory because there is always something that can be added based on the fundamentals.”