W4S4: WaLRUS Meets S4 for Long-Range Sequence Modeling

TL;DR

This paper introduces W4S4, a novel state space model built on wavelet frames, which enhances long-range sequence modeling by improving stability, efficiency, and information retention over existing models like S4.

Contribution

The paper presents W4S4, a wavelet-based SSM that offers stable diagonalization, fast kernel computation, and better long-term information retention, advancing the capabilities of sequence modeling architectures.

Findings

W4S4 outperforms HiPPO-based SSMs in long-horizon tasks.

W4S4 achieves consistent improvements on classification benchmarks.

W4S4 demonstrates enhanced long-range sequence modeling in various experiments.

Abstract

State Space Models (SSMs) have emerged as powerful components for sequence modeling, enabling efficient handling of long-range dependencies via linear recurrence and convolutional computation. However, their effectiveness depends heavily on the choice and initialization of the state matrix. In this work, we build on the SaFARi framework and existing WaLRUS SSMs to introduce a new variant, W4S4 (WaLRUS for S4), a new class of SSMs constructed from redundant wavelet frames. WaLRUS admits a stable diagonalization and supports fast kernel computation without requiring low-rank approximations, making it both theoretically grounded and computationally efficient. We show that WaLRUS retains information over long horizons significantly better than HiPPO-based SSMs, both in isolation and when integrated into deep architectures such as S4. Our experiments demonstrate consistent improvements across delay reconstruction tasks, classification benchmarks, and long-range sequence modeling, confirming that high-quality, structured initialization enabled by wavelet-based state dynamic offers substantial advantages over existing alternatives. WaLRUS provides a scalable and versatile foundation for the next generation of deep SSM-based models.