Radical AI社は、PyTorchネイティブで構築され、MLIP(機械学習による原子間ポテンシャル)時代のために設計された次世代原子論的シミュレーションエンジンTorchSimのリリースを発表した。材料科学における大きな技術的転換と位置づけられるTorchSimは、ASE(原子シミュレーション環境)やDFT(密度汎関数理論)といった従来のフレームワークと比較して、分子シミュレーションを桁違いに高速化することを約束する。
TorchSimはオープンソースであり、MACE、Fairchem、SevenNetのような機械学習モデルを活用する最新の材料ワークフローをサポートするように設計されている。Lennard-JonesやMorseのような古典的な相互作用ポテンシャルや、NVEやNVT Langevinのような積分スキームをサポートしている。このエンジンは、自動バッチ処理とGPUメモリ管理を特徴としており、複数のシステムで並行してシミュレーションを実行し、GPUハードウェアを最大限に活用できる。
速度の向上は著しい。社内ベンチマークによると、TorchSimは、一般的なMLIPモデルと組み合わせた場合、ASEの最大100倍のスピードアップを達成している。H100 GPU1台で、選択したモデルによっては数千個の原子を同時にシミュレーションできる。この性能向上は、学術的にも産業的にも、原子論的シミュレーションのスケールアップに不可欠である。
出典: TorchSim GitHubリポジトリ
TorchSimのアーキテクチャは、完全にPyTorchベースの実装を採用することで、従来のシミュレーションツールとは一線を画している。この設計の選択により、より広範な機械学習エコシステムとのシームレスな統合が可能になり、微分可能なシミュレーション、弾力的な特性予測、カスタム科学ワークフローの可能性が広がる。
初期のテスターたちは、APIの明快さと柔軟性を高く評価している。Preferred Networksの篠原耕平氏はこう語る:
このパッケージは、GPUバックエンドを使ったMLIP時代に分子動力学と構造緩和を行うための正しい方法だと感じます。APIのデザインはすっきりとしており、JAXのような関数型プログラミングのフレームワークからインスピレーションを得ています。
計算化学者であり、このプロジェクトに貢献しているOrion Archer Cohen氏は次のように述べている:
TorchSimは物事を正しく理解しています。C++とPythonのAPIを分割することも、cythonを使うことも、つらいファイル形式を使うことも、GPU使用率を5%にすることもないです。原子論的シミュレーションのプリミティブをPyTorchで書き換えることで、学習、理解、開発が驚くほど簡単になります。
TorchSimは、ASE、Phonopy、Pymatgenのような既存のツールとの統合をサポートし、一般的なシミュレーションワークフローを簡素化する高レベルAPIを含む。拡張可能なプロパティをサポートする新しいバイナリ軌道フォーマットにより、大規模かつハイスループットの研究に適している。
TorchSimは現在MITライセンスで利用可能で、Python 3.11+をサポートしている。開発者や研究者は、ソースコード、ドキュメント、サンプルワークフローが管理されているGitHubでプロジェクトを探索できる。