epifoundation-compass-dann-finetune

EpiFoundation — CompassDB DANN finetuned checkpoints (5 tissues)

Per-tissue finetuned EpiFoundation models that take paired single-cell ATAC + RNA input and produce, for each cell:

1. a 512-d cell embedding,
2. a cell-type prediction,
3. a per-gene expression prediction (scGPT-style 10-bin values, 0–9) with a zero-inflation gate.

All five are finetuned from the same backbone: Chtholly17/epifoundation-compass-dann.

Files

Each is 121 tensors / ~939.8 M parameters. Two tensors dominate the footprint: the gene-specific expression head value_decoder.gene_fc1_weight [36605, 16, 1088] (637 M params, 68%) — one weight block per gene in the full RNA vocabulary — and the ATAC peak embedding atac_emb.embedding.weight [524155, 512] (268 M, 29%). If you only need cell embeddings and cell-type calls, drop value_decoder.* after loading to free most of that.

Each file is a weights-only archive: {"model": state_dict, "epoch": 100, "steps": ...}. Optimizer / scheduler / scaler state has been stripped, so these are inference and finetune-init checkpoints, not training-resume points.

Held-out test performance

Against the same finetune recipe initialised from the public UCSC-VLAA checkpoint (only the pretrained backbone differs), accuracy improves on all five tissues — Blood +2.6 pp, Bone marrow +2.1, Brain +20.8, Kidney +2.3, T cells +0.1 — and cell-type silhouette improves everywhere (e.g. Blood 0.002 → 0.333).

Caveat: batch-mixing metrics are worse than that baseline on Blood, Bone marrow and Kidney (Blood kBET 0.324 → 0.052). This backbone retains more batch signal in the cell embedding. Prefer these checkpoints for cell-type annotation and expression prediction; if batch integration is the priority, apply a post-hoc correction to cell_emb.

Architecture

Per-tissue vocab sizes (these set cls_decoder / batch_emb output dims, so use the matching vocab files):

Getting cell embeddings and expression predictions

Code lives in the EpiFoundation repo (model/scTransformer.py, data/dataloader.py).

1. Build the model

import torch
from model import scCross
from tokenizer import GeneVocab

rna_vocab   = GeneVocab.from_file("vocabs/rna_vocab.json")
atac_vocab  = GeneVocab.from_file("vocabs/atac_vocab.json")
chr_vocab   = GeneVocab.from_file("vocabs/chr_vocab.json")
cell_vocab  = GeneVocab.from_file("vocabs/finetune_blood/cell_vocab.json")   # per tissue
batch_vocab = GeneVocab.from_file("vocabs/finetune_blood/batch_vocab.json")  # per tissue

model = scCross(
    num_class_cell=len(cell_vocab),
    num_rnas=len(rna_vocab),
    num_atacs=len(atac_vocab),
    num_values=2,
    num_chrs=len(chr_vocab),
    embed_dim=512, depth=6, heads=8, head_dim=1024,
    encoder="transformer", dropout=0.1,
    pad_token_idx_rna=rna_vocab["<pad>"],
    pad_token_idx_atac=atac_vocab["<pad>"],
    cell_emb_style="cls",
    mvc_arch_style="gene-specific concat query",
    use_batch_labels=True,
    batch_label_num=len(batch_vocab),
    use_chr_labels=True,
    transformer_backend="flash",     # or "pytorch"
    stage="value_finetune",
    use_zero_inflated=True,
    regression_max_value=9.0,
    use_dann=True, dann_lambda=0.5,
).cuda().eval()

ck = torch.load("epifoundation_dann_finetune_Blood.pth", map_location="cpu")
model.load_state_dict(ck["model"], strict=True)

If you hit unexpected mvc_decoder.* keys, drop them before loading — the value_finetune stage does not build that module:

sd = {k: v for k, v in ck["model"].items() if not k.startswith("mvc_decoder.")}
model.load_state_dict(sd, strict=True)

2. Run a forward pass

Inputs come from PairedSCDataset (paired ATAC + RNA h5ad, RNA pre-binned to 10 levels). One batch gives everything in a single forward:

pad_atac_id = atac_vocab["<pad>"]

with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
    out = model(
        atac=batch["atac_ids"].cuda(),
        rna=batch["rna_ids"].cuda(),
        src_key_padding_mask=batch["atac_ids"].cuda().eq(pad_atac_id),
        batch_id=batch["batch_ids"].cuda(),
        rna_chrs=batch["rna_chrs"].cuda(),
        atac_chrs=batch["atac_chrs"].cuda(),
    )

3. Read the outputs

Cell embedding — (n_cells, 512), the CLS-pooled representation. This is what you cluster / UMAP / integrate:

cell_emb = out["cell_emb"].float().cpu().numpy()

Cell-type prediction — logits over the tissue's cell vocab:

probs   = torch.softmax(out["cell_pred"].float(), dim=-1).cpu().numpy()
pred_id = probs.argmax(-1)
pred_names = [cell_vocab.lookup_token(int(i)) for i in pred_id]

Gene expression prediction — the zero-inflated head returns two branches; gate the regression output with the zero probability:

vp        = out["value_pred"]
regression = vp["regression"].float().cpu().numpy()          # (n_cells, n_genes), in [0, 9]
zero_prob  = torch.sigmoid(vp["zero_logits"]).float().cpu().numpy()

ZERO_THRESHOLD = 0.60                                        # value used at eval time
prediction = np.where(zero_prob > ZERO_THRESHOLD, 0.0, regression)

prediction is on the scGPT-style 10-bin scale (0–9), matching the binned ground truth — not raw counts or log-normalised expression. To compare against your own data, apply the same per-cell 10-bin discretisation.

scripts/eval_blood_full.py in the repo does exactly this end to end and writes a results.h5ad with .obsm["X_emb"], .layers["prediction"], .layers["regression"], .layers["zero_prob"] and predicted/ground-truth cell types in .obs.

Training recipe

Backbone (see the pretrain repo): CompassDB pretraining split, 208 paired ATAC+RNA samples, 1,546,146 cells; masked expression reconstruction with batch labels supplied to the decoders and a DANN batch discriminator on the cell embedding (gradient reversal λ = 0.1); Adam, lr 1e-4, effective batch 256, ~1.72 M steps.

Finetune (per tissue, identical settings across all five): joint cell-type classification + zero-inflated expression prediction. Encoder and the RNA/ATAC/chromosome embeddings transferred; cls_decoder, value_decoder, batch_emb, batch_disc re-initialised. Loss weights — cell type 1.0, MVC 1.0, zero-BCE 1.0, value-MSE 0.05, DANN 1.0; DANN λ = 0.5; class-weighted cross-entropy (sqrt inverse frequency). Adam, lr 1e-4, 100 epochs, effective batch 320 (32 × 2 GPUs × 5 accumulation), AMP, 2 × H200. RNA is HVG-selected per tissue (~2,360 genes) and 10-bin discretised.

Citation

Please cite the EpiFoundation work if you use these checkpoints.