epifoundation-compass-dann

EpiFoundation — CompassDB DANN pretrained backbone

Transformer backbone for paired single-cell ATAC + RNA data, pretrained on CompassDB with a domain-adversarial (DANN) batch discriminator on the cell embedding.

This repository contains the pretrained weights only — optimizer, scheduler and GradScaler states have been stripped from the training checkpoint (3.37 GB → 1.12 GB). It is meant as an initialization for downstream finetuning, not as a resumable training checkpoint.

Architecture

scTransformer encoder over ATAC peak tokens and RNA gene tokens.

Parameter budget is dominated by the ATAC peak embedding table:

Vocabulary sizes (fixed by the embedding tables)

Special tokens: <cls> = 0, <pad> = 2, <mask> = 3.

The vocabulary JSON files are not included in this repository. The peak/gene token IDs must match the CompassDB vocabularies these embedding tables were built against, otherwise the embeddings are meaningless. Ask the authors for atac_vocab.json (524k peaks), rna_vocab.json, chr_vocab.json and gene2chr.json before using the checkpoint on new data.

Pretraining data

CompassDB pretraining split — 208 paired ATAC + RNA samples, 1,546,146 cells, merged into a single shard per modality. RNA was preprocessed offline (normalize-total → log1p → quantile bin, bin_num = 2); ATAC is binarized peak accessibility. 212 distinct sample/batch labels supply the DANN discriminator targets.

Training

The exact pretraining config is in config.yml; a machine-readable summary is in config.json.

Downstream results

Five-tissue paired ATAC→RNA finetune on CompassDB, all initialized from this backbone (encoder + rna/atac/chr embeddings transferred; cls_decoder, value_decoder, batch_emb, batch_disc re-initialized). Joint cell-type classification + zero-inflated RNA value prediction, 100 epochs, finetune-side DANN λ = 0.5. Metrics on held-out test cells.

Against an otherwise identical finetune initialized from the earlier UCSC-VLAA checkpoint, cell-type 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, Kidney 0.273 → 0.569).

Limitations

• The DANN objective did not remove batch signal. A linear/kNN probe on the frozen pretrained embedding recovers the batch label with 95.5 % kNN-15 accuracy on held-out blood cells (batch silhouette 0.231), and the discriminator loss stayed pinned at ≈ ln(212) = 5.36 — chance level for 212 classes — for the whole run. The adversary never learned. Downstream batch-mixing metrics regress relative to the UCSC-VLAA-initialized baseline on Blood, Bone marrow and Kidney (e.g. Blood kBET 0.324 → 0.052, iLISI 0.231 → 0.041). Follow-up runs that raised λ to 8.0 and the loss weight to 1.0 did not change this. Treat the DANN component as ineffective, and the strong cell-type performance as coming from the reconstruction objective, not from adversarial batch removal.
• Not a resumable training checkpoint — optimizer/scheduler/scaler states were removed.
• The cls_decoder head (127 classes) was not trained during pretraining (task_weight.cell_type = 0.0); its weights are at initialization. Re-initialize it for downstream use.
• Token vocabularies are fixed and not shipped here (see above).
• Evaluated only on CompassDB tissues; no held-out external cohort.

Usage

The checkpoint is a torch.save dict with the state dict under key model, matching what finetune.py in the EpiFoundation codebase expects:

import torch

ckpt = torch.load("epifoundation_dann_pretrain.pth", map_location="cpu")
print(ckpt.keys())          # dict_keys(['model', 'epoch'])
state_dict = ckpt["model"]  # 119 tensors, no DDP 'module.' prefix

model.load_state_dict(state_dict)                 # full backbone

To transfer only the backbone (drop the heads that are re-initialized downstream):

skip = ("value_decoder", "mvc_decoder", "batch_emb", "cls_decoder", "batch_disc")
backbone = {k: v for k, v in state_dict.items() if not k.startswith(skip)}
model.load_state_dict(backbone, strict=False)

Or point a config at it directly:

train:
  model:
    pretrained: /path/to/epifoundation_dann_pretrain.pth

Note that pretrain.py also reads optimizer / scheduler / scaler from the checkpoint when resuming, so this file works as a finetune initialization but not as a pretrain resume point.

Files

License

MIT.