README.md

metadata

license: mit
datasets:
  - mims-harvard/SPATIA_MIST
tags:
  - ICML

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

Overview

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 them at niche and tissue levels, and generates cell morphology images conditioned on predicted state transitions using flow matching.

This repository hosts pretrained model weights, precomputed embeddings, evaluation results, and tutorial notebooks for reproducing the paper's prediction and generation experiments.

---

Repository Structure

mims-harvard/SPATIA/
├── embeddings/
│   └── spatia_embeddings.npy          # Precomputed cell embeddings (98.5 MB)
├── notebooks/
│   ├── prediction_tasks_showcase.ipynb # Clustering & annotation benchmarks
│   └── spatia_generation_inference.ipynb  # Flow-matching image generation
├── pretrained/
│   ├── best_model.pt                  # Pretrained SPATIA-scGPT weights (220 MB)
│   ├── args.json                      # Training configuration
│   ├── vocab.json                     # Gene vocabulary (30k+ genes)
│   └── all_dict_mean_std.csv          # Gene expression normalization statistics
├── results/
│   ├── table2_summary.csv             # Clustering benchmarks (Table 2)
│   └── table4_clustering_summary.csv  # scRNA-seq clustering (Table 4)
├── .gitattributes                     # Git LFS tracking rules
└── README.md

---

Components

pretrained/ — Model Weights & Config

Pretrained SPATIA-scGPT checkpoint for cell representation learning.

File	Description
best_model.pt	Model weights (12-layer transformer, 512-dim, trained on MIST atlas)
args.json	Full training hyperparameters (lr, batch size, masking ratios, etc.)
vocab.json	Gene-to-token vocabulary mapping
all_dict_mean_std.csv	Per-gene mean and std for input normalization

Usage:

import torch

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

embeddings/ — Precomputed Cell Embeddings

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)

notebooks/ — Tutorial Notebooks

Notebook	Description
prediction_tasks_showcase.ipynb	Reproduces clustering (Leiden) and cell-type annotation benchmarks from Tables 2 & 4. Runs on frozen embeddings — no GPU needed.
spatia_generation_inference.ipynb	Loads the generation model checkpoint, runs ODE-based flow matching to generate perturbation cell images, and evaluates with FID, SSIM, and morphology metrics.

Run locally:

pip install jupyter
jupyter notebook notebooks/prediction_tasks_showcase.ipynb

results/ — Evaluation Results

Precomputed benchmark results from the paper, ready for comparison.

File	Content
table2_summary.csv	Cell clustering (ARI, NMI) across Xenium & CosMx platforms for 6 models (PCA, scGPT, scFoundation, Nicheformer, UCE, SPATIA)
table4_clustering_summary.csv	scRNA-seq clustering benchmarks (ARI, NMI) for PCA, scGPT, Geneformer

.gitattributes — Git LFS Configuration

Configures Git Large File Storage for binary files (.pt, .npy, .h5, .ckpt, etc.). Required for cloning — install Git LFS before pulling:

git lfs install
git clone https://huggingface.co/mims-harvard/SPATIA

Or download individual files via the HuggingFace Hub:

from huggingface_hub import hf_hub_download

# Download pretrained weights
model_path = hf_hub_download(repo_id="mims-harvard/SPATIA", filename="pretrained/best_model.pt")

# Download embeddings
emb_path = hf_hub_download(repo_id="mims-harvard/SPATIA", filename="embeddings/spatia_embeddings.npy")

---

Quick Start

from huggingface_hub import snapshot_download

# Download everything
snapshot_download(repo_id="mims-harvard/SPATIA", local_dir="./SPATIA")

# Or download only what you need
from huggingface_hub import hf_hub_download

# For prediction tasks (no GPU needed):
hf_hub_download(repo_id="mims-harvard/SPATIA", filename="embeddings/spatia_embeddings.npy", local_dir="./SPATIA")
hf_hub_download(repo_id="mims-harvard/SPATIA", filename="notebooks/prediction_tasks_showcase.ipynb", local_dir="./SPATIA")

---

Related Resources

• Code: github.com/mims-harvard/SPATIA
• Dataset: mims-harvard/SPATIA_MIST (MIST atlas for spatial transcriptomics)
• Documentation: mims-harvard.readthedocs.io
• Paper: arXiv 2507.04704

---

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