简化序列建模的状态空间层.
0. TL;DR
本文介绍了一种新的序列建模方法——S5层，它在长序列建模任务中取得了最优性能。S5层在S4层的基础上进行了改进，通过采用多输入多输出（MIMO）的线性状态空间模型（SSM），提高了计算效率和性能。实验表明，S5层在多个长序列建模任务上表现优异，尤其在Long Range Arena（LRA）基准测试和Path-X任务上达到了最高准确率。
1. 背景介绍
长序列建模是机器学习领域中的一个重要挑战。传统的循环神经网络（RNN）在长序列上容易遭遇梯度消失或梯度爆炸的问题，导致模型性能下降。为了解决这些问题，研究者们提出了多种改进方法，如长短时记忆网络（LSTM）和门控循环单元（GRU）。然而，这些方法在处理极长序列时仍然面临计算复杂度和内存使用方面的限制。
最近，Gu等人提出了结构化状态空间序列（S4）层，该方法通过结合线性状态空间模型（SSM）、HiPPO框架和深度学习，实现了在长序列建模任务上的卓越性能。S4层利用多个独立的单输入单输出（SISO）SSM来处理序列数据，并通过非线性混合层来融合各个SSM的输出。尽管S4层在多个基准测试中取得了显著成绩，但其计算效率和内存使用仍有提升空间。
本文在S4层的基础上，引入了一种新的状态空间层——S5层。S5层通过采用多输入多输出（MIMO）SSM，简化了S4层的结构，并提高了计算效率。同时，S5层保持了与S4层相同的计算复杂度，但实现了更高的性能。
2. S5
S4层采用的离散SSM形式为：
\[\begin{aligned} x_{k+1} &= \overline{A}x_k + \overline{B} u_k \\ \end{aligned}\]该形式为单输入单输出，即针对标量输入$u_k$进行的。当输入向量$u=[u_1,u_2,…,u_k]$时，S4直接对每个分量独立地应用上述公式，然后将结果拼接起来。
S5层采用了多输入多输出（MIMO）的SSM形式，能够同时处理多个输入和输出，从而简化了模型结构。S5层通过以下两个步骤来生成输出序列：
- SSM计算：给定输入序列$u_{1:L} \in R^{L×H}$，MIMO SSM通过隐状态$x_k \in R^P$生成SSM输出（或预激活）$y_{1:L} \in R^{L×H}$。
- 非线性激活：将SSM输出通过非线性激活函数（如GELU）处理，生成最终的层输出。
S5层的性能对SSM状态矩阵的初始化非常敏感。为了确保高效的并行扫描，状态矩阵必须是对角矩阵。在S5层中采用了一种可学习的对角线状态矩阵初始化方法，并通过实验验证了其有效性。
此外，S5层还使用了一组可学习的离散化步长参数$Δ ∈ R^P$来控制SSM的动态行为。这些参数允许模型根据任务需求自适应地调整时间尺度。
与S4层相比，S5层不需要额外的位置线性层，因为MIMO SSM已经能够捕获序列中的位置信息。
3. 实验结果
为了评估S5层的性能，作者在多个长序列建模任务上进行了实验，包括Long Range Arena（LRA）基准测试和Path-X任务。LRA基准测试包含六个序列建模任务，序列长度从1024到超过16000不等。Path-X任务是一个特别困难的长序列分类任务，要求模型对复杂路径进行准确分类。
实验结果表明，S5层在多个长序列建模任务上取得了卓越性能。在LRA基准测试中，S5层的平均准确率为87.4%，在Path-X任务上的准确率为98.5%，均超过了S4层和其他基线方法。
此外作者还比较了S5层、S4层和S4D在运行时性能和内存使用方面的差异。实验结果表明，S5层在速度和内存使用方面均优于S4D，与S4层相当或更优。
