近年、時系列データの予測においては、再帰型ニューラルネットワークやAttention機構を用いた深層学習モデルが主流となっています。これらのモデルは、複雑な時系列データを捉える能力に優れていますが、一方でモデルの複雑化に伴う計算コストや過学習の問題も指摘されています。

しかし、最近の研究では、シンプルな線形モデルが、これらの深層学習モデルを凌駕する性能を示すケースがあることが明らかになってきました。そこで、この知見をもとに、時系列予測のための新たな線形モデルである TSMixer（Time-Series Mixer）が提案されました。TSMixerは、多層パーセプトロン（MLP）を積み重ねることで構成され、時間軸と特徴量の両方向の情報を効率的に抽出することができます。

今回は、このTSMixerについて、モデルの構造やPythonによる実装例を紹介します。

アーキテクチャ概観

TSMixerは、MLP-MixerのようにMLP（多層パーセプトロン）のみで構成されており、時系列データにおける時間依存性と変数間の関係を効果的に捉えるために設計されたモデルです。その核となるアイデアは、時間方向の混合と特徴量方向の混合を交互に行うMLPを組み合わせて情報を取り込むことです。

TSMixerの構成要素

TSMixerを構成するコンポーネントについては次の通りです:

• 時間方向の混合 (Time-mixing MLP)
  • すべての特徴量について共有されるMLP
  • 時間軸に沿って全結合層を適用し、時間的なパターンを学習
  • 一般的に単層のMLPで十分な性能を発揮
• 特徴量方向の混合 (Feature-mixing MLP)
  • すべての時間ステップについて共有されるMLP
  • 2層のMLPを使用して、変数間の複雑な関係を学習
• 混合の交互適用
  • 時間方向の混合と特徴量方向の混合を交互に繰り返すことで、時間的なパターンと変数間の関係を同時に考慮
  • この交互適用は、MLP-Mixerアーキテクチャに類似
• 時間投影 (Temporal Projection): 入力時系列の長さを予測対象の長さに変換する全結合層
• 残差接続 (Residual Connections): 学習を効率化し、不要な混合操作の影響を軽減
• 正規化 (Normalization): 時間軸と特徴量軸の両方に適用される2D正規化により、スケールを適切に調整

TSMixerの特徴

加えて、TSMixerの特徴についても言及しておきます:

• シンプルかつ効果的: 比較的シンプルな構造でありながら、高い予測精度を達成
• 計算量とパラメータ数の削減: 時間方向と特徴量方向の混合を交互に行うことで、計算量とパラメータ数を削減
• 柔軟性: 補助情報の活用など、さまざまな拡張が可能
• TSMixerとMLP-Mixerの違い:
  • 適用領域: TSMixerは時系列データに特化しており、MLP-Mixerは画像データ向けに設計されいる
  • 正規化手法: TSMixerは時間軸と特徴量軸の両方を考慮した2D正規化を採用

TSMixerの機能拡張

また、TSMixerは、時系列データに加えて、以下のような補助情報を扱うように拡張できます。

• 静的特徴量: 時間によらず一定の値をとる特徴量 (例: 地域情報)
• 動的特徴量: 時間とともに変化する特徴量 (例: 販促キャンペーン情報)
• 複数時系列データ: 複数の時系列データとその関連する特徴量を同時に扱う。

Extended TSMixer（拡張されたTSMixer）のアーキテクチャは次の:

• 調整ステージ（Align Stage）
  • 異なる形状の入力特徴量 (時系列データ、静的特徴量、動的特徴量) を同じ形状に変換
  • 時間投影 (Temporal Projection)と特徴量混合層 (Feature Mixing)を用いて、時系列データと動的特徴量を同じ長さのベクトルに変換
  • 静的特徴量を繰り返し適用して、動的特徴量と同様の長さのベクトルに変換
• 混合ステージ（Mixing Stage）
  • 時間方向の混合 (Time Mixing) と特徴量方向の混合 (Feature Mixing)を交互に適用して、すべての入力特徴量から時間的なパターンと変数間の関係を抽出し統合
• 出力層
  • 全結合層を用いて、各時間ステップの予測値を出力
  • 予測値は、実際の値そのもの (回帰問題) または、目標分布のパラメータ (例: 小売需要予測における負の二項分布) となる場合がある

TSMixerの性能

TSMixerは、多変量時系列予測において、特に大規模かつ複雑なデータセットに対して高い性能を示すことが報告されています。評価結果の要点は次の通りです。

• シンプルかつ効果的: MLPをベースにしたシンプルな構造でありながら、時間方向と特徴量方向の情報を効果的に統合することで、高度な性能を実現。
• 多様なタスクに対応: 気象、電力、小売など、様々な分野のデータセットで評価を行い、その汎用性の高さが示された。
• 大規模データセットへの強み: M5データセットのような大規模な小売データにおいて、補助情報を効果的に活用し、従来モデルを上回る性能を達成。
• 計算コストの低さ: RNNやTransformerベースのモデルと比較して、モデルサイズが小さく、高速な推論が可能。

TSMixerのPython実装

TSMixerには、PyPIで管理されているPythonパッケージがあるので、簡単に試すことができます。このパッケージは、PyTorchベースで実装されています。まずは、PyTorchをインストールします（インストールバージョンの指定については公式ページを参照してください）。

$ pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

続けて、pytorch-tsmixerをインストールします。

$ pip install pytorch-tsmixer

ではまず、TSMixerのモデルを構築する方法ですが、次のように記述します。

import torch
from torchtsmixer import TSMixer

model = TSMixer(
    sequence_length=10,     # 内部変換後のシーケンス長
    prediction_length=5,    # 出力シーケンス長
    input_channels=2,       # 入力データのチャンネル数
    output_channels=4,      # 出力データのチャンネル数
    activation_fn="relu",   # 活性化関数の指定
    num_blocks=2,           # Mixerブロックの数
    dropout_rate=0.1,       # ドロップアウト率
    ff_dim=64,              # モデルの隠れ層であるフィードフォワード層の次元
    normalize_before=True,  # 事前正規化を行うかどうか
    norm_type="batch",      # 正規化のタイプ ("batch" or "layer")
)

# (バッチサイズ, 入力シーケンス長, チャネル数)
x = torch.randn(3, 10, 2)
y = model(x)
print(y.shape)    # => (3, 5, 4)

また、Extended TSMixer（拡張されたTSMixer）については次のようになります。

import torch
from torchtsmixer import TSMixerExt

model = TSMixerExt(
    sequence_length=10,     # 内部変換後のシーケンス長
    prediction_length=5,    # 出力シーケンス長
    input_channels=2,       # 入力データのチャンネル数
    extra_channels=3,       # (将来既知の) 追加データの入力のチャネル数。
    hidden_channels=8,      # 隠れ層のチャンネル数
    static_channels=4,      # 静的情報のチャンネル数
    output_channels=4,      # 出力データのチャンネル数
    activation_fn="relu",   # 活性化関数の指定
    num_blocks=2,           # Mixerブロックの数
    dropout_rate=0.1,       # ドロップアウト率
    ff_dim=64,              # モデルの隠れ層であるフィードフォワード層の次元
    normalize_before=True,  # 事前正規化を行うかどうか
    norm_type="batch",      # 正規化のタイプ ("batch" or "layer")
)

x_hist = torch.randn(3, 10, 2, requires_grad=True)
x_extra_hist = torch.randn(3, 10, 3, requires_grad=True)
x_extra_future = torch.randn(3, 5, 3, requires_grad=True)
x_static = torch.randn(3, 4, requires_grad=True)

y = model.forward(
    x_hist=x_hist,
    x_extra_hist=x_extra_hist,
    x_extra_future=x_extra_future,
    x_static=x_static
)

最後に、モデルの訓練についての実装例を、以下に紹介します。

import torch
import torch.nn as nn
import torch.optim as optim
from torchtsmixer import TSMixer

# モデルのパラメータ
sequence_length = 10  # 入力シーケンスの長さ
prediction_length = 5  # 予測シーケンスの長さ
input_channels = 2  # 入力チャネル数
output_channels = 1  # 出力チャネル数

# TSMixerモデルの作成
model = TSMixer(sequence_length, prediction_length, input_channels, output_channels)
model.to("cuda:0")

# 損失関数と最適化アルゴリズム
criterion = nn.MSELoss()  # 平均二乗誤差 (MSE) を損失関数として使用
optimizer = optim.Adam(model.parameters(), lr=0.0001)  # Adamオプティマイザ

# ダミーデータセット (実際のデータに置き換える必要があります)
# フォーマット: (バッチサイズ, シーケンス長, 入力チャネル数)
X_train = torch.randn(10, 1024, sequence_length, input_channels).to("cuda:0")
y_train = torch.randn(10, 1024, prediction_length, output_channels).to("cuda:0")

# トレーニングループ
num_epochs = 10  # エポック数

for epoch in range(num_epochs):
    model.train()  # モデルをトレーニングモードに設定

    for X, y in zip(X_train, y_train):
        optimizer.zero_grad()  # パラメータの勾配を初期化

        # 順伝播
        outputs = model(X)
        loss = criterion(outputs, y)  # 損失を計算

        # 逆伝播とパラメータの更新
        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')  # エポックごとの損失を出力

print("トレーニング完了")

まとめ

今回は、従来のRNNやアテンション機構に代えてMLPを用いることで、シンプルなアーキテクチャでありながら優れた汎化性能を実現する、時系列予測のための新しいアーキテクチャであるTSMixerを紹介しました。

時系列予測タスクにおける様々な実験結果から、TSMixerが多変量時系列データの長期予測ベンチマークと、大規模な小売需要予測タスクの両方において高い有効性を示しています。特に、TSMixerは長期時系列予測ベンチマークにおいて、単変量モデルと同等の性能を達成できる唯一の多変量モデルです。

TSMixerアーキテクチャにはさらなる改善の可能性が大きく、さらに幅広い時系列予測タスクにおいて有効であると考えられます。

TSMixerのさらに詳しい情報については、以下の原論文を参照してください。また、Pythonの具体的な実装については、GitHubリポジトリを参照してください。

More Information:

• arXiv:2303.06053, Si-An Chen et al., 「TSMixer: An All-MLP Architecture for Time Series Forecasting」, https://arxiv.org/abs/2303.06053