[Paper Review] A Mamba Foundation Model for Time Series Forecasting

Haoyu Ma, Yushu Chen+8 more|arXiv (Cornell University)|Nov 5, 2024

Forecasting Techniques and ApplicationsDecision Sciences3 citations

TL;DR

TSMamba is a linear-complexity foundation model for time series forecasting based on the Mamba architecture, using dual forward and backward Mamba encoders to capture long-range temporal dependencies. It achieves state-of-the-art zero-shot and full-shot performance with significantly less training data by leveraging a two-stage transfer learning process and a channel-wise compressed attention module for cross-channel modeling.

ABSTRACT

Time series foundation models have demonstrated strong performance in zero-shot learning, making them well-suited for predicting rapidly evolving patterns in real-world applications where relevant training data are scarce. However, most of these models rely on the Transformer architecture, which incurs quadratic complexity as input length increases. To address this, we introduce TSMamba, a linear-complexity foundation model for time series forecasting built on the Mamba architecture. The model captures temporal dependencies through both forward and backward Mamba encoders, achieving high prediction accuracy. To reduce reliance on large datasets and lower training costs, TSMamba employs a two-stage transfer learning process that leverages pretrained Mamba LLMs, allowing effective time series modeling with a moderate training set. In the first stage, the forward and backward backbones are optimized via patch-wise autoregressive prediction; in the second stage, the model trains a prediction head and refines other components for long-term forecasting. While the backbone assumes channel independence to manage varying channel numbers across datasets, a channel-wise compressed attention module is introduced to capture cross-channel dependencies during fine-tuning on specific multivariate datasets. Experiments show that TSMamba's zero-shot performance is comparable to state-of-the-art time series foundation models, despite using significantly less training data. It also achieves competitive or superior full-shot performance compared to task-specific prediction models. The code will be made publicly available.

Motivation & Objective

• Address the limitations of specialized time series models in generalizing across domains and adapting to emerging patterns with scarce data.
• Overcome the quadratic complexity of Transformer-based models in long sequence forecasting by adopting the Mamba architecture.
• Reduce data and computational costs for time series foundation models by leveraging pretrained Mamba language models via transfer learning.
• Enable effective multivariate time series forecasting by introducing a compressed cross-channel attention module during fine-tuning.
• Demonstrate strong zero-shot and full-shot performance comparable to or better than state-of-the-art models with minimal fine-tuning data.

Proposed method

• Employ a dual Mamba encoder (forward and backward) to model long-range temporal dependencies in time series with linear complexity.
• Use a two-stage transfer learning strategy: first pretraining the Mamba backbone via patch-wise autoregressive prediction, then fine-tuning the prediction head and other components.
• Leverage pretrained Mamba large language models to initialize the time series foundation model, reducing reliance on large-scale time series datasets.
• Process each channel independently in the backbone to support variable channel counts across datasets, while introducing a compressed cross-channel attention module during fine-tuning to model inter-channel relationships.
• Apply a two-stage training pipeline: stage one optimizes the Mamba encoders on masked patch prediction; stage two fine-tunes the prediction head and refines the encoder features for long-term forecasting.
• Use a channel-independent Mamba backbone for scalability and adaptability across datasets with different numbers of channels, while enabling cross-channel modeling only during task-specific fine-tuning.

Experimental results

Research questions

• RQ1 Can a Mamba-based foundation model achieve competitive zero-shot forecasting performance across diverse time series datasets with minimal training data?
• RQ2 How does the two-stage transfer learning approach using pretrained Mamba LLMs improve data efficiency in time series forecasting?
• RQ3 To what extent can a linear-complexity Mamba architecture outperform quadratic-complexity Transformers in long-horizon time series forecasting?
• RQ4 How effective is the proposed compressed cross-channel attention module in capturing inter-channel dependencies during fine-tuning on multivariate time series?
• RQ5 Can a foundation model trained on diverse data generalize effectively to unseen time series tasks and domains without extensive retraining?

Key findings

• TSMamba achieves state-of-the-art zero-shot performance on benchmark datasets like ETTm2, Weather, and ILI, outperforming existing foundation models despite using significantly less training data.
• In full-shot settings, TSMamba achieves a 15% average performance gain over GPT4TS (Zhou et al., 2023), a GPT2-based foundation model, across multiple prediction horizons.
• TSMamba outperforms the state-of-the-art task-specific model PatchTST (Nie et al., 2023) on all evaluated datasets and prediction horizons.
• The model maintains strong performance across diverse data types, including univariate and multivariate time series with varying frequencies and channel counts.
• The two-stage transfer learning process enables effective knowledge transfer from Mamba LLMs, significantly reducing the need for large-scale time series pretraining.
• The compressed cross-channel attention module effectively captures inter-channel dependencies during fine-tuning, improving performance on multivariate datasets without increasing backbone complexity.