Machine Learning for neutron source distributions

本論文は、モンテカルロ粒子リストから中性子源分布を推定するために確率的生成モデルを用いた新規手法を提案・評価し、モデルの学習後に効率的かつメモリに依存しないサンプリングを可能にするものである。

要約

複雑で多層構造のケーキの完璧なレシピを再現しようとしていると想像してください。中性子科学の世界において、この「ケーキ」とは、源から放出される中性子（微小粒子）の流れであり、それぞれが独自の速度、方向、エネルギー、タイミングを持っています。

従来、科学者たちはこの流れを再現するために、主に 2 つの方法を試みてきました：

1. 「コピー＆ペースト」方式：巨大で低速なコンピュータシミュレーションを実行し、個々の中性子すべてを含む膨大なリストを生成します。このリスト（MCPL ファイルと呼ばれる）を保存し、繰り返し使用しようとします。問題点は、リストにある以上の中性子が必要になった場合、同じものを繰り返しコピー＆ペーストしてしまうことです。これにより、シミュレーション内に「不具合」や「ホットスポット」が生じ、まるで同じパン屑のパターンが無限に繰り返されて見えるようになります。
2. 「経験則」方式：成分を個別に観察して（例えば、「速い中性子はどれくらいあるか？」「遅い中性子はどれくらいあるか？」）、レシピを推測しようとします。問題点は、成分がどのように混ざり合うかを無視してしまうことです。実際には、速い中性子は常に特定の方向に移動している可能性がありますが、この方法はそれらを無関係であるかのように扱うため、実データの「風味」を失ってしまいます。

新しいアプローチ：「AI シェフ」
この論文は、機械学習を用いてこの問題を解決する新しい方法を紹介します。リストをコピーしたりルールを推測したりする代わりに、著者たちは 4 種類の異なる「AI シェフ」（生成モデル）を訓練し、中性子レシピの本質を学ばせました。

以下に、論文がどのように分解しているかを示します：

1. 訓練フェーズ（レシピの学習）

AI シェフたちは、元の低速なコンピュータシミュレーションのサンプル（「訓練データ」）を与えられます。彼らは単にリストを暗記するのではなく、すべての変数間の複雑な関係を学習します。

• 比喩：シェフに特定の種類の雲の写真を 1,000 枚見せると想像してください。彼らは写真を単に暗記するのではなく、その雲を「その雲」たらしめるものを学びます。つまり、縁の巻き方、密度、光の当たり方などです。これを学習すれば、彼らは一度も存在したことのない新しい雲を描くことができますが、それは完璧に正しい見た目になります。

2. 4 人の AI シェフ

著者たちは、どの AI モデルがレシピを最もよく学習したかを確認するために、4 種類の異なる AI モデルをテストしました：

• Normalizing Flows（NF）：これは、生地を完璧に引き伸ばしたり絞り込んだりするシェフだと考えてください。彼らは単純で均一な生地の玉（ランダムノイズ）から始め、それを中性子雲の正確で複雑な形状に引き伸ばします。論文によると、これが最高のシェフであり、元のデータと完璧に一致する最も正確な「新しい」中性子を作成しました。
• Variational Autoencoders（VAE）：このシェフは、レシピを要約に圧縮し、それを再構築しようとします。高速で複雑な形状には優れていますが、再構築されたケーキは元のものよりも少し「ぼやけて」いたり、鮮明さが欠けていたりすることがあります。
• Generative Adversarial Networks（GAN）：これは 2 人のシェフによる「綱引き」です。一人は偽物のケーキを焼き、もう一人は偽物を見つけようとします。偽物のケーキが本物と区別できなくなるまで、彼らは競争を続けます。この論文では、これらは訓練が少し難しく、同じいくつかのパターンを繰り返す「不正」を起こしやすいことがわかりました。
• Diffusion Models（DM）：このシェフは、ノイズの多い乱れたケーキから始め、完璧になるまで段階的にゆっくりと掃除していきます。これはうまく機能しますが、一粒の塵ずつ拾って部屋を掃除しようとするようなもので、非常に遅く、計算コストがかかります。

3. 結果：なぜ重要なのか

この論文は、これらの AI シェフを 2 つの現実世界のシナリオでテストしました：

• シナリオ A（TDR データセット）：複雑で高エネルギーの中性子源です。AI シェフたちはレシピを非常にうまく学習し、「コピー＆ペースト」による不具合なしに、元のシミュレーションと統計的に同一に見える数百万の新しい中性子を生成することができました。
• シナリオ B（ベンチマークデータセット）：実際のラボで測定されたデータと、AI が生成した中性子を比較した現実の実験です。AI（特に Normalizing Flow）は、実世界のデータとほぼ完璧に一致しました。

主な利点：
AI シェフがレシピを学習すれば、元の中性子の巨大で重いリストはもはや必要ありません。AI モデルは小さく（数キロバイト程度）、統計的に完璧な無制限の新しい中性子を瞬時に生成できます。これにより、莫大なコンピュータ時間とメモリが節約されます。

この論文が述べていないこと

著者たちは慎重に、これらのモデルがデータ駆動型であることを述べています。彼らは与えられたデータから厳密に学習します。

• 元のシミュレーションに特定の種類の中性子が欠けていた場合、AI はそれを発明しません（データの外を推測するようにモデルが特別に調整されている場合を除きますが、論文はこれが他の方法の特定の機能であり、ここでの主要な目的ではないと指摘しています）。
• この論文は、これらのモデルが新しい物理を予測したり、不良データを修正したりできるとは主張していません。これらは、中性子機器の設計に使用するために、既存のデータパターンを効率的に再現するためのツールです。

まとめ：
この論文は、不具合を起こしやすい中性子データの巨大なリストを、小さく賢い AI モデルに置き換えることができることを実証しています。これらのモデルは中性子流の「DNA」を学習し、必要に応じて新鮮で現実的な中性子を生成でき、これにより将来の中性子実験の設計をより速く、安価に、かつ正確に行うことを可能にします。テストされた 4 つのモデルの中で、Normalizing Flowが明確な勝者でした。

技術概要

技術概要：中性子源分布のための機械学習

問題定義
中性子光学および装置設計において、特定の表面における中性子の位相空間分布（位置、方向、時間、エネルギー、重み、偏極）の正確な特性評価は極めて重要です。従来、これらの分布はモンテカルロ（MC）シミュレーション（例：MCNP、PHITS、OpenMC）から導出され、モンテカルロ粒子リスト（MCPL）ファイルとして保存されます。MCPL ファイルはレイ・トレーシングソフトウェア（例：Vitess、McStas）の標準入力として機能しますが、その直接再利用には重大な限界があります：

1. 統計的アーティファクト： 元の MCPL ファイルにおけるサンプリング不足と統計的ノイズが下流のシミュレーションに伝播し、人工的な構造（例：「ホットスポット」）や収束の悪化を引き起こします。
2. 相関の喪失： 個々の変数に解析的または半解析的モデルを適合させる代替アプローチは、位相空間変数間の複雑な高次元相関を保持できず、特に高輝度源における現実的な輸送に不可欠な相関を失う傾向があります。
3. 計算コスト： 異なる装置構成に対して新しい源リストを生成するために高統計量の MC シミュレーションを再実行することは、時間がかかり、計算コストが高額です。

既存の機械学習ソリューション、例えばカーネル密度推定（KDSource）は、いくつかの問題に対処しますが、重要度重みの手動調整を必要とし、特に粒子重みが変化する際に完全な多変量位相空間推定に苦労することが多いです。

手法
著者は、モンテカルロ粒子リスト（MCPL）ファイルから直接中性子位相空間変数の同時確率分布を学習するための確率的生成モデル（PGM）を活用するフレームワークを提案します。一度訓練されれば、これらのモデルは元の粒子リストを必要とせずに、元の分布の相関を保持した新しい統計的に独立した中性子サンプルを生成できます。

• データセット： モデルの評価には 2 つの異なるデータセットが使用されました：

  1. TDR データセット： 高輝度源（HBS）パラ水素減速材の PHITS シミュレーション（676,558 個の中性子）。このデータセットには粒子重みを含み、広範なエネルギー範囲（コールド、熱、高速）をカバーしており、複雑な高次元（7 変数）のモデリング課題を提示します。
  2. ベンチマークデータセット： 2023 年の実験測定値と比較検証された、JULIC 中性子プラットフォームの PHITS シミュレーション（118,219 個の中性子）。このデータセットは一定の粒子重みを持ち（次元が 6 変数に削減）、実験スペクトルに対する直接検証を可能にします。
  3. 前処理： すべての変数は MinMax スケーリングを用いて正規化されました。時間とエネルギー変数には、その広いダイナミックレンジを処理するために対数変換が適用されました。

• 生成モデル： 4 つの異なる PGM 構造が比較されました：

  1. 正規化フロー（NF）： 特にカップリングフローとマスクド自己回帰フロー（MAF）です。これらのモデルは、単純な参照分布（ガウス分布）と複雑なデータ分布間の可逆変換を学習し、正確な尤度評価とサンプリングを可能にします。
  2. 変分オートエンコーダ（VAE）： 再構成損失とカルバック・ライブラー（KL）ダイバージェンス正則化項を通じて最適化される、データの潜在表現を学習するエンコーダ・デコーダ構造です。
  3. 生成敵対ネットワーク（GAN）： 実データと区別できないサンプルを生成するように、識別器ネットワークに対して訓練される生成器ネットワークです。
  4. 拡散モデル（DM）： 前方のノイズ付加プロセスを逆転させることを学習し、ランダムノイズからデータを再構築するモデルです。

• 実装： モデルは JUWELS ブースタークラスター上で PyTorch を使用して訓練されました。Vitess MC ソフトウェア用のカスタム C++ モジュール source AI が開発され、訓練された PyTorch モデル（TorchScript 経由）をロードし、オンザフライサンプリングを実行して、MCPL ファイル入力を事実上置き換えるように設計されました。

• 評価指標：

  • 最大平均不一致（MMD）： 再生核ヒルベルト空間（RKHS）において、学習された分布と元の MCPL 分布との距離を定量化するために使用されました。
  • サンプリング速度： 20,000 個の中性子を生成するために必要な時間。
  • 物理的検証： JULIC プラットフォームからの実験データとの比較による、シミュレートされた波長スペクトルの検証。

主要な結果

• 分布忠実度： すべてのモデルが基礎となる分布を正常に学習しましたが、正規化フロー（NF）が優れた性能を示しました。NF は、両方のデータセットで最低の MMD スコアを達成しました（例：TDR で $1.19 \times 10^{-4}$、ベンチマークで $1.20 \times 10^{-4}$）、VAE、GAN、拡散モデル、および KDSource ベースラインを上回りました。
• 相関保持： 2 次元ヒストグラムの視覚的分析により、NF が単純な適合手法で失われた変数間の複雑な相関（例：時間 - エネルギー、角度 - 位置）を正常に保持したことが確認されました。
• サンプリング効率： KDSource が最も高速なサンプリング時間を提供しましたが、NF は速度と精度の間の有利なバランスを提供しました。VAE と GAN も高速でしたが、反復的なノイズ除去の性質により、拡散モデルはサンプリングに著しく時間がかかりました。
• 実験的検証： ベンチマークデータセットにおいて、NF 生成源は実験的波長スペクトルを高い忠実度で再現し、KDSource と同等の性能を発揮しました。しかし、KDSource とは異なり、NF モデルはシグモイド活性化層によりトレーニングデータの範囲に厳密に従うため、明示的に設計されていない限り、物理的に非現実的な外挿（例：ガイド外にある中性子や観測されていない波長）の生成を防ぎます。
• 頑健性： 同じモデル構造は、2 つの異なるデータセットに適用された際、最小限の再調整しか必要とせず、異なる次元や分布の複雑さ全体にわたって頑健であることを示しました。

意義と主張
本論文は、現代の生成 ML モデルを中性子モンテカルロパイプラインに統合し、実験データに対して検証する初の比較研究であると主張しています。主な意義は、以下の点を示すことにあります：

1. PGM は代替可能である： 確率的生成モデルは、生の MCPL ファイルを置き換えることができ、メモリ効率（ギガバイトサイズのファイルに対しキロバイトサイズのモデル）に優れ、統計的に頑健な源表現手法を提供します。
2. 正規化フローが最適である： 試験された構造の中で、NF はこの特定のアプリケーションにとって最適なトレードオフを提供し、高い忠実度の分布学習と、ハードな物理的制約（例：有界な位相空間）をアーキテクチャに直接組み込む能力を提供します。
3. ワークフロー統合： Vitess 内で開発された source AI モジュールは、これらのモデルが既存の中性子シミュレーションワークフローにシームレスに統合可能であり、元のトレーニングデータを必要とせずに迅速かつ独立したサンプリングを可能にすることを証明しています。

著者は、生成モデルが初期の高忠実度 MC シミュレーションの必要性を置き換えるものではありませんが、直接的な MCPL 再利用が非現実的な、反復的なシミュレーションや高統計量サンプリングを必要とするシナリオにおいて、下流の装置設計と最適化のための効率的で柔軟かつ物理的に意味のあるツールを提供すると結論付けています。