技術的概要：torch-sla

問題定義

微分可能な疎行列線形代数は、グラフニューラルネットワーク（GNN）、非構造化メッシュ上のニューラル演算子、および微分可能なシミュレーションパイプラインを含む科学機械学習の基盤的な要件である。しかし、PyTorch エコシステムにはこれらのワークロードを支援する統合ライブラリが存在しない。ネイティブの PyTorch サポートは断片的である：torch.sparse は低レベルカーネルと微分不可能な CPU 専用（SuperLU ラッパー）の spsolve を提供する一方、torch.linalg は密行列に限定されている。

既存の回避策は、大規模または複雑なトレーニングシナリオには不十分である：

1. グラフの切断: 外部ソルバー（例：cuDSS、PETSc）を呼び出すと autograd グラフが切断され、ユーザーは勾配を破棄するか、すべての問題に対してカスタム torch.autograd.Function ラッパーを記述することを余儀なくされる。
2. 非効率的な微分: 反復ソルバーの $k$ 反復を素直に逆伝播させると、$O(k)$ ノードの計算グラフが生成される。自由度（DOF）が $10^6$ で 1,000 反復の問題の場合、中間ベクトルの保存には約 80 GB のメモリが必要となり、大規模トレーニングの実現性を損なう。
3. スケーラビリティの限界: $10^7$–$10^9$ の未知数を持つ問題は、分散ドメイン分解を必要とする。既存の PyTorch 分散プリミティブは、後方伝播における疎演算子の随伴に必要な転置ハロー交換をネイティブにサポートしていない。
4. バックエンドの断片化: 直接ソルバー（例：cuDSS）は小規模な問題（$<10^5$ DOF）では高速であるが、より大規模なシステムでは $O(n^{1.5})$ のフィルインとメモリ枯渇に悩まされる。反復ソルバーはスケーリング性が高いが、小規模な問題ではカーネル起動のオーバーヘッドを伴う。デバイスや問題サイズに基づいてこれらの異種バックエンド間でディスパッチする統合インターフェースは存在しない。

手法

著者は、3 つの中核コンポーネントを通じてこのギャップを埋めるオープンソースライブラリ torch-sla を提示する。

1. 統合バックエンド抽象化

torch-sla は、デバイスと問題サイズに基づいて 5 つの交換可能なバックエンド間でディスパッチする、単一の autograd 対応 API（.solve、.matvec、.eigsh、.det）を公開する。

• CPU: SciPy（SuperLU、UMFPACK、CG、BiCGStab）および Eigen（CG、BiCGStab）。
• GPU: NVIDIA cuDSS（LU、Cholesky、LDLT）、CuPy（反復および直接）、および PyTorch ネイティブの反復ソルバー。
• 自動ディスパッチ: システムはバックエンドを自動的に選択する。小〜中規模の問題（例：GPU 上で $<2 \times 10^6$ DOF）には直接ソルバーを、メモリ溢れを回避するために大規模な問題には反復ソルバーを選択する。対称性と正定値性は検出され、因数分解の最適化に利用される（例：LU から Cholesky へのアップグレード）。

このライブラリは、局所性とバッチングによって整理された 4 つの疎テンソルタイプをサポートする：SparseTensor（単一行列/共有パターン）、SparseTensorList（バッチごとの異なるパターン）、およびそれらの分散変種である DSparseTensor と DSparseTensorList。

2. 随伴微分フレームワーク

素直な逆伝播の $O(k)$ メモリおよびグラフ複雑性を回避するため、torch-sla は陰関数定理のアプローチを採用する。

• 線形システム: $Ax=b$において、後方伝播は単一の随伴システム$A^\top \lambda = \partial L / \partial x$を解く。行列に関する勾配は、疎パターン上でのみ$-\lambda_i x_j$ として計算される。
• 非線形システム: $F(u, \theta)=0$ において、後方伝播には 1 回の随伴線形ソルブ（$J^\top \lambda = \partial L / \partial u$）とベクトル - ヤコビアン積が必要であり、非線形反復の数に依存しない。
• 固有値問題: 固有値の勾配は Hellmann–Feynman 定理（$\partial \lambda / \partial A_{ij} = v_i v_j$）を利用し、疎パターン上での外積のみを必要とする。
• 複雑性: このフレームワークは autograd グラフを $O(1)$ ノードに、メモリをソルバー反復数（$k$）に依存しない $O(n + \text{nnz})$ に削減する。

3. ハロー交換を備えた分散レイヤー

マルチ GPU 実行のために、torch-sla は各プロセスが行ブロックとハローメタデータを所有するドメイン分解を実装する。

• 前方伝播: プロセスは局所的な疎行列ベクトル積（SpMV）操作の前にハロー値を交換する。
• 後方伝播: ライブラリは転置ハロー交換を実装する。前方伝播が所有値を隣接ハローにスキャッターするのに対し、後方伝播は入力ハロー勾配を所有プロセスの境界エントリに集約する。
• 実装: 前方および後方の通信の両方は、分散ソルブ全体でエンドツーエンドの微分可能性を確保するために torch.distributed を使用して torch.autograd.Function プリミティブとしてラップされている。

主要な貢献

1. 統合 API: 5 つのバックエンド（SciPy、Eigen、cuDSS、CuPy、PyTorch ネイティブ）にわたる直接、反復、非線形、および固有値ソルバーのための autograd 対応インターフェースを提供する単一ライブラリ。
2. スケーラブルな随伴フレームワーク: 反復ソルバーの逆伝播におけるメモリボトルネックを排除する $O(1)$ グラフ微分法により、数百万の未知数を持つ問題でのトレーニングを可能にする。
3. 分散 Autograd: 転置ハロー交換を介してエンドツーエンドの微分をネイティブにサポートする、PyTorch における分散疎ソルバーの最初の実装。
4. 拡張性: 共通インターフェースを実装することで、外部ライブラリ（例：PETSc、Trilinos、hypre）や学習済み前処理条件子を統合できるバックエンド抽象化。

実験結果

著者は NVIDIA H200 GPU を使用した 2 次元ポアソンベンチマークで torch-sla を評価した。

• スケーラビリティ: ライブラリは、単一 GPU で 1 億 6900 万 DOF、3 台の GPU で 4 億 DOF にスケーリングする。
• パフォーマンス: 100 万 DOF の問題において、torch-sla の CG バックエンド（190 ms）は cuDSS（7.3 秒）を 38 倍上回る。200 万 DOF 以上では、直接ソルバーはメモリ不足になるが、反復ソルバーはほぼ線形のスケーリングを維持する。
• メモリ効率: 1,000 回の CG 反復における素直な逆伝播と随伴法の比較において、素直なアプローチは約 64 GB を必要とし（2,000 反復で OOM）、一方、随伴法は一定の 約 328 MB の使用量を維持し、195 倍のメモリ削減を達成した。
• 勾配検証: 非線形および固有値問題に対して有限差分法との勾配検証を行ったところ、相対誤差は $10^{-5}$ 未満であった。
• エンドツーエンドの応用: ポアソン方程式における空間的に変化する導電率を復元する逆係数学習タスクは、標準的な PyTorch Adam ループを使用して 1 分未満で 0.23% の相対 L2 誤差に収束した。

意義

本論文は、torch-sla が PyTorch エコシステムにおける長年のギャップを埋め、JAX の jax.lax.custom_linear_solve や jax.scipy.sparse.linalg と同様の機能を提供しながら、統合された PyTorch ライブラリ内で実現していると主張する。これにより、研究者は以下のことが可能になる。

• 勾配を破棄したりコードを JAX に移植したりすることなく、疎線形ソルブを含むモデルのトレーニングが可能になる。
• 科学機械学習のワークロードを、分散 GPU クラスター上で数億の自由度にスケーリングできる。
• 疎ソルバーを微分可能なレイヤーとして扱い、学習済み前処理条件子の研究や物理シミュレーションのエンドツーエンド最適化を促進できる。

著者は、PyTorch ネイティブバックエンドにおける Jacobi 前処理条件子への現在の依存（非常に大規模なシステムでの収束速度を制限する）および 2 次元ポアソンベンチマークへの焦点を限界として指摘しており、3 次元 PDE や不定系は今後の課題として特定されている。