SPATIA: Multimodal Model for Prediction and Generation of Spatial Cell Phenotypes

Overview

Understanding how cellular morphology, gene expression, and spatial organization jointly shape tissue function is a central challenge in biology. Image-based spatial transcriptomics now provides high-resolution measurements of cell images and gene expression, but most methods analyze these modalities in isolation.

SPATIA is a multi-scale model for spatial transcriptomics that:

• Learns cell-level embeddings by fusing image-derived morphological tokens and transcriptomic tokens via cross-attention
• Aggregates embeddings at niche and tissue levels with transformer modules to capture spatial context
• Generates cell morphology images conditioned on predicted state transitions using flow matching

---

Installation

Requirements: Python 3.10, CUDA 12.1+, GPU with 40GB+ VRAM (A100/H100 recommended).

# 1. Create conda environment
conda env create -f environment.yml
conda activate spatia

# Or manually:
# conda create -n spatia python=3.10
# conda activate spatia
# pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 \
#     --index-url https://download.pytorch.org/whl/cu121

# 2. Install Flash Attention (required for the scPRINT transformer)
pip install flash-attn --no-build-isolation

# 3. Install data processing package (must be installed BEFORE scprint)
pip install -e data_processing

# 4. Install SPATIA-scprint
pip install -e gene_encoders/SPATIA-scprint

---

Usage: Models are saved to https://huggingface.co/mims-harvard/SPATIA

import torch
# Load pretrained weights
checkpoint = torch.load("pretrained/best_model.pt", map_location="cpu")
model.load_state_dict(checkpoint)

spatia_embeddings.npy contains cell embeddings extracted from the pretrained SPATIA model on the MIST atlas. These embeddings fuse morphological and transcriptomic information and can be used directly for downstream tasks without requiring a GPU.

Usage:

import numpy as np
embeddings = np.load("embeddings/spatia_embeddings.npy")
print(embeddings.shape)  # (num_cells, embedding_dim)

---

MIST Dataset Construction

See data_processing/README.md for the full MIST dataset construction pipeline (image cropping, h5ad annotation, LMDB merging, and data loading).

---

Representation Training (Stage 1)

Prerequisites

Before training, you need:

1. MIST dataset: constructed via the pipeline above (lamindb Collection + LMDB)
2. Base scPRINT checkpoint: download with cd gene_encoders/SPATIA-scprint/data/model && bash download.sh
3. Gene embeddings: data/generated/gene_embeddings.parquet (from the scPRINT pretrained model)
4. Biomart gene positions: data/main/biomart_pos.parquet (shipped in this repo)

Configure paths

Edit gene_encoders/SPATIA-scprint/config/0510_base_spatial_all_crossattention_medium.yml:

data:
  spatial_datadir: /your/path/to/lmdb/all.lmdb     # <- LMDB from Stage C
  collection_name: xenium_all_0212                   # <- lamindb Collection from Stage B
  gene_embeddings: ./data/generated/gene_embeddings.parquet
  do_gene_pos: ./data/main/biomart_pos.parquet
model:
  ckpt_path: ./data/model/scPRINT/medium.ckpt       # <- base scPRINT checkpoint
  clip_model_type: "facebook/vit-mae-base"           # ViT-MAE image encoder

Run training

# SPATIA-scprint (primary, requires lamindb + LMDB)
cd gene_encoders/SPATIA-scprint
scprint_spatial fit --config config/0510_base_spatial_all_crossattention_medium.yml \
    --model.ckpt_path data/model/scPRINT/medium.ckpt \
    --data.batch_size 8

# SPATIA-scgpt (alternative, uses accelerate for multi-GPU)
cd gene_encoders/SPATIA-scgpt
bash scgpt_spatial/0704_train_spatial_4h100.sh

Training uses PyTorch Lightning with bf16 mixed precision, DDP strategy, and Weights & Biases logging.

---

Prediction Tasks

Clustering and cell-annotation benchmarks (Tables 2 & 4). Runs on frozen embeddings — no GPU needed once embeddings are extracted.

See prediction_tasks/README.md for the full protocol: embedding extraction, dataset setup, resolution sweeps, and evaluation scripts.

---

Generation Task

OT-based data pairing (Stage 2) followed by flow-matching image generation (Stage 3).

See generative_tasks/spatia_flow/README.md for training config, data pairing, and inference details.

---

Citation

@article{kong2026spatia,
  title={Spatia: Multimodal model for prediction and generation of spatial cell phenotypes},
  author={Kong, Zhenglun and Qiu, Mufan and Boesen, John and Lin, Xiang and Yun, Sukwon and Chen, Tianlong and Kellis, Manolis and Zitnik, Marinka},
  booktitle = {Proceedings of the 43th International Conference on Machine Learning},
  year={2026}
}

Contact

• Zhenglun Kong — zhenglun_kong@hms.harvard.edu
• Marinka Zitnik — marinka@hms.harvard.edu

Zitnik Lab, Harvard Medical School