AMIST

Adaptive Multimodal Integration for Spatial Transcriptomics

A robust multimodal deep learning framework for spatial niche identification, seamlessly integrating gene expression and H&E histology via bidirectional cross-attention, adaptive gating, and masked graph autoencoding.

---

📖 Table of Contents

• Overview
• Key Contributions
• Architecture
• Installation
• Data Preparation
• Quick Start
• Repository Structure
• Citation
• License

---

🌟 Overview

Spatial Transcriptomics (ST) uniquely enables the measurement of gene expression while preserving spatial tissue context. However, relying on a single modality is challenging: transcriptomic profiles suffer from dropout noise, while morphological features (H&E images) often lack deep molecular specificity.

AMIST bridges this gap by harmonizing both modalities through an adaptive fusion mechanism. It leverages a spatial graph autoencoder to inherently preserve tissue topology, enabling highly accurate and biologically meaningful spatial niche identification.

Training Paradigm

Our framework utilizes a highly optimized Two-Stage Pipeline:

1. Contrastive Alignment: Aligns gene and histology embeddings in a shared latent space using a symmetric NT-Xent contrastive loss.
2. Spatial Graph Fine-Tuning: Fine-tunes the multi-modal embeddings using a Masked Graph Autoencoder (GAE) to reconstruct spatial neighborhoods, deeply capturing structural tissue morphology.

---

🚀 Key Contributions

• Foundation Model Synergy: Leverages robust pre-trained foundation models (scGPT-human for transcriptomics, UNI2-h or standard ViT for histology). By pre-computing embeddings offline, AMIST achieves massive computational savings during training.
• Bidirectional Cross-Attention: Facilitates context-aware, bidirectional feature aggregation between molecular profiles and morphological image patches.
• Adaptive Modality Gating: Dynamically assesses the reliability of each modality at every spatial spot. It shifts reliance automatically (e.g., heavily weighting histology in regions with sparse RNA reads).
• Spatial Consistency via GAE: Explicitly models tissue connectivity using a kNN graph, drastically improving the coherence of the identified spatial domains.

---

🏗️ Architecture

AMIST is designed for scalability and rapid iteration. Pre-computing foundation model features reduces per-epoch training time significantly while retaining deep representational power.

OFFLINE (once per sample — run scripts/extract_embeddings.py):
  scGPT(adata.X)  →  adata.obsm['X_scgpt']  (N, 512)  — frozen backbone
  ViT(patches)    →  adata.obsm['X_uni']    (N, 768)  — frozen backbone

┌─────────────────────────────────────────────────────────────────────────────────────────────┐
│                         AMIST Training Forward Pass (forward_from_features)                 │
│                                                                                             │
│  PRE-COMPUTED FEATURES  (loaded from adata.obsm, never recomputed during training)          │
│  ┌──────────────────────┐        ┌───────────────────────┐                                  │
│  │  feat_g (N, 512)     │        │  feat_i (N, 768)      │                                  │
│  │  scGPT CLS embedding │        │  ViT backbone output  │                                  │
│  └──────────┬───────────┘        └────────────┬──────────┘                                  │
│             │                                  │                                            │
│  ┌──────────▼───────────┐        ┌────────────▼──────────┐                                  │
│  │  proj_g  (trainable) │        │  proj_i  (trainable)  │
│  │  Linear(512→512)     │        │  Linear(768→768)      │
│  │  GELU + LayerNorm    │        │  GELU + LayerNorm     │                                  │
│  │  Linear(512→256)     │        │  Linear(768→256)      │                                  │
│  │  + LayerNorm         │        │  + LayerNorm          │                                  │
│  └──────────┬───────────┘        └────────────┬──────────┘                                  │
│             │  z_g (N, 256)                    │  z_i (N, 256)                              │
│             │                                  │                                            │
│  ┌──────────▼──────────────────────────────────▼──────────┐                                 │
│  │            BidirectionalCrossAttention                 │                                 │
│  │                                                        │                                 │
│  │   z_g_att = Attn(Q=z_g, K=z_i, V=z_i)  GE←Hist         │                                 │
│  │   z_i_att = Attn(Q=z_i, K=z_g, V=z_g)  Hist←GE         │                                 │
│  │   + residual + LayerNorm on each side                  │                                 │
│  └──────────┬──────────────────────────────────┬──────────┘                                 │
│             │  z_g_att (N, 256)                │  z_i_att (N, 256)                          │
│             │                                  │                                            │
│  ┌──────────▼──────────────────────────────────▼──────────┐                                 │
│  │                  AdaptiveGate                          │                                 │
│  │                                                        │                                 │
│  │   [modality dropout p=0.3 during training]             │                                 │
│  │   g = sigmoid(W [z_g_att ; z_i_att] + b)               │                                 │
│  │   z_f = g ⊙ z_g_att + (1-g) ⊙ z_i_att                  │                                 │
│  │   g→1: rely on transcriptomics (noisy histology)       │                                 │
│  │   g→0: rely on histology (strong morphology signal)    │                                 │
│  └────────────────────────┬────────────────────────────────┘                                 │
│                           │  z_f (N, 256)  +  gate_vals (N, 256)                            │
│                           │                                                                 │
│  ┌────────────────────────▼────────────────────────────────┐                                 │
│  │              MaskedGraphAutoEncoder  (Stage 2 only)     │                                 │
│  │                                                         │                                 │
│  │   Spatial kNN graph G=(V,E), k=6   — full-graph pass    │                                 │
│  │   Mask ratio=0.3 → replace with learnable mask token    │                                 │
│  │   GATConv × 3 layers (4 heads, residual + LayerNorm)    │
│  │   → h (N, 256)   gradient checkpointing per layer       │                                 │
│  │   Decoder → z_hat  →  L_recon on masked nodes only      │                                 │
│  └────────────────────────┬────────────────────────────────┘                                 │
│                           │  h (N, 256)  ← final embeddings                                 │
│                           │                                                                 │
│  OUTPUT                   ▼                                                                 │
│  ┌────────────────────────────────────────────────────────┐                                 │
│  │  Leiden Clustering on h  →  Spatial Niches             │                                 │
│  └────────────────────────────────────────────────────────┘                                 │
└─────────────────────────────────────────────────────────────────────────────────────────────┘


LOSS FUNCTIONS
──────────────
  Stage 1   L_contrastive   NT-Xent on (z_g, z_i) pairs — same spot = positive, others = negative
  Stage 2   L_recon         MSE reconstruction of masked nodes in GAE
            L_reg           L2 regularization on h
            ─────────────────────────────────────────────────────────
            L_stage2  =  λ1 · L_recon  +  λ2 · L_reg

TRAINING STAGES
───────────────
  Stage 1   Contrastive alignment  (mini-batch B=256, all samples per epoch)
              - Both scGPT and ViT backbones fully frozen (features pre-computed)
              - Train: proj_g, proj_i, cross-attention, adaptive gate
              - GAE: frozen and skipped (run_gae=False)
              - Loss: NT-Xent contrastive on (z_g, z_i)
              - feat_g / feat_i sliced by random mini-batch index each step;
                edge_index is passed but ignored since GAE does not run

  Stage 2   Spatial GAE fine-tune  (full-graph, all samples per epoch)
              - Both backbones remain fully frozen
              - proj_g, proj_i, cross-attention, adaptive gate: ALL FROZEN
              - Train: GAE only (learns spatial context from fixed z_f)
              - Loss: GAE masked node reconstruction + L2 reg on h
              - AMP (autocast + GradScaler) enabled to halve GATConv memory usage
              - Gradient checkpointing in GATConv layers to fit N~192K on GPU

WHY FULL-GRAPH FOR STAGE 2?
────────────────────────────
  The MaskedGraphAutoEncoder uses GATConv graph convolution on the spatial kNN graph.
  GATConv aggregates each cell's embedding from its k spatial neighbors, so all N cells
  and their edges must be present in a single forward pass. Mini-batching would produce
  edge indices pointing to cells outside the batch, making the graph convolution invalid.
  AMP (float16 autocast) + gradient checkpointing are used to fit the GATConv message
  tensor (E × heads × dim) on GPU at N=192K, k=6.

WHY PRE-COMPUTE?
────────────────
  Both scGPT and ViT backbones are fully frozen throughout training — their outputs
  never change. Running them every epoch wastes GPU compute (each backbone forward pass
  costs ~10–30s per sample). Pre-computing once and loading from disk saves this cost
  for every subsequent epoch, reducing training time from O(epochs × backbone_cost)
  to O(backbone_cost) + O(epochs × lightweight_cost).

---

🛠️ Installation

AMIST relies on robust open-source libraries. We recommend managing your environment with conda or venv.

# 1. Clone the repository
git clone https://github.com/biniyamlombe/AMIST.git
cd AMIST

# 2. Create and activate a virtual environment
python -m venv .venv
source .venv/bin/activate

# 3. Install AMIST and dependencies
pip install -e .

Pre-trained Foundation Models

AMIST leverages pretrained foundation models that need to be downloaded prior to training.

1. Transcriptomics: scGPT-human

pip install scgpt huggingface_hub

python - <<'EOF'
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id="bowang-lab/scGPT_human",
    local_dir="checkpoints/scGPT_human",
)
EOF

This downloads ~500 MB to checkpoints/scGPT_human/ containing weights, vocabulary, and configuration.

2. Histology: UNI2-h (Optional but Recommended)

Note: Requires accepting the model license on HuggingFace.

1. Visit MahmoodLab/UNI2-h to accept the license.
2. Generate an access token at huggingface.co/settings/tokens.
3. Download the model:

python - <<'EOF'
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id="MahmoodLab/UNI2-h",
    local_dir="checkpoints/UNI2-h",
    token="YOUR_HF_TOKEN",
)
EOF

Fallback: If UNI2-h is not accessible, you can set hist_encoder.backbone: "vit_base_patch16_224" in configs/default.yaml to use a standard ImageNet ViT.

---

📂 Data Preparation

Data is typically processed at the cell-level into .h5ad and .h5 objects. Example data structure for Prostate Adenocarcinoma (PRAD) samples:

/path/to/data/PRAD_cell/
├── TENX157/
│   ├── TENX157_cells.h5ad          # AnnData: (N, 2000) log-normalized HVGs + spatial coords
│   └── TENX157_cell_patches.h5     # HDF5: pre-extracted 128×128 H&E patches per cell
└── ...

Both files are identically ordered and row-aligned (adata.obs_names[i] matches patches['cell_ids'][i]). You can generate custom aligned data using scripts/preprocess_data_cell_level.py.

---

⚡ Quick Start

After installing dependencies and organizing your data, you can run AMIST:

# 1. Pre-compute offline backbone embeddings (Dramatically speeds up training)
python scripts/extract_embeddings.py \
    --data_dir /path/to/data/PRAD_cell \
    --scgpt_model_dir checkpoints/scGPT_human \
    --vit_backbone vit_base_patch16_224 \
    --batch_size 64 \
    --device cuda

# 2. Train the AMIST model end-to-end
python scripts/train.py \
    --config configs/default.yaml \
    --data_dir /path/to/data/PRAD_cell \
    --out_dir outputs/amist_run1

# 3. Test functionality with synthetic data (Smoke test)
python scripts/dry_run.py

For HPC clusters, SLURM scripts are readily available in scripts/slurm/.

---

📂 Repository Structure

AMIST is structured modularly. Here is the complete layout of the codebase:

amist/
├── README.md
├── requirements.txt
│
├── configs/                  # Configuration files
│   └── default.yaml          # Hyperparameters for all training stages
│
├── notebooks/                # Jupyter Notebooks for exploration and evaluation
│   ├── data_exploration.ipynb
│   ├── amist_eval_with_output.ipynb
│   └── processed_data_usage.md
│
├── src/                      # Core framework implementation
│   ├── data/
│   │   ├── dataset.py            # SpatialDataset: loads .h5ad, returns (gene, patch, coord)
│   │   ├── preprocessing.py      # HVG selection, normalization, log1p
│   │   ├── patch_extractor.py    # Per-spot H&E patch cropping + transforms
│   │   └── graph_builder.py      # Spatial kNN graph → edge_index COO tensor
│   │
│   ├── models/
│   │   ├── gene_encoder.py       # ScGPTGeneEncoder (default), GeneEncoder (ablation)
│   │   ├── hist_encoder.py       # HistologyEncoder (ViT + projection head)
│   │   ├── cross_attention.py    # BidirectionalCrossAttention implementation
│   │   ├── adaptive_gate.py      # AdaptiveGate (gating + modality dropout)
│   │   ├── graph_autoencoder.py  # MaskedGraphAutoEncoder (GATConv)
│   │   └── amist.py              # Core AMIST model + forward pass logic
│   │
│   ├── training/
│   │   ├── trainer.py            # AMISTTrainer: runs both stages, handles checkpoints
│   │   ├── losses.py             # NT-Xent contrastive, GAE recon, L2 regularizations
│   │   └── schedulers.py         # LR schedulers, staged optimizer setup
│   │
│   └── utils/
│       ├── io.py                 # Load/save AnnData, embeddings, checkpoints
│       └── logging.py            # Metric logging (console + wandb)
│
├── scripts/                  # Executable scripts
│   ├── train.py                  # CLI for stage 1 & 2 training
│   ├── extract_embeddings.py     # Offline backbone extraction
│   ├── run_baselines.py          # Baselines (gene_only, image_only, stagate)
│   ├── preprocess_data_cell_level.py # Data preparation pipeline
│   ├── blur.py                   # Histology blurring utility
│   ├── verify_blur.py            # Script to verify blurring efficacy
│   ├── dry_run.py                # Synthetic smoke test
│   └── slurm/                    # SLURM batch scripts for HPC clusters
│       ├── train.sh
│       ├── extract_embeddings.sh
│       ├── blur.sh
│       └── stagate.sh
│
├── logs/                     # Logging outputs (e.g., blur_verification.png)
└── results/                  # Generated outputs
    ├── checkpoints/          # Saved model weights
    ├── embeddings/           # Extracted and trained latent representations
    ├── figures/              # Generated visualizations
    └── metrics/              # Quantitative evaluation scores

---

📊 Evaluation & Baselines

AMIST naturally supports competitive benchmarking. You can execute models with ablated components directly:

python scripts/run_baselines.py --data_dir /path/to/data/PRAD_cell --model_type [gene_only|image_only|concat|stagate]

Extensive evaluation suites are available in the src/evaluation module and notebooks/ to validate clustering coherence, robustness, and biological relevance.

---

📝 Citation

If you use AMIST in your research, please cite it:

@article{AMIST2026,
  title={AMIST: Adaptive Multimodal Integration for Spatial Transcriptomics},
  author={Biniyam Lombe and Yundi Chen and Yifan Deng and Ethan Smith},
  year={2026}
}

---

📜 License

This project is licensed under the MIT License.