Nucleotide Transformer 500m 1000g

🚀 nucleotide-transformer-500m-1000g 模型

Nucleotide Transformers 是一系列基于全基因组 DNA 序列进行预训练的基础语言模型。与其他方法相比，我们的模型不仅整合了单一参考基因组的信息，还利用了来自 3200 多个不同人类基因组以及来自包括模式生物和非模式生物在内的广泛物种的 850 个基因组的 DNA 序列。通过稳健且广泛的评估，我们表明这些大模型与现有方法相比，能够提供极其准确的分子表型预测。

此系列模型中的 nucleotide-transformer-500m-1000g 是一个具有 5 亿参数的 Transformer 模型，它在 3202 个基因多样的人类基因组集合上进行了预训练。该模型同时提供了 Tensorflow 和 Pytorch 版本。

开发者： InstaDeep、NVIDIA 和 TUM

模型资源

• 仓库地址： Nucleotide Transformer
• 论文地址： The Nucleotide Transformer: Building and Evaluating Robust Foundation Models for Human Genomics

🚀 快速开始

在其下一次发布之前，若要使用该模型，需要通过以下命令从源代码安装 transformers 库：

pip install --upgrade git+https://github.com/huggingface/transformers.git

以下是一段示例代码，用于从一个虚拟 DNA 序列中获取对数几率和嵌入向量：

from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

# 导入分词器和模型
tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
model = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")

# 选择输入序列填充的长度。默认情况下，选择模型的最大长度，但您可以根据需要减小它，因为获取嵌入向量所需的时间会随着长度的增加而显著增加。
max_length = tokenizer.model_max_length

# 创建一个虚拟 DNA 序列并进行分词
sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]
tokens_ids = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]

# 计算嵌入向量
attention_mask = tokens_ids != tokenizer.pad_token_id
torch_outs = model(
    tokens_ids,
    attention_mask=attention_mask,
    encoder_attention_mask=attention_mask,
    output_hidden_states=True
)

# 计算序列嵌入向量
embeddings = torch_outs['hidden_states'][-1].detach().numpy()
print(f"嵌入向量形状: {embeddings.shape}")
print(f"每个标记的嵌入向量: {embeddings}")

# 添加嵌入维度轴
attention_mask = torch.unsqueeze(attention_mask, dim=-1)

# 计算每个序列的平均嵌入向量
mean_sequence_embeddings = torch.sum(attention_mask*embeddings, axis=-2)/torch.sum(attention_mask, axis=1)
print(f"平均序列嵌入向量: {mean_sequence_embeddings}")

✨ 主要特性

• 基于全基因组 DNA 序列预训练，整合多源信息，能更准确地进行分子表型预测。
• 利用 3200 多个不同人类基因组及 850 个其他物种基因组的 DNA 序列。
• 提供 Tensorflow 和 Pytorch 版本。

📦 安装指南

在其下一次发布之前，若要使用该模型，需要通过以下命令从源代码安装 transformers 库：

pip install --upgrade git+https://github.com/huggingface/transformers.git

💻 使用示例

基础用法

from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

# 导入分词器和模型
tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
model = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")

# 选择输入序列填充的长度。默认情况下，选择模型的最大长度，但您可以根据需要减小它，因为获取嵌入向量所需的时间会随着长度的增加而显著增加。
max_length = tokenizer.model_max_length

# 创建一个虚拟 DNA 序列并进行分词
sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]
tokens_ids = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]

# 计算嵌入向量
attention_mask = tokens_ids != tokenizer.pad_token_id
torch_outs = model(
    tokens_ids,
    attention_mask=attention_mask,
    encoder_attention_mask=attention_mask,
    output_hidden_states=True
)

# 计算序列嵌入向量
embeddings = torch_outs['hidden_states'][-1].detach().numpy()
print(f"嵌入向量形状: {embeddings.shape}")
print(f"每个标记的嵌入向量: {embeddings}")

# 添加嵌入维度轴
attention_mask = torch.unsqueeze(attention_mask, dim=-1)

# 计算每个序列的平均嵌入向量
mean_sequence_embeddings = torch.sum(attention_mask*embeddings, axis=-2)/torch.sum(attention_mask, axis=1)
print(f"平均序列嵌入向量: {mean_sequence_embeddings}")

📚 详细文档

训练数据

nucleotide-transformer-500m-1000g 模型在 3202 个基因多样的人类基因组上进行了预训练，这些基因组来自 27 个地理结构群体，涵盖非洲、美洲、东亚和欧洲血统，数据来源于 1000G 项目。这种多样性使得数据集能够更好地编码人类遗传变异。为了在输入模型的序列中进行单倍型重建，我们使用了 1000G 基因组项目的分相版本，其中总共包含 1.25 亿个突变，其中 1.11 亿个为单核苷酸多态性（SNP），1400 万个为插入缺失。数据集中核苷酸的总数为 192120 亿个，大约产生 32020 亿个标记。

训练过程

预处理

DNA 序列使用 Nucleotide Transformer 分词器进行分词。该分词器会尽可能将序列分词为 6 聚体标记，否则会将每个核苷酸单独分词，具体可参考相关仓库的 Tokenization 部分。该分词器的词汇表大小为 4105。模型的输入形式如下：

<CLS> <ACGTGT> <ACGTGC> <ACGGAC> <GACTAG> <TCAGCA>

分词后的序列最大长度为 1000。

使用的掩码过程是标准的 Bert 式训练方法：

• 15% 的标记被掩码。
• 80% 的情况下，被掩码的标记被替换为 [MASK]。
• 10% 的情况下，被掩码的标记被替换为一个与原标记不同的随机标记。
• 剩下 10% 的情况下，被掩码的标记保持不变。

预训练

该模型在 8 个 A100 80GB GPU 上对 3000 亿个标记进行了训练，有效批量大小为 100 万个标记。使用的序列长度为 1000 个标记。训练使用了 Adam 优化器 [38]，并采用了学习率调度，指数衰减率和 ε 常数的标准值分别为 β1 = 0.9、β2 = 0.999 和 ε=1e-8。在第一个热身阶段，学习率在 16000 步内从 5e-5 线性增加到 1e-4，然后按照平方根衰减直至训练结束。

🔧 技术细节

• 分词策略：使用 Nucleotide Transformer 分词器，优先将序列分词为 6 聚体标记，否则单独分词每个核苷酸，词汇表大小 4105。
• 掩码策略：标准 Bert 式训练掩码，15% 标记掩码，不同情况有不同替换规则。
• 训练配置：8 个 A100 80GB GPU，3000 亿标记，有效批量 100 万标记，序列长度 1000 标记，Adam 优化器及特定学习率调度。

📄 许可证

本项目采用 CC BY-NC-SA 4.0 许可证。

BibTeX 引用信息

@article{dalla2023nucleotide,
  title={The Nucleotide Transformer: Building and Evaluating Robust Foundation Models for Human Genomics},
  author={Dalla-Torre, Hugo and Gonzalez, Liam and Mendoza Revilla, Javier and Lopez Carranza, Nicolas and Henryk Grywaczewski, Adam and Oteri, Francesco and Dallago, Christian and Trop, Evan and Sirelkhatim, Hassan and Richard, Guillaume and others},
  journal={bioRxiv},
  pages={2023--01},
  year={2023},
  publisher={Cold Spring Harbor Laboratory}
}