spatialpeft 0.1.0

Parameter-efficient fine-tuning for spatial transcriptomics foundation models

SpatialPEFT

A Parameter-Efficient Fine-Tuning Framework for Spatial Transcriptomics Foundation Models

SpatialPEFT enables fine-tuning of spatial transcriptomics foundation models (Geneformer, Nicheformer, scGPT) on a single consumer-grade GPU with 16 GB VRAM, reducing peak training memory by up to 89% through Low-Rank Adaptation (LoRA), gradient checkpointing, and mixed-precision training.

Key Results

Dataset	Task	Method	Metric	Value
Xenium Breast Cancer (576,342 cells)	Cell type annotation	SpatialPEFT-LoRA r=4	Macro F1	0.9586
Xenium Breast Cancer (576,342 cells)	Cell type annotation	Zero-shot	Macro F1	0.7104
DLPFC Brain 12-slice (47,329 spots)	Spatial domain ID	SpatialPEFT-LoRA r=32	NMI	0.464
DLPFC Brain 12-slice (47,329 spots)	Spatial domain ID	Zero-shot	NMI	0.300

Peak VRAM: 2.17 GB (LoRA r=4, Geneformer-316M, batch=32, seq=256) Trainable parameters: 0.20% of total (633,227 / 316,995,840)

Installation

pip install spatialpeft

Or from source:

git clone https://github.com/applerplay/SpatialPEFT.git
cd SpatialPEFT
pip install -e .

Quick Start

import spatialpeft

# Load model with 4-bit quantization
model = spatialpeft.load("geneformer-316M", load_in_4bit=True)

# Inject LoRA adapters
model.inject_lora(r=4, target_modules=["query", "value"])

# Fine-tune on spatial AnnData
model.fit(adata, task="cell_type_annotation", spatial_key="spatial")

Framework Architecture

.h5ad / Zarr
    ↓
Data Engine Layer      — multi-platform preprocessing, rank tokenization, HDF5/Zarr lazy loading
    ↓
Model Hub Layer        — Geneformer / Nicheformer / scGPT, 4-bit/8-bit quantization
    ↓
PEFT Injection Layer   — LoRA on Q & V projections, spatial-aware coordinate adapter, gradient checkpointing
    ↓
Spatial Task Layer     — cell type annotation / spatial domain ID / neighborhood composition / cell density
    ↓
JSON Evaluation Report

Supported Models

Model	Parameters	Weight Format	LoRA Strategy	Peak VRAM (r=8)
Geneformer-316M	316.3M	HuggingFace .bin	HF PEFT standard	2.18 GB
Nicheformer-36M	36.4M	Lightning .ckpt	Manual QKV split	0.36 GB
scGPT-50M	50.3M	PyTorch .pt	Manual QKV split	0.41 GB

LoRA Rank Selection Guide

Based on empirical analysis across two benchmark datasets with different spatial resolutions:

Data type	Resolution	Recommended rank	Rationale
Xenium / MERFISH	True single-cell	r=4	Accuracy saturates at r=4; higher ranks waste parameters
Visium / spot-level	5–15 cells/spot	r=16–32	Accuracy improves continuously; no saturation observed

Reproducing Results

# Xenium breast cancer benchmark
python experiments/experiment_xenium.py

# DLPFC 12-slice brain benchmark
python experiments/experiment_dlpfc.py

# Nicheformer LoRA compatibility test
python experiments/test_nicheformer_lora.py

# scGPT LoRA compatibility test
python experiments/test_scgpt_lora.py

Hardware: Intel i9-14900KF · 64 GB DDR5 · NVIDIA RTX 4080 Super (16 GB VRAM) · WSL2 Ubuntu 22.04 · CUDA 12.6

Gradient Checkpointing VRAM Benchmark

Configuration	Batch	Peak VRAM	Feasible?
Full fine-tune (no PEFT)	2	~32 GB	No — OOM
LoRA, GC off	2	16.84 GB	Marginal
LoRA + GC on	2	2.15 GB	Yes — 87.2% saved
LoRA + GC on	4	3.64 GB	Yes — 89.0% saved

Citation

@article{spatialpeft2025,
  title={{SpatialPEFT}: A Parameter-Efficient Fine-Tuning Framework for Spatial Transcriptomics Foundation Models},
  author={[Author names]},
  journal={[Journal]},
  year={2025}
}

License

MIT License — see LICENSE for details.