Original: S INGLE L AYERS OF A TTENTION S UFFICE TO P REDICT P ROTEIN C ONTACTS
Revised: S INGLE L AYERS OF A TTENTION S UFFICE TO P REDICT P ROTEIN C ONTACTS

Original: A BSTRACT
Revised: A [Table]

Original: The established approach to unsupervised protein contact prediction estimates coevolving positions using undirected graphical models.
Revised: The established approach to unsupervised protein contact prediction estimates coevolving positions using undirected graphical models.

Original: This approach trains a
Revised: This approach trains a

Original: Potts model on a Multiple Sequence Alignment.
Revised: Potts model on a Multiple Sequence Alignment.

Original: Increasingly large Transformers are being pretrained on unlabeled, unaligned protein sequence databases but have demonstrated mixed results for downstream tasks, including contact prediction.
Revised: Increasingly large Transformers are being pretrained on unlabeled, unaligned protein sequence databases but have demonstrated mixed results for downstream tasks, including contact prediction.

Original: We argue that attention is a principled model of protein interactions, grounded in real properties of protein family data.
Revised: We argue that attention is a principled model of protein interactions, grounded in real properties of protein family data.

Original: We introduce an energy-based attention layer, factored attention , and show that it achieves comparable performance to
Revised: We introduce an energy-based attention layer, factored attention , and show that it achieves comparable performance to

Original: Potts models while sharing parameters both within and across families.
Revised: Potts models while sharing parameters both within and across families.

Original: We contrast factored attention with the Transformer to indicate that the Transformer leverages hierarchical signal in protein family databases not captured by our singlelayer models.
Revised: We contrast factored attention with the Transformer to indicate that the Transformer leverages hierarchical signal in protein family databases not captured by our singlelayer models.

Original: This raises the exciting possibility for the development of powerful structured models of protein family databases.
Revised: This raises the exciting possibility for the development of powerful structured models of protein family databases.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Inferring protein structure from sequence is a longstanding problem in computational biochemistry.
Revised: Inferring protein structure from sequence is a longstanding problem in computational biochemistry.

Original: Potts models, a particular kind of Markov Random Field (MRF), are the predominant unsupervised method for modeling interactions between amino acids.
Revised: Potts models, a particular kind of Markov Random Field (MRF), are the predominant unsupervised method for modeling interactions between amino acids.

Original: Potts models are trained to maximize pseudolikelihood on alignments of evolutionarily related proteins (Balakrishnan et al., 2011; Ekeberg et al., 2013; Seemayer et al., 2014).
Revised: Potts models are trained to maximize pseudolikelihood on alignments of evolutionarily related proteins (Balakrishnan et al., 2011; Ekeberg et al., 2013; Seemayer et al., 2014).

Original: Features derived from Potts models were the main drivers of performance at the CASP11 competition (Monastyrskyy et al., 2016) and have become standard in state-of-the-art supervised models (Wang et al., 2017; Yang et al., 2019; Senior et al., 2020).
Revised: Features derived from Potts models were the main drivers of performance at the CASP11 competition (Monastyrskyy et al., 2016) and have become standard in state-of-the-art supervised models (Wang et al., 2017; Yang et al., 2019; Senior et al., 2020).

Original: Inspired by the success of BERT (Devlin et al., 2018), GPT (Brown et al., 2020) and related unsupervised models in NLP, a line of work has emerged that learns features of proteins through self-supervised pretraining (Rives et al., 2020; Elnaggar et al., 2020;
Revised: Inspired by the success of BERT (Devlin et al., 2018), GPT (Brown et al., 2020) and related unsupervised models in NLP, a line of work has emerged that learns features of proteins through self-supervised pretraining (Rives et al., 2020; Elnaggar et al., 2020;

Original: Rao et al., 2019; Madani et al., 2020; Nambiar et al., 2020).
Revised: Rao et al., 2019; Madani et al., 2020; Nambiar et al., 2020).

Original: This new approach trains Transformer (Vaswani et al., 2017) models on large datasets of protein sequences.
Revised: This new approach trains Transformer (Vaswani et al., 2017) models on large datasets of protein sequences.

Original: Pretrained model performance raises questions about the importance of data and model scale (Lu et al., 2020; Elnaggar et al., 2020), whether neural features compete with evolutionary features extracted by established bioinformatic methods (Rao et al., 2019), and the beneﬁts of transfer learning (Shanehsazzadeh et al., 2020).
Revised: Pretrained model performance raises questions about the importance of data and model scale (Lu et al., 2020; Elnaggar et al., 2020), whether neural features compete with evolutionary features extracted by established bioinformatic methods (Rao et al., 2019), and the beneﬁts of transfer learning (Shanehsazzadeh et al., 2020).

Original: In this paper, we take the position that attention-based models can build on the strengths of both Potts trained on alignments and Transformers pretrained on databases.
Revised: In this paper, we take the position that attention-based models can build on the strengths of both Potts trained on alignments and Transformers pretrained on databases.

Original: Our contributions are as follows:
Revised: Our contributions are as follows:

Original: We analyze the assumptions made by the attention mechanism in the Transformer and show that a single attention layer is a well-founded model of interactions within protein families without any analogies to natural language.
Revised: We analyze the assumptions made by the attention mechanism in the Transformer and show that a single attention layer is a well-founded model of interactions within protein families without any analogies to natural language.

Original: We introduce a simpliﬁed energy-based attention model, factored attention , and show that it can successfully share parameters across positions within a family or share amino acid features across hundreds of families.
Revised: We introduce a simpliﬁed energy-based attention model, factored attention , and show that it can successfully share parameters across positions within a family or share amino acid features across hundreds of families.

Original: We carefully benchmark ProtBERT-BFD (Elnaggar et al., 2020) against an optimized Potts implementation and show that the pretrained Transformer extracts contacts competitively with Potts.
Revised: We carefully benchmark ProtBERT-BFD (Elnaggar et al., 2020) against an optimized Potts implementation and show that the pretrained Transformer extracts contacts competitively with Potts.

Original: We identify a large gap between the gains afforded by the assumptions of single-layer models and the performance of ProtBERT-BFD which suggests the existence of properties linking protein families that allow for effective modeling of thousands of families at once.
Revised: We identify a large gap between the gains afforded by the assumptions of single-layer models and the performance of ProtBERT-BFD which suggests the existence of properties linking protein families that allow for effective modeling of thousands of families at once.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Proteins are polymers composed of amino acids and are commonly represented as strings.
Revised: Proteins are polymers composed of amino acids and are commonly represented as strings.

Original: Along with this 1D sequence representation, each protein folds into a 3D physical structure.
Revised: Along with this 1D sequence representation, each protein folds into a 3D physical structure.

Original: Physical distance between positions in 3D is often a much better indicator of functional interaction than proximity in sequence.
Revised: Physical distance between positions in 3D is often a much better indicator of functional interaction than proximity in sequence.

Original: One representation of physical distance is a contact map C , a symmetric matrix in which entry C ij = 1 if the β carbons of i and j are within 8 ˚A of one another, andotherwise.
Revised: One representation of physical distance is a contact map C , a symmetric matrix in which entry C ij = 1 if the β carbons of i and j are within 8 ˚A of one another, andotherwise.

Original: Multiple Sequence Alignments.
Revised: Multiple Sequence Alignments.

Original: To understand structure and function of a protein sequence, one typically assembles a set of its evolutionary relatives and looks for patterns within the set.
Revised: To understand structure and function of a protein sequence, one typically assembles a set of its evolutionary relatives and looks for patterns within the set.

Original: A set of related sequences is referred to as a protein family , commonly represented by a Multiple Sequence
Revised: A set of related sequences is referred to as a protein family , commonly represented by a Multiple Sequence

Original: Alignment (MSA).
Revised: Alignment (MSA).

Original: Gaps in aligned sequences correspond to insertions from an alignment algorithm
Revised: Gaps in aligned sequences correspond to insertions from an alignment algorithm

Original: (Johnson et al., 2010; Remmert et al., 2012), ensuring that positions with similar structure and function line up for all members of the family.
Revised: (Johnson et al., 2010; Remmert et al., 2012), ensuring that positions with similar structure and function line up for all members of the family.

Original: After aligning, sequence position carries signiﬁcant evolutionary, structural, and functional information.
Revised: After aligning, sequence position carries signiﬁcant evolutionary, structural, and functional information.

Original: See Appendix A.1 for more information.
Revised: See Appendix A.1 for more information.

Original: Coevolutionary Analysis of Protein Families.
Revised: Coevolutionary Analysis of Protein Families.

Original: The observation that statistical patterns in MSAs can be used to predict couplings has been widely used to infer structure and function from protein families (Korber et al., 1993; G¨obel et al., 1994; Lapedes et al., 1999; Lockless & Ranganathan, 1999; Fodor & Aldrich, 2004; Thomas et al., 2008; Weigt et al., 2009; Fatakia et al., 2009).
Revised: The observation that statistical patterns in MSAs can be used to predict couplings has been widely used to infer structure and function from protein families (Korber et al., 1993; G¨obel et al., 1994; Lapedes et al., 1999; Lockless & Ranganathan, 1999; Fodor & Aldrich, 2004; Thomas et al., 2008; Weigt et al., 2009; Fatakia et al., 2009).

Original: Stateof-the-art approaches estimate interactions from MRF parameters, as described in Section 3.5.
Revised: Stateof-the-art approaches estimate interactions from MRF parameters, as described in Section 3.5.

Original:  M ETHODS
Revised:  M ETHODS

Original: Throughout this section, x = ( x 1 , . . .
Revised: Throughout this section, x = ( x 1 , . . .

Original: , x L ) is a sequence of length L from an alphabet of size A .
Revised: , x L ) is a sequence of length L from an alphabet of size A .

Original: This sequence is part of an MSA of length L with N total sequences.
Revised: This sequence is part of an MSA of length L with N total sequences.

Original: Recall that a fully-connected
Revised: Recall that a fully-connected

Original: Pairwise MRF over p variables X 1 , . . .
Revised: Pairwise MRF over p variables X 1 , . . .

Original: , X p speciﬁes a distribution
Revised: , X p speciﬁes a distribution

Original: [Equation]
Revised: [Equation]

Original: where Z is the partition function and E θ ( x i , x j ) is an arbitrary function of i , j , x i and x j .
Revised: where Z is the partition function and E θ ( x i , x j ) is an arbitrary function of i , j , x i and x j .

Original: For all models below, we can introduce an explicit functional E θ ( x i ) to capture the marginal distribution of X i .
Revised: For all models below, we can introduce an explicit functional E θ ( x i ) to capture the marginal distribution of X i .

Original: When introduced, we parametrize the marginal with E θ ( x i )
Revised: When introduced, we parametrize the marginal with E θ ( x i )

Original: =
Revised: =

Original: b i,x i for b ∈ R L × A .
Revised: b i,x i for b ∈ R L × A .

Original: 3. P OTTS M ODELS
Revised: 3. P OTTS M ODELS

Original: A Potts model is a fully-connected pairwise MRF with L variables, each representing a position in the MSA.
Revised: A Potts model is a fully-connected pairwise MRF with L variables, each representing a position in the MSA.

Original: An edge ( i, j ) is parametrized with a matrix W ij ∈
Revised: An edge ( i, j ) is parametrized with a matrix W ij ∈

Original: R
Revised: R

Original: A × A .
Revised: A × A .

Original: These matrices are organized into an order-4 tensor which form the parameters of a Potts model, see Figure 8.
Revised: These matrices are organized into an order-4 tensor which form the parameters of a Potts model, see Figure 8.

Original: Note that W ij =
Revised: Note that W ij =

Original: W ji .
Revised: W ji .

Original: The energy functional of a Potts model is given through lookups, namely
Revised: The energy functional of a Potts model is given through lookups, namely

Original: [Equation]
Revised: [Equation]

Original: 3. F ACTORED A TTENTION
Revised: 3. F ACTORED A TTENTION

Original: Factored attention has two advantages over Potts for modeling protein families: it shares a pool of amino acid feature matrices across all positions and it estimates O ( L ) parameters instead of O
Revised: Factored attention has two advantages over Potts for modeling protein families: it shares a pool of amino acid feature matrices across all positions and it estimates O ( L ) parameters instead of O

Original: ( L 2 ) .
Revised: ( L 2 ) .

Original: Sharing amino acid features.
Revised: Sharing amino acid features.

Original: Many contacts in a protein are driven by similar interactions between amino acids, such as many types of weakly polar interactions (Burley & Petsko, 1988; Jaenicke, 2000).
Revised: Many contacts in a protein are driven by similar interactions between amino acids, such as many types of weakly polar interactions (Burley & Petsko, 1988; Jaenicke, 2000).

Original: If two pairs of positions ( i, j ) and ( l, m ) are both in contact due to the same interaction, a
Revised: If two pairs of positions ( i, j ) and ( l, m ) are both in contact due to the same interaction, a

Original: Potts model must estimate completely separate amino acid features W ij and W lm .
Revised: Potts model must estimate completely separate amino acid features W ij and W lm .

Original: In order to share amino acid features, we want to compute all energies from one pool of A × A feature matrices.
Revised: In order to share amino acid features, we want to compute all energies from one pool of A × A feature matrices.

Original: The simplest way to accomplish this is by associating an L × L matrix A to every A ×
Revised: The simplest way to accomplish this is by associating an L × L matrix A to every A ×

Original: A feature matrix
Revised: A feature matrix

Original: W V .
Revised: W V .

Original: For H such pairs ( A , W V ) , we could introduce a factorized MRF:
Revised: For H such pairs ( A , W V ) , we could introduce a factorized MRF:

Original: [Equation]
Revised: [Equation]

Original: A row-wise softmax is taken to encourage sparse interactions and aid in normalization.
Revised: A row-wise softmax is taken to encourage sparse interactions and aid in normalization.

Original: This model allows the pairs ( i, j ) and ( l, m ) to reuse a single feature W hV , assuming A hij and A hlm are both large.
Revised: This model allows the pairs ( i, j ) and ( l, m ) to reuse a single feature W hV , assuming A hij and A hlm are both large.

Original: Scaling linearly in length.
Revised: Scaling linearly in length.

Original: Both Potts and the factorized model in Equation 2 have O
Revised: Both Potts and the factorized model in Equation 2 have O

Original: ( L 2 ) parameters.
Revised: ( L 2 ) parameters.

Original: However, contacts are observed to grow linearly over the wide range of protein structures currently available (Taylor & Sadowski, 2011; Kamisetty et al., 2013), which we examine in Figure 11.
Revised: However, contacts are observed to grow linearly over the wide range of protein structures currently available (Taylor & Sadowski, 2011; Kamisetty et al., 2013), which we examine in Figure 11.

Original: Given that the number of interactions we wish to estimate grows linearly in length, the quadratic scaling of these models can be greatly improved.
Revised: Given that the number of interactions we wish to estimate grows linearly in length, the quadratic scaling of these models can be greatly improved.

Original: One way to ﬁx this is by introducing the factorization A = W Q W TK , where W Q , W K ∈ R L × d .
Revised: One way to ﬁx this is by introducing the factorization A = W Q W TK , where W Q , W K ∈ R L × d .

Original: As before, we employ a row-wise softmax for sparsity and normalization.
Revised: As before, we employ a row-wise softmax for sparsity and normalization.

Original: Combining feature sharing with linear length scaling leads to factored attention , deﬁned in Equation 3.
Revised: Combining feature sharing with linear length scaling leads to factored attention , deﬁned in Equation 3.

Original: Like Potts, factored attention is a fully-connected pairwise MRF with L variables.
Revised: Like Potts, factored attention is a fully-connected pairwise MRF with L variables.

Original: The parameters of this model consist of H triples ( W Q , W K , W V ) , where W Q , W K ∈
Revised: The parameters of this model consist of H triples ( W Q , W K , W V ) , where W Q , W K ∈

Original: R
Revised: R

Original: L × d ; W V ∈ R
Revised: L × d ; W V ∈ R

Original: A × A ; and d is a hyperparameter.
Revised: A × A ; and d is a hyperparameter.

Original: Each such triple is called a head and d is the head size .
Revised: Each such triple is called a head and d is the head size .

Original: Unlike a Potts model, the parameters for each edge ( i, j ) are tied through the use of heads.
Revised: Unlike a Potts model, the parameters for each edge ( i, j ) are tied through the use of heads.

Original: The energy functional is
Revised: The energy functional is

Original: [Equation]
Revised: [Equation]

Original: where symm( M ) =
Revised: where symm( M ) =

Original: ( M + M T ) / 2 ensures the positional interactions are symmetric.
Revised: ( M + M T ) / 2 ensures the positional interactions are symmetric.

Original: Adding sequence-dependent interactions leads to standard attention, see Appendix A.2.
Revised: Adding sequence-dependent interactions leads to standard attention, see Appendix A.2.

Original: 3. S INGLE - LAYER ATTENTION
Revised: 3. S INGLE - LAYER ATTENTION

Original: Our single-layer attention model consists of a single Transformer encoder layer: an attention layer followed by a dense layer, with residual connections (He et al., 2016) and layer normalization (Ba et al., 2016) to aid in optimization.
Revised: Our single-layer attention model consists of a single Transformer encoder layer: an attention layer followed by a dense layer, with residual connections (He et al., 2016) and layer normalization (Ba et al., 2016) to aid in optimization.

Original: Transformer implementations typically use a sine/cosine positional encoding (Vaswani et al., 2017) or learned Gaussian positional encoding (Wolf et al., 2020), rather than the one-hot positional encoding used in our single-layer models.
Revised: Transformer implementations typically use a sine/cosine positional encoding (Vaswani et al., 2017) or learned Gaussian positional encoding (Wolf et al., 2020), rather than the one-hot positional encoding used in our single-layer models.

Original: Self-Supervised Losses.
Revised: Self-Supervised Losses.

Original: Given an MSA, state-of-the-art methods estimate Potts model parameters through pseudolikelihood maximization (Kamisetty et al., 2013; Ekeberg et al., 2013).
Revised: Given an MSA, state-of-the-art methods estimate Potts model parameters through pseudolikelihood maximization (Kamisetty et al., 2013; Ekeberg et al., 2013).

Original: On the other hand, BERT-like attention-based models are typically trained with variants of masked language modeling (Devlin et al., 2018).
Revised: On the other hand, BERT-like attention-based models are typically trained with variants of masked language modeling (Devlin et al., 2018).

Original: Pseudolikelihood is challenging to compute efﬁciently for generic models, unlike the masked language modeling loss.
Revised: Pseudolikelihood is challenging to compute efﬁciently for generic models, unlike the masked language modeling loss.

Original: Both of these losses require computing conditionals of the form p θ ( x i | x \ M ) , where M is a subset of { 1 , . . .
Revised: Both of these losses require computing conditionals of the form p θ ( x i | x \ M ) , where M is a subset of { 1 , . . .

Original: , L } \ { i } .
Revised: , L } \ { i } .

Original: The losses L PL and
Revised: The losses L PL and

Original: L MLM for pseudolikelihood and masked language modeling, respectively, are
Revised: L MLM for pseudolikelihood and masked language modeling, respectively, are

Original: [Equation]
Revised: [Equation]

Original: Regularization for Potts and factored attention are both based on MRF edge parameters, while single-layer attention is penalized using weight decay.
Revised: Regularization for Potts and factored attention are both based on MRF edge parameters, while single-layer attention is penalized using weight decay.

Original: More details can be found in Appendix
Revised: More details can be found in Appendix

Original: A.4.
Revised: A.4.

Original: 3. P RETRAINING ON S EQUENCE D ATABASES
Revised: 3. P RETRAINING ON S EQUENCE D ATABASES

Original: All single-layer models are trained on a set of evolutionarily related sequences.
Revised: All single-layer models are trained on a set of evolutionarily related sequences.

Original: Given a large database of protein sequences such as UniRef100 (Suzek et al., 2007) or BFD (Steinegger & S¨oding, 2018; Steinegger et al., 2019), these models cannot be trained until signiﬁcant preprocessing has been done: clustering, dereplication of highly related sequences, and alignment to generate an MSA for each cluster.
Revised: Given a large database of protein sequences such as UniRef100 (Suzek et al., 2007) or BFD (Steinegger & S¨oding, 2018; Steinegger et al., 2019), these models cannot be trained until signiﬁcant preprocessing has been done: clustering, dereplication of highly related sequences, and alignment to generate an MSA for each cluster.

Original: In contrast, the self-supervised approach taken by works such as (Elnaggar et al., 2020; Rives et al., 2020; Rao et al., 2019; Madani et al., 2020) applies BERT-style pretraining directly on the database of proteins with minimal preprocessing.
Revised: In contrast, the self-supervised approach taken by works such as (Elnaggar et al., 2020; Rives et al., 2020; Rao et al., 2019; Madani et al., 2020) applies BERT-style pretraining directly on the database of proteins with minimal preprocessing.

Original: More details on the differences between database-level models and single-family models can be found in Appendix A.5.
Revised: More details on the differences between database-level models and single-family models can be found in Appendix A.5.

Original: 3. E XTRACTING C ONTACTS
Revised: 3. E XTRACTING C ONTACTS

Original: Potts.
Revised: Potts.

Original: We follow standard practice and extract a contact map (
Revised: We follow standard practice and extract a contact map (

Original: cid:98) C ∈ R L × L from the order-4 interaction tensor W by setting (cid:98) C ij = (cid:107) W ij (cid:107) F .
Revised: cid:98) C ∈ R L × L from the order-4 interaction tensor W by setting (cid:98) C ij = (cid:107) W ij (cid:107) F .

Original: Factored Attention.
Revised: Factored Attention.

Original: Since factored attention is a pairwise MRF, we can compute its order-4 interaction tensor W and use the same procedure as Potts.
Revised: Since factored attention is a pairwise MRF, we can compute its order-4 interaction tensor W and use the same procedure as Potts.

Original: See Equation 4.
Revised: See Equation 4.

Original: Single-Layer Attention.
Revised: Single-Layer Attention.

Original: To produce contacts for an MSA, we compute attention maps from only the positional encoding (without sequence) and average attention maps from all heads.
Revised: To produce contacts for an MSA, we compute attention maps from only the positional encoding (without sequence) and average attention maps from all heads.

Original: Each singlelayer attention model is trained on one MSA, so the positional encoding is a feature shared by all sequences in the MSA.
Revised: Each singlelayer attention model is trained on one MSA, so the positional encoding is a feature shared by all sequences in the MSA.

Original: ProtBERT-BFD.
Revised: ProtBERT-BFD.

Original: We extract contacts from ProtBERT by averaging a subset of attention maps for an input sequence x .
Revised: We extract contacts from ProtBERT by averaging a subset of attention maps for an input sequence x .

Original: Of the 16 heads in 30 layers, we selected six whose attention maps had the top individual contact precisions over 500 families randomly selected from the Yang et al.
Revised: Of the 16 heads in 30 layers, we selected six whose attention maps had the top individual contact precisions over 500 families randomly selected from the Yang et al.

Original: (2019) dataset.
Revised: (2019) dataset.

Original: Predicted contacts for x are given by averaging the L × L attention maps from these six heads, then symmetrizing additively.
Revised: Predicted contacts for x are given by averaging the L × L attention maps from these six heads, then symmetrizing additively.

Original: See Table 2.
Revised: See Table 2.

Original: Average Product Correction (APC).
Revised: Average Product Correction (APC).

Original: Empirically, Potts models trained with Frobenius norm regularization have artifacts in the outputs (cid:98) C .
Revised: Empirically, Potts models trained with Frobenius norm regularization have artifacts in the outputs (cid:98) C .

Original: These are removed with the Average Product Correction (APC) (Dunn et al., 2008).
Revised: These are removed with the Average Product Correction (APC) (Dunn et al., 2008).

Original: Unless otherwise stated, we apply APC to all extracted contacts.
Revised: Unless otherwise stated, we apply APC to all extracted contacts.

Original:  R ESULTS
Revised:  R ESULTS

Original: Experimental Setup.
Revised: Experimental Setup.

Original: We use a set of 748 protein families from Yang et al.
Revised: We use a set of 748 protein families from Yang et al.

Original: to evaluate all models.
Revised: to evaluate all models.

Original: For Potts models and single attention layers, we train separate models on each individual
Revised: For Potts models and single attention layers, we train separate models on each individual

Original: MSA.
Revised: MSA.

Original: ProtBERT-BFD is frozen for all experiments.
Revised: ProtBERT-BFD is frozen for all experiments.

Original: We train models using PyTorchLightning (Falcon, 2019) and Weights and Biases (Biewald, 2020).
Revised: We train models using PyTorchLightning (Falcon, 2019) and Weights and Biases (Biewald, 2020).

Original: We extract contacts from each model following the procedure outlined in Appendix 3.5.
Revised: We extract contacts from each model following the procedure outlined in Appendix 3.5.

Original: We compare predicted contact maps (cid:98) C to true contact maps
Revised: We compare predicted contact maps (cid:98) C to true contact maps

Original: C using standard metrics based on precision.
Revised: C using standard metrics based on precision.

Original: A particularly important metric is precision at L , where L is the length of the sequence (Schaarschmidt et al., 2018;
Revised: A particularly important metric is precision at L , where L is the length of the sequence (Schaarschmidt et al., 2018;

Original: Shrestha et al., 2019).
Revised: Shrestha et al., 2019).

Original: This is computed by masking (cid:98) C to only consider positions ≥ 6 apart, predicting the top L entries to be contacts, and computing precision.
Revised: This is computed by masking (cid:98) C to only consider positions ≥ 6 apart, predicting the top L entries to be contacts, and computing precision.

Original: We provide more information on data and metrics in Appendix
Revised: We provide more information on data and metrics in Appendix

Original: A.7 and more information on model hyperparameters in Appendix A.8.
Revised: A.7 and more information on model hyperparameters in Appendix A.8.

Original: Attention assumptions reﬂected in 15,051 protein structures.
Revised: Attention assumptions reﬂected in 15,051 protein structures.

Original: We examine all 15,051 structures in the Yang et al.
Revised: We examine all 15,051 structures in the Yang et al.

Original: (2019) dataset for evidence of two key properties useful for single-layer attention models: few contacts per residue and the number of contacts scaling linearly in length.
Revised: (2019) dataset for evidence of two key properties useful for single-layer attention models: few contacts per residue and the number of contacts scaling linearly in length.

Original: Figureshows a linear trend of number of contacts versus length.
Revised: Figureshows a linear trend of number of contacts versus length.

Original: In Figure 12, we see that 80% of the 3,747,101 million residues in these structures have 4 or fewer contacts.
Revised: In Figure 12, we see that 80% of the 3,747,101 million residues in these structures have 4 or fewer contacts.

Original: Only 1.8% of residues have more than ten contacts.
Revised: Only 1.8% of residues have more than ten contacts.

Original: This shows that the row-wise softmax, which encourages each residue to attend to only a few other residues per-head, reﬂects structure found in the data.
Revised: This shows that the row-wise softmax, which encourages each residue to attend to only a few other residues per-head, reﬂects structure found in the data.

Original: Factored attention matches Potts performance on 748 families.
Revised: Factored attention matches Potts performance on 748 families.

Original: Figure 1 shows a representative sample of good quality contact maps extracted from all models.
Revised: Figure 1 shows a representative sample of good quality contact maps extracted from all models.

Original: Figure 2a summarizes the performance of all models over the set of 748 protein families.
Revised: Figure 2a summarizes the performance of all models over the set of 748 protein families.

Original: Factored attention, Potts, and ProtBERTBFD have comparable overall performance, with median precision at L of 0.46, 0.47, and 0.48, respectively.
Revised: Factored attention, Potts, and ProtBERTBFD have comparable overall performance, with median precision at L of 0.46, 0.47, and 0.48, respectively.

Original: Stratifying by number of sequences reveals that ProtBERT-BFD has higher precision on
Revised: Stratifying by number of sequences reveals that ProtBERT-BFD has higher precision on

Original: [Figure]
Revised: [Figure]

Original: P@L 0.
Revised: P@L 0.

Original: Factored Attn
Revised: Factored Attn

Original: P@L 0.
Revised: P@L 0.

Original: Single-Layer Attn
Revised: Single-Layer Attn

Original: P@L 0.
Revised: P@L 0.

Original: ProtBERT-BFD
Revised: ProtBERT-BFD

Original: P@L 0.
Revised: P@L 0.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: -range contacts, while
Revised: contacts, while

Original: ProtBERT-BFD has slightly better precision on short-range contacts.
Revised: ProtBERT-BFD has slightly better precision on short-range contacts.

Original: MSAs with fewer than 256 sequences.
Revised: MSAs with fewer than 256 sequences.

Original: For MSAs with greater than 1024 sequences, Potts, factored attention, and ProtBERT-BFD have comparable performance.
Revised: For MSAs with greater than 1024 sequences, Potts, factored attention, and ProtBERT-BFD have comparable performance.

Original: Single-layer attention is uniformly worse over all MSA depths.
Revised: Single-layer attention is uniformly worse over all MSA depths.

Original: Next, we evaluate the impact of sequence length on performance.
Revised: Next, we evaluate the impact of sequence length on performance.

Original: Figure 2b shows that factored attention and Potts achieve similar precision at L over the whole range of family lengths, despite factored attention having far fewer parameters for long families.
Revised: Figure 2b shows that factored attention and Potts achieve similar precision at L over the whole range of family lengths, despite factored attention having far fewer parameters for long families.

Original: This shows that factored attention can successfully leverage sparsity assumptions where they are most useful.
Revised: This shows that factored attention can successfully leverage sparsity assumptions where they are most useful.

Original: Long-range contacts are particularly important for downstream structure-prediction algorithms – long-range precision at L/ 5 is reported in both CASP12 and CASP

Revised: Long-range contacts are particularly important for downstream structure-prediction algorithms – long-range precision at L/ 5 is reported in both CASP12 and CASP

Original: (Schaarschmidt et al., 2018;
Revised: (Schaarschmidt et al., 2018;

Original: Shrestha et al., 2019).
Revised: Shrestha et al., 2019).

Original: Figure 3 breaks down contact precisions based on position separation into
Revised: Figure 3 breaks down contact precisions based on position separation into

Original: [Figure]
Revised: [Figure]

Revised: tt

Original: [Figure]
Revised: [Figure]

Revised: [Figure]

Original: short ( 6 ≤ sep < 12 ), medium ( 12 ≤ sep < 24 ), and long ( 24 ≤ sep ).
Revised: short ( 6 ≤ sep < 12 ), medium ( 12 ≤ sep < 24 ), and long ( 24 ≤ sep ).

Original: We see that ProtBERT-BFD performs best on short-range contacts, with a median increase of 0.068 precision at L/ 5 .
Revised: We see that ProtBERT-BFD performs best on short-range contacts, with a median increase of 0.068 precision at L/ 5 .

Original: On longrange contacts, there is no appreciable difference in performance to Potts and factored attention.
Revised: On longrange contacts, there is no appreciable difference in performance to Potts and factored attention.

Original: Across the range of contact bins, factored attention and Potts perform very similarly.
Revised: Across the range of contact bins, factored attention and Potts perform very similarly.

Original: Fewer heads can match Potts on L/ 5 contacts.
Revised: Fewer heads can match Potts on L/ 5 contacts.

Original: We probe the limits of parameter sharing by lowering the number of heads in factored attention and evaluating whether fewer heads can be used to precisely estimate contacts.
Revised: We probe the limits of parameter sharing by lowering the number of heads in factored attention and evaluating whether fewer heads can be used to precisely estimate contacts.

Original: Figure 4a shows that 128 heads can be used to estimate L/ 5 contacts as precisely as Potts over the full set of 748 families.
Revised: Figure 4a shows that 128 heads can be used to estimate L/ 5 contacts as precisely as Potts over the full set of 748 families.

Original: In Figure 4b, we see that factored attention with 32 and 64 heads is still able to achieve reasonable overall performance compared to Potts.and 64 heads have precision at L/ 5 at least as high as Potts for 329 and 348 families, respectively.
Revised: In Figure 4b, we see that factored attention with 32 and 64 heads is still able to achieve reasonable overall performance compared to Potts.and 64 heads have precision at L/ 5 at least as high as Potts for 329 and 348 families, respectively.

Original: If we wish to recover the top L contacts, 256 heads are required to match Potts across all families, as seen in Figure 13.
Revised: If we wish to recover the top L contacts, 256 heads are required to match Potts across all families, as seen in Figure 13.

Original: Having more heads than 256 does not further increase performance.
Revised: Having more heads than 256 does not further increase performance.

Original: Intriguingly,
Revised: Intriguingly,

Original: Figure 14 demonstrates that both Spearman and Pearson correlation between the order-4 interaction tensors of factored attention and Potts improve even when increasing to 512 heads.
Revised: Figure 14 demonstrates that both Spearman and Pearson correlation between the order-4 interaction tensors of factored attention and Potts improve even when increasing to 512 heads.

Original: We do not observe the same trends for increasing head size, as shown in Figure 15.
Revised: We do not observe the same trends for increasing head size, as shown in Figure 15.

Original: For some families, the number of heads can be reduced even further.
Revised: For some families, the number of heads can be reduced even further.

Original: We show an example on the
Revised: We show an example on the

Original: MSA built for PDB entry 3n2a .
Revised: MSA built for PDB entry 3n2a .

Original: In Figure 5a, we see that merely 4 heads are required to recover L/contacts nearly identical to those recovered by Potts.
Revised: In Figure 5a, we see that merely 4 heads are required to recover L/contacts nearly identical to those recovered by Potts.

Original: This shows that shared amino acid features and interaction parameters can enable identical performance with a 300x reduction in parameters.
Revised: This shows that shared amino acid features and interaction parameters can enable identical performance with a 300x reduction in parameters.

Original: The training dynamics of these models are shown in Figure 5b.
Revised: The training dynamics of these models are shown in Figure 5b.

Original: Both factored attention with 256 heads and Potts converge after roughly 100 gradient steps, whereas factored attention with 4 heads requires nearly 10,000 steps to converge.
Revised: Both factored attention with 256 heads and Potts converge after roughly 100 gradient steps, whereas factored attention with 4 heads requires nearly 10,000 steps to converge.

Original: In Figure 16, we show that the top L contacts are signiﬁcantly worse for 4 heads compared to Potts.
Revised: In Figure 16, we show that the top L contacts are signiﬁcantly worse for 4 heads compared to Potts.

Original: One set of amino acid features can be used for all families.
Revised: One set of amino acid features can be used for all families.

Original: Thus far we have only examined models that share parameters within single protein families.
Revised: Thus far we have only examined models that share parameters within single protein families.

Original: Since ProtBERT is trained on an entire database, it can leverage feature sharing across families to attain greater parameter efﬁciency and improved performance on small MSAs.
Revised: Since ProtBERT is trained on an entire database, it can leverage feature sharing across families to attain greater parameter efﬁciency and improved performance on small MSAs.

Original: [Figure]
Revised: [Figure]

Original: To explore the possibility that attention can share parameters across families, we train factored attention using a single set of frozen value matrices.
Revised: To explore the possibility that attention can share parameters across families, we train factored attention using a single set of frozen value matrices.

Original: We ﬁrst train factored attention normally on 3n2a with 256 heads, then freeze the learned value matrices for the remaining 747 families.
Revised: We ﬁrst train factored attention normally on 3n2a with 256 heads, then freeze the learned value matrices for the remaining 747 families.

Original: The query and key parameters are trained normally.
Revised: The query and key parameters are trained normally.

Original: In Figure 6, we compare the precision at L of factored attention with frozen 3n2a features to that of factored attention trained normally.
Revised: In Figure 6, we compare the precision at L of factored attention with frozen 3n2a features to that of factored attention trained normally.

Original: Using a single frozen set of features results in only 6 families seeing precision at L decrease by more than 0.05, with a maximum drop of 0.11.
Revised: Using a single frozen set of features results in only 6 families seeing precision at L decrease by more than 0.05, with a maximum drop of 0.11.

Original: Frozen values are also comparable to Potts performance, as expected – see Figure 17.
Revised: Frozen values are also comparable to Potts performance, as expected – see Figure 17.

Original: This suggests that, even for a single-layer model, a single set of value matrices can capture amino acid features across functionally and structurally distinct protein families.
Revised: This suggests that, even for a single-layer model, a single set of value matrices can capture amino acid features across functionally and structurally distinct protein families.

Original: Factored attention reduces total parameters estimated.
Revised: Factored attention reduces total parameters estimated.

Original: For an MSA of length L with alphabet size A , Potts models require (cid:0) L 2 (cid:1)
Revised: For an MSA of length L with alphabet size A , Potts models require (cid:0) L 2 (cid:1)

Original: A 2 parameters.
Revised: A 2 parameters.

Original: Factored attention with H heads and head size d requires H (2 Ld + A 2 ) parameters.
Revised: Factored attention with H heads and head size d requires H (2 Ld + A 2 ) parameters.

Original: In Figure 18, we plot number of parameters versus length for various values of H and d = 32 .
Revised: In Figure 18, we plot number of parameters versus length for various values of H and d = 32 .

Original: Potts requires a total of 12 billion parameters to model allfamilies.
Revised: Potts requires a total of 12 billion parameters to model allfamilies.

Original: Factored attention with 256 heads and head size 32 has 3.2 billion parameters; lowering to 128 heads reduces this to 790 million.
Revised: Factored attention with 256 heads and head size 32 has 3.2 billion parameters; lowering to 128 heads reduces this to 790 million.

Original: Half of this reduction comes from 107 families of length greater than 400.
Revised: Half of this reduction comes from 107 families of length greater than 400.

Original: ProtBERT-BFD is the most efﬁcient, with 420 million parameters.
Revised: ProtBERT-BFD is the most efﬁcient, with 420 million parameters.

Original: Ablations APC has a considerable impact on both Potts and factored attention, creating a median increase in precision at L of 0.1 and 0.07, respectively.
Revised: Ablations APC has a considerable impact on both Potts and factored attention, creating a median increase in precision at L of 0.1 and 0.07, respectively.

Original: The effect of APC is negligible for singlelayer attention and ProtBERT.
Revised: The effect of APC is negligible for singlelayer attention and ProtBERT.

Original: Replacing pseudolikelihood maximization with Masked Language
Revised: Replacing pseudolikelihood maximization with Masked Language

Original: Modeling did not appreciably change performance for either Potts or factored attention.
Revised: Modeling did not appreciably change performance for either Potts or factored attention.

Original: Addition of the single-site potential b i increases performance slightly for attention layers, but not enough to change overall trends.
Revised: Addition of the single-site potential b i increases performance slightly for attention layers, but not enough to change overall trends.

Original: To compare to ProtBERT-BFD, we train our single-layer attention models on unaligned families and found that performance degrades signiﬁcantly.
Revised: To compare to ProtBERT-BFD, we train our single-layer attention models on unaligned families and found that performance degrades signiﬁcantly.

Original: See Figures 19-22.
Revised: See Figures 19-22.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We have shown that single-layer factored attention models and the ProtBert-BFD Transformer achieve performance comparable to state-of-the-art Potts models on unsupervised contact extraction.
Revised: We have shown that single-layer factored attention models and the ProtBert-BFD Transformer achieve performance comparable to state-of-the-art Potts models on unsupervised contact extraction.

Original: We have also shown that the assumptions encoded by attention reﬂect important properties of protein families.
Revised: We have also shown that the assumptions encoded by attention reﬂect important properties of protein families.

Original: These results suggest that attention has a natural role in protein representation learning, without analogy to attention’s success in the domain of NLP.
Revised: These results suggest that attention has a natural role in protein representation learning, without analogy to attention’s success in the domain of NLP.

Original: Our results show that hierarchical structure, both within and across families, is a source of signal available to attention models.
Revised: Our results show that hierarchical structure, both within and across families, is a source of signal available to attention models.

Original: Understanding how such structure can be learned without the use of protein family labels could lead to the development of widely applicable modeling components for protein representation learning.
Revised: Understanding how such structure can be learned without the use of protein family labels could lead to the development of widely applicable modeling components for protein representation learning.

Original: We believe that unsupervised attention-based models have the potential to impact protein structure prediction as broadly as existing coevolutionary methods.
Revised: We believe that unsupervised attention-based models have the potential to impact protein structure prediction as broadly as existing coevolutionary methods.

Original: A. R ECOVERING F ACTORED A TTENTION FROM S TANDARD A TTENTION
Revised: A. R ECOVERING F ACTORED A TTENTION FROM S TANDARD A TTENTION

Original: Potts and Factored Attention estimate a single undirected graphical model from the training data.
Revised: Potts and Factored Attention estimate a single undirected graphical model from the training data.

Original: While a single graph can be a good approximation for the structure associated with a protein family, many families have subfamilies with different functional specializations and even different underlying contacts (Brown et al., 2007; Malinverni & Barducci, 2019).
Revised: While a single graph can be a good approximation for the structure associated with a protein family, many families have subfamilies with different functional specializations and even different underlying contacts (Brown et al., 2007; Malinverni & Barducci, 2019).

Original: Since subfamily identity is rarely known, allowing edge weights to be a function of sequence could enable the estimation of a family of highly related graphs.
Revised: Since subfamily identity is rarely known, allowing edge weights to be a function of sequence could enable the estimation of a family of highly related graphs.

Original: In the language of the Transformer, factored attention estimates a single graph because it computes queries and keys using only the positional encoding.
Revised: In the language of the Transformer, factored attention estimates a single graph because it computes queries and keys using only the positional encoding.

Original: We show more precisely that factored attention can be recovered from standard attention one of two ways
Revised: We show more precisely that factored attention can be recovered from standard attention one of two ways

Original: Computing queries and keys from one-hot positional encodings and values from one-hot sequence embeddings.
Revised: Computing queries and keys from one-hot positional encodings and values from one-hot sequence embeddings.

Original: Using a learned positional embedding and sequence embedding with identity query, key, and value matrices.
Revised: Using a learned positional embedding and sequence embedding with identity query, key, and value matrices.

Original: The symmetrization operator symm is not applied to Transformer multihead attention, so we present a slight variation of factored attention without symm in this section.
Revised: The symmetrization operator symm is not applied to Transformer multihead attention, so we present a slight variation of factored attention without symm in this section.

Original: Single attention layer.
Revised: Single attention layer.

Original: Given a sequence of dense vectors x = ( x 1 , . . .
Revised: Given a sequence of dense vectors x = ( x 1 , . . .

Original: , x L ) with x i ∈ R p , the attention mechanism of the Transformer encoder (multihead scaled dot product self-attention) produces a continuous representation
Revised: , x L ) with x i ∈ R p , the attention mechanism of the Transformer encoder (multihead scaled dot product self-attention) produces a continuous representation

Original: y ∈ R L × p .
Revised: y ∈ R L × p .

Original: If head size is d , this representation is computed using
Revised: If head size is d , this representation is computed using

Original: H heads ( W Q , W K , W V ) , where W Q , W K , W V ∈ R p × d .
Revised: H heads ( W Q , W K , W V ) , where W Q , W K , W V ∈ R p × d .

Original: Queries, keys, and values are deﬁned as
Revised: Queries, keys, and values are deﬁned as

Original: Q = xW Q , K = xW
Revised: Q = xW Q , K = xW

Original: K , V = xW V .
Revised: K , V = xW V .

Original: For a single head ( W Q , W K , W V ) , the output is given by
Revised: For a single head ( W Q , W K , W V ) , the output is given by

Original: [Equation]
Revised: [Equation]

Original: The full output in R dH is produced by concatenating all head outputs.
Revised: The full output in R dH is produced by concatenating all head outputs.

Original: A single Transformer encoder layer passes the output through a dense layer, applying layer-norms and residual connection to aid optimization.
Revised: A single Transformer encoder layer passes the output through a dense layer, applying layer-norms and residual connection to aid optimization.

Original: For the ﬁrst layer, the input x is a sequence of discrete tokens.
Revised: For the ﬁrst layer, the input x is a sequence of discrete tokens.

Original: To produce a dense vector combining sequence and position information, positional encodings and sequence embeddings are combined.
Revised: To produce a dense vector combining sequence and position information, positional encodings and sequence embeddings are combined.

Original: The positional encoding E pos ∈ R L × e produces a dense vector of dimension e for each position i .
Revised: The positional encoding E pos ∈ R L × e produces a dense vector of dimension e for each position i .

Original: The sequence embedding E seq ∈ R
Revised: The sequence embedding E seq ∈ R

Original: A × e maps each element of the vocabulary to a dense vector of dimension e .
Revised: A × e maps each element of the vocabulary to a dense vector of dimension e .

Original: Typically these are combined through summation to produce a dense vector ˜ x i =
Revised: Typically these are combined through summation to produce a dense vector ˜ x i =

Original: E seq ( x i ) + E pos ( i ) , which is input to the Transformer as described above.
Revised: E seq ( x i ) + E pos ( i ) , which is input to the Transformer as described above.

Original: For this paper, we use only multihead self-attention without the dense layer, layernorm, or residual connections, as these drastically hurt performance when employed for one layer.
Revised: For this paper, we use only multihead self-attention without the dense layer, layernorm, or residual connections, as these drastically hurt performance when employed for one layer.

Original: Factored attention from standard attention.
Revised: Factored attention from standard attention.

Original: Written explicitly, the input Transformer layer computes queries for a single head with Q =
Revised: Written explicitly, the input Transformer layer computes queries for a single head with Q =

Original: ( E pos + E seq ( x ))
Revised: ( E pos + E seq ( x ))

Original: W Q .
Revised: W Q .

Original: Keys and values are computed similarly.
Revised: Keys and values are computed similarly.

Original: To recover factored attention, we instead compute queries and keys via Q = E pos W Q and K = E pos W K , while values are given by V = E seq ( x ) W V .
Revised: To recover factored attention, we instead compute queries and keys via Q = E pos W Q and K = E pos W K , while values are given by V = E seq ( x ) W V .

Original: For simplicity, we one-hot encode both position and sequence, which corresponds using identity matrices E pos =
Revised: For simplicity, we one-hot encode both position and sequence, which corresponds using identity matrices E pos =

Original: I L ∈ R
Revised: I L ∈ R

Original: L × L and
Revised: L × L and

Original: E seq =
Revised: E seq =

Original: I A ∈ R A × A .
Revised: I A ∈ R A × A .

Original: Equivalently, one can view the positional encoding and sequence embedding as being learned while ﬁxing W Q , W K , and W V to be identity matrices.
Revised: Equivalently, one can view the positional encoding and sequence embedding as being learned while ﬁxing W Q , W K , and W V to be identity matrices.

Original: Implicit single-site term in single-layer attention.
Revised: Implicit single-site term in single-layer attention.

Original: For a single layer of attention, the product
Revised: For a single layer of attention, the product

Original: E pos W V is a matrix in R L × A .
Revised: E pos W V is a matrix in R L × A .

Original: This matrix does not depend on sequence inputs, thus allowing it to act as a single-site term.
Revised: This matrix does not depend on sequence inputs, thus allowing it to act as a single-site term.

Original: This suggests why inclusion of an explicit single-site term in Figure 22 had no effect for single-layer attention.
Revised: This suggests why inclusion of an explicit single-site term in Figure 22 had no effect for single-layer attention.

Original: A. C OMPUTATIONAL E FFICIENCY OF S INGLE L AYER M ODELS
Revised: A. C OMPUTATIONAL E FFICIENCY OF S INGLE L AYER M ODELS

Original: Model
Revised: Model

Original: [Table] [Table] [Table]
Revised: Length

Original: [Table]
Revised: Batches/s

Original:  Factored attention ( H = 256 ) [Table] [Table]
Revised: Potts 36. Factored attention ( H = 256 ) 32.

Original: Single-layer attention ( H = 128 ) 17.
Revised: Single-layer attention ( H = 128 ) 17.

Original: [Table]
Revised: [Table]

Original: In Table 1, we show the number of gradient steps per second for Potts, factored attention, and singlelayer attention.
Revised: In Table 1, we show the number of gradient steps per second for Potts, factored attention, and singlelayer attention.

Original: We ﬁx a batch size of 32 and report numbers for families of lengths 120, 440, and 904 in order to explore the impact of length on computational efﬁciency for all models.
Revised: We ﬁx a batch size of 32 and report numbers for families of lengths 120, 440, and 904 in order to explore the impact of length on computational efﬁciency for all models.

Original: Metrics are collected on a node with a single NVIDIA RTX 2080Ti, an Intel i7-7800x with 6 physical cores and 12 virtual cores, and 32 GB of system RAM.
Revised: Metrics are collected on a node with a single NVIDIA RTX 2080Ti, an Intel i7-7800x with 6 physical cores and 12 virtual cores, and 32 GB of system RAM.

Original: To get steady state batches per second, werun each model for and reportthe mean batchesper second between steps and 4500. and factored attentionhave similar throughput, while standard attentionscales much worse inlength due to its memory requirements necessitating the use of gradient accumulation.
Revised: To get steady state batches per second, we and report batches and factored attention similar throughput, while standard attention much worse in due to its memory requirements necessitating the use of gradient accumulation.

Original: This poor memory performance is because standard attention must compute L × L attention maps foreach batch element separately, whereas the L × L component of other models does not depend on sequence.
Revised: This poor memory performance is because standard attention must compute L × L attention maps for element separately, whereas the L × L component of other models does not depend on sequence.

Original: A. L OSSES
Revised: A. L OSSES

Original: The loss for all three models is of the form
Revised: The loss for all three models is of the form

Original: [Equation]
Revised: [Equation]

Original: where L is either pseudolikelihood or masked language modeling and R is a regularizer.
Revised: where L is either pseudolikelihood or masked language modeling and R is a regularizer.

Original: Potts regularization.
Revised: Potts regularization.

Original: Consider the order-4 interaction tensor W , where
Revised: Consider the order-4 interaction tensor W , where

Original: W ij ∈ R
Revised: W ij ∈ R

Original: A × A gives the parameters associated to edge ( i, j ) .
Revised: A × A gives the parameters associated to edge ( i, j ) .

Original: We regularize W by setting R ( θ ) =
Revised: We regularize W by setting R ( θ ) =

Original: (cid:80)
Revised: (cid:80)

Original: i<j (cid:107) W ij (cid:107) 2 F .
Revised: i<j (cid:107) W ij (cid:107) 2 F .

Original: This term is multiplied by λ · L · A , following Ovchinnikov et al.
Revised: This term is multiplied by λ · L · A , following Ovchinnikov et al.

Original: Factored attention regularization.
Revised: Factored attention regularization.

Original: Since factored attention is also a fully connected pairwise MRF, we use identical regularization to that of Potts.
Revised: Since factored attention is also a fully connected pairwise MRF, we use identical regularization to that of Potts.

Original: The order-4 tensor W is given by
Revised: The order-4 tensor W is given by

Original: [Equation]
Revised: [Equation]

Original: Single-layer attention regularization.
Revised: Single-layer attention regularization.

Original: Due the lack of an MRF interpretation for single-layer attention, we chose to use weight decay as is typically done for attention models.
Revised: Due the lack of an MRF interpretation for single-layer attention, we chose to use weight decay as is typically done for attention models.

Original: This corresponds to setting R ( θ ) to be the sum of Frobenius norm squared for all weights W Q , W K and W V .
Revised: This corresponds to setting R ( θ ) to be the sum of Frobenius norm squared for all weights W Q , W K and W V .

Original: Single-site term.
Revised: Single-site term.

Original: When any model has a single-site term, we follow standard practice and regularize its Frobenius norm squared.
Revised: When any model has a single-site term, we follow standard practice and regularize its Frobenius norm squared.

Original: A. D IFFERENCES BETWEEN D ATABASE - LEVEL AND S INGLE - FAMILY MODELS
Revised: A. D IFFERENCES BETWEEN D ATABASE - LEVEL AND S INGLE - FAMILY MODELS

Original: Difference in positional encoding.
Revised: Difference in positional encoding.

Original: For attention-based models, the contrast between single-family and database-level pretraining shows up entirely in the choice of positional encoding.
Revised: For attention-based models, the contrast between single-family and database-level pretraining shows up entirely in the choice of positional encoding.

Original: A singlefamily model, such as factored attention, uses a one-hot encoding of position in the MSA for that family.
Revised: A singlefamily model, such as factored attention, uses a one-hot encoding of position in the MSA for that family.

Original: As such, the positional encoding of factored attention relies on all preprocessing steps mentioned above.
Revised: As such, the positional encoding of factored attention relies on all preprocessing steps mentioned above.

Original: On the other hand, the positional encoding used by pretrained Transformers is computed on the raw input sequence, simplifying pretraining due to lack of clustering or alignment.
Revised: On the other hand, the positional encoding used by pretrained Transformers is computed on the raw input sequence, simplifying pretraining due to lack of clustering or alignment.

Original: While it is easier to implement positional encoding on raw sequence position, an ablation study in
Revised: While it is easier to implement positional encoding on raw sequence position, an ablation study in

Original: Section 4 shows that single-layer models fail on unaligned sequences.
Revised: Section 4 shows that single-layer models fail on unaligned sequences.

Original: This suggests that there is a tradeoff: clustering and alignment enables small, easily trained models to learn effectively, while large models can learn directly from the database if trained carefully.
Revised: This suggests that there is a tradeoff: clustering and alignment enables small, easily trained models to learn effectively, while large models can learn directly from the database if trained carefully.

Original: Simpler inference with pretrained Transformers.
Revised: Simpler inference with pretrained Transformers.

Original: Given a new sequence of interest and a database of sequences, single-family models require more steps for inference than pretrained Transformers.
Revised: Given a new sequence of interest and a database of sequences, single-family models require more steps for inference than pretrained Transformers.

Original: To apply a single-family model, one must query the database for related sequences, dereplicate the set, align sequences into an MSA, then train a model to learn contacts.
Revised: To apply a single-family model, one must query the database for related sequences, dereplicate the set, align sequences into an MSA, then train a model to learn contacts.

Original: On the other hand, a Transformer pretrained on the database simply computes a forward pass for the sequence of interest and its attention activations are used to predict contacts.
Revised: On the other hand, a Transformer pretrained on the database simply computes a forward pass for the sequence of interest and its attention activations are used to predict contacts.

Original: No explicit querying or aligning is performed.
Revised: No explicit querying or aligning is performed.

Original: A. P ROT BERT-BFD HEAD SELECTION
Revised: A. P ROT BERT-BFD HEAD SELECTION

Original: [Table]
Revised: [Table]

Original: A. D ATA AND M ETRICS A.7. S ELECTION OF P ROTEIN F AMILIES
Revised: A. D ATA AND M ETRICS A.7. S ELECTION OF P ROTEIN F AMILIES

Original: We used the following sets of families for model development:
Revised: We used the following sets of families for model development:

Original: A set of 748 families was chosen for performance evaluation.
Revised: A set of 748 families was chosen for performance evaluation.

Original: All metrics reported in the paper are on this set, with a single choice of hyperparameters for Potts models, factored attention, and standard attention.
Revised: All metrics reported in the paper are on this set, with a single choice of hyperparameters for Potts models, factored attention, and standard attention.

Original: The 748 families were chosen randomly from the Yang et al.
Revised: The 748 families were chosen randomly from the Yang et al.

Original: (2019) dataset, which consists of 15,051 MSAs generated from the databases UniClust30 and UniRef100 (Suzek et al., 2007), as well as metagenomic datasets.
Revised: (2019) dataset, which consists of 15,051 MSAs generated from the databases UniClust30 and UniRef100 (Suzek et al., 2007), as well as metagenomic datasets.

Original: Our random sample is representative of the range of MSA depths and protein lengths, see Figure 10.
Revised: Our random sample is representative of the range of MSA depths and protein lengths, see Figure 10.

Original: A set of six families from the 748 was chosen to choose hyperparameters for single-layer attention.
Revised: A set of six families from the 748 was chosen to choose hyperparameters for single-layer attention.

Original: They were chosen to span a range of MSA depth (large and small), as well as three different regimes of Potts performance (Good, Ok, Poor).
Revised: They were chosen to span a range of MSA depth (large and small), as well as three different regimes of Potts performance (Good, Ok, Poor).

Original: These families were used to tune hyperparameters as described in Section A.8.1.
Revised: These families were used to tune hyperparameters as described in Section A.8.1.

Original: See Table 3. 3.
Revised: See Table 3. 3.

Original: A set of ten families from the 748 was chosen where factored attention performed very poorly in our initial experiments.
Revised: A set of ten families from the 748 was chosen where factored attention performed very poorly in our initial experiments.

Original: Half were chosen to be long proteins and the other half to be short.
Revised: Half were chosen to be long proteins and the other half to be short.

Original: This set was used to optimize learning rate and regularization for factored attention to ensure reasonable model performance.
Revised: This set was used to optimize learning rate and regularization for factored attention to ensure reasonable model performance.

Original: See Table 4. 4.
Revised: See Table 4. 4.

Original: 500 entirely disjoint families were further selected randomly from Yang et al.
Revised: 500 entirely disjoint families were further selected randomly from Yang et al.

Original: and used to compute average precision for each head in ProtBERT-BFD (Elnaggar et al., 2020).
Revised: and used to compute average precision for each head in ProtBERT-BFD (Elnaggar et al., 2020).

Original: Performance on these families was used for selecting the top 6 heads, see Table 2.
Revised: Performance on these families was used for selecting the top 6 heads, see Table 2.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: A.7. P RODUCING C ONTACT M APS
Revised: A.7. P RODUCING C ONTACT M APS

Original: A PDB structure gives 3D coordinates for every atom in a structure.
Revised: A PDB structure gives 3D coordinates for every atom in a structure.

Original: We use Euclidean distance between the beta carbons to deﬁne distance between any pair of positions.
Revised: We use Euclidean distance between the beta carbons to deﬁne distance between any pair of positions.

Original: A pair of positions where this distance is less than 8 ˚ A is declared to be a contact.
Revised: A pair of positions where this distance is less than 8 ˚ A is declared to be a contact.

Original: A.7. S CORING C ONTACT P REDICTIONS
Revised: A.7. S CORING C ONTACT P REDICTIONS

Original: Given a predicted contact map (
Revised: Given a predicted contact map (

Original: cid:98) C ∈ R L × L and a true contact map C ∈ { 0 , 1 } L × L , we describe metrics for scoring (cid:98) C .
Revised: cid:98) C ∈ R L × L and a true contact map C ∈ { 0 , 1 } L × L , we describe metrics for scoring (cid:98) C .

Original: A sequence x = ( x 1 , . . .
Revised: A sequence x = ( x 1 , . . .

Original: , x L ) of length L has (cid:0) L 2 (cid:1) potential contacts.
Revised: , x L ) of length L has (cid:0) L 2 (cid:1) potential contacts.

Original: Since we see O ( L ) contacts, contact prediction a sparse prediction task.
Revised: Since we see O ( L ) contacts, contact prediction a sparse prediction task.

Original: Accordingly, we focus on precision-recall based quantitative analyses of ˆ C .
Revised: Accordingly, we focus on precision-recall based quantitative analyses of ˆ C .

Original: Common practice in the ﬁeld is to sort all (cid:0) L 2 (cid:1) entries of ˆ C in decreasing order and evaluate precision at various length thresholds, such as the top L or L/ 10 predictions (Shrestha et al., 2019).
Revised: Common practice in the ﬁeld is to sort all (cid:0) L 2 (cid:1) entries of ˆ C in decreasing order and evaluate precision at various length thresholds, such as the top L or L/ 10 predictions (Shrestha et al., 2019).

Original: Note that this analysis is similar to choosing recall cutoffs along a precision-recall curve, where sorted length index plays the role of recall on the x axis.
Revised: Note that this analysis is similar to choosing recall cutoffs along a precision-recall curve, where sorted length index plays the role of recall on the x axis.

Original: Unlike recall, length-based cutoffs do not rely on knowledge of the actual number of contacts.
Revised: Unlike recall, length-based cutoffs do not rely on knowledge of the actual number of contacts.

Original: In addition to the precision at various length (recall) cutoffs, we also computed Area Under the Precision-Recall Curve (AUC).
Revised: In addition to the precision at various length (recall) cutoffs, we also computed Area Under the Precision-Recall Curve (AUC).

Original: AUC is a widely used metric for comparing classiﬁers when the positive class is rare.
Revised: AUC is a widely used metric for comparing classiﬁers when the positive class is rare.

Original: A. H YPERPARAMETERS
Revised: A. H YPERPARAMETERS

Original: Potts.
Revised: Potts.

Original: We used λ = 0 .
Revised: We used λ = 0 .

Original: 5 , learning rate of 0 . 5 , and batch size 4096.
Revised: 5 , learning rate of 0 . 5 , and batch size 4096.

Original: Pad, gap, and mask were all encoded with the same token.
Revised: Pad, gap, and mask were all encoded with the same token.

Original: The Potts model is trained using a modiﬁed version of Adam presented in Dauparas et al.
Revised: The Potts model is trained using a modiﬁed version of Adam presented in Dauparas et al.

Original: This modiﬁcation was made to improve performance of Adam to match that of L-BFGS.
Revised: This modiﬁcation was made to improve performance of Adam to match that of L-BFGS.

Original: Factored attention.
Revised: Factored attention.

Original: We AdamW with a learning rate of 5e-3 and set λ = 0 .
Revised: We AdamW with a learning rate of 5e-3 and set λ = 0 .

Original: The default head size was set to 32 unless stated otherwise.
Revised: The default head size was set to 32 unless stated otherwise.

Original: Single-layer attention.
Revised: Single-layer attention.

Original: We set embedding size of 256, head size of 64, and number of heads 128.
Revised: We set embedding size of 256, head size of 64, and number of heads 128.

Original: The model is trained with AdamW using a learning rate of 5e-3 and weight decay of 2e-3.
Revised: The model is trained with AdamW using a learning rate of 5e-3 and weight decay of 2e-3.

Original: Attention dropout of 0.05 is also applied.
Revised: Attention dropout of 0.05 is also applied.

Original: The batch size is 32 and mask prob for masked language modeling is 0.15.
Revised: The batch size is 32 and mask prob for masked language modeling is 0.15.

Original: We use a separate mask token and pad,gap token.
Revised: We use a separate mask token and pad,gap token.

Original: ProtBERT-BFD.
Revised: ProtBERT-BFD.

Original: ProtBERT-BFD has 30 layers each with 16 heads and a hidden size of 1024 .
Revised: ProtBERT-BFD has 30 layers each with 16 heads and a hidden size of 1024 .

Original: The training dataset is a mixture of UniRef50 (Suzek et al., 2015) and BFD.
Revised: The training dataset is a mixture of UniRef50 (Suzek et al., 2015) and BFD.

Original: It has 2 , 122 million protein sequences.
Revised: It has 2 , 122 million protein sequences.

Original: See Elnaggar et al.
Revised: See Elnaggar et al.

Original: (2020) more information.
Revised: (2020) more information.

Original: A.8. H YPERPARAMETER S WEEP D ETAILS
Revised: A.8. H YPERPARAMETER S WEEP D ETAILS

Original: Potts.
Revised: Potts.

Original: The Potts model implementation using psuedolikelihood has been optimized by others, so we did not tune performance.
Revised: The Potts model implementation using psuedolikelihood has been optimized by others, so we did not tune performance.

Original: Since performance with MLM was comparable to pseudolikelihood, we did not sweep for MLM either.
Revised: Since performance with MLM was comparable to pseudolikelihood, we did not sweep for MLM either.

Original: Single-layer attention.
Revised: Single-layer attention.

Original: Standard attention is by far the most sensitive model to hyperparameters.
Revised: Standard attention is by far the most sensitive model to hyperparameters.

Original: To ﬁnd a reasonable set of hyperparameters, we ﬁrst swept over the six families in Table 3, performing a grid search over
Revised: To ﬁnd a reasonable set of hyperparameters, we ﬁrst swept over the six families in Table 3, performing a grid search over

Original: • H ∈ { 32 , 64 , 128 , 256 , 512 } • d ∈ { 32 , 64 , 128 , 256 , 512 } • e ∈ { 128 , 256 , 512 } • attention dropout in { 0 , 0 . 05 , 0 . 1 } • learning rate in { 5 e − 3 , 1 e − 2 } • weight decay in { 0 , 1 e − 3 , 2 e − 3 }
Revised: • H ∈ { 32 , 64 , 128 , 256 , 512 } • d ∈ { 32 , 64 , 128 , 256 , 512 } • e ∈ { 128 , 256 , 512 } • attention dropout in { 0 , 0 . 05 , 0 . 1 } • learning rate in { 5 e − 3 , 1 e − 2 } • weight decay in { 0 , 1 e − 3 , 2 e − 3 }

Original: We found that the choice H = 256 , d = 64 , e = 256 , attention dropout of 0 .
Revised: We found that the choice H = 256 , d = 64 , e = 256 , attention dropout of 0 .

Original: 05 , learning rate of 5 e − 3 and weight decay of 2 e − 3 performed well across all six families.
Revised: 05 , learning rate of 5 e − 3 and weight decay of 2 e − 3 performed well across all six families.

Original: Due to GPU memory constraints, we had to set H = 128 for further runs.
Revised: Due to GPU memory constraints, we had to set H = 128 for further runs.

Original: Factored attention.
Revised: Factored attention.

Original: We swept factored attention over the families in Table 4, performing a grid search over
Revised: We swept factored attention over the families in Table 4, performing a grid search over

Original: • learning rate in { 1 e − 3 , 5 e − 3 , 1 e − 2 , 5 e − 2 } • regularization coefﬁcient in { 1 e − 4 , 5 e − 4 , 1 e − 3 , 5 e − 3 , 1 e − 2 }
Revised: • learning rate in { 1 e − 3 , 5 e − 3 , 1 e − 2 , 5 e − 2 } • regularization coefﬁcient in { 1 e − 4 , 5 e − 4 , 1 e − 3 , 5 e − 3 , 1 e − 2 }

Original: We found that learning rate of 5e-3 and regularization of 0.01 were effective, but that other conﬁgurations such as regularization of 5e-3 also performed well.
Revised: We found that learning rate of 5e-3 and regularization of 0.01 were effective, but that other conﬁgurations such as regularization of 5e-3 also performed well.

Original: Both H and d are evaluated extensively in our results.
Revised: Both H and d are evaluated extensively in our results.

Original: A. A DDITIONAL F IGURES
Revised: A. A DDITIONAL F IGURES

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: P r o p o r t i o n o f
Revised: P r o p o r t i o n o f

Original: R e s i
Revised: R e s i

Original: d u e s i n t r
Revised: d u e s i n t r

Original: R o s e tt a
Revised: R o s e tt a

Original: D a t a s e t
Revised: D a t a s e t

Original: Potts
Revised: Potts

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] 
Revised: [Figure] [Figure] [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: 512 AttentionHeadSize 0 .0 .0 .0 .0 .0 .0 .0 .
Revised: 512 AttentionHeadSize 0 .0 .0 .0 .0 .0 .0 .0 .

Original: P r ec i s i on a t L      
Revised: P r ec i s i on a t L      512 AttentionHeadSize 0 .0 .0 . [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Revised: t

Original: [Figure]
Revised: [Figure]

Revised: has a

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: Potts
Revised: Potts

Original: P@L 0.
Revised: P@L 0.

Original: Potts 0 .0 .0 .0 .1 .
Revised: Potts [Figure] [Figure] [Figure]

Revised: [Figure]

Original: F ac t
Revised: F ac t

Original: o r e d
Revised: o r e d

Original: A tt n
Revised: A tt n

Original: F r o ze n
Revised: F r o ze n

Original: V a l u e s    # o f
Revised: V a l u e s    # o f

Original: M S
Revised: M S

Original: A s e qu e n ce s
Revised: A s e qu e n ce s

Original: [Figure]
Revised: [Figure]

Original: s (
Revised: of

Original: [Figure]
Revised: [Figure]

Original: illi

Original: [Figure]
Revised: [Figure]

Revised: m e t e r s (

Revised: M illi on )

Original: Potts
Revised: Potts

Original: H=
Revised: H=

Original: H= H=
Revised: H= H=

Revised: [Figure]

Revised: A tt

Original: [Figure]
Revised: [Figure]

Original: A
Revised: P

Original: [Figure]
Revised: [Figure]

Original: L
Revised: P

Original: [Figure]
Revised: [Figure]

Revised: L

Revised: APC

Revised: False

Original: True
Revised: True

Original: [Figure]
Revised: [Figure]

Original: e
Revised: s

Original: [Figure]
Revised: [Figure]

Revised: e n

Original: [Figure] [Figure] 
Revised: [Figure] [Figure] [Figure]