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.
What this notebook downloads. Everything comes from
bolerodata: the publishedBorzoi10Mmodel (its config carries the dataset registry; its checkpoint is the starting point), the frozen Flashzoi backbone, the published per-pseudobulk AP-1 score file, the 32 bp stores of two small datasets (about 4 GB together) and, at the end, the published AP-1 score model. A real run fine-tunes on the whole Bolero-10M corpus for 50 epochs on several GPUs. Here we run a few training steps on two datasets to watch the loop, then load the published AP-1 checkpoint and inspect what makes it a score model.
Setup¶
In [1]:
Copied!
import copy
import logging
import os
import joblib
import pandas as pd
from bolero import init
from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer
from bolerodata import MODELS, get_flashzoi_checkpoint, metadata
os.environ["WANDB_MODE"] = "disabled" # demo run: no logging
# quiet the per-file download checks
logging.getLogger("httpx").setLevel(logging.WARNING)
init(num_cpus=16, object_store_memory_gb=32)
import copy
import logging
import os
import joblib
import pandas as pd
from bolero import init
from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer
from bolerodata import MODELS, get_flashzoi_checkpoint, metadata
os.environ["WANDB_MODE"] = "disabled" # demo run: no logging
# quiet the per-file download checks
logging.getLogger("httpx").setLevel(logging.WARNING)
init(num_cpus=16, object_store_memory_gb=32)
2026-10-27 16:12:02,767 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 (Borzoi10M); 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]:
Copied!
SCORE_NAME = "AP1"
FOLD = 0
# The published multi-dataset ATAC model we fine-tune from (its config carries the registry).
pretrained = MODELS["Borzoi10M"]
PRETRAIN_CONFIG = pretrained.model_config
PRETRAIN_LORA_CKPT = str(pretrained.ckpt_path) # downloaded once
# The chromVAR score, per pseudobulk, per dataset (built from the cross-dataset chromVAR).
SHARED_DATA_PATHS = str(metadata.get_chromvar_score_path(SCORE_NAME))
# Frozen Flashzoi backbone (same base as every Bolero model).
BASE_CKPT = str(get_flashzoi_checkpoint(FOLD))
# The published AP-1 score model: the result of a real run, loaded at the end of the page.
SCORE_MODEL = MODELS["Bolero10M+AP1+260104"]
SAVENAME = f"Bolero10M_{SCORE_NAME}_atac_fold{FOLD}-demo"
SCORE_NAME = "AP1"
FOLD = 0
# The published multi-dataset ATAC model we fine-tune from (its config carries the registry).
pretrained = MODELS["Borzoi10M"]
PRETRAIN_CONFIG = pretrained.model_config
PRETRAIN_LORA_CKPT = str(pretrained.ckpt_path) # downloaded once
# The chromVAR score, per pseudobulk, per dataset (built from the cross-dataset chromVAR).
SHARED_DATA_PATHS = str(metadata.get_chromvar_score_path(SCORE_NAME))
# Frozen Flashzoi backbone (same base as every Bolero model).
BASE_CKPT = str(get_flashzoi_checkpoint(FOLD))
# The published AP-1 score model: the result of a real run, loaded at the end of the page.
SCORE_MODEL = MODELS["Bolero10M+AP1+260104"]
SAVENAME = f"Bolero10M_{SCORE_NAME}_atac_fold{FOLD}-demo"
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
2026-10-27 16:12:05 - WARNING - Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
The dataset registry comes from the pretrained model¶
The dataset_records (which datasets, their 32 bp stores + pseudobulks, their genome) come
straight out of the pretrained model's config — a real run trains on exactly that corpus.
The trainer reads every registered dataset's 32 bp store, and the whole corpus is about
1.6 TB. So this demo keeps two small datasets, bonobo and macaque cerebellum
development. It keeps the full genome list (GENOMES) all the same: the model's genome
embedding has one slot per genome in the corpus, and the pretrained weights expect all of
them.
In [3]:
Copied!
full_records = PRETRAIN_CONFIG["dataset_records"]
GENOMES = sorted({rec["genome"] for rec in full_records.values()})
print(f"{len(full_records)} datasets in the corpus, {len(GENOMES)} genomes: {GENOMES}")
# demo: two small datasets from the registry (a real run uses all of them)
DEMO_DATASETS = ["CBDevEvoBonobo", "CBDevEvoMacaque"]
dataset_records = {k: copy.deepcopy(full_records[k]) for k in DEMO_DATASETS}
for rec in dataset_records.values():
rec["dataset_sample_weight"] = 1
pd.DataFrame(dataset_records).T[["genome", "data_path"]]
full_records = PRETRAIN_CONFIG["dataset_records"]
GENOMES = sorted({rec["genome"] for rec in full_records.values()})
print(f"{len(full_records)} datasets in the corpus, {len(GENOMES)} genomes: {GENOMES}")
# demo: two small datasets from the registry (a real run uses all of them)
DEMO_DATASETS = ["CBDevEvoBonobo", "CBDevEvoMacaque"]
dataset_records = {k: copy.deepcopy(full_records[k]) for k in DEMO_DATASETS}
for rec in dataset_records.values():
rec["dataset_sample_weight"] = 1
pd.DataFrame(dataset_records).T[["genome", "data_path"]]
42 datasets in the corpus, 7 genomes: ['calJac4', 'hg38', 'mCalJac1.pat.X', 'mm10', 'monDom5.split', 'panPan3', 'rheMac10']
Out[3]:
| genome | data_path | |
|---|---|---|
| CBDevEvoBonobo | panPan3 | standard/metacell_dataset/parquet/CBDevEvoBono... |
| CBDevEvoMacaque | rheMac10 | standard/metacell_dataset/parquet/CBDevEvoMaca... |
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 — 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]:
Copied!
score_data = joblib.load(SHARED_DATA_PATHS)
print("datasets with scores:", len(score_data))
print("demo datasets scored:", all(k in score_data for k in DEMO_DATASETS))
example = score_data[DEMO_DATASETS[0]]
print(DEMO_DATASETS[0], "-> DataFrame", example.shape, "(pseudobulks x score)")
example.head()
score_data = joblib.load(SHARED_DATA_PATHS)
print("datasets with scores:", len(score_data))
print("demo datasets scored:", all(k in score_data for k in DEMO_DATASETS))
example = score_data[DEMO_DATASETS[0]]
print(DEMO_DATASETS[0], "-> DataFrame", example.shape, "(pseudobulks x score)")
example.head()
datasets with scores: 51 demo datasets scored: True CBDevEvoBonobo -> DataFrame (70, 1) (pseudobulks x score)
Out[4]:
| MA0478.2 | |
|---|---|
| group0-pseudobulk0 | -0.106890 |
| group0-pseudobulk1 | 0.669986 |
| group0-pseudobulk2 | 0.669986 |
| group0-pseudobulk3 | 0.677362 |
| group0-pseudobulk4 | 0.592079 |
Build the training config¶
The schedule below is a demo: one short epoch of a few steps. The real run used
max_epochs=50, lr_total_steps=2_000_000, save_state_every_n_epoch=20 and the default
batch_size of 4; the four demo lines are simply dropped for it.
In [5]:
Copied!
config = {
# run identity + outputs
"output_dir": "model",
"savename": SAVENAME,
"wandb_project": "BoleroScore",
"wandb_job_type": "lora",
"wandb_name": SAVENAME,
# data: the demo subset of the multi-dataset corpus
"dataset_records": dataset_records,
"fold_split_id": FOLD,
"_genomes": GENOMES, # the whole corpus' genomes, so the genome embedding matches
# 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,
# schedule: real run max_epochs=50, lr_total_steps=2_000_000, save_state_every_n_epoch=20
"max_epochs": 1,
"lr_total_steps": 2_000_000,
"save_state_every_n_epoch": 20,
# demo only: a few small steps (drop these four lines for a real run)
"warmup_steps": 10,
"train_batches": 20,
"val_batches": 5,
"batch_size": 1,
}
config = MultiBorzoiLoRATrainer.make_config(**config)
config = {
# run identity + outputs
"output_dir": "model",
"savename": SAVENAME,
"wandb_project": "BoleroScore",
"wandb_job_type": "lora",
"wandb_name": SAVENAME,
# data: the demo subset of the multi-dataset corpus
"dataset_records": dataset_records,
"fold_split_id": FOLD,
"_genomes": GENOMES, # the whole corpus' genomes, so the genome embedding matches
# 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,
# schedule: real run max_epochs=50, lr_total_steps=2_000_000, save_state_every_n_epoch=20
"max_epochs": 1,
"lr_total_steps": 2_000_000,
"save_state_every_n_epoch": 20,
# demo only: a few small steps (drop these four lines for a real run)
"warmup_steps": 10,
"train_batches": 20,
"val_batches": 5,
"batch_size": 1,
}
config = MultiBorzoiLoRATrainer.make_config(**config)
Build the model and fine-tune for a few steps¶
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. trainer.train() then builds the network,
loads the frozen base + the pretrained multi-dataset LoRA, and runs the fit loop.
The Borzoi10M checkpoint holds a conditioning encoder for every dataset of the corpus;
the model here has encoders for the two demo datasets only, so the load skips the others.
In [6]:
Copied!
# builds the corpus; registers the score dimension
trainer = MultiBorzoiLoRATrainer(config)
trainer.train() # builds model; loads base + pretrained multi-dataset LoRA; fits
# builds the corpus; registers the score dimension
trainer = MultiBorzoiLoRATrainer(config)
trainer.train() # builds model; loads base + pretrained multi-dataset LoRA; fits
Using single GPU for training Create BorzoiMultiDataset with config: dataset_records: 2 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: 1 output_prefix: pseudobulk shared_data_paths: /tmp/bolero-docs/bdhome/hf/datasets--arcinstitute--bolero-data/snapshots/29da7170e6b27b32751557824002f79769f60245/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: /tmp/bolero-docs/bdhome/flashzoi/fold_0.renamed.pt
Model does not have gene_count_output_head, skip Total model parameters 1201495140, trainable parameters 1024135972 Config contains lora_checkpoint_path, will load weights and perform fine-tuning. Loading LoRA model weights from: /tmp/bolero-docs/bdhome/hf/models--arcinstitute--bolero-models/snapshots/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS/Borzoi10M/checkpoint.pt
Input shape: x: torch.Size([1, 4, 524288]) embedding: torch.Size([1, 768]) signal: torch.Size([1, 1, 16384])
=========================================================================================================================================================== Layer (type (var_name)) Input Shape Output Shape Param # =========================================================================================================================================================== BorzoiLoRAMulti (BorzoiLoRAMulti) -- [1, 1, 16352] 437,254 ├─ConvDna (conv_dna) [1, 4, 524288] [1, 512, 262144] -- │ └─ConditionalLoRAConv (conv_layer) [1, 4, 524288] [1, 512, 524288] 31,232 │ │ └─EmbeddingMLP (lora_A_module) -- [1, 30, 60] 726,280 │ │ └─EmbeddingMLP (lora_B_module) -- [1, 512, 30] 4,211,200 │ │ └─Dropout (lora_dropout) [1, 512, 524288] [1, 512, 524288] -- │ └─MaxPool1d (max_pool) [1, 512, 524288] [1, 512, 262144] -- ├─SequentialwithArgs (res_tower) [1, 512, 262144] [1, 1280, 16384] -- │ └─ConvBlock (0) [1, 512, 262144] [1, 608, 262144] -- │ │ └─BatchNorm1d (norm) [1, 512, 262144] [1, 512, 262144] (1,024) │ │ └─GELU (activation) [1, 512, 262144] [1, 512, 262144] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 512, 262144] [1, 608, 262144] 26,509,728 │ └─MaxPool1d (1) [1, 608, 262144] [1, 608, 131072] -- │ └─ConvBlock (2) [1, 608, 131072] [1, 736, 131072] -- │ │ └─BatchNorm1d (norm) [1, 608, 131072] [1, 608, 131072] (1,216) │ │ └─GELU (activation) [1, 608, 131072] [1, 608, 131072] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 608, 131072] [1, 736, 131072] 31,878,496 │ └─MaxPool1d (3) [1, 736, 131072] [1, 736, 65536] -- │ └─ConvBlock (4) [1, 736, 65536] [1, 896, 65536] -- │ │ └─BatchNorm1d (norm) [1, 736, 65536] [1, 736, 65536] (1,472) │ │ └─GELU (activation) [1, 736, 65536] [1, 736, 65536] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 736, 65536] [1, 896, 65536] 39,106,496 │ └─MaxPool1d (5) [1, 896, 65536] [1, 896, 32768] -- │ └─ConvBlock (6) [1, 896, 32768] [1, 1056, 32768] -- │ │ └─BatchNorm1d (norm) [1, 896, 32768] [1, 896, 32768] (1,792) │ │ └─GELU (activation) [1, 896, 32768] [1, 896, 32768] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 896, 32768] [1, 1056, 32768] 47,941,856 │ └─MaxPool1d (7) [1, 1056, 32768] [1, 1056, 16384] -- │ └─ConvBlock (8) [1, 1056, 16384] [1, 1280, 16384] -- │ │ └─BatchNorm1d (norm) [1, 1056, 16384] [1, 1056, 16384] (2,112) │ │ └─GELU (activation) [1, 1056, 16384] [1, 1056, 16384] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 1056, 16384] [1, 1280, 16384] 57,864,640 ├─Sequential (signal_encoder) [1, 1, 16384] [1, 1280, 16384] -- │ └─Conv1d (0) [1, 1, 16384] [1, 128, 16384] 512 │ └─GroupNorm (1) [1, 128, 16384] [1, 128, 16384] 256 │ └─SiLU (2) [1, 128, 16384] [1, 128, 16384] -- │ └─Conv1d (3) [1, 128, 16384] [1, 512, 16384] 197,120 │ └─GroupNorm (4) [1, 512, 16384] [1, 512, 16384] 1,024 │ └─SiLU (5) [1, 512, 16384] [1, 512, 16384] -- │ └─Conv1d (6) [1, 512, 16384] [1, 1280, 16384] 1,967,360 │ └─GroupNorm (7) [1, 1280, 16384] [1, 1280, 16384] 2,560 ├─SequentialwithArgs (unet1) [1, 1280, 16384] [1, 1536, 8192] -- │ └─MaxPool1d (0) [1, 1280, 16384] [1, 1280, 8192] -- │ └─ConvBlock (1) [1, 1280, 8192] [1, 1536, 8192] -- │ │ └─BatchNorm1d (norm) [1, 1280, 8192] [1, 1280, 8192] (2,560) │ │ └─GELU (activation) [1, 1280, 8192] [1, 1280, 8192] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 1280, 8192] [1, 1536, 8192] 71,545,856 ├─MaxPool1d (_max_pool) [1, 1536, 8192] [1, 1536, 4096] -- ├─ConvBlock (horizontal_conv0) [1, 1280, 16384] [1, 1536, 16384] -- │ └─BatchNorm1d (norm) [1, 1280, 16384] [1, 1280, 16384] (2,560) │ └─GELU (activation) [1, 1280, 16384] [1, 1280, 16384] -- │ └─ConditionalLoRAConv (conv_layer) [1, 1280, 16384] [1, 1536, 16384] 1,967,616 │ │ └─EmbeddingMLP (lora_A_module) [1, 768] [1, 30, 1280] 10,132,480 │ │ └─EmbeddingMLP (lora_B_module) [1, 768] [1, 1536, 30] 12,106,240 │ │ └─Dropout (lora_dropout) [1, 1536, 16384] [1, 1536, 16384] -- ├─ConvBlock (horizontal_conv1) [1, 1536, 8192] [1, 1536, 8192] -- │ └─BatchNorm1d (norm) [1, 1536, 8192] [1, 1536, 8192] (3,072) │ └─GELU (activation) [1, 1536, 8192] [1, 1536, 8192] -- │ └─ConditionalLoRAConv (conv_layer) [1, 1536, 8192] [1, 1536, 8192] 2,360,832 │ │ └─EmbeddingMLP (lora_A_module) [1, 768] [1, 30, 1536] 12,106,240 │ │ └─EmbeddingMLP (lora_B_module) [1, 768] [1, 1536, 30] 12,106,240 │ │ └─Dropout (lora_dropout) [1, 1536, 8192] [1, 1536, 8192] -- ├─SequentialwithArgs (transformer) [1, 4096, 1536] [1, 4096, 1536] -- │ └─TransformerLayer (0) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (1) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (2) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (3) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (4) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (5) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (6) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 │ └─TransformerLayer (7) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,870,912 ├─SequentialwithArgs (upsampling_unet1) [1, 1536, 4096] [1, 1536, 8192] -- │ └─ConvBlock (0) [1, 1536, 4096] [1, 1536, 4096] -- │ │ └─BatchNorm1d (norm) [1, 1536, 4096] [1, 1536, 4096] (3,072) │ │ └─GELU (activation) [1, 1536, 4096] [1, 1536, 4096] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 1536, 4096] [1, 1536, 4096] 26,573,312 │ └─Upsample (1) [1, 1536, 4096] [1, 1536, 8192] -- ├─ConvBlock (separable1) [1, 1536, 8192] [1, 1536, 8192] -- │ └─Identity (norm) [1, 1536, 8192] [1, 1536, 8192] -- │ └─Identity (activation) [1, 1536, 8192] [1, 1536, 8192] -- │ └─Sequential (conv_layer) -- -- -- │ │ └─ConditionalLoRAConv (0) [1, 1536, 8192] [1, 1536, 8192] 12,397,658 │ │ └─ConditionalLoRAConv (1) [1, 1536, 8192] [1, 1536, 8192] 10,783,232 ├─SequentialwithArgs (upsampling_unet0) [1, 1536, 8192] [1, 1536, 16384] -- │ └─ConvBlock (0) [1, 1536, 8192] [1, 1536, 8192] -- │ │ └─BatchNorm1d (norm) [1, 1536, 8192] [1, 1536, 8192] (3,072) │ │ └─GELU (activation) [1, 1536, 8192] [1, 1536, 8192] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 1536, 8192] [1, 1536, 8192] 26,573,312 │ └─Upsample (1) [1, 1536, 8192] [1, 1536, 16384] -- ├─ConvBlock (separable0) [1, 1536, 16384] [1, 1536, 16384] -- │ └─Identity (norm) [1, 1536, 16384] [1, 1536, 16384] -- │ └─Identity (activation) [1, 1536, 16384] [1, 1536, 16384] -- │ └─Sequential (conv_layer) -- -- -- │ │ └─ConditionalLoRAConv (0) [1, 1536, 16384] [1, 1536, 16384] 12,397,658 │ │ └─ConditionalLoRAConv (1) [1, 1536, 16384] [1, 1536, 16384] 10,783,232 ├─SequentialwithArgs (final_joined_convs) [1, 1536, 16384] [1, 1920, 16384] -- │ └─ConvBlock (0) [1, 1536, 16384] [1, 1920, 16384] -- │ │ └─BatchNorm1d (norm) [1, 1536, 16384] [1, 1536, 16384] (3,072) │ │ └─GELU (activation) [1, 1536, 16384] [1, 1536, 16384] -- │ │ └─ConditionalLoRAConv (conv_layer) [1, 1536, 16384] [1, 1920, 16384] 56,769,920 │ └─Dropout (1) [1, 1920, 16384] [1, 1920, 16384] -- │ └─GELU (2) [1, 1920, 16384] [1, 1920, 16384] -- ├─TargetLengthCrop (crop) [1, 1920, 16384] [1, 1920, 16352] -- ├─OutputHead (final_output_head) [1, 1920, 16352] [1, 1, 16352] -- │ └─ConditionalLoRAConv (0) [1, 1920, 16352] [1, 1, 16352] 1,921 │ │ └─EmbeddingMLP (lora_A_module) -- [1, 1, 1920] 757,120 │ │ └─EmbeddingMLP (lora_B_module) -- [1, 1, 1] 263,937 │ │ └─Dropout (lora_dropout) [1, 1, 16352] [1, 1, 16352] -- │ └─Softplus (1) [1, 1, 16352] [1, 1, 16352] -- =========================================================================================================================================================== Total params: 1,201,495,140 Trainable params: 1,024,135,972 Non-trainable params: 177,359,168 Total mult-adds (Units.GIGABYTES): 36.49 =========================================================================================================================================================== Input size (MB): 8.46 Forward/backward pass size (MB): 2786.34 Params size (MB): 4087.35 Estimated Total Size (MB): 6882.15 =========================================================================================================================================================== EMA is not used. Current epoch: 1, max epochs: 1. Get dataloader with train mode
2026-10-27 16:12:34,543 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-11-59_384604_3552552/logs/ray-data
2026-10-27 16:12:34,544 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition]
2026-10-27 16:12:36,185 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-11-59_384604_3552552/logs/ray-data
2026-10-27 16:12:36,185 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> TaskPoolMapOperator[MapBatches(ReverseComplement)] -> LimitOperator[limit=22]
Data loader kwargs {'batch_size': 1, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}
- (Training) 0 [4/20] Ave Loss: 0.756; Ave Corr: 0.567 Last grad norm: 0.5554 Last batch Loss: 0.831
- (Training) 0 [9/20] Ave Loss: 0.745; Ave Corr: 0.542 Last grad norm: 0.3553 Last batch Loss: 0.394
- (Training) 0 [14/20] Ave Loss: 0.694; Ave Corr: 0.569 Last grad norm: 0.2808 Last batch Loss: 0.721
- (Training) 0 [19/20] Ave Loss: 0.719; Ave Corr: 0.548 Last grad norm: 0.2174 Last batch Loss: 0.785
2026-10-27 16:13:10,630 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-11-59_384604_3552552/logs/ray-data
2026-10-27 16:13:10,630 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition]
Get dataloader with eval mode
2026-10-27 16:13:10,802 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-11-59_384604_3552552/logs/ray-data
2026-10-27 16:13:10,803 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> LimitOperator[limit=6]
Data loader kwargs {'batch_size': 3, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}
- (Training) 1 Loss: 0.721; Mean Corr. : 0.544; Learning rate 4.999875000625004e-05. - (Validation) 1 Loss: 0.512; Mean Corr. : 0.308. Saving best checkpoint...
Load and update state dict from checkpoint file: /tmp/bolero-docs/bolero/docs/tutorials/score_model/model/Bolero10M_AP1_atac_fold0-demo.lora.best_checkpoint.pt
2026-10-27 16:14:13,067 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-11-59_384604_3552552/logs/ray-data
2026-10-27 16:14:13,068 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition]
Get dataloader with eval mode
2026-10-27 16:14:13,794 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-11-59_384604_3552552/logs/ray-data
2026-10-27 16:14:13,795 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> LimitOperator[limit=16]
Data loader kwargs {'batch_size': 3, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}
- (Validation) 1 [4/5] Mean Loss: 1.620; Mean Track Corr: 0.810;
Load the trained AP-1 score model¶
The demo trained for a few steps only. The published AP-1 score model
(Bolero10M+AP1+260104) is the result of the real run; we load it into the same network and
inspect it. Its checkpoint, too, holds encoders for datasets that are not in the demo
registry, and the load skips them.
In [7]:
Copied!
# unwrap DataParallel if present
model = getattr(trainer.model, "module", trainer.model)
model.load_checkpoint_from_path(SCORE_MODEL.ckpt_path, strict=False)
print("loaded trained checkpoint:", SCORE_MODEL.model_key)
# unwrap DataParallel if present
model = getattr(trainer.model, "module", trainer.model)
model.load_checkpoint_from_path(SCORE_MODEL.ckpt_path, strict=False)
print("loaded trained checkpoint:", SCORE_MODEL.model_key)
Loading LoRA model weights from: /tmp/bolero-docs/bdhome/hf/models--arcinstitute--bolero-models/snapshots/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS/Bolero10M+AP1+260104/checkpoint.pt
loaded trained checkpoint: Bolero10M+AP1+260104
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 [8]:
Copied!
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 byshared_data_paths/__shared_data__.
Model size¶
In [9]:
Copied!
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,201,495,140 trainable params: 1,024,135,972 (85.2%)
Real training vs. this demo¶
The demo ran a few steps on two datasets. A real run registers the whole corpus from the
pretrained config, drops the four demo lines, and fine-tunes for max_epochs=50 from the
pretrained multi-dataset LoRA, then saves model/{SAVENAME}.lora.best_checkpoint.pt.
The next page runs the trained AP-1 model forward.