——简介——
如何高效采样蛋白质尤其是天然无规蛋白质(IDP)的多种结构一直是计算结构生物学领域的一大难点。通常,此类问题是通过分子模拟(包括MC采样和MD采样)来解决,同时配合各种增强采样算法提高效率。近年来,得益于GPU算力的进步,各种基于深度学习的增强采样和构象生成算法被提出,但大都需要针对某一体系预先的MD轨迹数据,没有办法直接从序列生成具有多样性的构象。2024年1月,来自蒙特利尔大学Mila研究院Jian Tang课题组的Jiarui Lu等人开发了基于分数模型蛋白质zero-shot构象采样模型Str2str[1]。该工作目前在ICLR2024上作为Poster发布。该模型通过输入蛋白序列和结构(例如使用AlphaFold预测的结构),在一个分数模型上加噪再去噪,从而采样新的构象。该工具不同于之前的AI辅助增强采样方法和构象生成方法,不需要对结构提前进行MD采样,大大扩展了模型的使用范围。
——方法——
Str2str全称structure to structure,其基本逻辑如下图1所示:先利用蛋白质序列-结构数据集训练一个具有平移旋转等变性的分数模型;之后针对某一输入的蛋白序列通过结构预测模型如AlphaFold2等预测一个初始结构;将结构按一定步数不完全加噪;利用训练好的模型对结构进行去噪;最后得到的结构即为生成后的新结构。
这个先不完全加噪再去噪的过程被作者称为forward-backward过程(FB),具体公式如下:
其中,Tδ取值在0到T之间,并且可以通过控制其大小,来达到控制生成构象集合(ensemble)的结构多样性的目的。在训练时,分数模型的随机微分方程(SDE)如下所示:
为了更好的训练效果,作者仿照AlphaFold2的处理方法来表征蛋白结构:将蛋白质按氨基酸为单位划分为主链“frame”以及侧链各种扭转角。其中,对于蛋白主链的表征定义如下:其中,Xglobal代表蛋白的笛卡尔坐标,Xlocal代表在“frame”参考系下的坐标。Tframe可分解为一个旋转项Rframe和一个平移项vframe。而对于侧链,则仿照AF2的侧链对照表处理,本推文暂不赘述。至于去噪的神经网络架构,作者使用了一个更改后的IPA网络,被作者命名为DenoisingIPA,如下图所示:
可见,基本与AF2中的IPA模块相同,其中Si张量来源于positional encoding和分数模型的Time step encoding。此外,作者在其中加入了一个pair representation的transition模块,类似于AF2的Evoformer中所使用的edge transition模块,如下:值得注意的是,不同AF2的IPA模块,这个架构中其实并没有嵌入序列信息,所以其实蛋白序列的唯一作用是在一开始通过AF2或ESMFold等工具预测初始结构。在训练过程中,作者是用了两种loss函数来保证训练效果。其一是正常分数模型的loss函数:
与正常分数模型训练过程不同的是,由于在实际使用时,并不需要用到全部的扩散步数,因此可以只对0到Tm的区间算loss,其中Tm是一个提前设定的超参数,0<Tm<T。另一种loss被称为Auxiliary Loss,由主链MSEloss:之后,作者对Str2str的效果进行了benchmark。Benchmark的对象主要包括AF2-RVAE、EigenFold、idpGAN等。评价标准主要包括:1.有效性:包括原子clash和化学键键长。2.保真度:两两距离分布的JS散度(JS-PwD)和tICA降维图上的最近组分(JS-TIC)以及回旋半径Rg分布的JS散度(JS-Rg)。3.多样性:包括RMSD和TM-score的MAE。Benchmark结果如下表1所示,其中PF表示使用概率流进行采样的结果:
可见,与之前的EigenFold和idpGAN相比,Str2str不仅在有效性和保真度上有优势,在多样性上超过之前的方法很多,同时和MD的结果进行比较也能发现,单从采样效率上讲,Str2str的似乎媲美几十us级的MD采样。除了量化的benchmark之外,作者对很多MD采样领域比较著名的case进行了具体的分析,包括很多的fast-folding蛋白如:BPTI、Trp-cage等。其中BPTI的tICA结果如下所示:可见,Str2str的生成结果和长时间MD更类似,而其他方法则和短MD的结果更为类似,体现了Str2str更高的采样效率。Trp-cage的contact-map分布图如下图所示:和BPTI的结果类似,Str2str展现了更类似于长时间MD的采样效率。其他各种case的contact-map结果如下如所示:
图5. 各种fast-folding蛋白的benchmark结果此外,作者还统计了在不同去噪步数的选取情况下Str2str的表现,展现了模型可以通过参数控制生成结果多样性的特点,如下图所示:图6. 不同T_delta下的ensemble生成结果比较有意思的一点是随着T_delta的增大,模型似乎能更容易采样到不同的构象。下图的tICA结果也说明了这一点:
图7. 不同T_delta下的WW domain ensemble生成结果和MD结果比较Str2str的提出不仅为我们提供了一种高效的AI采样方法,同时这种不完全加噪方式也为将扩散模型迁移到flow上提供了一些启示。此外,据作者言,这种不完全加噪再去噪的过程实际上暗合MC采样或MD采样中跨越势能垒的过程,同样引人深思。图8. Str2str与常规采样方法的能垒跨越过程类比当然,就目前而言基于AI的增强采样方法仍然任重道远:比如,Str2str的结果显示,采样结构的比例和MD存在较大差距,这说明模型整体对于玻尔兹曼分布的感知还不足,只能作为初始结构采样器使用。其次,文章展现的都是一些fast-folding的蛋白,对于真正类coil的IDP蛋白缺乏足够的case study。Lu, Jiarui, et al. "Str2str: A score-based framework for zero-shot protein conformation sampling." The Twelfth International Conference on Learning Representations. 2024.