技術要約：拡散駆動型状態空間モデル (DDSSM)

問題提起

時系列モデリングに従事する実務家は、モデルの表現力と潜在的なシステムダイナミクスを捉える能力の間にあるトレードオフに直面することが多い。既存のアプローチは、通常、一方の目標を犠牲にしている：

• 深層状態空間モデル (Deep State Space Models: DSSMs)（例：Deep Kalman Filter）は、原理に基づいた推論とエンドツーエンドの学習を実現するが、通常、ガウス型遷移密度を仮定する。この仮定は表現力を制限し、マルチモーダルなダイナミクスを捉えることに失敗することが多く、ガウス事前分布を導入することで潜在表現を過度に正則化し、必要な構造的情報を破棄してしまう。
• 拡散モデル (Diffusion Models) は、非常に高い表現力を持ち、複雑でマルチモーダルな分布をモデル化することができる。しかし、時系列に適用する場合、過去の観測値に条件付けられた生成サンプラーとして使用されるか、あるいは2段階の学習（まずオートエンコーダを学習し、次に拡散モデルを学習する）に依存することが多い。この2段階のアプローチは、ダウンストリームの生成ダイナミクスに必要な情報を破棄してしまう、過度に正則化された潜在空間を生み出す可能性がある。

これら両方のフレームワークの強みを組み合わせるために、オートエンコーダと拡散モデルを逐次データに対してどのように原理的な方法で共同学習させるかという点には、依然として未解決の問題が存在する。

手法

著者らは、拡散モデルを状態空間モデル (SSM) の遷移密度の中に直接組み込んだ拡散駆動型状態空間モデル (Diffusion-Driven State Space Model: DDSSM) を提案する。

生成モデル

本モデルは、観測値 $x_{1:T}$ が、共変量 $u_{1:T}$ に条件付けられる可能性のあるマルコフ過程を通じて進化する潜在状態 $z_{1:T}$ から生成されると仮定する。結合尤度は以下のように分解される：
$$p(x_{1:T}, z_{1:T} | u_{1:T}) = p_{\eta,\theta}(z_{1:j}, x_{1:j} | u_{1:j}) \prod_{t=j+1}^T p^{(t)}_\psi(z_t | z_{t-j:t-1}, u_{t-j:t}) p^{(t)}_\theta(x_t | z_{t-j+1:t}, u_{t-j+1:t})$$

主な革新点は以下の通りである：

1. 拡散遷移 (Diffusion Transitions): パラメトリックなガウス型遷移 $p^{(t)}_\psi$ の代わりに、モデルは拡散プロセスを使用する。各タイムステップ $t$ において、遷移密度は中間的なノイズを含む潜在状態 $z^1_t, \dots, z^K_t$ を導入することによってモデル化される。遷移は以下のように定義される：
   $$p^{(t)}_\psi(z^0_t, z^{1:K}_t | z_{t-j:t-1}, u_{t-j:t}) = p(z^K_t) \prod_{k=1}^K p^{(t)}_\psi(z^{k-1}_t | z^k_t, z_{t-j:t-1}, u_{t-j:t})$$
   ここで $z^0_t = z_t$ であり、$z^K_t \sim \mathcal{N}(0, I)$ である。これにより、モデルは任意かつ、潜在的にマルチモーダルな遷移密度を学習することが可能になる。
2. 放出 (Emissions): 放出密度 $p^{(t)}_\theta$ は、ニューラルネットワークによってパラメータ化された標準的なガウス分布のままであり、残差の不確実性を処理する。一方で、拡散モデルは複雑な非ガウス構造を捉える。
3. 初期化: 初期状態分布 $p_\eta(z_{1:j}|u_{1:j})$ は、制限的な仮定を避けるために変分階層事前分布 (Variational Hierarchical Prior: VHP) を利用する。

推論と学習

真の事後分布は計算困難であるため、著者らは変分推論フレームワークを採用している。彼らは、真の事後分布の条件付き独立構造を反映した証拠下界 (ELBO) を導出している。

• 変分事後分布: 事後分布は、平滑化された遷移項（過去の潜在状態と将来の観測値に依存）と、拡散ノイズ過程項を含むように分解される。
• 共同学習: 2段階の手法とは異なり、DDSSMはエンドツーエンドで学習される。目的関数は以下の要素を含むELBOを最大化する：
  • 再構成誤差（放出の負の対数尤度）。
  • 初期化のためのKLダイバージェンス項。
  • 学習された遷移ダイナミクスの下で予測可能な潜在状態をエンコーダに生成させるための、拡散チェーン損失（拡散モデルの逆プロセスから導出）。
• プリコンディショニング (Preconditioning): 共同学習中の様々なノイズレベルにおける勾配の安定性を確保するため、著者らは入力と拡散ネットワークの出力を再スケーリングするためのプリコンディショニング (Karras et al., 2022) を採用している。

主な貢献

1. DDSSMフレームワーク: 遷移密度を拡散モデルに置き換えた状態空間モデルを導入することで、制限的なガウス仮定なしに、柔軟でマルチモーダルな潜在ダイナミクスを学習することを可能にした。
2. エンドツーエンド学習: 逐次データに対してオートエンコーダと拡散モデルを共同学習させる際の、原理的な解決策を提示した。これにより、一般的な潜在時系列拡散モデルに関連する過剰な正則化の問題を回避している。
3. 理論的導出: このハイブリッドモデルに対するELBOの厳密な導出を行い、最適な変分事後分布が真の事後分布の構造と一致するように分解され、近似ギャップを狭めることを示した。

実験結果

著者らは、バイモーダル・ノイズ（バイモーダルな遷移を持つ潜在ランダムウォーク）を特徴とするシミュレーション時系列データセットを用いてDDSSMを評価した。

• 設定: 彼らは、VHPを備えたDeep Kalman Filter (DKF) および直近観測値保持法 (Last Observation Carried Forward: LOCF) のベースラインと比較した。
• 指標: モデルの1ステップ予測分布と真の分布の間のイェンセン・シャノン・ダイバージェンス (JSD) を用いて性能を測定した。
• 知見:
  • ガウス型遷移に制約されたDKFは、真の2つのモードの間に確率質量を配置してしまい、バイモーダルな構造を捉えることに失敗した。
  • DDSSMは、真の分布の両方のモードを正常に捉えることができた。
  • 定量的には、DDSSMはDKFよりも有意に低いJSDを達成した。
  • 定性的には、DDSSMはデータに代表的な再構成および予測を行ったのに対し、DKFの予測は一貫したパターンを欠き、観測データをうまく再構成できなかった。

重要性と主張

本論文は、潜在時系列拡散を「拡散遷移分布を持つ状態空間モデル」として捉えることで、潜在ダイナミクスの学習を妨げることなく、エンドツーエンドの学習が可能になり、柔軟な再構成が得られることを主張している。

著者らは、DDSSMを、拡散モデルの表現力とSSMの構造的な厳密さの間の緊張関係を解消する手法として位置づけている。従来のDeep SSMにおけるガウス事前分布の仮定は、マルチモーダルな設定において潜在表現を学習する上で有害であり、彼らのアプローチは、潜在軌跡に対してより表現力のある事前分布を提供すると主張している。この研究は、2段階の学習パラダイムを超えて、時系列の潜在拡散モデルの文献を拡張するものである。

著者らは、今後の課題として、大規模なデータセットへのDDSSMの適用や、拡散遷移分布の力を最大限に活用するための深層SSM文献からのさらなるアイデアの組み込みを挙げている。