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

arXiv cs.LG Papers

Summary

BioM-JEPA introduces a joint-embedding predictive architecture that learns single-cell representations by predicting graph-connected gene blocks instead of individual genes, showing improved efficiency and downstream performance in perturbation-response tasks.

arXiv:2608.05928v1 Announce Type: new Abstract: Single-cell transcriptomes are sparse observations of coordinated biological programmes, yet most self-supervised models learn by reconstructing individual genes. Here we present BioM-JEPA, a joint-embedding predictive architecture that instead predicts aggregate representations of graph-connected gene blocks defined by protein-association and corpus-derived coexpression evidence. A student network infers each target-block representation from the remaining genes in a cell, while a slowly updated teacher supplies the corresponding target from the full observed gene set. Under the reported extraction procedure, block-level prediction produced embeddings with higher effective rank and weaker association with detected-gene depth in the tested diagnostics than token-prediction, random-block and reconstruction controls. Across CellBench tasks, frozen BioM-JEPA embeddings retained expression, pathway and neighbourhood information and achieved the lowest aggregate perturbation-response error among the evaluated models. Representation diagnostics were also consistent with canonical pancreatic programmes and compositional relationships between genetic perturbations. Linear attention avoids constructing a quadratic gene-by-gene attention matrix; in a matched one-epoch hPancreas experiment at batch size 8, BioM-JEPA provided 5.75-fold higher fine-tuning throughput and 3.76-fold higher held-out embedding throughput than scFoundation. Together, these results support graph-connected gene blocks as useful prediction units for JEPA-style representation learning in single-cell biology.
Original Article
View Cached Full Text

Cached at: 08/07/26, 07:53 AM

# BioM-JEPA: joint-embedding prediction of graph-connected gene blocks in single cells
Source: [https://arxiv.org/html/2608.05928](https://arxiv.org/html/2608.05928)
Yuhao Wang1,∗, Zelin Zang1,2,∗,†, Yuxuan Liu1, Zhen Lei2,3,4,†, Stan Z\. Li1,† 1Westlake University, Hangzhou, Zhejiang, China2Centre for Artificial Intelligence and Robotics \(CAIR\), Hong Kong Institute of Science & Innovation, Chinese Academy of Sciences \(HKISI\-CAS\), Hong Kong SAR, China3State Key Laboratory of Multimodal Artificial Intelligence Systems \(MAIS\), Institute of Automation, Chinese Academy of Sciences \(CASIA\), Beijing, China4School of Artificial Intelligence, University of Chinese Academy of Sciences \(UCAS\), Beijing, China∗These authors contributed equally to this work\.†Correspondence: Zelin Zang \([zangzelin@westlake\.edu\.cn](https://arxiv.org/html/2608.05928v1/mailto:[email protected])\), Zhen Lei \([zhen\.lei@ia\.ac\.cn](https://arxiv.org/html/2608.05928v1/mailto:[email protected])\) and Stan Z\. Li \([Stan\.ZQ\.Li@westlake\.edu\.cn](https://arxiv.org/html/2608.05928v1/mailto:[email protected])\)\.

###### Abstract

Single\-cell transcriptomes are sparse observations of coordinated biological programmes, yet most self\-supervised models learn by reconstructing individual genes\. Here we present BioM\-JEPA, a joint\-embedding predictive architecture that instead predicts aggregate representations of graph\-connected gene blocks defined by protein\-association and corpus\-derived coexpression evidence\. A student network infers each target\-block representation from the remaining genes in a cell, while a slowly updated teacher supplies the corresponding target from the full observed gene set\. Under the reported extraction procedure, block\-level prediction produced embeddings with higher effective rank and weaker association with detected\-gene depth in the tested diagnostics than token\-prediction, random\-block and reconstruction controls\. Across CellBench tasks, frozen BioM\-JEPA embeddings retained expression, pathway and neighbourhood information and achieved the lowest aggregate perturbation\-response error among the evaluated models\. Representation diagnostics were also consistent with canonical pancreatic programmes and compositional relationships between genetic perturbations\. Linear attention avoids constructing a quadratic gene\-by\-gene attention matrix; in a matched one\-epoch hPancreas experiment at batch size 8, BioM\-JEPA provided 5\.75\-fold higher fine\-tuning throughput and 3\.76\-fold higher held\-out embedding throughput than scFoundation\. Together, these results support graph\-connected gene blocks as useful prediction units for JEPA\-style representation learning in single\-cell biology\.

Single\-cell RNA sequencing measures each cell through an incomplete, depth\-dependent sample of its transcriptome\. Two cells in the same biological state can therefore contain different observed genes, while technical quantities such as library size and the number of detected genes can dominate their apparent similarity\[[13](https://arxiv.org/html/2608.05928#bib.bib1)\]\. The biological processes of interest are more stable than these individual measurements\. Cell identity, activation and response to perturbation are expressed through coordinated programmes in which many genes contribute partially redundant evidence\. A useful representation\-learning system must distinguish this programme\-level signal from the stochastic sampling process that produced the observed counts\.

Large pretrained models have made it possible to learn from transcriptomes collected across many tissues, studies and experimental conditions\. Geneformer represents cells as ranked gene contexts, scGPT uses generative masked modelling, and scFoundation and scMulan combine value\-aware representations with reconstruction or multitask training\[[20](https://arxiv.org/html/2608.05928#bib.bib4),[6](https://arxiv.org/html/2608.05928#bib.bib5),[9](https://arxiv.org/html/2608.05928#bib.bib6),[5](https://arxiv.org/html/2608.05928#bib.bib7)\]\. CellFM extends pretraining to 100 million cells, while Nicheformer incorporates both dissociated and spatial transcriptomes\[[23](https://arxiv.org/html/2608.05928#bib.bib8),[19](https://arxiv.org/html/2608.05928#bib.bib9)\]\. These models demonstrate the potential of broad transcriptomic pretraining\. They also expose an unresolved design question: what should a model predict in order to learn a reusable cell representation? Larger models and lower reconstruction error do not consistently produce better frozen embeddings, and recent evaluations have found substantial variation across tasks and datasets\[[11](https://arxiv.org/html/2608.05928#bib.bib10),[7](https://arxiv.org/html/2608.05928#bib.bib11),[1](https://arxiv.org/html/2608.05928#bib.bib12)\]\.

Joint\-embedding predictive architectures \(JEPAs\) provide an alternative to reconstructing the original observation\. A JEPA learns by predicting a target representation from related context\[[12](https://arxiv.org/html/2608.05928#bib.bib13)\]\. In I\-JEPA, the context surrounding a masked image region is used to predict the teacher representation of that region; V\-JEPA extends this principle to spatiotemporal regions in video\[[2](https://arxiv.org/html/2608.05928#bib.bib14),[4](https://arxiv.org/html/2608.05928#bib.bib15)\]\. In both settings, the target is a coherent part of the underlying scene\. A direct translation to transcriptomics would treat each masked gene as an independent target\. However, a single gene measurement is not the molecular equivalent of an image region: it is sparse, noisy and often biologically ambiguous when separated from its programme\.

These considerations led us to ask whether graph\-connected gene blocks could serve as the prediction units of a transcriptomic JEPA\. We developed BioM\-JEPA, a block\-level JEPA for single cells \(Fig\.[1](https://arxiv.org/html/2608.05928#S0.F1)\)\. A binary gene graph combines high\-confidence STRING v12 protein associations with transcriptome\-wide coexpression estimated from the unlabeled pretraining corpus\[[18](https://arxiv.org/html/2608.05928#bib.bib3)\]\. Connected sets sampled from this graph define candidate target blocks\. The student encoder observes the complementary genes in the cell and predicts one aggregate representation for each hidden block, while a slowly updated teacher receives the full observed gene set and provides the corresponding target representation\. Crucially, teacher states are aggregated across the genes from the block that are observed in that cell before the prediction error is evaluated\. Thus, BioM\-JEPA learns to infer a block\-scale representation instead of reproducing every target gene independently\. We reserve the terms biological programme and pathway for gene sets supported by an independent annotation or biological analysis; graph connectivity alone does not confer that status\.

The model uses linear attention throughout the student, teacher and predictor\. This avoids constructing a complete gene\-by\-gene attention matrix and allows computation to grow linearly with the number of observed genes at fixed model width\[[10](https://arxiv.org/html/2608.05928#bib.bib16)\]\. The biological objective and the computational architecture address separate problems: graph\-defined blocks determine what the representation is trained to preserve, whereas linear attention makes broad gene contexts practical\. The analysed representation was obtained after 10\.24 million cell presentations, less than half of the 22\.1\-million\-cell local collection and 2\.04% of the 502\-million\-cell scBaseCount resource on a presentation\-equivalent basis\[[22](https://arxiv.org/html/2608.05928#bib.bib2)\]\. This exposure is below the reported pretraining\-corpus sizes of Geneformer \(approximately 30 million cells\), scGPT \(more than 33 million\), scFoundation \(more than 50 million\) and CellFM \(102\.3 million\)\[[20](https://arxiv.org/html/2608.05928#bib.bib4),[6](https://arxiv.org/html/2608.05928#bib.bib5),[9](https://arxiv.org/html/2608.05928#bib.bib6),[23](https://arxiv.org/html/2608.05928#bib.bib8)\], showing that the representation emerged after a comparatively compact training exposure\.

We evaluated the resulting representations with shared\-backbone controls and across CellBench within\-dataset few\-shot annotation, reconstruction and perturbation tasks\[[21](https://arxiv.org/html/2608.05928#bib.bib17)\]\. Under the reported extraction procedure, block prediction produced higher effective rank and weaker association with detected\-gene depth in two diagnostics than token\-level prediction, random blocks or decoder\-only reconstruction\. Frozen BioM\-JEPA representations retained continuous expression, pathway and neighbourhood structure and provided the lowest aggregate perturbation\-response error among the evaluated models\. In a matched hPancreas timing experiment at batch size 8 with nearly identical trainable parameter counts, BioM\-JEPA increased one\-epoch fine\-tuning throughput by 5\.75\-fold and held\-out embedding throughput by 3\.76\-fold relative to scFoundation\. Finally, targeted analyses were consistent with canonical pancreatic cell programmes, directed predictive associations between graph blocks and compositional relationships between genetic perturbations\. These results support block\-level prediction as a practical construction for extending JEPA learning to molecular data\.

![Refer to caption](https://arxiv.org/html/2608.05928v1/figures/mainv2_nature/figure1_block_jepa_method.png)Figure 1:Graph\-connected block prediction with BioM\-JEPA\.a, A cell is represented as a sparse set of observed genes and their expression values\.b, A gene graph defines connected candidate target blocks\.c, Target\-block genes are removed from the student input in the dominant masking regime, whereas the EMA teacher receives the full observed gene set\. The student predicts the aggregate teacher representation of each target block from the complementary context; gradients do not pass through the teacher target\.d, BioM\-JEPA optimizes block\-level alignment and regularizes the predicted block representations\. It does not use an individual target\-gene reconstruction objective\. Detailed definitions and training pseudocode are provided in Supplementary Methods\.## 1Results

### 1\.1Token prediction can optimize without producing a robust cell representation

We first asked whether successful latent prediction at the gene level is sufficient to organize cells in a biologically useful representation space\. We trained BioM\-JEPA, pure token\-IJEPA and a decoder\-only control using the same vocabulary, hidden width, linear\-attention backbone and pretraining data\. Token\-IJEPA predicted individual teacher gene states, the decoder reconstructed expression, and BioM\-JEPA predicted one teacher representation for each graph\-defined target block\.

All three training losses decreased, while teacher\-target similarity increased for the two JEPA models \(Fig\.[2](https://arxiv.org/html/2608.05928#S1.F2)a\)\. Their pooled cell representations nevertheless developed differently\. Across matched hPancreas training snapshots, the effective rank of the BioM\-JEPA embedding rose from approximately 19 to nearly 40\. Token\-IJEPA remained near 19, whereas the decoder\-only representation contracted to fewer than ten effective dimensions \(Fig\.[2](https://arxiv.org/html/2608.05928#S1.F2)b\)\. Accurate token prediction or expression reconstruction therefore did not ensure that cell\-to\-cell variation was retained across many independent directions\.

The geometry also differed in its association with sequencing depth\. We correlated the leading embedding axis and the embedding norm with the number of detected genes\. BioM\-JEPA showed the weakest mean absolute association, token\-IJEPA the strongest and decoder\-only an intermediate value \(Fig\.[2](https://arxiv.org/html/2608.05928#S1.F2)c\)\. When the encoders were frozen and evaluated using the same CellBench Top\-5 probe, BioM\-JEPA achieved the highest macro\-F1F\_\{1\}on both hPancreas and cortex \(Fig\.[2](https://arxiv.org/html/2608.05928#S1.F2)d\)\. These results support a failure mode specific to pooled transcriptomic representations: a gene\-level predictive objective can converge while the resulting cell embedding remains low dimensional and coupled to technical depth\.

![Refer to caption](https://arxiv.org/html/2608.05928v1/figures/mainv2_nature/figure2_objective_failure_mode.png)Figure 2:Objective optimization separates from pooled representation quality\.a, Relative training losses for all three models \(solid lines\) and teacher\-target similarity for the two JEPA models \(dotted lines\)\. Loss is normalized within each objective family\.b, Effective rank of hPancreas cell embeddings across training snapshots\.c, Mean absolute association between embedding diagnostics and detected\-gene depth; open symbols show individual diagnostics and filled symbols their mean\.d, CellBench within\-dataset Top\-5 few\-shot macro\-F1F\_\{1\}for hPancreas and cortex \(mean and s\.d\. over five seeds\)\.
### 1\.2BioM\-JEPA improves the annotation–robustness balance

We next compared graph\-based target construction and objective choice using a shared encoder backbone and an identical embedding\-extraction and downstream\-evaluation pipeline\. Random block used target sets sampled from the same size range without graph expansion\. Token\-IJEPA retained the teacher–student architecture but predicted genes individually\. Decoder\-only retained graph blocks but replaced latent alignment with expression reconstruction\. BioM\-JEPA combined graph\-connected blocks with aggregate block\-level prediction \(Fig\.[3](https://arxiv.org/html/2608.05928#S1.F3)a\)\.

BioM\-JEPA achieved the highest mean Top\-5 macro\-F1F\_\{1\}across hPancreas and cortex \(0\.835\), compared with 0\.817 for token\-IJEPA, 0\.812 for random block and 0\.797 for decoder\-only \(Fig\.[3](https://arxiv.org/html/2608.05928#S1.F3)b\)\. The same ordering was observed for effective rank on both datasets \(Fig\.[3](https://arxiv.org/html/2608.05928#S1.F3)c\)\. Participation ratio, which measures how evenly variation is distributed across the representation, likewise placed BioM\-JEPA above random block and decoder\-only, with greater dataset dependence for token\-IJEPA\.

The complete block objective also reduced technical coupling in the tested diagnostics\. Correlations between detected\-gene count and either the leading embedding axis or embedding norm were lower for BioM\-JEPA than for the three controls \(Fig\.[3](https://arxiv.org/html/2608.05928#S1.F3)d\)\. Plotting within\-dataset annotation performance against depth association placed BioM\-JEPA in the favourable region of high annotation score and low technical coupling \(Fig\.[3](https://arxiv.org/html/2608.05928#S1.F3)e\)\. Random targets produced weaker annotation despite matching the block\-size range, while token prediction retained stronger depth association\. Because all models used the same encoder backbone, the observed differences were associated with the complete pretraining formulations rather than encoder width or depth\.

![Refer to caption](https://arxiv.org/html/2608.05928v1/figures/mainv2_nature/figure3_objective_controls.png)Figure 3:Controlled comparison of pretraining objective families\.a, Components used by the four pretrained variants\.b, Mean Top\-5 few\-shot macro\-F1F\_\{1\}over hPancreas and cortex\.c, Effective rank and participation ratio for hPancreas \(circles\) and cortex \(diamonds\)\.d, Absolute correlations between detected\-gene count and the leading embedding axis or embedding norm\.e, Mean within\-dataset macro\-F1F\_\{1\}versus mean depth association\. Higher annotation performance and lower association are preferred\.
### 1\.3Linear attention accelerates embedding extraction and fine\-tuning

The model contains twelve linear\-attention encoder layers and a four\-layer predictor\. Unlike dense softmax attention, the attention core does not construct a matrix containing every pair of observed genes\. Keys and values are summarized before they are combined with each query, so computation grows linearly with the number of observed genes at fixed model width\[[10](https://arxiv.org/html/2608.05928#bib.bib16)\]\. This allows BioM\-JEPA to preserve the full observed\-gene sequence rather than impose an attention\-driven token cap\.

We first isolated frozen encoder execution on identical ordered hPancreas cells at batch size 8\. BioM\-JEPA encoded 678\.26 cells per second, compared with 353\.35 cells per second for scFoundation, a 1\.92\-fold increase\. hPancreas contained a median of 319 and a maximum of 1,659 observed gene tokens, and no artificial length cap was used\.

We then measured a task\-realistic one\-epoch hPancreas run at the same batch size\. The final encoder block, leaf normalization parameters and an identical two\-layer prediction head were trainable, giving 26\.09 million trainable parameters for each model\. BioM\-JEPA processed 24\.37 training cells per second, compared with 4\.24 for scFoundation, a 5\.75\-fold increase \(Table[1](https://arxiv.org/html/2608.05928#S1.T1)\)\. During the no\-gradient validation and test loops, held\-out embedding throughput was 62\.78 versus 16\.68 cells per second, a 3\.76\-fold increase\. The complete train–validation–test epoch required 66\.08 s for BioM\-JEPA and 352\.68 s for scFoundation\. Within this matched protocol, BioM\-JEPA was faster than scFoundation in both frozen and partially tuned downstream use\.

Strong frozen\-probe performance emerged after 10\.24 million cell presentations, approximately 46% of the 22\.1\-million\-cell local collection\. The reported corpus sizes of Geneformer, scGPT, scFoundation and CellFM range from approximately 30 million to 102\.3 million cells\[[20](https://arxiv.org/html/2608.05928#bib.bib4),[6](https://arxiv.org/html/2608.05928#bib.bib5),[9](https://arxiv.org/html/2608.05928#bib.bib6),[23](https://arxiv.org/html/2608.05928#bib.bib8)\]\. BioM\-JEPA therefore learned a useful block\-aware representation after a comparatively compact training exposure\.

Table 1:Matched one\-epoch hPancreas efficiency\.Both models used batch size 8, bfloat16 precision, one A100\-SXM4\-80GB GPU, the same data order, pooling structure and prediction head\. Held\-out throughput is measured over the no\-gradient validation and test loops\.
### 1\.4Frozen embeddings retain expression, pathway and neighbourhood information

A representation can support cell\-type classification while discarding continuous transcriptional structure\. To test what remained accessible from each frozen embedding, we trained the same supervised decoder to predict a canonical expression vector \(Fig\.[4](https://arxiv.org/html/2608.05928#S1.F4)a\)\. Encoder weights were never updated, and every model used the same decoder architecture, supervision budget, held\-out split and stopping rule\.

Across eight CellBench datasets, BioM\-JEPA achieved the highest dataset\-mean cell\-wise Pearson correlation and the lowest normalized root\-mean\-square error among the evaluated frozen representation models \(Fig\.[4](https://arxiv.org/html/2608.05928#S1.F4)b\)\. It also recovered more of the 50 most highly expressed genes, gave the highest Reactome programme correlation and best preserved expression\-space nearest neighbours \(Fig\.[4](https://arxiv.org/html/2608.05928#S1.F4)c,d\)\. These metrics measure distinct properties: numerical expression fidelity, recovery of prominent genes, pathway\-level activity and local cell\-state geometry\.

Increasing the decoder training budget improved programme recovery for every model, but BioM\-JEPA remained strongest from 100 to 900 labelled cells per class \(Fig\.[4](https://arxiv.org/html/2608.05928#S1.F4)e\)\. In hPancreas, reconstructed marker\-programme matrices preserved the major endocrine and exocrine organization, including alpha, beta, ductal, acinar, delta and pancreatic\-polypeptide states \(Fig\.[4](https://arxiv.org/html/2608.05928#S1.F4)f\)\. Thus, block prediction did not improve annotation by reducing the transcriptome to a narrow class code; the frozen representation retained continuous and multiscale biological information\.

![Refer to caption](https://arxiv.org/html/2608.05928v1/figures/mainv2_nature/figure4_reconstruction_probe.png)Figure 4:Reconstruction probes of frozen single\-cell embeddings\.a, Shared\-decoder evaluation\.b, Cell\-wise expression fidelity across eight datasets\.c, Recovery of highly expressed genes and Reactome programmes\.d, Preservation of cell\-state neighbourhoods\.e, Reactome programme correlation as decoder supervision increases\.f, hPancreas marker\-programme reconstruction\. Small symbols inb–ddenote datasets; large points and intervals summarize the dataset mean and its uncertainty\.
### 1\.5Block representations support perturbation\-response prediction

We next evaluated whether frozen cell representations support prediction of responses that were not provided to the encoder\. The CellBench response model receives a control\-cell embedding and a perturbation identity and predicts the perturbed expression profile \(Fig\.[5](https://arxiv.org/html/2608.05928#S1.F5)a\)\. We evaluated Adamson, Norman and Dixit using matched Top\-kkcontrols and five fixed seeds\. The encoder remained frozen and the response predictor was identical across models\.

With Top\-5 controls, BioM\-JEPA achieved the lowest aggregate mean\-squared error of the log\-fold\-change response \(10\.31\), compared with 10\.87–10\.98 for scFoundation, scVI, scGPT and Geneformer \(Fig\.[5](https://arxiv.org/html/2608.05928#S1.F5)b\)\. Its Pearson correlation was highest for the aggregate and unseen\-single categories \(Fig\.[5](https://arxiv.org/html/2608.05928#S1.F5)c\)\. On the Norman combination subset, the only dataset containing combination perturbations, scVI obtained the highest correlation\.

BioM\-JEPA also recovered the largest fraction of the top 50 response genes \(0\.123\), followed by scFoundation \(0\.115; Fig\.[5](https://arxiv.org/html/2608.05928#S1.F5)d\)\. In a representative Adamson condition, its predictions reproduced the direction and relative magnitude of the dominant response programme \(Fig\.[5](https://arxiv.org/html/2608.05928#S1.F5)e\)\. Agreement across error, correlation and response\-gene recovery is consistent with the block representation preserving information relevant to both the magnitude and biological composition of perturbation responses in the tested protocol\.

![Refer to caption](https://arxiv.org/html/2608.05928v1/figures/mainv2_nature/figure5_perturbation_prediction.png)Figure 5:Perturbation prediction from frozen cell representations\.a, CellBench response\-prediction protocol\.b, Top\-5 MSE\-LFC \(mean±\\pms\.d\. over five seeds; lower is better\)\.c, Pearson\-LFC by perturbation type; combination perturbations are available for Norman only\.d, Top\-50 response\-gene overlap \(mean±\\pms\.d\. over five seeds; higher is better\)\.e, Representative Adamson Top\-5 observed and predicted response profiles\.
### 1\.6BioM\-JEPA representations reflect biological programmes and genetic relationships

We finally asked whether the frozen representation could expose recognizable biological structure without fitting a classifier\. In the hPancreas dataset\[[16](https://arxiv.org/html/2608.05928#bib.bib19)\], the genes that most frequently supplied the maximal token activation aligned with canonical cell identities:GCGin alpha cells,INSin beta cells,SSTin delta cells,PPYin pancreatic\-polypeptide cells,GHRLin epsilon cells andREG1Ain acinar cells \(Fig\.[6](https://arxiv.org/html/2608.05928#S1.F6)a\)\.

We then removed each of eight graph\-defined programmes and measured how far the frozen cell embedding moved relative to removal of an expression\- and detection\-matched random gene set\. Across the eight pre\-specified programme identities, the corresponding cell identity showed a larger displacement than the mean of the other identities \(mean excess shift, 5\.34 versus 0\.064; two\-sided paired Wilcoxon test across programmes,n=8n=8,P=0\.0078P=0\.0078; Fig\.[6](https://arxiv.org/html/2608.05928#S1.F6)b\)\. Directed ablations also produced an asymmetric matrix of changes in target\-block prediction error \(Fig\.[6](https://arxiv.org/html/2608.05928#S1.F6)c\)\. These results associate cell identity with coordinated graph programmes rather than only with isolated high\-activation genes\.

The Norman perturbation atlas provides a second test of biological organization\[[14](https://arxiv.org/html/2608.05928#bib.bib20)\]\. For each double perturbation, we compared its latent displacement with the sum of the displacements produced by its two constituent single perturbations\. All 80 evaluated combination conditions agreed more strongly with their matched constituent pair than with the mean of 200 expression\-matched random pairs per condition \(two\-sided paired Wilcoxon test across conditions,n=80n=80,P=7\.85×10−15P=7\.85\\times 10^\{\-15\}; Fig\.[6](https://arxiv.org/html/2608.05928#S1.F6)d\)\. The highest\-agreement examples included MAP2K3–MAP2K6, which converge on p38 signalling\[[15](https://arxiv.org/html/2608.05928#bib.bib21)\]; CEBPA–CEBPE, which coordinate granulocytic differentiation\[[3](https://arxiv.org/html/2608.05928#bib.bib22)\]; and the related developmental regulators TBX2–TBX3\[[17](https://arxiv.org/html/2608.05928#bib.bib23)\]\. These examples are consistent with established functional relationships being reflected in the geometry of the frozen representation and provide a starting point for testable biological hypotheses\.

![Refer to caption](https://arxiv.org/html/2608.05928v1/x1.png)Figure 6:Biological structure reflected in frozen BioM\-JEPA representations\.a, Cell\-type\-resolved max\-token winner fractions for canonical hPancreas genes\.b, Excess embedding shift after graph\-programme ablation relative to one expression\- and detection\-matched random set per programme\. Points summarize 24 cells per identity; the aggregate test uses the eight programme identities as paired statistical units\.c, Descriptive directed change in target\-block prediction error after ablating each context programme, corrected by the mean of three matched\-random ablations; each ordered pair was evaluated in 24 cells from the target\-matched identity\.d, Agreement between each of 80 Norman double\-perturbation displacements and the additive displacement of its matched constituent pair, compared with 200 matched\-random pairs per condition\.e, Selected literature\-supported gene pairs\. Teal points show the matched pair; grey points and intervals show the random\-pair mean and s\.d\. MAP2K3–MAP2K6, CEBPE–CEBPA, TBX3–TBX2 and CEBPB–MAPK1 represent p38 signalling, granulocytic differentiation, developmental regulation and ERK–C/EBPβ\\betasignalling, respectively\. Complete definitions and source\-data construction are provided in Supplementary Methods\.

## 2Discussion

BioM\-JEPA defines graph\-connected gene blocks as prediction units for joint\-embedding learning in single cells\. Under a shared backbone and extraction procedure, aggregate block prediction produced representations with higher effective rank, weaker association with detected\-gene depth in the tested diagnostics and stronger frozen\-probe performance than token prediction, random blocks and expression reconstruction\. These findings support prediction scale as a central design choice: convergence of a gene\-level objective does not by itself ensure a cell representation that retains broad biological variation\.

The block objective is well matched to sparse transcriptomic observations\. Failure to detect one gene is weak evidence that its biological programme is inactive, whereas multiple connected genes provide partially redundant evidence for a shared latent state\. Graph connectivity also gives each target an identity independent of gene order\. Graph\-connected blocks therefore play a role analogous to coherent spatial regions in I\-JEPA and spatiotemporal regions in V\-JEPA, while adapting the prediction unit to molecular data\[[2](https://arxiv.org/html/2608.05928#bib.bib14),[4](https://arxiv.org/html/2608.05928#bib.bib15)\]\.

The learned representation retained information across biological scales\. It supported few\-shot annotation, continuous expression reconstruction, Reactome programme recovery, local cell\-state neighbourhoods and perturbation\-response prediction\. Max\-token analysis aligned with canonical pancreatic genes, programme ablation selectively displaced the corresponding cell identities and latent perturbation arithmetic reflected established functional gene pairs\. Because the graph contains protein\-association and coexpression evidence but no cell\-type labels, regulatory direction or perturbation\-combination identities, these analyses reveal biological structure learned beyond the annotations directly supplied to the model\. They provide a representation\-level basis for prioritizing candidate relationships for experimental study\.

Linear attention makes this objective practical for broad observed\-gene contexts\. By contracting keys and values before applying each query, BioM\-JEPA avoids materializing all gene pairs and scales linearly with sequence length at fixed width\. In matched hPancreas measurements at batch size 8, BioM\-JEPA was 5\.75 times faster than scFoundation during one\-epoch partial fine\-tuning and 3\.76 times faster during held\-out embedding extraction, with nearly identical trainable parameter counts\. The representation was obtained after 10\.24 million cell presentations, below the published pretraining\-corpus sizes of Geneformer, scGPT, scFoundation and CellFM\. Together, these results show that a useful block\-aware representation can emerge from a comparatively compact training exposure\.

The pretraining collection lacks study accession identifiers, so study\-level overlap with CellBench cannot be excluded and annotation is reported as within\-dataset frozen\-probe evaluation\. Graph\-connected blocks are operational prediction targets rather than curated pathways, and the objective controls compare complete pretraining formulations\. Depth association was evaluated through the leading embedding axis and embedding norm\. These scope conditions motivate accession\-resolved pretraining corpora, component\-wise objective studies and broader technical\-covariate tests, while preserving the central observation that block\-level JEPA produced the strongest aggregate representation across the tested objectives\.

BioM\-JEPA extends joint\-embedding prediction from ordered visual regions to unordered molecular observations\. Its essential step is to define the latent target through relationships native to the modality rather than treating individual genes as transcriptomic image patches\. This principle provides a direct route to prediction units based on independently validated pathways, regulatory circuits and molecular complexes in future representation models for biological data\.

## 3Methods

### 3\.1Problem formulation

Let𝒱\\mathcal\{V\}denote a shared vocabulary ofG=19,264G=19\{,\}264genes\. A celliiis represented by its observed gene set𝒪i⊆𝒱\\mathcal\{O\}\_\{i\}\\subseteq\\mathcal\{V\}and normalized expression valuesxi​gx\_\{ig\}forg∈𝒪ig\\in\\mathcal\{O\}\_\{i\}\. Our aim is to learn a cell representation without cell\-type, perturbation or study labels\. Rather than reconstructing every expression value, BioM\-JEPA learns by predicting latent representations of structured subsets of the observed transcriptome\.

The central prediction unit is a*graph\-connected gene block*\. We use this term strictly for a connected set sampled from a gene graph; biological programme or pathway refers only to a set supported by an independent annotation or downstream analysis\. This distinction separates the mathematical object used for self\-supervision from the biological interpretation tested after training\.

### 3\.2Cell representation and gene graph

Raw counts were mapped to𝒱\\mathcal\{V\}, normalized to a constant library size and transformed asxi​g=log⁡\(1\+104​ci​g/∑g′ci​g′\)x\_\{ig\}=\\log\(1\+10^\{4\}c\_\{ig\}/\\sum\_\{g^\{\\prime\}\}c\_\{ig^\{\\prime\}\}\)\. Two technical values encoding a fixed resolution and log library depth were available to the encoders as context but were excluded from target sampling and downstream pooling\.

We constructed a sparse gene graph𝒢=\(𝒱,ℰ\)\\mathcal\{G\}=\(\\mathcal\{V\},\\mathcal\{E\}\)from two complementary sources\. High\-confidence STRING v12 protein associations supplied experimentally and computationally supported molecular links, and corpus\-level coexpression supplied transcriptomic proximity\[[18](https://arxiv.org/html/2608.05928#bib.bib3)\]\. For each gene, we retained at most 64 outgoing neighbours after taking the binary union of STRING associations with combined score at least 700 and the strongest approximate absolute coexpression neighbours\. The resulting adjacency contained 1,123,337 directed entries, with mean out\-degree 58\.31\. No cell\-type label, perturbation identity, pathway membership or transcription\-factor direction was used to construct the graph\.

### 3\.3Graph\-connected target blocks

For each cell and each ofK=4K=4targets, we sample an observed seed gene and a requested block size between 2,000 and 8,000 genes\. Breadth\-first expansion from the seed produces a candidate blockBi​k⊆𝒱B\_\{ik\}\\subseteq\\mathcal\{V\}\. Only genes detected in the current cell contribute to its target support,Ti​k=Bi​k∩𝒪iT\_\{ik\}=B\_\{ik\}\\cap\\mathcal\{O\}\_\{i\}\. The nominal student context is the complementary observed set𝒞i=𝒪i∖⋃kTi​k\\mathcal\{C\}\_\{i\}=\\mathcal\{O\}\_\{i\}\\setminus\\bigcup\_\{k\}T\_\{ik\}\. When this operation leaves fewer than 512 context tokens, the observed set is retained as context to avoid an uninformative short sequence\. The teacher always receives the full observed gene set\. Thus, target construction is shared across cells through the gene graph, whereas the realized target is specific to the genes observed in each cell\.

Each block is identified by the median vocabulary index of its candidate genes\. Its learned query combines the corresponding gene\-position embedding with a mask embedding and contains no target expression value\. On the student side, the query first selects visible context within a 64\-neighbour coexpression neighbourhood; when no such neighbour is visible, all visible context genes remain eligible\. A temperature\-controlled similarity weighting then summarizes this context into one student block state\.

### 3\.4Student–teacher prediction

The student encoderfθf\_\{\\theta\}and teacher encoderfξf\_\{\\xi\}share a 12\-layer, 768\-dimensional linear\-attention architecture\. The teacher is updated as an exponential moving average of the student,ξ←μ​ξ\+\(1−μ\)​θ\\xi\\leftarrow\\mu\\xi\+\(1\-\\mu\)\\theta, with momentum increasing from 0\.996 to 0\.9997\. It is never optimized by backpropagation\. A four\-layer predictorpϕp\_\{\\phi\}maps the pooled student context to the target\-representation space\.

For a valid cell–block pair\(i,k\)\(i,k\), the teacher target is the mean of the teacher token states overTi​kT\_\{ik\}, followed by EMA centring and gradient stopping\. Ifhi​gth^\{t\}\_\{ig\}is the teacher state of geneggandccis the running centre, the target is

ti​k=sg⁡\(1\|Ti​k\|​∑g∈Ti​khi​gt−c\)\.t\_\{ik\}=\\operatorname\{sg\}\\\!\\left\(\\frac\{1\}\{\|T\_\{ik\}\|\}\\sum\_\{g\\in T\_\{ik\}\}h^\{t\}\_\{ig\}\-c\\right\)\.The corresponding student statesi​ks\_\{ik\}is obtained by query\-conditioned pooling of the visible context, and the prediction ist^i​k=pϕ​\(si​k\)\\widehat\{t\}\_\{ik\}=p\_\{\\phi\}\(s\_\{ik\}\)\. Aggregation therefore precedes the prediction error: one graph\-connected block contributes one latent target, irrespective of how many of its genes are observed\. This is the defining difference from token\-level JEPA, which assigns a separate target to every masked gene\.

### 3\.5Block\-level objective

For the setℬ\\mathcal\{B\}of valid cell–block pairs in a minibatch, the alignment term is the mean squared distance between predicted and teacher block states\. We additionally regularize the batch of predictions to retain variation across embedding dimensions and reduce redundant covariance\. The complete objective is

ℒB=1\|ℬ\|​∑\(i,k\)∈ℬ‖t^i​k−ti​k‖22\+λvar​ℒvar\+λcov​ℒcov,\\mathcal\{L\}\_\{\\mathrm\{B\}\}=\\frac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\_\{\(i,k\)\\in\\mathcal\{B\}\}\\\|\\widehat\{t\}\_\{ik\}\-t\_\{ik\}\\\|\_\{2\}^\{2\}\+\\lambda\_\{\\mathrm\{var\}\}\\mathcal\{L\}\_\{\\mathrm\{var\}\}\+\\lambda\_\{\\mathrm\{cov\}\}\\mathcal\{L\}\_\{\\mathrm\{cov\}\},whereℒvar\\mathcal\{L\}\_\{\\mathrm\{var\}\}penalizes dimensions whose batch standard deviation falls below one andℒcov\\mathcal\{L\}\_\{\\mathrm\{cov\}\}penalizes off\-diagonal covariance\. We usedλvar=0\.05\\lambda\_\{\\mathrm\{var\}\}=0\.05andλcov=0\.01\\lambda\_\{\\mathrm\{cov\}\}=0\.01\. No individual target\-gene reconstruction term enters this objective\.

This construction has two consequences\. First, genes from the same connected block provide redundant evidence for a shared latent target, reducing the dependence of supervision on any one sparse measurement\. Second, the target dimension is fixed by the number of sampled blocks rather than the number of detected target genes, so cells of different sequencing depth contribute comparable numbers of prediction units\.

### 3\.6Linear\-attention encoder

All student, teacher and predictor layers use linear attention\. A positive feature map,φ​\(u\)=ELU⁡\(u\)\+1\\varphi\(u\)=\\operatorname\{ELU\}\(u\)\+1, permits keys and values to be contracted into a shared summary before each query is applied\. At fixed hidden width, this factorization grows linearly with the number of observed genes and avoids constructing a dense gene\-by\-gene attention matrix\[[10](https://arxiv.org/html/2608.05928#bib.bib16)\]\. It therefore permits the encoder to retain the observed gene sequence without an attention\-imposed length cap\.

### 3\.7Cell\-level embedding

After training, the encoder is frozen and each cell is represented by pooling its gene\-token states\. For the 768\-dimensional token matrix, we concatenate the element\-wise maximum with the mean of the five largest values in each feature dimension, producing a 1,536\-dimensional cell embedding\. Technical\-token positions are excluded from this operation\. The same cell embedding is used without encoder updates for annotation, reconstruction, perturbation prediction and representation diagnostics\.

### 3\.8Pretraining data and optimization

Pretraining used a scBaseCount snapshot containing records available no later than 25 February 2025\[[22](https://arxiv.org/html/2608.05928#bib.bib2)\]\. The local corpus comprised 22\.1 million cells\. We trained with an effective batch of 512 cells and analysed the representation after 20,000 optimizer steps, corresponding to 10\.24 million cell presentations\. This is 46\.33% of one traversal of the local collection and 2\.04% of the 502\-million\-cell scBaseCount resource on a presentation\-equivalent basis\. The local matrices do not retain source\-study accessions; this provenance boundary is considered in the Discussion\.

### 3\.9Evaluation framework

BioM\-JEPA is the method studied throughout this work\. To contextualize its representation behaviour, we evaluated three matched objective controls\. Token\-IJEPA predicted individual teacher gene states; random block replaced graph expansion with uniformly sampled target sets while preserving aggregate prediction; and decoder\-only replaced latent alignment with expression reconstruction\. Each control retained the vocabulary, encoder width, linear\-attention backbone and downstream extraction procedure used for BioM\-JEPA\.

### 3\.10Within\-dataset few\-shot annotation

We followed the CellBench\-LS within\-dataset protocol\[[21](https://arxiv.org/html/2608.05928#bib.bib17)\]\. For each cell type, one, five or nine labelled support cells were selected by label\-stratified sampling, an SVM was fitted to frozen embeddings and all remaining labelled cells formed the held\-out set\. Macro\-F1F\_\{1\}was the unweighted mean of class\-wiseF1F\_\{1\}values\. We used five fixed support\-set seeds; reported error bars show one standard deviation\.

### 3\.11Frozen\-embedding reconstruction

To measure information accessible from the representation, we trained the same multilayer perceptron on each model’s frozen embedding to predict a canonical 400\-gene expression vector\. We evaluated cell\-wise Pearson correlation, normalized root\-mean\-square error, overlap among the top 50 expressed genes, Reactome programme correlation\[[8](https://arxiv.org/html/2608.05928#bib.bib18)\]and overlap between 15\-nearest\-neighbour cell\-state graphs\. The decoder architecture, data split and optimization budget were shared across models\.

### 3\.12Perturbation\-response prediction

The CellBench response model received a frozen control\-cell embedding and a perturbation identity and predicted the corresponding perturbed expression profile\. Adamson, Norman and Dixit were evaluated with matched Top\-kkcontrol examples and five fixed seeds\. We quantified expression and log\-fold\-change error, Pearson correlation and overlap among the 50 strongest response genes\. Norman additionally provided single\- and combination\-perturbation strata\.

### 3\.13Efficiency evaluation

Embedding extraction and one\-epoch partial fine\-tuning were measured separately on hPancreas with batch size 8 on one NVIDIA A100\-SXM4\-80GB GPU using bfloat16 computation\. Both BioM\-JEPA and scFoundation received the same ordered cells\. For partial fine\-tuning, the final encoder block, leaf normalization parameters and a common two\-layer classification head were trainable, yielding 26,088,205 and 26,104,077 trainable parameters, respectively\. Training throughput includes forward propagation, backpropagation and the optimizer step; held\-out throughput measures no\-gradient encoding followed by the shared prediction head\.

### 3\.14Biological representation analyses

For the max\-token analysis, we identified the gene providing the largest token contribution in each held\-out hPancreas cell and calculated its winner fraction within each annotated cell type\. For programme ablation, genes from a graph\-defined set were removed and the embedding displacement was compared with expression\- and detection\-matched random removals\. Directed programme dependence was measured as the change in teacher\-target prediction error after ablating a context programme, corrected by the matched\-random effect\.

For Norman double perturbations, each perturbation was represented by its mean displacement from control cells in the frozen embedding space\. The sum of the two single\-perturbation displacements was compared with the observed double\-perturbation displacement, with expression\-matched random gene pairs defining the null distribution\.

### 3\.15Representation geometry and statistics

Effective rank and participation ratio were calculated from the eigenvalue spectrum of the centred embedding covariance\. Association with sequencing depth was summarized by the absolute Pearson correlation of detected\-gene count with the leading embedding axis and with embedding norm\. These diagnostics measure two specified technical axes rather than complete statistical independence from depth\.

Few\-shot and perturbation analyses used five fixed seeds unless stated otherwise\. Cross\-dataset intervals were obtained by nonparametric bootstrap resampling of datasets\. Programme\-ablation and double\-perturbation effects were evaluated with paired nonparametric tests using the statistical units specified in the corresponding figure legends\. Detailed metric definitions, pseudocode and complete evaluation protocols are provided in the Supplementary Information\.

## 4Acknowledgements

This work was supported by the National Key R&D Program of China \(No\. 2022ZD0115100\), the National Natural Science Foundation of China \(No\. U21A20427\), the Center of Synthetic Biology and Integrated Bioengineering of Westlake University \(No\. WU2022A009\), the “Pioneer” and “Leading Goose” R&D Program of Zhejiang \(No\. 2024C01140\), the Key Research and Development Program of Hangzhou \(No\. 2023SZD0073\), the InnoHK programme and the CAAI–Ant Research Fund\. We thank the Westlake University HPC Center for providing computational resources\.

## 5Data availability

The pretraining shards were derived from the public scBaseCount resource, and the evaluation datasets are available through the CellBench\-LS release \([https://github\.com/sky\-Yongjie\-Xu/2026\-CellBench](https://github.com/sky-Yongjie-Xu/2026-CellBench)\) and their original repositories\. The arXiv ancillary bundle contains the fixed dataset identities, benchmark protocol metadata, accession\-overlap audit and panel\-level numerical values underlying every quantitative figure and supplementary summary\. It does not redistribute count matrices, per\-cell embeddings or full prediction arrays\. No private patient\-level metadata were used\.

## 6Code availability

The BioM\-JEPA checkpoint, training and inference software, fixed graph, vocabulary and an AnnData\-to\-embedding interface will be released through GitHub\. The repository will also document data preparation, representation extraction and reproduction of the reported analyses\. Numerical source data for the figures accompany the arXiv submission\.

## References

- \[1\]\(2025\)Deep\-learning\-based gene perturbation effect prediction does not yet outperform simple linear baselines\.Nature Methods22,pp\. 1657–1661\.External Links:[Document](https://dx.doi.org/10.1038/s41592-025-02772-6)Cited by:[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1)\.
- \[2\]M\. Assranet al\.\(2023\)Self\-supervised learning from images with a joint\-embedding predictive architecture\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 15619–15629\.External Links:[Document](https://dx.doi.org/10.1109/CVPR52729.2023.01499)Cited by:[Relation to established JEPA objectives\.](https://arxiv.org/html/2608.05928#Ax2.SSx6.SSS0.Px1.p1.2),[§2](https://arxiv.org/html/2608.05928#S2.p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p3.1)\.
- \[3\]R\. Avellino, H\. J\. G\. van de Werken,et al\.\(2022\)Transcription factor\-driven coordination of cell cycle exit and lineage\-specification in vivo during granulocytic differentiation\.Nature Communications13,pp\. 3595\.External Links:[Document](https://dx.doi.org/10.1038/s41467-022-31332-1)Cited by:[Compositional perturbation geometry\.](https://arxiv.org/html/2608.05928#Ax2.SSx14.SSS0.Px4.p1.8),[§1\.6](https://arxiv.org/html/2608.05928#S1.SS6.p3.2)\.
- \[4\]A\. Bardeset al\.\(2024\)Revisiting feature prediction for learning visual representations from video\.Transactions on Machine Learning Research\.External Links:2404\.08471,[Link](https://arxiv.org/abs/2404.08471)Cited by:[Relation to established JEPA objectives\.](https://arxiv.org/html/2608.05928#Ax2.SSx6.SSS0.Px1.p1.2),[§2](https://arxiv.org/html/2608.05928#S2.p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p3.1)\.
- \[5\]H\. Bianet al\.\(2024\)scMulan: a multitask generative pre\-trained language model for single\-cell analysis\.InResearch in Computational Molecular Biology,pp\. 479–482\.External Links:[Document](https://dx.doi.org/10.1007/978-1-0716-3989-4%5F57)Cited by:[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1)\.
- \[6\]H\. Cuiet al\.\(2024\)scGPT: toward building a foundation model for single\-cell multi\-omics using generative AI\.Nature Methods21,pp\. 1470–1480\.External Links:[Document](https://dx.doi.org/10.1038/s41592-024-02201-0)Cited by:[Context relative to reported foundation\-model corpora\.](https://arxiv.org/html/2608.05928#Ax2.SSx2.SSS0.Px1.p1.1),[§1\.3](https://arxiv.org/html/2608.05928#S1.SS3.p4.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p5.1)\.
- \[7\]A\. DenAdelet al\.\(2026\)Evaluating the role of pretraining dataset size and diversity on single\-cell foundation model performance\.Nature Methods23,pp\. 1447–1457\.External Links:[Document](https://dx.doi.org/10.1038/s41592-026-03120-y)Cited by:[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1)\.
- \[8\]M\. Gillespieet al\.\(2022\)The Reactome pathway knowledgebase 2022\.Nucleic Acids Research50,pp\. D687–D692\.External Links:[Document](https://dx.doi.org/10.1093/nar/gkab1028)Cited by:[§3\.11](https://arxiv.org/html/2608.05928#S3.SS11.p1.1)\.
- \[9\]M\. Haoet al\.\(2024\)Large\-scale foundation model on single\-cell transcriptomics\.Nature Methods21,pp\. 1481–1491\.External Links:[Document](https://dx.doi.org/10.1038/s41592-024-02305-7)Cited by:[Context relative to reported foundation\-model corpora\.](https://arxiv.org/html/2608.05928#Ax2.SSx2.SSS0.Px1.p1.1),[§1\.3](https://arxiv.org/html/2608.05928#S1.SS3.p4.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p5.1)\.
- \[10\]A\. Katharopoulos, A\. Vyas, N\. Pappas, and F\. Fleuret\(2020\)Transformers are RNNs: fast autoregressive transformers with linear attention\.InProceedings of the 37th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.119,pp\. 5156–5165\.External Links:[Link](https://proceedings.mlr.press/v119/katharopoulos20a.html)Cited by:[§1\.3](https://arxiv.org/html/2608.05928#S1.SS3.p1.1),[§3\.6](https://arxiv.org/html/2608.05928#S3.SS6.p1.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p5.1)\.
- \[11\]K\. Z\. Kedzierska, L\. Crawford, A\. P\. Amini, and A\. X\. Lu\(2025\)Zero\-shot evaluation reveals limitations of single\-cell foundation models\.Genome Biology26,pp\. 101\.External Links:[Document](https://dx.doi.org/10.1186/s13059-025-03574-x)Cited by:[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1)\.
- \[12\]Y\. LeCun\(2022\)A path towards autonomous machine intelligence\.Note:OpenReviewExternal Links:[Link](https://openreview.net/forum?id=BZ5a1r-kVsf)Cited by:[Relation to established JEPA objectives\.](https://arxiv.org/html/2608.05928#Ax2.SSx6.SSS0.Px1.p1.2),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p3.1)\.
- \[13\]R\. Lopez, J\. Regier, M\. B\. Cole, M\. I\. Jordan, and N\. Yosef\(2018\)Deep generative modeling for single\-cell transcriptomics\.Nature Methods15,pp\. 1053–1058\.External Links:[Document](https://dx.doi.org/10.1038/s41592-018-0229-2)Cited by:[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p1.1)\.
- \[14\]T\. M\. Norman, M\. A\. Horlbeck, J\. M\. Replogle, A\. Y\. Ge, A\. Xu, M\. Jost, L\. A\. Gilbert, and J\. S\. Weissman\(2019\)Exploring genetic interaction manifolds constructed from rich single\-cell phenotypes\.Science365\(6455\),pp\. 786–793\.External Links:[Document](https://dx.doi.org/10.1126/science.aax4438)Cited by:[§1\.6](https://arxiv.org/html/2608.05928#S1.SS6.p3.2)\.
- \[15\]J\. Raingeaud, A\. J\. Whitmarsh, T\. Barrett, B\. Dérijard, and R\. J\. Davis\(1996\)MKK3\- and MKK6\-regulated gene expression is mediated by the p38 mitogen\-activated protein kinase signal transduction pathway\.Molecular and Cellular Biology16\(3\),pp\. 1247–1255\.External Links:[Document](https://dx.doi.org/10.1128/MCB.16.3.1247)Cited by:[Compositional perturbation geometry\.](https://arxiv.org/html/2608.05928#Ax2.SSx14.SSS0.Px4.p1.8),[§1\.6](https://arxiv.org/html/2608.05928#S1.SS6.p3.2)\.
- \[16\]Å\. Segerstolpe, A\. Palasantza, P\. Eliasson, E\. Andersson, A\. Andréasson, X\. Sun, S\. Picelli, A\. Sabirsh, M\. Clausen, M\. K\. Bjursell, D\. M\. Smith, M\. Kasper, C\. Ämmälä, and R\. Sandberg\(2016\)Single\-cell transcriptome profiling of human pancreatic islets in health and type 2 diabetes\.Cell Metabolism24\(4\),pp\. 593–607\.External Links:[Document](https://dx.doi.org/10.1016/j.cmet.2016.08.020)Cited by:[§1\.6](https://arxiv.org/html/2608.05928#S1.SS6.p1.1)\.
- \[17\]R\. Singh, W\. M\. Hoogaars, P\. Barnett, T\. Grieskamp, M\. S\. Rana, H\. Buermans, H\. F\. Farin, M\. Petry, T\. Heallen, J\. F\. Martin, A\. F\. M\. Moorman, P\. A\. C\. ’t Hoen, A\. Kispert, and V\. M\. Christoffels\(2012\)Tbx2 and Tbx3 induce atrioventricular myocardial development and endocardial cushion formation\.Cellular and Molecular Life Sciences69\(8\),pp\. 1377–1389\.External Links:[Document](https://dx.doi.org/10.1007/s00018-011-0884-2)Cited by:[Compositional perturbation geometry\.](https://arxiv.org/html/2608.05928#Ax2.SSx14.SSS0.Px4.p1.8),[§1\.6](https://arxiv.org/html/2608.05928#S1.SS6.p3.2)\.
- \[18\]D\. Szklarczyk, R\. Kirsch, M\. Koutrouli, K\. Nastou, F\. Mehryary, R\. Hachilif, A\. L\. Gable, T\. Fang, N\. T\. Doncheva, S\. Pyysalo, P\. Bork, L\. J\. Jensen, and C\. von Mering\(2023\)The STRING database in 2023: protein–protein association networks and functional enrichment analyses for any sequenced genome of interest\.Nucleic Acids Research51\(D1\),pp\. D638–D646\.External Links:[Document](https://dx.doi.org/10.1093/nar/gkac1000)Cited by:[STRING protein\-association adjacency\.](https://arxiv.org/html/2608.05928#Ax2.SSx5.SSS0.Px1.p1.1),[§3\.2](https://arxiv.org/html/2608.05928#S3.SS2.p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p4.1)\.
- \[19\]A\. Tejada\-Lapuertaet al\.\(2025\)Nicheformer: a foundation model for single\-cell and spatial omics\.Nature Methods22,pp\. 2525–2538\.External Links:[Document](https://dx.doi.org/10.1038/s41592-025-02814-z)Cited by:[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1)\.
- \[20\]C\. V\. Theodoriset al\.\(2023\)Transfer learning enables predictions in network biology\.Nature618,pp\. 616–624\.External Links:[Document](https://dx.doi.org/10.1038/s41586-023-06139-9)Cited by:[Context relative to reported foundation\-model corpora\.](https://arxiv.org/html/2608.05928#Ax2.SSx2.SSS0.Px1.p1.1),[§1\.3](https://arxiv.org/html/2608.05928#S1.SS3.p4.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p5.1)\.
- \[21\]Y\. Xu, Y\. Li, Y\. Yuan, C\. Yu, and Z\. Zang\(2026\)CellBench\-LS: benchmark evaluation of single\-cell foundation models for low\-supervision scenarios\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.64898/2026.04.01.714123)Cited by:[Temporal audit against the pretraining cutoff\.](https://arxiv.org/html/2608.05928#Ax2.SSx2.SSS0.Px3.p1.1),[§3\.10](https://arxiv.org/html/2608.05928#S3.SS10.p1.2),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p6.1)\.
- \[22\]N\. D\. Youngblut, C\. Carpenter, A\. Nayebnazar, A\. Adduri, R\. Shah, C\. Ricci\-Tam, J\. Prashar, R\. Ilango, N\. Teyssier, S\. Konermann, P\. D\. Hsu, A\. Dobin, D\. P\. Burke, H\. Goodarzi, and Y\. H\. Roohani\(2025\)scBaseCount: an AI agent\-curated, uniformly processed, and autonomously updated single\-cell data repository\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2025.02.27.640494)Cited by:[Pretraining corpus, sampling and exposure accounting](https://arxiv.org/html/2608.05928#Ax2.SSx2.p1.1),[§3\.8](https://arxiv.org/html/2608.05928#S3.SS8.p1.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p5.1)\.
- \[23\]Y\. Zenget al\.\(2025\)CellFM: a large\-scale foundation model pre\-trained on transcriptomics of 100 million human cells\.Nature Communications16,pp\. 4679\.External Links:[Document](https://dx.doi.org/10.1038/s41467-025-59926-5)Cited by:[Context relative to reported foundation\-model corpora\.](https://arxiv.org/html/2608.05928#Ax2.SSx2.SSS0.Px1.p1.1),[§1\.3](https://arxiv.org/html/2608.05928#S1.SS3.p4.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p2.1),[BioM\-JEPA: joint\-embedding prediction of graph\-connected gene blocks in single cells](https://arxiv.org/html/2608.05928#p5.1)\.

## Supplementary Information

## Supplementary Methods

### Study design and notation

We evaluated BioM\-JEPA as a frozen embedding model and used the same saved representation for annotation, clustering, reconstruction and perturbation prediction\. Letxi∈ℝGx\_\{i\}\\in\\mathbb\{R\}^\{G\}denote the canonical expression vector of cellii,yi∈\{1,…,C\}y\_\{i\}\\in\\\{1,\\ldots,C\\\}its cell\-type label andzi=f​\(xi\)∈ℝdz\_\{i\}=f\(x\_\{i\}\)\\in\\mathbb\{R\}^\{d\}the frozen model embedding\. All supervised heads were trained after freezingff\. We used fixed split identities and the same random seeds for every model within a task\. The reported benchmark models were UMAP, scVI, CellPLM, Geneformer, LangCell, scGPT, scMulan, scFoundation, Nicheformer and BioM\-JEPA where outputs were available\.

### Pretraining corpus, sampling and exposure accounting

The pretraining shards were prepared from a scBaseCount snapshot containing records available no later than 25 February 2025\. scBaseCount is an AI\-agent\-curated and uniformly processed single\-cell repository\[[22](https://arxiv.org/html/2608.05928#bib.bib2)\]; the current source repository contains more than 502 million cells across organisms, tissues and studies\. The present model was trained on a curated shard collection rather than for one full pass over the complete repository\.

The pretraining collection contained 442 CSR sparse matrices, each with 50,000 rows and 19,264 columns, for a total of

Nlocal=442×50,000=22,100,000N\_\{\\rm local\}=442\\times 50\{,\}000=22\{,\}100\{,\}000cells\. Shard order and row order within each shard were shuffled\. Shards were partitioned across eight distributed ranks and then across eight data\-loader workers per rank, so each row was assigned to one rank–worker pair during a traversal\.

Pretraining used eight GPUs, a local batch of eight cells per GPU and eight gradient\-accumulation batches\. IfBlocalB\_\{\\rm local\}is the local batch,DDthe number of distributed workers andAAthe accumulation factor, one optimizer step represents

Beff=Blocal​D​A=8×8×8=512B\_\{\\rm eff\}=B\_\{\\rm local\}DA=8\\times 8\\times 8=512cell presentations\. The representation analysed in this study was obtained at optimizer steps=20,000s=20\{,\}000, giving

Npresent=s​Beff=10,240,000N\_\{\\rm present\}=sB\_\{\\rm eff\}=10\{,\}240\{,\}000cell presentations\. One complete local pass requires approximately

Nlocal/Beff=43,164\.1N\_\{\\rm local\}/B\_\{\\rm eff\}=43\{,\}164\.1optimizer steps\. The analysed representation therefore precedes the first complete pass and covers 46\.33% of the local rows, with no intentional row resampling\. Apart from incomplete worker batches discarded by sortish batching,Nunique≃NpresentN\_\{\\rm unique\}\\simeq N\_\{\\rm present\}\. Relative to the 502\-million\-cell resource, 10\.24 million corresponds to 2\.04% on a presentation\-equivalent scale\.

#### Context relative to reported foundation\-model corpora\.

The 10\.24 million presentations are lower than the published pretraining\-corpus counts for Geneformer \(approximately 30 million cells\), scGPT \(more than 33 million\), scFoundation \(more than 50 million\) and CellFM \(102,304,686 cells\)\[[20](https://arxiv.org/html/2608.05928#bib.bib4),[6](https://arxiv.org/html/2608.05928#bib.bib5),[9](https://arxiv.org/html/2608.05928#bib.bib6),[23](https://arxiv.org/html/2608.05928#bib.bib8)\]\. Table[S1](https://arxiv.org/html/2608.05928#Ax2.T1)reports these quantities without converting them into nominal epochs\.

Table S1:Cell scale used to contextualize pretraining exposure\. The BioM\-JEPA value is the number of cell presentations used in this study; comparator values are published corpus sizes\.These quantities place the training exposure in context: the reported representation emerged after fewer cell presentations than the published corpus size of each comparator\. Because the models differ in architecture, vocabulary and preprocessing, Table[S1](https://arxiv.org/html/2608.05928#Ax2.T1)is an exposure comparison rather than a fitted scaling law\.

#### Evaluation\-overlap audit\.

No benchmark labels or task\-specific train–test splits were used during pretraining or graph construction\. The 442 local shards contain sparse expression matrices but no study\-accession or cell\-source manifest; study\-level overlap with the public datasets underlying CellBench\-LS therefore cannot be excluded\.

Table S2:Pretraining and graph\-overlap status\. Matrix\-only shards do not permit accession\-level exclusion of the public studies underlying the benchmark\.We therefore describe annotation as within\-dataset frozen\-probe evaluation\. The same terminology is applied to comparator models for which accession\-resolved pretraining manifests are unavailable\.

#### Temporal audit against the pretraining cutoff\.

The CellBench\-LS manuscript was posted in 2026, after the 25 February 2025 scBaseCount snapshot cutoff, but publication of the benchmark is distinct from generation of its constituent data\. We audited the dataset descriptions and original sources cited by CellBench\-LS\[[21](https://arxiv.org/html/2608.05928#bib.bib17)\], together with the additional Dixit perturbation dataset used by our CellBench adapter\. The latest source years for the general\-purpose datasets range from 2017 to 2024; the batch\-correction datasets trace to 2021, and the perturbation datasets trace to 2016–2019\. Thus, no benchmark dataset used here can be certified as generated after the pretraining cutoff \(Table[S3](https://arxiv.org/html/2608.05928#Ax2.T3)\)\. The audit establishes that the present benchmark suite is not a temporal holdout; it does not establish that its studies were included in the matrix\-only shards\.

Table S3:Benchmark temporal audit relative to 25 February 2025\. Years denote source\-publication dates reported by CellBench\-LS or, for the additional Dixit dataset, its original study; they are not necessarily the first date on which every raw file became public\.

### CellBench\-LS datasets and preprocessing

Annotation and clustering used the nine CellBench\-LS datasets PBMC12k, hPancreas, cortex, multiple sclerosis \(MS\), liver, Zheng68k, lung, COVID and immune\. Reconstruction used hPancreas, cortex, MS, liver, Zheng68k, lung, COVID and immune\. Perturbation prediction used Adamson, Norman and Dixit; combination perturbations occur in Norman\. Raw count matrices were retained where required\. Gene identifiers were mapped to each model’s published vocabulary, duplicated identifiers were resolved before inference and cells lacking a valid benchmark label were removed\.

Each model retained its documented tokenization and value transformation\. BioM\-JEPA mapped sparse raw counts to the 19,264\-gene vocabulary and appliedvi​g=log⁡\(1\+104​ci​g/∑g′ci​g′\)v\_\{ig\}=\\log\(1\+10^\{4\}c\_\{ig\}/\\sum\_\{g^\{\\prime\}\}c\_\{ig^\{\\prime\}\}\)\. Two scFoundation\-style control values represented a fixed high\-resolution value of 4 andlog10\\log\_\{10\}library size\. They were visible context entries for the student and teacher but were excluded from graph\-target eligibility\. Frozen BioM\-JEPA embeddings excluded these positions and concatenated the element\-wise maximum gene\-token state with the mean of the five largest gene\-token states per feature, producing a 1,536\-dimensional embedding\. The same pooling and control\-position policy was used for the objective controls\.

### Linear\-attention architecture

BioM\-JEPA used a 19,264\-gene vocabulary, hidden widthd=768d=768, twelve encoder layers and twelve heads per layer\. All encoder layers used flash linear attention\. The prediction head used four linear\-attention layers, twelve heads and dropout0\.050\.05\. The student was fully trainable\. The teacher shared the encoder architecture and was updated only by EMA with momentum annealed from0\.9960\.996to0\.99970\.9997\.

For one head, letQ,K,V∈ℝn×dhQ,K,V\\in\\mathbb\{R\}^\{n\\times d\_\{h\}\}andφ​\(u\)=ELU⁡\(u\)\+1\\varphi\(u\)=\\operatorname\{ELU\}\(u\)\+1\. Linear attention evaluates

yi=φ​\(qi\)⊤​\[∑jφ​\(kj\)​vj⊤\]φ​\(qi\)⊤​\[∑jφ​\(kj\)\]\+10−6,y\_\{i\}=\\frac\{\\varphi\(q\_\{i\}\)^\{\\top\}\\left\[\\sum\_\{j\}\\varphi\(k\_\{j\}\)v\_\{j\}^\{\\top\}\\right\]\}\{\\varphi\(q\_\{i\}\)^\{\\top\}\\left\[\\sum\_\{j\}\\varphi\(k\_\{j\}\)\\right\]\+10^\{\-6\}\},after setting padded keys and values to zero\. The numerator and normalizer are contracted once rather than materializing ann×nn\\times ntoken\-affinity matrix\. Withhhheads anddh=d/hd\_\{h\}=d/h, the attention\-core complexity is

𝒪​\(h​n​dh2\)=𝒪​\(n​d2h\),\\mathcal\{O\}\\\!\\left\(hnd\_\{h\}^\{2\}\\right\)=\\mathcal\{O\}\\\!\\left\(\\frac\{nd^\{2\}\}\{h\}\\right\),compared with𝒪​\(n2​d\)\\mathcal\{O\}\(n^\{2\}d\)for dense softmax attention\. The asymptotic statement assumes fixed width and does not include tokenization, feed\-forward layers or input/output transfer\. We therefore report it separately from measured end\-to\-end throughput\.

### Biological graph and target blocks

The implemented target graph was restricted to the shared gene vocabulary𝒱\\mathcal\{V\}, with\|𝒱\|=19,264\|\\mathcal\{V\}\|=19\{,\}264\. It used two sources\.

#### STRING protein\-association adjacency\.

STRING v12 human links were mapped to the vocabulary by preferred symbol and unambiguous aliases\[[18](https://arxiv.org/html/2608.05928#bib.bib3)\]\. Links with combined score below 700 and self\-edges were removed\. For each row, the 64 highest\-scoring mapped neighbours were retained and the result was symmetrized\. This graph mapped 17,892 STRING proteins and contained 356,772 nonzero adjacency entries\.

#### Corpus coexpression adjacency\.

All 22\.1 million local cells were transformed as

xi​g′=log⁡\(1\+104​xi​g/∑hxi​h\)\.x\_\{ig\}^\{\\prime\}=\\log\\\!\\left\(1\+10^\{4\}x\_\{ig\}/\\sum\_\{h\}x\_\{ih\}\\right\)\.Centred random projections with 256 projection dimensions approximated the absolute Pearson association between every gene pair\. The 64 strongest neighbours per gene were retained and symmetrized, producing 2,360,480 nonzero entries\. Thus, the coexpression graph was estimated from unlabeled expression matrices; it did not use a benchmark label, but accession overlap cannot be excluded for the reason above\.

The STRING and corpus\-coexpression adjacencies were binarized, combined by union and truncated row\-wise to 64 neighbours:

A=RowTop64⁡\(𝟏​\[ASTRING\>0\]∨𝟏​\[Acoexp\>0\]\)\.A=\\operatorname\{RowTop\}\_\{64\}\\left\(\\mathbf\{1\}\[A^\{\\rm STRING\}\>0\]\\lor\\mathbf\{1\}\[A^\{\\rm coexp\}\>0\]\\right\)\.Because the final row\-wise truncation was applied after the union,AAis a directed neighbourhood table rather than a guaranteed symmetric matrix\. It contains 1,123,337 entries and has mean out\-degree 58\.31\. No pathway\-membership table, transcription\-factor direction, cell\-type label or perturbation\-pair identity enteredAA\.

#### Operational target\-block definition\.

In this work a graph\-connected gene block is a*candidate target set*underAA\. Biological programme and pathway are reserved for independently supported gene sets\. For cellii, let𝒪i\\mathcal\{O\}\_\{i\}be its observed genes\. For each ofK=4K=4targets, we draw a seed

ri​k∼Uniform⁡\(𝒪i\)r\_\{ik\}\\sim\\operatorname\{Uniform\}\(\\mathcal\{O\}\_\{i\}\)and an integer requested sizeLi​k∼Uniform⁡\{2000,…,8000\}L\_\{ik\}\\sim\\operatorname\{Uniform\}\\\{2000,\\ldots,8000\\\}\. Breadth\-first expansion over outgoing neighbour lists gives

Bi​k=BFS⁡\(A,ri​k,Li​k\)\.B\_\{ik\}=\\operatorname\{BFS\}\(A,r\_\{ik\},L\_\{ik\}\)\.The cell\-specific target support is not all 2,000–8,000 candidate genes, but only the observed intersection

Ti​k=Bi​k∩𝒪i\.T\_\{ik\}=B\_\{ik\}\\cap\\mathcal\{O\}\_\{i\}\.The teacher block state is averaged overTi​kT\_\{ik\}\. The query identifier is the median vocabulary index inBi​kB\_\{ik\}, denotedci​kc\_\{ik\}, and contains a learned position embedding plus a learned mask vector but no target expression value\.

The nominal student context is

𝒞i=𝒪i∖⋃k=1KTi​k\.\\mathcal\{C\}\_\{i\}=\\mathcal\{O\}\_\{i\}\\setminus\\bigcup\_\{k=1\}^\{K\}T\_\{ik\}\.If\|𝒞i\|<512\|\\mathcal\{C\}\_\{i\}\|<512, the minimum\-context safeguard restores𝒞i=𝒪i\\mathcal\{C\}\_\{i\}=\\mathcal\{O\}\_\{i\}rather than presenting an uninformative short sequence\. This permits target values to remain visible in the affected cells\. On the student side, query\-conditioned pooling first restricts context candidates to

𝒞i​kDMT=𝒞i∩𝒩64coexp​\(ci​k\)\\mathcal\{C\}\_\{ik\}^\{\\rm DMT\}=\\mathcal\{C\}\_\{i\}\\cap\\mathcal\{N\}^\{\\rm coexp\}\_\{64\}\(c\_\{ik\}\)and falls back to𝒞i\\mathcal\{C\}\_\{i\}when the intersection is empty\. The teacher always receives the full observed gene set\.

#### Target\-visibility audit\.

We quantified the minimum\-context safeguard on 10,000 cells sampled uniformly from 16 of the 442 training shards, using four graph blocks, requested block sizes of 2,000–8,000 genes, the 64\-neighbour graph, the 512\-token threshold and random seed 42\. For every sampled cell, we reconstructed the realized target union, residual context and safeguard decision\.

Mean target coverage was 0\.6932, consistent with the training mean of 0\.6917 over 2,400 batches\. The safeguard restored the observed context in 289 of 10,000 cells \(2\.89%; 95% Wilson interval, 2\.58–3\.24%\)\. Those cells contained fewer observed genes than the remaining cells \(mean, 1,804\.8 versus 3,796\.4\) and had a mean residual context of 418\.6 tokens\. Weighting by target tokens, 401,208 of 25,821,660 targets were visible \(1\.55%\)\.

Table S4:Minimum\-context target\-visibility audit\. Intervals for counts show the median and interquartile range unless stated otherwise\.Target hiding governed 97\.11% of sampled cells and 98\.45% of target\-token exposures; the remaining minority followed the quantified minimum\-context safeguard\.

The random\-block control sampled candidate gene indices uniformly with requested ratio in\[0\.104,0\.415\]\[0\.104,0\.415\], matching the 2,000–8,000 candidate\-size interval, and used neither graph expansion nor DMT\. It was a size\-matched objective\-family control; it did not match graph degree or the realized number of observed target genes\.

### Block aggregation and BioM\-JEPA loss

#### Relation to established JEPA objectives\.

The general JEPA principle predicts a target state in representation space instead of reconstructing the raw input\[[12](https://arxiv.org/html/2608.05928#bib.bib13)\]\. In I\-JEPA, a context encoder and predictor estimate a collection of teacher patch embeddings for masked image regions\[[2](https://arxiv.org/html/2608.05928#bib.bib14)\]; V\-JEPA applies the same feature\-prediction principle to masked spatiotemporal regions\[[4](https://arxiv.org/html/2608.05928#bib.bib15)\]\. Denoting the patch or token targets in blockTkT\_\{k\}by\{hit:i∈Tk\}\\\{h\_\{i\}^\{t\}:i\\in T\_\{k\}\\\}, the corresponding token\-level objective has the generic form

ℒtoken=1∑k\|Tk\|​∑k∑i∈Tkℓ​\(h^i,sg⁡\(hit\)\)\.\\mathcal\{L\}\_\{\\rm token\}=\\frac\{1\}\{\\sum\_\{k\}\|T\_\{k\}\|\}\\sum\_\{k\}\\sum\_\{i\\in T\_\{k\}\}\\ell\\\!\\left\(\\widehat\{h\}\_\{i\},\\operatorname\{sg\}\(h\_\{i\}^\{t\}\)\\right\)\.BioM\-JEPA changes the supervision unit rather than the teacher–student principle\. It first maps each graph\-defined gene set to one teacher module state

tk=sg⁡\[Agg⁡\{hit:i∈Tk\}−c\],t\_\{k\}=\\operatorname\{sg\}\\\!\\left\[\\operatorname\{Agg\}\\\{h\_\{i\}^\{t\}:i\\in T\_\{k\}\\\}\-c\\right\],and then aligns one predicted statet^k\\widehat\{t\}\_\{k\}withtkt\_\{k\}\. Consequently, the implemented loss containsKKaggregate block\-level prediction terms, not∑k\|Tk\|\\sum\_\{k\}\|T\_\{k\}\|gene\-level prediction terms\.

The student and EMA teacher produced token matrices

His=fθ​\(xi,𝒞\),Hit=fθ¯​\(xi\),θ¯←μ​θ¯\+\(1−μ\)​θ\.H\_\{i\}^\{s\}=f\_\{\\theta\}\(x\_\{i,\\mathcal\{C\}\}\),\\qquad H\_\{i\}^\{t\}=f\_\{\\bar\{\\theta\}\}\(x\_\{i\}\),\\qquad\\bar\{\\theta\}\\leftarrow\\mu\\bar\{\\theta\}\+\(1\-\\mu\)\\theta\.For blockTkT\_\{k\}with representative vocabulary indexckc\_\{k\}, the query was

qk=e​\(ck\)\+m,q\_\{k\}=e\(c\_\{k\}\)\+m,whereckc\_\{k\}is the median vocabulary index of the sampled candidate set,e​\(ck\)e\(c\_\{k\}\)is its learned gene\-position embedding andmmis a learned mask embedding\. This query communicates a reproducible target\-block identifier but contains no target expression\. The reported DMT aggregator first restricted context candidates to𝒞i​k=𝒞i∩𝒩64coexp​\(ck\)\\mathcal\{C\}\_\{ik\}=\\mathcal\{C\}\_\{i\}\\cap\\mathcal\{N\}\_\{64\}^\{\\rm coexp\}\(c\_\{k\}\)and fell back to the complete visible context when the intersection was empty\. It then summarized the visible student context:

si​k=∑j∈𝒞i​kai​k​j​hi​js,ai​k​j=exp⁡\{cos⁡\(qk,hi​js\)/τ\}∑ℓ∈𝒞i​kexp⁡\{cos⁡\(qk,hi​ℓs\)/τ\}\.s\_\{ik\}=\\sum\_\{j\\in\\mathcal\{C\}\_\{ik\}\}a\_\{ikj\}h^\{s\}\_\{ij\},\\qquad a\_\{ikj\}=\\frac\{\\exp\\\{\\cos\(q\_\{k\},h^\{s\}\_\{ij\}\)/\\tau\\\}\}\{\\sum\_\{\\ell\\in\\mathcal\{C\}\_\{ik\}\}\\exp\\\{\\cos\(q\_\{k\},h^\{s\}\_\{i\\ell\}\)/\\tau\\\}\}\.The teacher block state was aggregated before any prediction error was formed,

ti​k=sg⁡\[\|Tk\|−1​∑j∈Tkhi​jt−c\],t^i​k=pϕ​\(si​k\),t\_\{ik\}=\\operatorname\{sg\}\\\!\\left\[\|T\_\{k\}\|^\{\-1\}\\sum\_\{j\\in T\_\{k\}\}h^\{t\}\_\{ij\}\-c\\right\],\\qquad\\widehat\{t\}\_\{ik\}=p\_\{\\phi\}\(s\_\{ik\}\),whereccis an EMA centre andsg\\operatorname\{sg\}stops gradients\. For the valid cell–block index setℬ\\mathcal\{B\}, alignment was

ℒalign=λ2\|ℬ\|​∑\(i,k\)∈ℬ‖t^i​k−ti​k‖22\+λs\|ℬ\|​∑\(i,k\)∈ℬ‖t^i​k‖t^i​k‖2−ti​k‖ti​k‖2‖22\.\\mathcal\{L\}\_\{\\rm align\}=\\frac\{\\lambda\_\{2\}\}\{\|\\mathcal\{B\}\|\}\\sum\_\{\(i,k\)\\in\\mathcal\{B\}\}\\\|\\widehat\{t\}\_\{ik\}\-t\_\{ik\}\\\|\_\{2\}^\{2\}\+\\frac\{\\lambda\_\{s\}\}\{\|\\mathcal\{B\}\|\}\\sum\_\{\(i,k\)\\in\\mathcal\{B\}\}\\left\\\|\\frac\{\\widehat\{t\}\_\{ik\}\}\{\\\|\\widehat\{t\}\_\{ik\}\\\|\_\{2\}\}\-\\frac\{t\_\{ik\}\}\{\\\|t\_\{ik\}\\\|\_\{2\}\}\\right\\\|\_\{2\}^\{2\}\.LetT^∈ℝ\|ℬ\|×d\\widehat\{T\}\\in\\mathbb\{R\}^\{\|\\mathcal\{B\}\|\\times d\}contain the predicted block states\. We used

ℒvar=1d∑r=1dmax\{0,γ−sd\(T^⋅r\)\},ℒcov=1d∑r≠sCov\(T^\)r​s2,\\mathcal\{L\}\_\{\\rm var\}=\\frac\{1\}\{d\}\\sum\_\{r=1\}^\{d\}\\max\\\{0,\\gamma\-\\operatorname\{sd\}\(\\widehat\{T\}\_\{\\cdot r\}\)\\\},\\qquad\\mathcal\{L\}\_\{\\rm cov\}=\\frac\{1\}\{d\}\\sum\_\{r\\neq s\}\\operatorname\{Cov\}\(\\widehat\{T\}\)\_\{rs\}^\{2\},We usedλ2=1\\lambda\_\{2\}=1,λs=0\\lambda\_\{s\}=0,λv=0\.05\\lambda\_\{v\}=0\.05andλc=0\.01\\lambda\_\{c\}=0\.01\. Thus, the active alignment was block\-level mean\-squared error and the optimized objective was

ℒBioM​\-​JEPA=ℒalign\+λv​ℒvar\+λc​ℒcov\.\\boxed\{\\mathcal\{L\}\_\{\\rm BioM\\text\{\-\}JEPA\}=\\mathcal\{L\}\_\{\\rm align\}\+\\lambda\_\{v\}\\mathcal\{L\}\_\{\\rm var\}\+\\lambda\_\{c\}\\mathcal\{L\}\_\{\\rm cov\}\}\.No individual target\-gene term is present\. The token\-IJEPA control instead optimized

ℒtoken=1\|𝒯\|​∑\(i,j\)∈𝒯ℓ​\(h^i​j,sg⁡\(hi​jt\)\)\.\\mathcal\{L\}\_\{\\rm token\}=\\frac\{1\}\{\|\\mathcal\{T\}\|\}\\sum\_\{\(i,j\)\\in\\mathcal\{T\}\}\\ell\\\!\\left\(\\widehat\{h\}\_\{ij\},\\operatorname\{sg\}\(h^\{t\}\_\{ij\}\)\\right\)\.The decoder\-only control retained graph blocks and optimized expression decoding withoutℒBioM​\-​JEPA\\mathcal\{L\}\_\{\\rm BioM\\text\{\-\}JEPA\}\.

Algorithm S1BioM\-JEPA pretraining1:Minibatch

XX; gene graph

𝒢\\mathcal\{G\}; student

fθf\_\{\\theta\}; teacher

fθ¯f\_\{\\bar\{\\theta\}\}; predictor

pϕp\_\{\\phi\}; EMA rate

μ\\mu
2:foreach minibatch

XXdo

3:

𝒯←SampleGraphBlocks​\(X,𝒢\)\\mathcal\{T\}\\leftarrow\\textsc\{SampleGraphBlocks\}\(X,\\mathcal\{G\}\)
4:

𝒞←𝒪​\(X\)∖⋃T∈𝒯T\\mathcal\{C\}\\leftarrow\\mathcal\{O\}\(X\)\\setminus\\bigcup\_\{T\\in\\mathcal\{T\}\}T
5:if

\|𝒞i\|<512\|\\mathcal\{C\}\_\{i\}\|<512for cell

iithen

6:

𝒞i←𝒪i\\mathcal\{C\}\_\{i\}\\leftarrow\\mathcal\{O\}\_\{i\}⊳\\trianglerightminimum\-context safeguard

7:endif

8:

Hs←fθ​\(X𝒞\)H^\{s\}\\leftarrow f\_\{\\theta\}\(X\_\{\\mathcal\{C\}\}\)
9:

Ht←StopGradient​\(fθ¯​\(X\)\)H^\{t\}\\leftarrow\\textsc\{StopGradient\}\(f\_\{\\bar\{\\theta\}\}\(X\)\)
10:

𝒵^←\[\]\\widehat\{\\mathcal\{Z\}\}\\leftarrow\[\\,\];

𝒵←\[\]\\mathcal\{Z\}\\leftarrow\[\\,\]
11:foreach valid target block

Tk∈𝒯T\_\{k\}\\in\\mathcal\{T\}do

12:

qk←e​\(ck\)\+mq\_\{k\}\\leftarrow e\(c\_\{k\}\)\+m
13:

sk←ContextPool​\(Hs,qk,𝒩64coexp​\(ck\)\)s\_\{k\}\\leftarrow\\textsc\{ContextPool\}\(H^\{s\},q\_\{k\},\\mathcal\{N\}^\{\\rm coexp\}\_\{64\}\(c\_\{k\}\)\)
14:

tk←MeanPool​\(HTkt\)−ct\_\{k\}\\leftarrow\\textsc\{MeanPool\}\(H^\{t\}\_\{T\_\{k\}\}\)\-c
15:

𝒵^←𝒵^∥pϕ​\(sk\)\\widehat\{\\mathcal\{Z\}\}\\leftarrow\\widehat\{\\mathcal\{Z\}\}\\mathbin\{\\\|\}p\_\{\\phi\}\(s\_\{k\}\)
16:

𝒵←𝒵∥StopGradient​\(tk\)\\mathcal\{Z\}\\leftarrow\\mathcal\{Z\}\\mathbin\{\\\|\}\\textsc\{StopGradient\}\(t\_\{k\}\)
17:endfor

18:

ℒ←MSE​\(𝒵^,𝒵\)\+λv​VariancePenalty​\(𝒵^\)\+λc​CovariancePenalty​\(𝒵^\)\\mathcal\{L\}\\leftarrow\\textsc\{MSE\}\(\\widehat\{\\mathcal\{Z\}\},\\mathcal\{Z\}\)\+\\lambda\_\{v\}\\textsc\{VariancePenalty\}\(\\widehat\{\\mathcal\{Z\}\}\)\+\\lambda\_\{c\}\\textsc\{CovariancePenalty\}\(\\widehat\{\\mathcal\{Z\}\}\)
19:

\(θ,ϕ\)←OptimizerStep​\(ℒ,θ,ϕ\)\(\\theta,\\phi\)\\leftarrow\\textsc\{OptimizerStep\}\(\\mathcal\{L\},\\theta,\\phi\)
20:

θ¯←μ​θ¯\+\(1−μ\)​θ\\bar\{\\theta\}\\leftarrow\\mu\\bar\{\\theta\}\+\(1\-\\mu\)\\theta
21:

c←UpdateEMACentre​\(c,𝒵\)c\\leftarrow\\textsc\{UpdateEMACentre\}\(c,\\mathcal\{Z\}\)
22:endfor

### Objective curves and target similarity

Raw objectives have different units; therefore, the curve in Fig\. 2a was normalized within run by the median loss during the first 400 optimization steps:

ℒ~r​\(s\)=ℒr​\(s\)medianu≤400⁡ℒr​\(u\)\.\\widetilde\{\\mathcal\{L\}\}\_\{r\}\(s\)=\\frac\{\\mathcal\{L\}\_\{r\}\(s\)\}\{\\operatorname\{median\}\_\{u\\leq 400\}\\mathcal\{L\}\_\{r\}\(u\)\}\.We plotted a centred rolling median and summarized local dispersion with the rolling standard deviation of residuals\. At every 800\-step training snapshot, we quantified target agreement with orthogonal\-Procrustes cosine similarity\. For centred prediction and teacher matricesP,T∈ℝn×dP,T\\in\\mathbb\{R\}^\{n\\times d\},

R⋆=arg⁡minR⊤​R=I⁡‖P​R−T‖F2,SProc=1n​∑i=1ncos⁡\(Pi​R⋆,Ti\)\.R^\{\\star\}=\\arg\\min\_\{R^\{\\top\}R=I\}\\\|PR\-T\\\|\_\{F\}^\{2\},\\qquad S\_\{\\rm Proc\}=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\\cos\(P\_\{i\}R^\{\\star\},T\_\{i\}\)\.This rotation\-invariant measure compares the geometry of the two latent sets without requiring their coordinate bases to coincide\. It was evaluated at the native supervision unit: blocks for BioM\-JEPA and genes for token\-IJEPA\.

### Embedding geometry and depth association

For centred embedding matrixZ∈ℝN×dZ\\in\\mathbb\{R\}^\{N\\times d\}with covariance eigenvaluesλr\\lambda\_\{r\}, definepr=λr/∑jλjp\_\{r\}=\\lambda\_\{r\}/\\sum\_\{j\}\\lambda\_\{j\}\. Effective rank and participation ratio were

reff=exp⁡\(−∑rpr​log⁡pr\),rPR=\(∑rλr\)2∑rλr2\.r\_\{\\rm eff\}=\\exp\\\!\\left\(\-\\sum\_\{r\}p\_\{r\}\\log p\_\{r\}\\right\),\\qquad r\_\{\\rm PR\}=\\frac\{\(\\sum\_\{r\}\\lambda\_\{r\}\)^\{2\}\}\{\\sum\_\{r\}\\lambda\_\{r\}^\{2\}\}\.Fig\. 2b evaluated effective rank every 800 steps on hPancreas\. Fig\. 3c evaluated both statistics on hPancreas and cortex\.

For each cell we retained total countsditotald\_\{i\}^\{\\rm total\}and the number of detected genesdinnzd\_\{i\}^\{\\rm nnz\}\. Letuiu\_\{i\}be the score along the leading embedding axis andni=‖zi‖2n\_\{i\}=\\\|z\_\{i\}\\\|\_\{2\}\. The four technical associations were

\|r​\(u,dtotal\)\|,\|r​\(u,dnnz\)\|,\|r​\(n,dtotal\)\|,\|r​\(n,dnnz\)\|,\|r\(u,d^\{\\rm total\}\)\|,\\quad\|r\(u,d^\{\\rm nnz\}\)\|,\\quad\|r\(n,d^\{\\rm total\}\)\|,\\quad\|r\(n,d^\{\\rm nnz\}\)\|,whererris Pearson correlation\. The main control figure displays the detected\-gene associations\. The failure\-mode summary reports their mean over the retained diagnostics and uses dataset values as the uncertainty unit\.

These diagnostics quantify association along the leading embedding axis and embedding norm\. Random\-block, token\-IJEPA and decoder\-only provide comparisons among complete pretraining formulations; their interpretation is therefore confined to the tested objective families\.

### Few\-shot annotation

For each dataset, top\-kkannotation usedk∈\{1,5,9\}k\\in\\\{1,5,9\\\}labelled cells per class as the support set\. Support indices were sampled independently for each class, and all remaining labelled cells formed the held\-out evaluation set\. Splitting was label\-stratified\. A support\-vector machine was trained on frozen embeddings with identical hyperparameters for all models\. We used five fixed seeds,4242–4646\.

For classcc, precision, recall andF1F\_\{1\}were

Pc=TPcTPc\+FPc,Rc=TPcTPc\+FNc,F1,c=2​Pc​RcPc\+Rc\.P\_\{c\}=\\frac\{\\mathrm\{TP\}\_\{c\}\}\{\\mathrm\{TP\}\_\{c\}\+\\mathrm\{FP\}\_\{c\}\},\\qquad R\_\{c\}=\\frac\{\\mathrm\{TP\}\_\{c\}\}\{\\mathrm\{TP\}\_\{c\}\+\\mathrm\{FN\}\_\{c\}\},\\qquad F\_\{1,c\}=\\frac\{2P\_\{c\}R\_\{c\}\}\{P\_\{c\}\+R\_\{c\}\}\.Accuracy and macro\-F1F\_\{1\}were

Acc=1N​∑i𝟏​\[y^i=yi\],F1macro=1C​∑c=1CF1,c\.\\operatorname\{Acc\}=\\frac\{1\}\{N\}\\sum\_\{i\}\\mathbf\{1\}\[\\widehat\{y\}\_\{i\}=y\_\{i\}\],\\qquad F\_\{1\}^\{\\rm macro\}=\\frac\{1\}\{C\}\\sum\_\{c=1\}^\{C\}F\_\{1,c\}\.Figures 2d and 3b use Top\-5 macro\-F1F\_\{1\}on hPancreas and cortex\. Points are means over five seeds and error bars are one standard deviation unless stated otherwise\.

### Matched efficiency benchmark

We used two complementary timing protocols\. BioM\-JEPA and scFoundation were evaluated on the same NVIDIA A100\-SXM4\-80GB GPU with PyTorch 2\.7\.1, CUDA 11\.8 and bfloat16 computation\. The same ordered raw\-count cells and 19,264\-gene vocabulary were used\. Gene mapping and disk input/output were completed before each timed interval\.

#### Encoder\-only extraction\.

We prepared 2,048 cells from hPancreas and used batch size 8, eight warm\-up iterations and 40 synchronized timed iterations\. No artificial token\-length cap was applied\. hPancreas contained 14,818 available cells and had observed\-gene lengths of 69–1,659 \(median, 319; mean, 386\.8; 95th percentile, 931\)\. BioM\-JEPA used its 1,536\-dimensional maximum\-plus\-top\-five\-mean embedding; scFoundation used its documented four\-part 3,072\-dimensional control\-token/max/mean representation\. These are model\-native frozen embeddings rather than dimension\-matched projections\. IfNNcells are encoded between synchronized timest0t\_\{0\}andt1t\_\{1\}, throughput is

Remb=Nt1−t0\.R\_\{\\rm emb\}=\\frac\{N\}\{t\_\{1\}\-t\_\{0\}\}\.
Table S5:Encoder\-only embedding throughput\. Values are cells s\-1; the final column is the BioM\-JEPA/scFoundation ratio\.
#### One\-epoch partial fine\-tuning\.

We used 1,279 hPancreas training cells, 427 validation cells and 427 test cells, with batch size 8 and one epoch\. Both models received the same 19,264 genes plus two scFoundation\-style control values and formed the same four\-part pooled vector \(last control, second\-last control, gene\-wise maximum and gene\-wise mean; 3,072 dimensions\)\. Each used the identical head

LayerNorm→Linear⁡\(3072,6144\)→GELU→Linear⁡\(6144,C\)\.\\operatorname\{LayerNorm\}\\rightarrow\\operatorname\{Linear\}\(3072,6144\)\\rightarrow\\operatorname\{GELU\}\\rightarrow\\operatorname\{Linear\}\(6144,C\)\.The head learning rate was10−310^\{\-3\}, the encoder learning rate was5×10−55\\times 10^\{\-5\}, weight decay was0\.010\.01and AdamW coefficients were\(0\.9,0\.999\)\(0\.9,0\.999\)\. We froze each decoder and updated only the final encoder block, direct parameters of leaf normalization layers and the common head\. Direct rather than recursive normalization\-parameter selection was used so that wrapper modules could not unintentionally unfreeze a decoder\. The resulting encoder\-trainable counts were 7,121,664 for BioM\-JEPA and 7,137,536 for scFoundation; total trainable counts including the common head were 26,088,205 and 26,104,077\.

Training throughput includes forward propagation, backward propagation and the optimizer step,

Rtrain=Nglobaltstep,R\_\{\\rm train\}=\\frac\{N\_\{\\rm global\}\}\{t\_\{\\rm step\}\},whereas held\-out throughput is

Rheldout=Nvalid\+Ntesttvalid\+ttest\.R\_\{\\rm heldout\}=\\frac\{N\_\{\\rm valid\}\+N\_\{\\rm test\}\}\{t\_\{\\rm valid\}\+t\_\{\\rm test\}\}\.The held\-out loop computed no\-gradient encoder embeddings followed by the identical head\. We therefore report its 3\.76\-fold ratio as task\-time embedding throughput and distinguish it from the 1\.92\-fold isolated encoder\-only measurement\.

Table S6:One\-epoch hPancreas timing\. One A100 GPU, bfloat16 and batch size 8 were used\.The train–validation–test epoch took 66\.08 s for BioM\-JEPA and 352\.68 s for scFoundation, corresponding to a 5\.34\-fold wall\-clock speed\-up\.

### Clustering

We clustered frozen embeddings withkk\-means using the true number of benchmark classes, while withholding labels during clustering\. Cluster accuracy used the optimal Hungarian assignmentπ\\pi:

ACCclust=maxπ⁡1N​∑i𝟏​\[yi=π​\(ci\)\]\.\\operatorname\{ACC\}\_\{\\rm clust\}=\\max\_\{\\pi\}\\frac\{1\}\{N\}\\sum\_\{i\}\\mathbf\{1\}\[y\_\{i\}=\\pi\(c\_\{i\}\)\]\.Normalized mutual information \(NMI\), adjusted Rand index \(ARI\) and mean silhouette width \(ASW\) were computed from the same assignment\. With contingency countsna​bn\_\{ab\}, row sumsna⁣⋅n\_\{a\\cdot\}, column sumsn⋅bn\_\{\\cdot b\}andN=∑a​bna​bN=\\sum\_\{ab\}n\_\{ab\},

NMI=2​∑a​bna​bN​log⁡N​na​bna⁣⋅​n⋅bH​\(Y\)\+H​\(C\)\.\\operatorname\{NMI\}=\\frac\{2\\sum\_\{ab\}\\frac\{n\_\{ab\}\}\{N\}\\log\\frac\{Nn\_\{ab\}\}\{n\_\{a\\cdot\}n\_\{\\cdot b\}\}\}\{H\(Y\)\+H\(C\)\}\.ARI used the standard chance\-corrected pair\-count statistic,

ARI=∑a​b\(na​b2\)−∑a\(na⁣⋅2\)​∑b\(n⋅b2\)\(N2\)12​\[∑a\(na⁣⋅2\)\+∑b\(n⋅b2\)\]−∑a\(na⁣⋅2\)​∑b\(n⋅b2\)\(N2\)\.\\operatorname\{ARI\}=\\frac\{\\sum\_\{ab\}\\binom\{n\_\{ab\}\}\{2\}\-\\frac\{\\sum\_\{a\}\\binom\{n\_\{a\\cdot\}\}\{2\}\\sum\_\{b\}\\binom\{n\_\{\\cdot b\}\}\{2\}\}\{\\binom\{N\}\{2\}\}\}\{\\frac\{1\}\{2\}\\\!\\left\[\\sum\_\{a\}\\binom\{n\_\{a\\cdot\}\}\{2\}\+\\sum\_\{b\}\\binom\{n\_\{\\cdot b\}\}\{2\}\\right\]\-\\frac\{\\sum\_\{a\}\\binom\{n\_\{a\\cdot\}\}\{2\}\\sum\_\{b\}\\binom\{n\_\{\\cdot b\}\}\{2\}\}\{\\binom\{N\}\{2\}\}\}\.For cellii, silhouette width wassi=\(bi−ai\)/max⁡\(ai,bi\)s\_\{i\}=\(b\_\{i\}\-a\_\{i\}\)/\\max\(a\_\{i\},b\_\{i\}\), whereaia\_\{i\}is its mean within\-cluster distance andbib\_\{i\}the smallest mean distance to another cluster\. CellBench AvgBio is

AvgBio=NMI\+ARI\+ASW3\.\\operatorname\{AvgBio\}=\\frac\{\\operatorname\{NMI\}\+\\operatorname\{ARI\}\+\\operatorname\{ASW\}\}\{3\}\.

### Frozen\-embedding reconstruction

Each model embedding was passed to the same512512–10241024–512512multilayer perceptron with identical optimization and early\-stopping rules\. We predicted a 400\-gene canonical expression vector\. Training budgets were 100, 500, 700 and 900 cells per class, capped by class availability\. Classwise rare\-safe stratification preserved every evaluable class in the held\-out set\. Five seeds,4242–4646, were used\.

For truthxix\_\{i\}and predictionx^i\\widehat\{x\}\_\{i\}, the cell\-wise Pearson score was

ρcell=1N​∑icorr⁡\(xi,x^i\)\.\\rho\_\{\\rm cell\}=\\frac\{1\}\{N\}\\sum\_\{i\}\\operatorname\{corr\}\(x\_\{i\},\\widehat\{x\}\_\{i\}\)\.Normalized RMSE was

NRMSE=∑i‖x^i−xi‖22∑i‖xi−x¯train‖22\.\\operatorname\{NRMSE\}=\\sqrt\{\\frac\{\\sum\_\{i\}\\\|\\widehat\{x\}\_\{i\}\-x\_\{i\}\\\|\_\{2\}^\{2\}\}\{\\sum\_\{i\}\\\|x\_\{i\}\-\\overline\{x\}\_\{\\rm train\}\\\|\_\{2\}^\{2\}\}\}\.ForS50​\(x\)S\_\{50\}\(x\), the indices of the 50 largest expression values,

Top50=150​N​∑i\|S50​\(xi\)∩S50​\(x^i\)\|\.\\operatorname\{Top50\}=\\frac\{1\}\{50N\}\\sum\_\{i\}\|S\_\{50\}\(x\_\{i\}\)\\cap S\_\{50\}\(\\widehat\{x\}\_\{i\}\)\|\.For Reactome programmmwith gene setGmG\_\{m\}, we formedqi​m=\|Gm\|−1​∑g∈Gmxi​gq\_\{im\}=\|G\_\{m\}\|^\{\-1\}\\sum\_\{g\\in G\_\{m\}\}x\_\{ig\}and its predicted analogueq^i​m\\widehat\{q\}\_\{im\}\. Program recovery was the median over programs ofcorr⁡\(q⋅m,q^⋅m\)\\operatorname\{corr\}\(q\_\{\\cdot m\},\\widehat\{q\}\_\{\\cdot m\}\)\.

Cell\-state geometry was measured bykk\-nearest\-neighbour overlap atk=15k=15:

kNN​@​15=115​N​∑i\|𝒩15​\(xi\)∩𝒩15​\(x^i\)\|\.\\operatorname\{kNN@15\}=\\frac\{1\}\{15N\}\\sum\_\{i\}\|\\mathcal\{N\}\_\{15\}\(x\_\{i\}\)\\cap\\mathcal\{N\}\_\{15\}\(\\widehat\{x\}\_\{i\}\)\|\.Large points in Fig\. 4b–d are means across eight datasets\. Their 95% intervals were obtained from 4,000 nonparametric bootstrap resamples of datasets\. Small symbols are individual datasets\.

For the hPancreas marker\-program matrix, we selected the six most abundant cell types and the five strongest positive marker genes per type from the training partition only\. Each gene was standardized with its training mean and standard deviation\. The displayed entry for cell typeaaand marker programbbis the held\-out mean standardized score of programbbamong cells of typeaa\. No test label or test expression value entered marker selection\.

### Perturbation\-response prediction

The CellBench response predictor received a frozen control\-cell embedding and a perturbation identity\. Top\-kkdenotesk∈\{1,5,9\}k\\in\\\{1,5,9\\\}control examples\. We used the benchmark perturbation splits and five seeds\. Letμp​g\\mu\_\{pg\}andμ0​g\\mu\_\{0g\}be mean expression of geneggunder perturbationppand control, and define the observed and predicted log\-fold changes

Δp​g=μp​g−μ0​g,Δ^p​g=μ^p​g−μ0​g\.\\Delta\_\{pg\}=\\mu\_\{pg\}\-\\mu\_\{0g\},\\qquad\\widehat\{\\Delta\}\_\{pg\}=\\widehat\{\\mu\}\_\{pg\}\-\\mu\_\{0g\}\.The reported error was

MSE​\-​LFC=1\|𝒫\|​G​∑p∈𝒫∑g=1G\(Δ^p​g−Δp​g\)2\.\\operatorname\{MSE\\text\{\-\}LFC\}=\\frac\{1\}\{\|\\mathcal\{P\}\|G\}\\sum\_\{p\\in\\mathcal\{P\}\}\\sum\_\{g=1\}^\{G\}\(\\widehat\{\\Delta\}\_\{pg\}\-\\Delta\_\{pg\}\)^\{2\}\.Pearson\-LFC was the Pearson correlation after concatenating\{Δp​g\}\\\{\\Delta\_\{pg\}\\\}and\{Δ^p​g\}\\\{\\widehat\{\\Delta\}\_\{pg\}\\\}within each evaluation category\. Categories were All, unseen single perturbations and unseen combinations; the combination category is defined only for Norman\. Response\-gene recovery was

Top50response=150​\|𝒫\|​∑p∈𝒫\|S50​\(\|Δp\|\)∩S50​\(\|Δ^p\|\)\|\.\\operatorname\{Top50\}\_\{\\rm response\}=\\frac\{1\}\{50\|\\mathcal\{P\}\|\}\\sum\_\{p\\in\\mathcal\{P\}\}\|S\_\{50\}\(\|\\Delta\_\{p\}\|\)\\cap S\_\{50\}\(\|\\widehat\{\\Delta\}\_\{p\}\|\)\|\.Fig\. 5b,d report mean±\\pmstandard deviation over five seeds\. Fig\. 5e shows one pre\-specified Adamson Top\-5 condition\. Genes were selected by the magnitude of the observed response, and the colour scale displays signed response divided by the corresponding observed response magnitude, clipped symmetrically to\[−1\.5,1\.5\]\[\-1\.5,1\.5\]\.

### Biological\-programme and genetic\-interaction analysis

All analyses in Fig\. 6 operate on frozen embeddings and token states\. Cell\-type and perturbation labels are used only after embedding extraction to organize the analysis; they do not update the encoder\.

#### Max\-token cell\-type programmes\.

Letui​gu\_\{ig\}denote the contribution of geneggto the pooled representation of celliiunder the fixed extraction rule\. For cell typeccwith held\-out cell setℐc\\mathcal\{I\}\_\{c\}, the max\-token winner fraction is

Ac​g=1\|ℐc\|​∑i∈ℐc𝟏​\[g=arg⁡maxj∈𝒪​\(xi\)⁡ui​j\]\.A\_\{cg\}=\\frac\{1\}\{\|\\mathcal\{I\}\_\{c\}\|\}\\sum\_\{i\\in\\mathcal\{I\}\_\{c\}\}\\mathbf\{1\}\\\!\\left\[g=\\arg\\max\_\{j\\in\\mathcal\{O\}\(x\_\{i\}\)\}u\_\{ij\}\\right\]\.Genes were ordered by their largest winner fraction and the displayed heat map retained canonical endocrine, exocrine, stromal and immune examples\. This analysis recovered, among others,GCG,INS,SST,PPY,GHRLandREG1Ain their expected hPancreas cell populations\.

#### Graph\-programme ablation\.

For a graph blockBBselected to represent an independently annotated cell programme, its cell\-type score is

Ac​B=1\|ℐc\|​∑i∈ℐc1\|B∩𝒪​\(xi\)\|​∑g∈B∩𝒪​\(xi\)ui​g\.A\_\{cB\}=\\frac\{1\}\{\|\\mathcal\{I\}\_\{c\}\|\}\\sum\_\{i\\in\\mathcal\{I\}\_\{c\}\}\\frac\{1\}\{\|B\\cap\\mathcal\{O\}\(x\_\{i\}\)\|\}\\sum\_\{g\\in B\\cap\\mathcal\{O\}\(x\_\{i\}\)\}u\_\{ig\}\.Scores are computed independently in every seed before aggregation\.

For cellii, letziz\_\{i\}be its unmodified embedding andzi\(−B\)z\_\{i\}^\{\(\-B\)\}the embedding after genes inBBare removed\. The mean displacement of cell typeccis

Dc​B=1\|ℐc\|​∑i∈ℐc‖zi−zi\(−B\)‖2\.D\_\{cB\}=\\frac\{1\}\{\|\\mathcal\{I\}\_\{c\}\|\}\\sum\_\{i\\in\\mathcal\{I\}\_\{c\}\}\\left\\\|z\_\{i\}\-z\_\{i\}^\{\(\-B\)\}\\right\\\|\_\{2\}\.Each of the eight programmes was paired with one random gene set matched for set size, mean expression and detection frequency\. IfD~c​B\\widetilde\{D\}\_\{cB\}is the mean displacement under this matched null, the excess programme effect is

Ec​B=Dc​B−D~c​B\.E\_\{cB\}=D\_\{cB\}\-\\widetilde\{D\}\_\{cB\}\.For each programme, the annotated cell identity corresponding to the programme was designated the matched identity before model inspection; all remaining identities formed the comparison group\. We sampled 24 cells from each of eight identities, giving 192 unique cells\. The plotted matched value for programmeBBis the mean excess displacement over its 24 corresponding cells; the comparison value is the mean over the 168 cells from the other seven identities\. The eight programme\-level matched and comparison means were the statistical pairs\. Their grand means were 5\.34 and 0\.064, respectively \(two\-sided paired Wilcoxon signed\-rank test,n=8n=8programmes, statisticW=0W=0,P=0\.0078125P=0\.0078125\)\. This programme\-level test replaces a cell\-level test that would treat cells from one dataset as independent biological replicates\. It is the only inferential comparison in Fig\. 6b, so no multiplicity correction was applied\.

#### Directed programme dependencies\.

Letεb​\(x\)\\varepsilon\_\{b\}\(x\)be the teacher\-target prediction error for target programmebbin cellxx\. After ablating context programmeaa, the directed effect was

Ia→b=𝔼x​\[εb​\(x\(−a\)\)−εb​\(x\)\]−𝔼x​\[εb​\(x\(−a~\)\)−εb​\(x\)\],I\_\{a\\rightarrow b\}=\\mathbb\{E\}\_\{x\}\\\!\\left\[\\varepsilon\_\{b\}\(x^\{\(\-a\)\}\)\-\\varepsilon\_\{b\}\(x\)\\right\]\-\\mathbb\{E\}\_\{x\}\\\!\\left\[\\varepsilon\_\{b\}\(x^\{\(\-\\widetilde\{a\}\)\}\)\-\\varepsilon\_\{b\}\(x\)\\right\],wherea~\\widetilde\{a\}is an expression\- and detection\-matched random set\. Positive values indicate that programmeaacontains context associated with prediction of programmebbbeyond the matched\-random expectation\. The matrix is directional because ablatingaawhile predictingbbis not equivalent to ablatingbbwhile predictingaa\. The matrix contains 56 ordered, non\-self programme pairs\. Each pair was evaluated on 24 cells from the target\-matched identity, and the null correction is the mean over three independently sampled matched random sets\. The matrix is presented descriptively; cell\-level unadjusted tests retained in the source table are not used as evidence for individual directed edges\.

#### Compositional perturbation geometry\.

For perturbationpp, letℐp\\mathcal\{I\}\_\{p\}be its cells andℐ0\\mathcal\{I\}\_\{0\}the controls\. We defined the mean latent displacement

δp=1\|ℐp\|​∑i∈ℐpzi−1\|ℐ0\|​∑i∈ℐ0zi\.\\delta\_\{p\}=\\frac\{1\}\{\|\\mathcal\{I\}\_\{p\}\|\}\\sum\_\{i\\in\\mathcal\{I\}\_\{p\}\}z\_\{i\}\-\\frac\{1\}\{\|\\mathcal\{I\}\_\{0\}\|\}\\sum\_\{i\\in\\mathcal\{I\}\_\{0\}\}z\_\{i\}\.For a measured double perturbation\(a,b\)\(a,b\), additive consistency was

Ca​b=cos⁡\(δa​b,δa\+δb\)\.C\_\{ab\}=\\cos\\\!\\left\(\\delta\_\{ab\},\\delta\_\{a\}\+\\delta\_\{b\}\\right\)\.The null distribution replaced\(a,b\)\(a,b\)with 200 expression\-matched random single\-perturbation pairs for each combination condition\. Each displacement used 24 cells from the combination or single condition and the same 128 control cells\. The statistical unit was one of the 80 evaluable combination conditions, and each observed cosine was paired with that condition’s random\-pair mean\. All 80 differences were positive \(two\-sided paired Wilcoxon signed\-rank test,n=80n=80conditions,P=7\.8495×10−15P=7\.8495\\times 10^\{\-15\}\)\. This was one pre\-specified aggregate comparison, so no multiplicity correction was applied\. Literature\-supported pairs highlighted in Fig\. 6 include MAP2K3–MAP2K6, CEBPA–CEBPE and TBX2–TBX3\[[15](https://arxiv.org/html/2608.05928#bib.bib21),[3](https://arxiv.org/html/2608.05928#bib.bib22),[17](https://arxiv.org/html/2608.05928#bib.bib23)\]; these examples are descriptive and were not assigned separatePPvalues\.

#### Null distributions and multiplicity\.

Expression\- and detection\-matched nulls were constructed by replacing each candidate gene with genes from the same bins of mean expression and detection frequency\. Figure 6b uses one matched set for each of eight programmes; Fig\. 6c uses three sets for each source programme and reports no edge\-level significance; and Fig\. 6d uses 200 random single\-perturbation pairs for each of 80 combination conditions\. The two reportedPPvalues correspond to one aggregate hypothesis in Fig\. 6b and one aggregate hypothesis in Fig\. 6d\. No family of individual programme, edge or gene\-pair claims was declared significant, and therefore no multiplicity correction was applied to these two aggregate tests\.

Algorithm S2Construction of the biological\-knowledge analyses in Fig\. 61:Frozen encoder

ff; hPancreas cells and labels; Norman controls, singles and doubles; graph modules

ℳ\\mathcal\{M\}
2:Extract token states and cell embeddings once with

ff
3:foreach hPancreas cell type

ccdo

4:Count the max\-token winner in each held\-out cell

5:Compute winner fractions

Ac​gA\_\{cg\}
6:endfor

7:foreach programme

B∈ℳB\\in\\mathcal\{M\}and cell type

ccdo

8:Re\-embed cells after removing

BB
9:Re\-embed cells after removing matched random gene sets

10:Compute excess displacement

Ec​BE\_\{cB\}
11:endfor

12:foreach ordered programme pair

\(a,b\)\(a,b\)do

13:Ablate context programme

aaand measure target error for

bb
14:Subtract the expression\-matched random\-set effect to obtain

Ia→bI\_\{a\\rightarrow b\}
15:endfor

16:foreach Norman perturbation

ppdo

17:Compute mean displacement

δp\\delta\_\{p\}from controls

18:endfor

19:foreach measured double perturbation

\(a,b\)\(a,b\)do

20:Compute

Ca​b=cos⁡\(δa​b,δa\+δb\)C\_\{ab\}=\\cos\(\\delta\_\{ab\},\\delta\_\{a\}\+\\delta\_\{b\}\)
21:Compare with expression\-matched random perturbation pairs

22:endfor

23:Apply the two pre\-specified aggregate paired tests; retain the directed matrix as descriptive

24:Export one source\-data table per panel and render Fig\. 6

### Figure source data

Figures 1–6 were generated from panel\-level source\-data tables\. The conceptual method schematic contains no numerical result and summarizes graph blocks, context removal, the EMA teacher, the stop\-gradient block target and aggregate block loss\. Figure 6 combines the max\-token, programme\-ablation, directed\-dependency and perturbation\-composition analyses described above\.

Each quantitative panel was assembled from the stated transformation of its source table\. Model colours, marker shapes, dataset order and axis limits were specified consistently across panels\. Error bars denote either seed standard deviations or dataset\-bootstrap intervals as identified in the captions\. The source\-data archive contains the numerical input for every panel\.

## Supplementary benchmark tables

The tables below report the complete CellBench\-style outputs used in this study\. Annotation is a within\-dataset few\-shot benchmark\. BioM\-JEPA rows use the frozen\-embedding and probe protocols defined above\. Tables are numbered consecutively as Table S1–S13\.

Table S7:Cross\-dataset summary of CellBench within\-dataset Top\-5 few\-shot annotation\. Macro\-F1F\_\{1\}values are averaged over the nine datasets in Tables[S8](https://arxiv.org/html/2608.05928#Ax3.T8)–[S10](https://arxiv.org/html/2608.05928#Ax3.T10); intervals are percentile 95% confidence intervals from 20,000 nonparametric bootstrap resamples of datasets\. Mean rank is calculated within each dataset among the ten displayed methods\. Datasets, rather than cells, are the uncertainty unit\.Table S8:Full CellBench within\-dataset few\-shot annotation benchmark with BioM\-JEPA\. The reported metrics are Acc/F1/Recall \(%\)\.Table S9:Full CellBench within\-dataset few\-shot annotation benchmark with BioM\-JEPA \(continued\)\. The reported metrics are Acc/F1/Recall \(%\)\.Table S10:Full CellBench within\-dataset few\-shot annotation benchmark with BioM\-JEPA \(continued\)\. The reported metrics are Acc/F1/Recall \(%\)\.Table S11:CellBench clustering performance with BioM\-JEPA \(AvgBio and accuracy, %\)\. AvgBio is the mean of NMI, ARI and ASW\.
Table S12:CellBench perturbation\-prediction error metrics with BioM\-JEPA\. Values are mean±\\pms\.d\. over five seeds; BioM\-JEPA ranks are computed among the nine displayed models\.Table S13:CellBench perturbation\-prediction directional and response\-gene metrics with BioM\-JEPA\. Values are mean±\\pms\.d\. over five seeds; BioM\-JEPA ranks are computed among the nine displayed models\.

Similar Articles

The Annotated JEPA

Hacker News Top

A step-by-step annotated implementation and explanation of Joint Embedding Predictive Architectures (JEPA) for self-supervised learning, covering I-JEPA, V-JEPA, and LeJEPA.