Coarse-Graining with Equivariant Neural Networks: A Path Towards Accurate and Data-Efficient Models

TL;DR

This paper explores how equivariant neural networks can significantly reduce data requirements in coarse-grained molecular modeling, achieving accurate force fields with minimal training data.

Contribution

It introduces equivariant convolutional operations into neural network force fields, enabling accurate coarse-grained models with drastically less data than traditional methods.

Findings

Equivariant networks outperform non-equivariant ones in data efficiency.

Models with equivariant operations can be trained on as little as one frame.

Traditional models require tens of microseconds of data.

Abstract

Machine learning has recently entered into the mainstream of coarse-grained (CG) molecular modeling and simulation. While a variety of methods for incorporating deep learning into these models exist, many of them involve training neural networks to act directly as the CG force field. This has several benefits, the most significant of which is accuracy. Neural networks can inherently incorporate multi-body effects during the calculation of CG forces, and a well-trained neural network force field outperforms pairwise basis sets generated from essentially any methodology. However, this comes at a significant cost. First, these models are typically slower than pairwise force fields even when accounting for specialized hardware which accelerates the training and integration of such networks. The second, and the focus of this paper, is the need for the considerable amount of data needed to train such force fields. It is common to use 10s of microseconds of molecular dynamics data to train a single CG model, which approaches the point of eliminating the CG models usefulness in the first place. As we investigate in this work, this data-hunger trap from neural networks for predicting molecular energies and forces can be remediated in part by incorporating equivariant convolutional operations. We demonstrate that for CG water, networks which incorporate equivariant convolutional operations can produce functional models using datasets as small as a single frame of reference data, while networks without these operations cannot.