HyenaDNA

重要链接：

• arxiv
• 博客
• Colab
• Hugging Face
• Discord
• YouTube（演讲）

简介

欢迎来到 HyenaDNA 仓库！HyenaDNA 是一个长序列基因组基础模型，在高达 100 万个 token 的上下文长度上进行了预训练，并且具有 单核苷酸分辨率。

该仓库目前仍在开发中，但我们非常期待将其交付给研究人员使用，请大家多多包涵 :)

这个仓库最适合希望预训练 HyenaDNA 模型或尝试论文中下游任务的研究人员。不过，最简单的入门方式是使用 HyenaDNA 的 Colab 笔记本，这是一个集成了 Hugging Face 的自包含笔记本。你可以加载预训练权重并在 GenomicBenchmarks 数据集上进行微调。此外，在免费层级下，你还可以对长达 45 万核苷酸的 DNA 序列进行推理并获取嵌入。对于 100 万长度的 DNA 序列，你可以在 Colab 的付费层级中申请 A100 卡，或者在自己的机器上运行该笔记本。

致谢： 代码大部分源自并扩展自 S4 和 Safari。

Discord

正在试用 Discord！也许它能促进大家分享关于如何以不同方式应用 HyenaDNA 的想法和技巧。欢迎大家在那里提问。

Hugging Face 预训练权重

请查看这些权重 :) 它们有不同的模型尺寸和可处理的最大训练序列长度。所有模型均基于单一人类参考基因组（hg38）进行预训练。

• tiny-1k
• tiny-1k-d256
• tiny-16k-d128
• small-32k
• medium-160k
• medium-450k
• large-1m

请参阅每个模型的建议 GPU 要求。

使用这些 Hugging Face 权重有几种方式，各有不同的特点：

1. Colab
2. 本仓库中的 PyTorch Lightning
3. 独立运行

依赖项

首先，我们来安装本仓库所需的依赖项。（如果你熟悉 Docker，可以跳过本节直接前往下方的 Docker 设置。）该仓库基于 PyTorch Lightning（一个训练框架）和 Hydra（一个面向配置的机器学习库）构建。熟悉这些工具将非常有帮助。

• 克隆仓库并进入目录：

git clone --recurse-submodules https://github.com/HazyResearch/hyena-dna.git && cd hyena-dna

• 创建一个 Conda 环境，Python 版本需为 3.8 或更高：

conda create -n hyena-dna python=3.8

• 本仓库使用 PyTorch 1.13 和 CUDA 11.7 开发：

conda install pytorch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 pytorch-cuda=11.7 -c pytorch -c nvidia

• 安装项目所需依赖：

pip install -r requirements.txt

• 安装 Flash Attention，以下 说明 将有所帮助：

cd hyena-dna
git submodule update --init
cd flash-attention
git submodule update --init
pip install -e . --no-build-isolation

• 可选：安装加速用的融合层（耗时较长）：

# 在 flash-attn 目录内
cd csrc/layer_norm && pip install . --no-build-isolation

Dockerfile

如果你熟悉 Docker，我们提供了一个预装所有依赖的镜像，这是最简单、最可靠的方式，但需要一定的 Docker 使用经验。

一个小插曲：你还需要克隆作为主 hyena-dna 仓库子模块的 flash-attn 仓库。这意味着你需要使用 --recurse-submodules 标志，以防你之前克隆时未启用该选项。

# 克隆主仓库及子模块
git clone --recurse-submodules https://github.com/HazyResearch/hyena-dna.git && cd hyena-dna

准备 Docker 容器：

# 在 hyena-dna 仓库中构建镜像（会自动读取 Dockerfile）。请将 $USER_NAME 替换为你自己的 Docker Hub 用户名。
docker build . -t $USER_NAME/hyena-dna

或者，

# 拉取已构建好的镜像（我们的 $USER_NAME 是 hyenadna）
docker pull hyenadna/hyena-dna:latest

# 运行容器：这将为你提供一个带有所有依赖的交互式 shell
docker run --gpus all -it -p80:3000 hyenadna/hyena-dna /bin/bash

更新：

我们实际上还有第二个 Docker 镜像，其中包含了所有的 Nucleotide Transformer 数据集、检查点以及用于复现 HyenaDNA 论文中最佳结果的精确命令和超参数设置。

docker pull hyenadna/hyena-dna-nt6:latest 
docker run --gpus all -it -p80:3000 hyenadna/hyena-dna-nt6 /bin/bash

这将带你进入 /wdr 目录，其中有一个名为 launch_commands_nucleotide_transformer 的文件，包含了针对 18 个 Nucleotide Transformer 数据集的所有启动命令及其相关超参数。

你可能会问，与第一个 Docker 镜像有什么区别呢？其实差别不大，只是部分依赖版本有所不同而已。

快速入门

快速开始使用这个仓库的方法是在一个小规模的基因组数据集上从头训练。让我们先试试看，确认一切设置是否正常。

下面的命令会自动将一个小数据集下载到 data/ 目录中。它使用一个小型的两层 HyenaDNA 模型，并配有一个线性解码器（头部），用于二分类任务。目前，该模型已经比 SotA 高出 7 个百分点（来自 GenomicBenchmarks 的一项任务），但如果我们使用预训练模型，还可以进一步提升性能。

python -m train wandb=null experiment=hg38/genomic_benchmark_scratch

下面我们来详细说明一下：

• -m 允许你以模块方式运行脚本（无需在名称中使用 .py）。
• train 调用主 train.py 脚本，启动所有训练和微调实验。
• wandb=null 表示连接到 WandB，但在快速测试时我将其设置为 null。如果你希望记录实验，可以使用类似 wandb.group=custom_name_here 的配置。
• experiment 传递实验配置，使用位于 configs/experiments/hg38/ 中的 genomic_benchmark_scratch.yaml 文件。
• 你也可以通过命令行传递其他配置，例如 dataset=your_custom_datset_name。不过这方面的内容我们稍后再详细介绍。

加载预训练权重

有两种方法可以使用 Hugging Face 上的预训练权重：

1. Hugging Face 集成（最佳示例），可通过 colab 实现。
2. 在本仓库中使用 PyTorch Lightning：
   • 你可以克隆 Hugging Face 的仓库，并将检查点路径传递给 PyTorch Lightning（实际的 .ckpt 文件来自 Lightning）。
   • 使用的标志是 train.pretrained_model_path=/path/to/ckpt。
   • 启动时需要确保模型配置与 Hugging Face 仓库中的配置一致。这些配置同样可以在 Hugging Face 仓库中找到。

独立代码（也基于 Hugging Face）

实际上我们还有第三种方式，但它只是将 Colab 笔记本的内容复制到本仓库中，保存为 .py 文件（如果你更喜欢这种方式的话）。这种方式集成的是 Hugging Face，而不是 PyTorch Lightning，因此不会提供所有高级功能，但它是一个独立的脚本，意味着更容易移植到你自己的代码库中。此方法假设你已经安装了所有依赖项。

• 参阅 huggingface.py 脚本，了解如何进行推理以及从 Hugging Face 加载预训练权重。
• 还有 standalone_hyenadna.py，其中包含了创建 HyenaDNA 模型所需的所有类。

实验

我们分享了用于在人类参考基因组（HG38）上进行预训练、在多个下游任务上进行微调，以及使用软提示标记和指令微调实现上下文学习变体的训练和数据加载代码。目前，你需要自行下载并预处理数据；我们稍后会分享这些步骤。

总的来说，建议你熟悉 configs/experiments/hg38 中的配置文件，因为我们所有的（示例）实验设置都保存在那里。

在人类参考基因组上进行预训练

第一步是下载人类参考基因组的数据。它由两个文件组成：一个是包含所有序列的 .fasta 文件，另一个是我们使用的区间文件 .bed。

文件结构应如下所示：

data
|-- hg38/
    |-- hg38.ml.fa
    |-- human-sequences.bed

• 将整个人类基因组的 .fa 格式文件下载到 hyena-dna/data/hg38 目录下。整个基因组大约有 24 条染色体，合并为一个文件，每条染色体基本上都是连续的序列。然后下载包含序列区间的 .bed 文件（包含染色体名称、起始位置、结束位置等信息，可用于从 .fasta 文件中提取相应序列）。

mkdir -p data/hg38/
curl https://storage.googleapis.com/basenji_barnyard2/hg38.ml.fa.gz > data/hg38/hg38.ml.fa.gz
curl https://storage.googleapis.com/basenji_barnyard2/sequences_human.bed > data/hg38/human-sequences.bed

启动预训练运行：

python -m train wandb=null experiment=hg38/hg38_hyena model.d_model=128 model.n_layer=2 dataset.batch_size=256 train.global_batch_size=256 dataset.max_length=1024 optimizer.lr=6e-4 trainer.devices=1

下面我们对这条命令做一些解释：

• experiment=hg38/hg38_hyena 传递了使用 Hyena(DNA) 模型的实验配置。
• model.d_model=128 和 model.n_layer=2 分别设置了模型的宽度和深度，这是关键的超参数。
• dataset.max_length=1024 设置了从数据集中采样的最大序列长度，模型的最大序列长度也会根据此值进行设定；或者你也可以手动设置 model.layer.l_max。
• model.d_inner 同样可以手动设置反向瓶颈维度（默认为 4 * d_model）。

还有很多其他可传递和自定义的命令，你可以查看 experiment=hg38/hg38_hyena 配置文件以获取更多细节。

在你自己的数据上进行预训练

要在你自己的数据上进行预训练，你只需要一个 .fasta 文件即可。不需要像我们用于 HG38 那样的 .bed 文件，我们使用 .bed 文件只是为了方便。你可以参考我们的物种分类数据加载器，学习如何构建一个通用的预训练数据加载器，它可以随机选择一条染色体，并从中抽取一段指定长度的序列。

示例预训练数据加载器：

src/dataloaders/datasets/species_dataset.py

通过替换 .fasta 文件为你的数据，并调整一些配置，物种分类数据加载器也可以用于预训练。此外，还需要修改代码以映射你 .fasta 文件中实际存在的染色体，因此你需要深入研究代码及其工作原理。总体而言，使用这个仓库的主要工作就是设置数据加载器和配置文件（虽然耗时，但非常值得）。

注意：如果你计划在自己的数据上进行预训练，请务必正确预处理数据，并确保数据样本符合数据加载器的预期。例如，大小写、未知字符等问题。另外，如果你的序列长度不固定（我们在实验中主要使用固定长度的序列，因为理论上下一个词预测应该能够处理变长序列），那么填充可能会变得显著或成为问题。比如，如果你的序列长度范围是 100 到 32,000，那么长度为 100 的序列将会有很多填充，这时你需要在损失函数中忽略这些填充标记，以避免训练不稳定。填充标记的默认值是 4，你可以在命令行中通过 +task.loss.ignore_index=4 来指定，或者直接修改配置文件中的 task.loss 部分。

基因组基准数据集

GenomicBenchmarks 是一套易于使用的序列级别分类数据集。我们将其作为尝试新方法的良好起点。

示例运行：

python -m train wandb=null experiment=hg38/genomic_benchmark dataset_name=human_enhancers_cohn train.pretrained_model_path=/path/to/ckpt dataset.max_length=500 model.layer.l_max=1024

此命令在其中一个数据集上运行 HyenaDNA 模型，该数据集会自动下载到 data/ 目录中。以下是其他数据集及其统计信息，您也可以将这些信息传递到配置中。configs/dataset/genomic_benchmark 中的配置已设置为正确加载数据集元数据（样本数、类别数等）。

与上述 快速入门 一样，您需要设置要使用的 dataset.max_length 标志，以及 model.layer.l_max，它指示模型您希望使用的最大序列长度。输入序列将被填充至 model.layer.l_max 长度。例如，如果数据样本长度为 500，而 l_max 设置为 1024，则会将 501 补齐至 l_max。

在此微调实验中新增的标志是通过 train.pretrained_model_path=/path/to/ckpt 传入预训练检查点。

该套件包含 8 个数据集，每次选择一个（通过传递 dataset.dataset_name 可自动设置类别数和序列数）。

# 名称                                序列数        类别数     中位数长度    标准差
# dummy_mouse_enhancers_ensembl       1210            2               2381          984.4  
# demo_coding_vs_intergenomic_seqs    100_000         2               200           0
# demo_human_or_worm                  100_000         2               200           0
# human_enhancers_cohn                27791           2               500           0
# human_enhancers_ensembl             154842          2               269           122.6
# human_ensembl_regulatory            289061          3               401           184.3
# human_nontata_promoters             36131           2               251           0
# human_ocr_ensembl                   174756          2               315           108.1

核苷酸 Transformer 数据集

您可以查阅 Nucleotide Transformer 论文附录，了解如何下载和处理这些数据集。

如果您想使用我们用于微调的预训练权重，您需要 Hugging Face 上的 tiny-1k-d256 权重。

更新：或者，您可以花一点时间学习如何使用 Docker，并直接使用我们的 预构建 Docker 镜像，其中包含了 HyenaDNA 论文中结果所用的确切核苷酸 Transformer 数据集/划分、预训练权重和超参数（这是迄今为止最便捷的复现结果的方式）。

示例运行：

# 从头开始训练
python -m train wandb=null experiment=hg38/nucleotide_transformer dataset_name=enhancer dataset.max_length=500 model.layer.l_max=1026

与 GenomicBenchmarks 类似，我们需要从 17 个核苷酸 Transformer 数据集中选择要使用的数据集。

请参阅 configs/dataset/nucleotide_transformer 中的数据集配置以获取更多元数据，以下是一些示例：

字段
名称 最大长度 类别数 样本数 指标

# enhancer 200   2  14968 MCC
# enhancer_types 200   3  14968 MCC
# H3 500   2  13468 MCC
# H3K4me1  500   2  28509 MCC
# H3K4me2  500   2  27614 MCC
# H3K4me3  500   2  33119 MCC
# H3K9ac   500   2  25003 MCC
# H3K14ac  500   2  29743 MCC
# H3K36me3 500   2  31392 MCC
# H3K79me3 500   2  25953 MCC
# H4 500   2  13140 MCC
# H4ac  500   2  30685 MCC
# promoter_all   300   2  53276 F1
# promoter_non_tata 300   2  47759 F1
# promoter_tata  300   2  5517  F1
# splice_sites_acceptor   600   2  19961 F1
# splice_sites_donor   600   2  19775 F1

数据文件结构应如下所示：

data
|-- nucleotide_transformer/
    |-- enhancer/
        |-- all_test_enhancer.fasta
        |-- all_train_enhancer.fasta
    |-- H3/
        |-- H3_test.fasta
        |-- H3_train.fasta
    |-- promoter_tata/
        |-- promoter_tata_test.fasta
        |-- promoter_tata_train.fasta
    |-- ...

上下文学习

我们使用 GenomicBenchmarks 来探索上下文学习 (ICL)。它应该会自动将数据下载到 data/ 目录中。

软提示示例运行：

python -m evals/soft_prompting_genomics

指令微调示例：

python -m evals/instruction_tuned_genomics

染色质谱

您需要参考 DeepSea 和 repo 以获取有关如何下载和预处理的信息。

染色质谱示例运行：

python -m train wandb=null experiment=hg38/chromatin_profile dataset.ref_genome_path=/path/to/fasta/hg38.ml.fa dataset.data_path=/path/to/chromatin_profile dataset.ref_genome_version=hg38

• dataset.ref_genome_path # 人类参考基因组文件路径（输入序列）
• dataset.ref_genome_version # 参考基因组版本（hg38 或 hg19，我们使用 hg38）
• dataset.data_path # 数据集标签文件路径

物种分类

您需要为每个要使用的物种下载 FASTA 文件（只需下载 .zip 文件，数据加载器会自动解压）。可以使用以下命令下载：

# 人类
wget -P human/ -r -nH --cut-dirs=12 --no-parent ftp://ftp.ncbi.nlm.nih.gov/genomes/all/GCA/009/914/755/GCA_009914755.4_T2T-CHM13v2.0/GCA_009914755.4_T2T-CHM13v2.0_assembly_structure/Primary_Assembly/assembled_chromosomes/FASTA/
# 狐猴
wget -P lemur/ -r -nH --cut-dirs=11 --no-parent ftp://ftp.ncbi.nlm.nih.gov/genomes/genbank/vertebrate_mammalian/Lemur_catta/latest_assembly_versions/GCA_020740605.1_mLemCat1.pri/GCA_020740605.1_mLemCat1.pri_assembly_structure/Primary_Assembly/assembled_chromosomes/FASTA/
# 家鼠
wget -P mouse/ -r -nH --cut-dirs=11 --no-parent ftp://ftp.ncbi.nlm.nih.gov/genomes/genbank/vertebrate_mammalian/Mus_musculus/latest_assembly_versions/GCA_921998355.2_A_J_v3/GCA_921998355.2_A_J_v3_assembly_structure/Primary_Assembly/assembled_chromosomes/FASTA/
# 猪
wget -P pig/ -r -nH --cut-dirs=11 --no-parent ftp://ftp.ncbi.nlm.nih.gov/genomes/genbank/vertebrate_mammalian/Sus_scrofa/latest_assembly_versions/GCA_002844635.1_USMARCv1.0/GCA_002844635.1_USMARCv1.0_assembly_structure/Primary_Assembly/assembled_chromosomes/FASTA/

# 河马
wget -P hippo/ -r -nH --cut-dirs=11 --no-parent ftp://ftp.ncbi.nlm.nih.gov/genomes/genbank/vertebrate_mammalian/Hippopotamus_amphibius/latest_assembly_versions/GCA_023065835.1_ASM2306583v1/GCA_023065835.1_ASM2306583v1_assembly_structure/Primary_Assembly/assembled_chromosomes/FASTA/

你的文件夹结构应该如下所示：

data
|-- species/
    |-- chimpanzee/
        |-- chr1.fna
        |-- chr2.fna
        |-- ...
    |-- hippo/
        |-- chr1.fna
        |-- chr2.fna
        |-- ...
    |-- human/
        |-- chr1.fna
        |-- chr2.fna
        |-- ...
    |-- mouse/
        |-- chr1.fna
        |-- chr2.fna
        |-- ...
    |-- orangutan/
        |-- chr1.fna
        |-- chr2.fna
        |-- ...
    |-- 其他物种 ...
|-- ...

示例物种运行：

python -m train wandb=null experiment=hg38/species dataset.species=[human,mouse,hippo,pig,lemur] train.global_batch_size=256 optimizer.lr=6e-5 trainer.devices=1 dataset.batch_size=1 dataset.max_length=1024 dataset.species_dir=/path/to/data/species/ model.layer.l_max=1026 model.d_model=128 model.n_layer=2 trainer.max_epochs=150 decoder.mode=last train.pretrained_model_path=null train.pretrained_model_state_hook=null

让我们来分解其中的一些参数：

• experiment=hg38/species # 本次实验的主要配置
• dataset.species # 你想要的物种列表（并且已经下载了它们的 .fasta 文件）
• decoder.mode=last # 使用最后一个 token 进行分类（而不是默认的池化操作）
• train.pretrained_model_path # 如果使用预训练模型，指向该模型；否则设置为 null
• train.pretrained_model_state_hook=null # 如果使用预训练模型，这将正确加载权重（而非头部）。如果不使用，则设置为 null

更高级的内容如下

设置下游任务（微调）

让我们看看设置下游任务需要什么。

主要组成部分是：

1. 模型权重和模型配置（这些通过顶部的 HuggingFace 提供）
2. 自定义数据集类和数据加载器类
3. experiment、dataset、pipeline、model 的配置。不用担心，我们为每个部分都提供了示例。

再次提供示例运行及启动命令分解：

python -m train wandb=null experiment=hg38/genomic_benchmark

模型配置：

我们在上面讨论了一些模型配置设置。为了方便起见，我们将模型配置放在了实验配置中（可以覆盖默认值），但如果你愿意，也可以将其放在 configs/model 目录下。还有一个单独的层配置位于 configs/model/layer。在这里理解 Hydra 配置系统会很有帮助。

使用超长上下文的标志（梯度检查点）

我们有一个检查点标志，可以在 GPU 上减少约 3 倍的内存占用（以支持更长的序列）。然而，这意味着如果你没有正确设置这些标志，可能会在加载检查点时遇到问题（如果模型是在启用这些标志的情况下预训练的，则应设置为 True；否则设置为 False）。

• model.checkpoint_mixer: True # 设置为 True 以减少内存占用
• model.checkpoint_mlp: True # 设置为 True 以减少内存占用

请注意，如果配置中没有这些选项，而你想在命令行中传递它们，需要在前面加上一个加号，例如：+model.checkpoint_mixer=True

如果你遇到类似以下的错误，即 state_dict 键不匹配，很可能是由于这些标志导致的，因此请尝试切换这些标志的开关：

预训练模型中缺少键！backbone.layers.0.mixer.layer.filter_fn.bias

设置数据集类

这里是一个用于 DNA 下游任务的示例数据集类。

src/dataloaders/datasets/genomic_bench_dataset.py

它基本上是一个标准的 PyTorch 数据集。将数据放在 data/ 目录下，例如 /data/your_custom_dataset_name，以便仓库能够找到它。

这里是一个用于 DNA 下游任务的示例数据加载器。实际上，还需要与 HyenaDNA 仓库进行一些连接。

src/dataloaders/genomic_bench_dataloader.py

请注意，类属性中设置了 _name_ = "genomic_benchmark"。这个名称是我们用来查找它的标识。此外，我们需要将数据加载器文件添加到 __init__.py 中，具体位置在该脚本的顶部，即 src/dataloaders/__init__.py。

我会模仿这个数据加载器文件。它基本上是让 Pytorch Lightning 在后台处理大量数据加载工作的机制。在初始化时传递所需的参数即可创建数据加载器。注意其中的 def setup() 方法，这是实例化数据集类的地方。setup() 会在训练脚本中被调用（稍后会详细介绍）。

这里有 3 个数据加载器函数用于创建训练、验证和测试的数据加载器。（在这个例子中，数据集只使用训练和测试数据加载器。）

创建配置文件

如上所述，主要配置是实验配置，对于我们的示例，位于 configs/experiment/hg38/genomic_benchmark.yaml。

你可以将这些部分中的每一项也视为独立的配置文件，例如 model、task、optimizer 等。你可以在主配置文件中直接编写这些内容，或者在顶部引用它们（作为默认值或覆盖值，两者之间存在细微差异）。

对于新的数据集，我们需要一个新的 dataset 配置文件和一个 pipeline 配置文件。这些配置文件会在实例化时被传递。

目前还没有提到流水线配置，但它用于定义一些不同的内容。请查看内部：

configs/pipeline/genomic_benchmark.yaml

你也应该尝试模仿这个配置文件，它会被引用到实验配置的顶部。我们在这里选择优化器、调度器、数据集名称以及任务类型（通常对于这些下游任务是分类任务，但我们也有其他解码器选项）。不用担心编码器。我们确实使用解码器，它只是一个简单的 MLP，用于将骨干网络映射到我们要预测的类别数量。当你创建数据集类时，它会要求指定输出维度 d_output，表示类别的数量，而解码器会在后台自动获取这一属性，同时也会从 d_model 中获取骨干网络的维度。解码器还可以有其他选项，比如 pool，即对 token 嵌入进行平均；或者 last 或 first，表示我们使用哪个 token 来让 MLP 学习。

如果你想在不同序列长度下进行训练，还需要在几个地方进行更改。主要是数据集配置和模型配置。你可以在实验配置中修改这些内容，也可以分别在独立的 dataset 和 dataloader 配置文件中设置默认值，具体取决于你的需求。

dataset 配置需要设置一个 max_length。

model.layer.l_max 也需要设置一个长度。通常设置为数据集的最大长度加 2。

启动微调实验

# 下游任务示例
python -m train wandb=null experiment=hg38/genomic_benchmark train.pretrained_model_path=<预训练模型权重路径>

数据集会自动下载到 data/ 目录（大概率如此），而且数据集不大，设置时间大约5-10分钟。你只需要从上面提到的 HuggingFace 下载权重，并调整配置以匹配模型设置和你想要使用的数据集序列长度即可。可能需要一些尝试和调整才能达到理想效果，但最终会值得的！

为了更详细地介绍这个 experiment 配置，我们来深入了解一下。该配置对 HyenaDNA 模型（类似 GPT）进行微调。我们重点关注 train 参数。

• remove_test_loader_in_eval: true # 在这个基准测试中没有测试集
  如果验证集和测试集相同，我们可以选择移除额外的测试加载器。

• pretrained_model_strict_load: False # 设置为 False 允许在新模型使用编码器/解码器时继续使用
  将此参数设置为 False 可以更好地兼容加载预训练权重。

要加载骨干网络而不加载头部，需要同时设置以下两个标志：

• pretrained_model_path: /home/workspace/eric/safari-internal/outputs/2023-03-23/07-10-41-239444/checkpoints/val/loss.ckpt
  这是用于指定作为骨干网络的预训练模型路径的地方。

• pretrained_model_state_hook

• _name_: load_backbone
  这是一个自定义钩子函数，它会正确加载骨干网络，并为下游任务添加一个新的 MLP 解码器头部。

• freeze_backbone: false # 如果设置为 False（即对整个模型进行微调），效果似乎更好
  这里可以选择是否冻结骨干网络。

加载微调后的模型

接下来我们将展示如何加载在下游任务上进行过微调的权重（尽管如此，模型仍将继续训练）。

• 权重来自上面提到的 HuggingFace。
• 这些权重适用于一个两层、d_model=128（宽度）、max_length=1024（序列长度）的模型。
• 将这些权重放置在仓库中的某个位置，通常我们会将其放在 outputs/ 目录下。

现在我们需要做的主要事情是更新配置文件中的相应参数。

# 微调后模型配置文件路径
safari-internal/configs/experiment/hg38/genomic_benchmark_load_finetuned_model.yaml

对于这个配置文件，选择你要用来训练的数据集 dataset.dataset_name，这里我们使用 human_nontata_promoters，因为上述权重就是在这个数据集上微调得到的。

接下来，你需要更新 train.pretrained_model_path 参数，将其指向你存放权重的具体路径。

现在我们可以用以下命令启动训练：

python -m train wandb=null experiment=hg38/genomic_benchmark_load_finetuned_model

这将运行主脚本 src/train.py。

让我们指出 train.py 脚本中几个关键的位置，因为代码中各部分的调用关系有些复杂。

• 权重加载是在 train.py 的 def load_state_dict() 函数中完成的。实际上，它会调用一个自定义的状态钩子来优雅地加载（位于 src/models/sequence/long_conv_lm.py 中的 load_backbone() 函数）。

• 前向传播是在 train.py 中的 SequenceLightning 模块内的 def forward() 函数中完成的，但实际上，它会调用 self.task.forward()，而后者才是真正调用模型的部分。也就是说，你需要进入 src/tasks/tasks.py，找到 class LMTask 及其 def forward() 函数。在这里你会看到对模型的实际调用。需要注意的是，解码器头部（一个用于分类的单层 MLP）与主模型的骨干网络（特征提取器）是分开的。

序列长度预热回调

我们实现了一个基于回调的序列长度预热调度器，它会在训练过程中分阶段逐步增加序列长度。基本上，脚本会检查当前处于哪个 epoch 和“阶段”，并根据该阶段的参数更新数据集和数据加载器。目前，你需要手动在配置文件中指定各个阶段，示例配置文件如下，相关部分位于配置文件底部，也在此处列出：

configs/experiment/hg38/hg38_hyena_seqlen_warmup_reload.yaml

注意事项： 你需要提前确保每个阶段的批次大小和序列长度都能适配显存。

为了让数据加载器与序列长度预热机制兼容，你需要实现一个接口，即 init_datasets()。它的具体实现如下所示：

关键点：

要使用这个回调，我们可以使用上面提供的示例配置文件：configs/experiment/hg38/hg38_hyena_seqlen_warmup_reload.yaml。

你需要手动设计各个阶段，即希望在哪些 epoch 逐步增加序列长度（同时降低批次大小）。请注意，每个阶段的 epochs 表示该阶段持续的时间（不是累计的）。

callbacks:
  seqlen_warmup_reload:
    # epochs 表示在该阶段运行多长时间（非累计！）
    # 这只是一个示例
    stage_params:
      - epochs: 2  # 表示在第 0 和第 1 个 epoch 运行此阶段
        seq_len: 1024
        batch_size: 256  # 后台中 grad accum = 1，因为 train.global_batch_size=256
      - epochs: 2  # 在第 2 和第 3 个 epoch 运行
        seq_len: 2048
        batch_size: 128
      - epochs: 2  # 在第 4 和第 5 个 epoch 运行
        seq_len: 4096
        batch_size: 64
      - epochs: 2  # 在第 6 和第 7 个 epoch 运行
        seq_len: 8192
        batch_size: 32
      - epochs: 4  # 在第 8、9、10 和第 11 个 epoch 运行
        seq_len: 16_384
        batch_size: 16
      - epochs: 4  # 在第 12、13、14 和第 15 个 epoch 运行
        seq_len: 32_768
        batch_size: 8

此外，在命令行中还有一些重要的参数需要注意：

在示例配置文件中，请注意以下内容：

• train.global_batch_size 切勿忘记设置！它会控制 accumulate_grad_batches，以保持每个阶段的学习率一致。例如，通常设置为 256 或 128（对于非常长的序列，可能为 64）。
• dataset.batch_size 现在指的是测试集的批次大小（或最终的序列长度和批次大小）。测试集始终保持不变。
• dataset.max_length 现在指的是测试集的序列长度和最大长度。测试集始终保持不变。
• model.layer.l_max 需要设置为最高的序列长度加 2（即测试集的大小）。

注意事项：

训练数据集会在训练过程中变化，但测试集始终固定不变。测试集的序列长度和批次大小通过常规方式在启动命令中设置，即 dataset.batch_size 和 dataset.。

从预训练模型获取 logits

这里有一个简单的 脚本，可以用来从预训练模型获取 logits。

这个过程并非自动化，因此你需要手动从 HF 下载权重，并将其保存在本地某个位置。要获取 logits，还需要模型的头部。

与 Huggingface 的区别在于：这个脚本旨在方便地获取嵌入表示，而不涉及模型头部。目前我们还没有使用 logits 的实际场景，因此如果你需要它们，还需要额外的步骤。

实验性功能

1. 我们实现了一个实验性的双向HyenaDNA模型。我们最近在GenomicBenchmarks数据集上进行了一项消融实验，其中我们从头开始训练，即未使用BERT的掩码语言建模进行预训练。我们将该模型与标准的因果HyenaDNA进行了对比，结果表明因果版本表现更好。然而，仍有一些用户非常希望使用双向HyenaDNA，因此我们提供了一种实现方式——当然，实现双向性的方法还有很多。

关于具体实现，我们只是简单地调整了FFT卷积中的填充方式。更多细节请参阅src/models/sequence/hyena.py脚本（例如搜索“bidirectional”）。

要使用双向模式，只需在启动时传入标志model.bidirectional=True即可，非常简单！

需要注意的是，当前代码库仅支持在下游任务上从头开始进行双向训练，即不支持掩码语言模型的预训练。使用双向性进行因果预训练并无意义，请谨慎使用！

2. 对于下游任务，我们新增了一个选项，允许仅对掩码标记进行处理或取平均。我们已将此功能更新到GenomicBenchmarks和Nuc Trans数据集的相关类中，具体实现方式请参考这些数据集类。使用方法如下：

• 您还需要设置正确的配置参数。请参考相关实验配置文件，例如/src/configs/experiment/hg38/genomic_benchmark.yaml，特别注意将dataset.return_mask=True和dataset.padding_side=right这两个参数启用。
• 您还需在实验配置中定义一个新的任务类型，名为masked_multiclass。以上设置的作用在于正确地将掩码传递给模型，这与之前的处理方式有所不同。

实践中，对于序列较短且填充较少的任务，我们发现这一改进效果并不显著。但如果您的序列包含大量填充，则该功能将带来明显帮助。不过，在论文中，我们并未采用此功能，默认使用左侧填充。

变更日志 / 更新内容：

• 在Hugging Face上增加了更多权重文件。
• 发布了包含核苷酸Transformer数据集、权重及精确超参数的Docker镜像，以便复现实验结果。
• 新增了一个实验性的双向选项。详情请参阅实验性功能。
• 我们还为下游任务添加了传递掩码并忽略填充标记的选项。详情请参阅实验性功能。
• 增加了一些关于如何使用您自己的数据进行预训练的提示。
• 提供了从预训练模型获取logits的示例。

引用

如果您觉得我们的工作有所帮助，欢迎引用！

@article{nguyen2023hyenadna,
      title={HyenaDNA: 高分辨率单核苷酸级长距离基因组序列建模}, 
      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}
}

HyenaDNA 快速上手指南

HyenaDNA 是一个长序列基因组基础模型，支持高达 100 万 token 的上下文长度和单核苷酸分辨率。本指南将帮助您快速搭建环境并运行模型。

1. 环境准备

系统要求

• 操作系统: Linux (推荐)
• Python: 3.8+
• GPU: 支持 CUDA 的 NVIDIA 显卡（预训练大模型建议 A100，小模型推理可用免费 tier Colab 或普通显卡）
• CUDA: 11.7 (官方开发环境)

前置依赖

• Git (需支持 submodule)
• Conda (推荐用于环境管理)
• Docker (可选，最简便的安装方式)

2. 安装步骤

您可以选择 Conda 手动安装 或 Docker 镜像 两种方式。

方式一：Conda 手动安装（推荐开发者）

1. 克隆仓库

   git clone --recurse-submodules https://github.com/HazyResearch/hyena-dna.git && cd hyena-dna
   

2. 创建 Conda 环境

   conda create -n hyena-dna python=3.8
   conda activate hyena-dna
   

3. 安装 PyTorch (CUDA 11.7) 注：国内用户可使用清华源加速

   conda install pytorch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 pytorch-cuda=11.7 -c pytorch -c nvidia
   # 或使用清华源
   # conda install pytorch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 pytorch-cuda=11.7 -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/nvidia
   

4. 安装基础依赖

   pip install -r requirements.txt
   

5. 安装 Flash Attention (核心组件)

   cd hyena-dna
   git submodule update --init
   cd flash-attention
   git submodule update --init
   pip install -e . --no-build-isolation
   

6. (可选) 安装融合层以加速

   cd csrc/layer_norm && pip install . --no-build-isolation
   

方式二：Docker 安装（最简便）

如果您熟悉 Docker，可直接拉取预构建镜像，包含所有依赖。

1. 拉取并运行镜像

   docker pull hyenadna/hyena-dna:latest
   docker run --gpus all -it -p80:3000 hyenadna/hyena-dna /bin/bash
   

   注：如需复现论文中 Nucleotide Transformer 数据集的结果，可使用 hyenadna/hyena-dna-nt6:latest 镜像。

3. 基本使用

快速入门：从头训练小模型

以下命令将自动下载一个小规模基因组数据集，并使用一个小型 HyenaDNA 模型（2 层）进行二分类任务训练。这是验证环境是否配置正确的最佳方式。

python -m train wandb=null experiment=hg38/genomic_benchmark_scratch

参数说明：

• wandb=null: 禁用权重与偏差跟踪，适合本地快速测试。
• experiment=hg38/genomic_benchmark_scratch: 加载位于 configs/experiments/hg38/ 下的配置文件。

使用预训练模型 (Hugging Face)

HyenaDNA 提供了多种尺寸的预训练权重（从 1k 到 1M 序列长度）。

方法 A：使用 Colab (最简单)

直接访问官方提供的集成 Notebook，无需本地配置即可加载权重并在 GenomicBenchmarks 数据集上进行微调或推理。

• HyenaDNA Colab Notebook

方法 B：本地加载权重 (PyTorch Lightning)

1. 从 Hugging Face 下载对应的 .ckpt 文件。
2. 确保本地配置文件中的模型参数与预训练模型一致。
3. 运行训练/微调命令并指定路径：

   python -m train train.pretrained_model_path=/path/to/your/model.ckpt experiment=your_config
   

方法 C：独立脚本推理

如果您希望将模型集成到自己的代码库中，可参考仓库内的独立脚本：

• huggingface.py: 演示如何加载 HF 权重并进行推理。
• standalone_hyenadna.py: 包含构建 HyenaDNA 模型所需的所有类定义。

预训练自定义数据

若需在自定义数据上预训练，只需准备 .fasta 文件。参考 src/dataloaders/datasets/species_dataset.py 修改数据加载器，并调整 configs/experiments/ 下的配置文件即可启动训练。