A simple approach to rotationally invariant machine learning of a vector quantity

Predicting potential energy surfaces with machine learning is now easier than ever. But once we move from scalar properties (like energy) to vectors or tensors (like dipole moments or polarizabilities), things get trickier. These properties must transform correctly under symmetry transformations - when a molecular geometry is rotated, the corresponding vector should rotate as well. While rotationally equivariant models can be used, we developed a simple and easy-to-use scheme (shown in the figure below), which we call rotate–predict–rotate (RPR).

The procedure consists of three steps:

1. Rotate each molecular geometry into a canonical orientation defined by its tensor of inertia.
2. Predict the vector or tensor property in that orientation using any machine learning model trained on data in the canonical orientation.
3. Rotate the predicted properties back to the original molecular orientation.

To test the method, we studied 1,2-dithioethane, a molecule whose dipole moment and polarizability vary with torsional motion. We trained a kernel ridge regression model using the relative-to-equilibrium descriptor, based on DFT/B3LYP/aug-cc-pVTZ calculations, and sampled the configurational space by running ground-state molecular dynamics. The selected points for the test set are shown in black:

The machine learning models produced accurate, rotationally covariant predictions of both properties across molecular dynamics trajectory frames. This was also achieved thanks to task-specific descriptor design. Because the approach relies only on elementary linear algebra, training is fast, prediction is inexpensive, and the models remain accurate, as shown in the scatter plots of dipole moment and polarizability components:

For more information, check out the paper: (Martinka et al., 2024).