[Paper Review] MambaTS: Improved Selective State Space Models for Long-term Time Series Forecasting

Xiuding Cai, Yaoyao Zhu+2 more|arXiv (Cornell University)|May 26, 2024

Stock Market Forecasting Methods11 citations

TL;DR

MambaTS extends selective state space models with variable-aware scanning and permutation strategies to achieve state-of-the-art long-term multivariate time series forecasting with linear complexity.

ABSTRACT

In recent years, Transformers have become the de-facto architecture for long-term sequence forecasting (LTSF), but faces challenges such as quadratic complexity and permutation invariant bias. A recent model, Mamba, based on selective state space models (SSMs), has emerged as a competitive alternative to Transformer, offering comparable performance with higher throughput and linear complexity related to sequence length. In this study, we analyze the limitations of current Mamba in LTSF and propose four targeted improvements, leading to MambaTS. We first introduce variable scan along time to arrange the historical information of all the variables together. We suggest that causal convolution in Mamba is not necessary for LTSF and propose the Temporal Mamba Block (TMB). We further incorporate a dropout mechanism for selective parameters of TMB to mitigate model overfitting. Moreover, we tackle the issue of variable scan order sensitivity by introducing variable permutation training. We further propose variable-aware scan along time to dynamically discover variable relationships during training and decode the optimal variable scan order by solving the shortest path visiting all nodes problem during inference. Extensive experiments conducted on eight public datasets demonstrate that MambaTS achieves new state-of-the-art performance.

Motivation & Objective

• Motivate improved long-term multivariate time series forecasting with efficient, scalable models beyond Transformers.
• Investigate limitations of Mamba in LTSF and develop targeted enhancements for better global and variable dependencies.
• Propose a suite of innovations (VST, TMB, dropout for selective parameters, VPT, VAST) to achieve state-of-the-art results.
• Demonstrate scalability and competitive performance across eight public LTSF datasets.

Proposed method

• Adopt Variable Scan along Time (VST) to form a global retrospective sequence by alternately mixing tokens from different variables at the same timestep.
• Replace causal convolution in Mamba with Temporal Mamba Block (TMB) by removing local convolution before SSMs.
• Introduce dropout on selective parameters within TMB to mitigate overfitting.
• Apply Variable Permutation Training (VPT) to reduce sensitivity to unknown variable order and enhance local interaction capacity.
• Develop Variable-Aware Scan along Time (VAST) to infer optimal variable scan order by estimating variable relationships and solving an asymmetric traveling salesman problem (ATSP) with a simulated annealing solver during inference.
• Maintain a channel-wise decoding head and use RevIN-inspired instance normalization for stability during training.

Experimental results

Research questions

• RQ1 Can Mamba-based architectures be adapted to efficiently model long-term dependencies in multivariate time series without quadratic complexity?
• RQ2 Do variable-order sensitivity and local-context interactions limit Mamba in LTSF, and can permutation and variable-aware strategies mitigate these issues?
• RQ3 Can a data-driven, inference-time optimal variable scan order improve LTSF performance over fixed or randomly shuffled orders?
• RQ4 What is the impact of removing causal convolution in temporal blocks for LTSF, and does dropout help counteract overfitting in selective Mamba blocks?
• RQ5 How do the proposed components interact to yield state-of-the-art results across diverse LTSF datasets?

Key findings

• MambaTS achieves new state-of-the-art results on most long-term multivariate forecasting tasks across eight public datasets.
• VST enables a global retrospective representation by aggregating histories from all variables.
• Temporal Mamba Block (TMB)—without causal convolution—improves temporal modeling efficiency and, with dropout, reduces overfitting.
• Variable Permutation Training (VPT) and Variable-Aware Scan along Time (VAST) substantially mitigate sensitivity to variable order and improve local interaction capacity.
• VAST uses a learned cost matrix to guide a simulated-annealing ATSP solver, yielding effective scan orders during inference.
• Overall, MambaTS attains linear (in sequence length) computational complexity and strong empirical performance compared to Transformer-based and other baselines.