[论文解读] Simplified State Space Layers for Sequence Modeling

Jimmy T. H. Smith, Andrew Warrington+1 more|arXiv (Cornell University)|Aug 9, 2022

Parallel Computing and Optimization Techniques被引用 77

一句话总结

本文提出 S5 层，一种多输入多输出的状态空间层，在效率上与 S4 相当，但采用完全递归的时域方法，在长距离序列任务上取得了最先进的结果。

ABSTRACT

Models using structured state space sequence (S4) layers have achieved state-of-the-art performance on long-range sequence modeling tasks. An S4 layer combines linear state space models (SSMs), the HiPPO framework, and deep learning to achieve high performance. We build on the design of the S4 layer and introduce a new state space layer, the S5 layer. Whereas an S4 layer uses many independent single-input, single-output SSMs, the S5 layer uses one multi-input, multi-output SSM. We establish a connection between S5 and S4, and use this to develop the initialization and parameterization used by the S5 model. The result is a state space layer that can leverage efficient and widely implemented parallel scans, allowing S5 to match the computational efficiency of S4, while also achieving state-of-the-art performance on several long-range sequence modeling tasks. S5 averages 87.4% on the long range arena benchmark, and 98.5% on the most difficult Path-X task.

研究动机与目标

• 以更简单、递归时域层为目标，推动并实现高效的长距离序列建模。
• 用一个多输入多输出的状态空间模型替换 S4 的一组独立的 SISO SSM，从而提高效率和可扩展性。
• 利用对角化的 HiPPO 启动初始化来捕捉长距离依赖。
• 在标准的长距离基准上展示具有竞争力或更优的性能，同时支持不规则采样和可变观测间隔。

提出的方法

• 用一个对角化的 MIMO SSM（潜在维度 P）替换 S4 的独立 SISO SSM。
• 对状态矩阵 A 进行对角化，以在时域中实现高效的并行扫描。
• 使用对角化的 HiPPO-启发矩阵（HiPPO-N）进行初始化，并在有 guiding 假设下将 S5 的动力学与 S4 联系起来。
• 使用零阶保持对连续时间 SSM 进行离散化，具备可学习时间尺度 Δ。
• 通过非线性激活和混合步骤，从 SSM前激活量产生层输出。
• 证明当潜在维度 P = O(H) 时，S5 在运行时和内存复杂度上与 S4 相匹配。

实验结果

研究问题

• RQ1 单个对角化动态的 MIMO SSM 能否实现与 S4 的一组 SISO SSM 相同的功能容量？
• RQ2 HiPPO 启动初始化（对角化的 HiPPO-N）是否在 MIMO 设置中为 S5 提供了有效的起点？
• RQ3 S5 层在处理不规则采样和时间可变参数时，能否在序列长度上保持线性复杂度？
• RQ4 相对于 S4 及其他基线，S5 在长距离序列基准上的表现如何？
• RQ5 放宽核心假设（如 A 相同、Δ 相同、P ≈ H）对 S5 设计的实际影响是什么？

主要发现

• S5 在长距离序列任务上实现最先进的性能，序列长度线性复杂度，效率与 S4 相当。
• 在 Long Range Arena 基准上，S5 的平均得分为 87.4%，在具有挑战性的 Path-X 任务上达到 98.5%。
• S5 在语音分类等任务上优于若干基线，并在多个任务上达到或超过 S4 的变体。
• 对角化的 HiPPO-N 初始化相对于 S4 使用的 HiPPO-LegS 初始化提供了具有竞争力的性能，使得动力学更加稳定且可对角化。
• S5 通过时域并行扫描支持不规则采样和时变 SSM，与基于卷积的 S4 内核不同。