对角状态空间和结构化状态空间一样有效.
1. 背景介绍
在深度学习领域，状态空间模型（State Space Models, SSMs）是一类用于序列数据建模的重要工具。传统SSMs通常具有复杂的结构和较高的计算成本，限制了其在长序列任务中的应用。近年来，随着Transformer模型的普及，各种变体不断涌现，旨在降低计算和内存需求。然而，这些变体在长范围依赖任务上的性能仍然不尽如人意。
本文提出了一种新的思路，即通过对角状态空间（DSS）来简化SSM的参数化，同时保持其表达能力。DSS的核心思想是利用对角矩阵来替代一般状态空间中的复杂矩阵，从而在保证性能的同时降低计算复杂度。
2. Diagonal State Space
在连续时间状态下，SSM由状态矩阵$A$、输入向量$B$和输出向量$C$参数化，定义了一个从输入信号$u$到输出信号$y$的函数映射。
\[\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}\]本文作者指出，对角状态空间与一般状态空间具有相同的表达能力。DSS通过参数化状态空间，避免了矩阵幂（$\overline{A}^k$）的计算，仅涉及结构化的矩阵-向量乘积。
根据SSM的连续-离散形式转换关系：
\[\begin{aligned} \overline{A} &= e^{\Delta \mathbf{A}} \\ \overline{B} &= (e^{\Delta A}-I ) A^{-1} B \end{aligned}\]则卷积核的第$k$个元素为：
\[\begin{aligned} \overline{K}_k &= C \overline{A}^k \overline{B} \\ &= C e^{A\cdot k\Delta} (e^{A\Delta}-I) A^{-1} B \\ \end{aligned}\]假设$A$可对角化：$A=V\Lambda V^{-1}$，特征值为$\lambda_1,…,\lambda_N$，记$(CV)(V^{-1}B)=\tilde{w}$，则上式可表示成：
\[\begin{aligned} \overline{K}_k &= C e^{A\cdot k\Delta} (e^{A\Delta}-I) A^{-1} B \\ &= (CV) e^{\Lambda\cdot k\Delta} (e^{\Lambda\Delta}-I) \Lambda^{-1} (V^{-1}B) \\ &= \sum_{i=1}^N \frac{e^{\lambda_i\cdot k\Delta} (e^{\lambda_i\Delta}-I)}{\lambda_i} \tilde{w}_i \\ &= \sum_{i=1}^N \frac{e^{\lambda_i\cdot k\Delta} (e^{\lambda_i\Delta}-I)}{\lambda_i (e^{L\lambda_i\Delta}-I)} (e^{L\lambda_i\Delta}-I)\tilde{w}_i \\ &= \sum_{i=1}^N \tilde{w}_i (e^{L\lambda_i\Delta}-I) \frac{1}{\lambda_i} \frac{e^{\lambda_i\cdot k\Delta} }{\sum_{l=0}^{L-1}e^{l\lambda_i\Delta}} \\ \end{aligned}\]记$P,P_{i,k}=\lambda_i k \Delta$，则卷积核可以表示为：
\[K = \tilde{w} \cdot \Lambda^{-1} (e^{\Lambda \Delta}-I) \cdot \text{elementwise-exp}(P)\]上式称为DSS-EXP。
另一方面，记$w_i=\tilde{w}_i(e^{L\lambda_i\Delta}-I)$，则上式也可表示成：
\[\begin{aligned} \overline{K}_k &= \sum_{i=1}^N \tilde{w}_i (e^{L\lambda_i\Delta}-I) \frac{1}{\lambda_i} \frac{e^{\lambda_i\cdot k\Delta} }{\sum_{l=0}^{L-1}e^{l\lambda_i\Delta}} \\ &= \sum_{i=1}^N w_i \frac{1}{\lambda_i} \frac{e^{\lambda_i\cdot k\Delta} }{\sum_{l=0}^{L-1}e^{l\lambda_i\Delta}} \\ \end{aligned}\]卷积核也可以表示为：
\[K = w \cdot \Lambda^{-1} \cdot \text{row-softmax}(P)\]上式称为DSS-SOFTMAX。
3. 实验分析
实验在Long Range Arena（LRA）基准测试上进行，该测试包含一系列序列级别的分类任务，输入长度多样（1K-16K），要求模型具备相似性、结构性和视觉空间推理能力。
DSS在LRA上的表现令人瞩目，其性能与S4相当，且在平均测试准确率上略有提升。此外，DSS保持了相对于最佳Transformer变体20个百分点的领先优势。
实验还表明，截断核长度会导致性能显著下降，进一步证实了状态空间模型在长范围依赖任务上的优势。
