Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation

Kevin Han Huang, Ni Zhan, Elif Ertekin, Peter Orbanz, Ryan P Adams

TL;DR: We study ways of incorporating diagonal symmetry into neural network wave functions, and show that post hoc averaging outperforms conventional in-training symmetrization.

Abstract: Incorporating group symmetries into neural networks has been a cornerstone of success in many AI-for-science applications. Diagonal groups of isometries, which describe the invariance under a simultaneous movement of multiple objects, arise naturally in many-body quantum problems. Despite their importance, diagonal groups have received relatively little attention, as they lack a natural choice of invariant maps except in special cases. We study different ways of incorporating diagonal invariance in neural network ansatze trained via variational Monte Carlo methods, and consider specifically data augmentation, group averaging and canonicalization. We show that, contrary to standard ML setups, in-training symmetrization destabilizes training and can lead to worse performance. Our theoretical and numerical results indicate that this unexpected behavior may arise from a unique computational-statistical tradeoff not found in standard ML analyses of symmetrization. Meanwhile, we demonstrate that post hoc averaging is less sensitive to such tradeoffs and emerges as a simple, flexible and effective method for improving neural network solvers.

Lay Summary:

We use neural networks to learn how electrons behave in crystals. Due to the large number of electrons, this can be very costly and difficult. Fortunately, crystals have many nice geometric symmetries, which constrains the behaviors of these electrons. There are many success stories in machine learning (ML) about designing neural networks with the right symmetries in mind, which helps to reduce the problem search space. We apply existing approaches from ML and investigate whether they can help with our task for crystals.

Surprisingly, the results are mixed. We find that employing symmetries during neural network training can actually hurt. On the other hand, adding symmetries to a trained network can improve the performance substantially. We also provide theoretical explanations for why these effects occur. These results are particularly interesting, in light of recent works in AI for science that show that symmetries may in fact be unnecessary for obtaining the best performance.

Link To Code: https://github.com/PrincetonLIPS/invariant-DeepSolid

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: quantum Monte Carlo, ab initio methods, Schrodinger equation, neural network wavefunctions, many-body methods, invariance and symmetry

Submission Number: 7816