ANISOTROPIC MESSAGE PASSING: GRAPH NEURAL NETWORKS WITH DIRECTIONAL AND LONG-RANGE INTERACTIONS

Abstract

Graph neural networks have shown great potential for the description of a variety of chemical systems. However, standard message passing does not explicitly account for long-range and directional interactions, for instance due to electrostatics. In this work, an anisotropic state based on Cartesian multipoles is proposed as an addition to the existing hidden features. With the anisotropic state, message passing can be modified to explicitly account for directional interactions. Compared to existing models, this modification results in relatively little additional computational cost. Most importantly, the proposed formalism offers as a distinct advantage the seamless integration of (1) anisotropic long-range interactions, (2) interactions with surrounding fields and particles that are not part of the graph, and (3) the fast multipole method. As an exemplary use case, the application to quantum mechanics/molecular mechanics (QM/MM) systems is demonstrated.

1. INTRODUCTION

Message passing graph neural networks (GNN) have shown great potential for the description of a wide range of chemical systems (Scarselli et al., 2009; Battaglia et al., 2016; 2018) . Particularly the description of quantum molecular (QM) systems with machine learning (ML) potentials has received a lot of interest (Gilmer et al., 2017; Unke et al., 2021b) . However, in its general form, message passing does not explicitly account for directionality, which plays an important role in many physical interactions (Glotzer & Solomon, 2007; Kramer et al., 2014) . In recent years, a range of models which include directional information have been proposed (Anderson et al., 2019; Klicpera et al., 2020; Miller et al., 2020; Schütt et al., 2021) . Especially models based on Clebsch-Gordan tensor products have shown superior data efficiency (Batzner et al., 2022; Batatia et al., 2022b; Musaelian et al., 2022) . However, the relatively high cost of these operations might hinder the application to larger systems such as biomolecules in solution. This difficulty is compounded by growing evidence that message passing cannot accurately resolve long-range interactions (Alon & Yahav, 2020; Dwivedi et al., 2022) . Note that we consider here interactions as long-range if convergence in real space is slow or non-existent, e.g. electrostatic interactions or polarization. In light of these challenges, which are exemplary illustrated in Figure 1 , we propose a model with the aim to (1) include directional information while (2) retaining computational efficiency and (3) incorporating (anisotropic) long-range interactions. Specifically, the addition of an anisotropic state to the existing hidden features is proposed. This anisotropic state is based on Cartesian multipoles and expressed as a linear combination of local frames. As a result, these multipoles are equivariant under rotations and can be used to describe anisotropy of interactions. The formalism is developed analogously to the concept of atomic multipoles commonly used in computational chemistry (Stone, 2013) . The proposed modification results in relatively little computational overhead compared to standard message passing models. Most importantly, the formulation based on multipoles allows for the hybrid treatment of particles and fields. This is of particular interest for two cases: (1) Systems of particles embedded in an external field, and (2) systems with large numbers of particles where long-range interactions may be treated with the fast multipole method (Rokhlin, 1985) . As a practical use case, we consider quantum mechanics/molecular mechanics (QM/MM) simulations (Warshel & Levitt, 1976) . In QM/MM simulations, a QM system is embedded in an external electrostatic field generated by MM particles, allowing for an efficient description of large systems, for instance protein-ligand complexes or enzymatic reactions in solution. As such, the QM/MM formalism might be ideally suited to apply ML potentials to extended systems. Consequently, a ML/MM formalism within the message passing framework is formulated analogous to QM/MM. 

2. RELATED WORK

ML potentials have sparked interest as a possible solution to the steeply scaling computational cost of quantum chemical methods (Schuch & Verstraete, 2009; Unke et al., 2021b) , with high-dimensional neural network potentials being some of the earliest proposed methods (Behler, 2011) . With the introduction of GNN (Scarselli et al., 2009; Battaglia et al., 2016; 2018) , the development of ML potentials has shifted away from handcrafted features to end-to-end differentiable models which learn features solely from distances and element types (Gilmer et al., 2017; Unke et al., 2021b) . Recognizing that GNN may fail to distinguish certain graphs has spurred the development of modified message passing schemes (Morris et al., 2019; Maron et al., 2019; Pozdnyakov & Ceriotti, 2022) . As a solution, the explicit integration of many-body interactions in the message passing (Kondor, 2018; Klicpera et al., 2020; Shui & Karypis, 2020; Zhao et al., 2021) , directional messages (Anderson et al., 2019; Schütt et al., 2021), and combinations thereof (Gasteiger et al., 2021; Batatia et al., 2022b) have been proposed. Building on the concept of equivariant CNNs (Cohen & Welling, 2016; Weiler et al., 2018) , equivariant message passing models based on Clebsch-Gordan tensor products have received increased attention (Anderson et al., 2019; Miller et al., 2020; Satorras et al., 2021; Brandstetter et al., 2021; Batzner et al., 2022; Musaelian et al., 2022) . A categorization based on employed architectural features was recently proposed by Batatia et al. (2022a) , offering a succinct overview. Deficiencies in the description of long-range interactions have led to models that include explicit interaction terms (Unke & Meuwly, 2019; Ko et al., 2021) , as well as models that integrate this information in the representation itself (Grisafi & Ceriotti, 2019; Grisafi et al., 2021; Unke et al., 2021a) . As an alternative to the ML potentials mentioned so far, ML within the QM/MM formalism might facilitate the description of condensed phase systems. Recent work has highlighted the challenges and potential of ML for the simulation of QM/MM systems (Zhang et al., 2018; Böselt et al., 2021; Pan et al., 2021; Hofstetter et al., 2022; Giese et al., 2022) .



Figure1: Exemplary situations arising in molecular systems for which message passing with distance labelled edges fails. (A): Rotation of one molecule by an angle α cannot be recognized without adding more edges. However, this change can be resolved through the addition of a dipole (blue arrow). (B): Similarly, a rotation by an angle β around the axis described by the aligned dipoles ('torsion') can be resolved by including a quadrupole (blue/red crossed arrows). (C): An electric field caused by charges (+) surrounding the molecule. While the electrostatic interaction is additive, the polarization on the black atom due to the opposing charges cancels out in the left case but adds up in the right case. (D): 'Long-range' interactions (dashed lines): Message passing can neither distinguish distance nor direction for the interactions between the red and the two blue-dotted particles.

