WaveSSM: Multiscale State-Space Models for Non-stationary Signal Attention
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
| 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 |
| 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 | 0.114 2 | 0.261 | 0.241 | 0.283 |
| Model | Raw(16kHz) | 0.5 (8kHz) |
|---|---|---|
| S4 | 96.08 | 91.32 2 |
| S4 | 95.27 | 91.59 1 |
| WaveSSM | 96.42 3 | 90.14 |
| WaveSSM | 96.47 2 | 90.84 3 |
| WaveSSM | 96.17 | 89.70 |
| WaveSSM | 96.09 | 89.61 |
| Method | ListOps | Text | Retrieval | Image | Pathfinder | PathX |
|---|---|---|---|---|---|---|
| Transformer | 36.37 | 64.27 | 57.46 | 42.44 | 71.40 | ✗ |
| S4 | 59.60 | 86.82 | 90.90 | 88.65 | 94.20 | 96.35 1 |
| S4 | 57.88 | 86.34 | 89.66 | 89.07 | 94.46 2 | ✗ |
| WaveSSM | 60.88 | 89.11 3 | 91.97 1 | 89.26 | 94.94 1 | ✗ |
| WaveSSM | 59.62 | 88.06 | 91.19 | 89.68 3 | ✗ | ✗ |
| WaveSSM | 60.13 | 89.33 1 | 91.88 2 | 89.35 | ✗ | ✗ |
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
| Task | Norm | Pre-norm | Dropout | LR | BS | Epoch | WD | ||||
| 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) |
| 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 |
| 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 |
| 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 | ✗ |
State Space Models
State space models represent sequences through evolving latent states, enabling efficient temporal modeling, system identification, and forecasting.
See all 295 papers