A quatum inspired neural network for geometric modeling

TL;DR

SpaTea introduces an equivariant, quantum-inspired Matrix Product State (MPS) based message-passing framework to surpass mean-field limitations of traditional geometric GNNs in SE(3)/E(3) settings. By replacing standard MP layers with a tensor-network contraction technology, it captures higher-order many-body interactions while preserving permutation and rotation symmetries. Empirical results on Newton dynamics and quantum Hamiltonian prediction demonstrate state-of-the-art accuracy, illustrating the method’s practical value for materials science, physics, and related domains. The work also highlights potential quantum-computing advantages and offers flexible integration paths with existing equivariant architectures. Overall, SpaTea broadens the expressive capacity of geometric graphs and opens avenues for quantum-informed learning on complex 3D systems.

Abstract

By conceiving physical systems as 3D many-body point clouds, geometric graph neural networks (GNNs), such as SE(3)/E(3) equivalent GNNs, have showcased promising performance. In particular, their effective message-passing mechanics make them adept at modeling molecules and crystalline materials. However, current geometric GNNs only offer a mean-field approximation of the many-body system, encapsulated within two-body message passing, thus falling short in capturing intricate relationships within these geometric graphs. To address this limitation, tensor networks, widely employed by computational physics to handle manybody systems using high-order tensors, have been introduced. Nevertheless, integrating these tensorized networks into the message-passing framework of GNNs faces scalability and symmetry conservation (e.g., permutation and rotation) challenges. In response, we introduce an innovative equivariant Matrix Product State (MPS)-based message-passing strategy, through achieving an efficient implementation of the tensor contraction operation. Our method effectively models complex many-body relationships, suppressing mean-field approximations, and captures symmetries within geometric graphs. Importantly, it seamlessly replaces the standard message-passing and layer-aggregation modules intrinsic to geometric GNNs. We empirically validate the superior accuracy of our approach on benchmark tasks, including predicting classical Newton systems and quantum tensor Hamiltonian matrices. To our knowledge, our approach represents the inaugural utilization of parameterized geometric tensor networks.

Figures (4)

• Figure 1: The initial step of spatial mixing: organizing neighborhood embeddings according to their (Euclidean) distance, e.g., $\phi^{s_3} \rightarrow \phi^{s_2}$ and $\phi^{s_1} \rightarrow \phi^{s}$. When $\phi^{s_2}$ and $\phi^{s_1}$ are equidistant from $\phi^s$, the spatial mixing process should be permutation-invariant when the order of $s_1$ and $s_2$ is exchanged.
• Figure 2: The second step of spatial mixing: constructing the effective Hamiltonian operator by contracting the embeddings of $\phi^s$ neighbor nodes with their matrix product state kernels. This process is inspired by the sweeping procedure of DMRG, and the resulting operator acts on $\phi^s$ for updating the representation. We note that the matrix-kernel of each node itself follows a matrix-product structure. For a comparison with a similar figure, see Figure \ref{['fig:DMRG']} in the appendix.
• Figure 3: Stacking each spatial and temporal layer will form our quantum inspired deep neural network architecture. We also add the standard residual and layer-normalization blocks as a node update between each layer vaswani2023attention.
• Figure 4: The renormalization in DMRG refers to the fact that the effective Hamiltonian for each node renormalizes (by contraction) the information from other nodes, as shown in this figure. Here, the effective Hamiltonian $\hat{H}_{\text{effective}}$ acting on the s-th state consists of the green renormalized states and the blue operator acting solely on state $s$.