\section{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. Thanks to their importance and difficulty, such problems, which we broadly refer to as \textit{learning from structure}, have recently developed into an exciting and promising application area for deep learning \citep{gainza2020deciphering, graves2020review, ingraham2019generative, pereira2016boosting, townshend2019end, won2019assessment}.

Successful applications of deep learning are often driven by techniques that leverage the problem structure of the domain---for example, convolutions in computer vision \citep{cohen2016inductive} and attention in natural language processing \cite{vaswani2017attention}. What are the relevant considerations in the domain of learning from structure? 
Using proteins as the most common example, the spatial \textit{arrangement} and \textit{orientation} of the amino acids govern the topology, dynamics and function of the molecule \citep{berg2002biochemistry}. These properties are in turn mediated by key pairwise and higher-order residue-residue \textit{interactions} \citep{del2006residue, hammes2006relating}. We refer to these as the \textit{geometric} and \textit{relational} aspects of the problem domain, respectively.


Current state-of-the-art methods for learning from structure are successful by leveraging one of these two aspects. Commonly, such methods either employ graph neural networks (GNNs), which are expressive in terms of relational reasoning \citep{battaglia2018relational} or convolutional neural networks (CNNs), which operate directly on spatial features. Here, we present a unifying architecture that bridges spatial and graph-based methods to leverage \textit{both} aspects of the problem domain. 

We do so by introducing \textit{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 \textit{geometric} features---features that transform as a vector under rotation of spatial coordinates. GVPs therefore allow for the embedding of geometric information at nodes and edges without reducing such information to hand-picked scalars that may not fully capture complex geometric relationships. We postulate that our approach makes it easier for a network 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 demonstrate our approach on two important problems connected to protein structure: model quality assessment and computational protein design. Model quality assessment (MQA) aims to select the best structural model of a protein from a large pool of candidate structures and is a crucial step in protein structure prediction \citep{cheng2019estimation}. Computational protein design (CPD) is the conceptual inverse of structure prediction, aiming to infer an amino acid sequence that will fold into a given structure. Our method outperforms existing methods on both tasks.


\section{Related Work}
Current methods for learning from protein structure generally use one of three classes of structural representations, which we outline below. We also describe representative and state-of-the-art examples for MQA and CPD to set the stage for our experiments later.


\paragraph{Sequential representation}
One way of representing a protein structure is as a sequence of feature vectors, one for each amino acid, that can be input into a 1D convolutional network. These representations intrinsically encode neither geometric nor relational information and instead rely on hand-selected
features to represent the 3D structural neighborhood of each amino acid. Such features can include contact-based features \citep{olechnovivc2017voromqa}, orientations or positions collectively projected to local coordinates \cite{karasikov2019smooth, wang2018computational}, and physics-inspired energy terms \citep{o2018spin2, uziela2017proq3d}. These methods were among the earliest developed for learning from structure, and a number of them (ProQ3D \citep{uziela2017proq3d}, VoroMQA \citep{olechnovivc2017voromqa}, SBROD \citep{karasikov2019smooth}) remain competitive in assessments of MQA.

\paragraph{Graph-based representations}
Graph-based methods represent amino acids or individual atoms as nodes and draw edges based on chemical bonds or spatial proximity, and are well-suited for complex relational reasoning \citep{battaglia2018relational}. These methods reduce the challenging task of representing the collective structural neighborhood of an amino acid to that of representing individual edges. The CPD method ProteinSolver \citep{strokach2020fast}, for example, uses pairwise Euclidean and sequence distance to represent edges. However, to encode more complex geometric information, a choice of scalar representation scheme is again necessary
\textit{i.e.}, vector-valued features whose components are intrinsically linked need to be decomposed into individual scalars. 
For example, the MQA method ProteinGCN \citep{sanyal2020proteingcn} projects pairwise relative positions into local coordinates. The CPD method Structured Transformer \citep{ingraham2019generative} additionally uses a quaternion representation of pairwise relative orientations.

\paragraph{Voxel-based representations} 
Many structural biology methods voxelize proteins or neighborhoods of amino acids into occupancy maps and then perform hierarchical 3D convolutions. These architectures circumvent the need for hand-picked representations of geometric information altogether and instead directly operate on approximate positions in 3D space. However, CNNs do not directly encode relational information within the structure. In addition, standard CNNs do not intrinsically account for the symmetries of the problem space. To address the latter, early methods, such as the MQA method 3DCNN \citep{derevyanko2018deep}, learned rotation invariance via data augmentation. More recent MQA methods such as Ornate \citep{pages2019protein} and many end-to-end CPD models \citep{anand2020protein, zhang2019prodconn} instead define voxels based on a local coordinate system. 

Voxel-based and graph-based architectures, catering to geometric and relational aspects of the problem domain, respectively, have complementary strengths and weaknesses. To address this, some recent work on small molecules aims to leverage geometric features in GNNs through voxelized convolutions at each graph node \citep{spurek2019geometric}, message transformation based on encoded angle information \citep{klicpera2020directional} or products with position vectors to account for vector features \citep{cho2019enhanced}. Our contribution is similar in spirit but more general and expressive in methodology. With the introduction of GVPs, we provide a conceptually simple augmentation of GNNs that allows for arbitrary scalar and vector features at all steps of graph propagation, and has the desired equivariance properties with respect to rotations of atomic structures in space.

\section{Methods}
In this section, we formally introduce the geometric vector perceptron, discuss datasets and evaluation metrics for quality assessment and protein design, and describe the protein representations and model architectures used for those tasks.

\subsection{Geometric vector perceptrons}
Given a tuple $(\mathbf{s}, \mathbf{V})$ of scalar features $\mathbf{s} \in \mathbb{R}^n$ and vector features $\mathbf{V}\in\mathbb{R}^{\nu\times 3}$,
the geometric vector perceptron computes new features $(\mathbf{s}',\mathbf{V}') \in \mathbb{R}^m \times \mathbb{R}^{\mu \times 3}$. As outlined in pseudocode in Algorithm~\ref{alg:gvp} and illustrated in Figure~\ref{fig:gvp}A, the computation is analogous to the linear combination of inputs in normal dense layers, with the addition of a norm operation and an intermediary hidden layer, which allow information in the vector channels to be extracted as scalars.  

\begin{algorithm}
\caption{Geometric vector perceptron}\label{alg:gvp}
\begin{algorithmic}[H]
\State \textbf{Input}: Scalar and vector features $(\textbf{s},\textbf{V})$ $\in$ $\mathbb{R}^n \times \mathbb{R}^{\nu \times 3}$ .
\State \textbf{Output}: Scalar and vector features $(\textbf{s}',\textbf{V}')$ $\in$ $\mathbb{R}^m \times \mathbb{R}^{\mu \times 3}$ .
\State $h \gets \max \left(\nu, \mu\right)$

\BState \textbf{GVP}: 
\State \ \ \ \ $\text{V}_h \gets \textbf{W}_h \textbf{V} \quad \in \mathbb{R}^{h \times 3} $ 
\State \ \ \ \ $\text{V}_\mu \gets \textbf{W}_\mu \text{V}_h \quad \in \mathbb{R}^{\mu \times 3} $ 
\State \ \ \ \ $\text{s}_h \gets \norm{\text{V}_h} \ (\text{row-wise}) \quad \in \mathbb{R}^{h}$ 
\State \ \ \ \ $\text{v}_\mu \gets \norm{\text{V}_\mu} \ (\text{row-wise}) \quad \in \mathbb{R}^{\mu}$
\State \ \ \ \ $\text{s}_{h+n} \gets \text{concat}\left(\text{s}_h, \textbf{s}\right) \quad \in \mathbb{R}^{h+n}$
\State \ \ \ \ $\text{s}_m \gets \textbf{W}_m \text{s}_{h+n} + \textbf{b}  \quad \in \mathbb{R}^{m}$
\State \ \ \ \ $\textbf{s}' \gets \sigma\left(\text{s}_m \right) \quad \in \mathbb{R}^{m}$
\State \ \ \ \ $\textbf{V}' \gets  \sigma^+\left(\text{v}_\mu\right) \odot \text{V}_\mu, (\text{row-wise multiplication}) \quad \in \mathbb{R}^{\mu \times 3}$ 
\BState  \textbf{return} $(\textbf{s}',\textbf{V}')$  
\end{algorithmic}
\end{algorithm}

$\mathbf{W}_h \in \mathbb{R}^{h \times \nu}$, $\mathbf{W}_\mu \in \mathbb{R}^{\mu \times h}$, and $\mathbf{W}_m \in \mathbb{R}^{m \times (n+h)}$ are weight matrices and $\mathbf{b}\in\mathbb{R}^m$ is a bias vector. 
$\mathbb{\sigma}$ and $\mathbb{\sigma^+}$ denote activation functions. In our experiments, we choose ReLU and sigmoid for $\mathbb{\sigma}$ and $\mathbb{\sigma^+}$, respectively, although other nonlinearities may be used. 
Additionally, $h$ may also be defined independently, but we use the maximum of $\nu, \mu$ for convenience. 


\begin{figure}
    \centering
    \includegraphics[width=0.8\linewidth]{schematic_v6.png}
    \caption{
     \textbf{(A)} Schematic of the geometric vector perceptron illustrating the computation shown in Algorithm~\ref{alg:gvp}. Given a tuple of scalar and vector input features $(\mathbf{s},\mathbf{V})$, the perceptron computes an updated tuple $(\mathbf{s}',\mathbf{V}')$. $\mathbf{s}'$ is a function of both $\mathbf{s}$ and $\mathbf{V}$. 
    \textbf{(B)} Illustration of the structure-based prediction tasks. In model quality assessment (top), the goal is to predict a quality score given the 3D structure of a candidate model. Individual atoms are represented as colored spheres. The quality score measures the accuracy of a candidate structure with respect to an experimentally determined structure (shown in gray). In computational protein design (bottom), the goal is to predict an amino acid sequence that would fold into a given protein backbone structure.
   }  
    \label{fig:gvp}
\end{figure}


Despite its conceptual simplicity, the GVP offers some desirable properties. First, the vector and scalar outputs of the GVP are equivariant and invariant, respectively, with respect to an arbitrary composition of rotations and reflections in 3D Euclidean space described by $R$ \textit{i.e.}, 

\begin{equation}
\label{eq:eq}
\text{GVP}(\left(\mathbf{s}, R(\mathbf{V}))\right) = (\textbf{s'}, R(\textbf{V}'))
\end{equation} 

This is due to the fact that the only operations on vector-valued inputs are scalar multiplication, linear combination, and the $L_2$ norm. We include a formal proof in Supplementary Information.

In addition, a GVP $G_s$ defined with $n, \mu=0$---that is, the part of a GVP that transforms vector features to scalar features---is able to $\epsilon$-approximate a function $F: \mathbb{R}^{\nu \times 3} \rightarrow \mathbb{R}$ that is invariant with respect to rotations and reflections in 3D under mild assumptions.

\newtheorem*{theorem}{Theorem}
\label{th:eps}
\begin{theorem}
    Let R describe an arbitrary rotation or reflection in $\mathbb{R}^3$.
    For $\nu \ge 3$ let $\Omega^\nu \subset \mathbb{R}^{\nu \times 3}$ be the set of all $\mathbf{V} = \begin{bmatrix}\mathbf{v}_1, & \ldots, & \mathbf{v}_\nu\end{bmatrix}^T\in \mathbb{R}^{\nu \times 3}$ such that $\mathbf{v}_1, \mathbf{v}_2, \mathbf{v}_3$ are linearly independent and $0<||\mathbf{v}_i||_2 \le b$ for all $i$ and some finite $b>0$.
   
    Then for any continuous $F: \Omega^\nu \rightarrow \mathbb{R}$ such that $F(R(\mathbf{V})) = F(\mathbf{V})$ and for any $\epsilon>0$, there exists a form $f(\mathbf{V})=\mathbf{w}^TG_s(\mathbf{V})$ such that $|F(\mathbf{V})-f(\mathbf{V})|< \epsilon$ for all $\mathbf{V}\in\Omega^\nu$.
\end{theorem}
\begin{proof} In Supplementary Information. \end{proof}

As a corollary, a GVP with nonzero $n,\mu$ is also able to approximate similarly-defined functions over the full input domain $\mathbb{R}^n \times \mathbb{R}^{\nu \times 3}$.

Finally, the GVP layer is nearly as fast as normal dense layers, incurring additional overhead only in reshaping and concatenation operations.

In addition to the GVP layer itself, we use a version of dropout that drops entire vector channels at random (as opposed to components within vector channels). We also introduce layer normalization for the vector features as  
\begin{equation}
\mathbf{v}_i \leftarrow \mathbf{v}_i / \sqrt{\frac{1}{|\{j\}|}\sum_j||\mathbf{v}_j||_2^2} \quad \forall \ i \quad ,
\end{equation}
where $\mathbf{v}_i \in \mathbb{R}^3$ are the individual row vectors of the vector feature matrix $\mathbb{R}^{\nu \times 3}$. 
This vector layer norm has no trainable parameters, but we continue to use normal layer normalization on scalar channels with trainable parameters $\gamma, \beta$.


\subsection{Dataset and evaluation}
We benchmark GVP-augmented graph neural networks on two distinct challenges in structural biology: model quality assessment and protein design. Figure~\ref{fig:gvp}B illustrates these tasks.     
\subsubsection{Model quality assessment}
Model quality assessment aims to select the best structural model of a protein from a large pool of candidate structures.\footnote{We refer to models of protein structure as "structural models" or "candidate structures" to avoid confusion with the term "model" as used in the ML community.} The performance of different MQA methods is regularly assessed in the biennial blind prediction competition CASP, of which the 13th was the most recent \citep{cheng2019estimation}. For a number of recently solved but unreleased structures, called \textit{targets}, structure generation programs produce a large number of candidate structures. MQA methods are ranked by how well they are able to predict the GDT-TS score of a candidate structure. The GDT-TS is a  measure of how similar two protein backbones are after global alignment; specifically, it is the mean of
\begin{equation}
    \text{GDT-N} = \frac{1}{\text{\# residues}}\sum_\text{all residues $i$} \mathbbm{1}\left[d(\text{target}_i, \text{candidate}_i)<N \text{ Angstroms}\right]   
\end{equation}
for $N=1,2,4,8$, where $d$ is the Euclidean distance. In addition to accurately predicting the absolute quality of a candidate structure, a good MQA method should also be able to accurately assess the \textit{relative} model qualities among a pool of candidates for a given target so that the best ones can be selected, perhaps for further refinement. Therefore, MQA methods are commonly evaluated on two metrics: a global correlation between the predicted and ground truth scores, pooled across all targets, and the average \textit{per-target} correlation among only the candidate structures for a specific target \citep{cao2016protein, derevyanko2018deep, pages2019protein}.

To fulfill this aim, we train our MQA model on an absolute loss and a pairwise loss. That is, for each training step we intake pairs $i,j$ where $i,j$ are candidate structures for the same target and compute
\begin{equation}
    \mathscr{L} = H(y^{(i)}-\hat{y}^{(i)}) + H(y^{(j)}-\hat{y}^{(j)}) +  H\left((y^{(i)}-y^{(j)})-(\hat{y}^{(i)}-\hat{y}^{(j)})\right)
\end{equation}
where $H$ is the Huber loss. When reshuffling at the beginning of each epoch, we also randomly pair up the candidate structures for each target. Interestingly, adding the pairwise term also improves global correlation, likely because the much larger number of possible pairs makes it more difficult to overfit. 


We train and validate on the set of candidate structures generated in the CASP 5-10 assessments, which collectively contain 528 targets and 79200 candidate structures. For testing, we predict model quality for a total of 20880 stage 1 and stage 2 candidate structures from CASP 11 (84 targets) and 12 (40 targets).

\subsubsection{Protein design}
\label{sec:dataset_cpd}
Computational protein design is the conceptual inverse of structure prediction, aiming to infer an amino acid sequence that will fold into a given structure. In comparison to model quality assessment, computational protein design is more difficult to unambiguously benchmark, as some structures may correspond to a large space of sequences and others may correspond to none at all.
Therefore, the proxy metric of \textit{native sequence recovery}---splitting the set of all known native structures in the PDB and attempting to design the sequences corresponding to held-out structures---is typically used to benchmark CPD models \citep{li2014direct, o2018spin2, wang2018computational}. Drawing an analogy between sequence design and language modelling, \citet{ingraham2019generative} also evaluate the model \textit{perplexity} on held-out native sequences. Both metrics rest on the implicit assumption that native sequences are optimized for their structures \citep{kuhlman2000native} and should be assigned high probabilities.

To best approximate real-world applications that may require design of novel structures, the held-out evaluation set should bear minimal similarity to the training structures. We use the CATH 4.2 dataset curated by \citet{ingraham2019generative} in which all available structures with 40\% nonredudancy are partitioned by their CATH (class, architecture, topology/fold, homologous superfamily) classification. The canonical training, validation, and test splits consist of 18204, 608, and 1120 structures, respectively.  


\subsection{Model architecture}

\paragraph{GVP-GNNs}
In a computation graph, a GVP can be placed at any node typically occupied by a dense layer. To most directly augment relational representations with geometric information, we modify GNNs by permitting all edge and node embeddings to have geometric vector and scalar features, and then use GVPs at all graph propogation and point-wise feed-forward steps.

\paragraph{Representations of proteins}
A protein structure is a sequence of amino acids, where each amino acid consists of four \textit{backbone} atoms\footnote{C$\alpha$, C, N, and O.}
and a set of \textit{sidechain} atoms located in 3D Euclidean space. Here we represent only the backbone  because our MQA benchmark corresponds to the assessment of backbone structure. In CPD, the sidechains are by definition unknown. 

Let X$_i$ be the position of atom X in the $i$th amino acid (where X is a C$\alpha$, C, N, or O atom).  We represent backbone structure as a graph $\mathcal{G} = (\mathcal{V}, \mathcal{E})$ where each node $\mathfrak{v}_i \in \mathcal{V}$ corresponds to an amino acid and has embedding $\mathbf{h}^{(i)}_\mathfrak{v}$ with the following features:
\begin{itemize}
    \item Scalar features $\{\sin, \cos\} \times {\phi, \psi, \omega}$, where $\phi, \psi, \omega$ are the dihedral angles computed from C$_{i-1}$, N$_{i}$, C$\alpha_{i}$, C$_{i}$, and N$_{i+1}$.
    \item The \textit{forward} and \textit{reverse} unit vectors in the directions of C$\alpha_{i+1} - \text{C}\alpha_i$ and C$\alpha_{i-1} - \text{C}\alpha_i$, respectively.
    \item The unit vector in the imputed direction of C$\beta_i-\text{C}\alpha_i$.\footnote{C$\beta$ is the second carbon from the carboxyl carbon C.} This is computed by assuming tetrahedral geometry and normalizing 
    \begin{equation*}
        \sqrt{\frac{1}{3}}(\mathbf{n} \times \mathbf{c})/||\mathbf{n} \times \mathbf{c}||_2-\sqrt{\frac{2}{3}}(\mathbf{n} + \mathbf{c})/||\mathbf{n}+\mathbf{c}||_2
    \end{equation*}
    where $\mathbf{n}=\text{N}_i-\text{C}\alpha_i$ and $\mathbf{c}=\text{C}_i-\text{C}\alpha_i$. The three unit vectors unambiguously define the orientation of each amino acid residue. 
   
    \item A one-hot representation of amino acid identity, when available.
\end{itemize}
The set of edges is $\mathcal{E}=\{\mathbf{e}_{i\rightarrow j}\}_{i\neq j}$ for all $i, j$ where $\mathfrak{v}_i$ is among the $k=30$ nearest neighbors of $\mathfrak{v}_j$ as measured by the distance between their C$\alpha$ atoms. This is motivated by the fact that if two entities have a relational dependency they are likely to be close together in space. 
Each edge has an embedding $\mathbf{h}^{(i\rightarrow j)}_e$ with the following features:
\begin{itemize}
    \item The unit vector in the direction of C$\alpha_i-\text{C}\alpha_j$.
    \item The radial basis encoding of the distance $\text{RBF}(||\text{C}\alpha_i-\text{C}\alpha_j||_2)$.
    \item A sinusoidal encoding of $i-j$ as described in \citet{vaswani2017attention}, representing distance along the backbone.
\end{itemize}
In our notation, each feature vector $\mathbf{h}$ is a concatenation of scalar and vector features as described above. Collectively, these features are sufficient for a complete description of the protein backbone. In particular, whereas previous graph-based representations depended on scalar edge embeddings to represent relative orientations and positions, we are able to directly encode the absolute orientations of each amino acid and relative positions in the equivalent of fewer real-valued channels.


\paragraph{Network architecture}
Our GVP-GNN takes as input the protein graph defined above and performs graph propagation steps that update the node embeddings according to:
\begin{eqnarray}
    \mathbf{h}^{(i\rightarrow j)}_m &:=& g\left(\text{concat}\left(\mathbf{h}_\mathfrak{v}^{(i)}, \mathbf{h}^{(i\rightarrow j)}_e\right)\right) \\
    \mathbf{h}^{(j)}_\mathfrak{v} &\leftarrow& \text{LayerNorm}\left(\mathbf{h}^{(j)}_\mathfrak{v} + \frac{1}{||\{i:\mathbf{e}_{i\rightarrow j}\in E\}||}\text{Dropout}\left(\sum_{i:\mathbf{e}_{i\rightarrow j}\in E} \mathbf{h}_m^{(i\rightarrow j)}\right)\right)
\end{eqnarray}
where $g$ is a composition of GVPs. We do not update edge embeddings and do not use a global graph embedding. Between node update steps, we also use a residual feed-forward point-wise update layer:
\begin{equation}
    \mathbf{h}^{(i)}_\mathfrak{v} \leftarrow \text{LayerNorm}\left(\mathbf{h}^{(i)}_\mathfrak{v} + \text{Dropout}\left(g\left(\mathbf{h}_\mathfrak{v}^{(i)}\right)\right)\right)
\end{equation}
In model quality assessment, the network performs regression against the true quality score of a candidate structure, a global scalar property. To obtain a single global representation, we apply a node-wise GVP to reduce all features to scalars after all graph propagation steps, and then average the representations across all nodes.\footnote{We also tried learning a weighted average of nodes, but this led to increased overfitting.} A final dense network with dropout and layer normalization then outputs the network's prediction.

In computational protein design, the network learns a generative model over the space of protein sequences conditioned on the given backbone structure. Following \citet{ingraham2019generative}, we frame this as an \textit{autoregressive} task and use a masked encoder-decoder architecture to capture the joint distribution over all positions: for each position $i$, the network models the distribution over amino acids at $i$ based on the complete structure graph, as well as the sequence information at positions $j < i$. The encoder first performs graph propagation on the structural information only. Then, sequence information is added to the graph, and the decoder performs further graph propagation where incoming messages $\mathbf{h}_m^{(i\rightarrow j)}$ for $i\ge j$ are computed only with the encoder embeddings. Finally, we use one last GVP with 20-way scalar output and softmax activation to output the probability of the amino acids.

The MQA network, structure encoder, and masked decoder each use 3 graph propagation layers. The feed-forward modules consist of $2$ GVPs and the message-gather function uses 3. 

\section{Experiments}

\subsection{Model quality assessment}
We compare our method against state-of-the-art methods on the CASP 11-12 test set in Table~\ref{tab:casp11-12}.


 These include representatives of voxel-based methods (3DCNN and Ornate), a graph-based method (GraphQA), and three methods that use sequential representations. All of these methods learn solely from protein structure,\footnote{There are two versions of GraphQA; we compare against the one using only structure information.} with the exception of ProQ3D, which in addition uses sequence information on related proteins that is not always available. We include ProQ3D because it is an improved version of the best single-model method in CASP 11 and CASP 12 \citep{uziela2017proq3d}. Our method outperforms all other structural methods in both global and per-target correlation, and even performs better than ProQ3D on all but one benchmark.
\begin{table}
  \caption{Comparison with state-of-the-art methods on CASP 11 and 12 in terms of global (Glob) and mean per-target (Per) Pearson correlation coefficients (higher is better). ProQ3D is set aside as the sole non-structure-only method. The top performing structure-only method for each metric is in bold, as is the top performing-method overall (if different). Our method generally improves over all other methods.}
  \label{tab:casp11-12}
  \centering
  
  \begin{tabular}{lccccccccccc}
    \toprule
    & \multicolumn{5}{c}{CASP 11} & & \multicolumn{5}{c}{CASP 12} \\
    \cmidrule(r){2-6} \cmidrule(r){8-12} 
     & \multicolumn{2}{c}{Stage 1} & & \multicolumn{2}{c}{Stage 2}      & & \multicolumn{2}{c}{Stage 1} & & \multicolumn{2}{c}{Stage 2}                  \\
    \cmidrule(r){2-3} \cmidrule(r){5-6} \cmidrule(r){8-9} \cmidrule(r){11-12}
    Method     & Glob     & Per & & Glob & Per & & Glob & Per & & Glob & Per \\
    \midrule
    Ours       & \textbf{0.84} & \textbf{0.66} & & \textbf{0.87} & \textbf{0.45} & & \textbf{0.79} & \textbf{0.73} & & \textbf{0.82} & \textbf{0.62} \\
    3DCNN \citep{derevyanko2018deep}   & 0.59 & 0.52 & & 0.64 & 0.40 & & 0.49 & 0.44 & & 0.61 & 0.51 \\ 
    Ornate \citep{pages2019protein}    & 0.64 & 0.47 & & 0.63 & 0.39 & & 0.55 & 0.57 & & 0.67 & 0.49 \\
    GraphQA \citep{baldassarre2020graphqa}             & 0.83 & 0.63 & & 0.82 & 0.38 & & 0.72 & 0.68 & & 0.81 & 0.61 \\
    VoroMQA \citep{olechnovivc2017voromqa} & 0.69 & 0.62 & & 0.65 & 0.42 & & 0.46 & 0.61 & & 0.61 & 0.56 \\
    SBROD \citep{karasikov2019smooth}      & 0.58 & 0.65 & & 0.55 & 0.43 & & 0.37 & 0.64 & & 0.47 & 0.61 \\
    \midrule
    ProQ3D \citep{uziela2017proq3d}        & 0.80 & \textbf{0.69} & & 0.77 & 0.44 & & 0.67 & 0.71 & & 0.81 & 0.60 \\
    \bottomrule
  \end{tabular}
\end{table}

The CASP 11-12 datasets have been the most well-benchmarked in the recent MQA literature. However, for completeness, we also evaluate our method on CASP 13 (Table \ref{tab:casp13}). Because of its recency, many target structures remain publicly unavailable. We use the stage 2 candidate structures of a subset of 20 targets previously used for benchmarking \citep{baldassarre2020graphqa}. Our method achieves improved results over all other methods, including ProteinGCN \citep{sanyal2020proteingcn}, a more recent graph-based method. Because of the small sample size, we emphasize that these results, although promising, should be considered preliminary until further structures for CASP 13 are available.

\begin{table}
  \caption{Performance on the 20 publicly available CASP 13 targets, stage 2 in terms of global and mean per-target Pearson correlation coefficient. As before, the top structure-only method is in bold, as is the top method overall (if different). Our method outperforms all other methods.}
  \label{tab:casp13}
  \centering
  \begin{tabular}{lccccc}
    \toprule
    Method     & Global     & Mean per-target \\
    \midrule
    Ours   &  \textbf{0.881} & \textbf{0.685} \\
    ProteinGCN \citep{sanyal2020proteingcn}    & 0.723 & 0.603 \\
    VoroMQA \citep{olechnovivc2017voromqa}     & 0.769 & 0.665 \\
    \midrule
    ProQ3D \citep{uziela2017proq3d}           & 0.849 & 0.671  \\
    \bottomrule
  \end{tabular}
\end{table}

\subsection{Protein design}
\label{sec:results_cpd}

Our method achieves state-of-the-art performance on CATH 4.2, representing a substantial improvement both in terms of perplexity and sequence recovery over Structured Transformer \citep{ingraham2019generative}, which was trained on the same training set (Table \ref{tab:ingraham}). Following \citet{ingraham2019generative}, we report evaluation on short (100 or fewer amino acid residues) and single-chain subsets of the CATH 4.2 test set, containing 94 and 103 proteins, respectively, in addition to the full test set. Although Structured Transformer leverages an attention mechanism on top of a graph-based representation of proteins, the authors note in ablation studies that removing attention appeared to increase performance. We therefore retrain and compare against a version of Structured Transformer with the attention layers replaced with standard graph propagation operations that we call Structured GNN. Our method also improves upon this model. 

\begin{table}
  \caption{Performance on the CATH 4.2 test set and its short and single-chain subsets in terms of per-residue perplexity (lower is better) and recovery (higher is better). Recovery is reported as the median over all structures of the mean recovery of 100 sequences per structure. Our method performs better than Structured Transformer and a variant of it, Structured GNN, in which we replaced the attention mechanisms with standard graph propagation operations (see \ref{sec:results_cpd}).}
  \label{tab:ingraham}
  \centering
  \begin{tabular}{lccccccc}
    \toprule
     & \multicolumn{3}{c}{Perplexity} & & \multicolumn{3}{c}{Recovery \%}  \\
    \cmidrule(r){2-4} \cmidrule(r){6-8}
    Method     & Short & Single-chain & All & & Short & Single-chain & All \\
    \midrule
    Ours & \textbf{7.10} & \textbf{7.44} & \textbf{5.29} & & \textbf{32.1} & \textbf{32.0} & \textbf{40.2} \\
    Structured GNN       & 8.31 & 8.88 & 6.55 & & 28.4 & 28.1 & 37.3 \\
    Structured Transformer \citep{ingraham2019generative} & 8.54 & 9.03 & 6.85 & & 28.3 & 27.6 & 36.4 \\
    \bottomrule
  \end{tabular}
\end{table}

We emphasize that the underlying architecture of Structured GNN is comparable to ours, with the exception of GVPs and geometric vector channels. In particular, the underlying protein graph is built in a similar manner, except with structural information encoded in solely \textit{scalar} channels. For example, \textit{relative} orientations are encoded in terms of quaternion coefficients (4 scalar channels per \textit{edge}), whereas we represent \textit{absolute} orientations with 3 vector channels per \textit{node}. Therefore, our improvement over Structured GNN is the most direct indication of the benefit of our approach that combines geometric and relational information. Additionally, due to the efficiency of our representation, we are able to achieve this performance gain with a modest but notable decrease in parameter count---1.01M in our model versus 1.38M in Structured GNN and 1.53M in Structured Transformer.



\section{Conclusion}
In this work, we developed the first architecture designed specifically for learning on dual relational and geometric representations of 3D macromolecular structure. At its core, our method, GVP-GNN, augments graph neural networks with computationally fast layers that perform expressive geometric reasoning over Euclidean vector features. Our method possesses desirable theoretical properties and empirically outperforms existing architectures on learning quality scores and sequence designs, respectively, from protein structure.

In further work, we hope to apply our architecture to other important structural biology problem areas, including protein complexes, RNA structure, and protein-ligand interactions. Our results more generally highlight the promise of domain-aware architectures in specialized applications of deep learning, and we hope to continue developing and refining such architectures in the domain of learning from structure.

\section*{Acknowledgements}
We thank Tri Dao and all members of the Dror group for feedback and discussions. 

\section*{Funding}
We acknowledge support from the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) program, and Intel Corporation. SE is supported by a Stanford Bio-X Bowes fellowship. RJLT is supported by the U.S. Department of Energy, Office of Science Graduate Student Research (SCGSR) program.

\medskip
\begin{small}
\bibliographystyle{abbrvnat}

\section*{Supplementary Information}


\subsection*{Equivariance and invariance}
The vector and scalar outputs of the GVP are equivariant and invariant, respectively, with respect to an arbitrary composition of rotations and reflections in 3D Euclidean space described by $R$ \textit{i.e.},  
\begin{equation*}
\text{GVP}(\left(\mathbf{s}, R(\mathbf{V}))\right) = (\mathbf{s'}, R(\mathbf{V}'))
\end{equation*} 
\begin{proof}
    
We can write the transformation described by $R$ as multiplying $\mathbf{V}$ with a unitary matrix $\mathbf{U} \in \mathbb{R}^{3 \times 3}$ from the right. The L$_2$-norm, scalar multiplications, and nonlinearities are defined row-wise as in Algorithm 1. We consider scalar and vector outputs separately. The scalar output, as a function of the inputs, is
\begin{equation*}
    \mathbf{s'} = \sigma\left( \mathbf{W}_m\begin{bmatrix}\norm{\mathbf{W}_h\mathbf{V}}\\ \mathbf{s}\end{bmatrix} + \mathbf{b}\right)
\end{equation*}
Since $\norm{\textbf{W}_h \mathbf{V} \mathbf{U}} = \norm{\textbf{W}_h \mathbf{V}}$, we conclude $\mathbf{s'}$ is invariant.
Similarly the vector output is
\begin{equation*}
    \mathbf{V'} = \sigma^+\left(\norm{\textbf{W}_\mu \textbf{W}_h \mathbf{V}}\right)\odot \textbf{W}_\mu \textbf{W}_h \mathbf{V}
\end{equation*}
The row-wise scaling can also be viewed as left-multiplication by a diagonal matrix $\mathbf{D}$. Since $\norm{\textbf{W}_\mu \textbf{W}_h \mathbf{V}} = \norm{\textbf{W}_\mu \textbf{W}_h \mathbf{VU}}$, $\mathbf{D}$ is invariant. Since
\begin{equation*}
    \mathbf{DW}_\mu\mathbf{W}_h\mathbf{(VU)} = \left(\mathbf{DW}_\mu\mathbf{W}_h\mathbf{V}\right)\mathbf{U}
\end{equation*}
we conclude that $\mathbf{V'}$ is equivariant. 
\end{proof}


\subsection*{The ability to approximate rotation-invariant functions}

The GVP inherits an analogue of the Universal Approximation property \cite{cybenko1989approximation} of standard dense layers. If $R$ describes an arbitrary rotation or reflection in 3D Euclidean space,\footnote{More precisely, if $R$ describes a unitary transformation.} we show that a GVP can approximate arbitrary scalar-valued functions invariant under $R$ and defined over $\Omega^\nu \subset \mathbb{R}^{\nu\times 3}$, the bounded subset of $\mathbb{R}^{\nu \times 3}$ whose elements can be canonically oriented based on three linearly independent vector entries. Without loss of generality, we assume the first three vector entries can be used.

The machinery corresponding to such approximations corresponds to a GVP $G_s$ with only vector inputs, only scalar outputs, and a sigmoidal nonlinearity $\sigma$, followed by a dense layer. This can also be viewed as the sequence of matrix multiplication with $\mathbf{W}_h$, taking the L$_2$-norm, and a dense network with one hidden layer. Such machinery can be extracted from any two consecutive GVPs (assuming a sigmoidal $\sigma$). The theorem is restated from the main text:

\begin{theorem}
    Let R describe an arbitrary rotation or reflection in $\mathbb{R}^3$.
    For $\nu \ge 3$ let $\Omega^\nu \subset \mathbb{R}^{\nu \times 3}$ be the set of all $\mathbf{V} = \begin{bmatrix}\mathbf{v}_1, & \ldots, & \mathbf{v}_\nu\end{bmatrix}^T\in \mathbb{R}^{\nu \times 3}$ such that $\mathbf{v}_1, \mathbf{v}_2, \mathbf{v}_3$ are linearly independent and $0<\norm{\mathbf{v}_i} \le b$ for all $i$ and some finite $b>0$. Then for any continuous $F: \Omega^\nu \rightarrow \mathbb{R}$ such that $F(R( \mathbf{V})) = F(\mathbf{V})$ and for any $\epsilon>0$, there exists a form $f(\mathbf{V})=\mathbf{w}^TG_s(\mathbf{V})$ such that $|F(\mathbf{V})-f(\mathbf{V})|< \epsilon$ for all $\mathbf{V}\in\Omega$.
\end{theorem}
\begin{proof} 
    The idea is to write $F$ as a composition $F=\Tilde{F} \circ \omega$ and $\omega = h \circ y$.  We show that multiplication with $\mathbf{W}_h$ and and taking the L$_2$-norm can compute $y$, and that the dense network with one hidden layer can approximate $\Tilde{F} \circ h$.
    
    Call an element $\mathbf{V} \in \Omega^\nu$ \textit{oriented} if $\mathbf{v}_1 = x_1 \mathbf{e}_x$, $\mathbf{v}_2 = x_2\mathbf{e}_x + y_2\mathbf{e}_y$, and $\mathbf{v}_3 = x_3\mathbf{e}_x + y_3\mathbf{e}_y+z_3\mathbf{e}_z$, with $x_1, y_2, z_3>0$.
    Define $\omega: \Omega^\nu \rightarrow \mathbb{R}^{3\nu-3}$ to be the \textit{orientation function} that orients its input and then extracts the vector of $3\nu-3$ coefficients, $[x_1, x_2, y_2, x_3, y_3, z_3, \ldots, x_i, y_i, z_i, \ldots]^T$. These elements can be written as
    \begin{eqnarray*}
        x_1 &=& \norm{\mathbf{v}_1}  \\
        x_i &=& \mathbf{v}_i \cdot \mathbf{v}_1/x_1, \quad i \ge 2\\
        y_2 &=& \sqrt{\norm{\mathbf{v}_2}^2 -x_2^2} \\
        y_i &=& (\mathbf{v}_i \cdot \mathbf{v}_2 - x_ix_2)/y_2, \quad i \ge 3\\
        z_3 &=& \sqrt{\norm{\mathbf{v}_3}^2 - x_3^2 - y_3^2}\\
        z_i &=& (\mathbf{v}_i \cdot \mathbf{v}_3 - x_ix_3 - y_iy_3)/z_3, \quad i \ge 4
    \end{eqnarray*}
    and are invariant under rotation and reflection, because they are defined using only the norms and inner products of the $\mathbf{v}_i$. Then $F = \tilde{F}\circ \omega$, where $\tilde{F}: [-b, b]^{3\nu-3}\rightarrow \mathbb{R}$.
    
    The key insight is that if we construct $\mathbf{W}_h$ such that the rows of $\mathbf{W}_h\mathbf{V}$ are the original vectors $\mathbf{v}_i, \forall i$ and all differences $\mathbf{v}_i - \mathbf{v}_j, \forall i, j\le \min(i,3)$, then we can compute $\omega(\mathbf{V})$ from the row-wise norms of $\mathbf{W}_h\mathbf{V}$. That is, $\omega = h \circ y$ where $\mathbf{y} = y(\mathbf{V}) = \norm{^.} \odot (\mathbf{W}_h\mathbf{V})\in\mathbb{R}^{4\nu-6}$ and $h$ is an application of the cosine law. The GVP precisely computes $\mathbf{y}$ as an intermediate step: $G_s(\mathbf{V}) = \sigma \odot (\mathbf{W}_m\mathbf{y} +\mathbf{b})$. 
    It remains to show that there exists a form $\tilde{f}(\mathbf{y}) = \mathbf{w}^T[\sigma \odot(\mathbf{W}_h\mathbf{y} + \mathbf{b})]$ that $\epsilon$-approximates $\tilde{F} \circ h : [-2b, 2b]^{4\nu-6}\rightarrow\mathbb{R}$. Up to a translation and uniform scaling of the hypercube, this is the result of the Universal Approximation Theorem \citep{cybenko1989approximation}.
\end{proof}


