MSGR: Multi-Scale Gene Refiner

This repository contains the official implementation of MSGR (Multi-Scale Gene Refiner), a hierarchical model for predicting spatial gene expression from histopathology images.

Overview

MSGR leverages biological priors — such as the Gene Ontology (GO) hierarchy — to organize genes into a tree structure and predict expression coarse-to-fine across multiple scales. The model uses a shared Transformer backbone with adaptive layer normalization (AdaLN), L2-normalized attention, and a cross-scale latent highway to refine predictions from root-level summaries down to individual genes.

Repository Structure

MSGR-pub/
├── README.md
├── doc/
│   └── msgr_arch.png               # Architecture diagram 
scripts
│
├── src_msgr/                       # Core MSGR implementation
│   ├── configs/                    # Training configurations
│   ├── dataset/
│   │   ├── __init__.py
│   │   └── hest_dataset.py         # Spatial transcriptomics dataset loader
│   ├── model/
│   │   ├── __init__.py             # Model registry
│   │   ├── model_interface.py      # PyTorch Lightning wrapper
│   │   ├── model_utils.py          # Model loading & training utilities
│   │   ├── model_metrics.py        # PCC / MSE / MAE metrics
│   │   ├── uni_baseline.py         # UNI embedding baseline
│   │   └── MSGR/
│   │       ├── __init__.py
│   │       ├── multi_scale_residual.py   # Main MSGR model
│   │       ├── backbone.py               # Transformer blocks
│   │       ├── spatial_transformer.py    # E(2)-equivariant neighborhood encoder
│   │       └── single_scale_target.py    # SSGT ablation baseline
│   ├── preprocess/
│   │   ├── __init__.py
│   │   ├── build_hierarchy.py      # Build GO/pathway/PPI/coexp/random hierarchies
│   │   └── extract_uni_embeddings.py
│   ├── train.py                    # Training entry point
│   ├── inference.py                # Inference entry point
│   ├── config_loader.py            # YAML configuration loader
│   └── utils.py                    # Logging, callbacks, seed helpers
│
└── src_ms/                         # Plug-in variants
    └── model/
        ├── EGN/
        ├── EGN_MS/                 # EGN + MS decoder
        ├── st_net/
        ├── st_net_ms/              # ST-Net + MS decoder
        ├── stflow/
        └── stflow_ms/              # STFlow + MS decoder

Installation

1. Clone the repository:

git clone https://github.com/NozomiMizore/MSGR
cd MSGR

2. Install dependencies:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install pytorch-lightning scanpy anndata pandas numpy scipy h5py pillow pyyaml addict

3. Optional: install gprofiler-official for GO hierarchy construction:

pip install gprofiler-official

Data Preparation

Expected Dataset Layout

/root/autodl-tmp/dataset/hest1k/{DATASET_NAME}/
    processed_data/
        selected_gene_list_{num_genes}.txt
        go_hierarchy_{num_genes}.json          # or other hierarchy prefix
        uni_embeddings/                          # pre-computed patch embeddings
    splits/
        train_{i}.csv
        test_{i}.csv
    st/
        {slide_id}.h5ad

Build the Hierarchy

Use src_msgr/preprocess/build_hierarchy.py to construct the hierarchy JSON:

from src_msgr.preprocess.build_hierarchy import HierarchyBuilder

builder = HierarchyBuilder()
builder.build_hierarchy_for_dataset(
    data_path="/root/autodl-tmp/dataset/hest1k/CCRCC/processed_data",
    num_genes=200,
    method="go",  # or "random", "pathway", "ppi", "coexp"
)

Supported hierarchy methods:

Method	Prefix	Description
go	go_hierarchy	Gene Ontology biological priors
random	random_hierarchy	Random hierarchy ablation
pathway	pathway_hierarchy	Pathway gene sets
ppi	ppi_hierarchy	Protein-protein interaction networks
coexp	coexp_hierarchy	Co-expression networks

Training

MSGR (Default)

export PYTHONPATH="/root/autodl-tmp/MSGR:$PYTHONPATH"
export CUDA_VISIBLE_DEVICES=0

python src_msgr/train.py --dataset CCRCC --gpus 1 --config ./src_msgr/configs/train_msgr.yaml

Multiple Datasets

python src_msgr/train.py --dataset COAD,HCC,LUNG,READ,SKCM \
    --gpus 1 --config ./src_msgr/configs/train_msgr.yaml

Available Configs

Config	Purpose
train_msgr.yaml	MSGR with GO hierarchy (default)
train_ssgt.yaml	Single-scale baseline
train_msgr_random.yaml	Random hierarchy ablation
train_msgr_pathway.yaml	Pathway hierarchy
train_msgr_ppi.yaml	PPI hierarchy
train_msgr_coexp.yaml	Co-expression hierarchy
train_ablation_direct.yaml	Direct prediction (no residual correction)
train_ablation_no_highway.yaml	No cross-scale latent highway

Training automatically detects k-fold splits under splits/train_{i}.csv and splits/test_{i}.csv.

Inference

python src_msgr/inference.py \
    --dataset CCRCC \
    --slide_id MEND144 \
    --ckpt_path logs/CCRCC/MSGR_GO/.../checkpoints/best-...ckpt \
    --config ./src_msgr/configs/infer_msgr.yaml \
    --gpu_id 0

Note: create an infer_msgr.yaml from your training config if it does not exist.

Model Registry

Key	Class	Description
MSGR	MultiScaleGeneResidual	Main hierarchical model
SSGT	SingleScaleGeneTarget	Single-scale ablation
UNI_BASELINE	UNIBaseline	2-layer MLP on frozen UNI embeddings

Plug-in Variants

The src_ms/ directory adapts the MSGR hierarchical decoder to other backbones:

• STFlow + MS: src_ms/model/stflow_ms/model/denoiser_no_gene.py
• EGN + MS: src_ms/model/EGN_MS/EGN_MS.py
• ST-Net + MS: src_ms/model/st_net_ms/st_net_ms.py

Each plug-in shares the same hierarchical target aggregation and default loss weights [0.0, 1.0, 2.0, ...] as MSGR.

Key Design

• Hierarchical decoding: predictions start at the virtual root and refine through GO domains/terms down to leaf genes.
• Residual correction: each scale predicts a residual added to the expanded parent prediction.
• Cross-scale latent highway: decoder blocks can attend to latent states from the previous scale.
• Multi-scale MSE loss: per-scale supervision with weights [0.0, 1.0, 2.0, 4.0] (normalized by default).
• Primary metric: mean Pearson Correlation Coefficient (PCC) per gene.

Citation

If you use this code in your research, please cite our paper:

@article{xu2026gene,
  title={Gene Ontology-Guided Hierarchical Spatial Gene Expression Prediction from Histopathology Images},
  author={Xu, Zhiwen and Yan, Xiaoming and Wu, Chengkun and Chen, Juan and Chi, Haoang and Xu, Liyang},
  journal={arXiv preprint arXiv:2608.00405},
  year={2026}
}

License

This project is released for research purposes. Please refer to the paper and code for usage terms.

Acknowledgements

• Histology embeddings are computed with the UNI encoder.
• Spatial transcriptomics data follows the HEST-1k benchmark layout.