13. Train the score model (Bolero-Score)¶

The multi-dataset model predicts cell-state-specific accessibility from DNA + a cell-state embedding. The score model ("Bolero-Score") adds one more conditioning input: a scalar chromVAR TF-activity score (e.g. AP-1, computed on the previous page). Given DNA + cell state + how active AP-1 is in that state, it predicts the accessibility that matches — letting you later dial the score and read out the TF's sequence-level effect.

Concretely it is the trained Bolero-10M ATAC model fine-tuned with the AP-1 score wired into the conditioning. This page shows how it is configured and what one extra input changes in the model.

This notebook runs in the score-model training directory, not in place under docs/. Like the multi-dataset page it needs the full Bolero-10M corpus, the pretrained multi-dataset checkpoint, and the per-pseudobulk chromVAR scores. A real run is heavy (multi-GPU, many epochs), so we do not call trainer.train() — we build the model exactly as training would, load the finished AP-1 checkpoint, and inspect what makes it a score model.

Setup¶

In [1]:

import json

import joblib

from bolero import init
from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer

import os

os.environ["WANDB_MODE"] = (
    "disabled"  # demo only builds/loads a model; no training, no logging
)
init(num_cpus=24, object_store_memory_gb=32)

import json import joblib from bolero import init from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer import os os.environ["WANDB_MODE"] = ( "disabled" # demo only builds/loads a model; no training, no logging ) init(num_cpus=24, object_store_memory_gb=32)

2026-07-17 12:14:35,031	INFO worker.py:1781 -- Started a local Ray instance.

Configuration — the score model is a fine-tune of the multi-dataset model¶

Two things distinguish this config from the plain multi-dataset one:

• shared_data_paths — the per-pseudobulk chromVAR score file (AP1.chromvar_score.scale.joblib.gz). This is what makes it a score model.
• lora_checkpoint_path — the finished Bolero-10M ATAC LoRA weights; training starts from them and only learns to use the new score input, rather than training from scratch.

Everything else (the dataset_records registry, the frozen Flashzoi base, signal_model) is inherited from the multi-dataset model.

In [2]:

SCORE_NAME = "AP1"
FOLD = 0

# The finished multi-dataset ATAC model we fine-tune from (config carries the dataset registry).
PRETRAIN_CONFIG = "/home/hanliu/data/250907-MultiDatasetModel/Bolero10M/model/251103-Bolero10M_atac_fold0-final.config.json"
PRETRAIN_LORA_CKPT = "/home/hanliu/data/250907-MultiDatasetModel/Bolero10M/model/251103-Bolero10M_atac_fold0-final.lora.best_checkpoint.pt"

# The chromVAR score, per pseudobulk, per dataset (built from the cross-dataset chromVAR).
SHARED_DATA_PATHS = f"/home/hanliu/data/260102-ScoreModel/score_data/{SCORE_NAME}.chromvar_score.scale.joblib.gz"

# Frozen Flashzoi backbone (same base as every Bolero model).
BASE_CKPT = (
    f"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_{FOLD}.renamed.pt"
)

# The trained score-model checkpoint we load for the demo (result of a real run).
SAVENAME = f"260102-Bolero10M_{SCORE_NAME}_atac_fold{FOLD}"
FINAL_CKPT = f"model/{SAVENAME}.lora.best_checkpoint.pt"

# genomes in the corpus (kept identical to init/final stages so the genome embedding matches)
GENOMES = [
    "calJac4",
    "hg38",
    "mCalJac1.pat.X",
    "mm10",
    "monDom5.split",
    "panPan3",
    "rheMac10",
]

SCORE_NAME = "AP1" FOLD = 0 # The finished multi-dataset ATAC model we fine-tune from (config carries the dataset registry). PRETRAIN_CONFIG = "/home/hanliu/data/250907-MultiDatasetModel/Bolero10M/model/251103-Bolero10M_atac_fold0-final.config.json" PRETRAIN_LORA_CKPT = "/home/hanliu/data/250907-MultiDatasetModel/Bolero10M/model/251103-Bolero10M_atac_fold0-final.lora.best_checkpoint.pt" # The chromVAR score, per pseudobulk, per dataset (built from the cross-dataset chromVAR). SHARED_DATA_PATHS = f"/home/hanliu/data/260102-ScoreModel/score_data/{SCORE_NAME}.chromvar_score.scale.joblib.gz" # Frozen Flashzoi backbone (same base as every Bolero model). BASE_CKPT = ( f"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_{FOLD}.renamed.pt" ) # The trained score-model checkpoint we load for the demo (result of a real run). SAVENAME = f"260102-Bolero10M_{SCORE_NAME}_atac_fold{FOLD}" FINAL_CKPT = f"model/{SAVENAME}.lora.best_checkpoint.pt" # genomes in the corpus (kept identical to init/final stages so the genome embedding matches) GENOMES = [ "calJac4", "hg38", "mCalJac1.pat.X", "mm10", "monDom5.split", "panPan3", "rheMac10", ]

The dataset registry comes from the pretrained model¶

We read the dataset_records (which datasets, their parquet stores + pseudobulks, their genome) straight out of the pretrained model's config — the score model trains on exactly the same corpus.

In [3]:

with open(PRETRAIN_CONFIG) as f:
    pretrain_cfg = json.load(f)

dataset_records = pretrain_cfg["dataset_records"]
for rec in dataset_records.values():
    rec["dataset_sample_weight"] = 1

print(f"{len(dataset_records)} datasets in the corpus")

with open(PRETRAIN_CONFIG) as f: pretrain_cfg = json.load(f) dataset_records = pretrain_cfg["dataset_records"] for rec in dataset_records.values(): rec["dataset_sample_weight"] = 1 print(f"{len(dataset_records)} datasets in the corpus")

45 datasets in the corpus

The conditioning score¶

shared_data_paths is a dict {dataset: DataFrame(pseudobulk × score)}. For AP-1 the score is the standardized cross-dataset chromVAR deviation of the AP-1 motif (MA0478.2) — one number per pseudobulk, comparable across datasets. That single column is the model's new __shared_data__ input (so shared_data_dim = 1).

In [4]:

score_data = joblib.load(SHARED_DATA_PATHS)
example_key = next(iter(score_data))
print("datasets with scores:", len(score_data))
print(
    "example dataset:",
    example_key,
    "-> DataFrame",
    score_data[example_key].shape,
    "(pseudobulks x score)",
)
score_data[example_key].head()

score_data = joblib.load(SHARED_DATA_PATHS) example_key = next(iter(score_data)) print("datasets with scores:", len(score_data)) print( "example dataset:", example_key, "-> DataFrame", score_data[example_key].shape, "(pseudobulks x score)", ) score_data[example_key].head()

datasets with scores: 51
example dataset: AIBSDev -> DataFrame (2595, 1) (pseudobulks x score)

Out[4]:

	MA0478.2
group0-pseudobulk0	-1.313889
group0-pseudobulk1	-1.352320
group0-pseudobulk2	-1.195127
group0-pseudobulk3	-1.340319
group0-pseudobulk4	-1.269483

Build the training config¶

In [5]:

config = {
    # run identity + outputs
    "output_dir": "model",
    "savename": SAVENAME,
    "wandb_project": "260102-BoleroScore",
    "wandb_job_type": "lora",
    "wandb_name": SAVENAME,
    # data: the same corpus as the multi-dataset model
    "dataset_records": dataset_records,
    "fold_split_id": FOLD,
    "_genomes": GENOMES,
    # model
    "base_checkpoint_path": BASE_CKPT,
    "signal_model": True,
    "lora_checkpoint_path": PRETRAIN_LORA_CKPT,  # fine-tune from the trained multi-dataset LoRA
    "cond_module_kwargs": {"norm_type": "layer"},
    # >>> the score conditioning <<<
    "shared_data_paths": SHARED_DATA_PATHS,
    # real training schedule
    "max_epochs": 50,
    "lr_total_steps": 2_000_000,
    "save_state_every_n_epoch": 20,
}
config = MultiBorzoiLoRATrainer.make_config(**config)

config = { # run identity + outputs "output_dir": "model", "savename": SAVENAME, "wandb_project": "260102-BoleroScore", "wandb_job_type": "lora", "wandb_name": SAVENAME, # data: the same corpus as the multi-dataset model "dataset_records": dataset_records, "fold_split_id": FOLD, "_genomes": GENOMES, # model "base_checkpoint_path": BASE_CKPT, "signal_model": True, "lora_checkpoint_path": PRETRAIN_LORA_CKPT, # fine-tune from the trained multi-dataset LoRA "cond_module_kwargs": {"norm_type": "layer"}, # >>> the score conditioning <<< "shared_data_paths": SHARED_DATA_PATHS, # real training schedule "max_epochs": 50, "lr_total_steps": 2_000_000, "save_state_every_n_epoch": 20, } config = MultiBorzoiLoRATrainer.make_config(**config)

Build the model, then load the trained checkpoint (no training)¶

Constructing the trainer builds the multi-dataset dataset eagerly — and because shared_data_paths is set, it also registers the extra __shared_data__ conditioning dimension and merges it into the model config. _setup_model() then builds the network and loads the frozen base + the pretrained multi-dataset LoRA. A real run would now call trainer.train(); instead we load the finished AP-1 checkpoint.

In [6]:

trainer = MultiBorzoiLoRATrainer(
    config
)  # builds the corpus; registers the score dimension
trainer._setup_model(
    print_model=False
)  # builds model; loads base + pretrained multi-dataset LoRA

model = getattr(
    trainer.model, "module", trainer.model
)  # unwrap DataParallel if present
model.load_checkpoint_from_path(
    FINAL_CKPT, strict=False
)  # <- the trained AP-1 score model
print("loaded trained checkpoint:", FINAL_CKPT)

trainer = MultiBorzoiLoRATrainer( config ) # builds the corpus; registers the score dimension trainer._setup_model( print_model=False ) # builds model; loads base + pretrained multi-dataset LoRA model = getattr( trainer.model, "module", trainer.model ) # unwrap DataParallel if present model.load_checkpoint_from_path( FINAL_CKPT, strict=False ) # <- the trained AP-1 score model print("loaded trained checkpoint:", FINAL_CKPT)

Using single GPU for training
Create BorzoiMultiDataset with config:
  dataset_records: 45 items <class 'dict'>
  dna_window: 524288
  train_region_step_sample: True
  use_regions: borzoi
  n_pseudobulks: 10
  paired_mode: ensemble
  reverse_complement: True
  max_jitter: 3
  batch_size: 4
  output_prefix: pseudobulk
  shared_data_paths: /home/hanliu/data/260102-ScoreModel/score_data/AP1.chromvar_score.scale.joblib.gz
  _genomes: 7 items <class 'list'>

Updated cond_module_kwargs in config with dataset specific information.

Setting up model from config

Loading base model weights from: /home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_0.renamed.pt

Model does not have gene_count_output_head, skip
Total model parameters 1210743216, trainable parameters 1033384048
Config contains lora_checkpoint_path, will load weights and perform fine-tuning.
Loading LoRA model weights from: /home/hanliu/data/250907-MultiDatasetModel/Bolero10M/model/251103-Bolero10M_atac_fold0-final.lora.best_checkpoint.pt

Loading LoRA model weights from: model/260102-Bolero10M_AP1_atac_fold0.lora.best_checkpoint.pt

loaded trained checkpoint: model/260102-Bolero10M_AP1_atac_fold0.lora.best_checkpoint.pt

What's special — the score joins the conditioning vector¶

Bolero's conditioning module turns each example's inputs into one embedding vector that generates every conditional-LoRA weight. For a multi-dataset model that vector is built from a dataset-specific part (the cell-state embedding, translated per dataset) plus a shared part. The shared part is where cross-dataset inputs live — and the score model simply adds __shared_data__ (the AP-1 score) alongside __genome__.

So the tell of a score model is one extra entry in dataset_shared_dims: a plain multi-dataset model has only {'__genome__': …}; here it also has '__shared_data__': 1.

In [7]:

cem = model.cond_emb_module
print("cond_emb_module:", type(cem).__name__)
print("dataset_shared_dims:", cem.dataset_shared_dims)  # note '__shared_data__': 1
print("shared conditioning vector width (output_dim):", cem.output_dim)

has_score = "__shared_data__" in cem.dataset_shared_dims
print(
    "\nconditions on a chromVAR score?",
    has_score,
    "| score width:",
    cem.dataset_shared_dims.get("__shared_data__"),
)

cem = model.cond_emb_module print("cond_emb_module:", type(cem).__name__) print("dataset_shared_dims:", cem.dataset_shared_dims) # note '__shared_data__': 1 print("shared conditioning vector width (output_dim):", cem.output_dim) has_score = "__shared_data__" in cem.dataset_shared_dims print( "\nconditions on a chromVAR score?", has_score, "| score width:", cem.dataset_shared_dims.get("__shared_data__"), )

cond_emb_module: ConditionEmbeddingModuleMulti
dataset_shared_dims: {'__genome__': 7, '__shared_data__': 1}
shared conditioning vector width (output_dim): 768

conditions on a chromVAR score? True | score width: 1

At training time each pseudobulk's AP-1 score is attached to its batch as __shared_data__, encoded (a small MLP), pooled with the genome embedding into the shared vector, and concatenated with the dataset-specific embedding — so the score ends up influencing every conditional-LoRA layer's generated weights. That is how a scalar TF activity reshapes the whole DNA→accessibility mapping.

Note on signal_model=True. This turns on the reference-coverage input branch and is on for all multi-dataset Bolero models — it is not the score. The chromVAR score is carried solely by shared_data_paths / __shared_data__.

Model size¶

In [8]:

total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"total params    : {total:,}")
print(f"trainable params: {trainable:,}  ({trainable/total:.1%})")

total = sum(p.numel() for p in model.parameters()) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"total params : {total:,}") print(f"trainable params: {trainable:,} ({trainable/total:.1%})")

total params    : 1,210,743,216
trainable params: 1,033,384,048  (85.4%)

Real training vs. this demo¶

We loaded a finished checkpoint. A real run fine-tunes for max_epochs=50 from the pretrained multi-dataset LoRA with the schedule above, then saves model/{SAVENAME}.lora.best_checkpoint.pt. To just watch the loop run, add the smoke-test overrides (max_epochs=3, warmup_steps=10, train_batches=50, val_batches=10) and call trainer.train().

The next page runs this trained AP-1 model forward.