SE(3)-equivariant prediction of molecular wavefunctions and electronic densities

TL;DR

This paper introduces SE(3)-equivariant deep learning architectures for predicting molecular wavefunctions and densities, achieving high accuracy and speedups over traditional methods, enabling efficient property derivation and improved initial guesses for quantum calculations.

Contribution

The authors develop SE(3)-equivariant neural network components for accurate wavefunction prediction, significantly outperforming previous models in speed and precision, and demonstrate applications in transfer learning and initial guess generation.

Findings

Achieves over 1000x speedup compared to ab initio methods.

Reduces prediction errors by up to 100x relative to prior state-of-the-art.

Enables direct derivation of energies and forces from predicted wavefunctions.

Abstract

Machine learning has enabled the prediction of quantum chemical properties with high accuracy and efficiency, allowing to bypass computationally costly ab initio calculations. Instead of training on a fixed set of properties, more recent approaches attempt to learn the electronic wavefunction (or density) as a central quantity of atomistic systems, from which all other observables can be derived. This is complicated by the fact that wavefunctions transform non-trivially under molecular rotations, which makes them a challenging prediction target. To solve this issue, we introduce general SE(3)-equivariant operations and building blocks for constructing deep learning architectures for geometric point cloud data and apply them to reconstruct wavefunctions of atomistic systems with unprecedented accuracy. Our model achieves speedups of over three orders of magnitude compared to ab initio methods and reduces prediction errors by up to two orders of magnitude compared to the previous state-of-the-art. This accuracy makes it possible to derive properties such as energies and forces directly from the wavefunction in an end-to-end manner. We demonstrate the potential of our approach in a transfer learning application, where a model trained on low accuracy reference wavefunctions implicitly learns to correct for electronic many-body interactions from observables computed at a higher level of theory. Such machine-learned wavefunction surrogates pave the way towards novel semi-empirical methods, offering resolution at an electronic level while drastically decreasing computational cost. Additionally, the predicted wavefunctions can serve as initial guess in conventional ab initio methods, decreasing the number of iterations required to arrive at a converged solution, thus leading to significant speedups without any loss of accuracy or robustness.