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 runs in the multi-dataset training directory on the full corpus + the RNA data, and — because a real run is heavy — it does not call trainer.train(). It builds the model with the real config and loads the trained gene checkpoint to inspect what the gene task adds.

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 — the gene task continues from there.
• _genomes is passed explicitly so the genome embedding stays identical to the ATAC model.

Setup¶

In [1]:

import json
import os
import pathlib

import pandas as pd

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

os.environ["WANDB_MODE"] = "disabled"
init(num_cpus=16, object_store_memory_gb=48)

import json import os import pathlib import pandas as pd from bolero import init from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer os.environ["WANDB_MODE"] = "disabled" init(num_cpus=16, object_store_memory_gb=48)

2026-07-13 12:11:47,384	INFO worker.py:1781 -- Started a local Ray instance.

In [2]:

# --- Configuration ---------------------------------------------------------------
MODEL_NAME = "Bolero10M"
FOLD = 0

INIT_CONFIG = f"model/250909-{MODEL_NAME}_atac_fold0-init.config.json"
# start the gene task from the trained ATAC-final checkpoint (previous page)
ATAC_FINAL_CKPT = f"model/250909-{MODEL_NAME}_atac_fold0-final.lora.best_checkpoint.pt"
# the trained gene checkpoint we load for the demo
GENE_CKPT = f"model/250909-{MODEL_NAME}_atac_fold0-gene.lora.best_checkpoint.pt"
BASE_CKPT = (
    f"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_{FOLD}.renamed.pt"
)

# per-dataset RNA data (log1p-CPM pseudobulk expression + highly-variable-gene lists)
GENE_DATA_DIR = "/home/hanliu/data/250901-GeneCountPred/prepare/gene_data"

# 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",
]

# --- Configuration --------------------------------------------------------------- MODEL_NAME = "Bolero10M" FOLD = 0 INIT_CONFIG = f"model/250909-{MODEL_NAME}_atac_fold0-init.config.json" # start the gene task from the trained ATAC-final checkpoint (previous page) ATAC_FINAL_CKPT = f"model/250909-{MODEL_NAME}_atac_fold0-final.lora.best_checkpoint.pt" # the trained gene checkpoint we load for the demo GENE_CKPT = f"model/250909-{MODEL_NAME}_atac_fold0-gene.lora.best_checkpoint.pt" BASE_CKPT = ( f"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_{FOLD}.renamed.pt" ) # per-dataset RNA data (log1p-CPM pseudobulk expression + highly-variable-gene lists) GENE_DATA_DIR = "/home/hanliu/data/250901-GeneCountPred/prepare/gene_data" # 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", ]

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; 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.

In [3]:

with open(INIT_CONFIG) as f:
    init_cfg = json.load(f)
dataset_records = init_cfg["dataset_records"]

# drop the species with no RNA data
del dataset_records["CBDevEvoOpossum"]

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

for dataset_key, records in dataset_records.items():
    dataset_name, *subset = dataset_key.split("+")
    subset = None if len(subset) == 0 else subset[0]
    key = (
        f"{dataset_name}+None"
        if subset in (None, "Control")
        else f"{dataset_name}+{subset}"
    )
    gene_data_path = f"{GENE_DATA_DIR}/{key}.cond_pseudobulk.log1pCPM.feather"
    deg_list_path = f"{GENE_DATA_DIR}/{key}.top15000_hvg.joblib"
    if pathlib.Path(gene_data_path).exists():
        records.update({"deg_list": deg_list_path, "gene_data_path": gene_data_path})
    else:
        records["deg_list"] = region_selection_for_atac_datasets[dataset_name]

# 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()
    }
)
print(task.value_counts())
task.to_frame("task").head(10)

with open(INIT_CONFIG) as f: init_cfg = json.load(f) dataset_records = init_cfg["dataset_records"] # drop the species with no RNA data del dataset_records["CBDevEvoOpossum"] # ATAC-only datasets borrow a related dataset's HVG list just to gene-center their regions region_selection_for_atac_datasets = { "Zu2023Nature": f"{GENE_DATA_DIR}/AIBSDev+None.top15000_hvg.joblib", "MouseAging": f"{GENE_DATA_DIR}/AIBSDev+None.top15000_hvg.joblib", "CBDevEvoMouse": f"{GENE_DATA_DIR}/AIBSDev+None.top15000_hvg.joblib", "MouseBodyAtlas": f"{GENE_DATA_DIR}/AIBSDev+None.top15000_hvg.joblib", "Li2023Science": f"{GENE_DATA_DIR}/HumanFirstTriBrain+None.top15000_hvg.joblib", "Zhang2021Cell": f"{GENE_DATA_DIR}/EncodeHumanAtlas+None.top15000_hvg.joblib", "CBDevEvoHuman": f"{GENE_DATA_DIR}/HumanFirstTriBrain+None.top15000_hvg.joblib", } for dataset_key, records in dataset_records.items(): dataset_name, *subset = dataset_key.split("+") subset = None if len(subset) == 0 else subset[0] key = ( f"{dataset_name}+None" if subset in (None, "Control") else f"{dataset_name}+{subset}" ) gene_data_path = f"{GENE_DATA_DIR}/{key}.cond_pseudobulk.log1pCPM.feather" deg_list_path = f"{GENE_DATA_DIR}/{key}.top15000_hvg.joblib" if pathlib.Path(gene_data_path).exists(): records.update({"deg_list": deg_list_path, "gene_data_path": gene_data_path}) else: records["deg_list"] = region_selection_for_atac_datasets[dataset_name] # 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() } ) print(task.value_counts()) task.to_frame("task").head(10)

ATAC + Gene    24
ATAC only      16
Name: count, dtype: int64

Out[3]:

	task
Zu2023Nature+NonN	ATAC only
Zu2023Nature+CNU	ATAC only
Zu2023Nature+IB	ATAC only
Zu2023Nature+CTXGlut	ATAC only
Zu2023Nature+CTXGaba	ATAC only
Zu2023Nature+CB	ATAC only
Zu2023Nature+HB	ATAC only
Zu2023Nature+MB	ATAC only
Li2023Science+Exc	ATAC only
Li2023Science+Inh	ATAC only

Build the gene-task config (real parameters)¶

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 training schedule is again the real one (max_epochs=200, lr_total_steps=4_000_000); a smoke test would shrink those as on the single-dataset page.

In [4]:

config = {
    "output_dir": "model",
    "savename": f"250909-{MODEL_NAME}_atac_fold{FOLD}-gene",
    "wandb_project": "250909-BorzoiMulti",
    "wandb_job_type": "lora",
    "wandb_name": f"250909-{MODEL_NAME}_atac_fold{FOLD}-gene",
    # 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 trained ATAC model
    # real training schedule
    "max_epochs": 200,
    "lr_total_steps": 4_000_000,
    "save_state_every_n_epoch": 25,
}
config = MultiBorzoiLoRATrainer.make_config(**config)

config = { "output_dir": "model", "savename": f"250909-{MODEL_NAME}_atac_fold{FOLD}-gene", "wandb_project": "250909-BorzoiMulti", "wandb_job_type": "lora", "wandb_name": f"250909-{MODEL_NAME}_atac_fold{FOLD}-gene", # 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 trained ATAC model # real training schedule "max_epochs": 200, "lr_total_steps": 4_000_000, "save_state_every_n_epoch": 25, } config = MultiBorzoiLoRATrainer.make_config(**config)

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

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 trained gene checkpoint instead of training.

In [5]:

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

model = getattr(trainer.model, "module", trainer.model)
model.load_checkpoint_from_path(GENE_CKPT, strict=False)  # <- the trained gene model
print("loaded trained checkpoint:", GENE_CKPT)

trainer = MultiBorzoiLoRATrainer(config) trainer._setup_model(print_model=False) model = getattr(trainer.model, "module", trainer.model) model.load_checkpoint_from_path(GENE_CKPT, strict=False) # <- the trained gene model print("loaded trained checkpoint:", GENE_CKPT)

Using single GPU for training
Create BorzoiMultiDataset with config:
  dataset_records: 40 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: 4
  output_prefix: pseudobulk
  shared_data_paths: None
  _genomes: 7 items <class 'list'>

Updated cond_module_kwargs in config with dataset specific information.

(_process_single_dataset pid=2475666) After deg_list filter, 14600 gene regions remain.

Setting up model from config

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

Total model parameters 1210501800, trainable parameters 1033140712
Config contains lora_checkpoint_path, will load weights and perform fine-tuning.
Loading LoRA model weights from: model/250909-Bolero10M_atac_fold0-final.lora.best_checkpoint.pt

Loading LoRA model weights from: model/250909-Bolero10M_atac_fold0-gene.lora.best_checkpoint.pt

loaded trained checkpoint: model/250909-Bolero10M_atac_fold0-gene.lora.best_checkpoint.pt

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: 40
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,210,501,800
trainable params: 1,033,140,712  (85.3%)

Real training vs. this demo¶

The trained gene checkpoint we loaded is the last stage of the full pipeline. Reproducing it means another 200-epoch, multi-GPU run continuing from the ATAC-final checkpoint — days of compute. To watch the loop, add the smoke-test overrides (max_epochs=3, warmup_steps=10, train_batches=50, val_batches=10) and call trainer.train().

---

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.