Fast Uncovering of Protein Sequence Diversity from Structure

## Abstract

Abstract We present InvMSAFold, a method for generating a diverse set of protein sequences that fold into a single structure. For a given structure, InvMSAFold defines a probability distribution over the space of sequences, capturing the amino acid covariances observed in Multiple Sequence Alignments (MSA) of homologous proteins. This allows for the generation of highly diverse protein sequences while preserving structural and functional integrity. We show that the higher diversity of sampled sequences translates into higher diversity in biochemical properties, pointing to exciting prospects for the applicability of our method in fields like protein design by providing diverse starting points.

## 1 Introduction

Inverse folding aims to predict amino acid sequences that fold into a given protein structure, and plays a fundamental role for example in the protein design pipeline of RFDiffusion []. Recent deep learning approaches such as ESM-1F [] or ProteinMPNN [] achieve remarkable accuracy in this task.
However, instead of predicting a single ground truth sequence, it is often desirable to have a method that is able to generate a variety of different sequences with the desired fold, i.e., solving a *many-to-one* problem, see Fig. 2. This diversity could be leveraged for example by starting from a source sequence [, ] and taking different molecular environments into consideration []. Such an approach would allow to expand the sequence design space while preserving structural consistency, allowing for a larger pool of sequences when selecting for additional properties like thermostability, solubility or toxicity. In drug discovery, for example, it would allow for the generation of large amount of diverse candidates, enabling further selection optimized for properties like bioavailability. Similarly, in biotechnology and enzyme engineering, it would facilitate the creation of enzymes with tailored properties, such as improved stability and activity under varying conditions.

In this work, we present a method that is able to generate diverse protein sequences given a structure, including sequences far away from the natural sequence, see Fig. 1.

Figure: Figure 1: InvMSAfolds expands the scope of inverse folding to the retrieval of the the entire landscape of homologous proteins with similar folds.
Refer to caption: x1.png

Recent architectures for inverse folding are based on encoder-decoder architectures, where a structure is encoded and a sequence decoded. Such models typically take into account only the native sequence of a given structure, training the model by maximizing its probability given the structure [, ].

Figure: Figure 2: *Many-to-one* nature of inverse folding. Top: Structure of 1KA0. Bottom: Some homologos sharing the fold.
Refer to caption: extracted/5669413/images/test_pymol.png

We show that the models we generate are able to capture the diversity of the protein family better than other models and are able to find sequences far away from the natural sequence that are predicted to still fold into the same structure. We also show that this increased diversity translates into a more spread distribution in other properties, enabling selection of promising sequences from a larger pool.

## 2 Methods

In this Section we describe the components of our architecture, InvMSAFold, and its training procedure. Given a structure-sequence pairing $(\bm{X},\bm{\sigma}_{X})$, inverse folding methods typically define a probability distribution on the space of sequences $\bm{\sigma}$ as $p(\bm{\sigma}|\bm{X},\bm{\sigma}_{X})$. Most deep learning based methods structure this distribution auto-regressively, using $p(\bm{\sigma}|\bm{X},\bm{\sigma}_{X})=\prod_{i=1}^{L}p(\sigma_{i}|\sigma_{i-1}
,\ldots,\sigma_{1},\bm{X},\bm{\sigma}_{X})$, where $L$ is the length, and train by minimizing the loss to the true sequence $\bm{\sigma}_{X}$, possibly after adding noise to the coordinates $\bm{X}$ [, ].

Both of these choices have drawbacks: Sampling amino acids auto-regressively requires a full forward pass through the neural network for every generated token, making it very expensive to use the models in a virtual-screening-like setting, where a large number of sequences are scanned for properties beyond folding into structure $\bm{X}$. Secondly, minimizing the loss on the true sequence only ignores the *many-to-one* property of inverse folding depicted in Figure 2. Even if the training is successful, one would expect the resulting distribution to be peaked around very few sequences and not capturing other parts of sequence space that might be interesting for the problem at hand.

##### The InvMSAFold architecture.

In this work, we propose InvMSAFold, which deviates from the methods described above in two important ways: Instead of directly returning a probability distribution over sequences of amino acids, we use the neural network to generate a set of parameters for a secondary probability distribution that is then used for sampling amino acids. With a little abuse of notation, we will consider $p(\bm{\sigma}|\bm{X},\bm{\sigma}_{X})=p(\bm{\sigma}|\bm{\theta}(\bm{X}),M_{X})$, where $M_{X}$ will be a subset of a multiple sequence alignment(MSA) and $\bm{\theta}(\bm{X})$ are parameters of a light-weight model generated by a neural network that is given the backbone coordinates $\bm{X}$ as input.

We explore different choices for parameterizing the light-weight model. In order to go beyond distributions that treat different positions in the protein as independent and to ensure sufficient expressivity, we focus on pairwise models [] of which a particular functional gives the well known Potts model [] from statistical physics. Such models have an experimentally validated ability to capture structure-sequence relationships []. We therefore train the neural network to learn the mapping $\bm{X}\rightarrow\Theta=(\bm{J},\bm{h})$, where $\bm{J}$ and $\bm{h}$ are tensors of parameters of size $L\times q\times L\times q$ and $L\times q$ respectively, $L$ is the sequence length and $q$ is the number of different amino acids. As is common in pairwise models, we call the quantities $h_{i}(a)$ the fields, indexed by position $i$ and amino acid type $a$, and the quantities $J_{ij}(a,b)$ the *couplings*, indexed by a pair of positions and a pair of amino acid types. The former describe the propensity of an amino acid to appear at a given position, while the latter describe the propensity of pairs of amino acids to appear at pairs of positions.

The complete InvMSAFold architecture is composed of two parts, the structure encoder and a decoder which outputs the fields and couplings, see Fig. 3.

Figure: Figure 3: Left: Decoder architecture up to the generation of the low-rank tensor $\mathcal{V}$. Right: The low-rank tensor $\mathcal{V}$ is used to generate the couplings $J$ and the fields $h$.
Refer to caption: x2.png

For the encoder, we use the pre-trained encoder from the ESM-IF1 model of [] which follows the GVP-GNN architecture proposed in []. This encoder gives a rotationally invariant representation of the input $\bm{X}$, which has been show to be crucial to properly model the geometric features coming from the protein’s 3D structure. During training we add Gaussian noise to these representations as commonly done in previous research [].

The decoder takes as input a batch of encoded structures, padded to a common length $L$, which each get embedded to a $L\times D$ tensor ($D=512$ in our experiments) and then pass through $6$ transformer layers with $8$ attention heads. The output is first embedded to into a $L\times DK$ dimensional space through a linear layer followed by an element-wise Relu activation function, which then gets reshaped to a $L\times K\times D$ tensor and is finally projected to a $L\times K\times q$ tensor $\mathcal{V}$ through another linear layer. By selecting $K\ll L$, we then compute a low-rank matrix coupling from $\mathcal{V}$ as follows

$$ $J_{ij}(a,b)=\frac{1}{\sqrt{K}}\sum_{k=1}^{K}v_{i}^{k}(a)v_{j}^{k}(b).$ (1) $$

By computing $\bm{J}$ as (1) we drastically reduce the number of parameters of our pairwise model from $\mathcal{O}(L^{2})$ to $\mathcal{O}(L)$, while the scaling by $\sqrt{K}$ ensures that the couplings remain $\mathcal{O}(1)$. We note that in other settings, low-rank decompositions of $J$ (1) have been shown to be as effective as the full-rank counterparts []. The fields $\bm{h}$ are computed by passing $\mathcal{V}$ through another linear layer, and then contracting the tensor by summing the entries across the latent dimension. Also here after summing we scale by $\sqrt{K}$ to ensure the fields to remain $\mathcal{O}(1)$.

##### InvMSAFold-PW and InvMSAFold-AR

Given the parameters $\bm{J}$ and $\bm{h}$ we explore two different approaches to define a probability distribution over sequences of amino acids. The first distribution, which we term InvMSAFold-PW, is a standard pairwise distribution [] that defines the negative log-likelihood of a sequence $\bm{\sigma}$ as

$$ $\displaystyle\log p^{pw}(\bm{\sigma}|\bm{h,J})\sim\sum_{i=1}^{L}\sum_{j=i+1}^{ L}J_{ij}(\sigma_{i},\sigma_{j})+\sum_{i}h_{i}(\sigma_{i}).$ (2) $$

While this distribution has been explored extensively for proteins [] as it possesses some nice theoretical properties [], it has the disadvantage that the normalization factor for Eq. 2 is intractable, and therefore also the likelihood. We cannot hence train by *maximum likelihood*, but have to resort to the approximation given by pseudo-likelihoods []. Moreover for sampling we have to resort to MCMC algorithms, which is known can struggle to navigate efficiently the landscape given by (2).

Based on the above practical limitations, we consider an alternative efficient auto-regressive variant of pairwise models [], which we term InvMSAFold-AR, that defines as autoregressive distribution $p^{ar}$ over amino acids having negative log-likelihood

$$ $\displaystyle\log p^{ar}(\sigma_{i}|\sigma_{1},\ldots,\sigma_{i-1},\bm{h,J}) \sim h_{i}(\sigma_{i})+\sum_{j=1}^{i-1}J_{ij}(\sigma_{i},\sigma_{j}).$ (3) $$

Since this parameterization decouples into a sequence of univariate distributions, it has the advantage that the normalization factor is tractable allowing for closed form computation of the likelihood. As a result we can train with *maximum likelihood* and also sample much more efficiently from the model. Since for InvMSAFold-AR the number of addends in the sum in Eq. 3 depends on the position $i$, we rescale the couplings as

$$ $J_{ij}\leftarrow\frac{J_{ij}}{max(i,j)},$ (4) $$

following what was done in []. We found this to be beneficial for training and note that without this scaling the neural network would have to generate couplings of significantly different magnitudes for different sites. This would not be a problem if we were to optimize the couplings directly, as in previous research [], but might be problematic if the couplings are generated by a neural network.

##### Training on Homologous Sequences

In most other works [, ], inverse folding models are trained to predict the ground truth sequences only. In this work, we aim to generate a distribution that captures the complete sequence space that is compatible with the input structure $\bm{X}$. To this end, we use the ground truth sequence corresponding to a structure $\bm{X}$ and extract an MSA $M_{X}$ from sequence database (see next section for details). We then use the pairs $(\bm{X},M_{X})$ for training by taking the mean negative pseudo log-likelihood of a random subsample of sequences in $M_{X}$ given the parameters generated by the network. For both InvMSAFold-PW and InvMSAFold-AR, we add a regularizing $L_{2}$ term for the fields and the couplings as is typical for these models [], []. For details on the training procedure our models, see Sec. 3.1.

##### Data and Train-Test Splits

In order to control the level of homology in our evaluation, we create three test sets, which we call the inter-cluster test set, the intra-cluster test set, and the MSA test set. We base these on the CATH database [], which classifies protein domains into superfamilies and then further clusters into the superfamilies based on sequence homology. We use the non-redundant dataset of domains at 40% similarity and associate to every domain a cluster as indicated in the CATH database. We then chose 10% of the sequence clusters uniformly at random and assign them to the inter-cluster test set, excluding these clusters from the training set. Because many superfamilies contain only one sequence cluster, there is a significant amount of superfamilies that appear only in the inter-cluster test set and not in the training set, making this a hard test set. We then create the less stringent intra-cluster test set by taking from every domain sequence cluster that is not in the inter-cluster test set and has at least two domains a single random domain. We then use the remaining domains as the training set. Finally, we create MSAs for all sequences in the datasets using the [MMseqs2](https://github.com/soedinglab/MMseqs2) software and the Uniprot50 database. We further split the sequences in the MSAs into 90% used for training and 10% for the MSA test set. This last test set is the least stringent one since it is based on domain sequences that are also in the training set.

## 3 Results

### 3.1 Model Training

For both training and inference we create embeddings for the structures using the ESM-IF1 encoder []. We then add independent Gaussian noise with standard deviation equal to $5\%$ of the standard deviation of the embeddings across all positions, dimensions and samples in the training set. To create the subsets of sequences necessary to compute the losses for the models we random subsample from the target MSA at each training step, where the number of sequences sampled is a hyperparameter. For InvMSAFold-PW, we train with a single structure in each batch, with a MSA subsample size for $M_{X}$ of $64$, a rank $K$ of 48, a learning rate of $10^{-4}$ and L2 regularization constants of $\lambda_{h}=\lambda_{J}=10^{-4}$ for fields and couplings. For InvMSAFold-AR, we select the above parameters through hypertuning. We refer the reader to Appendix A.3 for the details. Both models are trained with AdamW optimizer for a total of $94$ epochs. We monitor the the negative pseudo log-likelihood (npll) for InvMSAFold-PW and the negative likelihood (nll) for InvMSAFold-AR on the train and the different test sets. Results are shown in Figure 4.

Figure: Figure 4: Negative pseudo log-likelihood (npll) for the train and test sets by training epochs. The training loss is a rolling average over a window of $5$ epochs, which is why the plotted value is higher than the test losses initially.
Refer to caption: extracted/5669413/images/train_test_lossess_pw_ar.png

As can be seen from Fig. 4, for both models the ordering of the losses on the different datasets is consistent with the hardness reasoning behind the split in the last section: The *MSA* test set has sequences coming from domains which we have indeed seen during training, and the loss on this test set is similar to the training loss. The *intra-cluster* test set contains domains that we have not seen during training, but of which we have seen domains belonging to the same sequence cluster; the loss on this test set is higher than the training loss and seems to saturate during training. For the *inter-cluster* test dataset, which contains domains that belong to sequence clusters not present in the training set, the loss is significantly higher, even though still much better than random.

### 3.2 Covariance Reconstruction

The ability of a generative model to reproduce covariances between amino acids in MSAs has been shown to be a good metric for measuring how well the protein landscape is captured []. Models with this ability have been shown to enable
efficient sampling of experimentally validated, functional protein sequences [].

In order to test this ability, we created MSAs with synthetic samples from InvMSAFold-PW, InvMSAFold-AR and ESM-IF1, and compared the amino acid covariances in these MSAs with the covariances in the MSAs of natural sequences.
Given an MSA of sampled or natural sequences, we define the covariance between amino acids $a$ and $b$ at positions $i$ and $j$ as $C_{ij}(a,b)=f_{ij}(a,b)-f_{i}(a)f_{j}(b)$, where $f_{ij}(a,b)$ is the frequency of finding amino acids $a$ and $b$ at these positions in the same sequence in the MSA and $f_{i}(a)$ and $f_{j}(b)$ are the overall frequencies of the amino acids $a$ and $b$ at these positions. For a given MSA there are $q\times\binom{N}{2}$ covariances, where $q$ is the number of possible amino acids and $N$ the sequence length. In order to compare two sets of covariances, we calculate their Pearson correlation as in previous research []. For details on these experiments, we refer to the Appendix A.2.1.

As we can see from Figure 5 and the attached table, InvMSAFold-AR and InvMSAFold-PW capture the covariances better than ESM-IF1, especially at longer sequence lengths. In fact, InvMSAFold-AR shows the best performance overall, and this performance does not seem to deteriorate for longer sequences.

ESM-IF1 is the worst of the tested models for describing covariances; this result is not surprising since it was not trained on this task. The stronger performance of InvMSAFold-AR compared to InvMSAFold-PW is likely due to the latter being trained with pseudo-loglikelihoods, which do not capture covariances perfectly even when training a pairwise model directly [].

The result strengthens our confidence that InvMSAFold is able to model the sequence landscape of an unseen structure. We explore this in more detail in the next section.

Figure: Figure 5: Distribution of Pearson’s correlation coefficients between the covariances from sampled and natural sequences for domains having at least $2k$ natural sequences. Top: KDE plots representing how the correlations for InvMSAFold-AR (blue), InvMSAFold-PW (red) and ESM-IF1 (green) vary with the sequence length. The top row shows results for the Intra-Cluster test set, the bottom row for the Inter-Cluster test set. Bottom: Table with quartiles of Pearson correlations corresponding to the above plots for the different methods and Inter/Intra Cluster test set. Best model for each quartile is highlighted by boldface numbers.
Refer to caption: x3.png

### 3.3 Sequence Patterns

In this Section, we compare the PCA projections of MSAs of sampled sequences and natural sequences. To this end, we sampled 2000 sequences each from InvMSAFold-AR, InvMSAFold-PW, and ESM-IF1 and projected the sequences onto the first few principal components of the one-hot encoded MSA of natural sequences.

Results for an NDP Kinase (1xqi, []) are shown in Fig. 6, with more results in the Appendix. The most striking feature is the narrow focus around a single point in this space for ESM-IF1. This suggests that the distribution of ESM-IF1
is centered around a single sequence, with a very small coverage of the sequence space in general.

Figure: Figure 6: Sampled sequences projected onto the first two PCA components of natural sequences for 1xqiA00. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x4.png

**Table 1: Kullbach-Lieber Divervence between between Kernel Density of the natural sequence and sampled sequence. More precisely, for each method (esm, ports, ardca we sampled 2000 sequences for each backbone. We then project these sequences on the two main PCA components of the natural MSA. In this 2D space, we apply a Gaussian kernel density estimator of kernel size 1.0. We then use these densities to compute the KL divergence between the space of natural sequences and the one generated by the inverse folding approach.**
|  | ESM-IF1 | InvMSAFold-PW | InvMSAFold-AR |
| --- | --- | --- | --- |
| Inter-Cluster | 15.8 | 4.6 | $\bm{0.49}$ |
| Intra-Cluster | 11.9 | 11.2 | $\bm{0.67}$ |

Sequences sampled from InvMSAFold-PW show a broader coverage of sequence space but are still focused around the single dominant mode. This could be a result of the approximate training given the *pseudo likelihood*, or of the challenges of sampling from multimodal distributions for MCMC algorithms. While this increased diversity in the samples might already be useful in some settings, it is notable that secondary modes are not captured at all, and that even the dominant mode is still covered only in part.

In contrast, sequences sampled from InvMSAFold-AR cover the sequence space in this representation more comprehensively, demonstrating an ability to cover multiple distinct modes.

### 3.4 Predicted Structures of Sampled Sequences

Similarly to other works [, ], we test to what extent the sequences generated by our models are predicted to fold into the correct structure when fixing a desired sequence distance (number of different amino acids divided by the sequence length) to the native sequence.

In Fig. 7, we show a density plot for the sequence distances of sampled sequences of InvMSAFold-AR, InvMSAFold-PW and ESM-IF1 from the native sequences of $15$ structures in the inter-cluster test set.
InvMSAFold-AR and InvMSAFold-PW are able to generate sequences with low sequence similarity to the native sequence without further modifications. This is not the case for ESM-1F, which often generates sequences very close to the native sequence. We therefore sample sequences from ESM-IF1 at a higher temperature. We detail the sampling settings in the Appendix A.2.2

Figure: Figure 7: Sequence distances to native sequences, averaged over $15$ random structures rfom the inter-cluster test set.
Refer to caption: extracted/5669413/images/morekls/hamming_distances.png

Following the generation of these sequences, we predicted their structures with AlphaFold 2 [, ] with templates switched off. In Fig. 8 we compare the RMSD between the original structure and the refolded Alphafold structure on a series of examples and plot the results in dependence on the sequence distance from the natural sequence. As can be seen, for smaller distances the best-performing method is ESM-IF1. For larger distances, however, the performance of ESM-IF1 drops significantly. For both InvMSAFold-AR and InvMSAFold-PW, the performance deteriorates as well with increasing sequence distance, but more slowly than for ESM-IF1.

Figure: Figure 8: Comparison of the quality of the generated sequences under the constraint of increasing sequence distance from the native sequence. Shown is an average of 14 structures from the inter-cluster test set. We then refold the sequence with Alphafold and compare the refolded structure with the original one using RMSD. List of structures used: 1otjA00 2ytyA00 1p9hA00 2o0bA02 1b34B00 1f6mA02 2hs5A01 2bh8A02 2de6A03 1ia6A00 4gc1A01 3sobB02 1xqiA00 4yt9A01
Refer to caption: x5.png

Figure: Figure 9: Comparison of the quality of the generated sequences under the constraint of increasing sequence distance from the native sequence. Shown is an average of 14 structures from the inter-cluster test set. We then refold the sequence with Alphafold and compare the refolded structure with the original one using RMSD. List of structures used: 5i0qB02 4iulB00 2wfhA00 2egcA00 3oz6B02 3m8bA01 2pz0B00 2cnxA00 4k7zA02 1dl2A00 1vjkA00 1udkA00 3bs9A00
Refer to caption: x6.png

## 4 Protein Property Sampling

Possible applications for InvMSAFold are protein design tasks where other than structural properties of the designed protein are important. In a virtual screening like setting, one might want to generate a large number of sequences that obey certain structural constraints and then select from this set sequences that have other desirable properties. We test this setting by analyzing the range of (predicted) thermal stabilities and solubility of the sequences generated by the different models.

In order to predict thermal stabilities we use Thermoprot [], which takes into account whether the protein operates within a Mesophile or Thermophile species. For solubility, we use predictions from Protein-Sol [].
We generated sequences for domains from our test sets, aligned them, and predicted their properties.
Since the sequences have different numbers of aminoacids due to alignment, we verified that the number of gaps in the alignment was not discriminative for the two properties. We show the results for a single structure in Fig. 10, with more examples being shown in the Appendix B.2.

Figure: Figure 10: Comparison of the distribution of predicted solubility and thermostability of samples generated with InvMSAFold and ESM-IF1. This plot shows the results for domain 1ny1A00, with more examples available in the Appendix.
Refer to caption: x7.png

We observe that the support of InvMSAfold is significantly larger than of ESM-IF1, showing that the larger sequence diversity generated by both InvMSAFold-PW and InvMSAFold-AR translates into a wider range of protein properties. Note also that while we generated the same number of sequences for a fair comparison, the computational resources are very different. The InvMSAFold models need a single forward pass through the neural network to generate the parameters for the light-weight pairwise distributions. After this, diverse sequences can be sampled efficiently. For InvMSAFold-AR, for example, we can sample millions of sequences on CPUs in seconds.

## 5 Conclusion

InvMSAFold is an approach to inverse folding that combines the idea of modeling the complete sequence space conditioned on structural information and computational efficiency. The core idea is to approach the problem as a two-stage generation process, where we use an expressive large neural network for generating the parameters of a sufficiently expressive family-specific model. During training we compute the loss on a data augmentation obtained through an MSA that, by capturing the variability of sequences compatible with the input, ensures that the model does not focus narrowly on the native sequence. This latter idea is not specific to our model formulation, and could be exploited also by other architectures as ESM-1F. Once trained, such a model can be used to sample a large number of sequences efficiently. We validate InvMSAFold by showing that we are able to generate sequences far away from the native sequence while still retaining structural fidelity. We also show that the shift in distribution can enable the sampling of a broader range of protein properties of interest.

We believe InvMSAFold enables novel approaches to virtual screening, where one might for example optimize for properties like thermostability, altered substrate specificity, or reduced toxicity at the same time.

## Appendix A Appendix

### A.1 Fast pseudolikelihood computation

The main rationale behind the low rank approximation (1) is two-fold; first, it is known that, compared to the covariance matrix, low-rank
decompositons of the coupling matrix $\bm{J}$ have been shown to be as effective as the full-rank counterparts []. Secondly, for computational purposes we wanted to avoid bottleneck’s of quadratic costs in the length of the protein, both from a memory and a computational time point-of-view. Technically, once we enforce a low rank constraint of the matrix $J$, the number of active parameters indeed passes from $\mathcal{O}(L^{2})$ to $\mathcal{O}(L)$. If we though compute the pseudo-likelihood naively by pre-computing the full coupling matrix $J$ as in (1), we end up again into a quadratic cost. Luckily, with a careful implementation, we can indeed achieve a linear memory and computational cost. To see this, let us rewrite the crucial part(as the fields are by nature linear in cost):

$$ $\displaystyle\sum_{i<j,a,b}J_{i,j}[a,b]\delta_{i}[a]\delta_{j}[b]$ $\displaystyle=\frac{1}{2}\left(\sum_{i,j,a,b}\sum_{k=1}^{K}v^{k}_{i}[a]v_{j}^{ k}[b]\delta_{i}[a]\delta_{j}[b]-\sum_{k=1}^{K}\sum_{i,a}(v_{i}^{k}[a]\delta_{i }[a])^{2}\right)$ (5) $\displaystyle=\frac{1}{2}\left(\sum_{k=1}^{K}\left(\sum_{i,a}v^{k}_{i}[a] \delta_{i}[a]\right)^{2}-\sum_{k=1}^{K}\sum_{i,a}(v_{i}^{k}[a]\delta_{i}[a])^{ 2}\right).$ (6) $$

In the previous equation $\delta_{i}[a]$ is the Kronecker symbol, which equals $1$ is the amino-acid in the i-th position is equal to $a$, and zero otherwise. This has solved half of the problem, in the sense that we can now compute every term of the pseudo-likelihood summation in linear time, but again if we naively compute each term of the summation independently, we again end up with a cost of $\mathcal{O}(L^{2})$. Logically though the different terms of the pseudo-likelihood have a lot of shared structure, hence we can hope to recover a linear cost. Indeed, using Bayes’ theorem, we have that

$$ $p(\sigma_{p}|\bm{\sigma}_{-p})=\frac{p(\bm{\sigma})}{\sum_{c}p(\sigma_{p}=c, \bm{\sigma}_{-p})}=\frac{\exp\{-E(\bm{\sigma})\}}{\sum_{c}\exp\{-E(\sigma_{p}= c,\bm{\sigma}_{-p})\}}.$ (7) $$

we want to compute efficiently the energies and the calculate the logarithm of the loss. For the moment we will discard the fields contribution to the energy $E$ as its cost is linear and needs no careful implementation. The crucial part is that we should do minimal effort to calculate $E(\sigma_{p}=c,\bm{\sigma}_{-p})$ once we have already calculated $E(\bm{\sigma})$. To do this we note that, labelling as $\delta_{p}^{c}$ the indicator variable for the mutates sequence having amino acid $c$ in position $p$, we have that

$$ $\sum_{i,a}v^{k}_{i}[a]\delta_{i}^{c}[a]=\sum_{i\neq p,a}v_{i}^{k}\delta_{i}[a] +v_{p}^{k}[c]=\sum_{i,a}v_{i}^{k}\delta_{i}[a]-\sum_{a}v_{p}[a]\delta_{p}[a]+v _{p}^{k}[c].$ (8) $$

As we can see the first term is common to all energies, hence we can calculate it just once, while the other have a smaller cost. The thing to remember is that for any $p(\sigma_{p}|\bm{\sigma}_{-p})$ we have to calculate all the energies $E(\sigma_{p}=c,\bm{\sigma}_{-p})$ for all the dictionary letters $c$, and the we have to iterate across all positions $p$ to get the final loss.

Notice, that in case we also insert a quadratic penalty, we also have to make sure to calculate the latter in $\mathcal{O}(L)$. This though is rather simple, in fact one can just observe the following

$$ $\sum_{i,j,a,b}J_{i,j}[a,b]^{2}=\sum_{k,k^{\prime}}\left(\sum_{i,a}v_{i}^{k}[a] v_{i}^{k^{\prime}}[a]\right)^{2}-\sum_{k,k^{\prime}}\left(\sum_{i}\left(\sum_{ a}v_{i}^{k}[a]v_{i}^{k^{\prime}}[a]\right)^{2}\right),$ (9) $$

### A.2 Experiments details

#### A.2.1 Second order reconstruction

The experimental procedure for both datasets and can be organized in the following steps:

- 1.
We filtered the structures, keeping only those for which the MSA generated by [MMseqs2](https://github.com/soedinglab/MMseqs2) has at least $2k$ sequences. This because since we need to compute covariances, small MSAs could give very biased benchmarks, leading to wrong conclusions.
- 2.
For every sequence in the filtered dataset, we generate $10k$ syntetic samples for all of the three models under consideration. To generate the samples from Potts we leverage the efficient library [bmDCA](https://github.com/ranganathanlab/bmDCA), for ESM-IF1 we used the built in sampler(with some changes for speed improvements) and for ArDCA we built our own sampler. For bmDCA, we ran $10$ parallel chains and then pooled the results, for more details on the sampler parameters refer to the supplement.
- 3.
For ESM-IF1, given the samples, we re-aligned the samples using the full MSA to get a fair comparison. This because, while arDCA and Potts have seen many gaps during training, ESD-IF1 has never seen one since he focuses only on the native sequence. Hence he will produce un-aligned sequences which seldomly have gaps. To align we used the [PyHMMER](https://pyhmmer.readthedocs.io/en/stable/) library.
- 4.
Given the samples, we compute the covariance matrices of the generated samples and the one of the true MSA. We then compute the Pearson correlation between the flattening of the two.

#### A.2.2 Predicted Structures of Sampled Sequences

Synthetic sequences from both InvMSAFold-AR and InvMSAFold-PW, given the fact that they are trained on very diverse batches of a structure’s MSA at every training step, naturally produce sequences with the required diversity from the native sequence. On the other hand, ESM-1F1 produces sequences whose hamming distance from the native sequence is significantly lower compared to other two models. Given that we want to tests the different models’ ability to recover the native structure from sequences having high dissimilarity from the native sequence, to get such sequences from ESM-1F we leverage the built-in temperature parameter of the model. Increasing the temperature of the model allows to smooth the overall landscape of the model and sample more from the tails of the distribution, where the more diverse sequences will be.

Specifically, InvMSAFold-PW and InvMSAFold-AR generate sequences having normalized hamming distance between $0.65$ and $0.9$. We hence split the inverval $[0.65,0.9]$ into $5$ bins of equal width of $0.05$. Since we want to assess the three different models ability at different normalised hamming distances, we ideally would like to fill these bins with $10$ sequences having hamming distance within the specified ranges for each model. Unfortunately this will not be possible for all proteins, as some for some domains InvMSAFold-PW and InvMSAFold-AR might only be able to generate sequences with a having distance over $0.75$ for example, hence for those proteins the first two bins would remain empty. For ESM-1F the problem is the opposite, as at temperature $1$, the generated sequences do no fill the bins of highest dissimilarity even when a high number of sequences is generated; for ESM-1F we can overcome this by increasing the temperature of the model, while for InvMSAFold-PW and InvMSAFold-AR there is no way to overcome this problematic. Then againm our interest is to model highly dissimilar sequences, so this situation for InfMSAFold-PW and InvMSAFold-AR is not a problem. We hence iterate the following steps until all the bins are filled for ESM-1F

- 1.
Generate $1000$ samples from ESM-1F
- 2.
Fill the bins which have not yet been filled in previous steps
- 3.
If all the bins are full, return the selected sequences, otherwise increase the temperature of the model by $0.1$ and go back to step $1$.

While for InvMSAFold-PW and InvMSAFold-AR we just generate $10k$ sequences each for each domain, and then fill the bins relative to every domain.

### A.3 Hyper-tuning details

To train select some hyperparameters for arDCA we relayed on the library [Optuna](https://optuna.org/). The parameters we optimized for where Adam’s learning rate($lr$), the rank of the approximation for the Couplings matrix $J$($K$), the penalties for the couplings and the fields ($\lambda_{J},\lambda_{h}$), the batch size used for training($B$), the batch MSA size used in the loss($M$) and the value of dropout. To sample parameters we use the TPESampler, and we allowed for median pruning of bad trials to improve the speed of out hyper-tuning. We gave optuna $50$ trials of $90$ epochs each(irrespective of batch size), while as a selecting metric we used the average of the pseudo-likelihood on the inter- and intra-cluster test dataset. In the table below we recap the parameter values selected by the hyper-tuning

**Table 2: Parameters selected arDCA by hyperparamter optimization**
| Hypertuning results |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Model | Dropout | B | M | K | $(\lambda_{J},\lambda_{h})$ | lr |
| ArDCA | 0.1 | 8 | 32 | 48 | (3.2e-6, 5.0e-5) | 3.4e-4 |

We actually applied an identical strategy for those same parameters for the Potts model, but there the hyper-tuned values did not perform better(actually slightly worse) than those reported in the Methods 2 section. It is definitely something we would like to investigate further in subsequent research, but seemed out of scope for the current work.

## Appendix B Additional Plots

### B.1 PCA

Figure: Figure B.1: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x8.png

Figure: Figure B.2: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x12.png

Figure: Figure B.3: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x16.png

Figure: Figure B.4: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x20.png

Figure: Figure B.5: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x24.png

Figure: Figure B.6: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x28.png

Figure: Figure B.7: Sampled sequences projected onto the first two PCA components of natural sequences for various PDBs. We also used this density estimate to compute the Kullbach-Leiber (KL) divergence between the density of the natural data and the density of the sampled data. The values of these KLs are written in the title of each subplots.
Refer to caption: x32.png

### B.2 Bivariate Plots: Thermostability vs Solubility

Figure: Figure B.8: Comparison of the distribution of predicted solubility an thermostability of samples generated with InvMSAFold and ESM-IF1. These plots show the results for domain 1ia6A00, 1f6mA02, 1b34B00.
Refer to caption: x35.png

Figure: Figure B.9: Comparison of the distribution of predicted solubility an thermostability of samples generated with InvMSAFold and ESM-IF1. These plots show the results for domain 1ny1A00, 1otjA00, 1viuA00.
Refer to caption: x38.png

Figure: Figure B.10: Comparison of the distribution of predicted solubility an thermostability of samples generated with InvMSAFold and ESM-IF1. These plots show the results for domain 1m1sA00, 1xqiA00, 2hs5A01.
Refer to caption: x41.png

Figure: Figure B.11: Comparison of the distribution of predicted solubility an thermostability of samples generated with InvMSAFold and ESM-IF1. These plots show the results for domain 2m5jA00, 2x64A02, 3cx3A01.
Refer to caption: x44.png

Figure: Figure B.12: Comparison of the distribution of predicted solubility an thermostability of samples generated with InvMSAFold and ESM-IF1. These plots show the results for domain 3qwwA03, 3sobB02, 4atkB00.
Refer to caption: x47.png

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: * How is sampling of InvMSAFold-PW achieved? Which MCMC algorithm do you use? 
* By using a PCA projection you show that sequences generated by InvFoldMSA have a better coverage of sequence space. But why do you restrict the analysis to the first two principal components? 
* Have you tried AlphaFold3 to validate the sequences generated by InvFoldMSA?

__Typos / grammar__

* Line 117: "and whos outputs"
* Line 212: "can be reduce to"
* Line 248: "robsutly"
* Line 277: "chose" - should be present tense
* Line 302/303: What do you mean by "consistent with the hardness reasoning behind the split"
* Line 475: "becoming worse that both"
* The use of the symbol $\\propto$ to indicate quality up to an additive constant is a bit unusual.
soundness: 3
strengths: * An interesting idea to approach the inverse folding problem (i.e. the problem of generating sequences that fold into a given structure). 
* Proposes a low-rank approximation of the couplings and fields of the lightweight sequence model. 
* Fast generation of sequences that fit a well to a given structure.
summary: InvMSAFold is an inverse folding method that is optimized for diversity and speed. The general idea is to use a neural net to predict from an input structure and sequence a pairwise interaction model (a Potts model or Boltzmann machine) that captures the structure-sequence relationship and can be used to efficiently generate sequences that differ largely from the input sequence. To tame the number of parameters (fields and pairwise couplings), InvMSAFold predicts a low-rank approximation of the coupling matrix. The paper proposes two models: 1) InvMSAFold-PW is a full pairwise model that reduces the number of parameters significantly and also allows for efficient learning by using a maximum pseudo-likelihood. A drawback is that sequence generation requires MCMC. 2) InvMSAFold-AR is an autoregressive model whose likelihood is tractable thereby allowing for Maximum-likelihood parameter estimation as well as sampling of sequences in a straight forward fashion. Using various metrics the authors show that InvMSAFold, and in particular InvMSAFold-AR, outperforms current state of the art.
weaknesses: * The idea of generating a Potts model has already been proposed by Li et al. (2023).

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: I was not clear on the InvMSAFold-AR/-PW. I understand that PW requires MCMC sampling and AR does not but I wonder are there cases/tasks in which one of the two PW/AR models is more appropriate? 

What would be an example in which you could demonstrate preserved functional integrity that is not directly related to structural integrity in your model's generation of diverse protein sequences?  It seems an important question because when you want to design a protein to do some specific function (bind some small molecule or interact with another protein) you only care about structure to the extent that it acts as a proxy for function. But maybe it doesn't have to be? Do you think your models could get at function outside the restraint of the specific structure that is your input?
soundness: 3
strengths: The sampling speed of InvMSAFold is a lot faster than ESM-1F or ProteinMPNN, this is important when you want to generate millions of models, as I think could be reasonable for virtual screening/protein design applications. 

InvMSAFold seems able to sample more diverse regions of potential protein structure/function space than ESM-1F, again this is important when you are trying to select for particular properties (substrate specificity, thermostability).

That InvMSAFold is able to capture residue covariances in MSAs may also be useful for better backbone modeling that particular functions could then be engineered into.
summary: The paper presents InvMSAFold, an inverse folding model that generates the parameters of a probability distribution over the space of protein sequences with pairwise interwise interactions, allowing for efficient generation of diverse protein sequences while preserving structural and functional integrity. InvMSAFold is a neural network in which the inputs are the structure backbone coordinates X and the outputs are the parameters of a lightweight sequence model. The lightweight sequence model parameters are used to sample amino acid sequences compatible with the input structure. Training is based on the CATH database, which classifies protein domains into superfamilies and further into clusters based on sequence homology. The model is fast and has uses in protein design and virtual screening. Biologically, the model captures amino acid covariances observed in Multiple Sequence Alignments (MSA) of homologous proteins. The model expands the scope of inverse folding to retrieve a landscape of homologous proteins with similar folds (they say the 'entire' landscape, I don't think they have shown this).  I am overall very enthusiastic about this work.
weaknesses: There is not a specific example taken through to the conclusion that the model preserves "structural and functional integrity".  Functional integrity is what you want when you're designing new proteins/doing virtual screening. The authors should consider including such an example or clarifying this statement since that is a major claim of their paper. 

I was not clear on the InvMSAFold-AR/-PW. I understand that PW requires MCMC and AR does not but I wonder are there cases/tasks in which a PW vs AR model is more appropriate?

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: 1.	The symbols of Eq.(5) is not consistent with that of Eq.(3). It would be better to use consistent symbols.
2.	The proof in section 2.2.1 is incoherent. What is the function of Eq.(5)?
3.	In Eq.(7). It would be better to clarify that Eq.(7) is the L2 regularization term. 
4.	In section 3, it would be better to list the number of entries in each dataset.
5.	In section 4.1, what is the necessity of tuning the hyper-parameters of InvMSAFold-AR?
6.	It seems that InvMSAFold-PW performs better than InvMSAFold-AR at larger hamming distance. What is the probable cause?
soundness: 3
strengths: In their computational experiments, the authors demonstrated that the sequences generated by their models not only fold into the target structure but also exhibit greater diversity and more effectively capture the correlations between residues at different sites. Furthermore, the showed that this sequence diversity extends to other properties, such as predicted solubility and predicted and predicted thermostability. Overall, this paper represents a new methodological advancement.
summary: This paper proposes a neural network called InvMSAFold, which takes the protein structure as input, and outputs the parameters of two statistical models. These models are then used to generate a diverse set of protein sequences corresponding to the input structure. By utilizing these simple statistical models, the proposed pipeline effectively addresses two major challenges faced by other inverse-folding methods, such as ESM-IF: (1) the limited diversity of generated sequences and (2) slow sampling speed.
weaknesses: 1.	The authors only compare their method with ESM-IF1, and do not compare their method with other state-of-the-art inverse folding methods.
2.	In many places such as in section 1, "ESM-IF" was wrongly typed as "ESM-1F". This may lead readers to perceive the authors as lacking expertise.
3.	The article contains too many grammatical errors.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: I'd like to raise a few major points:
* It would strengthen the paper to benchmark against other methods as well. I would suggest for example a simple Potts model without the low-rank approximation and without the pre-trained ESM1-IF encoder, and an additional method such as ProteinMPNN beyond ESM1-IF. This would highlight the contributions of the paper more clearly, as currently, it may seem somewhat reliant on ESM1-IF.
* I highly recommend adding a plot that shows RMSD versus sequence recovery, as these metrics would provide valuable insights into the model’s performance.
* In Section 2.2.1, the explanation of Equation 7 and how it maintains linear scaling isn’t entirely obvious, at least to me. I suggest elaborating on this either within the main text or in the supplementary material to clarify the reasoning. It would be helpful to include 1-2 sentences explaining why the method or process is linear and how this linearity is established. This will provide clarity to the reader and strengthen the argument by highlighting the underlying reasoning behind the concept.
* To make the manuscript even stronger, it would be useful to include (1) an analysis of how the method scales with very large sequences or structures, and (2) a discussion of how the size of the MSA impacts model performance.

A minor point:
* In Section 2.2, I recommend including the formula that shows the normalization constant, as it is referenced in the text but not explicitly provided.

There are several typos throughout the manuscript that disrupt the flow. I have listed the ones I noticed while reading, but I recommend a re-read of the manuscript to specifically check for additional typos:
   - Line 17: The phrase “space of sequences with pairwise interwise interactions, capturing the amino acid…” contains the term “interwise,” which doesn’t seem correct or clear.
   - Lines 107-108: The word "Moreover" is used consecutively, which disrupts the flow.
   - Line 117: The word "Whos" should be corrected to "Whose."
   - Line 299: "We monitor the the negative..."—"the" is repeated.
   - Line 315: "A can be seen..." should likely be "As can be seen..."
soundness: 2
strengths: The authors present a novel and elegant approach to optimizing Potts model construction, training and sampling. The paper is well-structured, clearly outlining each crucial part of the methodology in a way that is easy to follow. The method is compared and benchmarked against a well-established approach, and performance metrics computed and reported.
summary: In this paper, the authors present an efficient method for designing protein backbones using a neural network that predicts a Potts model. The proposed architecture includes a pre-trained ESM1-IF encoder that encodes the protein backbone, generating rotation-invariant embeddings. These embeddings are then passed through a transformer-based decoder, which produces a low-rank matrix that is ultimately used to compute the fields and couplings. The low-rank approximation is a clever technique that helps mitigate the quadratic scaling cost typically associated with such computations. The neural network was trained using two distinct approaches: (1) a standard pseudo-likelihood loss, and (2) autoregressive sampling (over amino acids) with maximum likelihood training. To avoid training on single sequences, the model was trained on multiple sequence alignments (MSAs), with the mean pseudo-negative log-likelihood calculated over randomly sampled subsets of the MSA. Training and testing data were sourced from the CATH database, following its hierarchical classification to create test sets of varying difficulty, depending on the similarity between the training and test data. The authors demonstrate that their model better reconstructs covariance matrices compared to ESM1-IF, based on Pearson correlations. Moreover, the authors show that projected MSAs using PCA more closely reflect the natural sequence distributions, suggesting that their generated sequences, or predicted MSAs, are more diverse. When refolding designed sequences for test set structures, the InvMSAfold method proves to be more robust than ESM1-IF for sequences that deviate further from the native structure, and comparable to ESM1-IF for sequences that are highly similar to the native.
In conclusion, the paper demonstrates how a Potts model can be efficiently constructed, showing that the resulting model generates sequences that are plausible, diverse, refold successfully with AlphaFold2, and possess other promising biochemical attributes.
weaknesses: While the proposed methodology for improving the efficiency of Potts model construction is promising, there are a few areas where the paper could be strengthened. First, Potts models have long been used in fixed-backbone protein design, which makes it difficult to clearly identify the novelty and specific contributions of this work. Additionally, the method relies on components of ESM1-IF and then benchmarks against this model, which may limit the fairness or objectivity of the comparison. Another area for improvement is scalability. The paper does not provide any analysis on how the model handles large structures or long sequences, which could be useful for evaluating its broader applicability. Furthermore, there is no discussion on the significance of using MSAs for training versus single-sequence training, nor is there any exploration of how deep the MSAs need to be if they are indeed important.

Meta-review: The paper presents InvMSAFold, an inverse folding method designed to generate diverse protein sequences from a given structure.

The reviews identified strengths in the speed of the method and the greater diversity that more effectively capture the correlations between residues at different sites. 

Weaknesses include limited context including prior work and the fact that the authors compare only to ESM-IF1, and do not compare their method with other state-of-the-art inverse folding methods or to alphafold3.

The authors acknowledge the lack of comparisons and promise to add InvMSAFold-AR/PW with ProteinMPNN to a revised manuscript. Alphafold3 was not available at the scale needed to refold the sequences as suggested. The authors offered to correct typos.

Decision: Accept (Spotlight)