【阅读总结】Variant Effect Predictor: EVE 深度生成模型预测致病突变

省流：

本系列旨在整理Debora课题组一系列基于深度生成模型预测致病突变的工作，包括EVE、Tranception、TranceptEVE、EVEscape和popEVE，讨论数据来源与处理、模型架构与训练、性能测试与实例。

Model	Publish	Year	Available	简述
EVE	Nature	2021	https://github.com/OATML-Markslab/EVE https://evemodel.org/	VAE模型，输入MSA，输出突变的进化分数（野生型ELBO减突变型ELBO）
Tranception	ICML	2022	https://github.com/OATML-Markslab/Tranception
TranceptEVE	NIPS	2022		Tranception+EVE，weighted average
StructSeq	NIPS	2023		TranceptEVE+ESM-IF1
EVEscape	Nature	2023	https://github.com/OATML-Markslab/EVEscape	EVE+Biophysical
popEVE	medRxiv	2024	https://github.com/debbiemarkslab/popEVE https://pop.evemodel.org/	EVE+ESM-1v

EVE

数据

训练数据

3,851 genes in ClinVar (the April 2021 release) 存在有一个错义突变被标记为Pathogenic：

→ 每个基因对应到Uniprot/Swissprot中记录的单一蛋白

→ 每条蛋白序列通过jackhmmer在UniRef100 (April 20th 2020) 中搜索五次，构建不同MSA

→ MSA的构建与选择的原则：

1. retrieve sequences that align to at least 50% of the target protein sequence, and columns with at least 70% residue occupancy.

2. using 0.3 bits per residue as a reference; prioritize coverage $L_{cov}$ ≥ 0.8$L$ & 100, 000≥$N$≥10$L$; 达不到的话尝试$L_{cov}$≥0.7$L$ & $N$≤200, 000; 依然达不到的剩余3%不做限制。

   $L$: length of the target protein sequence; $N$: total number of sequences

   未发现模型表现与alignment depth $N/L_{cov}$之间存在强相关。

最终获得3,219 clinically relevant proteins with corresponding evolutionary training data

*MSA构建集成在Debora组之前的EVcouplings软件中，GitHub或者网站可用（2024.08.14查看网站，处于维护状态暂不可用）

模型

以下代码分析来自https://github.com/OATML-Markslab/EVE，ProteinGym-proteingym-baselines-EVE中代码基本一致。

EVE使用的模型架构为Variational Autoencoder (VAE)。简单来说，VAE学习样本高斯分布的均值和方差（先验分布），采样获得$p(x|z)$（$z$：潜在变量）。VAE的优化目标为 $\log P_{\theta }(x)$ 的下界Evidence Lower Bound (ELBO)，推导过程可见https://zhuanlan.zhihu.com/p/678864844，ELBO由重构损失（Reconstruction Loss）和KL散度（KL Divergence）组成：

$$\mathrm{ELBO}(\theta ,\varphi )={\mathbb{E}}_{z\sim q_{\varphi }(z\mid x)}\left[\log p_{\theta }(x\mid z)\right]-KL\left[q_{\varphi }(z\mid x)||p(z)\right]$$

Encoder/Decoder

具体代码实现中，EVE采用对称尺寸的encoder（2, 000 - 1, 000 - 300）和 decoder（300 - 1, 000 - 2, 000），用nn.init.constant_()函数初始化所有线性层bias为0.1。

对于encoder，经过三层线性层后分别映射到大小为50的潜在空间，获得z_mean和z_log_var。通过sampler函数从z_mean和z_log_var采样获得z，z输入decoder。

def sampler(self, mean, log_var):
    # torch.randn_like：创建与输入相同形状的新张量，该张量填充了均值为0、方差为1的正态分布随机数。
    eps = torch.randn_like(mean).to(self.device)
    z = torch.exp(0.5*log_var) * eps + mean
    return z

对于decoder，区别于standard VAE decoder的三层线性层相连，采用Bayesian Neural Network：对于第i层，用sampler函数来reparametrize对x进行线性操作的layer_i_weight和layer_i_bias。

经过三层处理后，再一次采样获得W_out和b_out用于最后一次线性操作。W_out做一次加权，权重为使用one-dimensional convolution捕获的对20种氨基酸的偏好性：依然用sampler函数，mean和log_var来自两个nn.Conv1d(40, 20, kernel_size=1,stride=1,bias=False).weights。

# 用加权W_out和b_out对x进行最后一次线性操作
x = F.linear(x, weight=W_out, bias=b_out)

# 温度缩放，mean初始化为1，log_var初始化为-10
if self.include_temperature_scaler:
    temperature_scaler = self.sampler(self.temperature_scaler_mean,self.temperature_scaler_log_var)
    x = torch.log(1.0+torch.exp(temperature_scaler)) * x

x = x.view(batch_size, self.seq_len, self.alphabet_size)
# 得到重构的x
x_recon_log = F.log_softmax(x, dim=-1) 

代码详见./EVE/VAE_encoder.py和./EVE/VAE_model.py，超参数设置详见./EVE/default_model_params.json。

*include_sparsity：DeepSequence (Nature Methods, 2018, Fig.5) 中假设蛋白质序列存在上位性，一个group内不同氨基酸与hidden units之间的权重，在不同groups中共享。EVE发现这一操作会使得模型在一些蛋白上表现显著下降，因此不再使用。

Loss function

损失函数为负ELBO：

def loss_function(self, x_recon_log, x, mu, log_var, kl_latent_scale, kl_global_params_scale, annealing_warm_up, training_step, Neff):
    BCE = F.binary_cross_entropy_with_logits(x_recon_log, x, reduction='sum') / x.shape[0]
    
    KLD_latent = (-0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())) / x.shape[0]
    if self.decoder.bayesian_decoder:
        # 对于Bayesian decoder，计算weights的KL散度
        # Neff：有效样本数，根据MSA的汉明距离进行加权
        KLD_decoder_params_normalized = self.KLD_global_parameters() / Neff
    else:
        KLD_decoder_params_normalized = 0.0
    # 默认参数中annealing_warm_up取0，warm_up_scale为1
    warm_up_scale = self.annealing_factor(annealing_warm_up,training_step)
    
    neg_ELBO = BCE + warm_up_scale * (kl_latent_scale * KLD_latent + kl_global_params_scale * KLD_decoder_params_normalized)
    
    return neg_ELBO, BCE, KLD_latent, KLD_decoder_params_normalized

BCE（Binary Cross Entropy）是二元交叉熵损失，即ELBO中的Reconstruction Loss

KLD_latent 是标准的VAE中的KL散度，计算潜在变量的后验分布$q(z\mid x)$和先验分布$p(z)$之间的散度。KL散度公式：

$$KLD=-\frac{1}{2}\sum _{i=1}^d\left(1+\log \left({\sigma }_i^2\right)-{\mu }_i^2-{\sigma }_i^2\right)$$

kl_latent_scale 和 kl_global_params_scale 用于分别控制潜在变量和解码器参数的KL散度对总损失的贡献，默认参数中两项均设置为1。

Prediction

一个蛋白质变体的$E_s$（evolutionary index）定义为野生型与突变型的ELBO之差:
$$E_s=-\log \frac{p\left(\mathbf{s}\mid {\theta }_{\mathbf{p}}\right)}{p\left(\mathbf{w}\mid {\theta }_{\mathbf{p}}\right)}\sim ELBO(\mathbf{w})-ELBO(\mathbf{s})$$

注意此处ELBO计算与损失函数不同，不引入scale系数。

def all_likelihood_components(self, x):
    mu, log_var = self.encoder(x)
    z = self.sample_latent(mu, log_var)
    recon_x_log = self.decoder(z)

    recon_x_log = recon_x_log.view(-1,self.alphabet_size*self.seq_len)
    x = x.view(-1,self.alphabet_size*self.seq_len)
    
    BCE_batch_tensor = torch.sum(F.binary_cross_entropy_with_logits(recon_x_log, x, reduction='none'),dim=1)
    KLD_batch_tensor = (-0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp(),dim=1))
    
    ELBO_batch_tensor = -(BCE_batch_tensor + KLD_batch_tensor)
    
    return ELBO_batch_tensor, BCE_batch_tensor, KLD_batch_tensor

对于每个变体，采样20k，取多个ELBO_batch_tensor的平均值，然后与野生型相减计算$E_s$

Gaussian Mixture Model (GMM)

观察到evolutionary index对致病标签有较强的相关性，因此，尝试用a Gaussian Mixture model with two components来拟合evolutionary index的分布，区分致病突变和良性突变。

GMM是一种概率模型，用来表示数据集中的多种分布。它假设数据由若干个高斯分布（也称为分量）混合而成。每个高斯分布对应于数据集中的一个子群体。GMM with two-component可以理解为这个模型由两个高斯分布的加权平均构成。每个高斯分布由mean和covariance matrix描述，并且每个分量都有一个权重表示该分量在整个混合分布中所占的比例。

对于给定的蛋白变体$\mathbf{s}$：

$$p\left({\mathbf{X}}_{\mathbf{s}}=1\mid {\mathbf{E}}_{\mathbf{s}}\right)=\alpha \ast p\left({\mathbf{X}}_{\mathbf{s}}=1\mid {\mathbf{E}}_{\mathbf{s}},{\theta }_p\right)+(1-\alpha )\ast p\left({\mathbf{X}}_{\mathbf{s}}=1\mid {\mathbf{E}}_{\mathbf{s}},{\theta }_o\right)$$

$X_s$代表是否致病（1 if pathogenic，0 if benign)；$\alpha$ 代表权重，grid search确定为0.3；${\theta }_o$、${\theta }_p$ 分别是overarching GMM and proteinspecific GMM的协方差矩阵；$p\left({\mathbf{X}}_{\mathbf{s}}=1\mid {\mathbf{E}}_{\mathbf{s}}\right)$ 即EVE score，量化给定$\mathbf{s}$是否致病。

可以借助sklearn来实现：

from sklearn import mixture

main_GMM = mixture.GaussianMixture(n_components=2, covariance_type='full',max_iter=1000,n_init=30,tol=1e-4)
main_GMM.fit(X_train)

不确定性用predictive entropy (PE) 来衡量：

$$PE=-\log p\left({\mathbf{X}}_{\mathbf{s}}=1\mid {\mathbf{E}}_{\mathbf{s}}\right)p\left({\mathbf{X}}_{\mathbf{s}}=1\mid {\mathbf{E}}_{\mathbf{s}}\right)-\log p\left({\mathbf{X}}_{\mathbf{s}}=0\mid {\mathbf{E}}_{\mathbf{s}}\right)p\left({\mathbf{X}}_{\mathbf{s}}=0\mid {\mathbf{E}}_{\mathbf{s}}\right)$$

本工作将PE最高的 25% 的变体标记为不确定，准确度可达 ∼ 90%。在实践中，具体基因的阈值可能根据precision/recall需求做调整。本部分代码见train_GMM_and_compute_EVE_scores.py。

性能

1. Fig.2：在3,219 proteins上评估致病性预测，以及至少3、5和10个良性标签和至少3、5和10个致病标签的蛋白子集。

   *用于比较的23种模型结果来自dbNSFP，排除了使用人口频率数据来训练的模型，因为该数据常常用于变体分类；有监督训练的模型无法获知具体的训练集/验证集划分，因此存在偏差。

2. case study：与TP53、PTEN19、MSH2、BRCA1四种蛋白的deep mutational scans (DMS) 数据计算AUC