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) .

