英文

nucleotide-transformer-500m-human-ref模型

核苷酸Transformer是一系列基础语言模型,它们是在全基因组的DNA序列上进行预训练的。与其他方法相比,我们的模型不仅整合了单一参考基因组的信息,还利用了来自3200多个不同人类基因组的DNA序列,以及来自各种物种(包括模型和非模型生物)的850个基因组的信息。通过强大而广泛的评估,我们展示了这些大型模型与现有方法相比提供了极其准确的分子表型预测。

这个系列中的一部分是 nucleotide-transformer-500m-human-ref,一个在人类参考基因组上进行预训练的500M参数Transformer模型。该模型提供了Tensorflow和Pytorch两种版本。

开发者:InstaDeep、NVIDIA和TUM

模型来源

如何使用

在下一次发布之前,需要使用以下命令从源代码中安装 transformers 库以使用这些模型:

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

这里提供了一个小的代码片段,以便从一个虚拟的DNA序列中检索logits和嵌入。

from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

# Import the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref")
model = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref")

# Create a dummy dna sequence and tokenize it
sequences = ['ATTCTG' * 9]
tokens_ids = tokenizer.batch_encode_plus(sequences, return_tensors="pt")["input_ids"]

# Compute the embeddings
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
)

# Compute sequences embeddings
embeddings = torch_outs['hidden_states'][-1].detach().numpy()
print(f"Embeddings shape: {embeddings.shape}")
print(f"Embeddings per token: {embeddings}")

# Compute mean embeddings per sequence
mean_sequence_embeddings = torch.sum(attention_mask.unsqueeze(-1)*embeddings, axis=-2)/torch.sum(attention_mask, axis=-1)
print(f"Mean sequence embeddings: {mean_sequence_embeddings}")

训练数据

nucleotide-transformer-500m-human-ref模型是在 GRCh38 human reference genome 上进行预训练的,该数据集可作为HuggingFace数据集 here 提供,包含30亿个核苷酸,约占500M个6-mer标记。

训练流程

预处理

DNA序列使用核苷酸Transformer分词器进行分词,当可能时,将序列分割为6-mer标记,否则将每个核苷酸分开分词,详细说明可参考相关存储库中的 Tokenization 部分。该分词器拥有一个词汇表大小为4105。然后,模型的输入形式为:

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

分词后的序列最大长度为1,000。

使用的掩码处理方法是Bert风格训练的标准方法:

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

预训练

使用8个A100 80GB GPU训练模型,总共使用300B个标记,有效的批大小为1M个标记。序列长度为1000个标记。使用Adam优化器[38],学习率调度,以及指数衰减率和epsilon常数的标准值,β1 = 0.9,β2 = 0.999,ε=1e-8。在第一个热身阶段,学习率在16k步内线性增加从5e-5到1e-4,然后按照平方根衰减到训练结束。

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}
}