ESM3——开启生物大模型时代

关于ESM3前身模型（如ESM-1b、ESM-MSA等）可以看之前文章：生物大模型——ESM模型（一）-CSDN博客

这篇文章主要讲解ESM3模型，部分图片来自小破站视频（关于该模型的详细视频讲解）：【ESM3: 生物大模型时代并未到来】https://www.bilibili.com/video/BV15SbYeSEzM?vd_source=dd65e5938f5f6dab5dc478dad590c5f9

论文地址：https://www.biorxiv.org/content/10.1101/2024.07.01.600583v1.full.pdf

代码：GitHub - evolutionaryscale/esm

官方网站：Evolutionary Scale · ESM3: Simulating 500 million years of evolution with a language model

概述

一个多模态的生成型语言模型：即除了语言模型外,还包括蛋白质的序列结构和功能信息。

采用了类似BERT的encoder-only架构,并加入了geometric attention（几何注意）等技术。该模型的输入包括序列、结构和功能等七种不同的部分,使用了大约10^24次方的计算资源进行训练,总共包含了98亿个参数。ESM3的发表标志着生物大模型时代的到来。

两大特点：多模态（序列、结构、功能），生成式模型

1、模型输入为三大类（7小类）

Sequence（蛋白质序列） 

Structure（蛋白质结构 包括：结构tokens   SS8：二级结构  SASA：溶剂可访问区域  coordinate 蛋白质绝对坐标）

Function（蛋白质功能 包括：Function keywords  Residue annotations）

（图片来源：钟二的伪电气白兰）

2、标记化（Tokenization）：

序列编码：蛋白质序列通过标记化转换为标记序列，其中每个氨基酸都被表示为一个标记。

结构编码：蛋白质结构通过离散自动编码器（Discrete Auto-Encoder）将三维原子结构压缩为离散标记。具体：使用 VQ-VAE（向量量化变分自编码器）来对蛋白质的三维结构进行编码。通过将高维的连续原子结构数据压缩为离散标记，ESM3 能够有效地表示和重建蛋白质的三维结构。ESM3提出了一种不变几何注意机制（Invariant Geometric Attention Mechanism），用于高效处理三维结构的标记。

（演示VQ-VAE，非论文结构编码流程流程）

功能编码：功能信息以每个序列位置的关键字集合的形式进行标记化编码。

ESM3 的功能编码通过将每个序列位置的生物学功能注释转化为离散标记来实现。这些功能注释来自预测的功能关键词集合，利用隐马尔可夫模型库对蛋白质序列进行注释。

总结多模态输入：

模型设计为能够接受多个模态输入，包括序列、结构和功能注释。这些输入作为独立的通道输入模型，并在模型中融合为单一的潜在空间。

这种方法允许ESM3根据不同的输入组合进行蛋白质生成，确保生成的新蛋白质能够符合指定的功能和结构要求。

3、模型结构

Transformer架构：

ESM3采用双向Transformer架构。输入的序列、结构和功能轨道被嵌入并融合，然后通过一系列Transformer块进行处理。

在模型的第一个Transformer块中包含了几何注意层，用于条件化原子结构坐标。

几何注意机制：

提出了一种不变几何注意机制来高效处理三维结构。该机制在每个氨基酸的键合几何定义的局部参考框架中运行，并通过转换到全局框架实现局部框架的全局交互。

这种机制通过注意机制的计算原语实现，并且具有良好的可扩展性。

输出头（Output Heads）：

模型输出端使用浅层MLP头将最终层的表示转换为每个轨道的标记概率。

4、模型训练

1.生成式掩码语言建模目标：

ESM3采用生成式掩码语言模型（Masked Language Model, MLM）进行训练。模型在每种模态的离散标记上进行掩码处理，并监督模型预测被掩码标记的身份。

在训练过程中，掩码标记会以噪声调度的方式进行采样，确保ESM3能够在不同的掩码组合下进行训练。这种方式区别于传统的掩码语言建模，允许模型从任何起点生成任意顺序的标记。

2.训练数据集：

ESM3在一个包含2.78亿个天然蛋白质的数据库上进行训练，这些数据来源于已知的序列和结构数据库，以及一些预测的结构。

为了增加训练数据，研究人员还使用逆折叠模型生成合成序列，并通过隐马尔可夫模型库预测功能注释。

3.模型参数规模：

ESM3模型在三个不同的参数规模下进行训练：1.4亿、7亿和98亿参数。模型的最终结构选择了相对较深的网络，其中98亿参数的模型包含216个Transformer块。

随着模型参数规模的增加，验证损失显著降低，尤其是在序列损失方面的改善最为明显。

5、输出与验证

1.生成与表示学习：

ESM3在生成过程中，标记以迭代的方式进行采样。从全掩码标记序列开始，标记可以一次或并行采样，直到所有标记完全解码。

模型的训练目标不仅有效地提升了生成能力，还提高了表示学习能力。噪声调度的选择使生成能力与表示学习达到平衡。

2.蛋白质生成：

ESM3能够在接受提示（prompts）的情况下生成新的蛋白质，并在生成蛋白质序列和结构上展示了高度的创新性（大模型生成的蛋白仅结构合理，未验证其功能性）。

模型可以基于给定的序列、结构或功能提示生成具有不同于天然蛋白质的结构和序列的蛋白质。

3.实验验证：

ESM3生成的新型绿色荧光蛋白（esmGFP）经过实验验证，显示出与自然界中的GFP相似的荧光特性。

esmGFP与最近的天然蛋白序列在序列上仅有36%的相似性，展示了该模型在探索自然界未发现的蛋白质设计空间中的能力。