[논문 리뷰] A Mamba Foundation Model for Time Series Forecasting

Haoyu Ma, Yushu Chen외 8명|arXiv (Cornell University)|2024. 11. 05.

Forecasting Techniques and ApplicationsDecision Sciences인용 수 3

한 줄 요약

TSMamba는 Mamba 아키텍처를 기반으로 한 선형 복잡도 기반 모델로, 장기적인 시간적 의존성을 포괄적으로 포착하기 위해 이중 전방 및 역방향 Mamba 인코더를 사용한다. 이중 전이 학습 프로세스와 채널 별 압축된 어텐션 모듈을 통해 다중 채널 모델링을 구현함으로써, 훈련 데이터가 훨씬 적은 상황에서도 최신 기준의 제로샷 및 풀샷 성능을 달성한다.

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.

연구 동기 및 목표

• 다양한 도메인 간 일반화 능력이 떨어지고 데이터가 부족한 상황에서의 패턴 적응 능력이 떨어지는 전문화된 시계열 모델의 한계를 해결한다.
• 장기 예측에서 트랜스포머 기반 모델의 이차 복잡도 문제를 해결하기 위해 Mamba 아키텍처를 채택한다.
• 전이 학습을 통해 미리 학습된 Mamba 언어 모델을 활용함으로써 시계열 기초 모델의 데이터 및 계산 비용을 절감한다.
• 세부 작업에 맞는 피팅 단계에서 압축된 다중채널 어텐션 모듈을 도입하여 효과적인 다변량 시계열 예측을 가능하게 한다.
• 최소한의 피팅 데이터로도 최신 기준 모델과 비교해 뛰어난 성능을 보이는 강력한 제로샷 및 풀샷 성능을 입증한다.

제안 방법

• 장기적인 시간적 의존성을 선형 복잡도로 모델링하기 위해 전방 및 역방향 Mamba 인코더를 이중으로 활용한다.
• 두 단계 전이 학습 전략을 사용한다: 첫 번째 단계에서는 패치 단위 자동회귀 예측을 통해 Mamba 백본을 사전학습하고, 두 번째 단계에서는 예측 헤드 및 기타 구성 요소를 피팅한다.
• 시계열 데이터셋의 대규모 요구량에 의존하지 않도록, 사전 학습된 Mamba 대규모 언어 모델을 활용해 시계열 기초 모델을 초기화한다.
• 백본에서 각 채널을 별도로 처리하여 다양한 데이터셋 간 채널 수의 변동성을 수용한다. 동시에 피팅 단계에서 압축된 다중채널 어텐션 모듈을 도입해 채널 간 관계를 모델링한다.
• 두 단계 훈련 파이프라인을 적용한다: 첫 번째 단계에서는 마스크된 패치 예측에 대해 Mamba 인코더를 최적화하고, 두 번째 단계에서는 예측 헤드를 피팅하고 장기 예측을 위한 인코더 특징을 정밀 조정한다.
• 채널 간 독립적인 Mamba 백본을 사용해 다양한 데이터셋에 대한 확장성과 적응성을 확보하면서, 다중채널 모델링은 작업에 특화된 피팅 단계에서만 활성화한다.

실험 결과

연구 질문

• RQ1 Mamba 기반 기초 모델은 훈련 데이터가 극히 적은 상황에서도 다양한 시계열 데이터셋에서 경쟁력 있는 제로샷 예측 성능을 달성할 수 있는가?
• RQ2 사전 학습된 Mamba LLM을 활용한 두 단계 전이 학습 전략은 시계열 예측에서 데이터 효율성을 어떻게 향상시키는가?
• RQ3 선형 복잡도를 가진 Mamba 아키텍처는 이차 복잡도를 가진 트랜스포머보다 장기 예측에서 얼마나 뛰어난 성능을 보일 수 있는가?
• RQ4 제안된 압축된 다중채널 어텐션 모듈은 다변량 시계열 피팅 단계에서 채널 간 의존성을 얼마나 효과적으로 포착하는가?
• RQ5 다양한 데이터로 학습된 기초 모델은 광범위한 재학습 없이도 예측 작업 및 도메인에 대해 효과적으로 일반화될 수 있는가?

주요 결과

• TSMamba는 ETTm2, Weather, ILI와 같은 벤치마크 데이터셋에서 기존 기초 모델보다 훨씬 적은 훈련 데이터로도 최신 기준의 제로샷 성능을 달성한다.
• 풀샷 설정에서는 GPT4TS(Zhou et al., 2023) — GPT2 기반 기초 모델 — 와 비교해 다중 예측 수준에서 평균 15%의 성능 향상을 기록한다.
• 모든 평가된 데이터셋과 예측 수준에서 최신 기준의 전용 모델인 PatchTST(Nie et al., 2023)를 능가한다.
• 단변량 및 다변량 시계열, 다양한 주파수와 채널 수를 가진 다양한 데이터 유형에서 강력한 성능 유지를 유지한다.
• 두 단계 전이 학습 프로세스를 통해 Mamba LLM에서의 지식 전이가 효과적으로 이루어져, 대규모 시계열 사전학습의 필요성을 크게 감소시킨다.
• 압축된 다중채널 어텐션 모듈은 백본의 복잡도를 증가시키지 않으면서도 피팅 단계에서 채널 간 의존성을 효과적으로 포착하여 성능 향상을 이룬다.