A Survey of Neural Network Variational Monte Carlo from a Computing Workload Characterization Perspective

本論文は、量子多体問題解決のためのニューラルネットワーク変分モンテカルロ法（NNVMC）の代表的な 4 つのアンサッツを対象に、GPU 上での実行特性を包括的に調査し、演算密度の低い要素別カーネルやデータ移動がボトルネックとなっていることを明らかにするとともに、スケーラブルな NNVMC システム構築に向けたアルゴリズムとハードウェアの共設計の指針を提示しています。

要約

🧪 1. 何をしているのか？（背景）

化学や材料科学の世界では、「新しい薬を作る」や「電池の性能を上げる」ために、電子がどう動き回っているかを計算する必要があります。
昔は、この計算には「超優秀な数学者（従来の計算手法）」が必要で、非常に正確ですが、計算コストが天文学的に高く、時間がかかりすぎていました。

そこで登場したのが、**「AI（ニューラルネットワーク）を使った計算」です。
これは、「経験則で電子の動きを学習する天才アシスタント」**のようなもので、従来の方法より圧倒的に速く、かつ正確な答えを出せる可能性があります。

🐢 2. 問題点はどこ？（発見）

「AI なら速いはず！」と思いきや、実際のところ**「とても重くて遅い」という問題がありました。
この論文は、その「重さ」の原因を、「料理の工程」**に例えて詳しく分析しました。

🍳 料理の例え：なぜ厨房（GPU）が混雑するのか？

この AI 計算は、大きく分けて 4 つの工程（ステージ）を繰り返します。

1. 材料の準備（Embedding）: 電子の位置情報を整理する。
2. 調理（Propagation）: 電子同士の関係性を AI が計算する。
3. 盛り付け（Readout）: 計算結果を形にする。
4. 味見と修正（Derivative/Laplacian）: 「これで正しいか？」を確認し、微調整する。

【ここが重要！】
従来の AI（画像認識やチャットボット）は、主に**「大きな鍋で大量の炒め物（行列計算）」**をするのが得意です。これは厨房（GPU）が最も得意とする仕事で、とても速く進みます。

しかし、この「電子シミュレーション」では、**「味見と修正（工程 4）」の段階で、「スプーンで 1 粒ずつ味見をする」「材料を 1 個ずつ並べ替える」ような、「細々とした作業」**が大量に発生してしまいます。

• 従来の AI = 大きな鍋で炒め物をする（得意！速い！）
• 電子シミュレーション = 炒め物もするが、**「1 粒ずつ味見して、材料を並べ替える」**作業が大半を占めてしまう。

今の GPU（厨房）は、「大きな鍋（行列計算）」を爆速で回せるように作られていますが、「スプーンで 1 粒ずつ扱う作業」には向いていません。
そのため、厨房が空回りしてしまい、**「材料を運ぶ時間（メモリー転送）」**がボトルネックになって、全体が遅くなってしまうのです。

🔍 3. 4 つの「レシピ」を比較した結果

論文では、4 つの異なる AI モデル（FermiNet, PauliNet, Psiformer, Orbformer）を比較しました。これらは「電子の動きをどう学習させるか」の異なるレシピです。

• FermiNet / PauliNet:
  • 「スプーンで味見する作業（細々した計算）」が非常に多い。
  • 結果: 厨房がパンクしやすく、メモリ（材料置き場）の容量不足に陥りやすい。
• Psiformer:
  • 「大きな鍋（行列計算）」の比率が増えた。
  • 結果: 以前より速くなったが、まだ「味見」の作業が邪魔をしている。
• Orbformer:
  • 「鍋」の作業は減ったが、「材料の並べ替え（データ移動）」が増えた。
  • 結果: 結局のところ、**「材料を運ぶ速度」**が全体のスピードを決定づけてしまう。

結論: どのレシピを使っても、「計算能力（CPU/GPU のパワー）」が余っているのに、「データ移動（メモリー帯域）」が追いついていないことが遅さの正体でした。

🚀 4. 今後の解決策（提案）

この「遅さ」を解消するために、論文では以下のような**「新しい厨房の設計図」**を提案しています。

1. 冷蔵庫の横で調理する（PIM: メモリ内処理）
   • 材料を冷蔵庫（メモリ）から運んでくるのが大変なら、冷蔵庫の横に小さな調理台を置いて、そこで「1 粒ずつの味見」をしてしまおう。運ぶ時間を減らす作戦です。
2. 作業を分担する（GPU と PIM の連携）
   • 「大きな鍋（行列計算）」は GPU で、「細々した味見（データ移動）」は冷蔵庫横の装置でやるように、作業を賢く分けよう。
3. 状況に合わせて厨房を変える（再構成可能ハードウェア）
   • 料理の工程（ステージ）によって、必要な道具が変わります。「炒め物モード」と「味見モード」で厨房の設備を切り替えられるようにしよう。
4. 倉庫（CPU/SSD）を活用する
   • 冷蔵庫（GPU メモリ）がパンパンになったら、大きな倉庫（CPU や SSD）に材料を一時預けて、必要な分だけ運ぶようにしよう。

📝 まとめ

この論文は、**「AI で化学計算をするには、単に AI を強くするだけではダメで、データの『運び方』と『作業の分け方』を根本から変える必要がある」**と教えてくれました。

まるで、「スーパーカー（GPU）で、渋滞する細い道（データ移動）を走らせようとしている」ような状態です。
これからは、「車（ハードウェア）」と「運転方法（アルゴリズム）」を一緒に設計し直して、細い道でもスムーズに走れるようにすることが、次世代の科学計算の鍵になるでしょう。

技術概要

論文「A Survey of Neural Network Variational Monte Carlo from a Computing Workload Characterization Perspective」の技術的サマリー

本論文は、量子化学における電子構造計算を目的とした**ニューラルネットワーク変分モンテカルロ法（NNVMC）**の計算負荷特性を、GPU 実行の観点から包括的に調査・分析したものです。従来の言語やビジョン分野の AI ワークロードとは異なり、NNVMC は物理法則に特化した実行ステージを持ち、そのパフォーマンスボトルネックも独特であることを実証的に明らかにしました。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

---

1. 問題定義 (Problem)

量子多体問題（特に電子シュレーディンガー方程式）の求解において、NNVMC は従来の手法（CCSD(T) や FCI など）と比較して、高い精度と有利なスケーリング（$O(N^4)$ 対 $O(N^7)$）を提供する有望なパラダイムとして注目されています。しかし、実用的な展開には以下の課題が存在します。

• 高い計算コストとメモリ制約: 現代の GPU 上での実行時間は長く、メモリ消費も大きいため、大規模な分子系へのスケーリングが制限されています。
• ワークロードの特殊性: 言語モデルや画像認識とは異なり、NNVMC は「マルコフ連鎖モンテカルロ（MCMC）サンプリング」「波動関数の構築」「微分/ラプラシアン評価」といった物理特有のステージを含みます。
• 単純な FLOP 数の限界: 総浮動小数点演算回数（FLOPs）だけでは、実行時間やメモリ挙動を予測できず、カーネルレベルの細かな挙動（特にデータ移動と演算のバランス）がボトルネックとなります。

2. 手法 (Methodology)

本研究では、代表的な 4 つの NNVMC アンスアツ（波動関数の仮定）である FermiNet, PauliNet, Psiformer, Orbformer を対象に、統一的なプロファイリングプロトコルを用いた実証分析を行いました。

• 対象モデルとコードベース:
  • DEEPQMC: FermiNet と PauliNet（JVP ベースのラプラシアン評価を使用）。
  • ONEQMC: Psiformer と Orbformer（Hutchinson 法に基づくラプラシアン評価を使用）。
• ハードウェア環境: NVIDIA RTX A5000, A100, H200 GPU。
• プロファイリング手法:
  • エンドツーエンド計測: 訓練・推論時の実行時間とメモリ使用量の測定。
  • カーネルレベル分析: NVIDIA Nsight Compute を用いた詳細な計測。
    • 算術強度 (Arithmetic Intensity, AI): 演算量とメモリアクセス量の比率。
    • ルーフラインモデル: カーネルが計算ボトルネックかメモリボトルネックかを特定。
    • ハードウェア利用率: SM（ストリーミングマルチプロセッサ）利用率、Tensor Core 活動、L2 キャッシュヒット率、メモリスループットなどの計測。
• 実験設定: 4 つの分子（LiH, CH4, C2H6, C4H4）に対して、FP32 精度、バッチサイズ 1024 で評価。

3. 主要な貢献 (Key Contributions)

1. NNVMC アンスアツのワークロード指向レビュー: DEEPQMC と ONEQMC の実装を含む、代表的な 4 つのモデルのアーキテクチャと実行パイプライン（ステージ A〜E）を体系的に整理しました。
2. 演算子・カーネルレベルの特性分析: 実測された算術強度とルーフライン分析を通じて、なぜ融合された要素ごとの演算（elementwise）やデータ移動カーネルが、一般的な行列積（GEMM）カーネルと同様に、あるいはそれ以上に実行時間を支配しているのかを説明しました。
3. ハードウェアレベルの挙動報告と共設計の示唆: SM 利用率、L2 キャッシュヒット率などの詳細な指標を報告し、スケーラブルな NNVMC システムに向けた「アルゴリズムとハードウェアの共設計」の方向性（フェーズ感知スケジューリング、メモリ中心最適化、異種アクセラレーション）を議論しました。

4. 主要な結果 (Key Results)

A. 実行時間とメモリの傾向

• モデル依存性: 実行時間のスケーリングはモデル設計に強く依存します。
  • PauliNet/FermiNet: 分子サイズが大きくなると実行時間が急激に増加（30〜42 倍）。これは、ラプラシアン評価のためにステージ A〜D の計算を多数回再実行する（JVP リプレイ）ため、微細な要素ごとの演算カーネルが支配的になるためです。
  • Psiformer/Orbformer: 比較的低いスケーリング（8〜9 倍）。Hutchinson 法により再実行オーバーヘッドが削減され、サンプリングや大規模な GEMM/Attention カーネルの割合が増加します。
• メモリ使用量: 実装スタック（メモリアロケータの挙動）に強く依存し、特に ONEQMC 系では分子サイズに応じたステップ的なメモリ増加が見られました。

B. カーネル特性とボトルネック

• 低算術強度の支配: どのモデルにおいても、融合された要素ごとの演算（fused elementwise）やデータ移動（layout/transpose）カーネルが実行時間の大部分を占め、算術強度が極めて低い（$10^{-2} \sim 10^{-1}$ FLOP/Byte）ことが判明しました。
• メモリバウンド: ルーフライン分析により、これらの低強度カーネルがメモリ帯域幅に制限されていることが示されました。GEMM カーネルが存在しても、エンドツーエンドのパフォーマンスはデータ移動に制約されます。
• モデルごとの違い:
  • PauliNet: ステージ E（微分評価）のリプレイにより、要素ごとの演算カーネルが支配的（52%）。
  • FermiNet: GEMM の割合が増加（30%）するが、依然としてメモリ敏感。
  • Psiformer: サンプリングフェーズで GEMM が支配的（62%）になり、計算集約度が高まるが、全体としては依然として不均一。
  • Orbformer: FlashAttention の採用により GEMM 割合が低下し、要素ごとの演算やデータ移動の割合が増加。メモリバウンド傾向が強まります。

C. ハードウェア利用率

• 多くのモデルで、ピーク命令スループットは 20〜40% 程度に留まり、L2 キャッシュヒット率は 50〜60% 程度です。これは、計算リソースが十分に活用されていないことを示唆しています。

5. 意義と将来の方向性 (Significance & Co-design Implications)

本研究は、NNVMC の加速には単一のカーネル最適化（例：Attention のみ）では不十分であり、フェーズ依存型かつメモリ中心の共設計が必要であることを示しました。

• PIM（メモリ内処理）の活用: 低算術強度の要素ごとの演算やデータ移動カーネルは、PIM によってメモリアクセスを削減することで大幅な性能向上が期待できます。
• GPU-PIM 協調システム: GEMM 集約フェーズは GPU で、メモリバウンドな微細カーネル集まりは PIM で処理するなど、フェーズに応じた動的なオフロードが有効です。
• 再構成可能アクセラレータ: 計算フェーズ（GEMM 重視）とメモリフェーズ（帯域幅重視）でアーキテクチャのバランスを動的に変化させることが重要です。
• Attention 以外の最適化: Orbformer のように Attention カーネルが実行時間の一部しか占めない場合、Attention 最適化だけでは全体性能は向上しません。要素ごとの演算やメモリアクセスパターンの最適化が不可欠です。
• メモリオフロード: 大規模分子系では GPU メモリ不足が発生するため、CPU DRAM や SSD への非同期オフロード戦略の検討が必要です。

結論

本論文は、NNVMC が単なる「大きな行列計算」ではなく、物理法則に制約された多段階の異種ワークロードであることを実証しました。将来的なスケーラブルな NNVMC システムの実現には、モデルのフェーズ特性を理解し、メモリ帯域幅の効率化とフェーズに応じたハードウェアリソースの割り当てを統合的に設計することが不可欠であると結論付けています。