PaperlayerView exact revision on arXiv

Source-linked AI summary

Diagonal State Spaces are as Effective as Structured State Spaces

Ankit Gupta, Albert Gu, Jonathan Berant

arXiv:2203.14343v3cs.LGcs.CL

TL;DR

Long-range reasoning remains difficult for attention-based models despite their effectiveness on short-range interactions. This paper introduces Diagonal State Space (DSS), a fully diagonal alternative to S4, and finds comparable performance on Long Range Arena while simplifying formulation, implementation, and analysis.

  • Problem

    Attention-based models have inadequate performance on benchmarks requiring long-range reasoning, despite their effectiveness for short-range interactions.

  • Method

    DSS removes S4’s low-rank correction by enforcing diagonal state matrices, yielding a simpler model that remains as expressive as general state spaces.

  • Results

    DSS achieves 81.88 average accuracy across six Long Range Arena tasks, comparable to S4’s 80.21 and above the best Transformer variant’s 61.41.

  • Takeaways & Limitations

    DSS provides a simple and effective approach for modeling long-range interactions across text, images, and audio.

  • Takeaways & Limitations

    The evaluation covers sequence-level classification tasks, leaving token-level generation and large-scale DSS pretraining for future work.

Abstract

from arXiv · show

Modeling long range dependencies in sequential data is a fundamental step towards attaining human-level performance in many modalities such as text, vision, audio and video. While attention-based models are a popular and effective choice in modeling short-range interactions, their performance on tasks requiring long range reasoning has been largely inadequate. In an exciting result, Gu et al. (ICLR 2022) proposed the $\textit{Structured State Space}$ (S4) architecture delivering large gains over state-of-the-art models on several long-range tasks across various modalities. The core proposition of S4 is the parameterization of state matrices via a diagonal plus low rank structure, allowing efficient computation. In this work, we show that one can match the performance of S4 even without the low rank correction and thus assuming the state matrices to be diagonal. Our $\textit{Diagonal State Space}$ (DSS) model matches the performance of S4 on Long Range Arena tasks, speech classification on Speech Commands dataset, while being conceptually simpler and straightforward to implement.

1 Introduction

Attention-based models struggle with long-range reasoning despite their success across modalities, while S4 improves long-range performance but is complex. DSS shows that a fully diagonal state-space parameterization can match S4 while being simpler and effective across several tasks.

  • Motivation: Omega(L^2) attention complexity becomes prohibitive for long inputs, and Transformer variants still report inadequate long-range reasoning performance.The cited benchmarks include Long Range Arena and SCROLLS.
  • Prior work: S4 uses linear state spaces instead of attention and advances Long Range Arena performance by 19 accuracy points over the best Transformer variant.Its reported benefits extend across text, images, time-series forecasting, speech recognition, and audio generation.
  • Prior work: S4's Diagonal Plus Low Rank parameterization requires reduction steps and linear algebraic techniques, making it difficult to understand, implement, and analyze.The design is centered around the HiPPO theory for long-range modeling.
  • Contribution: DSS removes the low-rank component, enforces diagonal state matrices, and remains provably as expressive as general state spaces.The model requires only basic linear algebra and can be implemented in a few lines of code.
  • Results: 81.88 average accuracy across six Long Range Arena tasks is comparable to S4's 80.21 and exceeds the best Transformer variant's 61.41.The evaluated tasks span input lengths from 1K to 16K and multiple modalities.
  • Results: 98.2 versus 98.1 performance on Speech Commands shows DSS is comparable to S4 for raw speech classification.The dataset uses raw audio samples of length 16K.
  • Conclusion: DSS is presented as a simple and effective method for modeling long-range interactions in text, images, and audio.The authors state that its effectiveness, efficiency, and transparency could contribute to adoption of state-space models over attention-based peers.

2 Background

State-space models map input signals to outputs through continuous-time dynamics, then discretize those dynamics for sequence processing. Their outputs can be represented as convolution with an SSM kernel, enabling parallel computation, while kernel construction remains the main challenge.

  • State Spaces: A continuous-time state-space model maps input u(t) to output y(t) through state dynamics governed by matrix A and vectors B and C.The state vector is x(t), with y(t)=C x(t).
  • State Spaces: A state space is diagonal when its state matrix is diagonal.This is the paper's operational definition of a diagonal state space.
  • Discretization: Discretization with sample time Delta and zero-order hold converts the continuous model into a sequence-to-sequence recurrence over inputs and outputs.The recurrence updates x_k from x_{k-1} and u_k, then computes y_k from x_k.
  • Discretization: The discretized parameters are obtained using matrix exponentiation for A and a corresponding transformation for B, while C remains unchanged.These substitutions define the discrete-time state-space recurrence.
  • SSM Kernel: Unrolling the recurrence expresses each output as a sum of transformed past inputs, defining a length-L SSM kernel K.The kernel collects the sequence of terms involving powers of the discretized state matrix.
  • Parallel Computation: Parallel evaluation computes y as the discrete convolution K * u instead of sequentially applying the recurrence.The convolution formulation produces all output elements from the input sequence and kernel.
  • Parallel Computation: FFT computes the convolution in O(L log(L)) time, improving on the naive O(L^2) multiplication cost.The polynomial-product view makes simultaneous computation of all y_k possible.
  • Kernel Construction: The main computational challenge is constructing K, because direct computation involves L distinct matrix powers.The paper therefore seeks alternate state-space parameterizations that make kernel computation easier.

3 Method

DSS replaces S4’s diagonal-plus-low-rank parameterization with diagonal state spaces while retaining expressive kernel representations and efficient computation. The method uses two kernel parameterizations, stabilized implementation choices, and recurrent-state computation for long-range dependencies.

  • 3.1 Diagonal State Spaces: Under mild technical assumptions, any state space with a well-behaved state matrix has a diagonal state space computing the same kernel.The resulting representation uses a structured matrix-vector product instead of matrix powers.
  • 3.1 Diagonal State Spaces: Proposition 1 expresses diagonal state-space kernels using either elementwise exponentials or row-softmax normalization of eigenvalue-dependent terms.The two forms correspond to the DSSEXP and DSSSOFTMAX variants.
  • 3.1 Diagonal State Spaces: Positive real parts in eigenvalues can make exponential terms grow with sequence length and destabilize training, motivating two alternative diagonal-state-space methods.The issue arises from terms as large as exp(λ_i∆(L − 1)).
  • 3.2 DSS Layer: DSS layers compute one kernel per feature coordinate, apply discrete convolution, add a residual GELU path, and use an output projection for information exchange.The implementation is described as requiring only a few lines of code.
  • 3.3 Initialization of DSS layer: The DSS layer requires O(NHL) kernel-computation time and space, O(BHL log(L)) convolution time, and O(BH^2L) projection time.Careful initialization is crucial for state-of-the-art performance; the experiments use Skew-Hippo initialization with N = 64.
  • 3.4 States of DSS via the Recurrent View: In the recurrent view, diagonal coordinates can be computed independently, and near-zero real eigenvalues allow information to persist across many timesteps.This recurrent formulation is more efficient than convolution for autoregressive decoding during inference.

4 Experiments

DSS is evaluated across long-sequence classification tasks in text, images, and audio, where its performance is comparable to S4 despite simpler state-space parameterization. Ablations indicate that initialization and long-range kernels are important, while learned kernels adapt their locality to task requirements.

  • Long Range Arena: 81.88 average accuracy across 6 LRA tasks compares with S4’s 80.21 and exceeds the best Transformer variant’s 61.41.The state-space models are left-to-right unidirectional, whereas other compared models could be bidirectional.
  • Raw Speech Classification: 98.2 raw-speech accuracy on Speech Commands is comparable to S4’s 98.1.The dataset uses raw audio samples of length 16000 in a 10-way classification task.
  • Ablations: Randomly initializing DSS parameters significantly degrades performance on most tasks and causes failure on PATH-X.Even with this degradation, DSS outperforms all non-state-space models on every task.
  • Ablations: Truncating DSSSOFTMAX kernels to length 128 significantly reduces performance across most tasks, supporting a role for long-range dependencies.On LISTOPS and IMAGE, truncated kernels still outperform all Transformer variants.
  • Analysis of Learned DSS Parameters: Learned kernels are mostly local for IMAGE and SC but often long-range for PATH-X, revealing task-dependent interaction patterns.This pattern helps explain why limiting kernel length affects PATH-X more strongly than IMAGE and SC.
  • Analysis of Learned DSS Parameters: DSS parameters change substantially during training, including real parts of Λ becoming positive and ∆log differing between short- and long-range tasks.A LISTOPS ∆log outlier of 22 could produce an extremely large ∆ after exponentiation and potentially destabilize training.

5 Discussion

DSS uses diagonal state-space parameterization to simplify S4-style sequence modeling, while retaining the broader state-space approach’s focus on efficient long-range contextualization. The discussion also identifies evaluation scope and parameter-sensitivity considerations for future work.

  • Discussion: DSS replaces S4’s diagonal-plus-low-rank parameterization with a diagonal one, avoiding Padé approximations, Woodbury reductions, and truncated-generating-function Fourier analysis.The resulting method is described as significantly simpler than S4.
  • Limitations and future work: DSS was evaluated on sequence-level classification tasks, with token-level generation and large-scale pretraining identified as future directions.The stated future tasks include language modeling and forecasting.
  • Limitations and future work: The initialization and learning rates of Λre, Λim, and ∆log affect DSS performance and convergence.The paper specifically identifies these DSS parameters as important training choices.

A.1 Diagonal State Spaces

The appendix proves that, under mild technical assumptions, diagonal state spaces can reproduce kernels generated by diagonalizable general state spaces. It then expresses the resulting kernels as structured matrix-vector products rather than matrix powers.

  • A.1 Diagonal State Spaces: Under nonzero-eigenvalue and non-periodicity assumptions, every kernel from a diagonalizable state space has an equivalent diagonal state-space representation.The proposition assumes e^(Lλ_i∆) ≠ 1 for every eigenvalue.
  • A.1 Diagonal State Spaces: The proof constructs transformed input and output parameters so the diagonal representation computes the same kernel as the original state space.The construction uses elementwise products involving CV and V^-1B.
  • A.1 Diagonal State Spaces: With elementwise exponentiation, the original kernel is written as rw · Λ^-1 · (e^(Λ∆) − I) · E.Here P_i,k = λ_i · k∆ and E = elementwise-exp(P).
  • A.1 Diagonal State Spaces: For the corresponding diagonal state space, the kernel can instead be written as w · Λ^-1 · S, where S is row-softmax(P).The appendix defines S as the matrix obtained by applying softmax to each row of P.

A.2 Numerically Stable softmax

The appendix addresses complex-valued softmax singularities with a numerically stable formulation that shifts inputs and uses regularized reciprocals. The resulting softmax function is bounded and differentiable, though the uncorrected computation has known singularities.

  • A.2 Numerically Stable softmax: Complex softmax can be undefined when the sum of exponentials vanishes, unlike its real-valued counterpart.The appendix gives softmax(0, iπ) as an example of an undefined input.
  • A.2 Numerically Stable softmax: The regularized reciprocal uses x̄/(x·x̄ + ϵ), keeping the denominator at least ϵ and bounding the reciprocal magnitude.This regularization supports stable complex-valued normalization.
  • A.2 Numerically Stable softmax: The stable computation subtracts the element with maximum real part before applying exponentials.The shifted values are defined as r_i = x_i − m, where m = max(x).
  • A.2 Numerically Stable softmax: The corrected softmaxϵ is always bounded and differentiable.The unregularized SSM softmax is numerically stable when scalars with negative real parts are exponentiated, but has singularities at specified c values.
  • A.2 Numerically Stable softmax: For inputs structured as c · (0, …, L − 1), the appendix derives an FFT-based softmax computation with separate cases based on the sign of Re(c).The derived form uses an inverse FFT and roots of unity.

A.3 Experimental Setup

DSS and S4 experiments use sequence-classification heads, AdamW training, and mostly S4-recommended hyperparameters. The setup reports task-specific DSS exceptions and evaluates test accuracy at the best validation checkpoint.

  • Experimental Setup: Both LRA and Speech Commands aggregate DSS sequence outputs into one vector, using mean pooling except for TEXT and PATHFINDER.TEXT and PATHFINDER use the rightmost token instead of mean pooling.
  • Experimental Setup: AdamW with a constant learning-rate schedule and validation-plateau decay is used across datasets.DSS parameters generally use an initial learning rate of 10^-3 without weight decay.
  • Experimental Setup: The experiments largely follow S4 authors’ recommended model sizes and update steps, with task-specific DSS exceptions.Exceptions include altered ∆log learning rates, training duration, patience, batch size, and initialization.
  • Experimental Setup: Reported test accuracy is measured at the checkpoint with the highest validation accuracy.This checkpoint-selection rule is applied to all reported experiments.
  • Experimental Setup: All experiments were conducted on a single A100 GPU with 40 GiB of memory.The hardware statement covers the reported experimental runs.

A.4 Learned Parameters of DSSSOFTMAX

This section presents trained Λ and ∆log parameters for DSSSOFTMAX across the tasks described in §4.

  • Figure 4 shows trained Λ values in DSSSOFTMAX for the tasks described in §4.
  • Together, the figures inspect how DSSSOFTMAX parameters are learned across the evaluated tasks.
  • Figure 5 shows trained ∆log values in DSSSOFTMAX for the tasks described in §4.

A.5 Implementation of DSSSOFTMAX

The DSSSOFTMAX implementation converts real-valued parameter storage into complex parameters, constructs position-dependent state terms, stabilizes exponentiation, and returns a mixture of softmaxes.

  • The layer computes reciprocal terms from conjugates and combines the resulting weights with exponentiated states through an einsum.The final output is the real part of a mixture of softmaxes over the state dimension.
  • The implementation stores complex Λ and W as pairs of real values before converting them to complex tensors.This storage accommodates the nonlinear behavior of ADAM moments.
  • It constructs P from exp(log_dt)Λ and sequence positions, then shifts P using positive-real-part Λ values before exponentiation.The shift supports a numerically stable softmax computation.
  • Figure 6 provides the core DSSSOFTMAX layer implementation in PyTorch.
Loading 2203.14343v3…