🚀 HyenaDNA

HyenaDNA是一个长距离基因组基础模型，它在单核苷酸分辨率下，以高达100万个标记的上下文长度进行了预训练。该模型可以解决基因组序列建模中的长距离依赖问题，为基因组学研究提供了强大的工具。

🚀 快速开始

资源链接

• arxiv
• 博客
• Colab
• GitHub

HuggingFace模型链接

我们已上传了所有预训练的HyenaDNA检查点的集合。你会看到不同大小和序列长度的模型。在LongSafari组织中还有每个模型的原始仅权重版本，这些版本设计为使用原始的GitHub仓库加载。这些模型的输出与上述集合中的模型相同，只是接口不同。

模型GPU要求

请参考GPU要求部分了解每个模型的建议硬件配置。

✨ 主要特性

• 长距离建模能力：能够处理高达100万个标记的上下文长度，远超以往的基因组Transformer模型。
• 单核苷酸分辨率：使用单字符分词器，实现单核苷酸分辨率的建模。
• 全局感受野：隐式长卷积使每一层都具有全局感受野。
• 优异的下游任务表现：在23个下游任务中达到了新的最优水平，包括预测调控元件、染色质图谱和物种分类等。

📦 安装指南

文档未提及具体安装步骤，暂不提供。

💻 使用示例

基础用法

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

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

# instantiate pretrained model
checkpoint = 'LongSafari/hyenadna-medium-160k-seqlen-hf'
max_length = 160_000

# bfloat16 for better speed and reduced memory usage
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)

# Generate some random sequence and labels
# If you're copying this code, replace the sequences and labels
# here with your own data!
sequence = 'ACTG' * int(max_length/4)
sequence = [sequence] * 8  # Create 8 identical samples
tokenized = tokenizer(sequence)["input_ids"]
labels = [0, 1] * 4

# Create a dataset for training
ds = Dataset.from_dict({"input_ids": tokenized, "labels": labels})
ds.set_format("pt")

# Initialize Trainer
# Note that we're using extremely small batch sizes to maximize
# our ability to fit long sequences in memory!
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)

# Now we can save_pretrained() or push_to_hub() to share the trained model!

其他参考资源

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

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

📚 详细文档

GPU要求（建议）

以下是我们认为每个模型可以使用的硬件（首选最低配置）建议。

模型与训练概述

HyenaDNA使用了一个简单的Hyena算子栈，这是一种次二次复杂度的算子，可替代Transformer中的注意力机制。Hyena算子通过使用修改后的输入投影、隐式卷积和门控（均为次二次操作），在语言建模中达到了与注意力机制相当的效果。

这使得HyenaDNA能够处理的上下文长度比使用密集注意力的以往基因组Transformer模型长500倍，并且在序列长度为100万时，训练速度比Flash Attention快160倍。

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

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

HyenaDNA在23个下游任务中达到了新的最优水平，包括预测调控元件、染色质图谱和物种分类等。我们还探索了在基因组学中长上下文带来的新能力，包括首次使用带有软提示可调标记的上下文学习和指令微调。

更多关于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

🔧 技术细节

• Hyena算子：一种次二次复杂度的算子，替代Transformer中的注意力机制，通过修改输入投影、隐式卷积和门控实现高效建模。
• 单字符分词器：实现单核苷酸分辨率的关键，使模型能够对单个核苷酸进行建模。
• 隐式长卷积：赋予模型全局感受野，增强长距离建模能力。
• 预训练任务：使用下一个核苷酸预测任务在人类参考基因组上进行预训练。

📄 许可证

本项目采用BSD 3 - 条款许可证。

📖 引用

如果你使用了本项目的相关内容，请引用以下文献：

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