BioM-JEPA: joint-embedding prediction of graph-connected gene blocks in single cells

Single-cell transcriptomics often suffers from sparsity and technical noise, where sequencing depth can dominate the apparent similarity between cells. Existing self-supervised models typically rely on reconstructing individual genes, which may not capture the coordinated biological programs that define cell identity. The authors introduce BioM-JEPA, a Joint-Embedding Predictive Architecture (JEPA) that shifts the prediction target from individual genes to aggregate representations of graph-connected gene blocks. These blocks are defined by integrating protein-association data (STRING) and corpus-derived coexpression evidence. The model uses a student-teacher framework where the student predicts the teacher's representation of a masked block based on the remaining observed genes, utilizing linear attention to maintain computational efficiency.

BioM-JEPA demonstrates superior performance compared to token-level prediction and reconstruction-based models. In diagnostic tests, BioM-JEPA embeddings exhibited higher effective rank and significantly lower correlation with sequencing depth, indicating that the model successfully disentangles biological signal from technical artifacts. Across CellBench tasks, the frozen embeddings retained high-fidelity information regarding continuous expression, pathways, and cell-state neighborhoods. Furthermore, BioM-JEPA achieved the lowest aggregate error in perturbation-response prediction tasks. The use of linear attention provided substantial throughput advantages, with 5.75-fold faster fine-tuning and 3.76-fold faster embedding extraction compared to existing models like scFoundation.

This study suggests that the unit of prediction in self-supervised transcriptomics is a critical design choice. By predicting graph-connected gene blocks, BioM-JEPA effectively captures the modular nature of biological programs. This approach not only improves the robustness of cell embeddings but also provides a more computationally efficient pathway for scaling to large-scale single-cell datasets. The ability of the frozen representation to reflect known biological relationships, such as synergistic genetic perturbations, highlights its potential as a versatile tool for hypothesis generation in systems biology.