Reciprocal Space Attention for Learning Long-Range Interactions

TL;DR

Long-range interactions are essential in heterogeneous environments but standard MLIPs, which rely on local cutoffs, fail to capture them. RSA introduces a reciprocal-space attention mechanism that maps linear-time attention into Fourier space using Fourier Positional Encoding (FPE) to represent long-range contributions, integrated with a short-range backbone like MACE to yield LR-MACE. The approach reproduces correct long-range behavior across SN2 reactions, dimer binding, random-charge systems, NaCl, phosphorene exfoliation, and bulk water, without predefined charges or empirical corrections, while maintaining end-to-end differentiability. This data-driven, scalable framework broadens the applicability of MLIPs to diverse materials and molecular systems by explicitly learning electrostatics, dispersion, and other nonlocal effects under periodic boundary conditions.

Abstract

Machine learning interatomic potentials (MLIPs) have revolutionized the modeling of materials and molecules by directly fitting to ab initio data. However, while these models excel at capturing local and semi-local interactions, they often prove insufficient when an explicit and efficient treatment of long-range interactions is required. To address this limitation, we introduce Reciprocal-Space Attention (RSA), a framework designed to capture long-range interactions in the Fourier domain. RSA can be integrated with any existing local or semi-local MLIP framework. The central contribution of this work is the mapping of a linear-scaling attention mechanism into Fourier space, enabling the explicit modeling of long-range interactions such as electrostatics and dispersion without relying on predefined charges or other empirical assumptions. We demonstrate the effectiveness of our method as a long-range correction to the MACE backbone across diverse benchmarks, including dimer binding curves, dispersion-dominated layered phosphorene exfoliation, and the molecular dipole density of bulk water. Our results show that RSA consistently captures long-range physics across a broad range of chemical and materials systems. The code and datasets for this work is available at https://github.com/rfhari/reciprocal_space_attention

Figures (6)

• Figure 1: (a) Overview of the short-range and long-range GNN architecture. Here, $Z_i$ denotes the atomic number of atom $i$, $\mathbf{C}$ are the crystal lattice vectors, $\mathbf{R}_i$ is the position of atom $i$, and $\mathbf{R}_{ij}$ is the displacement vector from atom $j$ to atom $i$ (i.e., $\mathbf{R}_{ij}=\mathbf{R}_j-\mathbf{R}_i$). Each interaction layer couples a short-range message-passing (SR-MP, blue) block with a long-range message-passing (LR-MP, green) block, resulting in a complete atomic representation that are passed through a readout head to predict energy and forces. (b) Schematic depiction of Reciprocal Space Attention (RSA) module, which provides the long-range message-passing channel. Atomic features $h_m^L$ and $h_n^L$ are first projected into query, key, and value vectors through feed-forward layers followed by nonlinearities $\sigma_{+}$ to obtain $\phi(Q_m)$ and $\phi(K_n)$. The transformed feature vectors are then rotated using Fourier Positional Encodings (FPE), parameterized by reciprocal lattice vectors $\mathbf{k}$, yielding the rotated $\tilde{\mathbf Q}_m$ and $\tilde{\mathbf K}_n$ vectors. The rotated keys, $\tilde{\mathbf K}_n$, are combined with the value vector $\mathbf V_n$ via an outer product, $\tilde{\mathbf K}_n \otimes \mathbf V_n$, to obtain a graph-level key-value cache. This cache is broadcast back to the node space and contracted (left-multiplied) with the rotated queries $\tilde{\mathbf Q}_m$ at the corresponding $k$-frequency, producing per-mode interactions for each node. Summation over reciprocal modes $\sum_k$ yields the long-range message. This framework provides a linear-scaling, end-to-end differentiable treatment of long-range interactions that naturally accounts for interactions across periodic images of an atom.
• Figure 2: Binding energy curves for the S$_\mathrm{N}$2 reaction complex comparing DFT reference data with predictions from the LR-MACE and SR-MACE models. The LR-MACE model closely reproduces the DFT energy profile along the reaction coordinate, capturing the correct long-range interaction behavior, whereas the SR-MACE model saturates beyond the cutoff due to the absence of explicit long-range interactions.
• Figure 3: Binding energy curves comparing density functional theory (DFT) reference data with predictions from the LR-MACE and SR-MACE models for (a) the CP dimer and (b) the water dimer. The LR-MACE model closely follows the DFT reference across all intermolecular separations, particularly in the long-range tail of the potential, whereas the SR-MACE model saturates beyond the cutoff range due to the absence of explicit long-range interactions.
• Figure 4: Inter-layer interaction dominated by dispersion. (a) Illustration of Phosphorene sheet exfoliation based on the original work of Deringer et al.Deringer2020 with the corresponding exfoliation curve shown in (b).
• Figure 5: (a) Snapshot of bulk water system at 300 K from a molecular dynamics trajectory. (b) Oxygen–oxygen radial distribution function ($\rm g_{\mathrm{OO}}(r)$) comparing the SR-MACE and LR-MACE models. Both models reproduce nearly identical structural features, indicating that the local structure of liquid water is primarily determined by short-range interactions, which the LR model successfully learns directly from the data. The black solid line represents the experimental reference for liquid water skinner2014structure, showing excellent agreement with both models.