使用结构化状态空间高效建模长序列

Efficiently Modeling Long Sequences with Structured State Spaces

阿尔伯特·顾 Albert Gu · Stanford · 2021-10-31 · arXiv:2111.00396 ↗ · 被引 3926

打开互动全文版(逐段中英对照 + 图/公式 + 论文问答)→

摘要 · Abstract

序列建模的一个核心目标是设计一个统一的原则性模型,能够处理跨多种模态和任务的序列数据,特别是长程依赖。尽管包括 RNN、CNN 和 Transformer 在内的传统模型有专门捕捉长依赖的变体,但它们仍然难以扩展到 10000 步或更长的序列。最近一种有前景的方法通过模拟基本状态空间模型(SSM)\( x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t) \) 来建模序列,并表明对于状态矩阵\( A \)的适当选择,该系统可以在数学上和经验上处理长程依赖。然而,该方法具有高昂的计算和内存需求,使其无法作为通用的序列建模解决方案。我们提出了结构化状态空间序列模型(S4),基于 SSM 的一种新参数化,并表明它可以在保留先前方法理论优势的同时,计算效率更高。我们的技术涉及对\( A \)进行低秩修正,使其能够稳定对角化,并将 SSM 简化为经过充分研究的柯西核计算。S4 在多种既定基准上取得了强大的实证结果,包括:(i) 在顺序 CIFAR-10 上达到 91%的准确率,无需数据增强或辅助损失,与更大的 2D ResNet 相当;(ii) 在图像和语言建模任务上大幅缩小与 Transformer 的差距,同时生成速度快 60 倍;(iii) 在 Long Range Arena 基准的每个任务上达到最先进水平,包括解决所有先前工作失败的长度 16k 的 Path-X 任务,同时与所有竞争对手一样高效。

A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies. Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of $10000$ or more steps. A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) \( x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t) \), and showed that for appropriate choices of the state matrix \( A \), this system could handle long-range dependencies mathematically and empirically. However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution. We propose the Structured State Space sequence model (S4) based on a new parameterization for the SSM, and show that it can be computed much more efficiently than prior approaches while preserving their theoretical strengths. Our technique involves conditioning \( A \) with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91\% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation $60\times$ faster (iii) SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efficient as all competitors.

核心贡献 · Key contributions

局限 · Limitations

论文章节 · Sections(共 18)

阅读逐段中英对照全文 →