Original: G ENERATIVE MODELS FOR  PROTEIN DESIGN
Revised: G ENERATIVE MODELS FOR GRAPH - BASED PROTEIN DESIGN

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Engineered proteins offer the potential to solve many problems in biomedicine, energy, and materials science, but creating designs that succeed is difﬁcult in practice.
Revised: Engineered proteins offer the potential to solve many problems in biomedicine, energy, and materials science, but creating designs that succeed is difﬁcult in practice.

Original: A signiﬁcant aspect of this challenge is the complex coupling between protein sequence and 3D structure, and the task of ﬁnding a viable design is often referred to as the inverse protein folding problem .
Revised: A signiﬁcant aspect of this challenge is the complex coupling between protein sequence and 3D structure, and the task of ﬁnding a viable design is often referred to as the inverse protein folding problem .

Original: We develop generative models for protein sequences conditioned on a graph-structured speciﬁcation of the design target.
Revised: We develop generative models for protein sequences conditioned on a graph-structured speciﬁcation of the design target.

Original: Our approach efﬁciently captures the complex dependency structures in proteins by focusing on dependencies that are long-range in sequence but local in 3D space.
Revised: Our approach efﬁciently captures the complex dependencies in proteins by focusing on those that are long-range in sequence but local in 3D space.

Original: Our framework signiﬁcantly improves upon prior parametric models of protein sequences given structure, and takes a step toward rapid and targeted biomolecular design with the aid of deep generative models.
Revised: Our framework signiﬁcantly improves upon prior parametric models of protein sequences given structure, and takes a step toward rapid and targeted biomolecular design with the aid of deep generative models.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: A central goal for computational protein design is to automate the invention of protein molecules with deﬁned structural and chemical properties.
Revised: A central goal for computational protein design is to automate the invention of protein molecules with deﬁned structural and functional properties.

Original: This ﬁeld has seen tremendous progess in the past two decades (Huang et al., 2016), including the design of novel 3D folds (Kuhlman et al., 2003), enzymes (Siegel et al., 2010), and complexes (Bale et al., 2016).
Revised: This ﬁeld has seen tremendous progess in the past two decades (Huang et al., 2016), including the design of novel 3D folds (Kuhlman et al., 2003), enzymes (Siegel et al., 2010), and complexes (Bale et al., 2016).

Original: However, the current practice often requires multiple rounds of trial-and-error, with ﬁrst designs frequently failing (Koga et al., 2012;
Revised: However, the current practice often requires multiple rounds of trial-and-error, with ﬁrst designs frequently failing (Koga et al., 2012;

Original: Rocklin et al., 2017).
Revised: Rocklin et al., 2017).

Original: Several of the challenges stem from the bottom-up nature of contemporary approaches that rely on both the accuracy of energy functions to describe protein physics as well as on the efﬁciency of sampling algorithms to explore the protein sequence and structure space.
Revised: Several of the challenges stem from the bottom-up nature of contemporary approaches that rely on both the accuracy of energy functions to describe protein physics as well as on the efﬁciency of sampling algorithms to explore the protein sequence and structure space.

Original: Here, we explore an alternative, top-down framework for protein design that directly learns a conditional generative model for protein sequences given a speciﬁcation of the target structure, which is represented as a graph over the sequence elements.
Revised: Here, we explore an alternative, top-down framework for protein design that directly learns a conditional generative model for protein sequences given a speciﬁcation of the target structure, which is represented as a graph over the sequence elements.

Original: Speciﬁcally, we augment the autoregressive self-attention of recent sequence models (Vaswani et al., 2017) with graph-based descriptions of the 3D structure.
Revised: Speciﬁcally, we augment the autoregressive self-attention of recent sequence models (Vaswani et al., 2017) with graph-based descriptions of the 3D structure.

Original: By composing multiple layers of structured self-attention, our model can effectively capture higher-order, interaction-based dependencies between sequence and structure, in contrast to previous parameteric approaches (O’Connell et al., 2018; Wang et al., 2018) that are limited to only the ﬁrst-order effects.
Revised: By composing multiple layers of structured self-attention, our model can effectively capture higher-order, interaction-based dependencies between sequence and structure, in contrast to previous parameteric approaches (O’Connell et al., 2018; Wang et al., 2018) that are limited to only the ﬁrst-order effects.

Original: The graph-structured conditioning of a sequence model affords several beneﬁts, including favorable computational efﬁciency, inductive bias, and representational ﬂexibility.
Revised: The graph-structured conditioning of a sequence model affords several beneﬁts, including favorable computational efﬁciency, inductive bias, and representational ﬂexibility.

Original: We accomplish the ﬁrst two by leveraging a well-evidenced ﬁnding in protein science, namely that long-range dependencies in sequence are generally short-range in 3D space (Morcos et al., 2011; Marks et al., 2011;
Revised: We accomplish the ﬁrst two by leveraging a well-evidenced ﬁnding in protein science, namely that long-range dependencies in sequence are generally short-range in 3D space (Marks et al., 2011; Morcos et al., 2011;

Original: Balakrishnan et al., 2011).
Revised: Balakrishnan et al., 2011).

Original: By making the graph and self-attention similarly sparse and localized in 3D space, we achieve computational scaling that is linear in sequence length.
Revised: By making the graph and self-attention similarly sparse and localized in 3D space, we achieve computational scaling that is linear in sequence length.

Original: Additionally, graph structured inputs offer representational ﬂexibility, as they accomodate both coarse, ‘ﬂexible backbone’ (connectivity and topology) as well as ﬁne-grained (precise atom locations) descriptions of structure.
Revised: Additionally, graph structured inputs offer representational ﬂexibility, as they accomodate both coarse, ‘ﬂexible backbone’ (connectivity and topology) as well as ﬁne-grained (precise atom locations) descriptions of structure.

Original: We demonstrate the merits of our approach via a detailed empirical study.
Revised: We demonstrate the merits of our approach via a detailed empirical study.

Original: Speciﬁcally, we evaluate our model at structural generalization to sequences of protein folds that were outside of the training set.
Revised: Speciﬁcally, we evaluate our model at structural generalization to sequences of protein folds that were outside of the training set.

Original: Our model achieves considerably improved generalization performance over the recent deep models of protein sequence given structure as well as structure-na¨ıve language models.
Revised: Our model achieves considerably improved generalization performance over the recent deep models of protein sequence given structure as well as structure-na¨ıve language models.

Original: 1. R ELATED W ORK
Revised: 1. R ELATED W ORK

Original: Generative models for proteins
Revised: Generative models for proteins

Original: A number of works have explored the use of generative models for protein engineering and design (Yang et al., 2018).
Revised: A number of works have explored the use of generative models for protein engineering and design (Yang et al., 2018).

Original: Recently O’Connell et al.
Revised: Recently O’Connell et al.

Original: (2018) and Wang et al.
Revised: (2018) and Wang et al.

Original: (2018) proposed neural models for sequences given 3D structure, where the amino acids at different positions in the sequence are predicted independently of one another.
Revised: (2018) proposed neural models for sequences given 3D structure, where the amino acids at different positions in the sequence are predicted independently of one another.

Original: Greener et al.
Revised: Greener et al.

Original: (2018) introduced a generative model for protein sequences conditioned on a 1D, context-free grammar based speciﬁcation of the fold topology.
Revised: (2018) introduced a generative model for protein sequences conditioned on a 1D, context-free grammar based speciﬁcation of the fold topology.

Original: Boomsma & Frellsen (2017) and Weiler et al.
Revised: Boomsma & Frellsen (2017) and Weiler et al.

Original: (2018) used deep neural networks to model the conditional distribution of letters in a speciﬁc position given the structure and sequence of all surrounding residues.
Revised: (2018) used deep neural networks to model the conditional distribution of letters in a speciﬁc position given the structure and sequence of all surrounding residues.

Original: In contrast to these works, our model captures the joint distribution of the full protein sequence while grounding these dependencies in terms of long-range interactions arising from the structure.
Revised: In contrast to these works, our model captures the joint distribution of the full protein sequence while grounding these dependencies in terms of long-range interactions arising from the structure.

Original: In parallel to the development of structure-based models, there has been considerable work on deep generative models for protein sequences in individual protein families with directed (Riesselman et al., 2018;
Revised: In parallel to the development of structure-based models, there has been considerable work on deep generative models for protein sequences in individual protein families with directed (Riesselman et al., 2018;

Original: Sinai et al., 2017) and undirected (Tubiana et al., 2018) latent variable models.
Revised: Sinai et al., 2017) and undirected (Tubiana et al., 2018) latent variable models.

Original: These methods have proven useful for protein engineering, but presume the availability of a large number of sequences from a particular family.
Revised: These methods have proven useful for protein engineering, but presume the availability of a large number of sequences from a particular family.

Original: In a complementary line of research, deep models of protein structure (Anand & Huang, 2018;
Revised: In a complementary line of research, deep models of protein structure (Anand & Huang, 2018;

Original: Ingraham et al., 2019; AlQuraishi, 2018) have been proposed recently that could be used to craft 3D structures for input to sequence design.
Revised: Ingraham et al., 2019; AlQuraishi, 2018) have been proposed recently that could be used to craft 3D structures for input to sequence design.

Original: Protein design For classical approaches to computational protein design, which are based on joint modeling of structure and sequence, we refer the reader to a review of both methods and accomplishments in Huang et al.
Revised: Protein design For classical approaches to computational protein design, which are based on joint modeling of structure and sequence, we refer the reader to a review of both methods and accomplishments in Huang et al.

Original: More recently, Zhou et al.
Revised: More recently, Zhou et al.

Original: (2018) proposed a non-parametric approach to protein design in which a target design is decomposed into substructural motifs that are then queried against a protein database.
Revised: (2018) proposed a non-parametric approach to protein design in which a target design is decomposed into substructural motifs that are then queried against a protein database.

Original: In this work we will focus on comparisons with direct parametric models of the sequence-structure relationship.
Revised: In this work we will focus on comparisons with direct parametric models of the sequence-structure relationship.

Original: Self-Attention
Revised: Self-Attention

Original: Our model extends the Transformer (Biswas et al., 2018) to additionally capture sparse, pairwise relational information between sequence elements.
Revised: Our model extends the Transformer (Biswas et al., 2018) to additionally capture sparse, pairwise relational information between sequence elements.

Original: The dense variation of this problem was explored in Shaw et al.
Revised: The dense variation of this problem was explored in Shaw et al.

Original: and
Revised: and

Original: Huang et al.
Revised: Huang et al.

Original: As noted in those works, incorporating general pairwise information incurs O ( L 2 ) memory (and computational) cost for sequences of length L , which can be highly limiting for training on GPUs.
Revised: As noted in those works, incorporating general pairwise information incurs O ( N 2 ) memory (and computational) cost for sequences of length N , which can be highly limiting for training on GPUs.

Original: We circumvent this cost by instead restricting the self-attention to the sparsity of the input graph.
Revised: We circumvent this cost by instead restricting the self-attention to the sparsity of the input graph.

Original: Given this graph-structured self-attention, our model may also be reasonably cast in the framework of message-passing or graph neural networks (Gilmer et al., 2017; Battaglia et al., 2018).
Revised: Given this graph-structured self-attention, our model may also be reasonably cast in the framework of message-passing or graph neural networks (Gilmer et al., 2017; Battaglia et al., 2018).

Revised: Our approach is similar to Graph Attention Networks (Veli ˇ ckovi ´ c et al., 2017), but augmented with edge features and an autoregressive decoder.

Original:  M ETHODS 2. R EPRESENTING STRUCTURE
Revised:  M ETHODS 2. R EPRESENTING STRUCTURE

Original: We represent molecular structure in terms of an attributed graph G =
Revised: We represent protein structure in terms of an attributed graph G =

Original: ( V , E ) with node features V = { v 1 , . . 
Revised: ( V , E ) with node features V = { v 1 , . . .

Original: , v N } and edge features E = { e ij } i (cid:54) =
Revised: , v N } and edge features E = { e ij } i (cid:54) =

Original: j .
Revised: j over the sequence residues (amino acids).

Original: This formulation can accommodate different variations on the macromolecular design problem, including both the ‘rigid backbone’ design where the precise coordinates of backbone atoms are ﬁxed, as well as the ‘ﬂexible backbone’ design where looser constraints such as blueprints of hydrogen-bonding connectivity (Koga et al., 2012) or 1D architectures (Greener et al., 2018) could deﬁne the structure of interest.
Revised: This formulation can accommodate different variations on the macromolecular design problem, including both the ‘rigid backbone’ design where the precise coordinates of backbone atoms are ﬁxed, as well as the ‘ﬂexible backbone’ design where looser constraints such as blueprints of hydrogen-bonding connectivity (Koga et al., 2012) or 1D architectures (Greener et al., 2018) could deﬁne the structure of interest.

Original: 3D considerations For a rigid-body design problem, the structure for conditioning is a ﬁxed set of backbone coordinates X
Revised: 3D considerations For a rigid-body design problem, the structure for conditioning is a ﬁxed set of backbone coordinates X

Original: = { x i ∈ R 3 : 1 ≤ i ≤ N } , where N is the number of positions 1 .
Revised: = { x i ∈ R 3 : 1 ≤ i ≤ N } , where N is the number of positions 1 .

Original: We desire a graph representation of the coordinates G ( X ) that has two properties: 
Revised: We desire a graph representation of the coordinates G ( X ) that has two properties: [Figure]

Original: • Invariance .
Revised: • Invariance .

Original: The features are invariant to rotations and translations.
Revised: The features are invariant to rotations and translations.

Original: [Figure]

Original: • Locally informative.
Revised: • Locally informative.

Original: The edge features incident to v i due to its neighbors N ( i ) , i e. { e ij } j ∈ N ( i ) , contain sufﬁcient information to reconstruct all adjacent coordinates { x j } j ∈ N ( i ) up to rigid-body motion.
Revised: The edge features incident to v i due to its neighbors N ( i ) , i.e. { e ij } j ∈ N ( i ) , contain sufﬁcient information to reconstruct all adjacent coordinates { x j } j ∈ N ( i ) up to rigid-body motion.

Original: While invariance is motivated by standard symmetry considerations the second property is motivated bythe limitations of current graph neural networks (Gilmer et al., 2017).
Revised: While invariance is motivated by standard symmetry considerations the second property is motivated by limitations of current graph neural networks (Gilmer et al., 2017).

Original: In these networks, updates to node features v
Revised: In these networks, updates to node features v

Original: i depend only on the edge and node features adjacent to v i .
Revised: i depend only on the edge and node features adjacent to v i .

Original: However, typically, these features are insufﬁcient to reconstruct the relative neighborhood positions { x j } j ∈ N ( i ) , so individual updates cannot fully depend on the ‘local environment’.
Revised: However, typically, these features are insufﬁcient to reconstruct the relative neighborhood positions { x j } j ∈ N ( i ) , so individual updates cannot fully depend on the ‘local environment’.

Original: For example, pairwise distances D ij and D il are insufﬁcient to determine if x j and x l are on the same or opposite sides of x i .
Revised: For example, pairwise distances D ij and D il are insufﬁcient to determine if x j and x l are on the same or opposite sides of x i .

Original: Structural encodings We develop invariant and locally informative features by ﬁrst augmenting the points x i with ‘orientations’ O i that deﬁne a local coordinate system at each point.
Revised: Structural encodings We develop invariant and locally informative features by ﬁrst augmenting the points x i with ‘orientations’ O i that deﬁne a local coordinate system at each point.

Original: We deﬁne these in terms of the backbone geometry as
Revised: We deﬁne these in terms of the backbone geometry as

Original: [Equation]
Revised: [Equation]

Original: where b i is the negative bisector of angle between the rays ( x i − 1 − x i ) and ( x i +1 − x i ) , and n i is a unit vector normal to that plane.
Revised: where b i is the negative bisector of angle between the rays ( x i − 1 − x i ) and ( x i +1 − x i ) , and n i is a unit vector normal to that plane.

Original: Formally, we have
Revised: Formally, we have

Original: [Equation]
Revised: [Equation]

Original: Finally, we derive the spatial edge features e ( s ) ij from the rigid body transformation that relates reference frame ( x i , O i ) to reference frame ( x j , O j ) .
Revised: Finally, we derive the spatial edge features e ( s ) ij from the rigid body transformation that relates reference frame ( x i , O i ) to reference frame ( x j , O j ) .

Original: While this transformation has 6 degrees of freedom, we decompose it into features for distance, direction, and orientation as
Revised: While this transformation has 6 degrees of freedom, we decompose it into features for distance, direction, and orientation as

Original: [Equation]
Revised: [Equation]

Original: Here r ( · ) is a function that lifts the distances into a radial basis 2 , the term in the middle corresponds to the relative direction of x j in the reference frame of ( x i , O i ) , and q ( · ) converts the 3 × 3 relative rotation matrix to a quaternion representation.
Revised: Here r ( · ) is a function that lifts the distances into a radial basis 2 , the term in the middle corresponds to the relative direction of x j in the reference frame of ( x i , O i ) , and q ( · ) converts the 3 × 3 relative rotation matrix to a quaternion representation.

Original: Quaternions represent rotations as four-element vectors that can be efﬁciently and reasonably compared by inner products Huynh (2009).
Revised: Quaternions represent rotations as four-element vectors that can be efﬁciently and reasonably compared by inner products Huynh (2009).

Original: Positional encodings Taking a cue from the original Transformer model, we obtain positional embeddings e ( p ) ij that encode the role of local structure around node i .
Revised: Positional encodings Taking a cue from the original Transformer model, we obtain positional embeddings e ( p ) ij that encode the role of local structure around node i .

Original: Speciﬁcally, we need to model the positioning of each neighbor j relative to the node under consideration i .
Revised: Speciﬁcally, we need to model the positioning of each neighbor j relative to the node under consideration i .

Original: Therefore, we obtain the position embedding as a sinusoidal function of the gap i − j .
Revised: Therefore, we obtain the position embedding as a sinusoidal function of the gap i − j .

Original: Note that this is instark contrast to the absolute positional encodings of the original Transformer, and instead matches the relative encodings in Shaw et al.
Revised: Note that this is in contrast to the absolute positional encodings of the original Transformer, and instead matches the relative encodings in Shaw et al.

Original: Node and edge features Finally, we obtain an aggregate edge encoding vector e ij by concatenating the structural encodings e ( s ) ij and the positional encodings e ( p ) ij.
Revised: Node and edge features Finally, we obtain an aggregate edge encoding vector e ij by concatenating the structural encodings e ( s ) ij and the positional encodings e ( p ) ij

Revised: and then linearly transforming them to have the same dimension as the model.

Original: We only include edges in the k -nearest neighbors graph of X , with k = 30 for all experiments.
Revised: We only include edges in the k -nearest neighbors graph of X , with k = 30 for all experiments.

Original: For node features, we compute the three dihedral angles of the protein backbone ( φ i , ψ i , ω i ) and embed these on the 3-torus as { sin , cos } × ( φ i , ψ i , ω i ) .
Revised: For node features, we compute the three dihedral angles of the protein backbone ( φ i , ψ i , ω i ) and embed these on the 3-torus as { sin , cos } × ( φ i , ψ i , ω i ) .

Revised: Flexible backbone features We also consider ’ﬂexible backbone’ descriptions of 3D structure based solely on topological binary edge features.

Revised: We combine the relative positional encodings with two binary edge features: contacts that indicate when the distance between C α residues at i and j are less than 8 Angstroms and hydrogen bonds which are directed and deﬁned by the electrostatic model of DSSP (Kabsch & Sander, 1983).

Revised: These features implicitly integrate over different 3D backbone conﬁgurations that are compatible with the speciﬁed topology.

Original: 2. S TRUCTURED T RANSFORMER
Revised: 2. S TRUCTURED T RANSFORMER

Original: In this work, we introduce a Structured Transformer model that draws inspiration from the selfattention based Transformer model (Vaswani et al., 2017).
Revised: In this work, we introduce a Structured Transformer model that draws inspiration from the selfattention based Transformer model (Vaswani et al., 2017) and is augmented for scalable incorporation of relational information.

Original: First, since L is often very large (several hundred amino acids being common), global all-vs-all attention leads to prohibitive computational and memory bottlenecks. We avert thesebottlenecks by restricting  attention for each node i to the set
Revised: While general relational attention incurs quadratic memory and computation costs, we avert these by restricting the attention for each node i to the set

Original: N ( i, k ) of its k -nearest neighbors in 3D space.
Revised: N ( i, k ) of its k -nearest neighbors in 3D space.

Original: Since our architecture is multilayered, iterated local attention can derive progressively more global estimates of context for each node i .
Revised: Since our architecture is multilayered, iterated local attention can derive progressively more global estimates of context for each node i .

Original: Second, unlike the standard Transformer, we also include edge features to embed the spatial and positional dependencies in deriving the attention.
Revised: Second, unlike the standard Transformer, we also include edge features to embed the spatial and positional dependencies in deriving the attention.

Original: Thus, our model generalizes Transformer to spatially structured settings.
Revised: Thus, our model generalizes Transformer to spatially structured settings.

Revised: Autoregressive decomposition We decompose the joint distribution of the sequence p ( s ) autoregressively as

Revised: [Equation]

Revised: where the conditional probability p ( s i | x , s <i ) of amino acid

Revised: s i at position i is conditioned on both the input structure x and the preceding amino acids s <

Original:  Encoder Our encoder module is designed as follows.
Revised: i = { s 1 , . . . s i − 1 } 4 . These conditionals are parameterized in terms of two sub-networks: an encoder that computes reﬁned node embeddings from structure-based node features V ( x ) and edge features E ( x ) and a decoder that autoregressively predicts letter s i given the preceding sequence and structural embeddings from the encoder. Encoder Our encoder module is designed as follows.

Original: A transformation W h : R d v (cid:55)→ R d produces initial embeddings
Revised: A transformation W h : R d v (cid:55)→ R d produces initial embeddings

Original: h i =
Revised: h i =

Original: W h ( v i ) from the node features v i pertaining to position i ∈
Revised: W h ( v i ) from the node features v i pertaining to position i ∈

Original: [ N ] (cid:44) { 1 , 2 , . . .
Revised: [ N ] (cid:44) { 1 , 2 , . . .

Original: , N } .
Revised: , N } .

Revised: We anticipate that alternative orderings for decoding the sequence may be favorable but leave this to future work

Original:  Each layer of the encoder implements a multi-head self-attention component, where head (cid:96) ∈
Revised: [Table] Each layer of the encoder implements a multi-head self-attention component, where head (cid:96) ∈

Original: [ L ] attends to a separate subspace of the embeddings. We write h i = Concat (cid:16) h (1) i , . . .
Revised: [ L ] can attend to a separate subspace of the embeddings via learned query, key and value transformations (Vaswani et al., 2017).

Original: , h ( L ) i (cid:17) , where each h ( (cid:96) ) i ∈ R d/L , to emphasize the subspaces. We reﬁne the embeddings via additional learned transformations.
Revised: The queries are derived from the current embedding at node i while the keys and values from the relational information r ij = ( h j , e ij ) at adjacent nodes j ∈ N ( i, k ) .

Original: Speciﬁcally, W ( (cid:96) )
Revised: Speciﬁcally, W ( (cid:96) )

Original: q maps h( (cid:96) ) i to query embeddings q ( (cid:96) )
Revised: q maps h i to query embeddings q ( (cid:96) )

Original: i , W ( (cid:96) )
Revised: i , W ( (cid:96) )

Original: z maps pairs ( h ( (cid:96) ) j , e ij) to key embeddings z ( (cid:96) )
Revised: z maps pairs r ij to key embeddings z ( (cid:96) )

Original: ij for j ∈ N ( i, k ) , and W ( (cid:96) ) v maps the same pairs ( h ( (cid:96) ) j , e ij) to value embeddings v ( (cid:96) ) ij for each i ∈ [ N ] , (cid:96) ∈
Revised: ij for j ∈ N ( i, k ) , and W ( (cid:96) ) v maps the same pairs r ij to value embeddings v ( (cid:96) ) ij for each i ∈ [ N ] , (cid:96) ∈

Original: [ L ] .
Revised: [ L ] .

Original: Decoupling the mappings for keys and values allows each to depend on different subspaces of the representation.
Revised: Decoupling the mappings for keys and values allows each to depend on different subspaces of the representation.

Original: We compute the similarity of key z ( (cid:96) ) ij to query q ( (cid:96) ) 
Revised: We compute the attention a ( (cid:96) ) ij between query q ( (cid:96) ) i and key z ( (cid:96) )

Original: i as a function of their scaled inner product:
Revised: ij as a function of their scaled inner product:

Original: [Equation]
Revised: [Equation]

Original: where m ( (cid:96)) ij = q ( (cid:96) )
Revised: where m ( (cid:96)

Original: i (cid:62) z ( (cid:96) ) ij √ d .
Revised: [Equation]

Original: This similarity is interpreted as the importance of v ( (cid:96) ) ij in reﬁning h ( (cid:96) )
Revised: The results of each attention head l are collected as the weighted sum h ( (cid:96) )

Original: i , and we obtain the updates
Revised: i = (cid:88) j ∈ N ( i,k

Original: [Equation]
Revised: [Equation]

Original: Thus, effectively, we reﬁne embeddings in each subspace separately via output of a residual connection at each layer.
Revised: a ( (cid:96) ) ij v ( (cid:96) ) ij and then concatenated and transformed to give the update ∆ h i = W o Concat (cid:16) h (1) i , . . , h ( L ) i (cid:17) .

Original: We  alternate between these self-attention layers and position-wise feedforward[Table] layers as in the original Transformer (Vaswani et al., 2017).
Revised: We update the embeddings with this residual and alternate between these self-attention layers and position-wise feedforward layers as in the original Transformer (Vaswani et al., 2017).

Original: We stack multiple layers atop each other, and thereby obtain continually reﬁned embeddings as we traverse the layers bottom up.
Revised: We stack multiple layers atop each other, and thereby obtain continually reﬁned embeddings as we traverse the layers bottom up.

Original: The encoder yields the embeddings produced by the topmost layer as its output.
Revised: The encoder yields the embeddings produced by the topmost layer as its output.

Original: Decoder Our decoder module is laid out similarly but for some important differences. First, we augment the edge information e ij with s j , i.e., the j th amino acid in  the output sequence. This combination is then propagated bottom up in the decoder stack analogous to the role of e ij during encoding.
Revised: Decoder Our decoder module has the same structure as the encoder but with augmented relational information r ij that allows access to the preceding sequence elements s <i in a causally consistent manner. Whereas the keys and values of the encoder are based on the relational information r ij = ( h j , e ij ) , the decoder can additionally access sequence elements s j as

Original: Second, unlike encoding, we perform decoding autoregressively , i.e., conditioned on the partial sequnence already decoded. Therefore, while training, we mask the information pertaining to the neighbor nodes subsequent to the position being decoded.
Revised: [Equation]

Original: Finally, as in the Transformer decoder (Vaswani et al., 2017), the embeddings output by the encoder are treated as keys that dictate, inpart, the embeddings output by the decoder.
Revised: Here h ( dec ) j is the embedding of node j in the current layer of the decoder, h ( enc ) j is the embedding of node j in the ﬁnal layer of the encoder, and g ( s j ) is a sequence embedding of amino acid s j at node j .

Original: The general layout of our structural decoder, modulo the inﬂux of augmented edge information and restriction to k -neighborhood, is similar to  the  decoder in the original Transformer setting (Vaswani et al., 2017) so we omit the details.
Revised: This concatenation and masking structure ensures that sequence information only ﬂows to position i from positions j < i , but still allows position i to attend to subsequent structural information. We stack three layers of self-attention and position-wise feedforward modules for the encoder and decoder with a hidden dimension of 128 throughout the experiments 5 .

Original: 2. T RAINING
Revised: 2. T RAINING

Original: Dataset To evaluate the ability of the models to generalize across different protein folds, we collected a dataset based on the CATH hierarchical classiﬁcation of protein structure (Orengo et al., 1997).
Revised: Dataset To evaluate the ability of the models to generalize across different protein folds, we collected a dataset based on the CATH hierarchical classiﬁcation of protein structure (Orengo et al., [Table]

Original: For all domains in the CATH 4.2 40% non-redundant set of proteins, we obtained full chains up to length 500 (which may contain more than one domain) and then cluster-split these at the 
Revised: For all domains in the CATH 4.2 40% non-redundant set of proteins, we obtained full chains up to length 500 (which may contain more than one domain) and then cluster-split these at the CATH topology level (i

Original: CATH topology level (i.e. fold level) into training, validation, and test sets at an 80/10/10 split.
Revised: [Table] e. fold level) into training, validation [Table] and test sets at an 80/10/10 split.

Original: Chains containing multiple CATH tpologies were purged with precedence for test over validation over train.
Revised: Chains containing multiple CATH tpologies were purged with precedence for test over validation over train.

Original: Our splitting procedure ensured that no two domains from different sets would share the same topologies (folds).
Revised: Our splitting procedure ensured that no two domains from different sets would share the same topologies (folds).

Original: The ﬁnal splits contained 18025 chains in the training set, 1637 chains in the validation set, and 1911 chains in the test set.
Revised: The ﬁnal splits contained 18025 chains in the training set, 1637 chains in the validation set

Original: Optimization We trained models using the learning rate schedule and initialization of (Vaswani et al., 2017), a dropout (Srivastava et al., 2014) rate of 10% , and early stopping based on validation perplexity.

Original: [Table]
Revised: [Table]

Original: models
Revised: and 1911 chains in the test set.

Original: [Table] models [Table]
Revised: Optimization We trained models using the learning rate schedule and initialization of (Vaswani et al., 2017), a dropout (Srivastava et al., 2014) rate of 10% , and early stopping based on validation perplexity.

Original:  R ESULTS
Revised:  R ESULTS

Original: Many protein sequences may reasonably design the same 3D structure (Li et al., 1996), and so we focus on likelihood-based evaluations of model performance.
Revised: Many protein sequences may reasonably design the same 3D structure (Li et al., 1996), and so we focus on likelihood-based evaluations of model performance.

Original: Speciﬁcally, we evaluate the perplexity per letter of test protein folds (topologies) that were held out from the training and validation sets.
Revised: Speciﬁcally, we evaluate the perplexity per letter of test protein folds (topologies) that were held out from the training and validation sets.

Original: Protein perplexities What kind of perplexities might be useful?
Revised: Protein perplexities What kind of perplexities might be useful?

Original: To provide context, we ﬁrst present perplexities for some simple models of protein sequences in Table 1.
Revised: To provide context, we ﬁrst present perplexities for some simple models of protein sequences in Table 1.

Original: The amino acid alphabet and its natural frequencies upper-bound perplexity at 20 and ∼ 17.8, respectively.
Revised: The amino acid alphabet and its natural frequencies upper-bound perplexity at 20 and ∼ 17.8, respectively.

Original: Random protein sequences under these null models are unlikely to be functional without further selection (Keefe & Szostak, 2001).
Revised: Random protein sequences under these null models are unlikely to be functional without further selection (Keefe & Szostak, 2001).

Original: First order proﬁles of protein sequences such as those from the Pfam database (El-Gebali et al., 2018), however, are widely used for protein engineering.
Revised: First order proﬁles of protein sequences such as those from the Pfam database (El-Gebali et al., 2018), however, are widely used for protein engineering.

Original: We found the average perplexity per letter of proﬁles in Pfam 32 (ignoring alignment uncertainty) to be ∼ 11.6.
Revised: We found the average perplexity per letter of proﬁles in Pfam 32 (ignoring alignment uncertainty) to be ∼ 11.6.

Original: This suggests that even models with high perplexities of this order have the potential to be useful models for the space of functional protein sequences.
Revised: This suggests that even models with high perplexities of this order have the potential to be useful models for the space of functional protein sequences.

Revised: [Table]

Original: The importance of structure We found that there was a signiﬁcant gap between unconditional language models of protein sequences and models conditioned on structure.
Revised: The importance of structure We found that there was a signiﬁcant gap between unconditional language models of protein sequences and models conditioned on structure.

Original: Remarkably, for a range of structure-independent language models, the typical test perplexities turned out to be ∼ 1617 (Table 2), which were barely better than null letter frequencies (Table 1). 
Revised: Remarkably, for a range of structure-independent language models, the typical test perplexities turned out to be ∼ 1617 (Table 2), which were barely better than null letter frequencies (Table 1). We emphasize that the

Revised: RNNs were not broken and could still learn the training set in these capacity ranges.

Revised: It would seem that protein language models trained on one subset of 3D folds (in our cluster-splitting procedure) generalize poorly to predict the sequences of unseen folds, which is important to consider when training protein language models for protein engineering and design.

Original: All structure-based models were observed to have a considerably lower perplexity.
Revised: All structure-based models had (unsurprisingly) considerably lower perplexities.

Original: In particular, our
Revised: In particular, our

Original: Structured Transformer model attained a perplexity of ∼ 7 on the full test set.
Revised: Structured Transformer model attained a perplexity of ∼ 7 on the full test set.

Revised: When we compared different graph features of protein structure (Table 3), we indeed found that using local orientation information was important.

Original: Improvement over proﬁle-based methods We also compared to a recent method SPIN2 that predicts, using deep neural networks, protein sequence proﬁles given protein structures (O’Connell et al., 2018).
Revised: Improvement over proﬁle-based methods We also compared to a recent method SPIN2 that predicts, using deep neural networks, protein sequence proﬁles given protein structures (O’Connell et al., 2018).

Original: Since SPIN2 is computationally intensive (minutes per protein for small proteins) and was trained on complete proteins rather than chains, we evaluated it on two subsets of the full test set: a ‘Small’ subset of the test set containing chains up to length 100 and a ‘Single chain’ subset containing only those models where the single chain accounted for the entire protein record in the Protein Data Bank.
Revised: Since SPIN2 is computationally intensive (minutes per protein for small proteins) and was trained on complete proteins rather than chains, we evaluated it on two subsets of the full test set: a ‘Small’ subset of the test set containing chains up to length 100 and a ‘Single chain’ subset containing only those models where the single chain accounted for the entire protein record in the Protein Data Bank.

Original: Both subsets discarded any chains with structural gaps.
Revised: Both subsets discarded any chains with structural gaps.

Original: Additionally, since

Original: SPIN2 sometimes output zero-probability predictions truncated by round-off error, we added a small pseudocount λ to its predictions at varying strengths to avoid inﬁnite perplexities.

Original: We found that our
Revised: We found that our

Original: Structured Transformer model considerably improved upon the perplexities of SPIN2 (speciﬁcally, an improvement on the order of ∼ 2-2.5).
Revised: Structured Transformer model considerably improved upon the perplexities of SPIN2 (Table 2).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We presented a new deep generative model to ‘design’ protein sequences given a graph speciﬁcation of their structure.
Revised: We presented a new deep generative model to ‘design’ protein sequences given a graph speciﬁcation of their structure.

Original: Our model augments the traditional sequence-level self-attention of Transformers (Vaswani et al., 2017) with relational 3D structural encodings and is able to leverage the spatial locality of dependencies in molecular structures for efﬁcient computation.
Revised: Our model augments the traditional sequence-level self-attention of Transformers (Vaswani et al., 2017) with relational 3D structural encodings and is able to leverage the spatial locality of dependencies in molecular structures for efﬁcient computation.

Original: When evaluated on unseen folds, the model achieves signiﬁcantly improved perplexities over the state-of-the-art parametric generative models.
Revised: When evaluated on unseen folds, the model achieves signiﬁcantly improved perplexities over the state-of-the-art parametric generative models.

Original: Our framework suggests the possibility of being able to efﬁciently design and engineer protein sequences with structurally-guided deep generative models, and underscores the central role of modeling sparse long-range dependencies in biological sequences. 
Revised: Our framework suggests the possibility of being able to efﬁciently design and engineer protein sequences with structurally-guided deep generative models, and underscores the central role of modeling sparse long-range dependencies in biological sequences. A CKNOWLEDGMENTS

Revised: We thank members of the MIT MLPDS consortium for helpful feedback and discussions

Original:  A PPENDIX
Revised:  A PPENDIX