Electrostatics from Laplacian Eigenbasis for Neural Network Interatomic Potentials

TL;DR

This work presents Phi-Module, a universal plugin that enforces Poisson's equation within neural interatomic potentials to learn electrostatics in a self-supervised manner. By representing the potential φ and charges ρ in the Laplacian eigenbasis and learning their spectral coefficients with a lightweight α-Net, the method adds a PDE-residual term and an electrostatic energy E_ES to improve total energy predictions while remaining memory-efficient. Theoretical results show convexity and monotone improvement in the objective, and experiments on OE62 and MD22 demonstrate consistent performance gains, stability in long MD runs, and favorable scaling compared to Ewald-based approaches. The approach offers a practical, data-efficient route to incorporate first-principles electrostatics into GNN potentials, with potential extensions to higher-order electrostatics and polarizability in future work.

Abstract

In this work, we introduce Phi-Module, a universal plugin module that enforces Poisson's equation within the message-passing framework to learn electrostatic interactions in a self-supervised manner. Specifically, each atom-wise representation is encouraged to satisfy a discretized Poisson's equation, making it possible to acquire a potential φ and corresponding charges \r{ho} linked to the learnable Laplacian eigenbasis coefficients of a given molecular graph. We then derive an electrostatic energy term, crucial for improved total energy predictions. This approach integrates seamlessly into any existing neural potential with insignificant computational overhead. Our results underscore how embedding a first-principles constraint in neural interatomic potentials can significantly improve performance while remaining hyperparameter-friendly, memory-efficient, and lightweight in training. Code will be available at https://github.com/dunnolab/phi-module.

Figures (10)

• Figure 1: Overview of the proposed $\Phi$-Module. $\Phi$-Module encodes electrostatic constraints based on Poisson's equation into hidden representations of any neural network interatomic potential. $\Phi$-Module is integrated at each step of message passing. It uses lightweight convolutional submodule which we refer to as $\boldsymbol{\alpha}$-Net to estimate coefficients of Laplacian eigenbasis directly from constantly updated atomic representations. Those eigenbasis coefficients are then used to optimize Poisson's equation residual $\|\textbf{L}\boldsymbol{\phi}-\boldsymbol{\rho}\|_2=0$ and compute electrostatic energy term $\textbf{E}^{\text{ES}}$ making an important contribution to predictions and leading to improved performance on computational chemistry problems. See \ref{['sec:phi-module']}.
• Figure 2: $\alpha$-Net. It transforms dense atomic representations into sparser coefficients of Laplacian's eigenbasis to acquire potentials and charges. See \ref{['par:ssl']}.
• Figure 3: Energy MAEs and computation time of baselines and their alternatives with $\Phi$-Module on OE62. $\Phi$-Module achieves comparable error to Ewald summation at the same time being almost as fast as clean baseline. See Section \ref{['par:oe62']}.
• Figure 4: Test expected validation MAE for $\boldsymbol{\Phi}$-$\text{E}_2$GNN against the baseline model on OE62. Any choice of selected hyperparameters leads to improved performance, underlining tuning stability of the $\Phi$-Module. See \ref{['par:hyper-stability']}.
• Figure 5: Total energy over 100 ps NVE simulation for (Left) AT-AT-CG-CG and (Right) Ac-Ala3-NHMe molecules obtained from baseline ViSNet and $\boldsymbol{\Phi}$-ViSNet. Energy drift is bounded at 0.0001% over the full trajectory for $\boldsymbol{\Phi}$-ViSNet in both cases. Moreover, it attains x4 and x2 smaller total magntitude of energy drift respectively compared to the baseline model. See Section \ref{['par:md-stability']}.

Theorems & Definitions (8)

• Theorem 3.1: Exact inner minimizer over $\rho$
• Theorem 3.2: Monotone objective decrease in optimization towards $\rho^\star$
• Proposition F.1: Symmetric vs. asymmetric gradients for the Poisson residual
• proof
• Theorem H.1: Exact inner minimizer over $\rho$
• proof
• Theorem H.2: Monotone objective decrease in optimization towards $\rho^\star$
• proof