Protein model quality assessment using rotation-equivariant, hierarchical neural networks

We train and test our method on candidate models submitted to multiple rounds of CASP, a biennial community experiment. CASP [11] addresses the protein structure prediction problem by withholding newly solved experimental structures (referred to as targets) and allowing computational groups to make predictions (referred to as models). Submitted models are released as sets in two stages (20 models per target for stage 1, 150 models per target for stage 2) for Model Quality Assessment (MQA), a specific subcategory of CASP that aims to assess the performance of scoring functions. Model quality is measured in terms of GDT_TS [12] based on the alignment of native structure and candidate model. GDT_TS ranges between 0 and 1, with higher GDT_TS value indicating better model quality.

We mirror the setup of the CASP experiment and split the CASP datasets based on protein target and release year. We train and validate our method on the set of models submitted to CASP 5-10 (500 targets for training, 58 for validation). For testing, we consider models submitted to stage 2 for CASP 11 (84 targets) and 12 (40 targets). We relaxed the structure of all models with the SCWRL4 software [13] to improve side-chain conformations prior to feeding the models into our network.

Our network builds on recent neural network architectures that are specifically designed to learn from 3D atomic structures [14, 15]. Figure 1 illustrates the architecture of our network. At its core are a point-based representation of atoms and multiple layers of rotation-equivariant convolutions. The goal is to predict a quality score for a given protein model. Each atom has an associated feature vector. At input this is simply the one-hot encoding of its element type¹¹1We represent carbon, oxygen, nitrogen, and sulfur atoms.. We then perform rotation-equivariant convolutions that result in a new feature vector associated with each atom. The convolution filters of each layer are constructed based on a truncated series of spherical harmonics (up to rotation order $l=2$), such that the filters are able to recognize structural motifs independent of their orientation or position in space. Convolutions are performed among a limited set of k-nearest neighbors (k=40) to account for the fact that the physical laws governing intra-molecular interactions are local. The subsequent convolution layer outputs features only for the alpha carbon of each amino acid residue. This subsampling operation aggregates information hierarchically and allows the network to recognize structural motifs at different scales. We average the alpha carbon feature vectors to obtain a fingerprint for the entire protein model²²2Instead, we could also choose to aggregate information at the level of a single point using further subsampling operations (see [15]). The optimal choice of hierarchy is likely application dependent.. From this fingerprint, we use two dense layers (250 and 150 units, ReLu activation) to calculate a single scalar quality score.

We formulate the training as a regression task aiming to predict the quality metric GDT_TS for each model. The Huber loss between the actual and the predicted GDT_TS is used as the loss function. We train with the Adam optimizer in TensorFlow [16] (learning rate $1.25\cdot 10^{-4}$) and monitor the loss on the validation set for every epoch. The weights of the best-performing network are then used to evaluate the predictions on the test set. We use Horovod [17] to distribute training across 4 NVIDIA Titan X GPUs.

We examine our results on model quality assessment, finding that we generally improve upon state-of-the-art methods (Table 1). We report multiple correlation metrics per method. Global correlation indicates the correlation between a method’s predictions and the GDT_TS scores of all protein models in a given set. Per target indicates the correlation with respect to a method’s predictions and the GDT_TS scores of the protein models for a given target (averaged over all targets). The two measures provide complementary information about a method’s performance. Good global correlation is desirable to judge the absolute quality of a set of candidate models. Per target correlation indicates a method’s ability to distinguish model quality among a set of models for one target (the main scoring challenge). We report Pearson, Kendall, and Spearman correlation coefficients for both global and per target correlations. Notably, our method also improves upon ProQ3D [10], which uses information on related proteins to make its predictions³³3GraphQA can also leverage information on related proteins, but here we compare against the version of GraphQA that only uses structural information..

In Figure 2, we explore whether the network has learned to encode certain structural motifs. We project the fingerprint (see Figure 1) of each protein model in the test set into a lower-dimensional space using Principal Component Analysis (PCA). Visual inspection reveals that the fingerprint contains information on both inter-atomic interactions and secondary protein structure.

In Figure 2A, we consider the prevalence of van der Waals interactions in each protein model. We use the software GetContacts [19] to identify van der Waals interactions and divide the total number of interactions by the number of amino acid residues per model. We note that protein models cluster based on the prevalence of van der Waals interactions in the plane of principal component (PC) 5 and PC 1. In Figure 2B, we consider the prevalence of alpha helices in a given protein model. We use DSSP [20] to assign each residue in a protein model to one type of secondary structure (’alpha helix’, ’beta sheet’ or ’other’). We then calculate the fraction of ’alpha helix’ residues over all residues in a model. Models cluster based on this fraction when plotted based on PC 4 and PC 7.