End-to-End Reaction Field Energy Modeling via Deep Learning based Voxel-to-voxel Transform

TL;DR

This work introduces PBNeF, a neural-field end-to-end solver for Poisson-Boltzmann electrostatics that operates on voxelized molecular representations and uses a neural field transformer to predict reaction-field energy. By combining a learnable voxel representation with Fourier embeddings and a learnable Gaussian interaction kernel, PBNeF captures spatial electrostatic interactions and produces EPB predictions without explicit PDE solving. On the AMBER PBSA benchmark, PBNeF achieves accuracy comparable to the Generalized Born model and delivers over 100× speedups relative to traditional PB solvers, with PBNeF-Lite offering even faster inference at a small accuracy trade-off. The approach enables real-time or large-scale biomolecular electrostatics calculations, with implications for drug design and protein engineering, while highlighting areas for improvement in atom-level energy precision.

Abstract

In computational biochemistry and biophysics, understanding the role of electrostatic interactions is crucial for elucidating the structure, dynamics, and function of biomolecules. The Poisson-Boltzmann (PB) equation is a foundational tool for modeling these interactions by describing the electrostatic potential in and around charged molecules. However, solving the PB equation presents significant computational challenges due to the complexity of biomolecular surfaces and the need to account for mobile ions. While traditional numerical methods for solving the PB equation are accurate, they are computationally expensive and scale poorly with increasing system size. To address these challenges, we introduce PBNeF, a novel machine learning approach inspired by recent advancements in neural network-based partial differential equation solvers. Our method formulates the input and boundary electrostatic conditions of the PB equation into a learnable voxel representation, enabling the use of a neural field transformer to predict the PB solution and, subsequently, the reaction field potential energy. Extensive experiments demonstrate that PBNeF achieves over a 100-fold speedup compared to traditional PB solvers, while maintaining accuracy comparable to the Generalized Born (GB) model.

Figures (6)

• Figure 1: The overview of the proposed PBNef method. The figure illustrates the workflow of the PBNef model for predicting electrostatic potential of biomolecules (EPB). The process begins with the voxel representation of a molecule, where atomic properties such as level-set values and charges are encoded into a 3D voxel grid. To address challenges in learning low-rank scalar features and capturing charge interactions, Fourier feature mapping is used to encode these features into high-frequency vectors. A Gaussian parameterized atom charge interaction model is employed to represent the interactions between atoms. The spatial interactions are understood through a neural field transformer that processes the 3D structure of the molecule. Given the potentially large size of voxel grids, a patch cropping strategy with random cropping and rotation augmentation is applied to improve model robustness. The final step involves the neural network predicting EPB values across the voxel grid, with results integrated into an atom charge map for EPB prediction. The lower portion of the figure highlights the key components, including the Gaussian parameterized interaction and high-frequency Fourier embedding, which enhance the model’s accuracy and robustness.
• Figure 2: Illustration of molecular features and high-frequency Fourier embedding. (a) Representation of atom charge and Solvent Excluded Surface (SES). Different colors correspond to different types of atoms, with charges indicated by "+" and "-" symbols. (b) Example of high-frequency Fourier embedding applied to scalar features like atom charge $q^{(i)}$.
• Figure 3: Visualization of Gaussian kernels for atom charge distribution. The figure illustrates the concept of using learnable Gaussian blur kernels, parameterized by a multivariate Gaussian distribution, to diffuse atom charge interactions to neighboring grid points in the molecular structure. The left panel shows a Gaussian kernel with a mean vector $\boldsymbol{\mu} = [0, 0]$ and a covariance matrix $\boldsymbol{\Sigma} = 1001$, resulting in an isotropic distribution. The right panel demonstrates a kernel with the same mean vector but a different covariance matrix $\boldsymbol{\Sigma} = 20.40.40.6$, leading to an anisotropic distribution.
• Figure 4: Illustration of Neural Field Transformer using U-Net like structure.
• Figure 5: (a) EPB Accuracy Comparison: Scatter plot showing the relationship between predicted and true energy values for the GB model (blue squares), PBNeF (green triangles), and PBNeF-Lite (red circles). The linear regression lines and corresponding equations illustrate the accuracy of each model, with $R^2$ and MAE values providing a quantitative measure of their performance. (b) Relative Error of Predicted EPB Values: Plot displaying the relative error per atom for each molecule across the three models, highlighting the consistency and variability in their predictions.