Importance Weighted Expectation-Maximization for Protein Sequence Design

The protein sequence design problem is to search for a sequence with desired property in the sequence space $\mathcal{V}^{L}$, where $\mathcal{V}$ denotes the vocabulary of amino acids and L denotes the desired sequence length. The target is to find a protein sequence with highest fitness given by a protein fitness function $f:\mathcal{V}^{L}\rightarrow\mathbb{R}$, which can be measured through wet-lab experiments. Wild-type refers to protein occurring in nature. Evolutionary search based methods are widely used (Bloom & Arnold, 2009; Arnold, 2018; Angermüller et al., 2020; Ren et al., 2022). They use wild-type sequence as starting point during iterative search. In this paper, we do not focus on the modification upon the wild-type sequence, but aim to efficiently generate novel and diverse sequences with improved protein functional properties.

We will develop our method based on Monte Carlo expectation-maximization (MCEM) (Bishop, 2006). A latent generative model assumes data x (e.g. a protein sequence) is generated from a latent variable z. To learn this latent model, the optimization procedure for maximizing the log marginal likelihood is to alternate between expectation step (E-step) and maximization step (M-step). EM directly targets the log marginal likelihood of an observation x by involving a variational distribution $q_{\phi}(z)$:

where $p_{\theta}(z|x)$ is the true posterior distribution and $p_{\theta}(x,z)=p_{\theta}(x|z)p_{\theta}(z)$ is the joint distribution, composed of the conditional likelihood $p_{\theta}(x|z)$ and the prior $p_{\theta}(z)$. In MCEM, E-step samples a set of z from $q_{\phi}(z)$ to estimate the expectation using Monte Carlo method, and then M-step fits the model parameters $\theta$ by maximizing the Monte Carlo estimation (Wei & Tanner, 1990). It can be proved that this process will never decease the log marginal likelihood (Bishop, 2006).

Our goal is to search over a space of discrete protein sequences $\mathcal{V}^{L}$ – $\mathcal{V}$ consists of 20 amino acids and $L$ is the sequence length – for sequence $x\in\mathcal{V}^{L}$ that maximizes a given fitness function $f:\mathcal{V}^{L}\rightarrow\mathbb{R}$. Let the fitness value $y=f(x)$, given a predefined threshold $\lambda$, we define a conditional likelihood function:

where $\mathcal{S}$ represents the event that the fitness of x is ideal ($y\geq\lambda$). We also assume a class of generative models $P_{\theta}(x)$ that can be trained to model the raw protein sequences and it will be kept fixed afterwards. Since the search space is exponentially large (O($20^{L}$)), random search would be time-intensive. Following Brookes et al. (2019), we formulate the protein design problem as generating satisfactory sequences from the posterior distribution $P_{\theta}(x|\mathcal{S})$:

where $P_{\theta}(\mathcal{S})=\int_{x}P_{\theta}(x)P(\mathcal{S}|x)dx$ is a normalization constant which does not rely on x. Protein sequences generated from $P_{\theta}(x|\mathcal{S})$ are not only more likely to be real proteins, but also have higher functional scores (i.e., fitness). The higher the $\lambda$ is, the higher fitness the discovered protein sequences have.

Directly generating satisfactory sequences from the posterior distribution $P_{\theta}(x|\mathcal{S})$ is highly efficient compared with randomly search over the exponentially discrete space. However, realizing this idea is difficult as computing $P_{\theta}(\mathcal{S})=\int_{x}P(\mathcal{S}|x)P_{\theta}(x)dx$ needs an integration over all possible x, which is intractable. Instead, we propose to learn a variational distribution $Q_{\phi}(x)$ with learnable parameter $\phi$ to approximate satisfactory proteins $P_{\theta}(x|\mathcal{S})$. Following Brookes et al. (2019), to find the optimal $\phi$ of the proposal distribution, we choose to minimize the KL divergence between the posterior distribution $P_{\theta}(x|\mathcal{S})$ and the variational distribution $Q_{\phi}(x)$:

where $\mathcal{H(P_{\theta})}=-E_{P_{\theta}(x|\mathcal{S})}\log P_{\theta}(x|\mathcal{S})$ is the entropy of $P_{\theta}(x|\mathcal{S})$ and can be dropped because it does not matter $\phi$.

Diversity is a key consideration in our protein design procedure, which not only satisfies the diverse nature of species, but also can reduce undesired inter-domain misfolding (Wright et al., 2005). In order to promote the diversity of the designed protein sequences, we introduce a latent variable z into our model to capture the high-order dependencies among amino acids in protein sequences. We assume the joint approximate distribution $Q_{\phi}(x,z)=Q_{0}(z)Q_{\phi}(x|z)$. Thus our final goal is to maximize the expected log-likelihood $\log Q_{\phi}(x)$ of the satisfactory proteins with respect to the ideal posterior distribution $P_{\theta}(x|\mathcal{S})$. By using the EM objective from Eq.(1), we derive

where $R_{\psi}(z|x)$ is another variational distribution with its parameter $\psi$ to approximate the intractable posterior $Q_{\phi}(z|x)$, and $\mathcal{F}(R_{\psi}(z|x),\phi)=E_{R_{\psi}(z|x)}[\log Q_{\phi}(x,z)-\log R_{\psi}(z|x)]+D_{KL}(R_{\psi}(z|x)||Q_{\phi}(z|x))$. Here $Q_{0}(z)$ is implemented as a standard normal distribution, and $Q_{\phi}(x|z)$ is a Transformer decoder (Vaswani et al., 2017) with $z$ as the first embedding input, which will be augmented by the combinatorial structure features (Sec. 3.4). We implement $R_{\psi}(z|x)$ as a normal distribution $\mathcal{N}(\mu_{\psi},\sigma^{2}_{\psi}\mathbb{I})$ with $\mu_{\psi},\sigma_{\psi}=\mathrm{Transformer}_{\psi}(x)$. The overall architecture is illustrated in Figure 3.

To maximize the objective defined in Eq.(5), we plan to learn the proposal distribution $Q_{\phi}(x,z)$ through Monte Carlo EM, because the sampling procedure and iterative optimization process can lead to a better estimate (Figure 1 (b)), resulting in novel and diverse proteins with higher fitness.

Since Eq.(5) involves the expectation over the intractable distribution $P_{\theta}(x|\mathcal{S})$ and the KL divergence of the intractable posterior distribution $Q_{\phi}(z|x)$, we use importance sampling to approximate its variational lower bound. We assume the original protein sequence is also generated from a same latent variable $z$, i.e. $P_{\theta}(x,z)=P_{\theta}(x|z)P_{0}(z)$, where $P_{0}(z)$ is a standard normal distribution and $P_{\theta}(x|z)$ is a Transformer with $z$ as the initial input embedding. We also assume the latent variable z in $P_{\theta}(x,z)$ and $Q_{\phi}(x,z)$ are defined on the same latent space. In practice, we learn $P_{\theta}(x|z)$ as the decoder in a pre-trained variational auto-encoder on raw protein sequences. We assume $\mathcal{S}$ is only conditioned on x, leading to $P_{\theta}(x,z|\mathcal{S})=\frac{P_{\theta}(x,z)P(\mathcal{S}|x)}{P_{\theta}(\mathcal{S})}$. Then the final training objective can be derived as:

Here we use the importance sampling fundamental identity to derive the equation (Robert & Casella, 2004). $\mathcal{C}=P_{\theta}(\mathcal{S})$ is a constant which does not rely on $\phi$ and $\psi$. The last inequality adopts the evidence lower bound (ELBO) (Bishop, 2006). We use importance sampling based EM to approximate the above objective (Robert & Casella, 2004) with joint samples $(x_{n},z_{n})\sim Q_{\phi}(x,z)$. Specifically, at each iteration, we perform,
E-step:

where $w(x_{n},z_{n})=\frac{P_{\theta}(x_{n}|z_{n})}{Q_{\phi^{(t)}}(x_{n}|z_{n})}P(\mathcal{S}|x_{n})$ is the unnormalized importance weight, $w^{\prime}(x_{n},z_{n})=\frac{w(x_{n},z_{n})}{\sum_{n=1}^{N}w(x_{n},z_{n})}$ is the normalized one, and $N$ is the sample size. Since $R_{\psi}(z|x)$ is assumed to be a normal distribution with its mean $\mu_{\psi}$ and variance $\sigma_{\psi}^{2}\mathbb{I}$ calculated from a Transformer encoder $\mu_{\psi},\sigma_{\psi}=\mathrm{Transformer}_{\psi}(x_{n})$, the KL term can be calculated in closed form (Eq. (17)).
M-step:

In E-step, we use two techniques to generate samples for $z$ and $x$ – using the real protein sequences $x$ with generated $z$ from $R_{\psi}(z|x)$ and using $z$ from a standard normal distribution with generated proteins from $Q_{\phi}(x|z)$. In M-step, we optimize the KL regularized data generation log likelihood with both high-fitness real and synthetic proteins. The procedure can be viewed as self-training with augmented synthetic protein data, which differs from prior approaches such as CbAS (Brookes et al., 2019).

As shown in previous work, the combinatorial structure of amino acids in protein sequences can be learned from a generative graphical model Markov random fields (MRFs) fitted on the sequences from the same family (Hopf et al., 2017; Luo et al., 2021). These structure constraints are the results of the evolutionary process under natural selection and may reveal clues on which amino-acid combinations are more favorable than others. Thus we incorporate these features into our model to guide it towards higher fitness landscape to faster find desired protein sequences.

Given a protein sequence $x=(x_{1},x_{2},..,x_{M})$ with $M$ amino acids, the generative model generates it with likelihood $P_{\epsilon}(x)=\frac{\exp(\Upsilon(x))}{Z}$ where $Z=\int_{x}\exp(\Upsilon(x))dx$ is a normalization constant and $\Upsilon(x)$ is the corresponding energy function, which is defined as the sum of all pairwise constraints and single-site constraints as follows:

where $\varepsilon_{i}(x_{i})$ denotes the single-site constraint of $x_{i}$ at position i and $\varepsilon_{ij}(x_{i},x_{j})$ denotes the pairwise constraint of $x_{i}$ and $x_{j}$ at position i, j. The above graphical model is illustrated in the upper half of Figure 3.

We train the model on protein sequences from the same family following CCMpred (Seemayer et al., 2014) using a pseudo-likelihood $\hat{P}_{\epsilon}(x)$ (provided in Appendix E) combined with $L_{2}$ regularization to make the learning of $P_{\epsilon}(x)$ easier. But different from them, we additionally add $L_{1}$ regularization to the training objective to make the graph sparse, of which the regularization coefficients are set to the same values as the $L_{2}$ regularization:

where $x$ denotes protein sequences from the same family, $\varepsilon_{i}=[\varepsilon_{i}(a_{1}),\varepsilon_{i}(a_{2}),...,\varepsilon_{i}(a_{20})]$ is the vector of the single-site constraints of the 20 amino acids at position i, and $\varepsilon_{ij}=[\varepsilon_{ij}(a_{1},a_{2}),\varepsilon_{ij}(a_{1},a_{3}),...,\varepsilon_{ij}(a_{M},a_{M-1})]$ is the vector of all possible pairwise constraints at position i, j.

After training the MRFs, we can encode a protein sequence $x$ with the learned constraints. Specifically, we first encode the i-th amino acid by concatenating its corresponding single-site constraint as well as the possible pairwise ones:

where $\varepsilon_{ij}(x_{i},a_{j\boldsymbol{\cdot}})$ gathers the 20 amino acids for any position $j\neq i$. Then we map $\boldsymbol{\varepsilon_{i}}(x_{i})$ to the amino-acid embedding space with trainable parameter $W_{\varepsilon}$, and add the mapped vector to the original amino-acid embedding $e(x_{i})$ to get the final feature vector as our model input:

where $H_{i}$ is the input embedding of the Transformer decoder, i.e., the first input is set to the sampled latent vector $z^{\prime}$ and the input for other position $i$ is set to the combinatorial structure augmented feature vector $\hat{e}(x_{i-1})$.

Combining Eq. (7) and (12), the learning process becomes:
E-step:

where $Q_{\phi}(x_{n}|z^{\prime}_{n};\boldsymbol{\varepsilon})$ is computed from $\mathrm{Transformer}(H_{0},H_{1:M})$. In practice, we also learn $P_{\theta}(x,z;\boldsymbol{\varepsilon})=P_{0}(z)P_{\theta}(x|z;\boldsymbol{\varepsilon})$ as a combinatorial structure feature enhanced latent generative model, where $P_{\theta}(x|z;\boldsymbol{\varepsilon})$ is a combinatorial structure feature enhanced Transformer with a learnable mapping matrix $W_{\eta}$ to transform the combinatorial structure features into the embedding space. $d$ is the dimensionality of $\sigma_{\psi}(x_{n})$ and $\sigma_{\psi,i}(x_{n})$ is its $i$-th dimension.
M-step:

The overall learning algorithm is given in Appendix B.4.

We first train a VAE model on raw protein sequences using a 6-layer Transformer as the encoder and a 2-layer Transformer as the decoder. We add the combinatorial structure features to the Transformer decoder input as described in Sec. 3.4 and its weight $W_{\eta}$ is learned jointly. The raw protein probability $P_{\theta}(x|z;\boldsymbol{\varepsilon})$ is then defined by the combinatorial structure feature augmented decoder in this VAE. The protein combinatorial structure constraints $\boldsymbol{\varepsilon}$ are learned on the training sequences for each dataset, rather than using real multiple sequence alignments (MSAs), to ensure a fair comparison. The approximate posterior probability $\tilde{P}_{\theta}(z|x)$ is defined by its encoder. Its embedding and feed-forward network sizes are 320 and 1280. The encoder is initialized with the pre-trained ESM-2 weights (Lin et al., 2023)¹¹1https://dl.fbaipublicfiles.com/fairesm/models/esm2_t6_8M_UR50D.pt. The latent variable $z$’s dimension is also 320. We use another 6-layer Transformer encoder followed by a linear mapping to learn $\mu_{\phi}$ and $\sigma_{\phi}$ for $R_{\psi}(z|x)=\mathcal{N}(\mu_{\psi},\sigma^{2}_{\psi}\mathbb{I})$ and another 2-layer Transformer decoder to learn $Q_{\phi}(x|z;\boldsymbol{\varepsilon})$. We initialize parameters $\psi^{(0)}$ and $\phi^{(0)}$ with $\theta$.

The number of iterations in the importance sampling-based MCEM is set to 10. The mini-batch size and learning rate are set to $4,096$ tokens and $1$e-$5$ respectively. The model is trained with $1$ NVIDIA RTX A$6000$ GPU card. We apply Adam algorithm (Kingma & Ba, 2015) as the optimizer with a linear warm-up over the first $4,000$ steps and linear decay for later steps. We randomly split each dataset into training/validation sets with the ratio of $9$:$1$. We run all the experiments for five times and report the average scores. More experimental settings are given in Appendix B.1. Following (Kamisetty et al., 2013), we set $\alpha_{\text{single}}=1$ and $\alpha_{\text{pair}}=0.2*(M-1)$ with $M$ equals to the protein sequence length.

In inference, we design protein sequences by taking the wild-type as encoder input and the latent vector is sampled from prior distribution $Q_{0}(z)=N(0,\mathbb{I})$. The sequences are decoded using sampling strategy with top-$5$. The candidate number is set to K=$128$ following the setting of Jain et al. (2022) on the GFP dataset.

Following Ren et al. (2022), we evaluate our method on the following eight protein engineering benchmarks:
(1) Green Fluorescent Protein (avGFP): The goal is to design sequences with higher log-fluorescence intensity values. We collect data following Sarkisyan et al. (2016). (2) Adeno-Associated Viruses (AAV): The target is to generate amino-acid segment (position $561-588$) for higher gene therapeutic efficiency. We collect data following Bryant et al. (2021). (3) TEM-1 $\beta$-Lactamase (TEM): The goal is to design high thermodynamic-stable sequences. We merge the data from Firnberg et al. (2014). (4) Ubiquitination Factor Ube4b (E4B): The objective is to design sequences with higher enzyme activity. We gather data following Starita et al. (2013). (5) Aliphatic Amide Hydrolase (AMIE): The goal is to produce amidase sequences with higher enzyme activity. We merge data following Wrenbeck et al. (2017). (6) Levoglucosan Kinase (LGK): The target is to optimize LGK protein sequences with improved enzyme activity. We collect data following Klesmith et al. (2015). (7) Poly(A)-binding Protein (Pab1): The goal is to design sequences with higher binding fitness to multiple adenosine monophosphates. We gather data following Melamed et al. (2013). (8) SUMO E2 Conjugase (UBE2I): We aim to find human SUMO E2 conjugase with higher growth rescue rate. Data are obtained following Weile et al. (2017). The detailed data statistics, including protein sequence length, data size and data source are provided in Appendix A.

We compare our method against the following representative baselines: (1) CMA-ES (Hansen & Ostermeier, 2001) is a famous evolutionary search algorithm. (2) FBGAN proposed by Gupta & Zou (2019) is a novel feedback-loop architecture with generative model GAN. (3) DbAS (Brookes & Listgarten, 2018) is a probabilistic modeling framework and uses adaptive sampling algorithm. (4) CbAS (Brookes et al., 2019) improves on DbAS by conditioning on the desired properties. (5) PEX proposed by Ren et al. (2022) is a model-guided sequence design algorithm using proximal exploration. (6) GFlowNet-AL (Jain et al., 2022) applies GFlowNet to design biological sequences. We use the implementations of CMA-ES, DbAS and CbAS provided in Trabucco et al. (2022) and for other baselines, we apply their released codes. To better analyze the influence of different components in our model, we also conduct ablation tests as follows: (1) IsEM-Pro-w/o-ESM removes ESM-2 as encoder initialization. (2) IsEM-Pro-w/o-ISEM removes iterative optimization process. (3) IsEM-Pro-w/o-MRFs removes MRFs features and iterative optimization process. (4) IsEM-Pro-w/o-LV removes latent variable, MRFs features and iterative optimization process. (5) ESM-Search samples sequences from the softmax distribution obtained by finetuning ESM-2 on the protein datasets and taking the wild-type as input.

We use three automatic metrics to evaluate the performance of the designed sequences: (1) MFS: Maximum fitness score. The oracle model adopted to evaluate MFS is described in Appendix B.1; (2) Diversity proposed by Jain et al. (2022) is used to evaluate how different the designed candidates are from each other; (3) Novelty proposed by Jain et al. (2022) is used to evaluate how different the proposed candidates are from the sequences in training data.

Table 1, 2 and 3 respectively report the maximum fitness scores, diversity scores and novelty scores of all models.

IsEM-Pro achieves the highest fitness scores on all protein families and outperforms the previous best method GFlowNet-AL by 55% on average (Table 1). The reasons are two-folds. On one hand, the importance sampling based MCEM can help our model to navigate to a better region instead of getting trapped in a worse local optima. On the other hand, the combinatorial structure features help to recognize the preferred mutation patterns which have higher success rate under the nature selection pressure, potentially leading to sequences with higher fitness scores.

IsEM-Pro achieves the highest average diversity score over the eight tasks (Table 2). Our model gains the highest diversity on $4$ out of $8$ tasks while GFlowNet-AL gains $2$ and CMA-ES gains $1$. It indicates that though involving combinatorial structure constraints can give the guidance for preferred protein patterns, it might also limit the sequence design to these patterns to some extent. The involved latent variable can capture complex inter-dependencies among amino acids, which benefits for more diverse protein design.

IsEM-Pro can design more novel protein sequences on all datasets (Table 3). Our model achieves higher novelty scores on all datasets due to the reason that more new samples are involved during the importance sampling based iterative optimization process, which is beneficial for more novel protein design.

Bottom halves of Table 1, 2 and 3 report the results of ablation tests. IsEM-Pro-w/o-MRFs improves the average diversity and novelty scores by as much as $33$x compared with IsEM-Pro-w/o-LV, which demonstrates that introducing a latent variable can significantly help to generate diverse proteins. IsEM-Pro-w/o-MRFs achieves higher maximum fitness scores than IsEM-Pro-w/o-ESM on all datasets, validating that adopting a pretrained protein language model as the encoder helps to design more satisfactory protein sequences. However, directly finetuning ESM-2 to sample candidates (ESM-Search) drops $1.3$ points on average fitness score compared with taking ESM-2 as an encoder (IsEM-Pro-w/o-MRFs), demonstrating that ESM-2 is not suitable for direct sequence design. Incorporating combinatorial structure features can further improve the fitness of the designed proteins (IsEM-Pro-w/o-ISEM V.S. IsEM-Pro-w/o-MRFs), based on which learning the proposal distribution by importance sampling based MCEM can better promote more desirable, diverse and novel protein generation. We also provide the model performance with a fully-connected decoder (non-autoregressive decoder) instead of the current autoregressive one in Appendix F.2. It shows the non-autoregressive decoder can not generate good candidates for protein with longer sequences such as LGK ($439$ amino acids).

To validate how close the proposal distribution $Q_{\phi}(x)$ is to the posterior distribution $P_{\theta}(x|\mathcal{S})$, we calculate the KL divergence through Monte Carlo approximation. Here we calculate the KL divergence between $Q_{\phi}(x)$ and $P_{\theta}(x|\mathcal{S})$ as we have proven in Lemma $C.1$ (provided in Appendix C.2) that the sampling difference between these two distributions can be bounded under this divergence. We leverage an unbiased and low-variance estimator (proof shown in Appendix C.1) to approximate the KL divergence as follows:

where $r(x)=\frac{P_{\theta}(x|\mathcal{S})}{Q_{\phi}(x)}$. The approximate KL divergence on eight protein datasets over $1$k-$10$k samples are illustrated in Figure 4. From the figure, we can see that the variance of KL divergence is very small over different sample size for all datasets. Besides, the KL divergence finally arrives at a small value, such as $0.018$ for E4B and $0.033$ for avGFP. It gives empirical evidence that when we sample from the ultimate $Q_{\phi}(x)$, it has minor difference compared with sampling from the posterior distribution $P_{\theta}(x|\mathcal{S})$.

Next, we study the effect of different implementation schemes of involving a latent variable with a pretrained encoder. Some works have tried adding the latent representation to the original embedding layer (Latent-Add) or using it as an additional memory (Latent-Memory) when adopting a pretrained language model as encoder (Li et al., 2020). We also implement our model with these two schemes, and evaluate the model performance on avGFP dataset. Table 4 shows that our method, which takes the latent representation as the first token input of decoder, achieves a higher fitness and novelty scores though with a mild decrease on diversity. We also provide the conditional VAE (CVAE) performance in Appendix F.1, which shows decreased performance. It demonstrates that iterative sampling inside the EM algorithm with explicit fitness constraints can lead to better proteins.

To gain an insight on how well the designed proteins are, we analyze the generated avGFP sequence with highest fitness in detail using Phyre2 tool (Kelley et al., 2015). Figure 5 (a) illustrates the generated variant can fold stably. According to the software, the most similar protein is Cytochrome b562 integral fusion with enhanced green fluorescent protein (EGFP) (Edwards et al., 2008). There are 227 residues (96% of the candidate sequence) have been modeled with 100.0% confidence by using this protein as template. Details are given in Appendix D. Figure 5(b) visualizes the superposition of the top-5 most similar templates to our sequence in the protein data bank, which are all fluorescent proteins and show highly consistent structure in most regions, validating that our model can design a real fluorescent protein.