mims-harvard/CASCADE-HLCA

CASCADE: context-aware single-cell modelling links cellular programmes to patient-level disease phenotypes

Model Card

This repository contains a CASCADE checkpoint pre-trained on the Human Lung Cell Atlas (HLCA) cohort.

Introduction

CASCADE integrates contextual information into both input representation and pre-training objectives, allowing the same cell to be interpreted through multiple biologically meaningful axes and enabling patient-level phenotype prediction from single-cell profiles.

1. Context-aware tokenisation. Each cell is encoded as context-dependent up- and down-regulated genes relative to a biologically defined reference group, producing multiple representations per cell across disease, tissue, cell type, and treatment contexts.

2. Context-specific representation learning. Shared cell embeddings are projected through separate context-specific projectors (disease, tissue, cell type, treatment), learning how molecular programmes vary across biologically meaningful contexts via contrastive objectives.

3. Patient representation & explainability. Cell-level embeddings are aggregated across all cells from a donor to produce a patient-level representation for multiscale phenotype prediction. CASCADE-Explainer identifies the cell types and genes most responsible for each prediction.

• Code: https://github.com/mims-harvard/CASCADE
• Project page: https://mims-harvard.github.io/CASCADE-website/
• Source dataset: https://cellxgene.cziscience.com/collections/edb893ee-4066-4128-9aec-5eb2b03f8287

Training Data

• Disease states (16): COVID-19, chronic obstructive pulmonary disease, chronic rhinitis, cystic fibrosis, hypersensitivity pneumonitis, interstitial lung disease, lung adenocarcinoma, lung large cell carcinoma, lymphangioleiomyomatosis, non-specific interstitial pneumonia, normal, pleomorphic carcinoma, pneumonia, pulmonary fibrosis, pulmonary sarcoidosis, squamous cell lung carcinoma
• Tissues (4): lung, lung parenchyma, nose, respiratory airway
• Cell types (51): identified by Cell Ontology (CL) ID — CL:0000037, CL:0000057, CL:0000077, CL:0000084, CL:0000097, CL:0000158, CL:0000186, CL:0000192, CL:0000236, CL:0000313, .... Look up terms at https://www.ebi.ac.uk/ols4/ontologies/cl.

Model Architecture

CASCADE's encoder (TransformerGenerator) is a shared transformer over context-aware gene token sequences, followed by context-specific projection heads (one per context in the table below) trained with a context-specific contrastive objective.

Model Hyperparameters

Hyperparameter	Value
Embedding dim (d_model)	384
Attention heads (nhead)	6
Transformer layers (nlayers)	12
Feedforward dim (dim_embedding)	384
Dropout	0.1
Vocabulary size	18711
Cell embedding style	avg-pool
Contexts	disease, cell_type, tissue
Context-specific projections	True
Domain adaptation (Sinkhorn)	True

Files Included

File	Purpose
model.safetensors	Model weights only (stripped of optimizer/scheduler/scaler state)
config.json	Architecture hyperparameters needed to reconstruct TransformerGenerator
tokenizer_dictionary_HLCA.pkl	Gene/context vocabulary (18711 tokens) used by the context-aware tokenizer
metadata_dictionary_HLCA.pkl	Obs-column metadata mapping preserved from the source AnnData
median_genes_*_all_HLCA.pkl	Per-context median expression reference used to derive up-/down-regulated gene tokens at tokenization time

Usage Instructions

The model architecture (TransformerGenerator) is not a standard transformers class, so loading it requires the cascade package from the GitHub repo rather than AutoModel:

pip install git+https://github.com/mims-harvard/CASCADE

import json, pickle
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from cascade.model.cascade_model import TransformerGenerator

repo_id = "<your-org>/CASCADE-HLCA"

config = json.load(open(hf_hub_download(repo_id, "config.json")))
vocab = pickle.load(open(hf_hub_download(repo_id, f"tokenizer_dictionary_HLCA.pkl"), "rb"))
weights = load_file(hf_hub_download(repo_id, "model.safetensors"))

model = TransformerGenerator(
    d_model=config["d_model"],
    nhead=config["nhead"],
    ntoken=config["vocab_size"],
    dim_embedding=config["dim_embedding"],
    nlayers=config["nlayers"],
    vocab=vocab,
    nclass=config["nclass"],
    dropout=config["dropout"],
    pad_token=config["pad_token"],
    cell_emb_style=config["cell_emb_style"],
    context_specific_projections=config["context_specific_projections"],
    constant_ctx=config["constant_ctx"],
    only_contrastive=config["only_contrastive"],
    DA=config["DA"],
    lambda_sinkhorn=config["lambda_sinkhorn"],
    merged_contexts=config["merged_contexts"],
)
model.load_state_dict(weights)
model.eval()

To tokenize new raw data for this model, see cascade/data/tokenizer.py in the GitHub repo, using the tokenizer_dictionary_HLCA.pkl and median_genes_*_all_HLCA.pkl files from this repo as the vocab and per-context median reference respectively.

Citation

Paper coming soon — see the project page for updates.

Contact

For any questions or feedback, please open an issue in the GitHub repository or contact Valentina Giunchiglia and Marinka Zitnik.