End-to-End Modeling of Reaction Field Energy Using Data-Driven Geometric Graph Neural Networks

Abstract

Electrostatic interactions are fundamental to the structure, dynamics, and function of biomolecules, with broad applications in protein–ligand binding, enzymatic catalysis, and nucleic acid regulation. The Poisson–Boltzmann (PB) equation provides a physically grounded framework for modeling these interactions. However, solving the PB equation for large and complex biomolecular systems remains computationally expensive as traditional numerical solvers scale poorly with system size. While the Generalized Born (GB) model offers a more computationally efficient approximation, it does so at the cost of reduced accuracy relative to full PB solutions. To overcome these limitations, we propose PBGNN, a novel end-to-end framework that uses data-driven geometric graph neural networks to directly approximate PB electrostatic energies without relying on the GB approximation. PBGNN incorporates sinusoidal embeddings of atomic charges and a message-passing architecture to efficiently capture long-range interactions in large biomolecules. To address training instability caused by high variance in atomic electrostatic potentials, we introduce a charge-weighted mean squared error (CMSE) optimization objective that improves convergence. We benchmark PBGNN on the AMBER PBSA suite and PBSMALL, a new dataset designed for rapid evaluation of small-molecule electrostatics in drug discovery contexts. The results demonstrate that PBGNN consistently achieves high accuracy in predicting the PB energy with linear computational complexity. Furthermore, it provides reliable and precise PB free energy predictions for both large biomolecular complexes and small-molecule datasets, showcasing its strong generalizability, scalability, and potential utility in drug discovery tasks requiring accurate electrostatic modeling of small molecules. Comprehensive ablation studies further reveal the impact of architectural components, such as geometric representation, objective design, and cutoff strategy, informing future research directions. Finally, we release PBGNN as an open-source, self-contained codebase, along with preprocessed datasets and a complete training and evaluation pipeline, to support scalable and accurate electrostatic analysis that facilitates future research in associated areas.

Graphical Abstract

1. INTRODUCTION

Electrostatic interactions are fundamental to understanding the structure, dynamics, and function of molecules across a wide range of scientific disciplines, including but not limited to computational biochemistry and biophysics. These interactions are critical not only for the stability and formation of biomolecular structures but also for regulating essential biological processes such as protein–ligand binding,^1,2 enzyme catalysis,^3–5 and DNA transcription.^6–8 A key tool for modeling these interactions is the Poisson–Boltzmann (PB) equation, which describes the electrostatic potential in and around charged biomolecules immersed in an ionic solution.^9,10 The accurate and efficient solution of the PB equation is vital for elucidating the underlying electrostatic mechanisms that govern biomolecular behavior. For instance, in drug discovery, precise electrostatic calculations enable the reliable prediction of binding affinity and specificity, which in turn impacts the efficacy and safety of therapeutic candidates,^11,12 although PB-based methods are known to exhibit accuracy degradation in the high-accuracy regime.^13–15

However, solving the PB equation for large and complex biomolecular systems is computationally demanding due to the intricate geometries involved and the need to model the mobile ionic environment. Traditional numerical methods, including finite difference^16–18 and finite element methods,^19,20 although accurate, suffer from poor scalability.^21–25 Among these, the finite difference method is commonly employed due to its relative simplicity and has been implemented in widely used solvers such as AMBER PBSA,^12,26–29 Delphi,^30,31 APBS,^10,32 MIBPB,^33,34 and CHARMM PBEQ.³⁵ While effective for small systems, these solvers become a computational bottleneck when applied to large biomolecular complexes or high-throughput tasks such as molecular dynamics simulations or virtual screening.³⁶ Thus, developing faster, yet still accurate alternatives to classical PB solvers is essential for advancing scalable and interactive modeling of biomolecular electrostatics.¹⁰ To address these limitations, several methods have been proposed. Among them, the Generalized Born (GB) model^21–24 offers a more computationally efficient approximation to long-range electrostatic interactions by modeling the solvation free energy as a superposition of effective spherical cavities. While the GB model improves computational efficiency and scalability, it sacrifices some accuracy compared to full PB solutions.

Recent advances in machine learning, particularly neural networks, offer new opportunities for data-driven modeling in biomolecular systems. These models leverage large datasets to learn complex spatial relationships, achieving notable success across chemistry and physics applications.^37–46 Additionally, with the growing availability of high-performance hardware such as Graphics Processing Units (GPUs) and Tensor Processing Units (TPUs), neural network-based PB energy prediction models have become increasingly feasible, owing to the fast inference speeds of neural networks on such accelerators. For instance, PBML²⁵ learns a correction model for the GB energy by modeling the residual between PB and GB predictions. While effective, this approach inherits the limitations of the GB model, as it relies on it as a backbone.

In this work, we present PBGNN, a novel, fully end-to-end framework that directly approximates PB electrostatic energies without relying on the GB model. PBGNN is motivated by recent advances in data-driven geometric deep learning and graph neural networks (GNNs). It constructs invariant geometric graphs that preserve fundamental physical symmetries, including invariance to translation, rotation, reflection, and discretization, thereby offering a robust representation of biomolecular systems. To enhance learning, we incorporate sinusoidal feature embeddings for scalar inputs and adopt a message-passing architecture to efficiently capture long-range interactions in large biomolecules. Furthermore, we introduce a charge-conditioned regression objective, termed charge-weighted mean squared error (CMSE), which mitigates optimization challenges arising from high variance in atomic electrostatic potentials, thereby improving training stability and convergence. Given the critical role of small molecules in drug discovery, accurate modeling of their electrostatic properties is essential for tasks such as ligand docking, lead optimization, and virtual screening. To facilitate evaluation in this regime, we introduce a new dataset, PBSMALL. PBSMALL is designed to enable the rapid prototyping and benchmarking of machine learning models for PB energy predictions and to support broader research within the community in this area.

We evaluated PBGNN on a diverse set of unseen biomolecules, including systems from the AMBER PBSA suite and PBSMALL. The results demonstrate that PBGNN consistently outperforms the GB model in PB energy prediction while maintaining robustness across variations in the molecular structure and grid resolution. Notably, PBGNN exhibits favorable scalability, achieving linear computational complexity in the PB energy prediction compared to the quadratic complexity of the GB model on CPUs. This enables more efficient electrostatic energy analysis for large-scale biomolecular systems. We further conducted comprehensive ablation studies to assess the individual contributions of architectural components, including geometric representation, loss function, and cutoff strategy, to the overall performance. These investigations offer insights into the key challenges and design considerations associated with modeling PB energy via an end-to-end approach, which informs future research directions. Finally, we have made our method available as an open-source, self-contained codebase, which includes both preprocessed datasets along with the complete training and evaluation pipeline. This release provides the research community with a scalable and accurate solution for the large-scale electrostatic analysis of biomolecular systems and facilitates future research in this area.

2. METHOD

In this section, we first briefly review the essential background knowledge of PB free energy calculations and then elaborate on the proposed PBGNN method. An overview of our proposed PBGNN is shown in Figure 1.

Figure 1.

Overview of the PBGNN framework for predicting electrostatic solvation free energy. (a) Input features extracted from PQR files include atom positions ${\mathbf{r}}_i$, and atom charges $q_i$. (b) The molecular graph is constructed using a cutoff-based approximate nearest neighbor search to form an adjacency matrix, followed by sinusoidal embedding to produce atom-wise features $\mathbf{H}$. These features, along with the graph structure and spatial distances, are passed through a message-passing graph neural network (PBGNN) to predict atomic potential contributions ${\hat{y}}_i$. (c) The predicted electrostatic solvation free energy $\mathrm{\Delta }\hat{G}$ is computed as a weighted sum of per-atom potentials by using the atomic charges.

2.1. PB Electrostatic Solvation Free Energy.

To analyze electrostatic potentials in biomolecular systems, the PB model is commonly employed. The PB model for the electrostatic potential of biomolecular systems is defined over the interior solute domain, ${\mathrm{\Omega }}_{\mathrm{i}\mathrm{n}\mathrm{t}}$, and the exterior solvent domain, ${\mathrm{\Omega }}_{\mathrm{e}\mathrm{x}\mathrm{t}}$, where dissolved ions are approximated by the Boltzmann distribution. In the interior solute domain ${\mathrm{\Omega }}_{\mathrm{i}\mathrm{n}\mathrm{t}}$, fixed charges $\left\{q_i\mid i=1:N_{\mathrm{a}}\right\}$ are located at their corresponding atomic centers $\left\{{\mathbf{r}}_i\mid i=1:N_{\mathrm{a}}\right\}$, where $N_{\mathrm{a}}$ is the number of atoms. These two domains are separated by a dielectric interface $\mathrm{\Gamma }$, with the solvent-excluded surface or molecular surface being the most commonly used interface.⁴⁷ Figure 2a illustrates the PB model. Based on this, the corresponding electrostatic potential described by the PB model¹⁶ is formulated as

$$\begin{gathered} \nabla \cdot \left(ϵ\right(\mathbf{r})\nabla \varphi (\mathbf{r}\left)\right)=-\rho \left(\mathbf{r}\right)-\lambda \left(\mathbf{r}\right)\sum _i{n}_i{q}_i \\ \mathrm{e}\mathrm{x}\mathrm{p}\left[-q_i\varphi \left(\mathbf{r}\right)/\mathrm{k}\mathrm{T}\right] \end{gathered}$$ (1)

where $\varphi \left(\mathbf{r}\right)$ represents the electrostatic potential at a given field position $\mathbf{r},ϵ\left(\mathbf{r}\right)$ is the spatially varying dielectric constant, and $\rho \left(\mathbf{r}\right)$ is the charge density. The second term on the right side represents the mobile ion charge density in the solvent, where $\lambda \left(\mathbf{r}\right)$ is the masking function for the Stern layer, $n_i$ is the number density of ion type $i,q_i$ is the charge of ion type $i,k$ is the Boltzmann constant, and $T$ is the temperature. The dielectric constant $ϵ\left(\mathbf{r}\right)$ is defined as

$$ϵ(\mathbf{r})=\left\{\begin{array}{ll}{ϵ}_1,& \mathbf{r}\in {\mathrm{\Omega }}_{\mathrm{i}\mathrm{n}\mathrm{t}},\\ {ϵ}_2,& \mathbf{r}\in {\mathrm{\Omega }}_{\mathrm{e}\mathrm{x}\mathrm{t}},\end{array}\right.$$ (2)

Figure 2.

Illustration of molecular features and sinusoidal encoding. (a) Schematic representation of the PB model. The solute domain ${\mathrm{\Omega }}_{\mathrm{i}\mathrm{n}\mathrm{t}}$ and solvent domain ${\mathrm{\Omega }}_{\mathrm{e}\mathrm{x}\mathrm{t}}$ are separated by the solvent-excluded surface Γ. Atoms are depicted as spheres with “+” and “−” symbols indicating partial charges, and the probe sphere represents the solvent-accessible boundary. (b) Visualization of sinusoidal positional encoding applied to the atom charge feature $q^{\left(i\right)}$, illustrating how different charge values are mapped across the embedding dimensions.

In many cases, particularly when the electrostatic potential is small relative to thermal energy (i.e., $\left|q_i\varphi \left(\mathbf{r}\right)/kT\right|\ll 1$), the nonlinear PB equation can be approximated by its linearized form.^27,48,49 This leads to the following expression:

$$\nabla \cdot \left(ϵ\right(\mathbf{r})\nabla \varphi (\mathbf{r}\left)\right)-\lambda \left(\mathbf{r}\right)\sum _i{n}_i{q}_i^2\varphi \left(\mathbf{r}\right)/\mathrm{k}\mathrm{T}=-\rho \left(\mathbf{r}\right)$$ (3)

Consequently, many partial differential equation (PDE) solving techniques, such as the finite difference method (FDM)^50–53 and preconditioned conjugate gradient methods,¹⁰ have been employed to calculate the electrostatic potential value $\varphi \left(\mathbf{r}\right)$. To this end, the PB electrostatic solvation free energy $\mathrm{\Delta }G$ is calculated as

$$\mathrm{\Delta }G=\frac{1}{2}\sum _{i=1}^{N_{\mathrm{a}}}{q}_i{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)$$ (4)

where ${\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)$ represents the reaction field potential, which can be computed directly via a singularity-free form of eq 1.²⁹ Efficiently finding numerical solutions to the PB equation for biomolecules is challenging due to the requirements of solving the PDE. To address this, several methods, such as the GB model,⁵⁴ have been proposed to improve computational efficiency by providing approximations to the PB model. While previous works have significantly advanced this area,^{16–18,21–24} the computationally intensive nature of accurately solving the PB equation remains a major limitation, particularly when scaling up to large biomolecules containing tens of thousands to millions of atoms.

2.2. PBGNN: End-to-End PB Energy Prediction Model Using Data-Driven Invariant Geometric Graph Neural Network.

A molecule $ℳ$ consisting of $N_{\mathrm{a}}$ atoms in Euclidean space is represented as $ℳ=\left\{{\mathbf{a}}_1,\dots ,{\mathbf{a}}_{N_{\mathrm{a}}}\right\}$, where each atom ${\mathbf{a}}_i$ is defined as a tuple ${\mathbf{a}}_i=\left(q_i,{\mathbf{r}}_i\right)$, with $q_i$ denoting the atomic charge and ${\mathbf{r}}_i$ representing the spatial position of the atom. Our objective is to develop an end-to-end, data-driven machine learning model $f$, based on a neural network architecture, that predicts the PB electrostatic solvation free energy, denoted as $\mathrm{\Delta }\hat{G}$, from the molecular structure $ℳ$. Designing such a model poses significant challenges that must be addressed to ensure physical fidelity and computational feasibility. Specifically, the model must: (1) efficiently capture long-range electrostatic interactions, which are an essential requirement dictated by the PB framework; (2) respect fundamental physical symmetries, including invariance to translations, rotations, permutations, and discretization artifacts; and (3) scale effectively to large biomolecular systems without relying on the GB model.

To address these challenges, we propose PBGNN, an end-to-end, data-driven PB energy prediction model grounded in graph neural networks. The details of our proposed method are provided in subsequent sections.

2.2.1. Invariant Geometric Graph for Biomolecules.

To account for long-range electrostatic interactions among atoms in a biomolecule, the PB model explicitly formulates the problem as a partial differential equation. In contrast, we propose representing the geometry and topology of a molecule using a graph-based approach in PBGNN. However, it is well established that GNNs struggle to preserve essential physical symmetries, such as invariance to translations, rotations, and grid discretization, which are critical for accurately modeling PB electrostatic solvation free energy.⁵⁵

To address this challenge, we propose an invariant geometric graph representation to more effectively capture the geometry and topology of biomolecules in modeling the PB electrostatic solvation free energy. An undirected geometric graph representation of a biomolecule is formally defined as

$$\overrightarrow{𝒢}:=(\mathbf{A},\mathbf{H},\overrightarrow{\mathbf{R}})$$ (5)

where $\mathbf{A}\in {[0,1]}^{N_{\mathrm{a}}\times N_{\mathrm{a}}}$ denotes the adjacency matrix that encodes the connectivity among all $N_{\mathrm{a}}$ atoms in a biomolecule $ℳ,\mathbf{H}\in {\mathbb{R}}^{N_{\mathrm{a}}\times C_{\mathrm{h}}}$ represents the atom feature matrix, with each row corresponding to a $C_{\mathrm{h}}$-dimensional feature vector describing atomic properties learned by neural networks, and $\overrightarrow{\mathbf{R}}$ denotes the matrix containing the spatial positions of all atoms. An illustration of the invariant geometric graph is presented in Figure 1b. The $i$-th row of $\mathbf{H}$ and $\overrightarrow{\mathbf{R}}$, denoted as ${\mathbf{h}}_i\in {\mathbb{R}}^{{𝒞}_{\mathrm{h}}}$ and ${\mathbf{r}}_i\in {\mathbb{R}}^3$, represent the feature vector and the spatial position of atom ${\mathbf{r}}_i$, respectively. Geometric graphs have been shown to preserve invariance under node permutations, orthogonal transformations, and translations.⁵⁶ Therefore, the geometric graph representation $\overrightarrow{𝒢}$ of a molecule $ℳ$ remains invariant under molecular translations, rotations, permutations, and discretization artifacts.

To model long-range electrostatic interactions, we consider the geometry and topology of the biomolecules. Specifically, for each atom ${\mathbf{a}}_i$, an interaction with atom ${\mathbf{a}}_j$ is defined if and only if the Euclidean distance between them is less than a predefined cutoff threshold $K$ Å, referred to as the cutoff distance. This interaction can be formally expressed as

$$\begin{array}{l}{d}_{ij}=‖{\mathbf{r}}_i-{\mathbf{r}}_j‖,\\ {\mathbf{A}}_{ij}=\mathbf{1}\left[d_{ij}<K\right],\end{array}$$ (6)

where $d_{ij}$ denotes the Euclidean distance between atoms ${\mathbf{a}}_i$ and ${\mathbf{a}}_j$, and $\mathbf{1}[\cdot ]$ is the indicator function that returns 1 if the condition is satisfied and 0 otherwise. It is important to note that this computation is performed over all possible atom pairs in the molecular system $ℳ$ to construct the adjacency matrix $\mathbf{A}$, resulting in $\left(\genfrac{}{}{0pt}{}{N_{\mathrm{a}}}{2}\right)$ pairwise combinations. We discuss strategies for computing this formulation efficiently in subsequent sections.

For each node in the graph, PBGNN utilizes the atomic charge $\left\{q_i\mid i=1:N_{\mathrm{a}}\right\}$ as the input feature. We adopt a sinusoidal feature embedding strategy to encode the scalar input, rather than directly using the scalar charge value as the node feature. Inspired by the neural tangent kernel literature,⁵⁷ incorporating features across multiple frequencies facilitates more efficient convergence in neural networks. Specifically, PBGNN encodes each atomic charge into a $C_{\mathrm{h}}$-dimensional vector, thereby enabling the neural network to better capture the underlying semantics of the scalar quantity. Figure 2b illustrates the core idea of this sinusoidal embedding approach using 32-dimensional vectors: it allows the PBGNN to distinguish between positive and negative scalar values through the flipped pattern observed in the first 16 dimensions and to interpret quantitative differences via variations in frequency components. In detail, sinusoidal feature embedding can be formulated as an embedding function that maps a scalar input (i.e., the atomic charge $q_i$) to a $C_h$-dimensional vector. This process is defined as follows:

$$\begin{array}{ll}{\mathbf{h}}_{ij}& :=\left\{\begin{array}{ll}\mathrm{s}\mathrm{i}\mathrm{n}\left({\omega }_j\cdot q_i\right),& j\mathrm{m}\mathrm{o}\mathrm{d}2=0\\ \mathrm{c}\mathrm{o}\mathrm{s}\left({\omega }_j\cdot q_i\right),& j\mathrm{m}\mathrm{o}\mathrm{d}2=1\end{array},{\omega }_j\right.\\ & =\mathrm{e}\mathrm{x}\mathrm{p}\left(-\frac{i}{C_{\mathrm{h}}-1}\mathrm{l}\mathrm{o}\mathrm{g}\theta \right)\end{array}$$ (7)

where ${\mathbf{h}}_i\in {\mathbb{R}}^{C_{\mathrm{h}}}$ denotes the embedding vector of the atomic charge $q_i$ corresponding to the $i$-th row of the atom feature matrix $\mathbf{H}$ in the geometric graph representation $\overrightarrow{𝒢}$. Each dimension $j$ encodes the value of $q_i$ at a specific frequency ${\omega }_j$. The hyperparameter $\theta$ regulates the range of frequencies used in the embedding, thereby influencing the representational capacity of the resulting vector.

2.2.2. Message Passing for Capturing Long-Range Interactions in PBGNN.

2.2.2.1. Overview of Message Passing for Long-Range Interactions.

Based on the defined adjacency matrix $\mathbf{A}$ and the atom feature matrix $\mathbf{H}$ in the geometric graph $\overrightarrow{𝒢}$, it is essential to guide the neural network in learning long-range charge interaction patterns during optimization. Inspired by the message-passing algorithm⁵⁸ in graph neural networks, we employ a graph message-passing framework to model charge interactions within the PB context. As illustrated in Figure 1b, the graph message-passing blocks are designed to iteratively update the atomic representations of $\mathbf{H}$ based on the molecular geometry encoded by $\mathbf{A}$. This update process can be formally expressed as

$${\mathbf{h}}_i^{\prime }=\sum _{j=1}^{N_{\mathrm{a}}}{\mathbf{A}}_{ij}^{\circ }{\left({\sigma }_{\mathrm{f}}\left(‖{\mathbf{r}}_i-{\mathbf{r}}_j‖\right)\right)}^{\circ }\left({\mathit{W}}_{\mathrm{h}}{\mathbf{h}}_j\right)$$ (8)

where ${\mathbf{h}}_i^{\prime }$ is the updated feature embedding, $°$ denotes element-wise multiplication, and ${\mathit{W}}_{\mathrm{h}}\in {\mathbb{R}}^{C_{\mathrm{h}}\times C_{\mathrm{h}}}$ is a parametrized weight matrix that transforms the interaction message. The function ${\sigma }_f$ is a parametrized filter that modulates the strength of message passing based on the interatomic distance. The detailed formulation of ${\sigma }_f$ is provided in the subsequent paragraph.

2.2.2.2. Parameterized Filter for Interaction Strength Control.

From eq 8, it is evident that the adjacency matrix $\mathbf{A}$ imposes a hard cutoff to define the atomic neighbors. Beyond this structural constraint, the parametrized filter ${\sigma }_f$ further modulates interaction behaviors in a data-driven manner. Given the wide range of biomolecular sizes, it is crucial for ${\sigma }_{\mathrm{f}}$ to effectively capture interatomic distances across different molecular scales. Inspired by continuous convolutional filters,^{59 ${\sigma }_{\mathrm{f}}$} is designed to incorporate $N_{\mathrm{k}}$ radial basis function (RBF) values $\left\{{\phi }_i\mid i=1:N_{\mathrm{k}}\right\}$ as intermediate features:

$${\phi }_{\mathrm{k}}(\Vert {\mathbf{r}}_i-{\mathbf{r}}_j\Vert )=\mathrm{e}\mathrm{x}\mathrm{p}({\Vert d_{ij}-c_{\mathrm{k}}\Vert }^2)$$ (9)

where the k-th radial basis function ${\phi }_{\mathrm{k}}$ is centered at $c_{\mathrm{k}}$. The centers $c_{\mathrm{k}}$ are uniformly distributed from 0 to $K$ Å at intervals of 0.1 Å, ensuring that the filters effectively cover the full range of interatomic distances observed in the dataset. Consequently, ${\sigma }_{\mathrm{f}}$ outputs the corresponding message-passing strength based on the encoded radial basis features:

$$\mathit{v}={\mathbf{W}}_{\mathrm{f}}{\left[{\phi }_1,\dots ,{\phi }_{N_{\mathrm{k}}}\right]}^T$$ (10)

$${\sigma }_{\mathrm{f}}\left(\Vert {\mathbf{r}}_i-{\mathbf{r}}_j\Vert \right)=\left[\frac{\mathrm{e}\mathrm{x}\mathrm{p}\left({\mathit{v}}_1\right)}{\sum _{\mathrm{k}=1}^{N_{\mathrm{k}}}\mathrm{e}\mathrm{x}\mathrm{p}\left({\mathit{v}}_{\mathrm{k}}\right)},\dots ,\frac{\mathrm{e}\mathrm{x}\mathrm{p}({\mathit{v}}_{N_{\mathrm{k}}})}{\sum _{\mathrm{k}=1}^{N_{\mathrm{k}}}\mathrm{e}\mathrm{x}\mathrm{p}\left({\mathit{v}}_{\mathrm{k}}\right)}\right]$$ (11)

where ${\mathbf{W}}_{\mathrm{f}}\in {\mathbb{R}}^{C_{\mathrm{h}}\times N_{\mathrm{k}}}$ is a parametrized matrix that projects the radial features to the corresponding message-passing strengths $\mathit{v}\in {\mathbb{R}}^{N_{\mathrm{k}}}$, and eq 11 normalizes the message-passing strengths to values between 0 and 1.

2.2.2.3. Geometric Graph Representation Update.

Finally, following eq 8, the atomic representations in the geometric graph are updated as

$$\mathbf{H}\leftarrow \mathbf{H}+\left[\begin{array}{c}{\mathbf{h}}_1^{\prime }\\ ⋮\\ {\mathbf{h}}_{N_{\mathrm{a}}}^{\prime }\end{array}\right]$$ (12)

Here, $[⋮]$ denotes the vertical (rowwise) concatenation operator.

The graph message passing procedure is iteratively applied to progressively capture long-range atomic interactions. In our experiments, performing five iterations of this procedure yields promising results in approximating the PB model.

2.2.3. Stabilize Atomic Reaction Field Potential Optimization Process.

After modeling atomic interactions through multiple iterations of the graph message-passing procedure in eq 8, the final atomic features $\mathbf{H}$ are utilized to optimize the prediction of atomic reaction field potentials:

$$\stackrel{\mathbf{ˆ}}{\mathbf{y}}={\mathbf{H}\mathbf{W}}_{\mathrm{o}}^{\mathrm{T}}$$ (13)

where ${\mathbf{W}}_{\mathrm{o}}\in {\mathbb{R}}^{1\times C_{\mathrm{h}}}$ is a parameter matrix that projects the atomic feature vector $\mathbf{h}$ in $C_{\mathrm{h}}$-dimensional space into the corresponding atomic potential, denoted as $\stackrel{\mathbf{ˆ}}{\mathbf{y}}=\left\{{\hat{y}}_i\mid i=1:N_{\mathrm{a}}\right\}$. A straightforward optimization objective is to minimize the discrepancy between the predicted potential $\left\{{\hat{y}}_i\mid i=1:N_{\mathrm{a}}\right\}$ and the ground truth potential $\left\{{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)\mid i=1:N_{\mathrm{a}}\right\}$ using the mean squared error (MSE) loss:

$${ℒ}_{\mathrm{M}\mathrm{S}\mathrm{E}}=\frac{1}{N_{\mathrm{a}}}\sum _{i=1}^{N_{\mathrm{a}}}{\left({\hat{y}}_i-{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)\right)}^2$$ (14)

However, in practice, this objective function fails to converge. This is attributed to the large variance in reaction field potential ${\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left(\mathbf{r}\right)$, which induces biased gradient estimates and leads to convergence failure during stochastic optimization. To address this issue, we propose a charge-conditioned regression objective, referred to as the Charge-weighted Mean Squared Error (CMSE), which demonstrates improved convergence behavior and effectiveness in optimizing the PBGNN:

$${ℒ}_{\mathrm{C}\mathrm{M}\mathrm{S}\mathrm{E}}=\frac{1}{N_a}\sum _{i=1}^{N_a}\left|q_i\right|{\left({\hat{y}}_i-{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)\right)}^2$$ (15)

where each term is weighted by the absolute value of the corresponding atomic charge. This design is motivated by the observation that atomic electrostatic potentials are strongly correlated with the magnitudes of the atomic charges. In the CMSE formulation, the gradient of ${ℒ}_{\mathrm{C}\mathrm{M}\mathrm{S}\mathrm{E}}$ with respect to the parameters $\mathbf{\Theta }=\left\{{\mathit{W}}_{\mathrm{b}},{\mathbf{W}}_{\mathrm{f}},{\mathbf{W}}_{\mathrm{o}}\right\}$ of PBGNN is derived as follows:

$$\frac{\partial {ℒ}_{\mathrm{C}\mathrm{M}\mathrm{S}\mathrm{E}}}{\partial \mathbf{\Theta }}=\frac{\partial {ℒ}_{\mathrm{C}\mathrm{M}\mathrm{S}\mathrm{E}}}{\partial {\hat{y}}_i}\times \frac{\partial {\hat{y}}_i}{\partial \mathbf{\Theta }}$$ (16)

$$=\frac{1}{N_{\mathrm{a}}}\times \frac{d}{d{\hat{y}}_i}[\left|q_i\right|{\left({\hat{y}}_i-{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)\right)}^2]\times \frac{\partial {\hat{y}}_i}{\partial \mathbf{\Theta }}$$ (17)

$$=\frac{1}{N_{\mathrm{a}}}\times 2\left|q_i\right|\left({\hat{y}}_i-{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)\right)\times \frac{\partial {\hat{y}}_i}{\partial \mathbf{\Theta }}$$ (18)

$$=\underset{\mathrm{f}\mathrm{i}\mathrm{r}\mathrm{s}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{r}\mathrm{m}}{\underbrace{\left|q_i\right|}}\underset{\mathrm{s}\mathrm{e}\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{d}\mathrm{t}\mathrm{e}\mathrm{r}\mathrm{m}}{\underbrace{\frac{2}{N_{\mathrm{a}}}\left({\hat{y}}_i-{\varphi }_{\mathrm{r}\mathrm{x}\mathrm{n}}\left({\mathbf{r}}_i\right)\right)\times \frac{\partial {\hat{y}}_i}{\partial \mathbf{\Theta }}}},$$ (19)

where the second term corresponds to the gradient of the MSE objective, while the CMSE objective introduces an additional weighting factor (the first term) based on the absolute atomic charge. This formulation implies that atoms with larger absolute charges $\left|q_i\right|$ contribute more significantly to the overall gradient, thereby aligning with the intended emphasis on high-charge atoms during training. In contrast, the standard MSE objective treats all atoms uniformly, regardless of the substantial variance in their electrostatic potentials, which often leads to unstable optimization. In our experiments, the proposed CMSE objective consistently enabled a more stable training process and demonstrated robust convergence.

Algorithm 1. Training Procedure of PBGNN for Approximating the PB Model.

Let $𝒟$ be the training dataset, consisting of $L$ molecules $\{({ℳ}^{(1)},{\varphi }_{rxn}^{\left(1\right)}(\mathbf{r})),\dots ,({ℳ}^{\left(L\right)},{\varphi }_{rxn}^{\left(L\right)}(\mathbf{r}))\}$.

Let $\mathbf{\Theta }=\{{\mathit{W}}_h,{\mathbf{W}}_f,{\mathbf{W}}_o\}$ be the set of trainable parameters of PBGNN.

Upon completion of training, the PB electrostatic solvation free energy predicted by PBGNN is computed as

$$\mathrm{\Delta }\hat{G}=\frac{1}{2}\sum _{i=1}^{N_{\mathrm{a}}}{q}_i{\hat{y}}_i$$ (20)

where $q_i$ denotes the atomic charge and ${\hat{y}}_i$ is the predicted atomic reaction field potential at atom $i$. For clarity, a detailed framework for optimizing the parameter set of PBGNN is provided in Algorithm 1.

2.3. Efficient Implementation for Large-Scale Biomolecular Systems.

As discussed above, a key computational challenge in PBGNN arises from constructing the adjacency matrix $\mathbf{A}$ for large-scale biomolecular systems. This process requires computing pairwise distances, as defined in eq 6, across all $\left(\genfrac{}{}{0pt}{}{N_{\mathrm{a}}}{2}\right)$ atom pairs to identify neighbors within a predefined cutoff distance $K$ for each atom. For large systems such as biomolecules with over 10,000 atoms, this becomes computationally intractable.

To accelerate neighbor list construction, an efficient approach is to leverage GPU-based computation. However, the limited memory capacity of GPUs significantly hinders the scalability of PBGNN for both training and inference on large-scale systems. In practice, out-of-memory (OOM) errors frequently occur when constructing neighbor lists for biomolecules containing more than 5000 atoms, particularly when the cutoff distance $K$ exceeds 10 Å. To address this bottleneck, we adopt an approximate nearest neighbor (ANN) method, inspired by the algorithm proposed by Maneewongvatana and Mount⁶⁰ and Malkov and Yashunin,⁶¹ which offers a more efficient alternative to the standard GPU-based exact nearest neighbor search. With ANN-based neighbor list construction, PBGNN can scale effectively to biomolecular systems with over 10,000 atoms, as demonstrated in our experiments.

In addition to efficient neighbor search, batchwise stochastic gradient optimization is essential for stable and scalable neural network training, as it reduces the gradient variance. However, packing a batch of geometric graphs $\overrightarrow{𝒢}$ corresponding to biomolecules poses a challenge due to variable atom counts and neighbor list sizes. A common approach is to pad each graph to match the largest one in the batch, but this leads to excessive memory overhead, often exceeding GPU limits. To overcome this, we implement a memory-efficient batch representation that avoids padding. Instead, we used an index vector to identify individual biomolecules within the concatenated data structures. This strategy significantly improves memory efficiency and enables parallel training across biomolecular systems of varying sizes, thereby enhancing the scalability of the PBGNN.

3. RESULTS AND DISCUSSIONS

3.1. Experimental Setup.

3.1.1. Dataset Curation.

To assess the performance of the proposed PBGNN in predicting PB free energy, we curated biomolecules from diverse systems included in the AMBER PBSA benchmark suite.⁶² This benchmark spans a wide range of biomolecular structures, including proteins, nucleic acids, and protein complexes. Specifically, the PBSA dataset comprises 570 proteins and 353 nucleic acid structures, selected for their distinct functional characteristics and greater structural rigidity relative to proteins. In addition, it includes 299 protein complex structures, which are substantially larger than individual proteins, to evaluate the model’s scalability. These datasets were selected to encompass a range of molecular scales, functional characteristics, and structural flexibilities distinct from typical proteins. The inclusion of systems varying in size, from small proteins comprising 268 atoms to large complexes containing up to 1,42,324 atoms, offers a stringent evaluation of PBGNN’s scalability and robustness across structurally heterogeneous biomolecular environments.

To further evaluate the generalization capability of PBGNN across different molecular scales, we additionally compiled a dataset of relatively small molecules (PBSMALL), consisting of 812 entries, to examine the model’s performance in predicting PB free energy in this regime. The PBSMALL dataset spans molecules ranging from 3 to 62 atoms, representing a regime where quantum mechanical effects and fine-grained charge interactions are particularly pronounced. This dataset was curated to reflect the critical role of small molecules in drug design, where accurate electrostatic modeling is essential for tasks such as ligand docking, lead optimization, and virtual screening. Details of both benchmarks are provided in the Supporting Information (Section S1).

Regarding the energy calculations, the PB energy for each molecule was computed by using the numerical solver provided in the AMBER/PBSA program, configured with the default atomic cavity radii specified in the topology files and a solvent probe radius of 1.4 Å. A uniform grid spacing of 0.35 Å was applied across all molecules, with all other simulation parameters aligned with those used in the AMBER 23 PBSA module.^63,64 The GB energy was calculated using the solver in the PMEMD program,^62–64 with the PBRadii parameter set to mbondi. Three versions of the igb parameter (1, 2, and 5) were evaluated.^22,65–67 These specific igb values were chosen because they are compatible with the mbondi radii set, which allows atomic radii smaller than 1 Å. Since the energy calculation is influenced by PBRadii, and the PBGNN model is trained using the mbondi radii, evaluating GB models under the same radii configuration ensures consistency and enables a fair comparison across methods. The detailed configurations for both PB and GB calculations are provided in Section S2.

We adopted an 80/20 split for training and test sets in both datasets, ensuring that no test-set biomolecules were included in the training phase. All reported results are evaluated on the held-out test set. The implementation details, including hyperparameter configurations and training procedures, are summarized in Table 1 and Section S4.

Table 1.

Detailed Training Setup of PBGNN

hyperparameter	value
optimizer	Adam
learning rate $\gamma$	0.0005
batch size $b$	4
cutoff threshold $K$	15 (PBSA), 30 (PBSMALL)
maximum training iterations $T$	32,000
loss function	${ℒ}_{\mathrm{C}\mathrm{M}\mathrm{S}\mathrm{E}}$
atom feature hidden dimension $C_{\mathrm{h}}$	128
number of radial basis functions $N_{\mathrm{k}}$	128
frequencies used in the embedding $\theta$	10,000
message passing iteration number $T_{\mathrm{m}\mathrm{s}\mathrm{g}}$	5

3.1.2. Evaluation Metrics.

To quantify the accuracy of surface modeling, we employed three evaluation metrics: the coefficient of determination ($R^2$), the mean absolute error (MAE), and the mean absolute percentage error (MAPE). Specifically, $R^2$ measures the proportion of variance in the ground truth PB electrostatic solvation free energy $\mathrm{\Delta }G$ that is explained by the predicted values $\mathrm{\Delta }\hat{G}$:

$$R^2=1-\frac{\sum _{i=1}^n{(\mathrm{\Delta }{G}_i-\mathrm{\Delta }{\hat{G}}_i)}^2}{\sum _{i=1}^n{\left(\mathrm{\Delta }{G}_i-\mathrm{\Delta }\bar{G}\right)}^2}$$ (21)

where $n$ is the number of molecules evaluated and $\mathrm{\Delta }\bar{G}$ is the mean of the ground truth values. The MAE quantifies the average absolute deviation between the predicted and ground truth PB free energy values:

$$\mathrm{M}\mathrm{A}\mathrm{E}=\frac{1}{n}\sum _{i=1}^n|\mathrm{\Delta }{\hat{G}}_i-\mathrm{\Delta }{G}_i|$$ (22)

The MAPE measures the average relative error, expressed as a percentage, between the predicted and ground truth values:

$$\mathrm{M}\mathrm{A}\mathrm{P}\mathrm{E}=100\times \frac{1}{n}\sum _{i=1}^n\left|\frac{\mathrm{\Delta }{\hat{G}}_i-\mathrm{\Delta }{G}_i}{\mathrm{\Delta }{G}_i}\right|$$ (23)

A detailed description of each metric is provided in the Supporting Information (Section S1.3).

3.2. Accuracy Analysis of PBGNN in Predicting PB Free Energy.

In this experiment, we evaluated the performance of PBGNN in predicting PB free energy across two benchmark datasets: AMBER PBSA and PBSMALL. Our evaluation focused on two main aspects: (1) whether PBGNN, based on the geometric graph architecture, achieves high predictive accuracy compared to the original PB solver; and (2) how its performance compares to that of a faster, approximate GB model. We employed three standard regression metrics to quantify prediction accuracy: the $R^2$, MAE, and MAPE. These metrics collectively assess both the absolute and relative accuracy of the predicted PB free energies, providing a comprehensive evaluation of the model’s performance.

For completeness, we also clarify the relationship between PBGNN and the recently proposed PBML model.²⁵ Unlike PBGNN, which provides an end-to-end prediction of PB free energies directly from molecular structures, PBML is formulated as a correction model that learns the residual between PB and GB energies. As a result, PBML requires GB energies as input, and its inference cost necessarily includes both the GB calculation and the PBML correction step. This design makes PBML fundamentally different from PBGNN, and prevents a one-to-one comparison in terms of accuracy and runtime within our current experimental setting. For this reason, we did not include PBML in our experimental evaluation. We highlight this distinction to help readers better understand the methodological positioning of PBGNN relative to existing ML-based approaches.

3.2.1. PBGNN Accurately Predicts PB Free Energy across Biomolecules of Various Sizes.

We first conducted experiments to evaluate the performance of PBGNN in predicting PB free energy using the AMBER PBSA benchmark. Specifically, we compared the predicted energies from PBGNN with those obtained using a traditional numerical PB solver across three categories of biomolecules. The results are presented in Figure 3a. As shown in the figure, PBGNN demonstrates strong agreement with the numerical PB solver for the energy prediction task. In particular, the regression analysis yields a slope of 0.98 and an $R^2$ of approximately 0.995, indicating a high predictive accuracy. This conclusion is further supported by the similarity in marginal distributions between the predicted energy values and the ground truth values from the numerical solver across the three biomolecular categories.

Figure 3.

Scatter plots illustrate the predicted PB free energies (kcal/mol) from PBGNN versus the ground truth values for both the AMBER PBSA benchmark and the curated PBSMALL dataset. (a) Accuracy on the AMBER PBSA benchmark: This scatter plot displays the correlation between predicted and reference energy values across three molecular categories: nucleic acids (purple), proteins (red), and protein complexes (green). The linear regression fit demonstrates strong agreement with an overall slope of 0.98 and an $R^2$ of 0.995. (b) Accuracy on the PBSMALL dataset: This plot shows the relationship between predicted and true energy values, highlighting both consistency and variability in the predictions. The linear fit yields a slope of 0.94 and an $R^2$ of 0.968, indicating robust performance on this diverse, curated dataset.

A detailed accuracy analysis was conducted for nucleic acid, protein, and protein complex systems on the AMBER PBSA benchmark, as presented in Figure 4. These three categories span a wide range of molecular sizes, from small nucleic acids (with a minimum of 268 atoms) to large protein complexes (up to 1,40,000 atoms), thereby providing a comprehensive spectrum of molecular geometries and scales to evaluate the robustness of PBGNN. The results indicate that PBGNN consistently demonstrates strong predictive performance across diverse biomolecular sizes and energy ranges. In particular, PBGNN achieves accurate PB free energy predictions on unseen molecules with $R^2$ values exceeding 0.98 across all three molecular categories. However, a slight performance degradation is observed in larger systems. For instance, the slope of the regression line for protein complexes is 0.85, compared to 0.98 for nucleic acids and 0.94 for proteins. This decline is primarily attributed to the limited cutoff distance $K$ used in PBGNN ($K=15$ for AMBER PBSA), which constrains the model’s ability to fully capture long-range atomic interactions in large biomolecular systems, thereby slightly reducing predictive accuracy. Moreover, PBGNN exhibits a strong generalization across varying energy scales. The three systems exhibit markedly different energy profiles, with proteins ranging from 0 to −5000 kcal/mol and nucleic acids extending from 0 to −60,000 kcal/mol. In summary, these experiments highlight the robustness and accuracy of PBGNN across a diverse range of biomolecular sizes and electrostatic energy scales.

Figure 4.

A detailed accuracy analysis on the AMBER PBSA benchmark was conducted for each biomolecular system, including nucleic acids, proteins, and protein complexes. The regression analysis results, including the slope and $R^2$, are annotated in each corresponding plot. All energy values are reported in kcal/mol.

As the PB free energy plays a critical role in drug discovery, particularly in estimating binding affinities where small ligands are key components, we further evaluated the accuracy of PBGNN on the PBSMALL benchmark. PBSMALL is specifically designed for small molecules, with all entries consisting of fewer than 65 atoms. The evaluation results are shown in Figure 3b. PBGNN demonstrates strong predictive performance even for small molecules, exhibiting a high correlation with PB energies computed using the numerical PB solver. Specifically, the regression analysis yields an $R^2$ of 0.968 and a slope of 0.94, confirming the accuracy of PBGNN in predicting PB free energies for small-molecule systems.

3.2.2. PBGNN Achieves Improved Accuracy Compared to the GB Model.

The GB model is an efficient and widely used method designed to accelerate free energy calculations while maintaining reasonable accuracy, making it a strong baseline for comparison with PBGNN. To assess their relative performance, we conducted an accuracy evaluation between PBGNN and the GB model, considering three versions of the igb parameter (1, 2, and 5), in approximating PB free energies on both the AMBER PBSA and PBSMALL benchmark datasets. Accuracy was evaluated using the MAPE, as illustrated by the molecule-wise scatter plots in Figures 5 and 6. Summary statistics of the results are presented in Table 2.

Figure 5.

Accuracy of PB free energy predictions by PBGNN and the GB model on the AMBER PBSA benchmark was evaluated using MAPE (%). Both the MAPE values and the atom count distributions are annotated along the respective axes.

Figure 6.

Accuracy of PB free energy predictions by PBGNN and the GB model on the PBSMALL benchmark was evaluated using MAPE (%). Both the MAPE values and the atom count distributions are annotated along the respective axes.

Table 2.

Statistical Comparison of PB Free Energy Prediction Errors between the GB Model and PBGNN on the AMBER PBSA and PBSMALL Benchmarks, Measured by MAPE^a

method	Q1 (%)	median (%)	Q3 (%)	$\mu \left(\mathrm{\%}\right)$	$\sigma \left(\mathrm{\%}\right)$
AMBER PBSA
GB (igb = 1)	1.94	3.18	5.65	4.19	3.21
GB (igb = 2)	2.35	3.40	5.17	3.97	2.32
GB (igb = 5)	2.09	3.54	4.66	3.57	1.98
PBGNN	1.10	2.59	5.17	3.82	3.71
PBSMALL
GB (igb = 1)	2.55	5.45	12.96	13.68	28.97
GB (igb = 2)	4.69	8.20	13.63	17.43	35.42
GB (igb = 5)	2.73	6.31	15.41	17.43	33.73
PBGNN	2.18	4.62	8.57	8.50	12.55

^aReported metrics include the first quartile (Q1), median, third quartile (Q3), mean ($\mu$), and standard deviation ($\sigma$).

On the AMBER PBSA benchmark, as shown in Figure 5, PBGNN achieved a lower median MAPE of 2.59%, compared to 3.18% for the GB model with igb = 1, 3.40% with igb = 2, and 3.54% with igb = 5, indicating improved consistency. Across the GB models from igb = 1 to igb = 5, we observed a steady improvement in performance, with the mean MAPE decreasing from 4.19% to 3.57%. The mean MAPE of PBGNN (3.82%) was also slightly lower than that of the GB model with igb = 1 (4.19%) and igb = 2 (3.97%), while remaining comparable to the GB model with igb = 5 (3.57%). Notably, as shown in Table 2, PBGNN maintained competitive error dispersion, with a first quartile (Q1) value of 1.10% and a third quartile (Q3) value of 5.17%, compared to Q1 of 1.94% and Q3 of 5.65% for the GB model with igb = 1, although its standard deviation was slightly higher (3.71% vs 3.21%). In terms of mean and standard deviation of the MAPE, GB (igb = 5) achieved the lowest mean percentage error and exhibited reduced variability, reflecting progressive improvements across the GB series. Nevertheless, PBGNN attains comparable performance to GB (igb = 5), with a lower Q1 and median percentage error, despite a slightly higher mean and standard deviation. This discrepancy may be attributed to PBGNN’s increased sensitivity to molecular size, which presents an interesting direction for future investigation. In summary, these results collectively demonstrate a modest improvement in both the accuracy and variability of the PB energy approximation by PBGNN across diverse biomolecular systems.

Meanwhile, on the PBSMALL benchmark, which presents a more challenging setting due to the small molecular sizes, the GB model with igb = 1 performs best among the three GB variants. It achieves a lower mean MAPE of 13.68% and a median MAPE of 5.45%, compared to a mean of 17.43% and a median of 6.31% for igb = 5. PBGNN again outperformed the GB model with igb = 1. It achieved a median MAPE of 4.62% versus 5.45% for GB with igb = 1, and a significantly lower mean MAPE of 8.50% compared to 13.68% for GB with igb = 1 (Figure 6). Additionally, the standard deviation of PBGNN’s error was substantially reduced (12.55% vs 28.97%). Collectively, these results demonstrate the enhanced robustness and generalization capability of PBGNN in small-molecule systems.

Overall, the experimental results indicate that PBGNN achieves reliable and precise PB free energy predictions for both large biomolecular complexes and small-molecule datasets. Its robust performance, especially in terms of median and mean MAPE on PBSMALL, demonstrates the model’s strong generalizability and its potential utility in drug discovery tasks that require accurate small-molecule electrostatics modeling.

3.3. Efficiency Analysis of PBGNN Compared with Numerical Solvers.

The numerical PB solver provides highly accurate electrostatic energy estimates but suffers from substantial computational overhead, particularly for large molecular systems. The GB model was developed as a faster alternative to approximate the PB energies at reduced computational cost. Both PB and GB solvers support CPU and GPU implementations; however, it remains unclear how their runtime performance scales across different molecular sizes and hardware configurations. Furthermore, although PBGNN has demonstrated strong accuracy and a high correlation with the PB solver, its runtime efficiency relative to that of these traditional methods has not been systematically assessed.

To address these questions, we conducted a comprehensive runtime evaluation on the AMBER PBSA benchmark using both CPU and GPU platforms. We recorded the elapsed time required for electrostatic energy computations across all molecules in the benchmark dataset using seven solver configurations: PB-CPU, PB-GPU, GB-SANDER-CPU, GB-PMEMD-CPU, GB-PMEMD-GPU, PBGNN-CPU, and PBGNN. The PB solver corresponds to the implementation in the AMBER/PBSA software,^{27,28,62,63,68} while the GB solvers include both the AMBER/SANDER^{22,62,63,65–67} and the highly optimized PMEMD implementations.^62,63,69,70 PBGNN is a neural network-based model designed for efficient execution on GPUs; however, to enable a comprehensive comparison, we also report its runtime under CPU-only execution (PBGNN-CPU). The number of atoms in the benchmark spans a wide range, from small molecules (<1000 atoms) to large biomolecular complexes (>15,000 atoms), enabling a thorough scalability analysis. For each molecule, we measured the elapsed time required to compute the electrostatic solvation energy. All CPU-based timing experiments were performed on a multi-socket machine featuring Intel(R) Xeon(R) Gold 6230 processors running at 2.10 GHz, comprising four sockets, each containing 20 physical cores. GPU-based experiments were conducted on a system equipped with an NVIDIA GeForce RTX 4090 GPU with 24 GB of memory.

As shown in Figure 7, PB-CPU exhibits the longest runtime across all molecule sizes, reaching up to 1 × 10⁴ s for the largest systems. PB-GPU reduces this runtime by roughly one to two orders of magnitude but remains computationally intensive for large molecules. GB methods significantly reduce runtime, with GB-SANDER-CPU, GB-PMEMD-CPU, and GB-PMEMD-GPU typically completing in under 10 s. In contrast, PBGNN achieves sub-second inference times for almost all molecules in the dataset. The regression curves in Figure 7b further underscore the scalability advantage of PBGNN. While the runtime of GB-SANDER-CPU shows quadratic growth with respect to atom count, both PBGNN and GB-PMEMD-GPU maintain nearly constant runtime profiles, indicative of linear or sublinear scaling behavior. Notably, GB-PMEMD-GPU relies on manually optimized CUDA kernels,^62,63,69,70 which require substantial development time and specialized expertise. In contrast, PBGNN achieves efficient inference without the need for platform-specific manual optimization, enabling a more accessible and flexible deployment for researchers without compromising runtime performance. For systems with more than 10,000 atoms, PBGNN demonstrates over 1,000 times speedup relative to PB-CPU and approximately 5 times speedup relative to GB-SANDER-CPU.

Figure 7.

Runtime performance comparison across PB, GB, and PBGNN. (a) Scatter plot showing the elapsed time (in seconds) versus the number of atoms for different solvers across CPU and GPU platforms. Each point represents a molecule from the benchmark dataset. Marginal distributions of atom counts (top) and elapsed times (right) are shown for each method. (b) Smoothed regression curves of runtime versus atom counts for GB-SANDER-CPU, GB-PMEMD-GPU, and PBGNN. Shaded areas represent the 95% confidence interval. PBGNN demonstrates substantially lower runtime, especially for large molecules, without requiring manual CUDA kernel optimization owing to its efficient neural approximation.

To further quantify these gains, we computed the average speedup across three biomolecular categories, namely, nucleic acids, proteins, and protein complexes, relative to the PB-CPU baseline (Figure 8a). Notably, the inference performance of PBGNN-CPU is comparable to that of PB-GPU, underscoring the effectiveness of multithreaded CPU computation for neural network inference. On average, PBGNN achieves a speedup exceeding 10⁴ times compared to PB-CPU across all categories. Additionally, Figure 8b shows that PBGNN outperforms GB-SANDER-CPU by an average factor of 4.74 times, with the largest improvements observed in protein complexes. We also observed that GB-PMEMD-GPU achieves substantial speedup on large protein complexes, whereas the acceleration on small nucleic acids is marginal. This discrepancy arises from the trade-off between data transfer time and GPU execution time. For small systems, the overhead associated with transferring data between the host (CPU) and the device (GPU) dominates the total runtime, effectively diminishing the benefits of parallel execution on the GPU. In contrast, for larger systems, the computational workload is sufficient to amortize the data transfer overhead, allowing the GPU’s parallelism to significantly accelerate the calculations. Notably, PBGNN exhibits a more consistent balance across different molecular sizes, maintaining a high efficiency without incurring significant data transfer overhead. These results underscore PBGNN’s capability to deliver real-time inference without compromising prediction accuracy.

Figure 8.

Speedup comparison across different solvers and molecular categories. (a) Bar plot showing the speedup of PB-GPU, GB-SANDER-CPU, GB-PMEMD-CPU, GB-PMEMD-GPU, PBGNN-CPU, and PBGNN relative to PB-CPU across three biomolecular categories: nucleic acids, proteins, and protein complexes. PBGNN achieves a speedup of over 10⁴ times relative to PB-CPU on average. (b) Zoomed-in comparison of speedup over the GB-SANDER-CPU baseline. The dashed line indicates the average speedup of PBGNN over GB-SANDER-CPU (Mean: 4.74×). PBGNN consistently outperforms both GB-SANDER-CPU and GB-PMEMD-GPU baselines across all categories, demonstrating its runtime efficiency.

In conclusion, our experiments demonstrate that PBGNN not only matches the numerical PB solver in accuracy but also delivers a dramatically faster runtime. Its superior scalability and low latency establish it as a promising tool for high-throughput and large-scale biomolecular modeling tasks.

3.4. Geometric Graph Representation Analysis of PBGNN.

PBGNN demonstrates a comparable ability to predict PB free energy to the numerical PB solver while offering a significantly more efficient approach to energy calculation tasks, particularly for large molecular systems. As a data-driven PB energy solver, a major advantage of PBGNN is its ability to substantially reduce reliance on domain-specific expertise for efficient PB energy computation. However, it relies on robust data-driven training and evaluation processes to achieve accurate predictions. Therefore, it is essential to investigate the key design properties of PBGNN to gain a deeper understanding of how the neural network performs the PB energy prediction.

3.4.1. Importance of Geometric Representations in the PB Free Energy Prediction.

The accurate modeling of electrostatic solvation free energy using PB theory inherently requires a consideration of complex atomic interactions governed by molecular geometry. To assess whether explicitly incorporating geometric information via graph-based neural networks improves prediction accuracy, we compared PBGNN, which leverages molecular geometry through a graph neural network, with a 3D convolutional neural network (ConvNet) baseline³⁷ that does not use a graph structure. While the ConvNet captures spatial context through voxelized molecular grids, it lacks the inductive bias required to model atom-wise interactions critical for PB energy computation. The detailed architecture of the ConvNet employed in this experiment is provided in Section S3. Notably, PBGNN achieves superior predictive performance with a significantly smaller parameter count (0.5 M vs 2.0 M), highlighting its architectural efficiency. We evaluated both models on two datasets: the AMBER PBSA benchmark and the curated PBSMALL dataset. The performance was assessed using three standard regression metrics: $R^2$, MAE, and MAPE. As shown in Table 3, PBGNN consistently outperforms ConvNet across all metrics and datasets. On the AMBER PBSA benchmark, PBGNN achieves an $R^2$ of 0.995, substantially higher than ConvNet’s 0.931, while reducing the MAE from 481.947 to 151.54 kcal/mol and the MAPE from 10.268% to 3.822%. Similarly, on the PBSMALL dataset, PBGNN obtains an $R^2$ of 0.968, outperforming ConvNet’s 0.915. The MAE is reduced from 1.204 to 0.475 kcal/mol, and the MAPE drops from 22.187% to 8.503%.

Table 3.

Performance Comparison of PB Free Energy Prediction with and without Geometric Representations, Evaluated Using R², MAE (kcal/mol), and MAPE (%) Metrics^a

method	# parameters	R2 (↑)	MAE (↓)	MAPE (↓)
AMBER PBSA
ConvNet	2.0M	0.931	481.947	10.268
PBGNN	0.5M	0.995	151.54	3.822
PBSMALL
ConvNet	2.0M	0.915	1.204	22.187
PBGNN	0.5M	0.968	0.475	8.503

^aResults are reported for two benchmark datasets, AMBER PBSA and PBSMALL. PBGNN, which encodes explicit geometric inputs, consistently achieves superior performance across all metrics and datasets compared to the non-geometry-based ConvNet baseline.

These results confirm that the incorporation of geometric graph representations significantly enhances the model performance in PB free energy prediction. Despite having 4 times fewer parameters, PBGNN achieves higher accuracy by more effectively modeling interatomic relationships relevant to the PB equation. This finding highlights the critical role of geometry-aware architectures in developing physically grounded and data-efficient neural predictors for biomolecular applications.

3.4.2. Effect of Optimization Objective on PB Free Energy Prediction.

The choice of training objective plays a critical role in determining the accuracy and generalization of neural predictors. While MSE is a widely used objective for regression tasks, it treats all prediction errors equally, regardless of their physical or structural context. In contrast, we propose the CMSE loss that incorporates domain-specific knowledge relevant to PB free energy modeling. This experiment aims to evaluate whether training with CMSE can improve predictive performance compared to that of the standard MSE.

We trained both PBGNN and ConvNet using MSE and CMSE on two benchmark datasets: AMBER PBSA and PBSMALL. Model performance was assessed using the $R^2$, MAE, and MAPE. As summarized in Figure 9, models trained with CMSE consistently outperform their MSE-trained counterparts across all evaluation metrics. On the AMBER PBSA benchmark, CMSE improves $R^2$ from 0.740 to 0.931 for ConvNet and from 0.878 to 0.995 for PBGNN. Correspondingly, MAE drops from 1253.364 to 481.947 kcal/mol (ConvNet) and from 964.613 to 151.540 kcal/mol (PBGNN), while MAPE is reduced from 34.565% to 10.268% (ConvNet) and from 29.500% to 3.822% (PBGNN). Similar trends are observed on the PBSMALL dataset, where CMSE increases $R^2$ from 0.763 to 0.915 (ConvNet) and from 0.884 to 0.968 (PBGNN), while substantially reducing MAE and MAPE (MAE drops from 1.204 and 0.475 kcal/mol; MAPE drops from 22.187% and 8.503% for PBGNN, respectively).

Figure 9.

Performance Comparison of PBGNN and ConvNet using different optimization objectives. Bar plots show the performance of models trained with MSE and the proposed CMSE on two benchmarks: AMBER PBSA (top row) and PBSMALL (bottom row). Evaluation metrics include $R^2$, MAE (kcal/mol), and MAPE (%). Models trained with CMSE consistently outperform those trained with MSE across all evaluation metrics, indicating that CMSE improves the prediction accuracy.

These results confirm that the CMSE objective leads to significant improvements in predictive accuracy and robustness across both models and datasets. By incorporating structural and physical context into the loss formulation, CMSE guides the model to focus on energetically meaningful discrepancies, resulting in better alignment with ground truth PB free energies. Therefore, CMSE serves as a principled and effective training objective for electrostatic energy prediction tasks.

3.4.3. Effect of Cutoff Distance on Prediction Accuracy.

In geometric graph representations, the choice of cutoff distance, $K$ determines the spatial range of interactions considered during message passing. In the context of PB free energy prediction, electrostatic interactions are inherently long-range, and a limited cutoff may hinder the model’s ability to accurately capture relevant dependencies accurately. This experiment aims to investigate the impact of varying the cutoff distance on the prediction performance of PBGNN. It is important to note that the cutoff distance cannot be increased arbitrarily due to constraints on GPU memory.

We trained PBGNN using different cutoff distances on two benchmark datasets: AMBER PBSA and PBSMALL. Changing the cutoff alters the number of neighboring atoms considered for each node during graph construction. We evaluated model performance using $R^2$ and MAPE. For both datasets, we selected the largest cutoff distances feasible on a consumer-grade GPU (NVIDIA RTX A6000 GPU) with 48 GB memory, resulting in $K=15$ for the AMBER PBSA benchmark and $K=30$ for the PBSMALL benchmark.

As shown in Figure 10, increasing the cutoff distance consistently improves the model performance. On the AMBER PBSA dataset, MAPE decreases from approximately 16% at a cutoff of 4 Å to below 5% at 15 Å, while $R^2$ increases from 0.80 to 0.99. A similar trend is observed for the PBSMALL dataset, where $R^2$ increases from 0.85 to 0.96 and MAPE decreases from approximately 20% to 9% as the cutoff grows from 5 to 30 Å. These results demonstrate that incorporating longer-range atomic interactions is essential for accurate electrostatic energy modeling. The improvement in both MAPE and $R^2$ with increasing cutoff distances highlights the importance of capturing long-range dependencies in PB-based energy estimation and suggests that careful tuning of this parameter can enhance the predictive capability of PBGNN without altering the model architecture.

Figure 10.

Impact of cutoff distance on prediction accuracy for PBGNN across two benchmark datasets: AMBER PBSA and PBSMALL. The left column presents the $R^2$, and the right column displays the MAPE (%) metric. As the cutoff distance increases, $R^2$ increases and MAPE decreases for both datasets, indicating improved predictive performance with longer-range interactions.

4. CONCLUSION

This study presents PBGNN, a data-driven framework that approximates PB electrostatic solvation free energies with high accuracy and computational efficiency. Leveraging a geometric graph neural network architecture and a charge-weighted mean squared error loss function for atom-wise potential training, PBGNN is designed to capture long-range electrostatic interactions in PB energy prediction across a broad range of molecular systems, from small drug-like molecules to large biomolecular complexes.

Comprehensive evaluations on the AMBER PBSA benchmark and the curated PBSMALL dataset demonstrate that PBGNN achieves predictive accuracy comparable to the numerical PB solver, with $R^2$ values exceeding 0.98 across all major biomolecular categories. It also outperforms widely used GB models in terms of median and mean MAPE, particularly on small-molecule systems, where accurate electrostatics is critical for drug discovery applications. Compared with 3D convolutional networks, PBGNN exhibits superior accuracy with a significantly smaller parameter count, underscoring the efficiency and expressiveness of its geometric graph-based design.

From a performance standpoint, PBGNN delivers substantial speedups compared to traditional solvers. It surpasses GB-SANDER-CPU by an average factor of 4.7 times and achieves up to 10,000 times speedup over PB-CPU, while maintaining near-linear scalability with respect to molecular size. Unlike GB-PMEMD-GPU, which requires manually optimized CUDA kernels, PBGNN achieves high inference efficiency without platform-specific engineering, enabling easy deployment across diverse computational environments.

Ablation studies further validate the importance of architectural choices, including geometric representation, optimization objectives, and the cutoff distance. Increasing the cutoff improves performance, although practical constraints imposed by GPU memory set an upper limit on feasible values. Future work may explore dynamic or adaptive graph construction techniques to more effectively capture long-range dependences in large systems.

One limitation of this work is that PBGNN was trained only on mbondi-based PB energies. This choice was made because the PB model adopts mbondi as the default setting. Accordingly, we trained PBGNN using PB energies calculated with mbondi and compared all methods (igb = 1, 2, and 5) under the same mbondi condition to ensure consistency. To further examine PBGNN’s zero-shot performance on unseen setups such as mbondi2 and mbondi3, we reported additional results in Section S5 of the Supporting Information. PBGNN demonstrated comparable performance even under these untrained radii settings. Looking forward, extending PBGNN to support additional radii sets or developing a unified model capable of handling multiple configurations represents a promising direction for future research.

Another limitation arises from the PB method itself. While the accuracy of PB/GB models is often benchmarked against experimental hydration free energies (HFEs),²³ success in this regime does not necessarily guarantee equivalent accuracy for the electrostatic component of binding free energy. Although PB-based methods have been extensively validated and are widely applied in biomolecular binding studies, their performance in the high-accuracy regime requires careful interpretation. Prior studies have shown that even modest errors in hydration free energy predictions can translate into much larger uncertainties in estimating the electrostatic component of binding,^13,15 motivating further refinement of PB solvers.¹⁴ Related work in explicit solvation has likewise demonstrated that widely used water models such as TIP3P and TIP4P can reproduce accurate hydration free energies still yield significantly larger difference in binding electrostatics.^71–73 These observations underscore the importance of recognizing the inherent limitations of both continuum and explicit approaches when applying them to predict binding energetics at a fine resolution.

In summary, PBGNN offers an accurate, scalable, and efficient solution for PB energy prediction, bridging the gap between physical rigor and computational tractability. Its performance on both small- and large-scale systems highlights its potential for broad applicability in molecular modeling, including high-throughput virtual screening, drug design, and electrostatics-driven simulations. The open-source release of the code, datasets, and training pipeline will facilitate reproducibility and foster further research in the area of neural approximations for biomolecular electrostatics.

Supplementary Material

ASSOCIATED CONTENT

Supporting Information

The Supporting Information is available free of charge at https://pubs.acs.org/doi/10.1021/acs.jctc.5c01193.

Benchmark dataset details, evaluation metrics, convolutional neural network architecture used as a baseline, and training configurations. To enable better reproducibility, code used to train and evaluate the models is available at: https://github.com/yxwu21/PBGNN, and the preprocessed data files for the AMBER/PBSA and PBSMALL datasets are available at: 10.5281/zenodo.15867553 (PDF)