解决问题
蛋白模拟
小小编程
机器学习
要看文献
成长思考
自娱自乐
ESM3蛋白大模型学习
type
status
date
slug
summary
tags
category
icon
password
ESM3 是一种生物学前沿生成模型,能够联合推理蛋白质的三种基本生物学特性:序列、结构和功能。 这三种数据模式在 ESM3 的输入和输出端以离散标记的轨迹表示。 可以向模型提供跨轨道的部分输入组合,ESM3 将为所有轨道提供输出预测
模型采用了类似BERT的encoder-only架构,并加入了geometric attention(几何注意)等技术。该模型的输入包括序列、结构和功能等七种不同的部分,使用了大约10^24次方的计算资源进行训练,总共包含了98亿个参数。ESM3的发表标志着生物大模型时代的到来。

特点:
- 多模态:除了语言模型,还有序列,结构和功能。SM3与不同于模拟多种生物分子的AlphaFold 3,只专注于蛋白质,但可以同时推理其序列、结构和功能,这种多模态能力属于领域首创。
- 生成式模型。
1. 模型输入(3大类7小类)
Sequence蛋白序列
20种氨基酸+5种特殊+修饰残基
Structure结构
- 结构tokens
- dssp二级结构8中结构
- SASA溶剂可及表面
- 蛋白质绝对坐标
蛋白功能
- Function keywords
- Residue annotations残基标注

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

encoder

decoder

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

3. 模型架构
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%的相似性,展示了该模型在探索自然界未发现的蛋白质设计空间中的能力。
补充知识
Scaling Laws for Neural Language Models:性能随模型大小、数据量、预算增加而提升。模型形状影响小。
NLP:自然语言处理Natural language processing (NLP)
为何要将子单元数值化(映射成向量):除了决策树模型,机器学习中绝大多数模型是不支持字符串数据的,想要模型能够顺利有效地学习,必须对字符串数据先数值化。另外,我们并不是直接对输入句子或者单词进行数值化,我们需要先将其切分成一个个有限的子单元,然后将这些子单元数值化。而这个将原始文本切分成子单元的过程就叫做Tokenization。
Tokenization:将原始文本切分成子单元的过程。
数值化:数值化的过程必然需要映射,而映射又需要一个目标集合或者说映射表。
3种token方法:
- 词粒度Tokenization

- 字粒度Tokenization

- subword(子词)粒度的Tokenization

基于计算的文本分析需要tokenization。对于文本(以英语为例),tokens可以是单词,单个字母或者子词(sub-word segmentation)。根据应用场景,这几种方法各有优缺点。对于蛋白质也是类似的,最简单、常用的tokens是单个氨基酸残基。
自编码器(Autoencoder, AE) ,包括编码和解码,将图形编码为短的向量z,再将向量解码为图形,训练目的让重建的图像和原图形相似。但不认识其他的向量,只能压缩图像。

损失函数:整个编码器-解码器架构在同一损失函数上共同训练,损失函数鼓励在输出端重建输入。 所以一般情况下使用的损失函数是编码器输入和解码器输出之间的均方误差。

Variational AutoEncoders - VAE变分自编码器(VAE)就是这样一种改进版的AE,约束了编码向量z,使z满足标准正态分布。解码器不仅认识编码器编出的向量,还认识其他来自标准正态分布的向量。训练完成后,我们就可以扔掉编码器,用来自标准正态分布的随机向量和解码器来实现随机图像生成了。

损失函数:由VAE两个目标定义。VAE两个目标是重建的输入和潜空间应为正态分布。因此,训练损失定义为重构损失和相似损失之和。重构损失还是是输入输出和重构输出的均方损失。相似损失是潜在空间分布与标准高斯(零均值和单位方差)之间的KL散度。损失函数是这两个损失的总和。

但是VAE编码出来的向量是连续向量。每一维都是浮点数,修改一点也能输出图片,但是不好看。VQ-VAE的作者认为,VAE的生成图片之所以质量不高,是因为图片被编码成了连续向量。而实际上,把图片编码成离散向量会更加自然。比如说好的美术作品是特征明确的。性别是男或者女也就是1或者0,而不是0.4.因此,VQ-VAE会把图片编码成离散向量。
因此,VQ-VAE会把图片编码成离散向量。但是也会面临问题。
神经网络会默认输入满足一个连续的分布,而不善于处理离散的输入。借鉴NLP中对于离散单词的处理方法。为了处理离散的输入单词,NLP模型的第一层一般都是词嵌入层,它可以把每个输入单词都映射到一个独一无二的连续向量上。这样,每个离散的数字都变成了一个特别的连续向量了。
把类似的嵌入层加到VQ-VAE的解码器前。这个嵌入层在VQ-VAE里叫做"embedding space(嵌入空间)",在后续文章中则被称作"codebook"。

但离散向量的另一个问题是它不好采样。现在和AE面临一样问题。
VQ-VAE和AE的唯一区别,就是VQ-VAE会编码出离散向量,而AE会编码出连续向量。
VQ-VAE不是一个VAE,而是一个AE。
参考资料:
Loading...
Last update: 2024-10-19
🎉终于搭建完成🎉
-- 感谢您的支持 ---
👏欢迎阅读👏