Caduceus ☤：双向等变长程DNA序列建模

[博客] | [arXiv] | [HuggingFace 🤗]

本仓库包含复现论文"Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling"Schiff et al. (2024)中结果的代码。

使用🤗的Caduceus

<a name="HF"></a> 我们已将预训练的Caduceus模型上传到Huggingface hub。可用的模型有：

Caduceus-Ph: kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16
- 在长度为131k的序列上训练，模型大小为256，16层。
- 训练50k步，批量大小为8。
- 使用反向互补(RC)数据增强进行训练。
Caduceus-PS: kuleshov-group/caduceus-ps_seqlen-131k_d_model-256_n_layer-16
- 在长度为131k的序列上训练，模型大小为256，16层。
- 训练50k步，批量大小为8。
- 模型具有RC等变性，因此不需要RC数据增强。

您可以直接在训练脚本中使用预训练模型，或修改初始化模型的配置。

要使用预训练模型进行掩码语言建模，请使用以下代码片段：

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

或者，您可以从头开始实例化模型，以在自己的数据上进行训练，如下所示：

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

在此仓库中开始使用

首先，创建一个包含所需依赖项的conda环境。

[object Object],[object Object],[object Object]

激活环境。

[object Object]

创建以下目录以存储保存的模型和slurm日志：

[object Object],[object Object],[object Object],[object Object]

复现实验

以下我们描述了复现论文中实验所需的步骤。整个过程中，运行实验的主要入口点是train.py脚本。我们还在slurm_scripts/目录中提供了用于启动预训练和下游微调实验的示例slurm脚本。

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

<a name="pretraining"></a> （数据下载说明复制自HyenaDNA仓库）

首先，下载人类参考基因组数据。它由2个文件组成，1个包含所有序列（.fasta文件），另一个包含我们使用的区间（.bed文件）。

文件结构应如下所示：

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

将fasta（.fa格式）文件（整个人类基因组）下载到./data/hg38。整个基因组中约有24条染色体（合并为1个文件），每条染色体基本上是一个连续序列。然后下载带有序列区间的.bed文件（包含染色体名称、开始、结束、分割，然后允许您从fasta文件中检索）。

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

使用命令行启动预训练运行

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

或者，如果使用安装了slurm的集群，请调整以下脚本：

slurm_scripts
|-- run_pretrain_caduceus.sh
|-- run_pretrain_hyena.sh
|-- run_pretrain_mamba.sh

并将训练作为批处理作业运行：

[object Object],[object Object]

GenomicBenchmarks

Grešová et al. (2023)中提出的GenomicBenchmarks由8个分类任务组成。

我们可以使用以下示例命令在其中一个任务上启动下游微调运行：

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

这个示例运行将在dummy_mouse_enhancers_ensembl任务上微调预训练的Caduceus-PS模型。请注意此处相对于上面的预训练命令存在一些额外的参数：

model.config_path包含预训练期间保存的模型配置路径。这将保存到预训练实验的运行目录中。
train.pretrained_model_path包含预训练模型检查点的路径。
dataset.conjoin_train决定数据集在下游微调训练期间是返回单个序列（dataset.conjoin_train=false）还是沿dim=-1连接序列及其反向互补。
dataset.conjoin_test与上面相同，但用于推理（例如，验证/测试）。
decoder.conjoin_train决定预测头（在Genomics Benchmark的情况下是平均池化和线性投影）在下游微调训练期间是否期望输入张量的形状为(batch_size, seq_len, d_model)或(batch_size, seq_len, d_model, 2)。当设置为true时，解码器在input[..., 0]和input[..., 1]上运行，并对结果进行平均以产生最终预测。
decoder.conjoin_test与上面相同，但用于推理（例如，验证/测试）。

请注意，此基准测试仅包含每个任务的训练和测试分割。因此，为了进行更有原则的评估，我们使用dataset.train_val_split_seed参数将训练数据随机分割为训练集和验证集（90/10）。我们对验证指标（准确率）进行早停，并重复5个随机种子。

与预训练一样，我们也可以使用提供的run_genomic_benchmark.sh脚本将微调运行作为批处理作业启动。我们还提供了一个辅助shell脚本wrapper_run_genomics.sh，可用于并行启动多个微调运行。

最后，run_genomics_benchmark_cnn.sh脚本可用于在下游任务上从头训练CNN基线。

核苷酸转换器数据集

<a name="nucleotidetransformer"></a> Nucleotide Transformer任务套件由Dalla-Torre等人（2023）提出。数据可在HuggingFace上获取：InstaDeepAI/nucleotide_transformer_downstream_tasks。

我们可以使用以下示例命令在其中一个任务上启动下游微调：

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

我们还可以以批处理作业方式启动（详情请参见run_nucleotide_transformer.sh和wrapper_run_nucleotide_transformer.sh）。

eQTL SNP变异效应预测

<a name="vep"></a> 这项任务来自最近由Kao等人，2023提出的长程基准（LRB）。数据可在HuggingFace上获取：InstaDeepAI/genomics-long-range-benchmark。对于这项任务，我们将模型拟合到DNA语言模型的预训练和冻结嵌入上。因此，为了进行评估，我们分两步进行：

第1步：提取嵌入从预训练模型中：运行vep_embeddings.py脚本从预训练模型中提取嵌入。请参见以下示例：

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

--rcps标志用于指示模型是反向互补等变的。使用其他模型时，请将此标志设置为false，即--no-rcps。为加快此步骤的速度，该脚本利用了torch分布式数据并行。

请参考slurm_scripts/dump_vep_embeddings.sh中提供的slurm脚本，以批处理作业方式启动此步骤。

第2步：将SVM模型拟合到嵌入使用此笔记本：vep_svm.ipynb。

引用

如果您觉得我们的工作有用，请使用以下方式引用我们的论文：

@article{schiff2024caduceus,
  title={Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling},
  author={Schiff, Yair and Kao, Chia-Hsiang and Gokaslan, Aaron and Dao, Tri and Gu, Albert and Kuleshov, Volodymyr},
  journal={arXiv preprint arXiv:2403.03234},
  year={2024}
}