Chtholly17/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

File	Tissue	Size	Test cells	Cell types
epifoundation_dann_finetune_Blood.pth	Blood	3.76 GB	1,435	13
epifoundation_dann_finetune_Bone_marrow.pth	Bone marrow	3.76 GB	939	9
epifoundation_dann_finetune_Brain.pth	Brain	3.76 GB	970	16
epifoundation_dann_finetune_Kidney.pth	Kidney	3.76 GB	1,275	7
epifoundation_dann_finetune_T_cells.pth	T cells	3.76 GB	1,210	3

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

Tissue	Accuracy	Balanced acc	Macro F1	silhouette (cell type)	per-cell Pearson	zero acc
Blood	0.894	0.720	0.702	0.333	0.556	0.925
Bone marrow	0.909	0.903	0.874	0.358	0.514	0.932
Brain	0.720	0.639	0.660	0.157	0.623	0.926
Kidney	0.948	0.833	0.836	0.569	0.563	0.942
T cells	1.000	1.000	1.000	0.921	0.471	0.899

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

Encoder	Transformer, embedding_dim=512, num_layers=6, head_num=8, FFN 1024, dropout 0.1
Inputs	ATAC peaks (max 8,000 tokens) + RNA genes (max 4,000), chromosome embeddings on
Cell embedding	CLS token, 512-d
Cell-type head	cls_decoder, input concat(cell_emb, batch_emb)
Expression head	value_decoder, mvc_arch_style="gene-specific concat query", zero-inflated (regression branch clipped to [0, 9] + zero-logit branch)
Vocab dims baked into weights	atac 524,155 · rna 36,605 · chr 44 · cell type & batch per tissue (see below)

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

Tissue	cell vocab	batch vocab
Blood	18	8
Bone marrow	14	8
Brain	21	7
Kidney	12	7
T cells	8	9

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.