jFoF: GPU Cluster Finding with Gradient Propagation

技術要約：「jFoF: 勾配伝播を用いたGPUクラスター探索」

問題提起
宇宙論的シミュレーションは、宇宙構造の基本構成要素であるダークマターハローを特定するために、Friends-of-Friends（FoF）アルゴリズムに大きく依存している。効果的な手法ではあるものの、従来のFoF実装は、現代の天体物理学のワークフローにおいて主に2つの制限に直面している：

1. パフォーマンスのボトルネック： 標準的な実装は、多くの場合、CPUベースの空間データ構造（例：$k$-d木）や分散メモリMPIフレームワークに依存している。これらのアプローチは、ホスト・デバイス間のデータ転送による大きなオーバーヘッドを招き、特にエンドツーエンドのGPUネイティブなシミュレーション・パイプライン（例：PMWD、GRAM-X）に統合される際、現代のGPUアーキテクチャの持つ膨大な並列性を活用することに苦慮する。
2. 非微分性： ハロー探索には、閾値設定、近傍選択、グラフ連結性といった離散的な操作が含まれており、これらは本質的に非微分的である。これにより、ハロー特定ステップを通じて勾配を伝播させることが困難となり、観測されたハローカタログから直接、宇宙論モデルの較正や初期条件の再構成を行うための勾配ベースの最適化および推論（ハミルトニアン・モンテカルロ法など）を阻害している。

手法
著者らは、JAXで実装された、完全にGPUネイティブなFriends-of-FriendsハローファインダーであるjFoFを提示している。このコードは完全にデバイス上で動作するように設計されており、コストのかかるCPU-GPU間のデータ転送を排除し、微分可能な天体物理学パイプラインへの統合を可能にしている。

• 近傍探索戦略： jFoFは、GPU実行に最適化された2つの相補的な近傍探索アルゴリズムを実装している：
  1. $k$-d木： 再帰やポインタ間接参照を避けるため、フラットな配列と固定形状を用いてGPU向けに適応させた階層的空間分割手法。jaxkdライブラリを利用しており、純粋なJAX版とCUDAコンパイル版の両方を提供する。
  2. 連結セル格子（Linked-Cell Grid）： シミュレーション体積を均一な立方体セルの格子に離散化する手法。このアプローチは、近傍探索を隣接するセルのローカルなステンシル内に限定することで、計算量を$O(N^2)$から$O(N)$へと低減し、最小像法（minimum-image wrapping）を介して周期境界条件を自然に処理する。
• グループ構築： 両方の手法は、反復的なラベル伝播スキームを採用している。粒子は、自身と連結された近傍との間で最小のラベルを採用し、最終的にハローを定義する連結成分へと収束する。
• 勾配伝播： 離散的なクラスタリングに伴う非微分性の問題に対処するため、jFoFは3つの勾配推定モードを導入している：
  1. 凍結割り当て（Frozen Assignment）： バックプロパゲーションの間、ハローのメンバーシップは固定されたものとして扱われる。勾配は連続変数（例：粒子の位置）を通じてのみ流れ、粒子の座標の解析関数（例：粒子の分散などの「装飾された」プロキシ）として定義されるハローの位置や特性の微分を可能にする。
  2. トポロジー最適化（REINFORCE）： 連結基準を確率的なベルヌーイ過程へと緩和する確率的アプローチ。REINFORCE（スコア関数）エスティメータを用いることで、ハローグラフの離散的なトポロジーを通じて勾配が伝播する。これにより、期待値におけるハローの連結性、メンバーシップ、および数の最適化が可能となり、連続的な場と離散的なカタログとの架け橋となる。

主な結果

• パフォーマンス・ベンチマーク： 著者らは、MPIを用いた並列CPU実装であるcFoFに対してjFoFのベンチマークを行った。粒子数が$64^3$から$512^3$の範囲のシミュレーションにおいて、GPUネイティブなjFoFの実装（$k$-d木および連結セル格子の両方）は、大規模な粒子数において、CPUベースラインと比較して最大で1桁のコア時間での高速化を達成した。
• 精度： jFoFによって生成されたハローカタログは、標準的なcFoF実装とほぼ完全に一致している。ハロー位置の相互相関係数は、すべてのスケールにおいて0.995を超えている。極めて低質量側（$\leq 40$粒子）においてのみ、浮動小数点精度の違いに起因するわずかな不一致が認められた。
• 勾配の検証：
  • 凍結割り当てのテストにおいて、著者らはハローの位置およびプロキシ質量特性（粒子の分散）の勾配を効率的に計算できることを示した。これらの勾配は、固定された初期位相シミュレーションから宇宙論パラメータ（$\sigma_8$）を回収するために、勾配降下法最適化器を効果的に誘導し、数回の関数評価内で真の値へと収束した。
  • トポロジー最適化のテストにおいて、REINFORCEベースのアプローチは、ランダムな初期条件からターゲットとなるハロー場分布を正常に再構成した。最適化器は、ターゲット領域の外側でハローを断片化し、領域内で粒子を結合することを学習し、高い相関係数で大規模構造の形態を回収した。

意義と主張
本論文は、jFoFが微分可能な構造特定の基礎を確立するものであると主張している。GPU加速されたパイプライン内で、構造形成モデルのエンドツーエンドの勾配ベースの最適化を可能にすることで、jFoFは以下を実現する：

• 直接推論： 近似的なバイアスモデルやエミュレータに頼ることなく、ハロー探索プロセスを直接微分することにより、宇宙論パラメータやサブグリッド物理モデルに対する推論を行う能力。
• 統一されたワークフロー： シミュレーション、解析、および推論を単一の微分可能な計算グラフへと統合し、物理パラメータに対する結合制約を容易にする。
• スケーラビリティ： 精度を損なうことなく、現代のアクセラレータ・アーキテクチャに対してFoFハロー探索を再実装できることを実証しており、粒子数$10^9$を超えるシミュレーションの解析への道を開くものである。

著者らは、トポロジー最適化のアプローチに関して控えめなトーンを維持しており、それは、手法が偏りのない勾配推定を提供する一方で、高い分散を持ち、緩和パラメータやサンプリング戦略の慎重な調整を必要とすることを指摘している。彼らは、jFoFをあらゆる動力学的問題に対する最終的な解決策としてではなく、離散的な観測カタログと連続的な物理場との間の溝を埋める概念実証ツールとして位置づけている。