Caduceus: Bi-directional Equivariant Long-range DNA Sequence Modeling

To reproduce pre-training experiments, follow these steps:

1. Prepare Data

Download the Human Reference Genome (hg38) files into ./data/hg38/. You need the .fasta file (the genome sequence) and the .bed file (sequence intervals).

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
curl https://storage.googleapis.com/basenji_barnyard2/sequences_human.bed > data/hg38/human-sequences.bed

2. Run Pre-training

Use the train.py script via the python -m train command. Key arguments for pre-training include dataset.mlm=true and model=caduceus.

Alternatively, if using a Slurm cluster, you can use the provided scripts in the slurm_scripts/ directory:

cd slurm_scripts
sbatch run_pretrain_caduceus.sh

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

Fine-tune a pre-trained Caduceus model on classification tasks from the GenomicBenchmarks suite.

Key Configuration Parameters:

• model.config_path: Path to the model config saved during pre-training.
• train.pretrained_model_path: Path to the pre-trained .ckpt file.
• dataset.conjoin_train / dataset.conjoin_test: Determines if the dataset returns a single sequence or the concatenation of a sequence and its reverse complement along dim=-1.
• decoder.conjoin_train / decoder.conjoin_test: Determines if the prediction head expects input of shape (batch_size, seq_len, d_model) or (batch_size, seq_len, d_model, 2). If true, the decoder averages results from input[..., 0] and input[..., 1].

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=256 \
    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

The BaseTask class is the foundation for defining model training objectives. It manages the loss function, arbitrary metrics, the forward pass, and optional encoder/decoder modules that interface with the dataset and model.

When implementing a custom task, you can leverage BaseTask to handle the orchestration of data through an encoder, a backbone model, and a decoder. It also provides built-in support for torchmetrics and custom metric functions.

Key components managed by BaseTask:

• loss: The primary loss function.
• metrics: A collection of output and loss-based metrics.
• encoder: (Optional) Interfaces dataset inputs with the model.
• decoder: (Optional) Interfaces model outputs with dataset targets.
• torchmetrics: A collection of torchmetrics for standardized evaluation (train, val, test).

from src.tasks.tasks import BaseTask

# Example conceptual usage
task = BaseTask(
    dataset=my_dataset,
    model=my_model,
    loss='cross_entropy',
    metrics=['accuracy', 'perplexity'],
    torchmetrics=['AUROC', 'F1']
)

After running the SVM evaluation, results are typically saved to a CSV file. You can visualize the AUROC performance across different models and TSS distance buckets using seab sns.catplot.

To ensure clean plots:

1. Rename model names using a mapping dictionary (e.g., model_name_replacement).
2. Group by model_name, Distance to TSS, use_tissue, and C to calculate the mean AUROC across seeds.
3. Select the best hyperparameter (C) for each model/bucket combination.
4. Use sns.catplot with kind="bar" and errorbar="sd" to show performance and standard deviation.

sns.set_style("whitegrid")
g = sns.catplot(
    data=df_plot,
    x="model_name",
    y="AUROC",
    col="Distance to TSS",
    hue="Distance to TSS",
    kind="bar",
    errorbar="sd",
    height=12,
    aspect=1,
    dodge=False,
    order=list(model_name_replacement.values()),
)
# ... (formatting and saving) ...
g.savefig(osp.join(PATH_TO_OUTPUTS, "SVM_results.png"))

The workflow for evaluating model embeddings via Support Vector Machine (SVM) involves:

1. Loading Data: Load train_embeds_combined.pt and test_embeds_combined.pt using fsspec and torch.load.
2. Preprocessing: Apply dataset_nan_filter and dataset_tss_filter to clean and segment the data.
3. Feature Engineering:
   • If conjoin_train or conjoin_test is True, average the forward and reverse complement embeddings: X = (X_fwd + X_rc) / 2.
   • Optionally concatenate tissue_embed to the feature vector.
4. Training: Use a sklearn.pipeline.make_pipeline containing StandardScaler and SVC.
5. Evaluation: Calculate the roc_auc_score on the test set.

Note: For large datasets, the training set is subsampled to 5000 samples using np.random.choice to speed up fitting.

svm_clf = make_pipeline(
    StandardScaler(),
    SVC(C=C, random_state=seed),
)

# ... (data preparation logic) ...

mask = np.random.choice(len(X), size=5000, replace= 5000 > len(X) )
if use_tissue: 
    X_train = X_with_tissue[mask]
else: 
    X_train = X[mask]
y_train = y[mask]

svm_clf.fit(X_train, y_train)
svm_y_pred = svm_clf.predict(X_test)
svm_aucroc = roc_auc_score(y_test, svm_y_pred)