ST-JEPA: Joint-Embedding Predictive Architecture for Spatial Transcriptomics

Sebastian Birk, Amirhossein Vahidi, Mohammad Vali Sanian, Arpit Merchant, Mohammad Lotfollahi

Keywords: spatial transcriptomics, graph foundation model, jepa, self-supervised learning

TL;DR: ST-JEPA tokenizes spatial cell graphs into multi-scale transformer sequences and uses a JEPA objective to learn hierarchical embeddings that beat all baselines on niche identification while integrating across technologies.

Abstract: Biological function emerges from the interplay of molecular programs within individual cells and the spatial organization of cells within tissues. Spatial transcriptomics captures both layers simultaneously, yet existing computational methods model them at a single scale: task-specific approaches focus on neighborhood-level clustering, while recent foundation models such as Nicheformer omit spatial context entirely during pretraining or, like Novae, lack hierarchical representations spanning multiple biological resolutions. We introduce ST-JEPA, the first joint embedding predictive architecture for spatial transcriptomics. ST-JEPA operates across three biological scales through a multi-scale graph tokenization that converts spatial cell graphs into structured transformer sequences at gene, cell, and cellular neighborhood resolution, producing hierarchical embeddings suitable for diverse downstream tasks. Trained with a JEPA objective on mouse brain data spanning two technologies with non-overlapping gene panels, ST-JEPA achieves the best niche identification (weighted NMI=0.67) and the best batch integration (iLISI) among methods that perform well on niche identification, without explicit integration objectives. A technology metatoken enables cross-platform generalization even when gene panels do not overlap. Systematic ablations across six design axes provide practical guidance for self-supervised learning on spatial transcriptomics data.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 13