[Paper Review] Simplifying and Understanding State Space Models with Diagonal Linear RNNs

Ankit Gupta, Harsh Mehta+1 more|arXiv (Cornell University)|Dec 1, 2022

Topic Modeling4 citations

TL;DR

This paper proposes Diagonal Linear RNNs (DLR), a simplified, continuous-time formulation of state space models that eliminates the need for discretization, enabling a cleaner, more interpretable architecture. Despite its simplicity, DLR matches or exceeds the performance of prior SSMs and attention models on long-range sequence tasks, including Long Range Arena and raw speech classification, while maintaining O(L) inference complexity.

ABSTRACT

Sequence models based on linear state spaces (SSMs) have recently emerged as a promising choice of architecture for modeling long range dependencies across various modalities. However, they invariably rely on discretization of a continuous state space, which complicates their presentation and understanding. In this work, we dispose of the discretization step, and propose a model based on vanilla Diagonal Linear RNNs ($\mathrm{DLR}$). We empirically show that, despite being conceptually much simpler, $\mathrm{DLR}$ is as performant as previously-proposed SSMs on a variety of tasks and benchmarks including Long Range Arena and raw speech classification. Moreover, we characterize the expressivity of SSMs (including $\mathrm{DLR}$) and attention-based models via a suite of $13$ synthetic sequence-to-sequence tasks involving interactions over tens of thousands of tokens, ranging from simple operations, such as shifting an input sequence, to detecting co-dependent visual features over long spatial ranges in flattened images. We find that while SSMs report near-perfect performance on tasks that can be modeled via $ extit{few}$ convolutional kernels, they struggle on tasks requiring $ extit{many}$ such kernels and especially when the desired sequence manipulation is $ extit{context-dependent}$. Despite these limitations, $\mathrm{DLR}$ reaches high performance on two higher-order reasoning tasks $\mathrm{ListOpsSubTrees}$ and $\mathrm{PathfinderSegmentation} ext{-}\mathrm{256}$ with input lengths $8K$ and $65K$ respectively, and gives encouraging performance on $\mathrm{PathfinderSegmentation} ext{-}\mathrm{512}$ with input length $262K$ for which attention is not a viable choice.

Motivation & Objective

• To simplify state space models by removing the discretization step that complicates understanding and implementation.
• To propose a direct, interpretable formulation based on vanilla Diagonal Linear RNNs (DLR) as a starting point for long-range sequence modeling.
• To empirically evaluate DLR's expressivity and performance on synthetic and real-world long-sequence tasks, comparing it to SSMs and attention-based models.
• To investigate the limitations of SSMs in handling context-dependent and high-kernel tasks, and to characterize their expressivity relative to attention mechanisms.
• To provide a clean, accessible benchmarking framework using 13 synthetic sequence-to-sequence tasks for evaluating long-range modeling capabilities.

Proposed method

• Propose DLR as a direct parameterization of diagonal linear RNNs, bypassing continuous state space discretization.
• Use complex-valued transition matrices with eigenvalues initialized near the unit circle to enable long-range signal propagation over millions of steps.
• Introduce a periodic initialization scheme that allows a DLR of size N to express arbitrary convolutional kernels of length N.
• Evaluate DLR on a suite of 13 synthetic tasks ranging from simple sequence shifts to complex spatial reasoning in flattened images.
• Benchmark DLR against S4, S4D, Transformers, and LocalAttention on Long Range Arena, speech classification, and causal language modeling.
• Use input lengths up to 262K tokens to test scalability and performance on extremely long sequences.

Experimental results

Research questions

• RQ1 Can a direct Diagonal Linear RNN (DLR) formulation match or exceed the performance of discretized SSMs without the complexity of continuous-time discretization?
• RQ2 How does DLR’s expressivity compare to that of SSMs and attention-based models on tasks requiring long-range dependencies and complex interactions?
• RQ3 What are the limitations of current SSMs in handling context-dependent sequence manipulations or tasks requiring many convolutional kernels?
• RQ4 Can DLR generalize to real-world long-sequence tasks such as raw speech classification and causal language modeling with competitive performance and efficiency?
• RQ5 Does the periodicity of DLR’s parameterization enable a principled initialization that supports arbitrary convolutional kernel expression?

Key findings

• DLR achieves performance on par with or better than S4 and S4D on all Long Range Arena tasks, including Path-X and ListOps, with a simpler and more interpretable formulation.
• On the PathfinderSegmentation-512 task with 262K input length, DLR achieves strong performance despite attention being computationally infeasible at this scale.
• DLR outperforms LocalAttention on ListOps-SubTrees and PathfinderSegmentation-256, especially at longer sequence lengths, and maintains higher throughput.
• DLR matches or exceeds S4 and S4D on the 10-way Speech Commands dataset, achieving 97.1% accuracy with O(L) decoding complexity.
• In causal language modeling on PG-19, DLR achieves test cross-entropy losses comparable to Transformers (2.44–2.88) with 1.2B parameters and 1.0× throughput.
• DLR is more robust to layer norm placement than Transformers, which are known to be highly sensitive to such architectural choices.