dl4h final project (NetID: yoheis2, paper:bulk rna bert)

[yshibata8513]

---

Contributor: Yohei Shibata (yoheis2@illinois.edu, NetID: yoheis2) — solo submission
Contribution Type: Option 4 — Full Pipeline (Dataset + Task + Model)
Paper: Gelard, M. et al. (2025). BulkRNABert: Cancer prognosis from bulk RNA-seq based language models. PMLR 259. https://proceedings.mlr.press/v259/gelard25a.html

---

Summary

This PR is the first PyHealth implementation of BulkRNABert (Gelard et al., 2025), submitted as an Option 4 full-pipeline contribution (Dataset + Task + Model + Example). The model was independently re-implemented in PyTorch by referencing the specifications and numerical values from the paper, its published hyperparameter specs, and the reference JAX/Haiku implementation instadeepai/multiomics-open-research (CC BY-NC-SA 4.0). No source code, config, or YAML from the reference implementation was copied　or pasted. Numerical values not published in the paper were identified from publicly available artifacts of the reference implementation and used only as numerical references, not as copied source code, in an independently written implementation.

What is novel here:

• The paper evaluates only the discrete mode (64-bin tokenization). This PR adds a continuous expression mode as a new ablation axis and benchmarks it. On TCGA 5-cohort downstream classification it improves f1_weighted from 0.9436 to 0.9641 (+2.05pt) — a cheap improvement axis that the paper did not test.

PyHealth convention compliance: The downstream pipeline fully conforms to BaseModel / BaseDataset / BaseTask / Trainer. Only the pretraining side uses a custom loop, because long step-based training requires facilities PyHealth's Trainer does not currently expose (see §Design decisions).

---

Files to Review

Core implementation

• pyhealth/models/bulk_rna_bert.py — BulkRNABert / BulkRNABertConfig / BulkRNABertClassifier + encode() + binning helpers
• pyhealth/datasets/tcga_rnaseq_embedding.py + configs/tcga_rnaseq_embedding.yaml — TCGARNASeqEmbeddingDataset(BaseDataset)
• pyhealth/tasks/tcga_cancer_classification_5cohort.py — TCGACancerClassification5Cohort(BaseTask) + __call__

Examples (follow the {dataset}_{task_name}_{model}.py naming convention)

• examples/bulk_rna_bert/tcga_rnaseq_mlm_bulk_rna_bert.py — MLM pretraining CLI
• examples/bulk_rna_bert/tcga_rnaseq_extract_embeddings_bulk_rna_bert.py — ckpt → .npy embedding extraction CLI
• examples/bulk_rna_bert/tcga_cancer_classification_5cohort_bulk_rna_bert.py — downstream classification + --ablation mode (discrete vs continuous) + --synthetic-demo (for CI / smoke tests)

Tests (CPU only, synthetic data only, tempfile.TemporaryDirectory + tearDown cleanup, 53 cases / ~6.75 s)

• tests/core/test_bulk_rna_bert.py — 30 cases (model / tokenizer / MLM / Trainer integration / gene_embedding .pt save→load round-trip)
• tests/core/test_bulk_rna_bert_downstream.py — 23 cases (encode / classifier / task __call__ / BaseDataset e2e)

Docs

• docs/api/models/pyhealth.models.bulk_rna_bert.rst
• docs/api/datasets/pyhealth.datasets.tcga_rnaseq_embedding.rst
• docs/api/tasks/pyhealth.tasks.tcga_cancer_classification_5cohort.rst
• Corresponding index RSTs (docs/api/{models,datasets,tasks}.rst) updated

---

End-to-end pipeline

Running the following three scripts in order produces pretrain → embedding extraction → downstream classification:

1. Pretrain — examples/bulk_rna_bert/tcga_rnaseq_mlm_bulk_rna_bert.py: tcga_preprocessed.csv → step_{N}/{params.pt, config.json}
2. Extract — examples/bulk_rna_bert/tcga_rnaseq_extract_embeddings_bulk_rna_bert.py: ckpt + CSV → (n_samples, 256) float32 .npy
3. Classify — examples/bulk_rna_bert/tcga_cancer_classification_5cohort_bulk_rna_bert.py: .npy + identifier CSV + mapping CSV → test metrics

Each step is independently re-runnable because data is passed via CSV / .npy / ckpt directory, never in-memory handoff.

---

Ablation: discrete vs continuous expression mode

examples/bulk_rna_bert/tcga_cancer_classification_5cohort_bulk_rna_bert.py --ablation mode compares the two encodings while keeping the head MLP, split, and seed identical.

Setup: TCGA 5-cohort, 11,504 samples, seed=42, stratified 70/10/20, head MLP [256, 128] with SELU, Adam lr=1e-3, 1500 epochs, early stopping disabled, best ckpt selected on validation loss.
At pretraining time, following the paper, the initial weights of gene_embedding during pretraining were initialized from the values saved in the checkpoint of the reference implementation repository.

Setting	f1_weighted	f1_macro	accuracy
paper: MLP + IA3 fine-tune (5 seeds)	0.942 ± 0.004	0.918 ± 0.006	—
ours: discrete, head-only (1 seed)	0.9436	0.9306	0.9442
ours: continuous, head-only (1 seed)	0.9641	0.9526	0.9642

This PR's downstream setup is encoder-frozen + head-only MLP; IA3 fine-tuning is out of scope (see §Design decisions). Under the same recipe, continuous yields +2.05pt f1_weighted / +2.20pt f1_macro / −16% loss over discrete. This ablation is not evaluated in the paper.

Fairness caveats (not an apples-to-apples comparison with the paper) — treat the numbers above as a same-recipe discrete-vs-continuous comparison, not an absolute comparison against the paper:

• Sample count: paper 11,274 vs this PR 11,504 (a 230-sample difference from preprocessing filter differences; class distributions do not match exactly).
• Split: the paper does not disclose its train/val/test ratios or the seed set beyond "5 seeds", so the stratified 70/10/20 (seed=42) split used here cannot be reproduced as an identical subset.
• Evaluation pipeline: paper = MLP head + IA3 fine-tuning; this PR = head-only (no IA3).
• Seeds: this PR reports a single seed=42; the paper averages 5 seeds.

---

Design decisions

• Pretraining runs outside Trainer. Long step-based training needs gradient accumulation, step-based checkpoints, and a SIGTERM-safe stop path; Trainer(epochs=...) is too coarse for this. The custom loop is confined to examples/bulk_rna_bert/tcga_rnaseq_mlm_bulk_rna_bert.py, and the forward() → {"loss": ...} contract is honored on both the pretrain and downstream sides.
• Downstream is a two-stage pipeline: frozen encoder + cached .npy + head-only MLP. Encoder forward is an O(L²) attention over L = 19,062, which is heavy. Running the encoder once via the extract_embeddings CLI and caching embeddings to .npy lets the head be re-trained across multiple ablations / seeds without re-encoding. On-the-fly encode + head was deliberately not adopted.
• Weight initialization. A custom MultiHeadSelfAttention (Q/K/V/O = He-uniform) + FFN TruncatedNormal(std = 1/sqrt(fan_in)) + zero bias. With PyTorch's default nn.MultiheadAttention initialization, per-dim std collapses to ≈ 0.01 and produces a representation collapse where all samples yield near-identical embeddings, so downstream cannot separate classes. The custom initialization reproduces an equivalent scale to the reference JAX/Haiku implementation.
• Split. TCGA 5-cohort classification is a one-sample-one-prediction task where patient and sample are 1:1 in this dataset, so split_by_patient is not structurally applicable. The self-contained stratified_split_indices produces a train/val/test 3-way split (default 0.7 / 0.1 / 0.2).

---

Reproducibility caveats

• Seeds: single seed=42 (paper averages 5 seeds); no variance estimate reported.
• Hyperparameters: aligned with the public ckpt's actual values — dropout=0, layer_norm=False (differ from the paper text). Optimizer is AdamW (equivalent to the paper's Adam given weight_decay=0).
• gene_embedding initial values: only the gene_embedding layer (gene-ID embedding, analogous to positional encoding) is initialized from the three tensors gene_embedding.{embed.weight, proj.weight, proj.bias} inside the reference public ckpt params.joblib (CC BY-NC-SA 4.0), after which this PR's MLM pretraining is run on top. The mechanism is the pretrain CLI's --init-gene-embedding-from <path.pt> flag (in examples/bulk_rna_bert/tcga_rnaseq_mlm_bulk_rna_bert.py); the internal helper _load_gene_embedding_from_pt() copies only those three keys and explicitly ignores any other keys in the state_dict (attention / FFN / LM-head weights). All other parameters are trained from random initialization. The paper's reported numbers likewise assume gene_embedding is initialized from pretrained embeddings (gene2vec in the paper's case) before MLM pretraining, so the "initialize gene_embedding from pretrained values" precondition is shared between the paper and this PR. The ckpt file itself is not redistributed by this PR — users fetch it from the reference repository.
• Out of scope: IA3 fine-tuning / first-stage gene_embedding training from a gene2vec corpus / redistribution of pretrained weights.

---

Tests

.venv/bin/python -m pytest tests/core/test_bulk_rna_bert.py \
    tests/core/test_bulk_rna_bert_downstream.py -q
# 53 passed in 6.75s  (30 pretrain + 23 downstream, all CPU, synthetic only)

Coverage:
pretrain (forward/backward/mask ratio/binning), state save→load, round-trip, downstream (encode/classifier/stratified split/Trainer smoke), task & dataset e2e (__call__ edge cases / merged CSV / label ordering).
All tests use tempfile.TemporaryDirectory(), no real data / network, 2–5 patients per fixture.

---

[github-actions]

This PR has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs.