Agro Nucleotide Transformer 1b

🚀 AgroNT - 植物基因組DNA語言模型

AgroNT是一個主要基於可食用植物基因組訓練的DNA語言模型。它運用Transformer架構和自注意力機制，通過掩碼語言建模目標，利用48種不同植物的基因型數據，學習核苷酸序列的通用表示。該模型擁有10億個參數，上下文窗口為1024個標記。

許可證

本項目採用CC BY-NC-SA 4.0許可證。

模型信息

🚀 快速開始

模型概述

AgroNT是一個主要在可食用植物基因組上訓練的DNA語言模型。具體而言，AgroNT使用具有自注意力機制的Transformer架構和掩碼語言建模目標，利用來自48種不同植物物種的高可用性基因型數據，學習核苷酸序列的通用表示。AgroNT包含10億個參數，上下文窗口為1024個標記。

AgroNT使用非重疊的6聚體分詞器將基因組核苷酸序列轉換為標記。因此，1024個標記大約對應6144個鹼基對。

📦 安裝指南

文檔未提及具體安裝步驟，可參考代碼示例從Hugging Face倉庫加載模型和分詞器。

💻 使用示例

基礎用法

from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch


model_name = 'agro-nucleotide-transformer-1b'

# fetch model and tokenizer from InstaDeep's hf repo
agro_nt_model = AutoModelForMaskedLM.from_pretrained(f'InstaDeepAI/{model_name}')
agro_nt_tokenizer = AutoTokenizer.from_pretrained(f'InstaDeepAI/{model_name}')

print(f"Loaded the {model_name} model with {agro_nt_model.num_parameters()} parameters and corresponding tokenizer.")

# example sequence and tokenization
sequences = ['ATATACGGCCGNC','GGGTATCGCTTCCGAC']

batch_tokens = agro_nt_tokenizer(sequences,padding="longest")['input_ids']
print(f"Tokenzied sequence: {agro_nt_tokenizer.batch_decode(batch_tokens)}")

torch_batch_tokens = torch.tensor(batch_tokens)
attention_mask = torch_batch_tokens != agro_nt_tokenizer.pad_token_id

# inference
outs = agro_nt_model(
    torch_batch_tokens,
    attention_mask=attention_mask,
    encoder_attention_mask=attention_mask,
    output_hidden_states=True
)

# get the final layer embeddings and language model head logits
embeddings = outs['hidden_states'][-1].detach().numpy()
logits = outs['logits'].detach().numpy()

📚 詳細文檔

預訓練

數據

我們的預訓練數據集主要基於Ensembl Plants數據庫中的可食用植物參考基因組構建。該數據集包含48個不同物種的約1050萬個基因組序列。

處理

每個物種的所有參考基因組被組裝成一個單一的FASTA文件。在這個FASTA文件中，除A、T、C、G之外的所有核苷酸都被替換為N。使用分詞器將字母字符串轉換為標記序列。

分詞器的字母表由通過組合A、T、C、G獲得的4⁶ = 4096種可能的6聚體組合，以及代表獨立A、T、C、G和N的五個額外標記組成。它還包括三個特殊標記：填充標記[PAD]、掩碼標記[MASK]和分類標記[CLS]。這導致詞彙表包含4104個標記。為了對輸入序列進行分詞，分詞器從分類標記開始，然後從左到右轉換序列，儘可能匹配6聚體標記，必要時使用獨立標記（例如，當存在字母N或序列長度不是6的倍數時）。

分詞示例

核苷酸序列: ATCCCGGNNTCGACACN 標記: <CLS> <ATCCCG> <G> <N> <N> <TCGACA> <C> <N>

訓練

使用掩碼語言建模（MLM）目標以自監督方式對AgroNT進行預訓練。在自監督學習設置中，不需要每個序列的註釋（監督），因為我們可以對序列的一部分進行掩碼，並使用未掩碼部分的信息來預測掩碼位置。

這使我們能夠利用大量未標記的基因組測序數據。具體而言，輸入序列中15%的標記被選擇進行增強，其中80%被掩碼標記替換，10%被詞彙表中的另一個標記隨機替換，最後10%保持相同的標記。

分詞後的序列通過模型，對掩碼標記計算交叉熵損失。預訓練使用1024個標記的序列長度和150萬個標記的有效批量大小進行315000次更新步驟，導致模型總共訓練了4725億個標記。

硬件

模型預訓練使用Google TPU - V4加速器進行，具體是包含512個設備的TPU v4 - 1024。我們總共訓練了大約四天。

BibTeX引用

@article{mendoza2023foundational,
  title={A Foundational Large Language Model for Edible Plant Genomes},
  author={Mendoza-Revilla, Javier and Trop, Evan and Gonzalez, Liam and Roller, Masa and Dalla-Torre, Hugo and de Almeida, Bernardo P and Richard, Guillaume and Caton, Jonathan and Lopez Carranza, Nicolas and Skwark, Marcin and others},
  journal={bioRxiv},
  pages={2023--10},
  year={2023},
  publisher={Cold Spring Harbor Laboratory}
}