[論文レビュー] $ξ$-torch: differentiable scientific computing library

Muhammad Kasim, S. M. Vinko|arXiv (Cornell University)|Oct 5, 2020

Model Reduction and Neural Networks参考文献 14被引用数 4

ひとこと要約

本稿では、$ olcc och$ と呼ばれる、PyTorch をベースにしたライブラリを紹介する。このライブラリは、根の探索、常微分方程式（ODE）の数値解法、数値積分法などの科学的計算用の微分可能関数型を提供し、数値的安定性とメモリ効率を向上させるために解析的導出による勾配を備えている。これにより、物理を組み込んだ機械学習におけるエンド・トゥ・エンドの微分可能シミュレーションが可能となり、1次および高次導関数の計算が可能であることが、ミラー設計や分子動力学の最適化において高速収束を示した。

ABSTRACT

Physics-informed learning has shown to have a better generalization than learning without physical priors. However, training physics-informed deep neural networks requires some aspect of physical simulations to be written in a differentiable manner. Unfortunately, some operations and functionals commonly used in physical simulations are scattered, hard to integrate, and lack higher order derivatives which are needed in physical simulations. In this work, we present $ξ$-torch, a library of differentiable functionals for scientific simulations. Example functionals are a root finder and an initial value problem solver, among others. The gradient of functionals in $ξ$-torch are written based on their analytical expression to improve numerical stability and reduce memory requirements. $ξ$-torch also provides second and higher order derivatives of the functionals which are rarely available in existing packages. We show two applications of this library in optimizing parameters in physics simulations. The library and all test cases in this work can be found at https://github.com/xitorch/xitorch/ and the documentation at https://xitorch.readthedocs.io.

研究の動機と目的

• 物理を組み込んだディープラーニングのための既存のライブラリに統合的で微分可能な科学的関数型が不足している問題に対処すること。
• 前向きアルゴリズムへの自動微分ではなく、解析的導関数を用いることで、科学的シミュレーションにおける数値的安定性の向上とメモリ使用量の削減を図ること。
• 振動モード周波数の計算などの高度な科学的シミュレーションに必要な高次導関数（例：2次導関数）を効率的かつ正確に計算できること。
• 純関数と線形演算子の抽象化を提供することで、深層学習と物理シミュレーションの統合を容易にすること。
• ODEソルバーや根の探索器などの複雑なシミュレーション部品を経由した勾配伝搬を可能にすることで、物理的システムのパラメータ最適化を支援すること。

提案手法

• 根の探索、ODEソルバー（solve_ivp）、数値積分法（mcquad, SQuad）などの関数型を、解析的に導出した勾配を備えて実装することで、数値的安定性の向上とメモリ消費量の削減を図る。
• 高次導関数の正しい計算を保証するため、入力に純関数を用いる。出力に影響する状態変数を追跡することで、状態を持つオブジェクトメソッドを純関数に変換するメカニズムを導入する。
• 出力に影響するパrameter名を手動で指定できる EditableModule 基底クラスを導入し、複雑なクラスにおける正確な勾配計算を可能にする。
• 高次元問題における行列構築のコストを回避するため、スパース線形代数演算のための LinearOperator 抽象化を提供する。
• PyTorch の自動微分エンジンを活用し、解析的導関数を組み合わせることで、gradgradcheck の検証を経て 2次以降の高次導関数を計算する。
• SciPy をインspirationとして、最適化、統合、補間、線形代数の各モジュールに分かれたモジュラーなライブラリ構造を設計し、微分可能な科学的計算に特化した形にしている。

実験結果

リサーチクエスチョン

• RQ1 解析的勾配を用いた微分可能な科学的関数型を実装することで、科学的シミュレーションにおける数値的安定性の向上とメモリ使用量の削減が達成できるか？
• RQ2 ODEソルバーや根の探索器のような複雑な関数型に対して、2次導関数などの高次導関数を効率的かつ正確に計算する方法は何か？
• RQ3 統合された微分可能な関数型ライブラリが、科学的応用分野における物理を組み込んだ機械学習の加速にどの程度寄与できるか？
• RQ4 状態を持つオブジェクトメソッドを、微分可能な科学的計算における正しい勾配計算を可能にするために、どのように純関数に変換できるか？
• RQ5 エンド・トゥ・エンドの微分可能なシミュレーションを通じて、$ olcc och$ がミラー設計や分子動力学のような物理的システムの効果的なパrameter最適化を可能にするか？

主な発見

• $ olcc och$ ライブラリは、アルゴリズム的実装への逆誤差伝搬ではなく、解析的表現を用いて科学的関数型の1次および高次導関数を正確に計算でき、数値的安定性が向上し、メモリ使用量が削減された。
• ライブラリは、PyTorch の gradgradcheck を用いた検証を通じて、2次導関数の正しく計算されることを確認し、複雑な微分可能な関数型における堅牢性を示した。
• ミラー設計の事例では、根の探索関数型により、ニューラルネットワークベースのミラー表面が200イテレーションで最適化され、目的の点に光線を正確に集光できた。
• 分子動力学の事例では、solve_ivp 関数型により、初期速度の最適化が可能となり、t=1 時点で「OX」という文字を形成するシナリオが実現され、Adam を用いて7500イテレーションで収束した。
• ライブラリのモジュラー構造と純関数、編集可能な状態変数のサポートにより、複雑な物理シミュレーションパイプラインにおける信頼性の高い勾配計算が可能になった。
• モンテカルロ積分や固定サンプル数積分などの微分可能な関数型の提供により、標準的な ODE や根の探索を越えた微分可能な科学的計算の範囲が拡張された。