Learning from Protein Structure with Geometric Vector Perceptrons

1 INTRODUCTION . Many efforts in structural biology aim to predict , or derive insights from , the structure of a macromolecule ( such as a protein , RNA , or DNA ) , represented as a set of positions associated with atoms or groups of atoms in 3D Euclidean space . These problems can often be framed as functions mapping the input domain of structures to some property of interest—for example , predicting the quality of a structural model or determining whether two molecules will bind in a particular geometry . Thanks to their importance and difficulty , such problems , which we broadly refer to as learning from structure , have recently developed into an exciting and promising application area for deep learning ( Graves et al. , 2020 ; Ingraham et al. , 2019 ; Pereira et al. , 2016 ; Townshend et al. , 2019 ; Won et al. , 2019 ) . Successful applications of deep learning are often driven by techniques that leverage the problem structure of the domain—for example , convolutions in computer vision ( Cohen & Shashua , 2017 ) and attention in natural language processing ( Vaswani et al. , 2017 ) . What are the relevant considerations in the domain of learning from structure ? Using proteins as the most common example , we have on the one hand the arrangement and orientation of the amino acid residues in space , which govern the dynamics and function of the molecule ( Berg et al. , 2002 ) . On the other hand , proteins also possess relational structure in terms of their amino-acid sequence and the residue-residue interactions that mediate the aforementioned protein properties ( Hammes-Schiffer & Benkovic , 2006 ) . We refer to these as the geometric and relational aspects of the problem domain , respectively . Recent state-of-the-art methods for learning from structure leverage one of these two aspects . Commonly , such methods employ either graph neural networks ( GNNs ) , which are expressive in terms of relational reasoning ( Battaglia et al. , 2018 ) , or convolutional neural networks ( CNNs ) , which operate directly on the geometry of the structure . Here , we present a unifying architecture that bridges these two families of methods to leverage both aspects of the problem domain . We do so by introducing geometric vector perceptrons ( GVPs ) , a drop-in replacement for standard multi-layer perceptrons ( MLPs ) in aggregation and feed-forward layers of GNNs . GVPs operate directly on both scalar and geometric features—features that transform as a vector under a rotation of spatial coordinates . GVPs therefore allow for the embedding of geometric information at nodes and ∗Equal contribution edges without reducing such information to scalars that may not fully capture complex geometry . We postulate that our approach makes it easier for a GNN to learn functions whose significant features are both geometric and relational . Our method ( GVP-GNN ) can be applied to any problem where the input domain is a structure of a single macromolecule or of molecules bound to one another . In this work , we specifically demonstrate our approach on two problems connected to protein structure : computational protein design and model quality assessment . Computational protein design ( CPD ) is the conceptual inverse of protein structure prediction , aiming to infer an amino acid sequence that will fold into a given structure . Model quality assessment ( MQA ) aims to select the best structural model of a protein from a large pool of candidate structures and is an important step in structure prediction ( Cheng et al. , 2019 ) . Our method outperforms existing methods on both tasks . 2 RELATED WORK . ML methods for learning from protein structure largely fall into one of three types , operating on sequential , voxelized , or graph-structured representations of proteins . We briefly discuss each type and introduce state-of-the-art examples for MQA and CPD to set the stage for our experiments later . Sequential representations In traditional models of learning from protein structure , each amino acid is represented as a feature vector using hand-crafted representations of the 3D structural environment . These representations include residue contacts ( Olechnovič & Venclovas , 2017 ) , orientations or positions collectively projected to local coordinates ( Karasikov et al. , 2019 ) , physicsinspired energy terms ( O ’ Connell et al. , 2018 ; Uziela et al. , 2017 ) , or context-free grammars of protein topology ( Greener et al. , 2018 ) . The structure is then viewed as a sequence or collection of such features which can be fed into a 1D convolutional network , RNN , or dense feedforward network . Although these methods only indirectly represent the full 3D structure of the protein , a number of them , such as ProQ4 ( Hurtado et al. , 2018 ) , VoroMQA ( Olechnovič & Venclovas , 2017 ) , and SBROD ( Karasikov et al. , 2019 ) , are competitive in assessments of MQA . Voxelized representations In lieu of hand-crafted representations of structure , 3D convolutional neural networks ( CNNs ) can operate directly on the positions of atoms in space , encoded as occupancy maps in a voxelized 3D volume . The hierarchical convolutions of such networks are easily compatible with the detection of structural motifs , binding pockets , and the specific shapes of other important structural features , leveraging the geometric aspect of the domain . A number of CPD methods ( Anand et al. , 2020 ; Zhang et al. , 2019 ; Shroff et al. , 2019 ) and the MQA methods 3DCNN ( Derevyanko et al. , 2018 ) and Ornate ( Pagès et al. , 2019 ) exemplify the power of this approach . Graph-structured representations A protein structure can also be represented as a proximity graph over amino acid nodes , reducing the challenge of representing a collective structural neighborhood in a single feature vector to that of representing individual edges . Graph neural networks ( GNNs ) can then perform complex relational reasoning over structures ( Battaglia et al. , 2018 ) —for example , identifying key relationships among amino acids , or flexible structural motifs described as a connectivity pattern rather than a rigid shape . Recent state-of-the-art GNNs include Structured Transformer ( Ingraham et al. , 2019 ) on CPD , ProteinSolver ( Strokach et al. , 2020 ) on CPD and mutation stability prediction , and GraphQA ( Baldassarre et al. , 2020 ) on MQA . These methods vary in their representation of geometry : while some , such as ProteinSolver and GraphQA , represent edges as a function of their length , others , such as Structured Transformer , indirectly encode the 3D geometry of the proximity graph in terms of relative orientations and other scalar features . 3 METHODS . Our architecture seeks to combine the strengths of CNN and GNN methods in learning from biomolecular structure by improving the latter ’ s ability to reason geometrically . The GNNs described in the previous section encode the 3D geometry of the protein by encoding vector features ( such as node orientations and edge directions ) in terms of rotation-invariant scalars , often by defining a local coordinate system at each node . We instead propose that these features be directly represented as geometric vectors—features in R3 which transform appropriately under a change of spatial coordinates—at all steps of graph propagation . This conceptual shift has two important ramifications . First , the input representation is more efficient : instead of encoding the orientation of a node by its relative orientation with all of its neighbors , we only have to represent one absolute orientation per node . Second , it standardizes a global coordinate system across the entire structure , which allows geometric features to be directly propagated without transforming between local coordinates . For example , representations of arbitrary positions in space—including points that are not themselves nodes—can be easily propagated across the graph by Euclidean vector addition . We postulate this allows the GNN to more easily access global geometric properties of the structure . The key challenge with this representation , however , is to perform graph propagation in a way that simultaneously preserves the full expressive power of the original GNN while maintaining the rotation invariance provided by the scalar representations . We do so by introducing a new module , the geometric vector perceptron , to replace dense layers in a GNN . 3.1 GEOMETRIC VECTOR PERCEPTRONS . The geometric vector perceptron is a simple module for learning vector-valued and scalar-valued functions over geometric vectors and scalars . That is , given a tuple ( s , V ) of scalar features s ∈ Rn and vector features V ∈ Rν×3 , we compute new features ( s′ , V′ ) ∈ Rm ×Rµ×3 . The computation is illustrated in Figure 1A and formally described in Algorithm 1 . At its core , the GVP consists of two separate linear transformations Wm , Wh for the scalar and vector features , followed by nonlinearities σ , σ+ . However , before the scalar features are transformed , we concatenate the L2 norm of the transformed vector features Vh ; this allows us to extract rotation-invariant information from the input vectors V. An additional linear transformation Wµ is inserted just before the vector nonlinearity to control the output dimensionality independently of the number of norms extracted . The GVP is conceptually simple , yet provably possesses the desired properties of invariance/equivariance and expressiveness . First , the vector and scalar outputs of the GVP are equivariant and invariant , respectively , with respect to an arbitrary composition R of rotations and reflections in 3D Euclidean space — i.e. , if GVP ( s , V ) = ( s′ , V′ ) then GVP ( s , R ( V ) ) = ( s′ , R ( V′ ) ) ( 1 ) Algorithm 1 Geometric vector perceptron Input : Scalar and vector features ( s , V ) ∈ Rn × Rν×3 . Output : Scalar and vector features ( s′ , V′ ) ∈ Rm × Rµ×3 . h← max ( ν , µ ) GVP : Vh ←WhV ∈ Rh×3 Vµ ←WµVh ∈ Rµ×3 sh ← ‖Vh‖2 ( row-wise ) ∈ Rh vµ ← ‖Vµ‖2 ( row-wise ) ∈ R µ sh+n ← concat ( sh , s ) ∈ Rh+n sm ←Wmsh+n + b ∈ Rm s′ ← σ ( sm ) ∈ Rm V′ ← σ+ ( vµ ) Vµ ( row-wise multiplication ) ∈ Rµ×3 return ( s′ , V′ ) This is due to the fact that the only operations on vector-valued inputs are scalar multiplication , linear combination , and the L2 norm.1 We include a formal proof in Appendix A . In addition , the GVP architecture can approximate any continuous rotation- and reflection-invariant scalar-valued function of V. More precisely , let Gs be a GVP defined with n , µ = 0—that is , one which transforms vector features to scalar features . Then for any function f : Rν×3 → R invariant with respect to rotations and reflections in 3D , there exists a functional form Gs able to -approximate f , given mild assumptions . Theorem . Let R describe an arbitrary rotation and/or reflection in R3 . For ν ≥ 3 let Ων ⊂ Rν×3 be the set of all V = [ v1 , . . . , vν ] T ∈ Rν×3 such that v1 , v2 , v3 are linearly independent and 0 < ||vi||2 ≤ b for all i and some finite b > 0 . Then for any continuous F : Ων → R such that F ( R ( V ) ) = F ( V ) and for any > 0 , there exists a form f ( V ) = wTGs ( V ) such that |F ( V ) − f ( V ) | < for all V ∈ Ων . We include a formal proof in Appendix A . As a corollary , a GVP with nonzero n , µ is also able to approximate similarly-defined functions over the full input domain Rn × Rν×3 . In addition to the GVP layer itself , we use a version of dropout that drops entire vector channels at random ( as opposed to coordinates within vector channels ) . We also introduce layer normalization for the vector features as V← V/ √ 1 ν ‖V‖22 ∈ R ν×3 ( 2 ) That is , we scale the row vectors of V such that their root-mean-square norm is one . This vector layer norm has no trainable parameters , but we continue to use normal layer normalization on scalar channels with trainable parameters γ , β . We study our hypothesis that GVPs augment the geometric reasoning ability of GNNs on a synthetic dataset ( Appendix B ) . The synthetic dataset allows us to control the function underlying the groundtruth label in order to explicitly separate geometric and relational aspects in different tasks . The GVP-augmented GNN ( or GVP-GNN ) matches a CNN on a geometric task and a standard GNN on a relational task . However , when we combine the two tasks in one objective , the GVP-GNN does significantly better than either a GNN or a CNN .

Summary: The challenge of predicting the structure of biological macro-molecules is widely relevant in many applications and difficult to address. This paper divides the types of approaches taken to address this challenge into those that use "geometric" information (i.e. positions of molecules in space), and those that utilize "relational methods" mainly through graphs (how different parts of molecule relate). This study is an attempt to integrate the two source of information by a novel network architecture. They introduce geometric vector perceptrons as a way of summarizing geometric information for graph layers without loss of information as it happens in dense layers. They evaluate the performance of these architectures on MQA and CPD tasks, both relevant and standard benchmarks in the field.