Lightweight and Effective Tensor Sensitivity for Atomistic Neural Networks

TL;DR

This paper introduces HIP-NN-TS, an extension of the HIP-NN model that incorporates tensor sensitivity to improve accuracy in atomistic predictions with minimal additional parameters.

Contribution

The paper presents a novel tensor sensitivity framework for HIP-NN that enhances accuracy while maintaining a low parameter count, especially effective on complex datasets.

Findings

HIP-NN-TS outperforms HIP-NN in accuracy across multiple datasets.

Achieves a record MAE of 0.927 kcal/mol on the COMP6 benchmark.

Tensor sensitivities provide greater benefits as dataset complexity increases.

Abstract

Atomistic machine learning focuses on the creation of models which obey fundamental symmetries of atomistic configurations, such as permutation, translation, and rotation invariances. In many of these schemes, translation and rotation invariance are achieved by building on scalar invariants, e.g., distances between atom pairs. There is growing interest in molecular representations that work internally with higher rank rotational tensors, e.g., vector displacements between atoms, and tensor products thereof. Here we present a framework for extending the Hierarchically Interacting Particle Neural Network (HIP-NN) with Tensor Sensitivity information (HIP-NN-TS) from each local atomic environment. Crucially, the method employs a weight tying strategy that allows direct incorporation of many-body information while adding very few model parameters. We show that HIP-NN-TS is more accurate than HIP-NN, with negligible increase in parameter count, for several datasets and network sizes. As the dataset becomes more complex, tensor sensitivities provide greater improvements to model accuracy. In particular, HIP-NN-TS achieves a record mean absolute error of 0.927 kcal/mol for conformational energy variation on the challenging COMP6 benchmark, which includes a broad set of organic molecules. We also compare the computational performance of HIP-NN-TS to HIP-NN and other models in the literature.