Jinze Xue

Machine Learning Interatomic Potentials (MLIPs), trained with Quantum Mechanics data, can model potential energy surfaces for molecular systems with very high accuracy and extreme speedups compared to reference quantum calculations, offering a powerful tool for studying complex chemical and biological systems. This work presents the LAMMPS-ANI interface, which scales our ANI neural network potential models for large systems, demonstrated with up to 100 million atoms across up to 128 NVIDIA A100 GPUs. The high performance of LAMMPS-ANI was achieved through a comprehensive code redesign, in-depth performance profiling, and advanced GPU performance optimizations. Our benchmarks show that ANI is 30 to 60 times faster than the state-of-the-art Allegro Model, emphasizing its speed and efficiency. We highlight our work in large-scale molecular dynamics using ANI potentials, presenting benchmark results for water boxes (up to 100 million atoms) and a solvated HIV capsid (44 million atoms). We also present results for accurately simulating complex reaction processes at unprecedented scales, such as methane combustion (300 thousand atoms) and early Earth chemistry experiment (228 thousand atoms) demonstrating the spontaneous formation of glycine.

In this work we introduce TorchANI-Amber, an interface for routine molecular dy- namics simulations of biomolecular systems using ANI-style machine learning poten- tials. TochANI-Amber incorporates the ANI neural network potentials into the Amber software suite, supporting Amber’s two engines: sander and pmemd. In addition to implementing all published ANI models, the interface is extensible to other energy predicting potentials through a simple mechanism requiring no knowledge of Amber’s codebase. To illustrate this versatility, we implement extensions to the AIMNet2 and Nutmeg potentials. The interface is integrated with Amber’s neighborlists, and it also supports an optimized CUDA implementation for computing the ANI models’ feature vectors, enabling simulations of systems with hundreds of thousands of atoms at the neural network’s level of theory (approaching DFT accuracy). The interface is designed so that all amber capabilities can be used with ANI potentials instead of force fields. To evaluate the energy conservation, stability, and performance of the ANI potentials as used through the interface, we run MD simulations on different biomolecular systems, including ubiquitin and Trp-cage proteins in explicit solvent. Additionally, we demon- strate the use of these potentials in the context of enhanced sampling techniques, such as temperature replica-exchange molecular dynamics.

In this work we introduce TorchANI 2.0, a significantly improved version of the free and open source TorchANI software package for training and evaluation of ANI (ANAKIN-ME) deep learning models. TorchANI 2.0 builds upon the foundation of its predecessor, while addressing its limitations and introducing new features. These changes greatly enhance its extensibility, performance, and suitability as a framework for developing models ready for molecular dynamics applications. These improvements include the introduction of a modular system to add arbitrary pairwise potentials to models, CUDA-accelerated optimization for faster and more memory-efficient calculation of local atomic features, and a batched system for better performance of network ensembles, among others. Our benchmarks demonstrate that TorchANI 2.0 achieves significant speedup over previous versions in both training and inference, and the library enhancements allow users to train physically constrained models that better represent important qualities of chemical systems. We demonstrate this by introducing three new ANI models that incorporate these features and evaluating their capabilities

This work introduces LEGOLAS, a fully open source TorchANI-based neural network model designed to predict NMR chemical shifts for protein backbone atoms. LEGOLAS has been designed to be fast, and without loss of accuracy, as our model is able to predict backbone chemical shifts with root-mean-square errors of 2.69 ppm for N, 0.95 ppm for Ca, 1.40 ppm for Cb, 1.06 ppm for C’, 0.52 ppm for amide protons, and 0.29 ppm for H. The program predicts chemical shifts at least one order of magnitude faster than the widely utilized SHIFTX2 model. This breakthrough allows us to predict NMR chemical shifts for a very large number of input structures, such as frames from a molecular dynamics trajectory. In our simulation of the protein BBL from E. coli, we observe that averaging the chemical shift predictions for a set of frames of an MD trajectory substantially improves the agreement with experiment with respect of using a single frame of the dynamics. We also show that LEGOLAS can be successfully applied to the problem of recognizing the native states of a protein among a set of decoys.

Training stability of large language models(LLMs) is an important research topic. Reproducing training instabilities can be costly, so we use a small language model with 830M parameters and experiment with higher learning rates to force models to diverge. One of the sources of training instability is the growth of logits in attention layers. We extend the focus of the previous work and look not only at the magnitude of the logits but at all outputs of linear layers in the Transformer block. We observe that with a high learning rate the L2 norm of all linear layer outputs can grow with each training step and the model diverges. Specifically we observe that QKV, Proj and FC2 layers have the largest growth of the output magnitude. This prompts us to explore several options: 1) apply layer normalization not only after QK layers but also after Proj and FC2 layers too; 2) apply layer normalization after the QKV layer (and remove pre normalization). 3) apply QK layer normalization together with softmax capping. We show that with the last two methods we can increase learning rate by 1.5x (without model divergence) in comparison to an approach based on QK layer normalization only. Also we observe significant perplexity improvements for all three methods in comparison to the baseline model.