CrossCat は、MIT の Probabilistic Computing Project (probcomp) が開発した、探索的データ解析のためのベイズ手法です。行列データの列をいくつかの「ビュー(view)」に分割し、各ビューは行をさらに「クラスタ」に分割して、連続値・カテゴリ値が混在するデータの背後にある依存構造を教師なしで発見します。これを卒業研究として確率的プログラミング言語 GenJAX で再実装し、オリジナルの Python 2 実装と比較しました。
卒業研究としてやったこと
研究の目的は次の 2 点を実証することでした。
- 統計的妥当性: GenJAX 版がオリジナル版と同等の推論精度・挙動を示し、真の生成モデルの構造を復元できること
- 計算性能: 実用的なシナリオ(行数・列数の増加)に対して、オリジナル版よりも高速に収束し、高いスケーラビリティを持つこと
実装
- 混合型 SBP(Stick-Breaking Process)マルチビュー CrossCat の生成モデルを GenJAX で記述
- 行クラスタ・列ビュー・クラスタパラメータ・SBP スティック・ハイパーパラメータをそれぞれ更新するトレースベースの Gibbs サンプリングカーネル
jitted_updateによる JIT コンパイルで、全行・全ビューの一括並列更新をベクトル化- Docker 化された 2 つの環境(GenJAX / オリジナル Python 2.7)で再現性を確保
- Synthetic / DHA / Adult の 3 データセットで、前処理→ベンチマーク→可視化を共通パイプラインに統一
実験結果
速度測定の比較(1 イテレーションあたりの平均時間、CPU):
| データセット | Original | GenJAX |
|---|---|---|
| Adult (10,000 行 × 15 列) | 0.75 s | 0.067 s |
| Synthetic (1,000 × 50) | 0.104 s | 0.032 s |
| Synthetic (1,000 × 10) | 0.030 s | 0.022 s |
| DHA (307 × 64) | 0.061 s | 0.045 s |
データが大きくなるほど高速化率が上がり、Adult では約 11 倍になりました。JIT コンパイルの初期コストは反復回数・データ規模に応じて償却され、特に行数に対するスケーラビリティが大きく改善されました。
リポジトリ
- ソース: https://github.com/shohei81/GenJAX_CrossCat
- Docker Compose で両環境を起動でき、
examples/dha_genjax.pyで GenJAX 版の DHA デモを再現できます