08. Add the gene-expression task (multi-dataset + gene)¶

The multi-dataset ATAC model predicts cell-state–specific chromatin accessibility across the whole corpus. Bolero's full model also predicts transcript abundance. This page shows the final stage: take the trained ATAC model and continue training with a gene-count head and per-dataset RNA data, so one model outputs both accessibility tracks and gene expression.

Like the previous page, this notebook builds the model from the published registry and runs a short demo on two of its datasets: ChromiumPBMC, which has matched transcriptomes, and Li2023Science+Inh, which has none. The full gene task trains on 41 datasets, and its training data is the whole corpus plus the RNA data. The notebook loads the published gene checkpoint (Borzoi10M-Gene) to inspect what the gene task adds, then trains a few steps.

What changes from the ATAC page:

• output_head_type="gene_count" — adds a Decima-style gene-count head (and a gene-mask conv) on top of the shared trunk.
• use_regions="borzoi_gene" — training regions are centered on genes rather than tiled across the genome.
• Per-dataset RNA data — each dataset with matched transcriptomes gets a gene_data_path (log1p-CPM pseudobulk expression) and a deg_list (highly variable genes) attached to its record.
• lora_checkpoint_path = the trained ATAC-final checkpoint (published as Borzoi10M) — the gene task continues from there.
• _genomes is passed explicitly so the genome embedding stays identical to the ATAC model.

Setup¶

In [1]:

import copy
import os

import pandas as pd
import torch

from bolero import init
from bolero.tl.model.borzoi.model_lora import filter_dataset_encoders
from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer
from bolerodata import DATASETS, MODELS, get_flashzoi_checkpoint

# log offline, as on the previous pages; for a real run, use your own W&B project
os.environ["WANDB_MODE"] = "offline"
init(num_cpus=16, object_store_memory_gb=48)

import copy import os import pandas as pd import torch from bolero import init from bolero.tl.model.borzoi.model_lora import filter_dataset_encoders from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer from bolerodata import DATASETS, MODELS, get_flashzoi_checkpoint # log offline, as on the previous pages; for a real run, use your own W&B project os.environ["WANDB_MODE"] = "offline" init(num_cpus=16, object_store_memory_gb=48)

2026-10-27 16:25:02,403	INFO worker.py:1781 -- Started a local Ray instance.

In [2]:

# --- Configuration ---------------------------------------------------------------
FOLD = 0
# the two hg38 datasets the demo trains on: one with RNA, one ATAC-only
DEMO_DATASETS = ["ChromiumPBMC", "Li2023Science+Inh"]

# the previous page's published ATAC model: its registry, and its checkpoint to continue from
ATAC_MODEL = MODELS["Borzoi10M"]
ATAC_FINAL_CKPT = str(ATAC_MODEL.ckpt_path)
# the published gene model we load for the demo
GENE_MODEL = MODELS["Borzoi10M-Gene"]
GENE_CKPT = GENE_MODEL.ckpt_path
# frozen Flashzoi backbone, as on the previous pages
BASE_CKPT = str(get_flashzoi_checkpoint(FOLD))

# genomes, in the exact order used by the ATAC model (keeps the genome embedding consistent)
GENOMES = [
    "calJac4",
    "hg38",
    "mCalJac1.pat.X",
    "mm10",
    "monDom5.split",
    "panPan3",
    "rheMac10",
]
assert GENOMES == GENE_MODEL.model_config["_genomes"]

# --- Configuration --------------------------------------------------------------- FOLD = 0 # the two hg38 datasets the demo trains on: one with RNA, one ATAC-only DEMO_DATASETS = ["ChromiumPBMC", "Li2023Science+Inh"] # the previous page's published ATAC model: its registry, and its checkpoint to continue from ATAC_MODEL = MODELS["Borzoi10M"] ATAC_FINAL_CKPT = str(ATAC_MODEL.ckpt_path) # the published gene model we load for the demo GENE_MODEL = MODELS["Borzoi10M-Gene"] GENE_CKPT = GENE_MODEL.ckpt_path # frozen Flashzoi backbone, as on the previous pages BASE_CKPT = str(get_flashzoi_checkpoint(FOLD)) # genomes, in the exact order used by the ATAC model (keeps the genome embedding consistent) GENOMES = [ "calJac4", "hg38", "mCalJac1.pat.X", "mm10", "monDom5.split", "panPan3", "rheMac10", ] assert GENOMES == GENE_MODEL.model_config["_genomes"]

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/arcinstitute/bolero-models/resolve/main/BorzoiSignalMultiDS/Borzoi10M/config.json "HTTP/1.1 307 Temporary Redirect"

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:25: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.

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/arcinstitute/bolero-models/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS%2FBorzoi10M%2Fconfig.json "HTTP/1.1 200 OK"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/arcinstitute/bolero-models/resolve/main/BorzoiSignalMultiDS/Borzoi10M/checkpoint.pt "HTTP/1.1 302 Found"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/arcinstitute/bolero-models/resolve/main/BorzoiSignalMultiDS/Borzoi10M-Gene/config.json "HTTP/1.1 307 Temporary Redirect"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/arcinstitute/bolero-models/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS%2FBorzoi10M-Gene%2Fconfig.json "HTTP/1.1 200 OK"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/arcinstitute/bolero-models/resolve/main/BorzoiSignalMultiDS/Borzoi10M-Gene/checkpoint.pt "HTTP/1.1 302 Found"

Attach RNA data to the dataset registry¶

We start from the same dataset_records registry as the ATAC model, then annotate each dataset with its gene data. A dataset with matched transcriptomes gets a gene_data_path (its expression table) and a deg_list, both published per dataset and fetched through bolerodata; a dataset with no RNA still trains on ATAC, but gets a deg_list from a related dataset so its regions are gene-centered and it is not over-represented. Species with no transcriptomes at all (here CBDevEvoOpossum) are dropped.

The full gene task annotates all 41 remaining datasets. The demo keeps its two datasets first, so it downloads only their files; the loop is the same for any subset.

In [3]:

full_records = ATAC_MODEL.model_config["dataset_records"]
print(len(full_records) - 1, "datasets in the full gene task (all but CBDevEvoOpossum)")

# the demo's datasets
dataset_records = {k: copy.deepcopy(full_records[k]) for k in DEMO_DATASETS}
for rec in dataset_records.values():
    rec["dataset_sample_weight"] = 1


def published_rna(dataset_name):
    """The dataset's published RNA files, or None if it has no matched transcriptome."""
    dataset = DATASETS[dataset_name]
    try:
        gene_data_path = (
            dataset.get_pseudobulk_gene_data_path()
        )  # log1p-CPM pseudobulks
    except KeyError:
        return None
    deg_list_path = dataset.get_pseudobulk_hvg_list_path()  # top 15,000 HVGs
    return {"deg_list": str(deg_list_path), "gene_data_path": str(gene_data_path)}


# ATAC-only datasets borrow a related dataset's HVG list just to gene-center their regions
region_selection_for_atac_datasets = {
    "Zu2023Nature": "AIBSDev",
    "MouseAging": "AIBSDev",
    "CBDevEvoMouse": "AIBSDev",
    "MouseBodyAtlas": "AIBSDev",
    "Li2023Science": "HumanFirstTriBrain",
    "Zhang2021Cell": "EncodeHumanAtlas",
    "CBDevEvoHuman": "HumanFirstTriBrain",
}

for dataset_key, records in dataset_records.items():
    dataset_name, *subset = dataset_key.split("+")
    subset = None if len(subset) == 0 else subset[0]
    # RNA is per dataset; of the subsets, only a "Control" subset reuses it
    rna = published_rna(dataset_name) if subset in (None, "Control") else None
    if rna is not None:
        records.update(rna)
    else:
        donor = DATASETS[region_selection_for_atac_datasets[dataset_name]]
        records["deg_list"] = str(donor.get_pseudobulk_hvg_list_path())

# which datasets carry RNA (ATAC + Gene) vs. ATAC-only
task = pd.Series(
    {
        k: ("ATAC + Gene" if "gene_data_path" in v else "ATAC only")
        for k, v in dataset_records.items()
    }
)
task.to_frame("task")

full_records = ATAC_MODEL.model_config["dataset_records"] print(len(full_records) - 1, "datasets in the full gene task (all but CBDevEvoOpossum)") # the demo's datasets dataset_records = {k: copy.deepcopy(full_records[k]) for k in DEMO_DATASETS} for rec in dataset_records.values(): rec["dataset_sample_weight"] = 1 def published_rna(dataset_name): """The dataset's published RNA files, or None if it has no matched transcriptome.""" dataset = DATASETS[dataset_name] try: gene_data_path = ( dataset.get_pseudobulk_gene_data_path() ) # log1p-CPM pseudobulks except KeyError: return None deg_list_path = dataset.get_pseudobulk_hvg_list_path() # top 15,000 HVGs return {"deg_list": str(deg_list_path), "gene_data_path": str(gene_data_path)} # ATAC-only datasets borrow a related dataset's HVG list just to gene-center their regions region_selection_for_atac_datasets = { "Zu2023Nature": "AIBSDev", "MouseAging": "AIBSDev", "CBDevEvoMouse": "AIBSDev", "MouseBodyAtlas": "AIBSDev", "Li2023Science": "HumanFirstTriBrain", "Zhang2021Cell": "EncodeHumanAtlas", "CBDevEvoHuman": "HumanFirstTriBrain", } for dataset_key, records in dataset_records.items(): dataset_name, *subset = dataset_key.split("+") subset = None if len(subset) == 0 else subset[0] # RNA is per dataset; of the subsets, only a "Control" subset reuses it rna = published_rna(dataset_name) if subset in (None, "Control") else None if rna is not None: records.update(rna) else: donor = DATASETS[region_selection_for_atac_datasets[dataset_name]] records["deg_list"] = str(donor.get_pseudobulk_hvg_list_path()) # which datasets carry RNA (ATAC + Gene) vs. ATAC-only task = pd.Series( { k: ("ATAC + Gene" if "gene_data_path" in v else "ATAC only") for k, v in dataset_records.items() } ) task.to_frame("task")

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/index/pseudobulk.gene_and_hvg_data.joblib "HTTP/1.1 302 Found"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/gene_data/ChromiumPBMC%2BNone.cond_pseudobulk.log1pCPM.feather "HTTP/1.1 302 Found"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/gene_data/ChromiumPBMC%2BNone.top15000_hvg.joblib "HTTP/1.1 302 Found"

2026-10-27 16:25:05 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/gene_data/HumanFirstTriBrain%2BNone.top15000_hvg.joblib "HTTP/1.1 302 Found"

41 datasets in the full gene task (all but CBDevEvoOpossum)

Out[3]:

	task
ChromiumPBMC	ATAC + Gene
Li2023Science+Inh	ATAC only

Build the gene-task config¶

Same registry, now with output_head_type="gene_count", use_regions="borzoi_gene", the explicit _genomes list, and lora_checkpoint_path pointing at the trained ATAC-final checkpoint. The schedule is a demo: one short epoch of 20 training steps at batch size 1. The real run used max_epochs=200, lr_total_steps=4_000_000, save_state_every_n_epoch=25 and the default batch_size of 4; the demo lines are simply dropped for it.

In [4]:

config = {
    "output_dir": "model",
    "savename": f"Bolero10M_atac_fold{FOLD}-gene-demo",
    "wandb_project": "BorzoiMulti",
    "wandb_job_type": "lora",
    "wandb_name": f"Bolero10M_atac_fold{FOLD}-gene-demo",
    # data
    "dataset_records": dataset_records,
    "fold_split_id": FOLD,
    # model
    "base_checkpoint_path": BASE_CKPT,
    "signal_model": True,
    # gene task
    "use_regions": "borzoi_gene",
    "output_head_type": "gene_count",
    "_genomes": GENOMES,  # keep the genome embedding aligned with the ATAC model
    "lora_checkpoint_path": ATAC_FINAL_CKPT,  # continue from the published ATAC model
    # schedule: real run max_epochs=200, lr_total_steps=4_000_000, save_state_every_n_epoch=25
    "max_epochs": 1,
    "lr_total_steps": 4_000_000,
    "save_state_every_n_epoch": 25,
    # 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 = { "output_dir": "model", "savename": f"Bolero10M_atac_fold{FOLD}-gene-demo", "wandb_project": "BorzoiMulti", "wandb_job_type": "lora", "wandb_name": f"Bolero10M_atac_fold{FOLD}-gene-demo", # data "dataset_records": dataset_records, "fold_split_id": FOLD, # model "base_checkpoint_path": BASE_CKPT, "signal_model": True, # gene task "use_regions": "borzoi_gene", "output_head_type": "gene_count", "_genomes": GENOMES, # keep the genome embedding aligned with the ATAC model "lora_checkpoint_path": ATAC_FINAL_CKPT, # continue from the published ATAC model # schedule: real run max_epochs=200, lr_total_steps=4_000_000, save_state_every_n_epoch=25 "max_epochs": 1, "lr_total_steps": 4_000_000, "save_state_every_n_epoch": 25, # 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 — then load the published gene checkpoint¶

As before, constructing the trainer builds the multi-dataset dataset (now also wiring up the gene data), and _setup_model() builds the model — this time with a gene-count head — and loads the frozen base plus the ATAC-final checkpoint. We then load the published gene checkpoint with strict=True, after dropping the encoders of the datasets the demo leaves out, so every weight of the demo model comes from the gene model.

In [5]:

trainer = MultiBorzoiLoRATrainer(config)
trainer._setup_model(print_model=False)

model = getattr(trainer.model, "module", trainer.model)
state = torch.load(GENE_CKPT, map_location="cpu", weights_only=False)["state_dict"]
state = filter_dataset_encoders(
    state, dataset_records
)  # keep the demo datasets' encoders
model.load_state_dict(state, strict=True)  # <- the trained gene model
del state
print("loaded trained checkpoint: Borzoi10M-Gene")

trainer = MultiBorzoiLoRATrainer(config) trainer._setup_model(print_model=False) model = getattr(trainer.model, "module", trainer.model) state = torch.load(GENE_CKPT, map_location="cpu", weights_only=False)["state_dict"] state = filter_dataset_encoders( state, dataset_records ) # keep the demo datasets' encoders model.load_state_dict(state, strict=True) # <- the trained gene model del state print("loaded trained checkpoint: Borzoi10M-Gene")

2026-10-27 16:25:06 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/metacell/parquet/ChromiumPBMC/ChromiumPBMC-MetaCell-32bp.tar "HTTP/1.1 302 Found"

2026-10-27 16:25:06 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/metacell/parquet/Li2023Science%2BInh/Li2023Science-MetaCell-32bp.tar "HTTP/1.1 302 Found"

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_gene
  n_pseudobulks: 10
  paired_mode: ensemble
  reverse_complement: True
  max_jitter: 3
  batch_size: 1
  output_prefix: pseudobulk
  shared_data_paths: None
  _genomes: 7 items <class 'list'>

2026-10-27 16:25:09 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/pseudobulk/condition/ChromiumPBMC/pseudobulk_records_and_cond.cov5000000.joblib "HTTP/1.1 302 Found"

2026-10-27 16:25:09 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/pseudobulk/condition/Li2023Science%2BInh/pseudobulk_records_and_cond.cov5000000.joblib "HTTP/1.1 302 Found"

Updated cond_module_kwargs in config with dataset specific information.

(_process_single_dataset pid=3611427) After deg_list filter, 14789 gene regions remain.

Setting up model from config

Loading base model weights from: /mnt/work/bdhome/flashzoi/fold_0.renamed.pt

Total model parameters 1202745955, trainable parameters 1025384867
Config contains lora_checkpoint_path, will load weights and perform fine-tuning.
Loading LoRA model weights from: /mnt/work/bdhome/hf/models--arcinstitute--bolero-models/snapshots/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS/Borzoi10M/checkpoint.pt

Skipping checkpoint encoders of datasets not in the config: ['AIBSDev', 'CBDevEvoBonobo', 'CBDevEvoHuman', 'CBDevEvoMacaque', 'CBDevEvoMarmoset', 'CBDevEvoMouse', 'CZIBodyAtlas', 'EE+Control', 'EncodeHumanAtlas', 'Fear+Control', 'HMBABasalGangliaHuman', 'HMBABasalGangliaMacaque', 'HMBABasalGangliaMarmoset', 'HYSocialRebound', 'Harris2025Nature', 'HumanBodyDev', 'HumanBrainAging', 'HumanBrainDev', 'HumanCBDev', 'HumanFirstTriBrain', 'KA+Control', 'Li2023Science+Exc', 'Li2023Science+NonN', 'MouseAging', 'MouseBodyAtlas', 'MultiSpeciesHuman', 'MultiSpeciesMacaque', 'MultiSpeciesMarmoset', 'MultiSpeciesMouse', 'MultipleSclerosis+Control', 'NeurIPSBMMC', 'Stress+Control', 'Torpor+Control', 'Zhang2021Cell', 'Zu2023Nature+CB', 'Zu2023Nature+CNU', 'Zu2023Nature+CTXGaba', 'Zu2023Nature+CTXGlut', 'Zu2023Nature+HB', 'Zu2023Nature+IB', 'Zu2023Nature+MB', 'Zu2023Nature+NonN']

loaded trained checkpoint: Borzoi10M-Gene

What's special — the gene-count head¶

Setting output_head_type="gene_count" adds two modules on top of the shared trunk that the ATAC model does not have: a gene_count_output_head (predicts per-gene transcript abundance) and a gene_mask_conv (encodes which positions belong to the gene being scored). The accessibility signal path is still present — the model predicts both tasks — but training regions are now gene-centered (use_regions="borzoi_gene"), and the extra gene-count loss is added for datasets that carry RNA.

In [6]:

print("has gene_count_output_head:", hasattr(model, "gene_count_output_head"))
print("has gene_mask_conv        :", hasattr(model, "gene_mask_conv"))
print("gene_count_output_head:")
print(model.gene_count_output_head)

# the conditioning machinery (per-dataset encoders + genome embedding) is unchanged from the ATAC model
cem = model.cond_emb_module
print("\nper-dataset conditioning encoders:", len(cem.cell_encoder_dict))
print("dataset_shared_dims:", cem.dataset_shared_dims)

print("has gene_count_output_head:", hasattr(model, "gene_count_output_head")) print("has gene_mask_conv :", hasattr(model, "gene_mask_conv")) print("gene_count_output_head:") print(model.gene_count_output_head) # the conditioning machinery (per-dataset encoders + genome embedding) is unchanged from the ATAC model cem = model.cond_emb_module print("\nper-dataset conditioning encoders:", len(cem.cell_encoder_dict)) print("dataset_shared_dims:", cem.dataset_shared_dims)

has gene_count_output_head: True
has gene_mask_conv        : True
gene_count_output_head:
OutputHead(
  (0): ConditionalLoRAConv(
    (conv): Conv1d(1920, 1, kernel_size=(1,), stride=(1,))
    (lora_dropout): Dropout(p=0.01, inplace=False)
    (lora_A_module): EmbeddingMLP(
      (mlp): Sequential(
        (0): Sequential(
          (0): Linear(in_features=768, out_features=256, bias=True)
          (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
          (2): GELU(approximate='none')
        )
        (1): Sequential(
          (0): Linear(in_features=256, out_features=256, bias=True)
          (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
          (2): GELU(approximate='none')
        )
        (2): Linear(in_features=256, out_features=1920, bias=True)
      )
    )
    (lora_B_module): EmbeddingMLP(
      (mlp): Sequential(
        (0): Sequential(
          (0): Linear(in_features=768, out_features=256, bias=True)
          (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
          (2): GELU(approximate='none')
        )
        (1): Sequential(
          (0): Linear(in_features=256, out_features=256, bias=True)
          (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
          (2): GELU(approximate='none')
        )
        (2): Linear(in_features=256, out_features=1, bias=True)
      )
    )
  )
  (1): AdaptiveAvgPool1d(output_size=1)
  (2): Softplus(beta=1.0, threshold=20.0)
)

per-dataset conditioning encoders: 2
dataset_shared_dims: {'__genome__': 7}

In [7]:

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,202,745,955
trainable params: 1,025,384,867  (85.3%)

Train the gene task on the demo datasets¶

trainer.train() builds the model again from the config (base + ATAC-final weights), trains the gene task on the demo schedule, and writes checkpoints under model/. The gene-count loss is added for ChromiumPBMC, which has RNA; Li2023Science+Inh trains on ATAC only.

In [8]:

# train() builds its own copy of the model; free the one inspected above first
del model
trainer.model = None
torch.cuda.empty_cache()

trainer.train()

# train() builds its own copy of the model; free the one inspected above first del model trainer.model = None torch.cuda.empty_cache() trainer.train()

wandb: Tracking run with wandb version 0.28.0

wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
wandb: Run data is saved locally in /mnt/work/runs/wandb/offline-run-20261027_162528-bzafcsdv
wandb: View this run in the terminal with `wandb leet`

wandb: WARNING URL not available in offline run

Setting up model from config

Loading base model weights from: /mnt/work/bdhome/flashzoi/fold_0.renamed.pt

Total model parameters 1202745955, trainable parameters 1025384867
Config contains lora_checkpoint_path, will load weights and perform fine-tuning.
Loading LoRA model weights from: /mnt/work/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]             1,688,069
├─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,202,745,955
Trainable params: 1,025,384,867
Non-trainable params: 177,361,088
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:25:40,562	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-25-00_026732_3610717/logs/ray-data

2026-10-27 16:25:40,563	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition] -> TaskPoolMapOperator[MapBatches(_process_mask_coords)]

2026-10-27 16:25:41,666	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-25-00_026732_3610717/logs/ray-data

2026-10-27 16:25:41,667	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> TaskPoolMapOperator[MapBatches(ReverseComplementMinusStrand)] -> LimitOperator[limit=22]

Reverse complement minus strand gene regions
Data loader kwargs {'batch_size': 1, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}

 - (Training) 0 [4/20] Ave Loss: 1.454; Ave Corr: 0.795 Last grad norm: 0.6488 Last batch Loss: 1.250 

 - (Training) 0 [9/20] Ave Loss: 2.344; Ave Corr: 0.875 Last grad norm: 2.3165 Last batch Loss: 0.809 

/mnt/work/bolero/.pixi/envs/dev/lib/python3.11/site-packages/torch/optim/lr_scheduler.py:232: UserWarning: The epoch parameter in `scheduler.step()` was not necessary and is being deprecated where possible. Please use `scheduler.step()` to step the scheduler. During the deprecation, if epoch is different from None, the closed form is used instead of the new chainable form, where available. Please open an issue if you are unable to replicate your use case: https://github.com/pytorch/pytorch/issues/new/choose.
  warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)

 - (Training) 0 [14/20] Ave Loss: 2.198; Ave Corr: 0.881 Last grad norm: 1.6326 Last batch Loss: 4.066 

 - (Training) 0 [19/20] Ave Loss: 2.066; Ave Corr: 0.887 Last grad norm: 4.5300 Last batch Loss: 1.382 

2026-10-27 16:26:10,406	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-25-00_026732_3610717/logs/ray-data

2026-10-27 16:26:10,407	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition] -> TaskPoolMapOperator[MapBatches(_process_mask_coords)]

Get dataloader with eval mode

2026-10-27 16:26:10,810	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-25-00_026732_3610717/logs/ray-data

2026-10-27 16:26:10,811	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: 1.972; Mean Corr. : 0.884; Learning rate 4.9999375001562504e-05.
 - (Validation) 1 Loss: 1.148; Mean Corr. : 0.879.
Saving best checkpoint...

Load and update state dict from checkpoint file: /mnt/work/bolero/docs/tutorials/model_training/model/Bolero10M_atac_fold0-gene-demo.lora.best_checkpoint.pt

2026-10-27 16:27:11,228	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-25-00_026732_3610717/logs/ray-data

2026-10-27 16:27:11,228	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition] -> TaskPoolMapOperator[MapBatches(_process_mask_coords)]

Get dataloader with eval mode

2026-10-27 16:27:12,307	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-25-00_026732_3610717/logs/ray-data

2026-10-27 16:27:12,307	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: 0.494; Mean Track Corr: 0.714; 

wandb: 
wandb: Run history:
wandb:              train/learning_rate ▁▆██
wandb:            train/train_corr_data ▁▇██
wandb:                 train/train_loss ▂▁█▂
wandb: train/train_loss_gene_count_loss ▁▁█▆
wandb:     train/train_loss_multinomial ▂▁█▁
wandb:         train/train_loss_poisson ▇█▁█
wandb:      train/train_total_grad_norm ▁▄▃█
wandb:                val/val_corr_data ▁
wandb:                     val/val_loss ▁
wandb:     val/val_loss_gene_count_loss ▁
wandb:                               +2 ...
wandb: 
wandb: Run summary:
wandb:                  final_test_corr 0.71431
wandb:                  final_test_loss 0.49449
wandb:  final_test_loss_gene_count_loss 0
wandb:      final_test_loss_multinomial 0.68138
wandb:          final_test_loss_poisson -0.48669
wandb:                 final_valid_corr 0.87925
wandb:                 final_valid_loss 1.14781
wandb: final_valid_loss_gene_count_loss 0.06367
wandb:     final_valid_loss_multinomial 1.53189
wandb:         final_valid_loss_poisson -1.2666
wandb:                              +13 ...
wandb: 

wandb: You can sync this run to the cloud by running:
wandb: wandb sync /mnt/work/runs/wandb/offline-run-20261027_162528-bzafcsdv

wandb: Find logs at: /mnt/work/runs/wandb/offline-run-20261027_162528-bzafcsdv/logs

Real training vs. this demo¶

The published gene checkpoint is the last stage of the full pipeline. Reproducing it means another 200-epoch, multi-GPU run on all 41 datasets, continuing from the ATAC-final checkpoint — days of compute. For that run, annotate every dataset of the registry and drop the demo lines.

---

That completes Bolero's model side: a single frozen Flashzoi backbone, conditioned per cell state and per dataset, trained jointly across the Bolero-10M corpus to predict cell-state–specific accessibility and transcript abundance across species — the model behind the paper's results.