对角状态空间模型的参数化和初始化.
0. TL;DR
本文系统地研究了对角状态空间模型（Diagonal State Space Models，DSSMs）的参数化和初始化方法。文章首先介绍了状态空间模型（SSMs）在深度学习中的应用背景，特别是S4模型在处理长序列依赖任务中的优势。作者进一步提出了S4D模型，这是一种简化版的对角SSM，通过特定的初始化方法，能够在多种任务中取得与S4模型相媲美的性能。S4D模型的计算更为简单高效，且在图像、音频和医疗时间序列等领域的基准测试中表现出色。
1. 背景介绍
状态空间模型（SSMs）是一类在现代深度学习中用于处理序列数据的模型，它们在计算效率和建模能力上超越了传统的RNN、CNN和Transformer模型。S4模型是SSMs的一个重要版本，它通过使用特殊的HiPPO矩阵来建模长距离依赖关系，但在实际计算中需要复杂的算法和线性代数技术。为了简化这一过程，研究者们探索了将状态矩阵限制为对角矩阵的可能性，这导致了DSSMs的出现。尽管几乎所有SSMs都有等价的对角形式，但找到有效的初始化方法对于性能至关重要。
2. S4D模型
S4D模型是本文提出的一种对角SSM，它结合了S4的计算和参数化优势以及DSS的初始化方法。S4D模型的核心在于其对角状态矩阵的参数化和初始化，这使得模型在保持简单性的同时，也具备了强大的性能。
SSMs可以表示为线性时不变系统，通过状态矩阵$A$、输入矩阵$B$和输出矩阵$C$来参数化。
\[\begin{aligned} \mathbf{x}^\prime(t) &= A \mathbf{x}(t) + B \mathbf{u}(t) \\ \mathbf{y}(t) &= C \mathbf{x}(t) \\ \end{aligned}\]形式上，构造下列形式的卷积核$\overline{K}$，即可将序列运算转化为卷积运算：
\[\begin{aligned} \overline{K} &= (C\overline{B}, C\overline{A} \overline{B},...,C \overline{A}^k \overline{B}, ...) \\ y &= \overline{K} * u \end{aligned}\]在对角SSMs中，状态矩阵$A$是对角的。作者指出，S4D模型的卷积核$K$可以通过Vandermonde矩阵乘法来高效计算：
\[\begin{aligned} \overline{K} &= (C\overline{B}, C\overline{A} \overline{B},...,C \overline{A}^k \overline{B}, ...) \\ &= \begin{bmatrix}\overline{B}_0C_0 & \cdots & \overline{B}_{N-1}C_{N-1} \end{bmatrix} \begin{bmatrix} 1 & \overline{A}_0 & \overline{A}_0^2 & \cdots & \overline{A}_0^{N-1} \\ 1 & \overline{A}_1 & \overline{A}_1^2 & \cdots & \overline{A}_1^{N-1} \\ \cdots & \cdots & \cdots & \cdots & \cdots \\ 1 & \overline{A}_{N-1} & \overline{A}_{N-1}^2 & \cdots & \overline{A}_{N-1}^{N-1} \end{bmatrix} \\ &= \overline{B}^T \cdot C \cdot \mathcal{V}_L(\overline{A}) \end{aligned}\]
此外，S4D模型提出了若干参数化策略：
- $A$的参数化：为了保证模型的稳定性，$A$的实部被限制为负数，通过指数函数来参数化。
- $B$和$C$的参数化：S4D模型中$B$和$C$被单独参数化，并在训练中更新。
- 共轭对称性：为了确保输出为实数，参数化复数SSM时，参数以共轭对出现，从而减少参数数量。
3. 实验分析
在长距离arena基准测试中，S4D变体在几乎所有数据集上都具有竞争力，并且在几个数据集上优于S4变体。并在整个LRA上平均达到85%，比原始Transformer高出30多个点。
在Path-X数据集上，文章进一步消融了参数化选择，包括离散化方法、对$R(A)$的约束、$B$的可训练性和归一化方法。结果表明，默认的S4参数化选择是一个强大的基线，大多数其他选择差异不大。然而，在Path-X任务上，DSS的softmax归一化实际上会损害性能，因此一般不推荐使用。
