Graph Tokenization Meets JEPA: Self-Supervised Learning on Spatial Cell Graphs
Abstract
Applying transformers to graph-structured biological data requires converting irregular topologies into sequential token representations—a challenge that current graph foundation models address through various tokenization and positional encoding schemes. We introduce ST-JEPA, a joint-embedding predictive architecture that operates on spatial cell graphs derived from spatial transcriptomics data. Our key contribution is a multi-scale graph tokenization that converts local subgraph neighborhoods in a k-nearest neighbor graph into structured transformer sequences at three biological resolutions—neighborhood, cell, and gene—using fixed sinusoidal segment embeddings to encode graph topology. Unlike graph contrastive methods that require augmentation design or masked autoencoders that reconstruct in data space, ST-JEPA predicts representations of masked graph regions entirely in latent space via block masking across the subgraph sequence. Trained on mouse brain data spanning two technologies with non-overlapping feature sets, ST-JEPA achieves the best niche identification (weighted NMI=0.67) and the strongest cross-technology integration among high-performing methods, without explicit alignment objectives. Systematic ablations across six design axes— graph tokenization strategy, feature encoding, positional encoding, sequence length, masking ratio, and normalization—provide practical guidance for self-supervised learning on attributed spatial graphs.