WaveSSM: Multiscale State-Space Models for Non-stationary Signal Attention

Ruben Solozabal; Velibor Bojkovic; Hilal Alquabeh; Klea Ziu; Kentaro Inui; Martin Takac

cs.LGFeb 25, 2026

WaveSSM: Multiscale State-Space Models for Non-stationary Signal Attention

Authors: Ruben Solozabal, Velibor Bojkovic, Hilal Alquabeh, Klea Ziu, Kentaro Inui, Martin Takac

Organizations: MBZUAI · RIKEN AIP

Abstract

State-space models (SSMs) have emerged as a powerful foundation for long-range sequence modeling, with the HiPPO framework showing that continuous-time projection operators can be used to derive stable, memory-efficient dynamical systems that encode the past history of the input signal. However, existing projection-based SSMs often rely on polynomial bases with global temporal support, whose inductive biases are poorly matched to signals exhibiting localized or transient structure. In this work, we introduce WaveSSM, a collection of SSMs constructed over wavelet frames. Our key observation is that wavelet frames yield a localized support on the temporal dimension, useful for tasks requiring precise localization. Empirically, we show that on equal conditions, WaveSSM outperforms orthogonal counterparts as S4 on real-world datasets with transient dynamics, including physiological signals on the PTB-XL dataset and raw audio on Speech Commands.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison between global polynomial bases and localized wavelet frames. Left: Legendre orthogonal-basis functions are utilized in the HiPPO framework to construct an online projection of the input signal with global support. Right: wavelet-based frames (Morlet in the example) exhibit more localized time support, promoting selective representations with the ability to attend transient events more effectively.

Panel of Figure 2 — Figure 2: Left: Overview of the WaveSSM framework. The input signal is projected onto a wavelet frame, inducing a continuous-time state-space representation via differentiation of the projection. The resulting SSM maintains a compact latent state with addressable, time-local components, which can be decoded independently or processed for sequence modeling. Right: Visualization of the Jacobian $\partial h_{T}/\partial x_{t}$ , illustrating the influence of inputs at different time steps on the final hidden state. Results are shown for Morlet wavelets of order $N=100$ under the scaled measure $\mu_{sc}$ ; see Appendix D.1 for details.

Panel of Figure 3 — Figure 3: Visualization of the transition matrices $A$ obtained for the wavelet families for $N=40$ . From left to right: Morlet, Gaussian-derivative (Gauss), Mexican hat (Mexhat), Discrete Prolate Spheroidal Slepians (DPSS), and Daubechies (db6).

Figure 4: Condition number of the discretized frame operator

\kappa(S)

as a function of the state dimension

N

. Raw frames (dashed) can become ill-conditioned as

N

grows, while tightening

F

(in solid lines) mitigates this degradation. The improvement is especially pronounced for Morlet, Gaussian, Daubechies, and Mexican-hat frames.

Figure 5: Approximation errors on various frames for the two-step piecewise constant function, see Appendix C for details.

Panel of Figure 6 — Figure 6: Continuous Window Copying Task. Left: A white-noise signal of length- $L{=}4000$ is streamed to the model together with sparse random markers defining non-overlapping windows of interest $W\in[2,20]$ . After a variable delay, the model must selectively retrieve the marked segments and apply a prescribed operation (i.e., addition) to generate the output. Right: Results reported as reconstruction MSE in log scale as a function of the number of windows $W$ . For WaveSSM, we report the mean with confidence intervals across the five wavelet variants introduced in this work under the scaled measure.

Model	Diag	Sub-Diag	Super-Diag	Form	Rhythm	Overall
Transformer ( Vaswani et al., 2017b )	0.876	0.882	0.887	0.771	0.831	0.849
Spacetime ( Zhang et al., 2023 )	0.941	0.933	0.929	0.883	0.967	0.931
S4 ( Gu et al., 2022b )	0.939	0.929	0.931 3	0.895	0.977 2	0.934
InceptionTime ( Fawaz et al., 2020 )	0.931	0.930	0.921	0.899	0.953	0.927
Wavelet features ( Strodthoff et al., 2021 )	0.855	0.859	0.874	0.757	0.890	0.847
Mamba ( Gu & Dao, 2023 )	0.929	0.905	0.912	0.876	0.952	0.915

Table 1: AUROC (

\uparrow

) on the PTB-XL dataset for ECG multi-label/multi-class classification. The three best performing methods are highlighted in red 1 (First), blue 2 (Second), and violet 3 (Third). Result reproduced from ( Karami et al., 2025 ) .

Model	ETTh 1	ETTm 1	Weather	ECL
S4	0.116 3	0.292	0.245	0.432
Informer	0.269	0.512	0.359	0.582
LogTrans	0.273	0.598	0.388	0.624
Reformer	2.112	1.793	2.087	7.019
Prophet	2.735	2.747	3.859	6.901
WaveSSM ${}_{\text{MorletS}}$	0.114 2	0.261	0.241	0.283

Table 2: Univariate long sequence time-series forecasting results on Informer dataset reported as MSE (↓). Highlighted in red 1 (First), blue 2 (Second), and violet 3 (Third). Full results in Appendix G .

Model	Raw(16kHz)	0.5 $\times$ (8kHz)
S4 ${}_{\text{LegS}}$	96.08 $\scriptstyle\pm 0.15$	91.32 $\scriptstyle\pm 0.17$ 2
S4 ${}_{\text{FouT}}$	95.27 $\scriptstyle\pm 0.20$	91.59 $\scriptstyle\pm 0.23$ 1
WaveSSM ${}_{\text{MorletS}}$	96.42 $\scriptstyle\pm 0.26$ 3	90.14 $\scriptstyle\pm 0.20$
WaveSSM ${}_{\text{GaussS}}$	96.47 $\scriptstyle\pm 0.16$ 2	90.84 $\scriptstyle\pm 0.13$ 3
WaveSSM ${}_{\text{MexhatS}}$	96.17 $\scriptstyle\pm 0.16$	89.70 $\scriptstyle\pm 0.31$
WaveSSM ${}_{\text{dpssS}}$	96.09 $\scriptstyle\pm 0.28$	89.61 $\scriptstyle\pm 0.15$

Table 3: SC35 classification. Classification accuracy (

\uparrow

) in unprocessed audio signals (

L=16{,}000

) and test-time frequency shift at half the training frequency. Highlighted in red 1 (First), blue 2 (Second), and violet 3 (Third). All results are reproduced in this work for 3 different seeds.

Method	ListOps ${\scriptstyle(2048)}$	Text ${\scriptstyle(4096)}$	Retrieval ${\scriptstyle(4000)}$	Image ${\scriptstyle(1024)}$	Pathfinder ${\scriptstyle(1024)}$	PathX ${\scriptstyle(16,384)}$
Transformer	36.37	64.27	57.46	42.44	71.40	✗
S4 ${}_{\text{LegS}}$	59.60	86.82	90.90	88.65	94.20	96.35 1
S4 ${}_{\text{FouT}}$	57.88	86.34	89.66	89.07	94.46 2	✗
WaveSSM ${}_{\text{MorletS}}$	60.88	89.11 3	91.97 1	89.26	94.94 1	✗
WaveSSM ${}_{\text{GaussS}}$	59.62	88.06	91.19	89.68 3	✗	✗
WaveSSM ${}_{\text{MexhatS}}$	60.13	89.33 1	91.88 2	89.35	✗	✗

Table 4: Accuracy (

\uparrow

) on the Long Range Arena (LRA) benchmark. “✗” indicates infeasible to learn a non-trivial solution. Highlighted in red 1 (First), blue 2 (Second), and violet 3 (Third). Extended comparison can be found in Table S 8 and hyperparameter settings are detailed in Table S 5 .

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Comparison of approximations of $f_{\text{one}}$ target with db6 and Legendre frames on various budgets $N$ .

Panel of Figure 8 — Figure 8: Comparison of approximations of $f_{\text{one}}$ target with gaus2 and Legendre frames on various budgets $N$ .

Panel of Figure 9 — Figure 9: Comparison of approximations of $f_{\text{one}}$ target with gaus3 and Legendre frames on various budgets $N$ .

Panel of Figure 10 — Figure 10: Comparison of approximations of $f_{\text{one}}$ target with mexh and Legendre frames on various budgets $N$ .

Panel of Figure 11 — Figure 11: Comparison of approximations of $f_{\text{one}}$ target with morl and Legendre frames on various budgets $N$ .

Panel of Figure 12 — Figure 12: Comparison of approximations of $f_{\text{two}}$ target with db6 and Legendre frames on various budgets $N$ .

Panel of Figure 13 — Figure 13: Comparison of approximations of $f_{\text{two}}$ target with gaus2 and Legendre frames on various budgets $N$ .

Panel of Figure 14 — Figure 14: Comparison of approximations of $f_{\text{two}}$ target with gaus3 and Legendre frames on various budgets $N$ .

Panel of Figure 15 — Figure 15: Comparison of approximations of $f_{\text{two}}$ target with mexh and Legendre frames on various budgets $N$ .

Panel of Figure 16 — Figure 16: Comparison of approximations of $f_{\text{two}}$ target with morl and Legendre frames on various budgets $N$ .

Panel of Figure 17 — Figure 17: Approximation errors on various frames for piecewise constant functions. The $L^{2}$ approximation error is shown as a function of the number of retained coefficients $N$ on a log–log scale for a one-step function (left) and a two-step function (right). Results are reported for Legendre polynomials and several wavelet-based frames, illustrating differences in convergence behavior and the effect of increasing discontinuity complexity.

Figure 18: Representative atoms for Morlet, Gaussian-derivative, Mexican hat, Daubechies, and DPSS (rows) at fixed scales

s\in\{1,\,0.315,\,0.125\}

(columns), which control temporal width. All atoms are

\ell_{2}

-normalized; therefore, as a consequence of energy normalization, peak magnitudes vary across scales.

Task	$L$	$H$	$N$	Norm	Pre-norm	Dropout	LR	BS	Epoch	WD	$(\Delta_{\min},\Delta_{\max})$
ListOps	6	256	64	BN	False	0	0.01	50	40	0.05	(0.001, 0.1)
Text	6	256	64	BN	True	0	0.01	16	32	0.05	(0.001, 0.1)
Retrieval	6	256	64	BN	True	0	0.01	64	20	0.05	(0.001, 0.1)
Image	6	512	64	LN	False	0.1	0.01	50	200	0.05	(0.001, 0.1)
Pathfinder	6	256	128	BN	True	0	0.001	64	200	0.03	(0.001, 0.1)
PathX	6	256	128	BN	True	0	0.0005	32	50	0.05	(0.001, 0.1)

Appendix

Table 5: Hyperparameters used for the reported results, shared across WaveSSM and S4. L is the number of layers; H, the embedding size; N, the order of the SSM; Dropout, the dropout rate; LR, the learning rate; BS, the batch size; Epochs, the training epochs; WD, weight decay; and

(\Delta_{\min}

\Delta_{\max})

, the range for the discretization step-size.

	ETTh 1					ETTh 2					ETTm 1
Model	24	48	168	336	720	24	48	168	336	720	24	48	96	288	672
S4	0.061 1	0.079 1	0.104 2	0.080 1	0.116 3	0.095 3	0.191	0.167	0.189	0.187	0.024	0.051	0.086	0.160 2	0.292
Informer	0.098	0.158	0.183	0.222	0.269	0.093 1	0.155	0.232	0.263	0.277	0.030	0.069	0.194	0.401	0.512
LogTrans	0.103	0.167	0.207	0.230	0.273	0.102	0.169	0.246	0.267	0.303	0.065	0.078	0.199	0.411	0.598
Reformer	0.222	0.284	1.522	1.860	2.112	0.263	0.458	1.029	1.668	2.030	0.095	0.249	0.920	1.108	1.793
LSTMa	0.114	0.193	0.236	0.590	0.683	0.155	0.190	0.385	0.558	0.640	0.121	0.305	0.287	0.524	1.064

Appendix

Table 6: Univariate Informer’s long sequence time-series forecasting results (MSE ↓).

	ETTh 1					ETTh 2					ETTm 1
Model	24	48	168	336	720	24	48	168	336	720	24	48	96	288	672
S4	0.525	0.641	0.980	1.407	1.162	0.871	1.240	2.580	1.980 1	2.650 3	0.426	0.580	0.699	0.824	0.846
Informer	0.577	0.685	0.931	1.128	1.215	0.720	1.457	3.489	2.723	3.467	0.323 1	0.494	0.678	1.056	1.192
LogTrans	0.686	0.766	1.002	1.362	1.397	0.828	1.806	4.070	3.875	3.913	0.419	0.507	0.768	1.462	1.669
Reformer	0.991	1.313	1.824	2.117	2.415	1.531	1.871	4.660	4.028	5.381	0.724	1.098	1.433	1.820	2.187
LSTMa	0.650	0.702	1.212	1.424	1.960	1.143	1.671	4.117	3.434	3.963	0.621	1.392	1.339	1.740	2.736

Appendix

Table 7: Multivariate Informer’s long sequence time-series forecasting results (MSE ↓).

Model	ListOps	Text	Retrieval	Image	Pathfinder	PathX-128
	(2,048)	(4,096)	(4,000)	(1,024)	(1,024)	(16,384)
Transformer ( Vaswani et al., 2017a )	36.37	64.27	57.46	42.44	71.40	✗
Reformer ( Kitaev et al., 2020 )	37.27	56.10	53.40	38.07	68.50	✗
BigBird ( Zaheer et al., 2020 )	36.05	64.02	59.29	40.83	74.87	✗
Linear Trans. ( Katharopoulos et al., 2020 )	16.13	65.90	53.09	42.34	75.30	✗
Performer ( Choromanski et al., 2021 )	18.01	65.40	53.82	42.77	77.05	✗

Appendix

Table 8: Extended comparison on LRA.

State Space Models

State space models represent sequences through evolving latent states, enabling efficient temporal modeling, system identification, and forecasting.

See all 295 papers

MuonSSM: Orthogonalizing State Space Models for Sequence Modeling
Jun 29, 2026Thai-Khanh Nguyen, Ngoc-Bich-Uyen Vo, Thieu N. Vo +2Long-Context Modeling Sequence Modeling
Faculty of Information Technology, Dainam University, Hanoi University of Science and Technology, Hanoi 10000, Vietnam · Faculty of Artificial Intelligence, Posts and Telecommunications Institute of Technology, Hanoi 10000, Vietnam · Department of Computer Science, University of Bath, Bath BA2 7AY, UK +1
LogSig-SSM: Time-Series Modelling with Multi-Scale Log-Signature Compression for State-Space Models
Oct 4, 2026Felix Oury, Nicolas Calvo Peiro, Reiko J. TanakaIrregular Time-Series Modeling Multivariate Time Series Forecasting
Department of Computing Imperial College London · Department of Bioengineering Imperial College London
DSSMs: State Space Models with Explicit Memory via Delay Differential Equations
Jul 11, 2026Yixiao Qian, Song Chen, Jiaxu Liu +2Dynamical Systems Long-Context Modeling
College of Control Science and Engineering, Zhejiang University · Department of Mathematics, National University of Singapore · School of Mathematical Sciences, Zhejiang University

Jun 29, 2026J/K move · Enter open · S save

MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

Thai-Khanh Nguyen, Ngoc-Bich-Uyen Vo, Thieu N. Vo, Tan M. Nguyen+1

Faculty of Information Technology, Dainam University, Hanoi University of Science and Technology, Hanoi 10000, Vietnam · Faculty of Artificial Intelligence, Posts and Telecommunications Institute of Technology, Hanoi 10000, Vietnam · Department of Computer Science, University of Bath, Bath BA2 7AY, UK+1

Long-Context Modeling Sequence Modeling State Space Models

Open paper AskPDF arXiv

State space models (SSMs) have emerged as efficient linear-time alternatives to attention for long-sequence modeling. However, existing SSMs often suffer from instability and memory degradation over extended horizons due to poorly conditioned first-order updates and unbalanced update geometry. We introduce MuonSSM, a general framework that stabilizes SSM training by explicitly conditioning the geometry of memory updates rather than the recurrent transition matrix. MuonSSM augments SSMs with a momentum-based pathway and a lightweight Newton Schulz transformation on low-rank input injections, yielding bounded and spectrally conditioned updates while preserving parallel scan complexity. Theory shows that MuonSSM improves gradient propagation, mitigates spectral amplification, and enriches memory representations over long horizons. Extensive experiments across language, vision, and time-series benchmarks show consistent gains in accuracy, robustness, and long-context performance when integrated into diverse SSM backbones. These results establish geometric conditioning of updates as a principled pathway to stable, scalable sequence modeling.

Abstract

State Space Models

Explore similar work

MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

LogSig-SSM: Time-Series Modelling with Multi-Scale Log-Signature Compression for State-Space Models

DSSMs: State Space Models with Explicit Memory via Delay Differential Equations