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

Checkpoint	Epoch_58_Step_1715504 (run epi_compass_small_arch_dann)
Parameters	301,185,880 (301.2 M) across 119 tensors
File	epifoundation_dann_pretrain.pth (1.12 GB, fp32)
Cell embedding	512-d, CLS-pooled

Architecture

scTransformer encoder over ATAC peak tokens and RNA gene tokens.

Component	Value
Encoder	Transformer, flash-attention backend
Layers / hidden dim / heads / FFN dim	6 / 512 / 8 / 1024
Dropout	0.2
Max tokens	ATAC 8,000, RNA 8,000
Embedding method	id_only
Cell embedding style	cls
MVC decoder	concat query
Chromosome embeddings	on (use_chr_labels: True)
Batch embeddings	on (use_batch_labels: True)
Batch adversary	DANN via gradient reversal, λ = 0.1, loss weight 0.1

Parameter budget is dominated by the ATAC peak embedding table:

Submodule	Tensors	Params	Share
atac_emb	3	268,368,384	89.10 %
rna_emb	3	18,742,784	6.22 %
encoder	72	12,616,704	4.19 %
cls_decoder	10	854,655	0.28 %
batch_disc	6	251,604	0.08 %
mvc_decoder	7	213,507	0.07 %
batch_emb	3	109,568	0.04 %
chr_emb	3	23,552	0.01 %
norm, bn_atac, bn_rna	12	5,122	0.00 %

Vocabulary sizes (fixed by the embedding tables)

Vocab	Size
ATAC peaks	524,155
RNA genes	36,605
Batch (DANN classes)	212
Chromosome	44
Cell type (cls_decoder output)	127

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

Objective	Masked value/expression reconstruction (MVC, binary) + adversarial batch CE
Task weights	mvc: 1.0, dann: 0.1, cell_type: 0.0 (classification head not trained during pretraining)
Optimizer	Adam, lr 1e-4, cosine annealing with warm restarts
Precision	AMP
Effective batch	256
Trained to	epoch 58, global step 1,715,504
Hardware	NVIDIA H200 (Duke Compute Cluster)

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.

Tissue	Cells	Types	Acc	Bal. acc	Macro F1	sil(CT)	NMI	ARI	kBET	iLISI	sil(batch)	Pearson	Zero acc
Blood	1,435	13	0.894	0.720	0.702	0.333	0.613	0.338	0.052	0.041	0.232	0.556	0.925
Bone marrow	939	9	0.909	0.903	0.874	0.358	0.630	0.387	0.060	0.120	0.195	0.514	0.932
Brain	970	16	0.720	0.639	0.660	0.157	0.530	0.344	0.934	0.012	0.087	0.623	0.926
Kidney	1,275	7	0.948	0.833	0.836	0.569	0.523	0.186	0.035	0.037	0.435	0.563	0.942
T cells	1,210	3	1.000	1.000	1.000	0.921	0.295	0.076	0.840	0.050	0.716	0.471	0.899

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

File	Description
epifoundation_dann_pretrain.pth	Model weights (1.12 GB)
config.json	Machine-readable architecture + training summary
config.yml	Original pretraining config

License

MIT.