Source-linked AI summary
Long Range Language Modeling via Gated State Spaces
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, Behnam Neyshabur
TL;DR
Long-range sequence modeling needs alternatives to attention's O(L^2) cost, especially for autoregressive language modeling. The paper introduces GSS, a gated state-space layer, and reports faster training than DSS, competitive fixed-compute performance against Transformer baselines, longer-input generalization, and gains from local-attention hybrids.
Problem
The paper addresses autoregressive long-range language modeling under attention's O(L^2) sequence-length cost.
Method
GSS extends state-space models with gating that reduces FFT dimensionality and targets TPU training bottlenecks.
Results
GSS trains 2–3× faster than DSS, is fairly competitive with Transformer baselines in fixed-compute comparisons, generalizes to inputs up to 65k, and improves with local-attention hybrids.
Takeaways & Limitations
GSS provides a straightforward state-space approach for long-range language modeling with efficient training, longer-input generalization, and optional local-attention improvements.
Takeaways & Limitations
Perplexity comparisons at sequence length 4k are slightly unfair because block-recurrent baselines cache state across segments and access almost the entire past.
Abstract
from arXiv · showhide
State space models have shown to be effective at modeling long range dependencies, specially on sequence classification tasks. In this work we focus on autoregressive sequence modeling over English books, Github source code and ArXiv mathematics articles. Based on recent developments around the effectiveness of gated activation functions, we propose a new layer named Gated State Space (GSS) and show that it trains significantly faster than the diagonal version of S4 (i.e. DSS) on TPUs, is fairly competitive with several well-tuned Transformer-based baselines and exhibits zero-shot generalization to longer inputs while being straightforward to implement. Finally, we show that leveraging self-attention to model local dependencies improves the performance of GSS even further.
1 Introduction
The paper develops GSS as a gated state-space alternative for autoregressive long-range language modeling, targeting efficient training and inference. It reports faster training than DSS, competitive performance with Transformer baselines, improved hybrid performance, and generalization to much longer inputs.
- Motivation: State-space models reduce sequence-length complexity from Ω(L^2) to O(L log L) using FFT-based convolutions and remain fully parallelizable during training.Their convolution–recurrence equivalence also supports recurrent inference without approximation.
- GSS: GSS modifies state-space layers with gating to reduce FFT dimensionality, addressing TPU training bottlenecks while simplifying initialization requirements.The paper motivates gating as an empirically supported way to reduce dimensionality in token-mixing operations.
- Results: GSS is fairly competitive with well-tuned Transformer-based baselines in fixed-compute comparisons across English books, Github code, and ArXiv mathematics.It lags on some tasks in the fixed-parameter setting, but is competitive when comparing exact TPUv4 training hours.
- Results: Interleaving local-attention Transformer layers with GSS improves performance at roughly no additional parameter or training-compute cost.The hybrid adds short-range interaction modeling to a GSS stack.
- Results: GSS generalizes from training sequences of at most 4k tokens to evaluated lengths up to 65k, with performance improving rather than degrading.The reported behavior indicates that the model uses additional context despite not being trained at those lengths.
- Inference: Recurrent inference makes state-space models efficient per decoded token, although adding local attention complicates this advantage in the hybrid model.The paper reports as much as 60× better inference efficiency for S4 as an example.
2 Related Work
Related work addresses the quadratic cost of attention through linearization or sparsification, while state-space layers offer recurrent and convolutional alternatives for sequence modeling. Recent efforts seek to retain these efficiency advantages while approaching attention-model performance.
- Attention: Attention models avoid the vanishing-gradients issue associated with RNNs but require O(L^2) time and space to construct QK^⊤.This quadratic cost limits the feasible sequence length.
- Attention: Linear-attention methods remove softmax to rearrange attention as Q(K^⊤V), reducing the operation's dependence on sequence length from quadratic to linear.Examples include Performer, Linear Attention, Random Feature Attention, and cosFormer.
- Attention: Sparse-attention methods avoid computing the entire QK^⊤ matrix by calculating only selected elements.BigBird and GMAT are cited as examples of this approach.
- State-space models: Classical recursive or state-space layers update states sequentially and emit outputs, requiring O(L) time when their transition and emission functions are efficient.Their sequential structure provides a direct alternative to attention-based sequence processing.
- Hybrid approaches: Transformer-XL and Block-Recurrent Transformer incorporate state into attention-based models through sliding windows or parameterized transition and emission functions.These approaches seek state-space benefits while retaining attention mechanisms.
- State-space models: Linear dynamical-system models such as S4 and DSS recast state-space computation as convolution with a very large kernel.This representation supports efficient parallel computation while preserving recurrent decoding properties.
3 Method
The method develops Gated State Space layers by combining gated reduced-dimensional contextualization with diagonal state-space computation, then extends them into a hybrid with interleaved Transformer blocks. State-space convolution enables efficient parallel training, while the hybrid targets complementary local and longer-range interactions.
- State Space Preliminaries: State-space recurrences unroll into convolution with kernel K, enabling autoregressive recurrence at inference and parallel computation during training.Sequential recurrence is slow for long training inputs, whereas convolution computes all outputs in parallel once K is available.
- State Space Preliminaries: FFT computes the convolution in O(L log L) time instead of the naive O(L^2) operation.The outputs are coefficients of the product of two degree L − 1 polynomials, which can be multiplied efficiently with FFT.
- GSS Layer: GSS replaces the quadratic attention component in GAU with a simplified DSS layer, retaining O(L log L) contextualization.The layer uses gating to reduce the dimensionality processed by the state-space component.
- GSS Layer: GSS maps X ∈ R^L×E to O ∈ R^L×E by elementwise multiplying a DSS-contextualized representation with V before projection.The implementation uses GELU activation; typical settings include H = E/4 = 256, N = 512, and F = 4096.
- GSS Layer: Fixing ∆ = 1 and randomly initializing Λ simplify kernel computation and perform comparably to the proposed specialized initialization in the authors’ experiments.The fixed-∆ choice removes separate ∆ values across coordinates, while Λ parameters are represented in log space.
- GSS-Transformer-Hybrid: A hybrid sparsely interleaves Transformer blocks with GSS layers, using chunked Transformer inputs to add local modeling alongside GSS’s broader interactions.Transformer inputs are divided into non-overlapping chunks of length 512 and processed independently.
4 Results
Across four varied language-modeling datasets, GSS was evaluated against DSS and strong Transformer-based baselines for perplexity, throughput, compute efficiency, and length generalization. GSS was consistently faster than DSS, competitive with baselines under matched compute, and generally benefited from longer evaluation contexts, though robustness varied by dataset and model size.
- Datasets: The evaluation covered LM1B, PG19, ArXiv, and Github, spanning short sentences, books, mathematics articles, and source code.These datasets differ in modality and typical document length; Github documents concatenate repository files across boundaries.
- DSS comparison: GSS outperformed DSS while being 2-3× faster on all tasks, using 4k training sequences except 512 for LM1B.The comparison used fixed-parameter models and reported perplexity alongside steps per second.
- Evaluation setup: The study used consistent baseline vocabularies and token-level held-out perplexity, while adjusting batch size with sequence length to keep tokens per batch constant.DSS used 12 layers and GSS 16 layers to match parameter count; training used 125k steps and 219 tokens per batch.
- Baseline comparisons: In fixed-parameter comparisons, GSS variants came close to but did not beat the strongest block recurrent baseline while running faster through full layer parallelization.GSS used roughly 192M parameters versus approximately 196M for the baselines.
- Baseline comparisons: In fixed-compute comparisons, larger GSS-L models matched baseline training time, and GSS-Hybrid outperformed the best baseline on PG19 and Github.ArXiv also improved substantially, but remained weaker than the baseline, possibly because its vocabulary was not tailored to mathematical symbols.
- Length generalization: GSS generalized from 4k-token training to evaluation lengths of 512, 4k, 16k, and 65k, with performance improving on PG19 as context increased.Smaller models generalized on ArXiv and Github, but generalization could degrade for larger models or noisier data.
A.1 Fast convolution via FFT
The method reformulates the computation as a circulant matrix-vector product, enabling fast evaluation through FFT and inverse FFT operations.
- The matrix-vector product is represented using a circulant matrix constructed from the input sequence.
- FFT and inverse FFT reduce the circulant matrix-vector product to O(L log L) time.
- Zero-padding the vectors to length 2L yields an equivalent circulant formulation for the computation.
A.2 Implementation of GSS
The implementation provides pseudocode for simplified DSS kernels and the DSS operation, including kernel construction, FFT convolution, and a residual parameter term.
- The simplified DSS kernel takes H, L, and N=512 as inputs.
- The kernel code constructs complex Λ and C parameters, evaluates state-space terms across sequence length L, and returns an H×L real kernel.
- The DSS operation normalizes u, computes state-space kernels, performs padded real FFT convolution, and adds D*u.