BioM-JEPA

BioM-JEPA is a self-supervised model for single-cell transcriptomics that learns cell representations by predicting the latent representation of a block of related genes rather than reconstructing individual expression values. It was developed at Westlake University with Chinese Academy of Sciences collaborators and posted as a preprint in August 2026.

The design starts from a measurement problem. Single-cell RNA sequencing samples each transcriptome incompletely, so two cells in the same biological state can contain different detected genes while technical factors such as library size dominate their apparent similarity. A missing count for one gene is weak evidence that its programme is inactive, but several connected genes carry partially redundant evidence for the same latent state. BioM-JEPA therefore builds a gene graph from protein-association and coexpression evidence, samples connected sets from it, and trains a student encoder to predict one aggregate representation per hidden block from the remaining genes in the cell — while an exponential-moving-average teacher, which sees the full observed gene set, supplies the target.

This places it alongside GeneJEPA as a joint-embedding predictive architecture for transcriptomes, but with a different prediction unit: GeneJEPA predicts the latent states of individually masked genes, BioM-JEPA a block-scale summary. The paper positions that choice against reconstruction-trained peers such as Geneformer, scGPT, scFoundation, scMulan and CellFM, arguing that what a model is asked to predict matters more for frozen-embedding quality than how much it reconstructs.

#Key Features

• Graph-connected blocks as prediction units: Targets are connected gene sets grown by breadth-first expansion over a graph built from STRING v12 associations and corpus-level coexpression, giving each target an identity independent of gene order.
• Cell-specific target support: The teacher target is pooled only over block genes actually detected in that cell, so supervision never rests on a single sparse measurement.
• No reconstruction decoder: Training optimizes block-level latent alignment with variance and covariance regularization, with no per-gene expression reconstruction loss.
• Linear attention throughout: Student, teacher and predictor contract keys and values before applying each query, so cost grows linearly with the number of observed genes and the full gene sequence is preserved rather than truncated.
• Label-free graph construction: No cell-type label, perturbation identity, pathway membership or regulatory direction enters the graph, so biological structure found in the embedding is learned rather than supplied.

#Technical Details

Student and teacher share a 12-layer, 768-dimensional linear-attention encoder with 12 heads per layer over a 19,264-gene vocabulary, followed by a four-layer predictor; the teacher is updated only by EMA, with momentum annealed from 0.996 to 0.9997. The gene graph retains at most 64 neighbours per gene from the union of STRING associations scoring at least 700 and the strongest coexpression neighbours, giving 1,123,337 directed edges at mean out-degree 58.31. Each cell draws four target blocks with requested sizes between 2,000 and 8,000 genes. Pretraining used a scBaseCount snapshot forming a 22.1-million-cell local corpus, analysed after 20,000 steps at an effective batch of 512 cells — 10.24 million cell presentations, below the reported corpus sizes of its reconstruction-trained peers. Cells are pooled into a 1,536-dimensional frozen embedding. Across CellBench, block prediction reached a mean Top-5 few-shot macro-F1 of 0.835 on hPancreas and cortex against 0.817 for token-level prediction, 0.812 for random blocks and 0.797 for decoder-only reconstruction, with hPancreas effective rank rising from about 19 to nearly 40. On Adamson, Norman and Dixit it gave the lowest aggregate log-fold-change MSE (10.31 versus 10.87–10.98) and recovered 0.123 of the top 50 response genes. In a matched one-epoch hPancreas run with 26.09 million trainable parameters on one A100, it processed 24.37 training cells per second versus 4.24 for scFoundation.

#Applications

The intended use is as a frozen encoder: computational biologists embed a dataset once and reuse the representation for few-shot cell-type annotation, expression and Reactome-programme reconstruction, neighbourhood-preserving analysis, and perturbation-response prediction. Two analyses point further — ablating a graph programme displaces the matching cell identity far more than an expression-matched random gene set, and latent displacements for double perturbations agree with the sum of their single constituents, recovering pairs such as MAP2K3–MAP2K6 and CEBPA–CEBPE — giving a representation-level route to prioritizing candidate gene relationships for experimental follow-up.

#Impact

The contribution is a concrete answer to what a transcriptomic JEPA should predict, plus evidence that the answer matters: holding encoder, vocabulary and pretraining data fixed, changing only the prediction unit moved effective rank, depth coupling and probe accuracy together. The authors are explicit about scope: the pretraining collection does not retain study accessions, so study-level overlap with the evaluation benchmarks cannot be excluded and annotation is reported as within-dataset frozen-probe performance; graph-connected blocks are operational targets, not curated pathways. The work is a preprint awaiting peer review, and the checkpoint, training and inference software, fixed graph and AnnData interface are announced for GitHub release but not yet public — independent replication and adoption both wait on it.