Long Range Language Modeling via Gated State Spaces

TL;DR

This paper introduces Gated State Space (GSS), a new autoregressive model that efficiently captures long-range dependencies in sequences, outperforming some baselines and generalizing well to longer inputs.

Contribution

The paper proposes Gated State Space (GSS), a novel layer that trains faster and generalizes better for long sequences compared to existing state space models.

Findings

GSS trains significantly faster than diagonal S4 on TPUs.

GSS is competitive with well-tuned Transformer baselines.

GSS exhibits zero-shot generalization to longer inputs.

Abstract

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.