Neural Quantum States in Mixed Precision

技術要約：混合精度におけるニューラル量子状態

問題提起
科学計算は、特に現実世界の現象のシミュレーションにおいて、精度を確保するために伝統的に倍精度（64ビット）演算に依存してきました。しかし、GPUのようなハードウェアアクセラレータの普及に伴い、優れたパフォーマンス、メモリ使用量の削減、およびエネルギー効率の向上を実現する低精度フォーマット（単精度 f32、半精度 f16、bfloat16 など）が魅力的になっています。現代の科学計算、特に量子多体系を解くための変分モンテカルロ（VMC）法における重要なボトルネックは、十分に非相関なマルコフ連鎖モンテカルロ（MCMC）サンプルを生成するための計算コストです。低精度演算は高速化をもたらす一方で、MCMCサンプリングの精度（具体的には、バイアスの導入やエルゴード性の破壊の可能性）に与える影響については、まだ十分に研究されていません。本研究は、MCMC駆動型の科学的アプリケーションにおける混合精度演算の適用可能性を厳密に評価する必要性に対処するものです。特に、ニューラル量子状態（NQS）に焦点を当てています。

手法
著者らは、理論的導出と経験的検証の組み合わせを通じてこの問題にアプローチしています。

1. 理論的枠組み:

   • 本研究では、低精度をターゲット分布 $\pi(x)$ の対数確率に対する加法的摂動 $\delta(x)$ としてモデル化し、摂動を受けた分布 $\tilde{\pi}(x)$ を定義します。
   • 著者らは、この摂動によってメトロポリス・ヘイスティングス（MH）MCMCに導入される誤差に関する解析的な境界を導出しています。正確な分布と摂動を受けた分布の差を定量化するために、全変動（Total Variation; TV）距離を利用しています。
   • マルコフ連鎖の混合特性（具体的には、収縮率 $r$ を確立するためのドブリン（Doeblin）マイノリゼーション条件）を組み込むことで、標準的な粗い境界よりもタイトなバイアスの境界を導出しています。
   • 定理 III.3 は、ガウスノイズを用いた局所的な MH プロポーザルに関する具体的な境界を提供し、バイアスがノイズレベル $\sigma$ と収縮率 $r$ に比例することを示しています。決定的なのは、連鎖が急速に混合する場合（$r \ll 1$）、バイアスは $O(\sigma)$ ではなく $O(\sigma^2)$ でスケールすることです。

2. 経験的検証:

   • トイモデル: 著者らは、精度誤差をシミュレートするために対数密度にガウスノイズを明示的に注入した「ノイジー」な制限ボルツマンマシン（RBM）を用いて、理論的境界を検証しています。
   • 混合精度 VMC フレームワーク: 著者らは、VMC のサンプリングコンポーネントを低精度（f16 または bf16）で実行し、勾配計算とパラメータ更新（フォワード／バックワードパス）は倍精度（f64）で維持するという実用的な混合精度スキームを実装しています。パラメータは高精度で更新され、その後、サンプリングステップのためにダウンキャストされます。
   • ベンチマーク: 本手法は、反強磁性相および常磁性相の両方の横磁場イジングモデル（TFIM）、およびハイゼンベルクモデルにおいてテストされています。アーキテクチャには、RBM および 残差畳み込みニューラルネットワーク（ResCNN）が使用されています。

主な貢献

1. 誤差境界の導出: 本論文は、離散 MH MCMC において低精度演算によって導入されるサンプリング誤差に関する一般的な解析的境界を導出しています。これは、ノイズレベル、連鎖の混合時間、および結果として生じる定常分布のバイアスの間の理論的なつながりを確立しています。
2. 混合精度 VMC フレームワーク: 著者らは、勾配更新の精度要件を維持しつつ、MCMC サンプリングを低精度で行うフレームワークを導入しています。これにより、サンプリングのボトルネックと勾配更新の精度要件を効果的に分離しています。
3. スケーラビリティの経験的実証: 本研究は、NQS の最適化において、学習性能や基底状態近似の精度を損なうことなく、大幅な高速化（最大 3.5 倍）が可能であることを示しています。

結果

• 理論的一致: ノイジーな RBM トイモデルによる経験的結果は、導出された理論的境界と非常によく一致しています。相対期待値におけるモンテカルロサンプリング誤差は、ノイズパラメータ $\sigma$ が 1 のオーダーになるまで、摂動に対して鈍感なままです。
• 精度の保持: 2D TFIM およびハイゼンベルクモデルにおける完全な VMC 最適化において、倍精度参照解に対する相対エネルギー誤差は、f64、f32、f16、および bf16 のデータ型間でほぼ同一です。
• ノイズ分析: 低精度評価における数値誤差（$\delta$）の分布は、近似的にガウス分布に従うことが判明し、理論的仮定を裏付けました。この誤差の標準偏差は一貫して小さく（$\sigma < 1$）、MCMC 連鎖が信頼性高く収束することを保証しました。
• 高速化の傾向: 高速化係数は、サンプル数とモデルパラメータ数の両方とともに系統的に増加します。これは、GPU メモリ利用率の向上および算術ユニットの飽和に起因しています。
• 勾配ノイズ: 本研究は、f32 および f16 レジームにおいて、確率的なモンテカルロサンプリングノイズが低精度数値誤差よりも支配的であることを確認しており、これが最適化ダイナミクスが安定している理由を説明しています。

意義と主張
本論文は、混合精度演算が、特にニューラル量子状態のための MCMC サンプリングを加速するための、実行可能かつ効果的な戦略であることを主張しています。著者らは以下の通り述べています。

• 量子状態のサンプリングを含む VMC アルゴリズムの大部分は、精度を損なうことなく半精度で実行できます。
• 導出された理論的枠組みは、MCMC に依存する機械学習の手法（NQS を超えて、ベイズ学習やエネルギーベースモデルにも適用可能）における混合精度の適用可能性を評価するための厳密な方法を提供します。
• 一般的な機械学習では混合精度は標準的ですが、その MCMC 駆動型の科学計算への適用は十分に探索されていません。本研究は、これらの技術をよりスケーラブルでエネルギー効率の高い量子多体系のシミュレーションに採用するために必要な、理論的正当性と実践的な証拠を提供します。
• 著者らは、サンプリングの高速化は顕著であるものの、局所エネルギー計算や前処理ステップに混合精度を適用することで、さらなる利得が得られる可能性があると指摘していますが、これには追加の実践的な課題が伴います。

本研究は、基礎となるマルコフ連鎖の混合特性が導入された数値ノイズを処理するのに十分な堅牢性を備えている限り、混合精度は GPU 依存の科学的ワークフローの未来にとってインパクトの大きいツールであると結論付けています。