WaLRUS: Wavelets for Long-range Representation Using SSMs

TL;DR

The paper tackles online representation of long-range dependencies in sequences using state-space models (SSMs) and identifies limits of fixed bases in HiPPO-based approaches. It introduces WaLRUS, a wavelet-based SSM constructed via the SaFARi framework from Daubechies wavelets, enabling scalable, memory-efficient representations for non-smooth signals. WaLRUS provides scaled and translated variants whose updates follow the diagonalizable A structure, allowing fast online computation; the typical dynamic is $\frac{d}{dT} \vec{c}(T) = -A_{(T)} \vec{c}(T) + B_{(T)} u(T)$ with a reduced effective dimension $N_{\text{eff}}$. In experiments on M4, Speech Commands, and wavelet benchmarks, WaLRUS achieves consistently higher reconstruction fidelity and more robust spike-detection than Legendre and Fourier HiPPO variants. The work suggests WaLRUS as a practical, frame-agnostic initialization and online representation tool for SSM-based ML, with future work exploring other wavelets and frame choices and the associated trade-offs.

Abstract

State-Space Models (SSMs) have proven to be powerful tools for modeling long-range dependencies in sequential data. While the recent method known as HiPPO has demonstrated strong performance, and formed the basis for machine learning models S4 and Mamba, it remains limited by its reliance on closed-form solutions for a few specific, well-behaved bases. The SaFARi framework generalized this approach, enabling the construction of SSMs from arbitrary frames, including non-orthogonal and redundant ones, thus allowing an infinite diversity of possible "species" within the SSM family. In this paper, we introduce WaLRUS (Wavelets for Long-range Representation Using SSMs), a new implementation of SaFARi built from Daubechies wavelets.

Figures (8)

• Figure 1: An input signal comprising three random spikes is sequentially processed by SSMs and reconstructed after observing the entire input. Only the wavelet-based SSM constructed using WaLRUS can clearly distinguish adjacent spikes.
• Figure 2: Left: Elements of a Daubechies-$22$ wavelet frame, with father wavelet $\phi$, mother wavelet $\psi$, and two scales. Right: The scaled and translated $A$ matrices for WaLRUS with $N=21$.
• Figure 3: The kernel generated by HiPPO-LegT with window size $W=2000$ and representation size $N=500$. Three key non-ideal aspects of the kernel are noticeable. A) poor localization due to substantial non-zero values outside $W$, B) coefficient loss from at bottom left of the kernel, and C) coefficient loss at the bottom right of the kernel for $t \in (1500,2000)$.
• Figure 4: Left: The ideal kernels, which yield zero representation error, are shown for Translated-WaLRUS (using the D22 wavelet), HiPPO-LegT, and HiPPO-FouT. Right: The corresponding kernels generated by the translated models are presented for comparison. WaveT has superior localization within the window of interest compared to HiPPO-LegT and HiPPO-FouT.
• Figure 5: Comparing reconstruction MSE between WaveS, LegS, and FouS. Error bars represent the first and third quantile of MSE. WaveS produces the lowest MSE in each dataset.