Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration

arXiv:2501.19179 · physics.chem-ph, cs.LG, physics.comp-ph · Submitted 2025-01-31 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration".

Jane: The gist This work proposes the Charge Equilibration Layer for Long-range Interactions (CELLI) to efficiently model non-local interactions in equivariant Graph Neural Network potentials by generalizing the classical charge equilibration…

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at this paper titled "Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration". It sounds like they are tackling the problem of how to make these graph neural network potentials better at handling things that don't just happen right next to each other.

Jane: Exactly. The authors, Paul Fuchs, Michal Sanocki, Julija Zavadlav, they’re from TUM and the Munich Data Science Institute. They are proposing this new layer called CELLI to fix a problem where current models can't see beyond a certain radius.

Lu: It's about moving past strictly local assumptions in these models. They are using equivariant graph neural network potentials like Allegro or MACE, but they realize that those local methods struggle with things like charge transfer or electrostatic interactions that happen across larger distances.

Meng: So the main idea is to get long-range effects into a model that is supposed to be strictly local, which seems a bit contradictory at first glance.

Lalam: Basically, they are introducing a method called CELLI as a model agnostic building block. It’s designed to let these GNN potentials learn non-local interactions by explicitly incorporating charge information derived from the environment into the prediction process itself >

The paper's summary: Tom: Okay, so they introduce this Charge Equilibration Layer for Long-range Interactions, or CELLI. They break down the total potential energy into two parts, a local part and an electrostatic part.

Jane: That’s right. They use the established charge equilibration method called Qeq to compute partial charges and that electrostatic energy component using features already present in the GNN.

Lu: The key is that after computing those equilibrated charges, the GNN then learns the correction term, which is dependent on non-local features provided by those new, equilibrated charges >

Meng: So instead of treating short-range and long-range stuff separately—like subtracting them from training data—they want a single model to handle both things together through this equilibration process.

Lalam: It’s about integrating the Qeq method into the GNN structure so that the network can learn how to use those equilibrated charges to predict those non-local interactions >

The paper's improvements: Tom: The paper points out a few ways this approach is better than what we see now. They show that integrating CELLI into models like Allegro and MACE actually pushes their performance beyond just making the model bigger.

Jane: They found that replacing just one interaction layer with CELLI reduced the energy and force mean absolute errors from fifteen point five meV/atom down to nine point four meV/atom on molecular dynamics simulations >

Lu: They also tested this on datasets like OE62, where they showed that CELLI can generalize across a more diverse chemical space while only marginally increasing computational costs >

Meng: And they mentioned that the small variant of CELLI outperformed the largest Allegro baseline model, which suggests applying this correction scheme is more effective than just scaling up the model size >

Lalam: They compare it directly to other state-of-the-art models, showing that when you integrate CELLI with Allegro, you get results comparable to models that use long-range correction schemes like Neural P3M and Ewald corrections >

Conclusion: Tom: So, what’s the big picture here? They’re showing that we can effectively model those crucial long-range electrostatic effects and charge transfer in MLPs using these equivariant GNNs. It addresses a major limitation of strictly local models.

Jane: The implication is that we can use these highly expressive equivariant GNN architectures to capture things like ionic liquids or protein folding better because they are no longer limited by a strict cutoff radius >

Lu: They’re not just adding an extra term; they are fundamentally changing how the network learns to incorporate environmental information, making it robust for chemically diverse systems >

Meng: From an engineering standpoint, this means we can run these simulations on parallel GPU setups with only a marginal increase in computational cost compared to the original models >

Lalam: Ultimately, "Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration" gives us a way to efficiently model long-range electrostatic interactions and charge transfer by integrating the established Qeq method into a model agnostic building block >

Technical University of Munich

physics.chem-ph, cs.LG, physics.comp-ph

Submitted: 2025-01-31

Updated: 2025-06-27

DOI: 10.1038/s41524-025-01790-4

Code: https://github.com/tummfm/chemtrain

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 90/100

The gist: The gist This work proposes the Charge Equilibration Layer for Long-range Interactions (CELLI) to efficiently model non-local interactions in equivariant Graph Neural Network potentials by

Key concepts

Strictly Local MLPs: These are machine learning potentials that assume atoms only interact with their immediate neighbors. They are fast but cannot capture long-range forces like electrostatic interactions or charge transfer, limiting their ability to model complex molecular behaviors.
Charge Equilibration Method (Qeq): Qeq is a mathematical technique used to redistribute atomic charges within a system to minimize the total energy while keeping the total charge conserved. It calculates partial charges based on electronegativities and chemical hardnesses, providing accurate long-range electrostatic information.
CELLI Architecture: CELLI is a novel component inserted into GNN potentials that splits the total potential energy into a Coulombic part and a correction term. It uses Qeq to compute partial charges and then embeds these equilibrated charges back into the network features, allowing the GNN to learn non-local interactions effectively.

Terminology

GNNs — 2× in this paper · explained in 2 episodes
Accuracy — 1× in this paper · explained in 1 episode
Benchmark — 1× in this paper · explained in 2 episodes
Charge Equilibration Layer for Long-range Interactions — 1× in this paper
Graph Neural Network — 1× in this paper
Graph Neural Network (GNN) — 1× in this paper · explained in 4 episodes
Interpretability — 1× in this paper · explained in 15 episodes
Model-Agnostic — 1× in this paper · explained in 2 episodes

Summary

The gist

This work proposes the Charge Equilibration Layer for Long-range Interactions (CELLI) to efficiently model non-local interactions in equivariant Graph Neural Network potentials by generalizing the classical charge equilibration (Qeq) method, thereby extending GNNs' capability to capture long-range effects while maintaining high interpretability through explicit charges.

Motivation and Problem Statement

Strictly local MLPs, which assume that atomic interactions are dominated by their immediate environment, cannot model interactions beyond a strict effective cutoff radius, such as charge transfer or electrostatic interactions<ref:2501.19179#pg3>. These effectively local MLPs can accurately capture short-range interactions but fail to capture long-range effects like electrostatic interactions, charge transfer, and dispersion effects<ref:2501.19179#pg3>. Long-range interactions are crucial in several key physical phenomena, including molecular aggregation, protein folding, or the behavior of ionic liquids<ref:2501.19179#pg3>. Traditional methods often treat short-range interactions with MLPs separately from long-range interactions by subtracting them from training data and adding them to predictions<ref:2501.19179#pg4>. However, capturing long-range electrostatic effects requires atomic charges that depend on the dynamic chemical environment, which traditional charge prediction methods often fail to account for<ref:2501.19179#pg4>.

CELLI Architecture and Mechanism

CELLI introduces a novel architectural building block that generalizes the Qeq method to a model-agnostic building block for modern equivariant GNN potentials<ref:2501.19179#pg3>. The approach splits the total potential energy U = UCoul + ∆U into an electrostatic component UCoul and a correction ∆U<ref:2501.19179#pg5>. CELLI leverages the Qeq method to compute partial charges Q and electrostatic energy UCoul using features of the GNN<ref:2501.19179#pg5>. Subsequently, the GNN learns the correction ∆U dependent on non-local features provided through the equilibrated partial charges embedded by CELLI<ref:2501.19179#pg5>. This enables flexible integration of Qeq-based charge prediction and non-local interaction modeling within a single model<ref:2501.19179#pg5>.

Implementation Details

The architecture is demonstrated by inserting one instance of CELLI between the Interaction Layers in an equivariant GNN like Allegro<ref:2501.19179#pg6>. The latent features from the previous layers encode scalar descriptions of local edge environments, which are then used by a multi-layer perceptron MLPR to predict a per-edge contribution to electronegativity<ref:2501.19179#pg6>. Species embedding is also used to obtain particle hardnesses Ji using a generalized soft plus activation function<ref:2501.19179#pg6>. The Charge Equilibration Method (Qeq) then takes these environment-dependent electronegativities, species or environment-dependent hardnesses, and charge radii to predict partial charges and the coulombic potential<ref:2501.19179#pg7>. These equilibrated charges are then embedded into the scalar features to provide non-local information to the network<ref:2501.19179#pg7>.

Performance and Generalization

On benchmark systems, CELLI achieves state-of-the-art results for strictly local models<ref:2501.19179#pg3>. When integrated with Allegro and MACE, CELLI can overcome the inherent locality of state-of-the-art MLPs<ref:2501.19179#pg7>. The method generalizes to diverse datasets and large structures while providing high computational efficiency and robust predictions<ref:2501.19179#pg3>. Experiments on benchmark systems show significant improvements over models presented by Ko et al., Shaidu et al., and Kim et al.<ref:2501.19179#pg8>. Furthermore, CELLI can be applied to strictly local GNNs, which are highly parallelizable across multiple GPUs<ref:2501.19179#pg9>.

Conclusion and Future Work

CELLI promises efficient and accurate MD simulations of large, complex structures under strict computational cost constraints<ref:2501.19179#pg9>. The method addresses crucial limitations by leveraging highly expressive equivariant GNNs to model long-range electrostatic effects and charge-state dependence<ref:2501.19179#pg9>. The authors plan to interface CELLI with the large-scale molecular dynamics simulation framework LAMMPS and extend CELLI with other physics-based priors<ref:2501.19179#pg9>.

--- Page 23 ---

The paper presents CELLI, a model-agnostic building block introducing the established Qeq method into highly descriptive equivariant GNN MLPs<ref:2501.19179#pg9>. Using equivariant GNNs, CELLI can propose accurate parameters for chemically highly diverse environments<ref:2501.19179#pg9>. Through the Qeq method, CELLI integrates information about long-range electrostatic interactions and charge transfer into effectively local MLPs<ref:2501.19179#pg9>. Therefore, CELLI offers a solution to the long-standing challenge of accurately modeling long-range interactions with MLPs for chemically diverse systems and applications<ref:2501.19179#pg9>.

How it works

The Charge Equilibration Method (Qeq) proposes to redistribute charges in the system to minimize the total energy while maintaining charge conservation PN i=1 Qi = Qtot<ref:2501.19179#pg5>. The contribution of charges to the total energy UQeq(R, Q) consists of the coulombic interaction between charges UCoul given in equation (5) and a second-order approximation of the charge-core interaction determined by the electronegativities χi and chemical hardnesses Ji<ref:2501.19179#pg5>. The minimum of UQeq is the solution of the linear system ∂2UCoul/∂Qi∂Qj + Jii Qj = −χi, subject to the charge conserving equality constraint 1 T Q = Qtot<ref:2501.19179#pg5>.

Results and Benchmarks

In a series of experiments with crucial charge-transfer and charge-state dependence [3], CELLI integrates with Allegro and MACE to overcome the inherent locality of state-of-the-art MLPs<ref:2501.19179#pg7>. On the OE62 dataset, CELLI can generalize across a more diverse chemical space while only marginally increasing computational costs<ref:2501.19179#pg7>. In molecular dynamics simulations, replacing one interaction layer by CELLI reduced the energy and force mean absolute errors from 15.5 meV/atom to 9.4 meV/atom<ref:2501.19179#pg9>.

Scalability and Stability

CELLI outperformed the baseline Allegro models with only a marginal increase in computational cost<ref:2501.19179#pg9>. The small CELLI variant outperformed the largest Allegro baseline model, demonstrating that applying the Qeq scheme is more effective than merely increasing model size<ref:2501.19179#pg9>. In MD simulations, none of the 16 selected structures suffered from instabilities such as broken bonds or overlapping particles for both Allegro variants<ref:2501.19179#pg9>.

Data Availability

The datasets used in this study are publicly available to download<ref:2501.19179#pg9>. The software chemtrain used to train the models and perform MD simulations is publicly available at https://github.com/tummfm/chemtrain<ref:2501.19179#pg9>.

--- Page 24 ---

The results show that CELLI outperformed the baseline Allegro models (Table 3) with only a marginal increase in computational cost<ref:2501.19179#pg9>. Notably, even the small CELLI variant outperformed the largest Allegro baseline model, demonstrating that applying the Qeq scheme is more effective than merely increasing model size<ref:2501.19179#pg9>. In fact introduction of CELLI to strictly local Allegro model makes its performance comparable to MPNN architectures with state-of-the-art long-range correction schemes<ref:2501.19179#pg9>. Moreover, CELLI combined with Allegro not only improves upon baseline models but also achieves results comparable with state-of-the-art Dimnet++ with Neural P3M and significantly outperforms PaiNN models with Ewald and Neural P3M corrections<ref:2501.19179#pg9>.

Improvements for AI systems

Bold Header: charge transfer modeling via CELLI application

This allows equivariant GNN potentials to model long-range electrostatic interactions and charge transfer by integrating the established Qeq method into a model-agnostic building block.

Bold Header: generalization across diverse chemical spaces

CELLI enables MLPs to generalize across a more diverse chemical space while only marginally increasing computational costs, moving beyond the limitations of models trained on smaller, simpler datasets.

Bold Header: enhanced simulation stability via MD validation

The integration of CELLI into Allegro and MACE provides robust predictions for samples unseen in training, ensuring that models are stable for many timesteps during molecular dynamics simulations without introducing artifacts like broken bonds or overlapping particles.

Bold Header: improved scalability and efficiency

CELLI allows the implementation of long-range corrections to strictly local Allegro model makes its performance comparable to MPNN architectures with state-of-the-art long-range correction schemes, potentially increasing compatibility for running simulations across multiple GPUs with only a marginal increase in computational cost.

Abstract

Graph Neural Network (GNN) potentials relying on chemical locality offer near-quantum mechanical accuracy at significantly reduced computational costs. Message-passing GNNs model interactions beyond their immediate neighborhood by propagating local information between neighboring particles while remaining effectively local. However, locality precludes modeling long-range effects critical to many real-world systems, such as charge transfer, electrostatic interactions, and dispersion effects. In this work, we propose the Charge Equilibration Layer for Long-range Interactions (CELLI) to address the challenge of efficiently modeling non-local interactions. This novel architecture generalizes the classical charge equilibration (Qeq) method to a model-agnostic building block for modern equivariant GNN potentials. Therefore, CELLI extends the capability of GNNs to model long-range interactions while providing high interpretability through explicitly modeled charges. On benchmark systems, CELLI achieves state-of-the-art results for strictly local models. CELLI generalizes to diverse datasets and large structures while providing high computational efficiency and robust predictions.