M3-Gen

Official repository for the paper: "M3-Gen: Interpretable Multimodal Generation of Gene Expression Profiles Using Clinical and Imaging Data"

M3-Gen learns to generate patient-level RNA-seq gene expression profiles conditioned on multimodal information extracted from TCGA:

• Whole-slide histopathology images (WSIs), tiled into patches and encoded with a foundation image encoder.
• Clinical/textual metadata, turned into short natural-language descriptions with an LLM and encoded with a text encoder.

A contrastive model first aligns image, text and RNA-seq representations in a shared embedding space; the resulting image/text embeddings are then used to condition a Wasserstein GAN (with gradient penalty) that generates synthetic gene expression profiles. Attention-based conditioning also yields an interpretable attention map over WSI patches for each generated sample.

The paper can be accessed here

Repository structure

data/                             GDC manifests + TCGA cohort table used to build the dataset
src/
├── dataset/
│   ├── download_data.py          Download WSIs + RNA-seq from GDC and assemble the dataset
│   ├── utils.py                  GDC API / cleaning helpers used by download_data.py
│   └── contrastive_dataset.py    Dataset/Dataloader for the contrastive model
├── preprocessing/
│   ├── patch_preprocessing.py    Otsu tissue segmentation + tiling of WSIs into patches
│   └── text_preprocessing.py     LLM-based generation of clinical text descriptions
│
├── contrastive_model.py, rnaseq_contrastive_model.py   Contrastive model architectures
├── train_contrastive.py                                Train image-text-RNA contrastive model
├── train_contrastive_rna_seq.py                        Train the RNA-seq branch only
├── train_contrastive_rna_seq_with_vae.py               Same, with a VAE bottleneck (see src/vae/)
├── preprocess_embeddings.py                             Precompute WSI patch embeddings with a trained image encoder
│
├── vanilla_generative_model.py            Unconditional WGAN-GP baseline
├── benchmark_generative_model.py          Class-conditional (tumor type) WGAN-GP baseline
├── generative_model_conditional_concat.py Multimodal GAN, text/image/mean-conditioning via concatenation
├── generative_model_conditional_attention.py   Multimodal GAN, attention-based conditioning (main M3-Gen model)
│
├── display_attention_weights.py, display_attention_weights_all.py   Attention map visualization
├── compare_de_real_vs_generated.py        Differential-expression + pathway-enrichment analysis
├── test_robustness_domain_shift.py        Robustness to WSI staining/blur/folding augmentations
├── visualize_augmentations.py             Visualize the augmentation pipeline used above
│
├── data_loader.py, multi_patch_gan_dataloader.py,
├── augmented_multi_patch_gan_dataloader.py, benchmark_gan_dataloader.py   Dataloaders
├── generative_model_utils.py, img_manip.py                                Shared model/image utilities
└── corr_score.py, distribution_distances.py, unsupervised_metrics.py,
    utility_evaluation.py, de_genes.py                                     Evaluation metrics (imported by the scripts above, not run standalone)

All entry points below are plain Python scripts with argparse CLIs — run them directly, there are no notebooks or shell wrappers in this repository.

1. Setup

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Notes:

• Python 3.10+ is required.
• requirements.txt pins CPU-agnostic versions of torch/torchvision. If you need CUDA wheels (recommended for training), install them from the official index for your CUDA version instead, e.g.:

  pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
  

• openslide-python needs the native OpenSlide library. openslide-bin (in requirements.txt) ships it as a wheel; on Linux you can alternatively install it via your package manager (apt install openslide-tools).
• text_preprocessing.py downloads a gated model from the Hugging Face Hub. Authenticate first — never hardcode tokens in scripts:

  huggingface-cli login   # or: export HF_TOKEN=...
  

2. Build the dataset

2.1 Download WSIs and RNA-seq data from GDC

python src/dataset/download_data.py \
    -output_path /path/to/dataset \
    -max_tissue_size 70000000 \
    -num_workers 8

This reads the manifests/cohort table in data/ (override with -data_dir <dir>), downloads matching tissue slides and RNA-seq files from the GDC API, and writes dataset_info.pkl, metainfos.pkl, case_ids.txt, rna_seq.parquet and slides_info.parquet under /path/to/dataset.

2.2 Tile whole-slide images into patches

python src/preprocessing/patch_preprocessing.py --dataset_path /path/to/dataset

Uses Otsu thresholding to find tissue regions and writes valid_patches.json (candidate 128/256/512/1024px tiles per case) into the dataset directory.

2.3 Generate clinical text descriptions

python src/preprocessing/text_preprocessing.py \
    --dataset_path /path/to/dataset \
    --model_name ContactDoctor/Bio-Medical-Llama-3-8B

Loads (4-bit quantized) an instruction-tuned LLM and turns each case's clinical metadata into a short descriptive paragraph, used later as the "text" modality.

3. Contrastive pretraining

Aligns image, text and RNA-seq embeddings in a shared space (image encoder: timm UNI foundation model; text encoder: a Hugging Face clinical language model).

python src/train_contrastive.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/contrastive_model \
    --num_epochs 10 \
    --batch_size 16 \
    --multi_patch --num_patches 8

RNA-seq-only variants (used for ablations / the VAE bottleneck experiments):

python src/train_contrastive_rna_seq.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/contrastive_rna_seq

python src/train_contrastive_rna_seq_with_vae.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/contrastive_rna_seq_vae

Run python src/<script>.py --help for the full list of hyperparameters (learning rate, embedding dims, freezing encoders, etc.).

Precompute WSI patch embeddings

Once a contrastive model is trained, precompute per-case image embeddings (used to condition the generative models below):

python src/preprocess_embeddings.py \
    --dataset_path /path/to/dataset \
    --model_path /path/to/experiments/contrastive_model/<run_id> \
    --num_workers 8

4. Train the generative models

All generative-model scripts share the same dataset/experiment layout and most CLI flags. Common flags:

Main model — attention-based multimodal conditioning

python src/generative_model_conditional_attention.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/attention_latent512 \
    --epochs 500 --batch_size 8 --latent_dims 512 --hidden_dims 256 \
    --num_patches 64 --num_workers 8 --efficient

To evaluate an already-trained checkpoint instead of retraining, add --test_only (the checkpoint is expected at <output_path>/generator_last_epoch.pt).

Ablation — concatenation-based conditioning

python src/generative_model_conditional_concat.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/concat_mean \
    --epochs 500 --batch_size 8 --latent_dims 256 \
    --concat_type mean \
    --efficient

--concat_type selects which modality is concatenated to the noise vector: text, image, or mean (average of the two).

Baselines

Class-conditional (tumor type only, no imaging/text):

python src/benchmark_generative_model.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/benchmark \
    --epochs 500 --batch_size 8 --latent_dims 256

Fully unconditional:

python src/vanilla_generative_model.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/vanilla \
    --epochs 500 --batch_size 64 --latent_dims 64

5. Interpretability — attention maps

python src/display_attention_weights.py \
    --dataset_path /path/to/dataset \
    --case_id <case_id> \
    --model_path /path/to/experiments/contrastive_model/<run_id> \
    --generator_path /path/to/experiments/attention_latent512/generator_last_epoch.pt \
    --output_path /path/to/experiments/attention

Omit --case_id and use display_attention_weights_all.py with the same other arguments to render attention maps for every case in the dataset.

6. Evaluation

6.1 Differential expression: Real vs. Generated

Exploratory biological-plausibility check: computes DE genes between two tumor types in real data and in data sampled from a trained generator, then compares the overlap (and, optionally, enriched pathways via Enrichr).

python src/compare_de_real_vs_generated.py \
    --dataset_path /path/to/dataset \
    --model_path /path/to/experiments/attention_latent512/generator_last_epoch.pt \
    --output_dir /path/to/experiments/de_real_vs_generated \
    --group_a BRCA --group_b LUAD \
    --latent_dims 512 --hidden_dims 256 --num_patches 64 \
    --device cuda \
    --run_enrichment --gene_set_library KEGG_2021_Human

Drop --run_enrichment if gseapy is not installed or you don't need the pathway-overlap step. --device cpu works too (just slower).

6.2 Robustness to WSI domain shift

Evaluates the attention-based generator when patches are corrupted with synthetic staining/focus/tissue-folding artifacts, and compares against a clean baseline.

python src/test_robustness_domain_shift.py \
    --dataset_path /path/to/dataset \
    --model_path /path/to/experiments/attention_latent512/generator_last_epoch.pt \
    --image_encoder_path /path/to/experiments/contrastive_model/<run_id> \
    --output_path /path/to/experiments/robustness_test \
    --latent_dims 512 --hidden_dims 256 --num_patches 64 --efficient \
    --embeddings_in_main --embedding_batch_size 64 \
    --blur_mode local --augmentation_probability 0.33 \
    --apply_stain_variation --apply_out_of_focus --apply_tissue_folding

Enable any subset of --apply_stain_variation, --apply_out_of_focus, --apply_tissue_folding to test that corruption in isolation or combined; a "clean" and an "augmented" run are always produced side by side under --output_path.

To preview what the augmentation pipeline actually does to a patch before running the full test:

python src/visualize_augmentations.py \
    --dataset_path /path/to/dataset \
    --output_dir /path/to/out \
    --num_patches 8 --augmentation_probability 1.0

Running on a SLURM cluster

Every script above is a regular python foo.py --args invocation, so submitting it to SLURM only requires a small generic wrapper — there is no repository-specific job script to maintain. For example:

#!/usr/bin/bash
#SBATCH --job-name m3gen
#SBATCH --gpus 1
#SBATCH --time 24:00:00
#SBATCH --output logs/%x_%j.out
#SBATCH --error  logs/%x_%j.err

module load CUDA/12.4.0
module load Python/3.12.3-GCCcore-14.2.0
source /path/to/venv/bin/activate

python src/generative_model_conditional_attention.py \
    --dataset_path /path/to/dataset \
    --output_path /path/to/experiments/attention_latent512 \
    --epochs 500 --batch_size 8 --latent_dims 512 --efficient

Save this as job.sh and submit with sbatch job.sh. If you need a Hugging Face token on the compute node (text_preprocessing.py), export it from your shell environment or ~/.bashrc — do not commit it to a script.

Citation

If you use this code, please cite:

@article{m3gen,
  title={M $\^{} 3$-Gen: Interpretable Multimodal Generation of Gene Expression Profiles Using Clinical and Imaging Data},
  author={Panaccione, Francesca Pia and Sgaravatti, Carlo and Venere, Marco},
  journal={arXiv preprint arXiv:2607.21343},
  year={2026}
}