probml/dynamax

解決する課題

Dynamaxは、確率状態空間モデル（SSM）を扱うための高性能なライブラリを提供します。隠れた状態が時間の経過とともに変化し、観測データを生成するシステムにおいて、状態推定（推論）とパラメータ推定（学習）を行うプロセスを簡素化します。

仕組み

JAX上に構築されたこのライブラリは、隠れマルコフモデル（HMM）、線形ガウスSSM、非線形/一般化ガウスSSMを含む、さまざまなSSMを実装しています。APIは2つのレイヤーを提供します。機能的に純粋な低レベルの推論アルゴリズムのセットと、モデルクラスのためのユーザーフレンドリーなオブジェクト指向インターフェースです。JAXエコシステムと統合されており、確率的勾配降下法にはOptaxを、ハミルトニアンまたは逐次モンテカルロ法にはBlackjaxを使用します。

対象者

状態およびパラメータ推定に効率的なベクトル化演算を必要とする、時系列データや確率モデリングに取り組む研究者や開発者向けに設計されています。

ハイライト

• JAX駆動: 高性能コンピューティングと vmap による演算のベクトル化を実現するためにJAXを活用しています。
• 多様なモデルサポート: HMM、線形ガウス、非線形ガウス、および一般化ガウス状態空間モデルが含まれます。
• 包括的なツール群: オンラインフィルタリング、オフラインスムージング、および将来の状態と観測の予測をサポートしています。
• エコシステムとの統合: 高度なパラメータ推定とベイズ推論のために、OptaxおよびBlackjaxと互換性があります。