Probabilistic Computers for Neural Quantum States

本論文は、スパースボルツマンマシンアーキテクチャと確率計算ハードウェア（FPGA）を組み合わせることで、ニューラル量子状態におけるモンテカルロサンプリングのボトルネックを克服し、最大 6400 スピンの 2 次元横磁場イジングモデルの正確な基底状態エネルギー計算と、900 スピンの深層モデルの効率的な学習を可能にすることを示す。

要約

大勢の人々の集団の挙動を予測しようとする様子を想像してみてください。そこでは、一人ひとりが隣人に対して複雑で目に見えない方法で絶えず反応しています。物理学の世界では、これを「量子多体系」と呼びます。これを通常のコンピュータでシミュレーションしようとするのは、風が砂を吹き飛ばしている間に砂浜の砂粒をすべて数えようとするようなもので、非常に遅く、大規模な集団の場合にはしばしば不可能です。

本論文は、この問題を解決するための新たな手法として、高度なソフトウェアと専用ハードウェアを組み合わせるアプローチを紹介しています。以下に、彼らのアプローチをシンプルなアナロジーを用いて解説します。

1. 課題：シミュレーションの「交通渋滞」

科学者たちは、これらの量子集団をモデル化するために「ニューラル量子状態（NQS）」と呼ばれる手法を用いています。ニューラルネットワークは、集団がどのように振る舞うかを予測する非常に賢い地図のようなものと考えることができます。しかし、この地図を更新するためには、コンピュータは「全員が一歩左に動いたらどうなるか？」といった数百万回ものランダムなシミュレーションを実行して、その結果を確認する必要があります。

標準的なコンピュータ（CPU）では、このサンプリング過程が巨大な交通渋滞となります。コンピュータはこれらのランダムなシナリオを生成することにあまりにも多くの時間を費やしてしまい、実際には答えを学習することができません。これが著者たちが解決しようとした「ボトルネック」です。

2. 解決策：専用「確率的」エンジン

汎用コンピュータにランダム性をシミュレーションさせる代わりに、著者たちはFPGA（再プログラム可能で専用ハードウェアのように動作するチップ）を用いたカスタムマシンを構築しました。

• アナロジー: 標準的なコンピュータは、手作業で図書館を整理しようとする一人の非常に賢い司書のようなものです。正確ですが、遅いです。一方、著者たちの確率コンピュータは、2,200 人の小さく速い労働者（p-bitsと呼ばれる）を雇い、全員が同時に本を整理させるようなものです。
• 仕組み: これらの p-bits は、隣接する要素に基づいて 2 つの状態（コインが表か裏かに落ちるような）の間を行き来する単純な単位です。これらはハードウェアに直接組み込まれているため、「ランダムであること」を考えなくてもよく、本質的にランダムです。これにより、シミュレーションに必要な数百万のシナリオをほぼ瞬時に生成することができます。

3. 最初のブレークスルー：巨大な集団のシミュレーション

チームはこの新しいハードウェアを用いて、量子スピン（小さな磁石の網の目）の 2 次元格子をシミュレーションしました。

• 結果: 彼らは80 行×80 列（スピン 6,400 個）の格子のシミュレーションに成功しました。
• 重要性: 従来の手法では、この規模に達しようとするとクラッシュするか、時間がかかりすぎてしまうという課題がありました。彼らの専用ハードウェアにより、高い精度でこの規模に到達することができ、標準的なコンピュータでは扱えないほど大規模な量子シミュレーションを専用「確率的」チップが処理できることを実証しました。

4. 2 番目のブレークスルー：「ディープ」ラーニングのトリック

著者たちはまた、複雑なパターンを理解するのに優れているため、「より深い」ニューラルネットワーク（論理の層をさらに積み重ねたもの）を使用したいと考えていました。しかし、深いネットワーク通常は「周辺化」と呼ばれる数学的ステップを必要とします。これは、一人ひとりの個人を個別に測定して集団の平均身長を計算しようとするようなもので、深いネットワークにとっては計算上不可能です。

• イノベーション: 彼らは**「デュアル・サンプリング・アルゴリズム」**を発明しました。
• アナロジー: 集団全体を一度に測定しようとする代わりに、外側の人々（可視層）を固定し、中央の人々（隠れ層）だけを整理させます。この「条件付きサンプリング」を行うことで、不可能な数学計算を行わずに答えを導き出すことができます。
• 結果: 彼らは、30 行×30 列（スピン 900 個）のシステムに対して、単一の FPGA チップ上でこれらの深いネットワークの学習に成功しました。その結果、これらの深いネットワークは実際にはより効率的であり、単純で浅いネットワークと同じ精度の結果を得るために必要な「設定」（パラメータ）の数が少ないことがわかりました。

まとめ

要約すると、本論文は主に 2 つのことを主張しています。

1. ハードウェアの高速化: ランダムなコイン投げの巨大な軍隊のように動作するカスタムチップ（FPGA）を構築することで、量子シミュレーションの規模拡大を妨げていた速度制限を取り除きました。彼らは 6,400 個の粒子のシステムをシミュレーションし、これはこの種の手法では以前は到達不可能だった規模です。
2. より賢いアルゴリズム: 彼らは、不可能な数学計算を回避する量子物理学向けの「ディープ」ニューラルネットワークの新しい学習手法を開発しました。これにより、より強力かつ効率的なモデルが可能になりました。

著者たちは、この専用ハードウェアと新しいアルゴリズムを組み合わせることで、これまでにないほど大規模で複雑な量子システムのシミュレーションが可能になり、以前は研究が難しすぎた物質や物理学の理解への扉が開かれると結論付けています。

技術概要

技術的概要：ニューラル量子状態のための確率論的コンピュータ

1. 問題定義

量子多体系の正確な古典シミュレーションは、凝縮系物理学および量子化学における根本的な課題である。量子モンテカルロ（QMC）やテンソルネットワークといった確立された手法は高精度を達成しているが、本質的な限界に直面している。すなわち、QMC は一般的な系において符号問題に悩まされ、テンソルネットワークは二次元および臨界点近傍における不利なエンタングルメントのスケーリングに苦しむ。

ニューラル量子状態（NQS）は、ニューラルネットワークを用いて多体波動関数をパラメータ化することで、スケーラブルな代替手段を提供する。しかし、NQS の変分モンテカルロ（VMC）トレーニングは、マルコフ連鎖モンテカルロ（MCMC）サンプリングの計算コストによってボトルネックとなっている。系サイズが増大するにつれ、サンプリングを通じて観測量や確率的パラメータ勾配を推定するために必要な時間が、制限付きボルツマンマシン（RBM）のような比較的単純なアーキテクチャであっても、許容し得ないものとなる。このボトルネックは、複雑な量子相を探索するために必要な大規模な系サイズ（例えば、$>10^3$ スピン）へのスケーリングを阻害している。

2. 手法

著者らは、スパースなボルツマンマシンアーキテクチャを確率論的コンピューティングハードウェアに直接マッピングすることにより、サンプリングのボトルネックを克服するハードウェアとソフトウェアの共設計アプローチを提案する。

A. 確率論的ハードウェアアーキテクチャ

この手法の中核は、フィールドプログラマブルゲートアレイ（FPGA）を用いた確率論的コンピュータ（p-computer）の実装である。

• P-bits: ハードウェアは、論理状態 $\{-1, +1\}$ 間で揺らぐ古典的確率単位である確率ビット（p-bits）を利用する。これらの単位は、サンプリングに必要なボルツマン分布を自然に実装する。
• スパース接続（FRBM）: 密結合ネットワークのルーティング混雑と $O(N^2)$ の配線複雑さを回避するため、著者らは**さらに制限されたボルツマンマシン（FRBM）**を採用する。このアーキテクチャは厳密な局所接続（ユークリッド距離 $k=2$、スピンあたり 13 個の隣接スピンに対応）を強制し、配線複雑さを $O(N)$ に低減する。
• ハイブリッド実行: ホスト CPU がパラメータ最適化（確率的再構成を使用）を処理する一方、FPGA は高スループットなサンプリャとして機能する。FPGA は並列 p-bit 更新を通じてスピン構成を生成し、これらは勾配集積とパラメータ更新のために CPU に転送される。
• 精度: FPGA 実装は、p-bit 密度と並列性を最大化するために 10 ビット固定小数点演算を使用し、ホスト CPU は最適化における数値的安定性のために単精度浮動小数点（FP32）を使用する。

B. 深層モデルのための双対サンプリングアルゴリズム

浅い RBM よりも表現力が高いが、隠れユニット上の非自明な周辺化に苦しむ深層ボルツマンマシン（DBM）のトレーニングを可能にするため、著者らは双対サンプリングアルゴリズムを導入する。

• 概念: 計算的に高価な補助変数上の周辺化の代わりに、このアルゴリズムは条件付きサンプリングに置き換える。
• プロセス:
  1. 外側ループ: 物理層から可視構成（$v$）をサンプリングする。
  2. 内側ループ: 各固定された可視構成に対して、可視ユニットをクランプし、補助（隠れおよび深層）層上でギブスサンプリングを実行する。
  3. 推定: 局所エネルギー計算に必要な波動関数の比率は、固定された可視状態に対する補助変数上の条件付き期待値として推定される。
• 効率性: このアプローチは、物理スピンサンプリングと補助層サンプリングを分離し、分散を低減するとともに、単一スピン反転のたびに再サンプリングする必要性を回避する。これにより、厳密な局所性制約下でのスパースな深層アーキテクチャのトレーニングが可能となる。

C. スケーラビリティ戦略

• マルチ FPGA クラスタリング: 大規模系（例えば、$80 \times 80$ 格子）の場合、FRBM グラフは METIS グラフ分割ツールを用いて複数の FPGA に分割される。境界 p-bit は高速 FMC リンクを介して非同期に交換され、局所 p-bit は同期して更新される。これにより、システムは単一チップのリソースを超えてスケーリング可能となる。

3. 主要な貢献

1. ハードウェア加速サンプリング: 著者らは、スパースなボルツマンマシンをマルチ FPGA クラスタにマッピングすることを示し、CPU および GPU ベースラインと比較して圧倒的なサンプリング速度向上を達成した。
2. 双対サンプリングアルゴリズム: 彼らは、非自明な周辺化を条件付きサンプリングに置き換えることで、変分モンテカルロにおけるスパースな深層ボルツマンマシンのトレーニングを可能にする新規アルゴリズムを導入した。
3. パラメータ効率: 彼らは、スパースな深層アーキテクチャ（DBM）が、同程度の精度を達成するために必要な浅いネットワーク（RBM）と比較して、著しく少ないパラメータ数でより低い変分エネルギーを達成することを示し、パラメータ効率を向上させた。

4. 結果

この手法は、臨界点における二次元横磁場イジングモデル（TFIM）で検証された。

• 単一 FPGA パフォーマンス:

  • $35 \times 35$ 格子（1,225 スピン）において、システムは約 100 回の最適化反復内で化学的精度（相対誤差 $|\Delta E/E_{ref}| \le 1.6 \times 10^{-3}$）に達した。
  • サンプリングは FPGA 上の総ウォールクロック時間の 5% 未満を消費したが、CPU ベースラインははるかに少ないサンプル数であっても、その時間の 20〜30% をサンプリングに費やした。
  • 基底状態エネルギーは、強磁性体と磁場分極の極限の間を滑らかに補間し、連続時間経路積分モンテカルロのベンチマークと一致した。

• マルチ FPGA スケーリング:

  • 相互接続された 6 個の FPGA クラスタを使用することで、著者らは最大**$80 \times 80$（6,400 スピン）**の格子をシミュレートした。
  • 系サイズが増大しても、境界通信のオーバーヘッドを最小化（$L=80$ でカット割合 5.6%）しつつ、システムは化学的精度内で収束を維持した。
  • 非同期通信により、局所 p-bit は 15 MHz にオーバークロック可能となり、厳密なグローバル同期に必要なクロック速度を大幅に上回った。

• 深層モデルトレーニング:

  • $10 \times 10$ 格子において、双対サンプリングアルゴリズムはスパースな DBM のトレーニングに成功し、化学的精度を達成した。
  • パラメータ効率: スパースな DBM は、同程度の精度を達成するために必要なスパースな RBM（$N_p \approx 3100$）と比較して、約半分のパラメータ数（$N_p \approx 1300$）でより低い変分エネルギーを達成した。
  • スケーラビリティ: このアルゴリズムは、単一 FPGA 上で**$30 \times 30$ 格子（900 スピン）**に適用され、これまで深層 NQS では扱いが困難だった系に対する深層モデルのトレーニングの実現可能性を示した。
  • GPU 上でのアルゴリズム的スケーリング分析によると、固定スパース性下では反復時間が線形次元の 2 乗に比例して増加（$t_{iter} \propto L^2$）し、これは総スピン数 $N=L^2$ と整合的であった。

5. 意義と主張

本論文は、確率論的ハードウェアが量子多体系の変分シミュレーションにおけるサンプリングボトルネックを効果的に緩和すると主張する。スパースなボルツマンマシンアーキテクチャと p-bit ハードウェアを組み合わせることで、著者らは以下を実証する。

• スケーラビリティ: 最大 6,400 スピンの量子系をシミュレートする能力は、現在の CPU および GPU ベースの NQS 実装の限界を超えている。
• アーキテクチャの深さ: 双対サンプリングの導入により、より良いパラメータ効率と、浅いネットワークでは表現できない複雑な相関（体積則エンタングルメントなど）を表現する能力を持つ、深層かつスパースなモデルのトレーニングが可能となる。
• 将来の道筋: この研究は、確率論的コンピューティングを、量子物質を古典的にシミュレートするためのスケーラブルな経路として位置づける。著者らは、p-bit アーキテクチャが FPGA プロトタイプから専用 CMOS 回路へと成熟するにつれ、サンプリング、局所エネルギー評価、勾配集積を単一ダイ上でさらに統合することで、レイテンシとエネルギー消費を桁違いに削減でき、今日アクセス可能なものよりもはるかに大規模な量子系に対する VMC を実用的にする可能性があると示唆している。

著者らは非ストークラスティック系については控えめな立場をとっており、非自明な符号構造を持つ系への手法の拡張には、複素パラメータや位相ネットワークが必要となり、それは現在の範囲を超えると指摘している。同様に、サンプリングのボトルネックは解決されたものの、ホストベースの確率的再構成更新により、全体のトレーニングコストは系サイズに対して線形のまま残っており、これが将来のハードウェア加速のターゲットであると特定している。