W4S4: WaLRUS Meets S4 for Long-Range Sequence Modeling

2025-06-09

AI Summary

To address instability in state matrix initialization and rapid decay of long-range information in state space models (SSMs) for long-sequence modeling, this paper proposes WaLRUS—a novel SSM variant named W4S4—that pioneers the integration of redundant wavelet frames into SSM state dynamics design. W4S4 enables theoretically guaranteed stable diagonalization and efficient kernel computation without low-rank approximations, overcoming key limitations of mainstream initialization schemes such as HiPPO in capturing long-term dependencies. Leveraging wavelet-structured initialization, W4S4 consistently outperforms baselines including S4 on delayed reconstruction, time-series classification, and long-range modeling tasks, while significantly enhancing both modeling capacity and computational efficiency of deep SSMs. The core contribution lies in establishing a rigorous theoretical connection between wavelet analysis and SSM initialization, empirically validating its substantial performance gains across diverse sequence modeling benchmarks.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Intelligent Robots: State Estimation
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Improving State Space Models via wavelet-based initialization
• Enhancing long-range dependency handling in sequence modeling
• Optimizing SSM performance without low-rank approximations

Innovation

• Uses redundant wavelet frames for SSMs
• Enables stable diagonalization without low-rank
• Improves long-range information retention significantly