TorchANI 2.0: An Extensible, High-Performance Library for the Design, Training, and Use of NN-IPs.

Area of Science:

• Computational Chemistry
• Machine Learning in Chemistry

Background:

• TorchANI is a popular open-source software for developing deep learning models.
• Existing versions have limitations in extensibility and performance for molecular dynamics.

Purpose of the Study:

• Introduce TorchANI 2.0, a significantly improved version of the TorchANI software package.
• Enhance extensibility, performance, and suitability for molecular dynamics applications.
• Enable training of physically constrained deep learning models.

Main Methods:

• Introduced a modular system for arbitrary pairwise potentials.
• Implemented CUDA-accelerated optimization for atomic feature calculations.
• Developed a batched system for improved network ensemble performance.

Main Results:

• TorchANI 2.0 demonstrates significant speedup in training and inference compared to previous versions.
• New features enhance memory efficiency and computational speed.
• Enabled training of three new ANI models with improved physical constraints.

Conclusions:

• TorchANI 2.0 provides substantial performance and feature enhancements for deep learning in chemistry.
• The improved framework facilitates the development of accurate, physically constrained models for molecular dynamics.
• The software is well-suited for advancing computational chemistry research.