Unifying O(3) Equivariant Neural Networks Design with Tensor-Network Formalism

TL;DR

This paper proposes using fusion diagrams, a technique widely employed in simulating SU(2)-symmetric quantum many-body problems, to design new spatial equivariant components for neural networks, and incorporates a fusion block into pre-existing equivariant architectures (Cormorant and MACE).

Abstract

Many learning tasks, including learning potential energy surfaces from ab initio calculations, involve global spatial symmetries and permutational symmetry between atoms or general particles. Equivariant graph neural networks are a standard approach to such problems, with one of the most successful methods employing tensor products between various tensors that transform under the spatial group. However, as the number of different tensors and the complexity of relationships between them increase, maintaining parsimony and equivariance becomes increasingly challenging. In this paper, we propose using fusion diagrams, a technique widely employed in simulating SU($2$)-symmetric quantum many-body problems, to design new equivariant components for equivariant neural networks. This results in a diagrammatic approach to constructing novel neural network architectures. When applied to particles within a given local neighborhood, the resulting components, which we term "fusion blocks," serve as universal approximators of any continuous equivariant function defined in the neighborhood. We incorporate a fusion block into pre-existing equivariant architectures (Cormorant and MACE), leading to improved performance with fewer parameters on a range of challenging chemical problems. Furthermore, we apply group-equivariant neural networks to study non-adiabatic molecular dynamics of stilbene cis-trans isomerization. Our approach, which combines tensor networks with equivariant neural networks, suggests a potentially fruitful direction for designing more expressive equivariant neural networks.

Figures (2)

• Figure 1: Schematic illustration of the implementation of fusion blocks in the MACE architecture. For each atom the fusion block first fuses all the neighboring atoms for a given radius cut-off by pre-selected fusion diagram templates. Specifically, for each neighboring atom, we fuse the information from the root, neighbor atom, and their connecting edge. Then the fusion block applies an aggregation method: in the present work, we simply sum all the neighbors.
• Figure 2: (a) Illustration of photo-induced cis-trans isomerization of stilbene (b) Initial and end configurations of three representative trajectories, which are Wigner-sampled.dahl1988morse (c) the one-dimensional cut of stilbene ground/ excited-state PESs by rotating the carbon-carbon bond as illustrated in (a), which illustrates the energetic diagram of stilbene isomerization process.