Broad Obesity-2 Perturbation Prediction

CrunchDAO Hackathon Submission — Broad Obesity-2 Challenge (Phase 2)

Extends the Phase 1 approach to handle double-gene perturbations (e.g., FOXC1+SOX6). Uses the same DeltaInterpolationNetV9 cross-attention transformer with scGPT embeddings and Retrieve-and-Shift cell generation, adapted for the Crunch 2 dataset and scoring format.

See the Phase 1 repository (Crunch_1) for full architectural details.

---

What Changed from Phase 1

Aspect	Phase 1 (Crunch 1)	Phase 2 (Crunch 2)
Perturbation types	Single-gene only	Single + double-gene (GENE_A+GENE_B)
Control baselines	NC only	NC (single-guide) + NC+NC (double-guide), kept separate
Query/key features	Single gene features	Double perts: average of constituent gene features
Gene embeddings	Single gene embedding	Double perts: average of constituent gene embeddings
infer() parameter	genes_to_predict	predict_genes
Proportion output	pre_adipo, adipo, lipo, other	Same + lipo_adipo = lipo / adipo
Checkpoint name	ensemble_v9.pt	ensemble_v9_crunch2.pt
MMD² implementation	Naive O(N·M·D) — crashes at scale	Kernel-trick + row-batching, capped at 2k cells

---

Method Overview

The model predicts expression deltas (shift from control) rather than absolute expression. For each perturbation:

1. Feature construction — cross-perturbation expression profile + control statistics per gene. For double perts, features are averaged across the two constituent genes.
2. Cross-attention transformer — attends over all training perturbation deltas to predict the mean expression shift. Query and key paths both fuse expression features with scGPT embeddings.
3. Confidence calibration — attention entropy + embedding outlier distance calibrate predictions toward the population mean for out-of-distribution genes.
4. Retrieve-and-Shift — samples cells from KNN-selected training perturbations and affine-transforms them to match the predicted mean and HVG standard deviations.
5. Proportion prediction — neural proportion head blended with KNN lookup on predicted expression deltas.
6. Ensemble of 16 models — 8 on full data + 8 with leave-one-out validation, weighted by held-out Pearson correlation.

See Method description.md for the full write-up.

---

Repository Structure

.
├── main.py                  # All model code: train(), infer(), evaluate()
├── Method description.md    # Full method and rationale write-up
├── notebook.ipynb           # Development notebook
├── requirements.txt         # Python dependencies
└── resources/               # Created at runtime
    └── embeddings/
        ├── gene_embeddings_combined.npy   # [NOT TRACKED] 512-dim scGPT embeddings
        └── gene_embedding_order.txt       # [NOT TRACKED if large]

---

Data

Requires the Broad Obesity-2 Challenge dataset (not included — available from CrunchDAO):

File	Description
data/obesity_challenge_2.h5ad	scRNA-seq matrix with single- and double-gene perturbations + NC / NC+NC controls
data/program_proportion.csv	Cell program annotations (pre_adipo, adipo, lipo, other) per perturbation
data/predict_perturbations_2.txt	List of perturbations to predict at inference time
data/predict_genes_2.txt	List of output genes (defaults to all if absent)
data/obesity_challenge_2_local_gtruth.h5ad	Local ground truth for evaluate() (optional)
data/program_proportion_local_gtruth.csv	Ground truth proportions for evaluate() (optional)

---

Setup

pip install -r requirements.txt

Gene embeddings — extract from the scGPT whole-human checkpoint before training:

python extract_foundation_embeddings.py \
    --gene_list data/obesity_challenge_2.h5ad \
    --out_dir resources/embeddings/

This produces resources/embeddings/gene_embeddings_combined.npy and gene_embedding_order.txt.

---

Usage

Train

from main import train
train(
    data_directory_path  = "data/",
    model_directory_path = "resources/",
)
# Saves: resources/ensemble_v9_crunch2.pt

Infer

from main import infer
import pandas as pd, anndata as ad

perts = pd.read_csv("data/predict_perturbations_2.txt", header=None)[0].tolist()
genes = ad.read_h5ad("data/obesity_challenge_2.h5ad").var_names.tolist()

infer(
    data_directory_path              = "data/",
    prediction_directory_path        = "outputs/",
    prediction_h5ad_file_path        = "outputs/prediction.h5ad",
    program_proportion_csv_file_path = "outputs/predict_program_proportion.csv",
    model_directory_path             = "resources/",
    predict_perturbations            = perts,
    predict_genes                    = genes,   # renamed from genes_to_predict in Phase 1
)

Evaluate (requires local ground truth)

from main import evaluate
results = evaluate(
    data_directory_path   = "data/",
    model_directory_path  = "resources/",
    output_directory_path = "outputs/",
)

---

Key Hyperparameters

All hyperparameters are in the Config dataclass at the top of main.py:

Parameter	Value	Description
d_model	256	Transformer hidden dim
n_heads	8	Attention heads
n_layers	4	Transformer encoder layers
epochs	1200	Max training epochs
n_ensemble	16	Ensemble size (8 full + 8 LOO)
k_retrieve	10	KNN neighbours for Retrieve-and-Shift
cells_per_pert	100	Synthetic cells generated per perturbation
ctrl_types	["NC", "NC+NC"]	Both control types recognised

---

Official Leaderboard Scores

Scores from the CrunchDAO Broad Obesity-2 leaderboard (held-out test set — not reproducible locally):

Metric	Leaderboard Score
Pearson Delta	0.094
MMD	0.120
L1 Distance	0.160