BindCraft算法解析和测试

发表于 2024-10-23 | 更新于 2024-12-09

antibody design | ai | protein design | alphafold

这篇文章pipeline的核心工具是colabdesign（Sergey Ovchinnikov团队研发，故而Sergey作为BindCraft的共同通讯），利用colabdesign提供的afdesign模块和proteinMPNN模块，进行binder design。

原文链接：https://www.biorxiv.org/content/10.1101/2024.09.30.615802

代码链接：https://github.com/martinpacesa/BindCraft

算法部署

算法是基于colabdesign开发的，需要部署colabdesign环境，下载colabfold提供的alphafold multimer权重（可以各个weight都下载，反复用）。

基于代码的算法解析

算法主过程采用四阶段优化，同时可以额外用proteinMPNN（soluble版）进行binder性质改善。然后生成的结果进行综合的打分评估。算法的核心在于下列自定义hallucination loss组合，其中给contact loss（intra+inter）最高的权重，从而诱导生成具有一定靶向能力和二级结构的binder。

Binder confidence pLDDT (weight 0.1)
Interface confidence i_pTM (weight 0.05)
Normalized predicted alignment error (pAE) within the binder (weight 0.4)
Normalized predicted alignment error (pAE) between binder and target (weight 0.1)
Residue contact loss within binder (weight 1.0)
Residue contact loss between the target and binder - if hotspots are specified, the rest of the target is masked from this loss (weight 1.0)
Radius of gyration of binder (weight 0.3)
“Helicity loss” - penalize or promote backbone contacts every in a 3 residue offset to promote the hallucination of helical or non-helical designs (weight -0.3)

与训练时BP过程相似，不同的是，训练过程利用权重的梯度信息更新模型权重，而这里数据的梯度信息更新序列embedding logits。同时多批数据训练类似，这里用上五个权重以减少设计偏差。因此序列初始化也同样会影响优化结果。

class mk_af_model(design_model, _af_inputs, _af_loss, _af_prep, _af_design, _af_utils):
  def __init__(self,
               protocol="fixbb", 
               use_multimer=False,
               use_templates=False,
               debug=False,
               data_dir=".", 
               **kwargs):

colabdesign的模型对象是通过多继承的实现：

design_model实现了一个通用的处理逻辑（不止用在afdesign），
_af_design实现了基于alphafold的design过程，
_af_inputs实现对输入结构的处理，
_af_prep实现对binder序列的初始化，
_af_loss接入了loss函数实现

本身这个实现架构还是可以的，但colabdesign的代码问题在于，上述每个父类模块之间会存在相互的方法调用，而这些父类本身没有继承关系。这就导致在代码阅读时经常找不到某个方法。同时里面有大量函数式编程的痕迹（如闭包），从而使得代码阅读和修改存在较大挑战。

初始序列logits生成

1 2	print("Stage 1: Test Logits") af_model.design_logits(iters=50, e_soft=0.9, models=design_models, num_models=1, sample_models=advanced_settings["sample_models"], save_best=True)

这一阶段是初始序列logits迭代优化的一个阶段，默认设置迭代50次。其核心是利用反向传播算法进行序列logits的梯度下降优化，这里的grad和params的shape为（1，L，20）。

论文中对此部分的序列logits的公式描述如下，即分为未曾归一化的logits和softmax归一化后的logit的一个线性组合:

$logits=(1-\lambda)\times logits + \lambda\times softmax(\frac{logits}{temperature})$ $\lambda=\frac{(step+1)}{iterations}$

# colabdesign.shared.model.soft_seq
# 相关代码部分内容
seq["soft"] = jax.nn.softmax(seq["logits"] / opt["temp"])
seq["pseudo"] = opt["soft"] * seq["soft"] + (1-opt["soft"]) * seq["input"]

也就是说随着迭代过程，logits会更加偏向softmax的部分。这样设计的一种意义是可以让前期快速梯度下降，然后平滑的将logits过度到softmax之后的概率分布（softmax之后的logits其实就是一个PSSM矩阵）。不过实际测试代码可以看到系数lambda的值应当成上e_soft=0.9，这与论文有一点不同。

1 2	lr = self.opt["learning_rate"] * lr_scale self._params = jax.tree_map(lambda x,g:x-lr*g, self._params, self.aux["grad"])

如果初始迭代后得到的最佳plddt>0.65（后续也是这样），才会继续后续优化，否则会重新初始化。

if advanced_settings["optimise_beta"]:
    # temporarily dump model to assess secondary structure
    af_model.save_pdb(model_pdb_path)
    _, beta, *_ = calc_ss_percentage(model_pdb_path, advanced_settings, 'B')
    os.remove(model_pdb_path)

    # if beta sheeted trajectory is detected then choose to optimise
    if float(beta) > 15:
        advanced_settings["soft_iterations"] = advanced_settings["soft_iterations"] + advanced_settings["optimise_beta_extra_soft"]
        advanced_settings["temporary_iterations"] = advanced_settings["temporary_iterations"] + advanced_settings["optimise_beta_extra_temp"]
        af_model.set_opt(num_recycles=advanced_settings["optimise_beta_recycles_design"])
        print("Beta sheeted trajectory detected, optimising settings")

正如论文中补充说明的，算法在检测到trajectory存在明显的beta-sheet时，会自动增加recycle的数量以盼提高预测准确性。这可能是因为beta-sheet结构更加稳定规律，因此需要更高的准确性。但其实也可以同时考虑helix。

由于是alphafold进行结构预测，因此在输入预测模型前会根据当前分布得到一个序列氨基酸信息，具体代码如下：

1 2	# colabdesign.af.model.mk_af_model:_get_model:_model line 167 update_aatype(seq['pseudo'][0].argmax(-1), inputs)

序列logits的SoftMax优化

这一个阶段实际上是从上一阶段平滑过渡下来的。设置lambda=1(代码中soft=1)，使得上述soft_seq方法只考虑softmax之后的logits。softmax优化实际上就是得到一个最为合理的PSSM矩阵。

1
2

af_model.design_soft(advanced_settings["temporary_iterations"], e_temp=1e-2, models=design_models, num_models=1,
    sample_models=advanced_settings["sample_models"], ramp_recycles=False, save_best=True)

序列Onehot优化（ArgMax）

然后onehot优化是在softmax优化的基础上，对每个位置概率最高（argmax）的氨基酸进行概率微调。这里具体实现同样是在上面的soft_seq方法中。

1
2

af_model.design_hard(advanced_settings["hard_iterations"], temp=1e-2, models=design_models, num_models=1,
    sample_models=advanced_settings["sample_models"], dropout=False, ramp_recycles=False, save_best=True)

由于argmax运算是不可微的，因此这里通过选择性裁切softmax logits的梯度信息来实现反向传播更新。

seq["pssm"] = jax.nn.softmax(seq["logits"])
seq["soft"] = jax.nn.softmax(seq["logits"] / opt["temp"])
seq["hard"] = jax.nn.one_hot(seq["soft"].argmax(-1), seq["soft"].shape[-1])
seq["hard"] = jax.lax.stop_gradient(seq["hard"] - seq["soft"]) + seq["soft"]

# create pseudo sequence
seq["pseudo"] = opt["soft"] * seq["soft"] + (1-opt["soft"]) * seq["input"]
seq["pseudo"] = opt["hard"] * seq["hard"] + (1-opt["hard"]) * seq["pseudo"]

ProteinMPNN非接触序列优化

这一步通过固定接触面的设计残基，重新设计其他部位进行性质改善，筛选保留得分最高的两个。

1	mpnn_trajectories = mpnn_gen_sequence(trajectory_pdb, binder_chain, trajectory_interface_residues, advanced_settings)

设计结果的指标筛选

AF2 confidence pLDDT score of the predicted complex (> 0.8)
AF2 interface predicted confidence score (i_pTM) (> 0.5)
AF2 interface predicted alignment error (i_pAE) (> 0.35)
Rosetta interface shape complementarity (> 0.55)
Number of unsaturated hydrogen bonds at the interface (< 3)
Hydrophobicity of binder surface (< 35%)
RMSD of binder predicted in bound and unbound form (< 3.5 Å)

生成测试

PDL1

本文作者: GCS-ZHN

本文链接: https://gcszhn.top/2024/10/23/BindCraft算法解析和测试/