技術要約：中性子反射率測定におけるリアルタイムなサロゲートフリー・ベイズ逆問題解析に向けて

問題提起

中性子反射率測定（NR）は、ナノスケールでの表面および界面の特性評価において極めて重要な技術である。しかし、NRは間接的な測定手法であり、物理的特性（層の厚さ、散乱長さ密度、粗さ）を抽出するには逆問題を解く必要がある。現在のアプローチは、特に不確実性定量化（UQ）と計算効率に関して、以下のような重大な限界に直面している。

• 複雑な系における非効率性： マルコフ連鎖モンテカルロ法（MCMC）のような従来のサンプリング手法は、未知パラメータが増加する場合や、複雑な多層構造（例：有機ELデバイス）を扱う場合、サンプリング効率が悪く、収束時間が長くなるという課題がある。
• 物理的直観の喪失： 近年提案されているサロゲート機械学習モデルは、ニューラルネットワークを用いてデータからパラメータを直接予測することで、この逆問題を回避しようとするものである。これらは高速ではあるが、支配的な物理方程式をブラックボックス近似に置き換えてしまうため、必然的に物理的直観を放棄することになり、また大規模で代表的な訓練データセットを必要とする。
• 勾配利用の欠如： 自動微分（AD）や勾配ベースの最適化が機械学習の他の分野で革命をもたらしているにもかかわらず、NRソフトウェアは依然として勾配フリーの手法（例：レーベンバーグ・マーカート法、差分進化法）やヒューリスティックな手法に依存しており、フォワードモデルの正確な勾配を活用できていない。

手法

著者らは、自動微分（AD）を利用して、未知パラメータに関するフォワード・スペキュラ反射率モデルの正確な勾配を計算する、サロゲートフリー・ベイズ逆問題解析フレームワークを提案している。これにより、これまでNRでは活用されてこなかった高性能な勾配ベースの推論スキームの適用が可能となる。

コア構成要素

1. フォワードモデルと自動微分：

   • フォワード反射率は、薄膜の標準的な物理モデルであるアベレス行列形式を用いて計算される。
   • 著者らは、このモデルをニューラルネットワークで近似する代わりに、JITコンパイル、並列化、および自動微分をサポートするPythonライブラリであるJAXを用いて計算を実装している。
   • これにより、予測反射率 $\hat{R}(Q, \theta)$ のパラメータ $\theta$（スラブおよび計器パラメータ）に対する勾配 $\nabla_\theta \hat{R}(Q, \theta)$ を正確に計算することが可能となる。

2. 勾配ベースの最適化：

   • 本フレームワークでは、誤差指標（例：$\chi^2$）を最小化するために、確率的勾配降下法（具体的にはADAMオプティマイザ）を採用している。
   • この手法は、高次元空間において最適パラメータを見つけるために非常に効果的であることが示されており、従来の勾配フリー手法よりも効果的に局所解を回避できる。

3. 推論のためのハミルトニアン・モンテカルロ法（HMC）：

   • 完全なベイズ推論を行うために、著者らはHMC（具体的にはNo-U-Turn Sampler, NUTS）を利用している。
   • HMCは、対数事後分布の勾配がサンプリングの軌道を導く擬似的な動的システムを構築する。
   • フォワードモデルからの正確な勾配を活用することで、HMCは勾配フリーのMCMCと比較して、大幅に高いサンプル効率（有効サンプルサイズ：Effective Sample Size）を実現し、収束に必要な反復回数を大幅に削減できる。

4. 近似的UQのための変分推論（VI）：

   • 大規模なデータセットや複雑なモデルを必要とするシナリオでは、**変分推論（VI）**を採用している。
   • VIは、真の事後分布 $p(\theta|R)$ をパラメータ化された分布 $q(\theta; \phi)$ で近似することにより、推論問題を最適化問題へと変換する。
   • パラメータ $\phi$ は、フォワードモデルの勾配を用いて、エビデンス下界（ELBO）を最大化するように最適化される。これにより、数時間ではなく数秒単位での近似的なUQが可能となる。

ソフトウェア実装

著者らは、JAXエコシステム上に構築されたオープンソースのPythonライブラリであるrefjaxを公開した。これには以下の機能が含まれる。

• フォワード反射率計算（スムアリングカーネルを含む）の標準化されたインターフェース。
• 自動微分のネイティブサポート。これにより、ユーザーは任意のJAX互換オプティマイザやサンプラー（例：ADAM、HMC、VI）をNRデータに直接適用できる。

主な結果

本手法は、2つのケーススタディを通じて検証された。

1. 結晶質石英膜（ベンチマーク）：

   • 最適化： ADAMを使用することで、すべてのパラメータを自由にした既存の報告値（1.32）をわずかに上回る、1.302の $\chi^2$ 誤差を達成した。
   • 推論（HMC）： NUTSサンプラーは、1チェーンあたり2,000サンプルという効率的な設定で収束した。事後分布の平均は $\chi^2 = 1.22$ を達成した。
   • 効率の比較： 勾配フリーのSample-Adaptive MCMC（SA-MCMC）と比較して、HMCは優れた収束性を示した。HMCの有効サンプルサイズ（ESS）は単調に増加したが、SA-MCMCは初期探索後に混合（mixing）の低下とESSの減少を示した。HMCは、SA-MCMCよりもはるかに速く「良好に収束した」閾値（ESS $\approx$ 200）に到達した。

2. 有機EL（OLED）多層デバイス（高複雑性）：

   • シナリオ： 4種類の異なるOLEDデバイス（アニール温度が異なる）を含む、59個の未知パラメータ（層の厚さ、SLD、粗さ、および共有計器パラメータ）の同時フィット。
   • 変分推論： VIアプローチは、結合事後分布の推定に成功した。
   • パフォーマンス： 各推論実行は、標準的なハードウェア上で20秒未満で完了した。
   • 物理的洞察： モデルは、アニールされたサンプルにおけるTFB層のSLD増加（溶媒の除去と密度の平衡化を示す）や、界面幅の変化といった物理的傾向を正確に特定した。また、推論されたシリコン裏面のSLD（2.09 $\pm$ 0.0178）は文献値（2.07）と密接に一致しており、解の物理的な妥当性を検証した。

意義と主張

本論文は、物理的な忠実度を損なうことなく、NR推論の速度と効率における「ステップチェンジ（飛躍的進歩）」を主張している。

• サロゲートフリーの正確性： MLサロゲートとは異なり、本手法は正確な支配方程式を使用するため、物理的直観を保持し、大規模なラベル付き訓練データセットを必要としない。
• 計算効率： 自動微分（AD）の統合により、以下が可能となった。
  • HMCによる、従来のMCMCと比較したサンプル効率の桁違いの向上。
  • 変分推論による迅速な近似的UQ。これにより、推論時間を数時間から数秒へと短縮した。
• スケーラビリティ： 本フレームワークは、高次元のパラメータ空間や、従来の計算手法では困難であった複雑なマルチデバイスの同時フィットを扱えるよう設計されている。
• 将来への展望： 著者らは、オープンソースライブラリ（refjax）を提供することで、NRと現代的な機械学習・最適化技術との間の溝を埋め、高スループットかつ動的な実験におけるリアルタイム解析を促進することを目指している。

著者らは、本研究を、サロゲートフリーのリアルタイム・ベイズ逆問題解析を可能にするための基礎的な一歩として位置づけており、これにより、NRのデータ解析が高スループット化や実験の進展に即応できることを保証している。