SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Daiwen Sun $^{1}$ , He Huang $^{2,1}$ , Yao Li $^{2,3}$ , Xinqi Gong $^{1,2*}$ , Qiwei Ye $^{2*}$
$^{1}$ Institute for Mathematical Sciences, School of Mathematics, Renmin University of China, Beijing, China
$^{2}$ Beijing Academy of Artificial Intelligence, Beijing, China
$^{3}$ School of Life Sciences, Tsinghua University, Beijing, China
sundaiwen@ruc.edu.cn, huanghe@baai.ac.cn, liyao17@tsinghua.org.cn, xinqigong@ruc.edu.cn, qwye@baai.ac.cn
We propose a novel neural network-based approach to modeling protein dynamics using an implicit representation of a protein's surface in 3D and time. Our method utilizes the zero-level set of signed distance functions (SDFs) to represent protein surfaces, enabling temporally and spatially continuous representations of protein dynamics. Our experimental results demonstrate that our model accurately captures protein dynamic trajectories and can interpolate and extrapolate in 3D and time. Importantly, this is the first study to introduce this method and successfully model large-scale protein dynamics. This approach offers a promising alternative to current methods, overcoming the limitations of first-principles-based and deep learning methods, and provides a more scalable and efficient approach to modeling protein dynamics. Additionally, our surface representation approach simplifies calculations and allows identifying movement trends and amplitudes of protein domains, making it a useful tool for protein dynamics research. Codes are available at https://github.com/Sundw-818/DSR, and we have a project webpage that shows some video results, https://sundw-818.github.io/DSR/.
Proteins, which are made up of chains of amino acids, are essential molecules in living organisms. Protein with specific three-dimensional structures are responsible for a wide range of biological processes such as enzymatic reactions, transport of molecules, and structural support. However, proteins are not static objects in vivo, but constantly in motion, with individual amino acids vibrating, rotating, and shifting in response to the change of environment. The biological functions of proteins are fulfilled by their dynamical behaviors, and researchers study the dynamics and functions of proteins by molecular dynamics(MD) simulations[1, 2].
First-principles molecular dynamics use the forces calculated from the electronic states by solving the Schrödinger equation (namely, density functional theory) to compute the coordinate trajectories[3], with programs such as VASP[4] and GPAW[5]. However, the computational expense of MD typically increases exponentially in relation to the number of electronic degrees of freedom. In order to reduce the computational cost, empirical force field (classical force field) models with simplified functional forms are employed in MD simulation, such as AMBER[6, 7], CHARMM[8, 9], and GROMOS[10-12], which are designed to capture the bonded and non-bonded interactions between atoms. Besides, machine learning techniques can be employed to fit force fields or potential energy surfaces[13-16]. These force field-based methods are much faster than ab initio methods without exhaustive QM
*Correspondence to: Xinqi Gong<xinqigong@ruc.edu.cn>, Qiwei Ye <qwye@baai.ac.cn>.
calculations. However, the integration time step in MD cannot be too large to resolve the fastest motion in the system, for example the vibration of hydrogen-involved bonds whose timescales are as short as a few femtoseconds. This restricts MD from having more speed advantages. To address this problem, coarse-grained methods are being developed to perform molecular dynamics simulations faster and on larger systems by eliminating some unimportant degrees of freedom[17-20]. Fast as it is, coarse-grained simulations are only able to reproduce molecular properties at low resolution, such as  $Rg^2$ , sacrificing the all-atom details. Given this perspective, existing molecular dynamics approaches suffer from either unaffordable computational costs or insufficient precision. Moreover, inherent limitations such as the considerable random fluctuation of simulated trajectories and the consequential low signal-to-noise ratio pose significant challenges in detecting crucial protein functional motion. Even for the current most advanced Anton 3, under the parallel processing of 512 nodes, for a large system of 1 million atoms, only 100 microseconds of MD simulation can be performed each day[21].
These aforementioned methods give thorough understanding of the natures of proteins by modeling protein molecules as discrete particles. But when studying protein-protein interaction or protein-ligand binding, knowledge about surfaces of proteins is what researchers need. To fulfill such demands, a continuous, high-level representation of molecular shape known as the molecular surface is proposed, as shown in Figure 1.Based on the the concept of "key-lock pairs" emphasizing complementarity of molecular surface[22], the molecular surface models a protein as a continuous shape with geometric and chemical features. One repre
sentative method, molecular surface interaction fingerprinting (MaSIF), is a geometric deep learning method that captures fingerprints important for specific biomolecular interactions, enabling accurate and efficient analysis of protein-ligand complexes and potential applications in drug discovery and molecular design[23, 24]. Furthermore, some works investigated the application of end-to-end learning techniques for the analysis of protein surfaces and the identification of potential functional sites, such as drug-target binding sites[25].
Previously, surface representation was mainly used for studying inter-molecular interactions or assisting molecular docking, and it had not been applied to analysing molecular dynamics simulation data. Inspired by the above works, we propose Dynamical Surface Representation (abbreviated as DSR), a method for modeling protein dynamics by representing protein surfaces with implicit neural networks. Implicit neural representations[26, 27] are a class of deep learning models that can represent complex geometric shapes using a continuous function without requiring an explicit surface representation. It's well-suited for modeling the dynamic shapes of proteins which have complex and varied conformations. Currently, deep learning applications in molecular dynamics (MD) mainly target small molecules, while our method can be applied to very large proteins.
In this paper, we explore the use of implicit neural representations for modeling the dynamic shapes of proteins. To be specific, protein surfaces are represented as a zero-level set of signed distance function (SDF) in DSR. The value of SDF is defined as the minimum distance from a point in Euclidean space to the surface of the object, where the positive value indicates that the point is outside the object, otherwise, it is inside. In addition, we introduce a time variable in DSR model to simulate the change of the protein surface over time, which is the dynamic simulation of the protein. In summary, we use the idea of implicit neural representation and use the neural network as the implicit neural representation of points on the 3D and time region, so as to achieve the purpose of modeling protein dynamic changes in continuous space and time domain.
In summary, our paper's main contribution is two-fold:
- We build an implicit neural network that learns SDF from raw point cloud data in 3D + time domain, which is a both temporally and spatially continuous protein dynamic representation.
- We have achieved dynamic modeling of large proteins by using surface representations, which overcomes the limitations of previous atomic models or coarse-grained models that were restricted to small molecules or proteins.
Figure 1: The protein surface representation. Showing the protein surface (PDB ID: 1PGA) at different time steps in a molecular dynamics trajectory.
2 DSR
In this section, we present DSR, a dynamical surface representation as implicit neural networks for protein. We first introduce the definition of SDF, and then describe how we use SDF to model proteins. The architecture of the whole method is shown in Figure 2.
2.1 SDF Representation
Let  $\Omega = [-1,1]^3$  be a spatial domain,  $\tau = [-1,1]$  a temporal domain, and  $\mathcal{M}_t$  be a 3D manifold embedded in  $\Omega$  at time  $t\in \tau$ . For any point  $\pmb {x} = (x,y,z)\in \Omega$  at time  $t$ , the  $f_{\mathcal{M}_t}:\Omega \to \mathbb{R}$  is defined as
$$
f _ {\mathcal {M} _ {t}} (\boldsymbol {x}) = \left\{ \begin{array}{l l} \mathrm {d} (\boldsymbol {x}, \mathcal {M} _ {t}) & \text {i f x o u t s i d e} \mathcal {M} _ {t} \\ 0 & \text {i f x b e l o n g i n g t o} \mathcal {M} _ {t} \\ - \mathrm {d} (\boldsymbol {x}, \mathcal {M} _ {t}) & \text {i f x i n s i d e} \mathcal {M} _ {t} \end{array} \right. \tag {1}
$$
where  $\mathrm{d}(x,\mathcal{M}_t)\coloneqq \inf_{y\in \mathcal{M}_t}\mathrm{d}(x,y),\mathrm{d}(x,y) = \| x - y\| _2$
The zero-level set, which is the surface of the protein at time  $t$ , is presented by the points where  $f_{\mathcal{M}_t}(\cdot) = 0$ . That is, the points on the protein surface at time  $t$  can be represented as set  $\{\pmb{x} \mid f_{\mathcal{M}_t}(\pmb{x}) = 0, \pmb{x} \in \Omega\}$ .
2.2 Modeling Protein Dynamical Surface with SDFs
The protein surface representation is in the form of a 3D shape. In order to utilize the implicit function form of SDF, a typical approach is to partition the space into a grid, and then compute the SDF value for each grid point. However, for irregular objects, the accurate SDF values cannot be computed and thus approximation algorithms, such as the 8SSEDT algorithm, are usually employed. This method is associated with two limitations: firstly, protein surfaces are often intricate and rugged, and estimation errors can result in larger biases after modeling; secondly, the SDF calculation for dynamic models is linearly dependent on time and cubically dependent on resolution, and pre-computed SDFs will incur high computational costs. Hence, an alternative strategy is used herein, whereby SDF is learned directly from the raw point clouds, rather than utilizing pre-computed SDF for supervised learning. This completely circumvents the need for pre-computing SDF, significantly reduces computational costs, and enhances computational efficiency. In the following, we will elaborate on how to learn SDF directly from the raw point clouds.
For a given input point cloud  $\mathcal{X} = \{\pmb{x}_i\}_{i\in I}\subset \mathbb{R}^3$  at time  $t$ , with point normal vector data (optional),  $\mathcal{N} = \{\pmb{n}_i\}_{i\in I}\subset \mathbb{R}^3$ , our objective is to find the optimal parameters  $\theta$  of an MLP  $f(\pmb {x},t;\theta)$ , where  $f:\mathbb{R}^{4 + d_z}\to \mathbb{R}$ , that can accurately estimate a signed distance function to a surface  $\mathcal{M}_t$  defined by the point cloud  $\mathcal{X}$  and normals  $\mathcal{N}$  at time  $t$ .
The form of our loss function is as follows:
$$
\ell (\theta) = \ell_ {\mathcal {X}} (\theta) + \lambda \mathbb {E} _ {\boldsymbol {x}} \left(\| \nabla_ {\boldsymbol {x}} f (\boldsymbol {x}, t; \theta) \| - 1\right) ^ {2} + \alpha \| z \| \tag {2}
$$
where  $\lambda > 0$  is a parameter,  $\| \cdot \| = \| \cdot \|_2$  is the Euclidean 2-norm, and
$$
\ell_ {\mathcal {X}} (\theta) = \frac {1}{| I |} \sum_ {i \in I} \left(| f \left(\boldsymbol {x} _ {i}, t; \theta\right) | + \tau \| \nabla_ {\boldsymbol {x}} f \left(\boldsymbol {x} _ {i}, t; \theta\right) - \boldsymbol {n} _ {i} \|\right) \tag {3}
$$
promotes the vanishing of  $f$  on  $\mathcal{X}$ , and in the presence of normal data (i.e.,  $\tau = 1$ ),  $\nabla_{\pmb{x}}f$  to the given normals  $\mathcal{N}$ .
The first term in the summation part of Eq. (3) indicates that the SDF value of the surface points should be as small as possible. The second term represents the loss between the point normal vector and the ground truth.
The second term in Eq. (2) called the Eikonal term promotes the gradients  $\nabla_{\pmb{x}}f$  to possess a 2-norm of unity, i.e.  $\| \nabla_{\pmb{x}}f\| = 1$ , which is the gradient property of SDFs. Note that this property is crucial to the loss design, allowing us to avoid using pre-computed SDF values for supervised learning. While its sufficiency has not been fully established in previous literature, we provide a more comprehensive proof in the Appendix B.
The third term  $z$  in Eq. (2) is a latent code that we introduce into the model to specify different protein types, and its parameters are learnable. We claim that with sufficient data, we can explore the latent space for greater generalization.
The global minimum of the loss in Eq. (2) will be the solution of the Eikonal partial differential equation, i.e.,
$$
\| \nabla_ {\boldsymbol {x}} f (\boldsymbol {x}) \| = 1, \tag {4}
$$
which will also be a signed distance function, where  $f$  approaches 0 on  $\mathcal{X}$ , with gradients  $\mathcal{N}$ . The calculation of the expectation is based on some probability distribution  $x \sim \mathcal{D}$  in  $\mathbb{R}^3$ , which will be introduced in Section 3.2.
Throughout the optimization process, we did not use the pre-computed SDF value to supervise the training, which is an end-to-end learning architecture from raw point clouds data to SDF value. In the next section, we will introduce the experiment setup in detail, including the datasets, implementation details, and model evaluation metrics.
Figure 2: Overview of DSR. The upper part is the training process. First, the manifold and space points are sampled at different times, added the time dimension and be input into the DSR Model to obtain SDF. And then the Gradient Module is used to calculate the derivative of SDF with respect to  $(x,t)$ . The lower part is the inference process. We first divide the space into a grid based on the desired resolution, then adds the time dimension and inputs them into the DSR Model to predict the SDF. Finally, the Marching Cube algorithm is utilized to generate different level sets, with the zero-level set representing the protein dynamical surface.
3 DSR on large-scale proteins simulation data
3.1 Dataset.
500ns Trajectory Data. This data set[28] currently contains two 100ns atomistic molecular dynamics trajectories of Abeta (40 residues), one wild type and one E22G, following the protocol of the 500ns trajectories published in [29]. The simulation was performed at  $310\mathrm{K}$  in a generalized Born implicit solvent. In this work, we used the first 2000 frames of wild type trajectories.
MDAnalysisData. MDAnalysisData collects a set of data resources pertaining to computational biophysics, primarily focusing on molecular dynamics (MD) simulations and the structural and dynamic attributes of biomolecules. We used five of these datasets,
(1) AdK equilibrium dataset[30] of 4187 frames,
(2) AdK conformation transition[31] sampled with two methods:
- Dynamic importance sampling molecular dynamics (DIMS MD) of 102 frames[32],
- Framework Rigidity Optimized Dynamics Algorithm (FRODA) of 142 frames[33],
(3) A short MD trajectory of I-FABP (intestinal fatty acid binding protein) of 500 frames[34],
(4) A trajectory of the NhaA membrane protein of 5000 frames[35],
(5) Two MD trajectories of YiiP membrane protein, the short one has 900 frames, and the long one has 9000 frames[36]. We opted for the longer one here.
For these 6 trajectories data, we use a fixed size 2000 to crop the frames due to the computational resources, and we keep the full length for those shorter than the cropping size.
GPCRmd. The GPCRmd (http://gpcrmd.org/) [37] is an online platform that incorporates web-based visualization capabilities and shares data. The GPCRmd database includes at least one representative structure from each of the four structurally characterized GPCR classes. The GPCRmd platform holds more than 600 GPCR MD simulations from GPCRmd community and individual contributions. Each system was simulated for 500 ns in three replicates (total time  $1.5\mu s$ ). We selected one system per GPCR family, and a total of 4 trajectories were used to train the model, of which trajectory ids are 10792_trj_81, 10840_trj_87, 10912_trj_95 and 15711_trj_791 respectively.
DRYAD_MD. This dataset contains the trajectory data of 23 proteins simulated by Jumper JM[38]. It is worth noting that the protein trajectories within this dataset exhibit significant oscillations, posing a challenge for model training and resulting in a high failure rate. For this reason, we screened out 14 proteins with mean  $\mathrm{RMSD}^2$  less than 1.0 for our experiments.
3.2 Implementation Details.
Problem setup. In this paper, we use PyMol[39] to obtain the protein surface representation, the points and normal vectors, to train the model. The goal is to obtain a protein surface dynamic model with respect to continuous space and time.
Architecture. For representing shapes we used level sets of MLP  $f(x, t; \theta) = \mathbb{R}^4 \times \mathbb{R}^{d_z} \to \mathbb{R}$ , called DSR Model, with 8 layers, each containing 512 hidden units, and a single skip connection from the input to the middle layer. We set the loss parameters to  $\lambda = 0.1$ ,  $\tau = 1$ ,  $\alpha = 1e - 3$ . And the model architecture is as Figure 2. The activation function between fully connected layers is softplus activation:  $x \mapsto \frac{1}{\beta} \ln \left(1 + e^{\beta x}\right)$ , where  $\beta = 100$ . The initial latent code vector  $\mathbf{z}$  of size 192, were sampled from  $\mathcal{N}(0, 1.0^2)$ . The Gradient Module was implemented by PyTorch Autograd to calculate  $\nabla_{\boldsymbol{x}} f(\boldsymbol{x})$ .
Distribution  $\mathcal{D}$ . For all experiments, we utilize the distribution  $\mathcal{D}$  as defined in the IGR[40] for the expectation in Eq. (2). This distribution is determined by taking the average of a uniform distribution and a sum of Gaussians that are centered at  $\mathcal{X}$  with a standard deviation equal to the distance to the  $k$ -th nearest neighbor (where we set  $k = 50$ ).
Level set extraction. We extract the zero (or any other) level set of a trained model  $f(x, t; \theta)$  using the Marching Cubes algorithm[41] implemented in the python scikit-image package, which can use any large-size grid to achieve any high resolution.
3.3 Evaluation Metrics.
We use three evaluation metrics commonly used in 3D modeling to evaluate the similarity between two 3D shapes from three aspects: volume, distance, and normal vectors, which are Volumetric Intersection over Union (IoU), Chamfer distance and Normal Consistency (NC). These three metrics are all normalized to a range of 0-1, providing a comprehensive evaluation of the model's performance from different perspectives.
IoU. IoU compares the reconstructed volume with the ground truth shape (higher is better). For two arbitrary shapes  $A, B \subseteq \mathbb{S} \in \mathbb{R}^n$  is attained by:
$$
I o U = \frac {\left| A \cap B \right|}{\left| A \cup B \right|} \tag {5}
$$
The RMSD here was calculated between consecutive frames to reflect the smoothness of the protein trajectory.
Chamfer distance. Chamfer distance is a standard metrics to evaluate the distance between two point sets  $\mathcal{X}_1, \mathcal{X}_2 \subset \mathbb{R}^n$  (lower is better).
$$
\mathrm {d} _ {\mathrm {C}} \left(\mathcal {X} _ {1}, \mathcal {X} _ {2}\right) = \frac {1}{2} \left(\mathrm {d} _ {\vec {\mathrm {C}}} \left(\mathcal {X} _ {1}, \mathcal {X} _ {2}\right) + \mathrm {d} _ {\vec {\mathrm {C}}} \left(\mathcal {X} _ {2}, \mathcal {X} _ {1}\right)\right) \tag {6}
$$
$$
\mathrm {d} _ {\vec {\mathrm {C}}} \left(\mathcal {X} _ {1}, \mathcal {X} _ {2}\right) = \frac {1}{\left| \mathcal {X} _ {1} \right|} \sum_ {\boldsymbol {x} _ {1} \in \mathcal {X} _ {1}} \min  _ {\boldsymbol {x} _ {2} \in \mathcal {X} _ {2}} \| \boldsymbol {x} _ {1} - \boldsymbol {x} _ {2} \| \tag {7}
$$
NC NC evaluate estimated surface normals (higher is better). Normal consistency between two normalized unit vectors  $n_i$  and  $n_j$  is defined as the dot product between the two vectors. For evaluating the surface normals, given the object surface points and normal vectors:  $X_{pred} = \{(\pmb{x}_i,\overrightarrow{n_i})\}$ , and the ground truth surface points and normal vectors:  $X_{gt} = \{(\pmb{y}_j,\overrightarrow{m_j})\}$ , the surface normal consistency between  $X_{pred}$  and  $X_{gt}$ , denoted as  $\Gamma$ , is defined as:
$$
\Gamma \left(X _ {g t}, X _ {\text {p r e d}}\right) = \frac {1}{| X _ {g t} |} \sum_ {j \in | X _ {g t} |} \left| \vec {n _ {j}} \cdot \vec {m} _ {\theta \left(\boldsymbol {y} _ {j}, X _ {\text {p r e d}}\right)} \right| \tag {8}
$$
$$
\theta \left(\boldsymbol {y} _ {j}, X _ {\text {p r e d}} := \left\{\left(\boldsymbol {x} _ {i}, \overrightarrow {n _ {i} ^ {\prime}}\right) \right\}\right) = \underset {i \in \left| X _ {\text {p r e d}} \right|} {\arg \min } \left\| \boldsymbol {y} _ {j} - \boldsymbol {x} _ {i} \right\| _ {2} ^ {2} \tag {9}
$$
4 Results and Analysis
We design experiments from three aspects to verify the ability to learn SDF from the raw point clouds, the ability to reconstruct the protein dynamical surface, and the generalization of the model in terms of temporal interpolation and extrapolation. To our knowledge, we are the first to model long-term dynamics of large proteins through surface representation, with no other methods currently available for comparison. Therefore, we conduct the following analysis to verify the learning ability of the DSR model and the effectiveness of protein surface representation. Additionally, to demonstrate that our model is not doing something trivial, we compare our method with linear interpolation on interpolation tasks.
4.1 The Ability to Learn SDF
We show the reconstruction results of training on Abeta protein in Figure 3, which demonstrates the ability of our model to learn the signed distance field. The first row in the figure is the reconstruction results of our model, and the second row is the ground truth. It can be seen that our model has the ability to reconstruct the dynamic shape of the protein at different moments, and can restore a clearer and more delicate shape than the ground truth. Furthermore, we show some videos of this protein on the webpage and discuss the effect of normal vectors on learning process in the Appendix D.
Figure 3: Given different times as input, the zero-level set visualization of the SDF value predicted by the model for protein abeta. The first row is the reconstruction result of our model, and the second row is the ground truth.
Note that SDF represents the shortest distance (signed) from the surface of the object, which means that different level sets form surfaces at certain distances from the manifold surface. As shown
in Figure 4 (a), the visualization of different level sets of SDF values predicted by our model is consistent with the meaning expressed by SDF. In addition, our model is continuous with respect to spatial and can use to reconstruct and predict for arbitrary space resolutions. Figure 4(b) shows the reconstruction of a protein at different resolutions at a certain time. We also calculated three evaluation metrics, namely IoU, Chamfer distance and NC, which have values of  $0.9361 \pm 0.0265$ ,  $0.0009 \pm 0.0008$  and  $0.9967 \pm 0.0011$  respectively.
Resolution
Figure 4: (a) The visualization of different level-set of the SDF predicted by our model; positive level sets are in red; negative ones are in blue; the zero level set, represents the approximated surface in white. (b) The reconstruction of a protein at different resolutions.
[128, 128, 128]
[256, 256, 256]
[512, 512, 512]
4.2 Reconstruction across Time
In this experiment, we validate the ability of our trained model to reconstruct the protein shapes in training frames (see Table 1). For data AdK equilibrium, NhaA, YiiP, and DRYAD_MD, we used the first 1500 frames of each trajectory and diluted them by ten times. That is, selecting 150 frames for each trajectory for training. Similarly, every second frame was selected within the first 80 frames for DIMS, every third frame was chosen within the first 120 frames for FRODA, and every fifth frame was selected within the first 400 frames for I-FABP.