license: bsd-3-clause tags:

• dna
• 生物学
• 基因组学
• 鬣狗

HyenaDNA

欢迎！HyenaDNA是一个长距离基因组基础模型，在单核苷酸分辨率下预训练了长达100万个标记的上下文长度。

查看下方的模型概述和训练信息。更棒的是，看看这些资源。

资源:

• arxiv
• 博客
• colab
• github

所有HuggingFace模型链接:

我们上传了一个集合，包含了所有预训练的HyenaDNA检查点。

你会看到不同大小和序列长度的模型。在LongSafari组织中，还有每个模型的原始仅权重版本，这些模型设计为通过原始github仓库加载。这些模型与上述集合中的模型输出相同，只是接口不同。

查看每个模型的GPU需求。

使用HyenaDNA

在这个简短的代码示例中，我们展示了如何在序列分类任务上微调HyenaDNA。这个示例使用了medium检查点，最大序列长度为16万个核苷酸。请注意，如果你使用的序列长度超过了所选检查点的最大支持长度，训练将会失败。

在测试中，我们能够在Colab T4 GPU（16GB显存）上训练长达约25万个核苷酸的序列。对于更长的序列长度，需要更多的内存。

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers import TrainingArguments, Trainer, logging
import torch

# 实例化预训练模型
checkpoint = 'LongSafari/hyenadna-medium-160k-seqlen-hf'
max_length = 160_000

# 使用bfloat16以获得更好的速度和减少内存使用
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)

# 生成一些随机序列和标签
# 如果你复制这段代码，请将这里的序列和标签
# 替换为你自己的数据！
sequence = 'ACTG' * int(max_length/4)
sequence = [sequence] * 8  # 创建8个相同的样本
tokenized = tokenizer(sequence)["input_ids"]
labels = [0, 1] * 4

# 创建训练数据集
ds = Dataset.from_dict({"input_ids": tokenized, "labels": labels})
ds.set_format("pt")

# 初始化Trainer
# 注意我们使用了非常小的批次大小以最大化
# 在内存中容纳长序列的能力！
args = {
    "output_dir": "tmp",
    "num_train_epochs": 1,
    "per_device_train_batch_size": 1,
    "gradient_accumulation_steps": 4,
    "gradient_checkpointing": True,
    "learning_rate": 2e-5,
}
training_args = TrainingArguments(**args)

trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()

print(result)

# 现在我们可以使用save_pretrained()或push_to_hub()来分享训练好的模型！

你可能还会发现这些笔记本有用。虽然它们不是专门针对HyenaDNA的，但它们包含了训练DNA和序列分类模型的额外示例。

• 如何微调一个核苷酸Transformer模型
• 如何在文本分类上微调模型

GPU需求（建议）

以下是针对每个模型的硬件（建议最低配置）建议。

GPU在：预训练、微调、推理时

• tiny-1k: (T4, T4, T4)
• small-32k: (A100-40GB, T4, T4)
• medium-160k: (A100-40GB, T4, T4)
• medium-450k: (A100-40GB, A100-40GB, T4)
• large-1m: (A100-80GB, A100-80GB, A100-40GB)

模型与训练概述

HyenaDNA使用了一堆简单的Hyena操作符，这些操作符是Transformer中注意力的次二次替代品。Hyena操作符通过修改输入投影、隐式卷积和门控（所有次二次操作）来匹配语言建模的质量。

这使得HyenaDNA能够达到比之前使用密集注意力的基因组Transformer模型长500倍的上下文长度，并且在1M序列长度下训练速度快160倍（相比Flash Attention）。

我们使用了一个单字符分词器，主要词汇表包含4个核苷酸（加上特殊标记），实现了单核苷酸分辨率，这在基因组基础模型中尚属首次。此外，隐式长卷积使得每一层都具有全局感受野。

我们在人类参考基因组（HG38）上使用下一个标记（核苷酸）预测进行预训练。

HyenaDNA在23个下游任务上设定了新的SotA，包括预测调控元件、染色质谱和物种分类。我们还探索了长上下文在基因组学中开启的新能力，包括首次使用上下文学习与软提示可调标记和指令微调。

查看我们的博客了解更多关于HyenaDNA的细节！

作者

Eric Nguyen*, Michael Poli*, Marjan Faizi*, Armin Thomas, Callum Birch-Sykes, Michael Wornow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, Stefano Ermon, Stephen Baccus, Chris Re.

联系方式

Eric Nguyen, etnguyen@stanford.edu
Michael Poli, poli@stanford.edu
Marjan Faizi, Marjan_Faizi@hms.harvard.edu

引用

欢迎引用我们 :)

@article{nguyen2023hyenadna,
      title={HyenaDNA: Long-Range Genomic Sequence Modeling at Single Nucleotide Resolution}, 
      author={Eric Nguyen and Michael Poli and Marjan Faizi and Armin Thomas and Callum Birch-Sykes and Michael Wornow and Aman Patel and Clayton Rabideau and Stefano Massaroli and Yoshua Bengio and Stefano Ermon and Stephen A. Baccus and Chris Ré},
      year={2023},
      eprint={2306.15794},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}