Modeling Multimodal

Multimodal Model

The multimodal model trains on gene expression + clinical features + MRI/mammography imaging using three encoders fused via a gated late-fusion module:

Encoder	Class	Input
Gene	RNABERTEncoder (RNA BERT 110M)	mRNA expression matrix
Clinical	FTTransformerEncoder	Clinical feature vector
Image (optional)	MRMGHierarchicalImageEncoder (FM-BCMRI 3D ViT)	DICOM MRI series

Patient overlap is small: only patients present in all modalities (gene ∩ clinical ∩ image) are included (inner join on patient_id).

---

Checkpoints

RNA BERT (gene encoder)

Property	Value
Model	ibm-research/biomed.rna.bert.110m.mlm.multitask.v1
Source	HuggingFace (gated)
Required	No — falls back to a linear projection automatically

Set HF_TOKEN in your environment or in .docker.env to enable RNA BERT. If unavailable, training continues with a linear projection in place of the transformer.

FM-BCMRI (image encoder)

Property	Value
Model	FM-BCMRI pretrained 3D ViT checkpoint
Source	Google Drive (FM-BCMRI project)
Required	No — falls back to random-initialised weights automatically
Default path	/workspace/models/breast_MR_checkpoint.pth.tar

Download breast_MR_checkpoint.pth.tar from the link above and place it in the models/ directory:

OncoLearn/
└── models/
    └── breast_MR_checkpoint.pth.tar

If the checkpoint is missing, the image encoder logs a warning and continues with random weights.

---

Config Files

Two ready-made multimodal configs are provided, differing only in the label task:

Config	Pipeline	Labels	Classes
tcga_brca_cbioportal_pam50.yaml	tcga_brca_cbioportal_pam50.py	PAM50 subtype (Basal/Her2/LumA/LumB/Normal)	5
tcga_brca_cbioportal_stage.yaml	tcga_brca_cbioportal_stage.py	AJCC pathologic stage (I/II/III/IV)	4

Both use cBioPortal as the data source (clinical + mRNA) and TCIA for imaging.

PAM50 config (tcga_brca_cbioportal_pam50.yaml)

model:
  name: oncolearn.model.multimodal.gated_late_fusion
  num_stage_classes: 5   # PAM50: Basal / Her2 / LumA / LumB / Normal
  num_subtype_classes: 0
  freeze_encoders: true
  dropout: 0.2
  modality_dropout_prob: 0.3
  encoders:
    - name: oncolearn.encoder.multimodal.RNABERTEncoder
      modality: oncolearn.modality.gene
      output_dim: 128
    - name: oncolearn.encoder.multimodal.FTTransformerEncoder
      modality: oncolearn.modality.clinical
      output_dim: 64
    - name: oncolearn.encoder.multimodal.MRMGHierarchicalImageEncoder
      modality: oncolearn.modality.image
      output_dim: 256
      checkpoint_path: /workspace/models/breast_MR_checkpoint.pth.tar

data:
  pipeline: data/configs/modeling/multimodal/preprocessing/tcga_brca_cbioportal_pam50.py
  splits_dir: data/configs/modeling/multimodal/splits/pam50/kfold/fold_0

training:
  max_epochs: 50
  batch_size: 8              # smaller batch recommended — imaging tensors are large
  num_workers: 4
  accelerator: auto
  devices: 1
  early_stopping_patience: 10
  use_class_weights: true
  seed: 42
  optimizer:
    name: torch.optim.AdamW
    params:
      lr: 0.0001
      weight_decay: 0.00001

output:
  dir: outputs
  experiment_name: tcga_brca_cbioportal_pam50
  save_every_n_epochs: 5

The data.pipeline field points to a Python pipeline file that defines which datasets to load and how to join them. See Pipeline DSL for details.

Key options

Option	Description
model.encoders[image].checkpoint_path	Path to the FM-BCMRI .pth.tar checkpoint
model.freeze_encoders	true to freeze RNA BERT and FM-BCMRI backbones
model.modality_dropout_prob	Randomly drop a modality during training; at least one modality is always kept
training.batch_size	Keep low (8–16) when imaging is enabled — image tensors are memory-intensive
training.use_class_weights	Inverse-frequency class weighting — recommended given the small patient overlap
data.splits_dir	Path to a folder with train.txt, test.txt, validation.txt for fixed K-fold splits

---

Data Prerequisites

# mRNA expression (cBioPortal)
oncolearn cbioportal download --cohorts BRCA

# Imaging (TCIA)
oncolearn tcia download --cohorts BRCA --yes

Expected layout after downloading:

data/
├── sources/
│   ├── cbioportal/
│   │   └── TCGA-BRCA/
│   └── tcia/
│       └── TCGA-BRCA/
│           └── TCIA_TCGA-BRCA_*/
│               └── TCGA-BRCA/

---

Generate K-Fold Splits

Before training, generate stratified K-fold patient splits. The command intersects gene ∩ clinical ∩ image patient sets and stratifies on the chosen label:

# 5-fold stratified splits — PAM50
oncolearn preprocess multimodal kfold 5 --stratified --label pam50

# 5-fold stratified splits — AJCC stage
oncolearn preprocess multimodal kfold 5 --stratified --label stage

Output is written to data/configs/modeling/multimodal/splits/<label>/kfold/fold_N/.

Then set data.splits_dir in your YAML config to one of the generated fold directories:

data:
  pipeline: data/configs/modeling/multimodal/preprocessing/tcga_brca_cbioportal_pam50.py
  splits_dir: data/configs/modeling/multimodal/splits/pam50/kfold/fold_0

See oncolearn preprocess multimodal kfold for the full reference.

---

Training

With a config file (recommended)

# PAM50
oncolearn train --config data/configs/modeling/multimodal/tcga_brca_cbioportal_pam50.yaml

# AJCC stage
oncolearn train --config data/configs/modeling/multimodal/tcga_brca_cbioportal_stage.yaml

In Docker (WSL2/AMD)

docker compose --profile prod-rocm-wsl run --rm prod-rocm-wsl \
  python -m oncolearn.trainer \
    --config data/configs/modeling/multimodal/tcga_brca_cbioportal_pam50.yaml

---

Expected Output

Epoch 2/50  ━━━━━━━━━━━━━━━━━ 4/4
  train_loss: 1.512
  val_loss:   1.601
  val_acc:    0.200

The effective dataset is small because only patients present in all modalities are included (inner join on patient_id).

Checkpoints are saved to outputs/tcga_brca_cbioportal_pam50/ (or tcga_brca_cbioportal_stage/):

• best_model.ckpt — best validation accuracy
• epoch_N.ckpt — periodic snapshots

---

See the Training guide for a full breakdown of arguments and output format, and the Docker guide for all available services.