🌊

【Julia×FastAPI】JSON疎結合・インデックス契約とDDM／Heunで長期ロールアウト安定化

に公開
Julia
Python
PyTorch
FastAPI
Graph neural network
tech

はじめに

教師信号に沿った短期予測だけを見ると動いてしまう一方で、自己回帰で数十〜数百ステップを積み重ねると数値拡散や発散が顕在化する——グラフニューラルネットによる偏微分方程式サロゲートでよく見られる落とし穴です。原因は単一ではなく、(i) 時間ディスクリゼーションが粗くエネルギー構造を保持しない、(ii) 全格子を一度に畳み込む設計では GPU メモリがボトルネック、(iii) 局所メッセージパッシングだけでは長距離相関を効率的に伝播できない、という三層に分解できます。

本稿ではオープンソースで公開した Physics GNN Surrogate · Long Rollout Stabilization において、これらをどう突破したかをまとめます。下文では、DDM／マルチグリッドの動く概念描画と長期ロールアウトの波形重合も順に示します。本プロジェクトの目的は次の四点です。

  • 長期ロールアウトでもハミルトニアン様の量が暴れにくい時間方向の安定化
  • DDM（領域分割）＋ Halo 同期による空間方向のスケールアウト
  • マルチグリッド的 Restriction / Prolongation とテンソル型メッセージパッシングによる長距離相互作用と、**Julia 参照ソルバ対比の ROI（Speedup）**の定量化
  • FastAPI と HTTP.jl による JSON 疎結合・E2E 検証で、インデックス契約を運用経路でも再現すること

ここでいう「DEC-GNN」は、離散勾配・離散発散がグラフエッジ上で結合する離散ハミルトニアン構造をニューラルネットで近似する、離散外微分形式（DEC）のスピリットを GNN に載せた設計を指します。産業応用では、同様のパイプラインを 格子／メッシュ上のスカラー場・ベクトル場の時間発展サロゲートに転用でき、設計変更やメッシュ更新が頻繁なケースで「何度もフル CFD を回さない」判断材料になります。

時間方向の安定化：Heun 法と Symplectic Loss

潜在状態 h に対し、GNN がベクトル場 f_\theta(h) \approx dh/dt を与えるとき、**Heun 法（改良オイラー）**は次のように書けます。中間状態 \tilde{h}^{(t+1)} = h^{(t)} + \Delta t\, f_\theta(h^{(t)}) を経由し、滑らかさと安定性のバランスが改善します。

h^{(t+1)} = h^{(t)} + \frac{\Delta t}{2} \left( f_{\theta}(h^{(t)}) + f_{\theta}(\tilde{h}^{(t+1)}) \right)

さらにデータ一致項だけでは長期的なエネルギー偏移が抑えきれない場合があるため、離散ハミルトニアン \mathcal{H}(h)（運動エネルギー項とエッジポテンシャル項の組）について、時間方向のジャンプをペナルティ化する Symplectic Loss を足します。

\mathcal{L}_{total} = \mathcal{L}_{data} + \lambda_{symp} \sum_{t} \left\| \mathcal{H}(h^{(t+1)}) - \mathcal{H}(h^{(t)}) \right\|^2

ハイパーパラメータ \lambda_{symp} は、ロールアウト長をカリキュラム学習で伸ばす過程とセットで調整すると効果的です。その結果、初期モデルで観測されやすかった「十数ステップ以降のエネルギー漂移」が 構造的にペナルティを受けるようになり、検証セットでの長い自己回帰 MSE の劣化が緩やかになります。

空間方向の大規模化：DDM と Halo 同期

数百万ノード級では、密な全結合や巨大バッチは現実的ではありません。そこで Metis 等によるグラフ分割 でサブドメインを構成し、各パッチはローカルノードと 1-hop Halo（ゴースト層） を保持します。時間ステップごとにメッセージパッシングと Heun 更新を パッチ内で実行したあと、境界値を sync_halo_features で全局整合させます。

この設計のメリットは二つあります。

  1. GPU メモリ：パッチ単位テンソルでピーク使用量を抑制できる。
  2. 合成可能性：「パッチ内动力学」と「境界同期」が別モジュールになり、テストとデバッグが独立する。

読み手には「単なるデータ並列」ではなく、離散ドメイン分解と整合条件を明示したモジュール分割と捉えてもらうと、既存の並列線形ソルバ文化との対応がイメージしやすいです。

DDMとHalo同期の概念図

サブドメイン境界に載せた Halo／ゴースト層 を同タイムステップで揃え直す過程です。境界付近のバッファを誇張して、通信と整合の順序が目に見えるようにしています。

長距離相互作用：マルチグリッドとテンソルネットワーク（MPS）的思考

局所 MP だけでは、情報がエッジホップで伝わるため深さ・ステップに依存します。細格子 \to 粗格子への Restriction R と 粗 \to 細への Prolongation P を sparse COOとして Julia 側 IR に埋め込み、Python では torch.sparse.mm で合成します。粗スケールでテンソル型メッセージパッシングを行うことで、メッシュ上の「ボンド次元」をもつ縮約付きメッセージとして長距離モードをキャプチャします（行列積状態／テンソルネットワークの「縮約」との対応は設計上のメタファとして有用です）。

マルチグリッドの概念図

細格子での restriction と粗格子での伝播、そこへ戻す prolongation が時系列として追えるよう分割したものです。色分けした矢印は、細⇄粗スケールの往復によって長波長モードを拾う視覚的な手がかりです。

時間積分・DDM・マルチグリッドはそれぞれ 関手的可能なモジュールとして分割されており、別物理（異なるエッジ属性・次元）へは IR と変換層を差し替えるだけで応用できます—これが次節の評価設計にも効いてきます。

評価と ROI：ゼロショットと Speedup

学習時と異なる格子サイズ・初期条件の JSON を Julia で生成し、Python で ゼロショット自己回帰ロールアウトを評価します。累積誤差は次の RMSE で定義します（\hat{u} は予測、u は参照、V はノード集合）。

\text{RMSE}_{rollout} = \sqrt{ \frac{1}{T \cdot |V|} \sum_{t=1}^{T} \sum_{i \in V} \left\| \hat{u}_i^{(t)} - u_i^{(t)} \right\|^2 }

併せて、離散ハミルトニアンに基づく エネルギー漂移（Drift） を記録し、推論 1 ステップあたりの GPU 時間（CUDA Event 等）と Julia ソルバのマクロステップ時間を比較して Speedup を算出します。規模・ハードウェア・バッチ条件により数値は変動しますが、「設計・検証ループを短縮できるサロゲート」としての 投資対効果（ROI） を定量で説明する材料になります。数千倍級の Speedup は、細かい時間刻みの参照積分対比でサロゲート 1 ステップを置換するシナリオで報告されうるオーダーであり、実運用では精度要件とドリフト許容でトレードオフを取ります。

ロールアウト波形の可視化と評価

ゼロショット自己回帰ロールアウトでは、学習に使っていない条件で Julia が生成した Ground Truth の時系列 \{ u_i^{(t)} \} に対し、GNN サロゲートが自分の出力を入力として再利用しながら得た \{ \hat{u}_i^{(t)} \} を同一プロットに重ね、そのまま時間を進めます。Python 側で次の (\text{RMSE}_{rollout}) を算出しつつ、代表的ノードの時系列を重ねて可視化し、動かしながら追跡しました。

\text{RMSE}_{rollout} = \sqrt{ \frac{1}{T \cdot |V|} \sum_{t=1}^{T} \sum_{i \in V} \left\| \hat{u}_i^{(t)} - u_i^{(t)} \right\|^2 }

ロールアウト推論波形の比較

青い実線が Ground Truth（Julia の参照波形）、橙色の破線が GNN サロゲート（自己回帰） による位移 u^{(t)} の推移です。数十マクロステップ先でも振幅が異常増幅せずに参照と追走し続ける様子から、離散ウェーブのエネルギー様構造が自己回帰連鎖のなかでも破綻しにくいことが読み取れます。

可視化の狙いは、短期の単発誤差では隠れる 位相ズレ・振幅歪み・ドリフトを、時間伸長でも追える形に落とすことです。(\text{RMSE}_{rollout}) とともに見ることで、Heun・Symplectic 損失・空間モジュールを束ねた設計がロールアウト鎖での意味を保つことが確認できます。

実運用 API と E2E テスト：Julia と FastAPI の疎結合

本プロジェクトの実装では、Julia 側がデータ生成・参照シミュレーション・IR 構築（ノード／エッジは 1-based 慣習）を担い、Python 側が PyTorch／PyTorch Geometric による学習・推論（テンソル・COO は 0-based）を担う、という適材適所の分割をそのまま運用レイヤまで延ばしています。

両言語は HTTP と JSON で疎結合に接続しています。その境界でグラフ構造やスパース行列のインデックスが無言のうちにズレると、単体テストが通っても本番だけ破綻する典型パターンに落ちます。そこで FastAPI を API 層に置き、ペイロードに含まれるノード ID・エッジ番号・COO などを 1-based (\leftrightarrow) 0-based のラウンドトリップ変換として集約し、Julia クライアント（HTTP.jl）から送受信される JSON と PyTorch 側のテンソル表現のあいだを仕様として固定しました。

この 言語間境界でのインデックス安全保障こそ、再現性のあるサロゲート運用にとって地味だが決定的な魅力です。実装リポジトリでは、Julia から API へ実データを投げ、変換後の推論結果が往復で破綻しないことを確認する E2E テストまで通しており、「ローカル JSON ファイルを手で差し替える評価」と同じ契約が、実運用想定の HTTP 経路でも守られることを検証済みです。

おわりに：インデックス契約と OSS

本プロジェクトでは、次の三つのコア思想を README とコードの両方において設計原則として徹底しました。

  1. Julia / Python の適材適所：参照シミュレーションと IR 生成は Julia、学習・推論は PyTorch / PyG。
  2. インデックス安全保障：Julia の 1-based を JSON に明示し、Python 側は convert_julia_to_python_indices（およびスパース COO 用変換）に一元集約。見えないオフバイワンバグを仕様で封じる。
  3. 応用圏論的モジュール性：時間積分・空間同期・射影を独立コンポーネントとして合成し、別モデル評価にも同じ評価関手（メトリクス・プロファイラ）を流用できる構造にした。

参考文献・リンク

Discussion