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

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

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

使用🤗的Caduceus

我们已将预训练的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数据增强。

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

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

from transformers import AutoModelForMaskedLM, AutoTokenizer

# 查看hub上的`Caduceus`集合页面以获取可用模型列表。
model_name = "kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)

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

from transformers import AutoConfig, AutoModelForMaskedLM

# 在此处添加任何配置覆盖，详情请参见hub上的`config.json`文件。
config_overrides = {}
# 查看hub上的`Caduceus`集合页面以获取可用模型列表。
config = AutoConfig.from_pretrained(
 "kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16",
 **config_overrides,
)
model = AutoModelForMaskedLM.from_config(config)

在此仓库中开始使用

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

conda env create -f caduceus_env.yml

激活环境。

conda activate caduceus_env

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

mkdir outputs
mkdir watch_folder

复现实验

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

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

（数据下载说明复制自HyenaDNA仓库）

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

文件结构应如下所示：

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

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

mkdir -p data/hg38/
curl https://storage.googleapis.com/basenji_barnyard2/hg38.ml.fa.gz > data/hg38/hg38.ml.fa.gz
gunzip data/hg38/hg38.ml.fa.gz  # 解压fasta文件
curl https://storage.googleapis.com/basenji_barnyard2/sequences_human.bed > data/hg38/human-sequences.bed

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

python -m train \
  experiment=hg38/hg38 \
  callbacks.model_checkpoint_every_n_steps.every_n_train_steps=500 \
  dataset.max_length=1024 \
  dataset.batch_size=1024 \
  dataset.mlm=true \
  dataset.mlm_probability=0.15 \
  dataset.rc_aug=false \
  model=caduceus \
  model.config.d_model=128 \
  model.config.n_layer=4 \
  model.config.bidirectional=true \
  model.config.bidirectional_strategy=add \
  model.config.bidirectional_weight_tie=true \
  model.config.rcps=true \
  optimizer.lr="8e-3" \
  train.global_batch_size=1024 \
  trainer.max_steps=10000 \
  +trainer.val_check_interval=10000 \
  wandb=null

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

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

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

cd slurm_scripts
sbatch run_pretrain_caduceus.sh

GenomicBenchmarks

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

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

python -m train \
    experiment=hg38/genomic_benchmark \
    callbacks.model_checkpoint_every_n_steps.every_n_train_steps=5000 \
    dataset.dataset_name="dummy_mouse_enhancers_ensembl" \
    dataset.train_val_split_seed=1 \
    dataset.batch_size=128 \
    dataset.rc_aug=false \
    +dataset.conjoin_train=false \
    +dataset.conjoin_test=false \
    loader.num_workers=2 \
    model=caduceus \
    model._name_=dna_embedding_caduceus \
    +model.config_path="<path to model_config.json>" \
    +model.conjoin_test=false \
    +decoder.conjoin_train=true \
    +decoder.conjoin_test=false \
    optimizer.lr="1e-3" \
    trainer.max_epochs=10 \
    train.pretrained_model_path="<path to .ckpt file>" \
    wandb=null

这个示例运行将在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基线。

核苷酸转换器数据集

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

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

python -m train \
    experiment=hg38/nucleotide_transformer \
    callbacks.model_checkpoint_every_n_steps.every_n_train_steps=5000 \
    dataset.dataset_name="${task}" \
    dataset.train_val_split_seed=${seed} \
    dataset.batch_size=${batch_size} \
    dataset.rc_aug="${rc_aug}" \
    +dataset.conjoin_test="${CONJOIN_TEST}" \
    loader.num_workers=2 \
    model._name_=dna_embedding_caduceus \
    +model.config_path="<model_config.json的路径>" \
    +model.conjoin_test=false \
    +decoder.conjoin_train=true \
    +decoder.conjoin_test=false \
    optimizer.lr="1e-3" \
    trainer.max_epochs=10 \
    train.pretrained_model_path="<.ckpt文件的路径>" \
    trainer.max_epochs=20 \
    wandb=null

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

eQTL SNP变异效应预测

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

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

torchrun \
    --standalone \
    --nnodes=1 \
    --nproc-per-node=8 \
    vep_embeddings.py \
      --num_workers=2 \
      --seq_len=131072  \
      --bp_per_token=1  \
      --embed_dump_batch_size=1 \
      --name="caduceus-ps_downstream-seqlen=131k"  \
      --model_name_or_path="kuleshov-group/caduceus-ps_seqlen-131k_d_model-256_n_layer-16" \
      --rcps

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

致谢

本仓库改编自HyenaDNA仓库，并利用了其中定义的大量训练、数据加载和日志记录基础设施。 HyenaDNA最初源自S4和Safari仓库。

我们要感谢Evan Trop和InstaDeep团队就Nucleotide Transformer排行榜 和长程基准任务进行的有益讨论。

最后，我们要感谢MosaicML为部分预训练实验提供计算资源。