跳转到内容

ESM-2

模型概述

描述

ESM-2 是一个预训练的双向编码器 (BERT 风格模型)，用于氨基酸序列。 ESM-2 模型为氨基酸提供嵌入，这些嵌入在下游任务（如结构和功能预测）上实现了最先进的性能。 ESM-2 已经以多种不同的模型尺寸进行了训练。 BioNeMo2 包括 650M 和 3B 参数变体的转换检查点。 650M 模型有 33 层、20 个注意力头和 1280 的隐藏空间维度。 3B 模型有 36 层、40 个注意力头和 2,560 的隐藏空间维度。

这些模型已准备好用于商业用途。

第三方社区考量

此模型并非 NVIDIA 拥有或开发。 此模型是根据第三方对该应用和用例的要求开发和构建的 [1]；请参阅 ESM-2 3B 模型的非 NVIDIA 模型卡 和 ESM-2 650M 模型的非 NVIDIA 模型卡 的链接

参考文献

[1] Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y. 和 dos Santos Costa, A., 2023. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637), pp.1123-1130.

[2] "UniProt: the universal protein knowledgebase in 2021." Nucleic acids research 49, no. D1 (2021): D480-D489.

[3] Devlin, J., Chang, M.W., Lee, K. and Toutanova, K., 2018. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.

模型架构

架构类型: BERT

网络架构: ESM-2

输入

输入类型: 文本 (蛋白质序列)

输入参数: 1D

与输入相关的其他属性: 蛋白质序列表示为规范氨基酸字符串，最大长度为 1022。 更长的序列会自动截断为该长度。

输出

输出类型: 嵌入 (氨基酸和序列级别)

输出参数: 1D

与输出相关的其他属性: 数值向量，具有浮点值，对应于输入蛋白质序列中每个氨基酸的嵌入。 最大输出长度为 1022 个嵌入 - 每个氨基酸一个嵌入向量。

软件集成

运行时引擎

  • BioNeMo, NeMo, Megatron, TransformerEngine

支持的硬件微架构兼容性

  • [Ampere]
  • [Hopper]
  • [Volta]

[首选/支持的] 操作系统

  • [Linux]

模型版本

训练与评估

训练数据集

来自 HuggingFace 的原始 ESM-2 检查点使用 UniProt 2021_04 序列数据库进行训练。 有关训练数据集的更多详细信息，请参阅 Linet al. 2023。 原始作者使用的训练/测试拆分未分发。 NVIDIA 编译的预训练数据库遵循类似的方法，详见 UniProt 数据集。

推理

引擎: BioNeMo, NeMo

测试硬件

  • [Ampere]
  • [Hopper]
  • [Volta]

许可证

ESM-2 按照 Apache 2.0 许可证提供。

竞争性基准测试

准确率

从 UniRef 2024_03（参见 UniProt 数据集）中随机选择了 328,360 个 UniRef50 代表性序列的验证集。 此验证集用于确保 BioNeMo 转换的检查点的输出与其在使用 HuggingFace Transformers 库评估时的输出一致。

检查点 HuggingFace BioNeMo2 Lin et al. 2023
650M 7.001 7.002 6.95
3B 6.003 6.004 6.49

不同的验证集

HuggingFace 和转换后的 BioNeMo2 检查点在一个新策划的验证集上进行了评估。 Lin et al. 2023 报告了困惑度以供比较，但原始的训练/测试拆分不可用。

训练性能

单节点训练性能

ESM-2 Single-Device Training Performance

对于 ESM2-650M 模型尺寸，纯 pytorch 基线（使用 torch.compile() 编译）在批处理大小大于 16 时引发了内存不足错误。 bionemo2 模型可以处理 46 的批处理大小，在 NVIDIA A100 上达到 59.2% 的模型 FLOP 利用率。

模型扩展

ESM-2 Model Scaling

在 650M、3B 和 15B 模型变体上训练 ESM-2 表明，与纯 pytorch 基线相比，BioNeMo2 框架具有改进的性能。 这些实验在分布在两个节点上的 16 个 NVIDIA A100 或 16 个 NVIDIA H100 GPU 上进行。 注意:* 15B 模型变体在 64 个 GPU 上使用 BioNeMo2 框架进行训练。

设备扩展

ESM-2 Device Scaling

在 32 个节点上的 256 个 NVIDIA A100 上训练 ESM-3B 实现了 96.85% 的理论线性吞吐量，该吞吐量预计从单节点（8 个 GPU）性能外推，表示在 256 个设备上模型 FLOP 利用率为 60.6%。