TLDR: This research paper investigates the ‘spectral bias’ in diagonal State Space Models (SSMs), revealing that current initialization methods are overly sensitive to the discretization step and often lead to inefficient frequency coverage. It proposes S4D-DFouT, a novel initialization scheme that directly places poles in the discrete Fourier domain, ensuring uniform frequency coverage and decoupling decay from frequency. This approach improves robustness and scalability, achieving state-of-the-art results on benchmarks like Long Range Arena, including training from scratch on PathX-256.
State Space Models (SSMs) have emerged as a powerful tool for understanding and modeling long sequences across various fields, from image processing to natural language understanding. At their core, SSMs process sequences by using a long-range convolution kernel, which is essentially a continuous-time linear dynamical system. This allows them to capture dependencies over extended periods and offers strong stability guarantees.
Traditionally, initializing the parameters of these models often relies on the HiPPO framework, which uses an online approximation of orthogonal polynomials. However, these traditional SSMs can be computationally expensive, especially for very long sequences, due to the complexity of propagating a dense state through the sequence.
The Rise of Diagonal SSMs and a Hidden Challenge
Recently, a simpler alternative, diagonal SSMs, has gained traction. These models achieve similar performance levels while being significantly more efficient. This efficiency comes from simplifying the kernel computation, often by restricting the state matrix to be diagonal. Despite their practical success, the theoretical reasons behind the effectiveness of these diagonal variants, particularly how the HiPPO framework applies to them, haven’t been thoroughly explored.
This research paper, titled Uncovering the Spectral Bias in Diagonal State Space Models, takes a crucial step to investigate these diagonal SSM initialization schemes from a frequency perspective. The authors, Ruben Solozabal, Velibor Bojkovic, Hilal AlQuabeh, Kentaro Inui, and Martin Takáˇc, aimed to systematically understand how to parameterize these models and uncover the inherent learning biases within them.
Unveiling the Spectral Bias
The key insight from their analysis is that existing initialization schemes for diagonal SSMs, whether based on HiPPO inverse-frequency laws or linear grids, introduce a significant challenge: a ‘spectral bias’. These methods typically define the state matrix in continuous time, which is then discretized with a learnable parameter, Delta (∆). This discretization process creates an entanglement between the decay rate and the oscillation frequency. As a result, adjusting the temporal resolution (∆) inadvertently alters both the decay and resonant frequencies of the system, making the model’s behavior highly sensitive to this parameter.
This sensitivity often leads to models compensating by spreading poles (which determine the system’s frequency response) across a wide range, resulting in over-parameterization, non-uniform spectral sensitivity, and even aliasing artifacts. Essentially, the model’s ability to capture long-range dependencies becomes highly dependent on a coincidental alignment between the dominant timescales in the input data and the poles in the system.
Introducing S4D-DFouT: A Universal Initialization
To address these limitations, the researchers propose a novel initialization scheme called S4D-DFouT (Diagonal State Space Model – Discrete Fourier Transform). This approach directly constructs the discrete-time state matrix, placing all poles uniformly around the unit circle in the complex plane, modulated by a shared exponential decay. The imaginary components ensure complete and uniform coverage of the frequency spectrum, while a learnable damping factor governs memory retention.
A significant advantage of S4D-DFouT is that it decouples the interdependency between the decay rate and the oscillation frequency. This design makes the initialization robust and less sensitive to the choice of the discretization step (∆). In a special case where the damping factor is zero, S4D-DFouT effectively reduces to the Discrete Fourier Transform (DFT), allowing the model to perfectly represent any circular convolutional kernel.
Furthermore, S4D-DFouT can be applied layer-wise, synchronizing the initialization across multiple SSMs in a layer to ensure a uniform grid of resonant frequencies, preventing redundancy and competition among different parts of the model.
Also Read:
- STDiff: A New Diffusion Framework for Robust Time Series Imputation in Industrial Systems
- Enhancing Speech Clarity: The Power of Straight Paths in AI Models
Experimental Validation and Key Findings
The paper demonstrates the effectiveness of S4D-DFouT through several experiments:
- Continuous Copying Task: This task, which tests long-term memory, showed that previous initializations failed when the discretization step was inappropriate. S4D-DFouT, however, successfully reconstructed the ideal delay kernel regardless of the initialization.
- Pixel-level Image Classification (sCIFAR): On serialized image data, the study revealed that kernels learned by S4D initializations often exhibit a ‘local attention’ profile, focusing on nearby pixels rather than truly global interactions. This suggests that SSMs can exploit local biases in data.
- Long Range Arena (LRA) Benchmark: S4D-DFouT achieved state-of-the-art results on this benchmark, which is designed to test long-range context capture. Notably, it was the first work to successfully train from scratch on PathX-256, a task previously requiring extensive self-pretraining. This highlights S4D-DFouT’s ability to scale to harder tasks without problem-specific tuning.
- Raw Speech Classification: The method also showed performance gains on the Speech Commands dataset, particularly in zero-shot resampling scenarios.
The research concludes that while LRA tasks are considered challenging, SSMs often succeed by leveraging local biases. S4D-DFouT’s uniform, alias-free spectral support eliminates the need for complex tuning of the discretization step, making it a more robust and scalable solution for diagonal SSMs. These insights also suggest that the perceived difficulty of some LRA tasks might be overestimated, calling for new benchmarks that truly challenge models to capture long-range dependencies without relying on local shortcuts.
