MambaTS: Improved Selective State Space Models for Long-term Time Series Forecasting

系列文章目录

MambaTS：改进的用于长期时间序列预测的选择性状态空间模型

---

---

摘要

近年来，Transformers 已成为长期序列预测（LTSF）事实上的架构，但面临着二次复杂度和排列不变偏差等挑战。 最近的模型 Mamba 基于选择性状态空间模型 (SSM)，已成为 Transformer 的竞争替代品，提供可比的性能、更高的吞吐量和与序列长度相关的线性复杂度。 在本研究中，我们分析了当前 Mamba 在 LTSF 中的局限性，并提出了四项有针对性的改进，从而形成了 MambaTS。 我们首先引入沿时间的变量扫描，将所有变量的历史信息排列在一起。 我们建议 Mamba 中的因果卷积对于 LTSF 来说不是必需的，并提出了临时 Mamba 块（TMB）。 我们进一步结合了 TMB 选择性参数的 dropout 机制，以减轻模型过度拟合。 此外，我们通过引入变量排列训练来解决变量扫描顺序敏感性的问题。 我们进一步提出了随时间的变量感知扫描，以在训练期间动态发现变量关系，并通过解决推理期间访问所有节点的最短路径问题来解码最佳变量扫描顺序。 在八个公共数据集上进行的大量实验表明，MambaTS 实现了新的最先进的性能。

---

一、引言

长期时间序列预测（LTSF）在天气、金融、医疗、能源和交通等各个领域有着广泛的应用[1,2,3]。 随着深度学习的快速发展，当前时间序列预测的方法已从传统的统计学习方法转向基于深度学习的方法，例如循环神经网络（RNN）和时间卷积神经网络（TCN）；[4, 5] ）。 自从 Transformer [6] 推出以来，基于 Transformer 的方法已成为主流 LTSF 方法 [7, 8]，利用其自注意力机制有效捕获时间序列数据中的长期依赖性。

然而，最近的研究强调了变形金刚在 LTSF 中面临的挑战。 首先，这些方法受到二次复杂度的困扰，其中计算成本随着上下文的长度而迅速增加[9,10,11]。 此外，一些研究发现，基于 Transformer 的 LTSF 方法的性能并不一定会随着回溯窗口的增加而提高 [8, 12]。 值得注意的是，最近一项名为 DLinear [13] 的研究对 Transformer 在 LTSF 中排列不变偏差的有效性提出了质疑，取得了令人惊讶的结果，并超越了大多数使用单层馈电的最先进的 (SOTA) 基于 Transformer 的方法。 转发网络。

表 1：案例研究：用于 LTSF 任务的 Mamba。

图 1：训练速度和 GPU 比较。

最近流行的状态空间模型（SSM；[14,15,16]）引起了我们的注意。 SSM 是有原则的顺序模型，使用常微分方程描述状态随时间的演变，这使得它们自然适合时间序列建模。 最近，Mamba [17] 通过引入选择机制来过滤掉不相关的信息并重置状态，从而增强了传统的 SSM。 它还采用了硬件感知设计，可实现高效的并行训练。 与 Transformer 相比，Mamba 在各个领域都表现出了具有竞争力的性能 [17,18,19]，提供快速推理和序列长度的可扩展性。

Mamba 固有的顺序性和选择性本质促使我们在 LTSF 中验证其性能并将其与 Transformer 进行比较。 最初，我们用 Mamba Block 替换了两种基于 SOTA Transformer 的方法 PatchTST [8] 和 iTransformer [12] 中的 Transformer Block，从而产生了相应的 Mamba 版本2。 随后，我们将它们与 PatchTST 和 iTransformer 进行了比较，结果如表 1 和图 1 所示。我们的研究结果表明，与 PatchTST 和 iTransformer 相比，Mamba 版本的训练速度提高了 1.3 倍，内存消耗减少了 5.3 倍 和 7.0 倍。 然而，尽管有这些改进，我们在这些 LTSF 数据集上没有观察到 Mamba 相对于 Transformer 的性能优势。

这似乎与 Mamba 天然更适合时间序列建模的直觉相矛盾。 我们对 Mamba 进行了四项改进，并引入了 MambaTS，这是一个基于改进的选择性 SSM 专门为 LTSF 设计的框架。 首先，我们介绍沿时间可变扫描（VST）。 与 PatchTST [8] 类似，我们最初将每个变量分割成补丁，并将它们线性映射到标记。 然而，与PatchTST不同的是，我们采用变量混合的方式，在同一时间步交替组织不同变量的token，从而随着时间的推移将它们扩展以形成全局历史回顾窗口。 然而，变量之间缺乏确定的因果关系导致 VST 后出现局部非连续跨越结果。 在这种情况下，原始 Mamba 块中的因果卷积可能会影响模型性能。 此外，一些研究认为回溯窗口的因果卷积是不必要的，并且可能会限制时间特征提取[20, 21]。 为了解决这些问题，我们通过在 SSM 之前删除局部卷积来引入 Temporal Mamba Block (TMB)。 此外，为了减轻过度拟合，我们将 dropout 机制 [22] 纳入 TMB 的选择性参数中，基于过度信息集成可能导致过度拟合的发现 [8]，正如我们在 Mamba 实验中观察到的那样。

虽然 TMB 是为时间建模而定制的，但它仍然对扫描顺序敏感，突出了变量组织的重要性。 为了解决这个问题，我们引入了变量排列训练（VPT）。 通过在每次迭代中打乱变量顺序，我们减轻了未定义变量顺序的影响并提高了模型的局部交互能力。 随之而来的问题是我们如何确定最佳的通道扫描顺序？ 受图论中拓扑排序[23, 24]的启发，我们进一步提出了一种新颖的沿时间变量感知扫描（VAST）。 在VAST中，我们考虑线性序列中节点的位置，例如a→c，b→c，其中合理的拓扑顺序可能是a→b→c。 具体来说，在训练期间，我们更新一个距离矩阵，该距离矩阵根据每次排列后的前向传播结果捕获变量关系。 该公式将最优扫描顺序问题转换为密集图中的非对称旅行商问题 (ATSP) 的实例。 为了解决这个问题，我们利用基于模拟退火的 ATSP 求解器来导出最终的通道扫描顺序。 据我们所知，我们是第一个在 SSM 背景下研究变量扫描顺序以进行时间序列建模的人。

通过上述四种设计，MambaTS以线性复杂度实现了跨时间和变量的全局依赖关系的高效建模。 对八个流行的公共数据集进行的广泛实验表明，MambaTS 在大多数 LSTF 任务和设置中实现了 SOTA 性能。

总之，我们的贡献如下：
• 我们引入了MambaTS，这是一种基于改进的选择性SSM 构建的新型时间序列预测模型。 通过合并VST，我们有效地组织所有变量的历史信息以创建全局回顾序列。
• 我们提出了TMB，它解决了LTSF 中普通Mamba 因果卷积的可有可无性。 此外，我们针对 TMB 的选择性参数采用了 dropout 机制，以减轻模型过度拟合。
• 我们实施VPT策略，消除未定义变量顺序的影响，进一步增强模型的局部上下文交互能力。
• 我们提出VAST，它在训练期间发现不同变量之间的关系，并利用ATSP 求解器来确定推理期间的最佳变量扫描顺序。

二、 Related Work

长期时间序列预测传统的 LTSF 方法利用时间序列数据的统计特性和模式进行预测 [25]。 近年来，LTSF 已转向深度学习方法，利用各种神经网络来捕获复杂的模式和依赖性，从而提高 LTSF 的性能。 这些方法可以大致分为两类：变量混合和变量无关。 变量混合方法采用不同的架构来对跨时间和变量的依赖关系进行建模。 由于序列建模的性质，RNN [26,27,28] 最初被引入 LTSF。 TCN 以其局部偏差而闻名，可有效捕获时间序列数据中的局部模式，并在 LTSF 中显示出有希望的结果 [4,21,29]。 随后引入了 Transformer，通过 self-attention 来完成远程依赖建模，并已成为 LSTF 的主流方法[2]。 然而，由于二次复杂度，基于 Transformer 的方法在优化效率方面一直存在困难 [7,30,9,10,11]。 最近，这些方法通过基于补丁的技术得到了显着的改进 [31, 8]。 MLP 也常用于 LTSF，并以其简单直接的架构取得了令人印象深刻的结果 [32]。 图神经网络已被用来建模变量之间的关系[33]。 FourierGNN [34] 将整个时间序列信息表示为超变量图，并采用傅里叶图神经网络进行全局依赖建模。 另一方面，变量无关方法仅专注于在变量独立性假设下对时间依赖性进行建模[13,8,35]。 这些方法以其简单性和高效性而闻名，通常能够减轻模型过度拟合并取得显着的成果 [13, 8]。 然而，这种假设可能会使问题过于简单化，并可能导致不适定的情况[31]。

状态空间模型 最近，一些工作 [25] 将 SSM [36] 与深度学习相结合，并展示了解决远程依赖问题的巨大潜力。 然而，状态表示的计算和内存要求过高，常常阻碍其实际应用[15]。 人们提出了 SSM 的几种有效变体，例如 S4 [15]、H3 [14]、门控状态空间 [37] 和 RWKV [38]，以提高实际任务中的模型性能和效率。 Mamba [17] 通过引入基于 S4 的数据相关选择机制来有效过滤特定输入并捕获随序列长度缩放的远程上下文，解决了传统 SSM 方法的关键限制。 Mamba 在长序列建模方面展示了线性时间效率，并在基准评估中超越了 Transformer 模型 [17]。

Mamba 还成功扩展到图像[39,18,19]、点云[40]、表[41]和图[42,43]等非序列数据，以增强其捕获远程依赖关系的能力。 为了解决 Mamba 的扫描顺序敏感性，一些研究引入了双向扫描 [18]、多向扫描 [44, 39]，甚至自动方向扫描 [45]。 然而，目前考虑时间问题中可变扫描顺序问题的工作有限。 为了应对这一挑战，我们引入了VAST策略来进一步增强MambaTS的表达能力。 与我们的方法类似，Graph-Mamba [42] 也提出了类似的排列策略来扩展图上的上下文感知推理。 然而，它基于节点优先级，引入了专门为图设计的偏置策略。

三、 Preliminaries

状态空间模型 SSM 通常被视为线性时不变 (LTI) 系统，它通过状态表示 h(t) 将连续输入信号 x(t) 映射到相应的输出 y(t)。 该状态空间描述了状态随时间的演变，可以使用常微分方程表示如下：

这里，$h^{\prime }(t)=\frac{dh(t)}{dt}$，A、B、C 和 D 是与时间无关的 SSM 的参数。

离散化 由于 SSM 的连续性，寻找 SSM 的分析解决方案非常具有挑战性。 离散化通常用于促进离散域中的分析和求解，这涉及将连续时间状态空间模型近似为离散时间表示。 这是通过以固定时间间隔对输入信号进行采样以获得其离散时间对应物来完成的。 由此产生的离散时间状态空间模型可以表示为：

这里，$h_k$表示时刻k的状态向量，$x_k$表示时刻k的输入向量。 矩阵$\overline{A}\mathrm{and}\overline{B}$ 是使用适当的离散化技术（例如 Euler 或 ZOH（零阶保持）方法）从连续时间矩阵 A 和 B 导出的。 在这种情况下，$\overline{\mathbf{A}}=\exp (\Delta \mathbf{A}),\overline{\mathbf{B}}=(\Delta \mathbf{A}{)}^{-1}(\exp (\Delta \mathbf{A})-\mathbf{I})\cdot \Delta \mathbf{B}.$。

选择性扫描机制 Mamba 通过允许参数以上下文相关的方式影响序列上的交互，进一步引入了选择性 SSM。 这种选择性机制使 Mamba 能够过滤掉时间序列任务中不相关的噪声，同时选择性地传播或忘记与当前输入相关的信息。 这与之前具有静态参数的SSM方法不同，但它确实打破了LTI特性。 因此，Mamba采取了硬件优化的方法，并实现了并行扫描训练来应对这一挑战。

四、 模型架构

对于多元时间序列预测问题，给定一个具有 K 个变量$({\mathbf{x}}_1,{\mathbf{x}}_2,\cdots ,{\mathbf{x}}_L),$的 L​​ 长度回溯窗时间序列数据，其中${\mathbf{x}}_i\in {\mathbb{R}}^K$表示时间步 i 处 K 个变量的值。 我们的目标是预测未来 T 个时间步的值，表示为$({\mathbf{x}}_{L+1},\cdots ,{\mathbf{x}}_{L+2},\cdots ,{\mathbf{x}}_{L+T}).$。

4.1 整体架构

MambaTS 的架构如图 2 所示。它主要由嵌入层、实例归一化层、N× 时间 Mamba 块和预测头组成。

修补和标记化单个时间点的语义信息往往是稀疏的。 遵循 PatchTST [8]，我们每 s 个时间步将每个变量分割成大约 M = L/s 的补丁，聚合来自相邻点的信息，然后通过线性映射将它们映射到 D 维标记。 此过程有效地减少了内存使用并消除了冗余信息。

沿时间进行变量扫描通过嵌入 K 个变量，我们得到 K × M 个令牌。 为了充分利用 Mamba 的线性复杂性和选择性优势，同时建立回溯窗口的全面表示，我们引入了沿时间可变扫描机制。 VST 在每个时间步以时间上交替的方式排列变量标记。 这种结构化的组织使模型能够更准确地捕获时间序列数据的长期依赖性和动态变化。 随后，我们将 VST 的结果输入编码器。

图 2：MambaTS 的整体架构。

编码器 编码器由 N 个堆叠的 TMB 组成，每个 TMB 包含两个分支。 右侧 SSM 分支专注于序列建模，而左侧分支包含门控非线性层。 原始Mamba块的计算过程如下：

这里，因果卷积作为移位 SSM，被插入到 SSM 模块之前，以增强相邻标记之间的连接，如图 2 (b) 所示。 然而，VST 生成的结果可能不是局部因果关系。 因此，TMB（参见图 2 ©）删除了该组件。 此外，为了防止过度拟合，我们为选择性参数引入了 dropout 机制 [22]，如下所示：

预测头由于编码器能够捕获全局依赖关系，出于效率原因，类似于 PatchTST [8]，我们在解码阶段采用与通道无关的解码方法。 每个通道均使用简单的线性头单独解码。

实例归一化为了减轻训练和测试数据之间的分布变化，按照 RevIN [46]，我们将每个输入通道标准化为零均值和单位标准差，并跟踪这些统计数据以对模型预测进行反归一化。

损失函数我们选择均方误差（MSE）损失作为主要损失函数，由下式给出：

4.2 变量排列训练

为了减轻未定义通道顺序的影响并增强局部上下文交互，我们引入了变量排列训练（VPT）策略。 具体来说，对于由K×M个token组成的编码器输入，VPT在每个时间步以一致的顺序对K个token进行混洗，并在解码后恢复混洗状态，以确保序列的正确输出。

4.3 Variable-Aware Scan along Time

为了确定最佳变量扫描顺序，必须考虑变量之间的关系。 然而，由于这些关系是未知的，我们提出了 VAST，一种简单而有效的方法来估计变量关系，指导推理阶段的扫描顺序。

训练 对于任意 K 个变量，我们维护一个有向图邻接矩阵 $P\in {\mathbb{R}}^{K\times K}$，其中 $p_{i,j}$表示从节点 i 到节点 j 的成本。 值得注意的是，通过在训练期间引入变量排列，我们可以探索扫描顺序的各种组合并使用等式（1）评估其有效性。 （5）。 经过置换，得到节点索引序列 V = { v 1 , v 2 , ⋯   , v K } \mathbf{V}=\{v_{1},v_{2},\cdots,v_{K}\} V={v1​,v2​,⋯,vK​}，其中$v_k$表示打乱后的序列中的新索引。 随后，导出K-1个转移元组$\{(v_1,v_2),(v_2,v_3),\cdots (v_{K-1},v_K)\}$。 对于每个样本，在网络的第 t 次迭代期间计算训练损失 $l^{(t)}$。 因此，我们用指数移动平均线更新 P：

其中 β 是一个超参数，控制移动平均线的速率，确定新估计对变量关系的影响。 为了促进有效的训练，我们扩展了方程。 (6)批量版本。 通过简单的中心化操作，消除不同样本批次的影响，即 ${\bar{l}}^{(t)}=l^{(t)}-\mu (l^{(t)})$，其中 µ 为均值函数，将方程修改为：

推理 在整个训练过程中，P 用于确定最佳变量扫描顺序。 这涉及解决非对称旅行商问题，该问题寻求访问所有节点的最短路径。 给定 P 表示的密集连通性，找到最优遍历路径是 NP 困难的。 因此，我们引入了一种基于启发式的模拟退火[47]算法来进行路径解码。

五、 Experiments

数据集 我们对八个公共数据集进行了广泛的实验，如表 2 所示，包括两个 ETT 数据集 [7]、天气、电力、交通 [9]、太阳能 [28]、Covid-19 [48] 和 PEMS [20] ，涵盖电力、能源、交通、天气、健康等领域。

基线和指标为了证明 MambaTS 的有效性，我们将其与 LTSF 的 SOTA 模型进行了比较，包括五种流行的基于 Transformer 的方法：PatchTST [8]、iTransformer [12]、FEDformer [10]、Autoformer [9] 和三种竞争方法 非基于 Transformer 的方法：DLinear [13]、MICN [21] 和 FourierGNN [34]。 遵循 PatchTST [8]，我们主要使用均方误差 (MSE) 和平均绝对误差 (MAE) 来评估模型。

实施细节实验是在 NVIDIA RTX 3090 Ti 24 GB GPU 上使用 Adam 优化器 [49] 进行的，测试版为 (0.9, 0.999)。 训练运行了 10 个 epoch，并使用 3 的耐心提前停止以避免过度拟合。 所有比较模型的最佳参数选择都是在验证集上仔细调整的。

5.1 主要结果

表3显示了多元长期预测的结果。 总体而言，MambaTS 在大多数数据集的不同预测范围内取得了新的 SOTA 结果（以红色粗体突出显示）。 虽然 DLinear 和 PatchTST 假定变量独立性，并且在具有少量变量（例如 ETTh2/m2 (K = 7) 和天气 (K = 21)）的数据集上表现良好，但它们的性能在具有大量变量（例如流量）的复杂数据集上表现不佳 (K = 862) 和 Covid19 (K = 948)，凸显了变量无关假设的局限性。 相反，iTransformer 表现出截然不同的性能，在复杂的数据集上表现出色，但在变量较少的数据集上表现不佳。 其他基线在某些预测场景下展示了特定数据集的竞争结果。

表 3：多元长期序列预测结果。 对于 Covid-19 数据集，所有模型均采用 L = 96 的回溯窗口长度，对于其余数据集，采用 L = 720 的回溯窗口长度。
表 4：组件消融。 VST：随时间变化的扫描。 TMB：颞曼巴区块。 VAST：随时间变化的感知扫描。 此处列出了所有预测长度的平均结果。

5.2 消融研究与分析

为了验证所提出组件的合理性和有效性，我们进行了广泛的消融实验，如表4、表5、图3和图4所示。

组件消融 表 4 列出了组件消融。 最初，将 VST 与基于 Mamba 的基线 PatchTST 集成可以提高大多数数据集的性能，展示了考虑所有变量的好处。 单独替换 TMB 会带来显着的性能提升，强调了 TMB 对于时间建模的功效。 结合 VST 和 TMB 产生的性能优于单独的 VST，但略低于原始 TMB，这归因于 VST 包括所有变量，而 TMB 消除了局部偏差因果卷积，使其对变量顺序更加敏感。 不过，这个问题已通过引入 VAST 得到解决（参见表 4 第 4 行）。 通过这些组件，MambaTS 实现了最佳性能。

图 3：TMB 的脱落消融。 左：TMB 在天气数据集上具有不同的丢失率。 右：训练时的损失曲线。

Dropout Ablation 我们进一步分析了TMB 中dropout 的作用。 图 3（左）显示了 MambaTS 在 Weather 数据集上不同 dropout 率（0.1-0.5）的结果。 与无dropout相比，MambaTS的MSE降低随着dropout率的增加而增加，在0.2和0.3时达到最佳性能，超过0.4时性能下降。 图 3（右）展示了训练期间相应的损失曲线，表明 TMB 中的 dropout 有助于防止过早收敛和过拟合。 此外，我们观察到 dropout 有助于降低验证损失。

VAST消融在表5中，我们对VAST策略进行了广泛的消融研究，重点关注路径解码策略的设计和选择。 表 5 中的“W/o VPT”表示 MambaTS 在没有 VPT 作为基线的情况下进行训练。 “随机 (100x)”表示在使用 VPT 训练后对 100 次测试运行进行采样并对结果求平均值。 可以看出，“Random (100x)”显着优于“W/o VPT”，这强调了 VPT 的有效性。 图 4 中的进一步视觉比较表明，在大多数情况下，即使是随机变量扫描也优于“W/o VPT”。 然后，我们探索了不同的启发式解码策略，包括贪婪策略（GD）、局部搜索（LS）、Lin 和 Lernighan（LK）以及模拟退火（SA）。 我们默认采用 SA 作为求解器。 作为效率和性能之间的权衡，由于其指数复杂性，我们没有采用精确的 ATSP 求解器。 如表 5 和图 4 所示，SA 在相对性能一致性方面始终优于其他求解器。

5.3 模型分析

增加回溯窗口 先前的研究表明，基于 Transformer 的方法不一定会受益于不断增长的回溯窗口 [13, 8]，这可能是由于长时间输入分散了注意力。 在图 5 中，我们评估了 MambaTS 在这种情况下的表现，并将其与几个基线进行了比较。 可以观察到，MambaTS 始终表现出从不断增长的输入序列中受益的能力。 iTransformer、PatchTST 和 DLinear 也表现出这种优势，但 MambaTS 的整体曲线低于 PatchTST 和 DLinear。 与 iTransformer 相比，MambaTS 受益于更长的回溯窗口。 此外，我们注意到 iTransformer 似乎在单个数据集任务上表现出不连续的增益。

图 5：MambaTS 在具有不同长度回溯窗口的不同数据集中的性能。

效率分析 MambaTS 使用 VST 集成所有变量的历史信息，并通过 TMB 进行全局依赖性建模。 MambaTS 的计算复杂度为 O(KL P )，其中 K 表示变量数量，L 表示回溯窗口的长度，P 表示补丁步长。 表 6 概述了其他模型的计算复杂度。

其中，Autoformer和FEDformer采用传统的point-wise tokenization，计算复杂度依赖于L。MICN和FourierGNN专注于建模变量之间的依赖关系，复杂度较高。 DLinear 实现了 O(L) 复杂度，但缺乏变量混合。 PatchTST 和 iTransformer 的复杂度分别为 O(L /P 2 ) 和 O(K2)。 相比之下，MambaTS 的 O(K · L /P ) 复杂度在这些方法之间取得了平衡，如表 3 所示，突出了有意义且有效的平衡。

六、 Conclusion

在这项工作中，我们提出了 MambaTS，这是一种基于改进的选择性 SSM 构建的新型多元时间序列预测模型。 我们首先引入VST来组织所有变量的历史信息，形成全局回顾序列。 认识到 Mamba 中因果卷积对于 LTSF 的可有可无，我们提出了 Temporal Mamba Block。 此外，我们通过 dropout 正则化增强了 Mamba 的选择性参数，以防止过度拟合并增强模型性能。 为了缓解扫描顺序敏感性问题，我们实施变量排列训练来应对未定义变量顺序的影响。 最后，我们提出 VAST 在训练期间动态发现变量之间的关系，并使用 ATSP 求解器解决最短路径访问问题，以确定最佳变量扫描顺序。 通过这些设计，MambaTS 实现了具有线性复杂度的全局依赖建模，并在多个数据集和预测设置中建立了新的最先进的结果。

---