Pith. sign in

REVIEW 3 major objections 4 minor 12 references

CellWorld: From Gene-Level Reconstruction to Latent Cell Prediction in Spatial Transcriptomics Foundation Models

T0 review · 3 major / 4 minor · reviewed 2026-08-10 · deepseek-v4-flash

Pith's one-line read Latent-space prediction of masked cells, not gene reconstruction, makes spatial transcriptomics foundation models that beat larger baselines at a fraction of the parameters.

desk verdict The paper makes a real contribution—latent cell prediction for spatial transcriptomics with controlled scaling experiments—but the SOTA claims need error bars and the closest reconstruction-based baselines are missing. read the letter →

arxiv 2608.06659 v1 pith:CLA5IGPE submitted 2026-08-07 cs.AI

classification cs.AI
keywords spatialtranscriptomicsfoundationmodellatentpredictionjoint-embeddingpredictivearchitectureself-supervisedlearningcellannotationrepresentationcontext
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that the best way to pretrain a foundation model for spatial transcriptomics is not to reconstruct masked gene measurements but to predict latent representations of masked cells from their visible spatial neighbors, aided by a small partial-expression hint. The authors build CellWorld, a joint-embedding predictive architecture with a context encoder, an EMA-updated target encoder, and a hint-conditioned spatial predictor, and pretrain it on 46 million human cells. They report that even the smallest CellWorld, with 5.74M trainable parameters, outperforms existing spatial transcriptomics foundation models on all 11 linear-probe benchmarks and all seven fine-tuned spatial benchmarks, and that a frozen Large model pretrained on only 5% of the corpus beats every fully fine-tuned baseline on all seven spatial benchmarks. The paper's broader claim is that latent prediction avoids reproducing assay-specific technical variation and that, under a fixed cell budget, broad biological source diversity and enough optimization matter more than raw cell count for spatial transfer.

What carries the argument

The load-bearing mechanism is the hint-conditioned spatial predictor with a 10% partial-expression hint. Each masked cell's query is built from a random 10% subset of its nonzero gene expression (plus organ and platform embeddings), projected into predictor space, and combined with a learnable mask token; the predictor then recovers the EMA target encoder's latent for that cell from the visible context. The hint is used only by the predictor, not by either encoder, and it resolves what the paper calls cell-level target ambiguity -- the fact that a cell's identity and state are not uniquely determined by its spatial neighborhood. The objective is plain mean-squared error in the latent space, and the paper shows that without the hint the Large model's predictor degenerates and training collapses around 3,000 steps, while with the hint the model stays stable and outperforms its no-hint counterpart.

What would settle it

Train CellWorld-Large with hint ratio 0.10 but with the spatial coordinates of all cells randomly permuted during pretraining; if downstream spatial benchmarks (region and niche) remain at the same level as the unpermuted model, the claim that spatial context drives the gains would be falsified. Alternatively, a matched model pretrained with the hint but without any visible context (using only the hint to predict the target latent) that matches CellWorld's spatial benchmark scores would show the spatial context is unnecessary.

Watch

Extended reading notes

Core claim

CellWorld's central discovery is that shifting the pretraining objective from gene-level reconstruction to latent cell prediction turns spatial transcriptomics into a tractable foundation-model setting. Given a local patch of cells, CellWorld masks 60% of cells, encodes the remaining visible cells with a spatial Transformer using 2D-ALiBi relative attention biases, and trains a spatial predictor to recover, from the visible context and a 10% random partial-expression hint, the latent representations that an EMA target encoder produces for the masked cells when it sees the complete patch. The target encoder and predictor are discarded after pretraining; only the cell tokenizer and context encoder transfer to downstream tasks such as cell annotation, region prediction, and niche-composition prediction. The paper shows across four held-out datasets that this objective yields state-of-the-art linear-probe and fine-tuning results at every model scale, and that a frozen encoder pretrained on a broadly sampled 5% of the corpus outperforms all fully fine-tuned baselines on spatial tasks. The authors interpret these results as evidence that predicting abstract cell states, not raw expression values, avoids encouraging the model to reproduce assay-specific technical variation and thus produces more transferable spatial representations.

Load-bearing premise

The 10% partial-expression hint has to be strong enough to prevent predictor collapse at scale but weak enough that spatial context still determines the prediction; if this balance is fragile, the method's success rests on a hyperparameter rather than on a robust principle.

Editorial extensions

If this is right

  • At every tested scale, CellWorld outperforms existing ST foundation models on all 11 linear-probe benchmarks and all seven fine-tuned spatial benchmarks, with performance generally improving as model size grows from 5.74M to 94.56M trainable parameters.
  • A frozen CellWorld-Large trained on 5% of the corpus (64 slides, 11 organs) beats every fully fine-tuned baseline across all seven spatial benchmarks, indicating that enough optimization and broad source diversity can substitute for more cells.
  • Under the fixed 10,000-step schedule, the 25% corpus subset matches the full corpus, while proportionally scaling steps with data size preserves cell annotation but degrades spatial transfer, implying spatial transfer needs repeated exposure to the same local contexts.
  • Concentrating the same cell budget in fewer slides and organs lowers spatial-task performance, so biological source diversity matters for spatial transfer more than raw cell count.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • A natural next test is whether the same latent-prediction objective transfers to temporal or perturbational spatial data, where the target is not a masked present cell but a future cell state; CellWorld's authors hint at this direction in their conclusion, but the experiments to check it have not been run.
  • Because the hint controls the trade-off between target specificity and reliance on spatial context, one could test whether the optimal hint ratio changes with tissue type or platform, or whether an adaptive hint that shrinks as training progresses yields better spatial representations.
  • The finding that a 5% broadly sampled subset matches the full corpus suggests that pretraining data might be selected by diversity rather than volume; a direct comparison would be to hold out entire organs and measure how quickly spatial-transfer performance drops as source diversity is removed.
  • If the predictor's degeneration at scale is the main obstacle, then techniques that stabilize joint-embedding training without hints, such as centering or variance regularization, could remove the need for the hint entirely and may further improve spatial transfer.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper proposes CellWorld, a spatial transcriptomics foundation model pretrained by predicting latent representations of masked cells from visible spatial context plus a 10% partial-expression hint. The architecture uses a context encoder, an EMA-updated target encoder, and a hint-conditioned spatial predictor with 2D-ALiBi attention. Pretraining is performed on 46 million human cells spanning three platforms and 11 organs, at four model scales from 5.74M to 94.56M trainable parameters. Controlled experiments compare masking strategies, hint ratios, prediction objectives, and spatial permutation; scaling experiments vary model size and data fraction/coverage; comparisons are made against PCA, CellPLM, Nicheformer, and scGPT-spatial on four held-out datasets with FOV-level splits. The paper claims state-of-the-art results on all 11 linear-probe benchmarks and all seven fine-tuned spatial benchmarks, with a frozen 5% corpus pretraining variant outperforming fully fine-tuned baselines on the spatial benchmarks.

Significance. If the results hold, CellWorld provides a credible demonstration that latent-space predictive pretraining can scale for spatial transcriptomics, and the controlled scaling analysis is a useful contribution to how data diversity and optimization interact. The paper has real strengths: FOV-level splits reduce spatial leakage, three-seed averages are used, multiple ablations isolate design choices, and the scaling experiments explicitly probe confounds such as repeated exposure and source diversity. Code is promised. The main limitations are the unresolved role of the partial-expression hint, the absence of the closest reconstruction-based baselines from the comparison, and the lack of error bars on headline claims; these temper the strength of the state-of-the-art conclusion.

major comments (3)
  1. [Method, Eqs. (4)-(7); Controlled Design Analysis, Table 3] The hint in Eq. (5) is a random 10% sample of the masked cell's own expressed genes, and the target in Eq. (4) is the EMA encoder's output on the full patch containing that cell; therefore the MSE in Eq. (7) can in principle be minimized by a near-direct hint-to-latent mapping, with spatial context playing only a minor role. Table 3 shows that removing the hint improves both spatial tasks at Base scale (Region 0.571 vs 0.539, Niche 0.551 vs 0.531), and Figure 4 shows that no-hint Large training collapses after about 3,000 steps, so the 0.10 hint is load-bearing. The spatial-permutation control in Table 4 demonstrates some use of the expression-coordinate correspondence, but it does not quantify how much the predictor relies on visible context when the hint is present. Please add a context-free control (hint only, no visible context) and/or a direct measure of context contribution, such as attention or gradient analysis; without this, the conceptual claim that CellWorld learns from visible spatial context rather than performing hint-to-latent autoencoding is not established.
  2. [Related Work; Comparison with Existing Methods, Table 6] The Related Work identifies SToFM and HEIST as closely related existing methods, and the Introduction states that even CellWorld-Small outperforms every existing method on all benchmarks, yet neither SToFM nor HEIST (nor BrainBeacon) appears in Table 6 or Figure 1. The actual comparison is limited to PCA, CellPLM, Nicheformer, and scGPT-spatial. Because the central claim is state-of-the-art status among ST foundation models, the omission of the closest reconstruction-based baselines is load-bearing. Please include these methods in the comparison, or precisely qualify the claim to the evaluated baselines.
  3. [Experimental Setup; Tables 3, 5, 6] All reported scores are point estimates averaged over three downstream seeds, with no standard deviations, confidence intervals, or significance tests shown in any table or figure. The claims that CellWorld outperforms every baseline on all benchmarks require evidence that the margins are not within seed variability; for example, some linear-probe margins are small and the fine-tuning comparison on Cell Brain shows NicheFormer at 0.958 versus CellWorld-Huge at 0.953. Please report per-seed values or error bars for at least the headline comparisons in Tables 3, 5, and 6, and indicate which differences are significant.
minor comments (4)
  1. [Experiments, benchmark definition] The phrase 'all seven fine-tuned spatial benchmarks' is used repeatedly but never explicitly defined; please state that these are the Region and Niche tasks across the four held-out datasets, excluding the Cell annotation task, or list them explicitly.
  2. [Table 5] The 'Steps (k)' column uses values '10', '2.5', and '0.5'; writing '10k', '2.5k', and '0.5k' would avoid ambiguity with the percentage column in the same table.
  3. [Figure 1 caption] The caption says 'Comparison with existing methods' while the Introduction claims 'every existing method'; please align the wording with the actual set of evaluated baselines, or add the missing methods.
  4. [Related Work] HEIST is described as incorporating spatial and contrastive objectives, but no downstream comparison or discussion of its performance on the chosen benchmarks is provided; a brief note on why it was not included would improve clarity.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the transfer claims rest on held-out benchmarks, and the EMA/hint pretraining loop is standard self-supervised prediction rather than a fitted-label derivation.

full rationale

CellWorld's central claims are empirical: that latent-space predictive pretraining transfers to held-out cell annotation, region, and niche tasks, and that scaling and data-diversity trends behave as described. These claims are tested on four held-out datasets excluded from pretraining, under linear probing and fine-tuning, against PCA, Nicheformer, scGPT-spatial, and CellPLM. The pretraining target (Eq. 4) is an EMA encoder output on the complete patch, and Eq. 7 minimizes MSE against that latent; this is the standard BYOL/I-JEPA-style self-supervised loop, not a fitted parameter renamed as a prediction. The 10% hint (Eq. 5) is an input feature subsampled from the target cell, not a downstream label; the paper openly reports in Table 3 and Figure 4 that the no-hint Large model collapses and that the hint restores target-specific variation, which is a robustness limitation of the design, not a circular derivation. The only self-citation is SToFM (Zhao et al. 2025, overlapping author H. Zhou), used in Related Work to distinguish reconstruction-based embedding models from CellWorld; it is not load-bearing for any result or benchmark. No equation in the paper reduces by construction to a previously fitted value, and no downstream number is generated from the training objective alone.

Assumptions & free parameters 6 free parameters · 5 assumptions · 0 invented entities

The ledger captures the tuned hyperparameters, modeling assumptions, and absence of new physical entities. The free parameters are pretraining design choices selected by ablation or recipe rather than derived quantities; the axioms are the empirical assumptions about stability, leakage, and transfer without which the benchmark results would not support the central claim.

free parameters (6)
  • Mask ratio r = 0.6
    Fraction of cells withheld as prediction targets; chosen by ablation across {0.4, 0.6, 0.8}. Performance varies marginally, so the exact value is not load-bearing but is a tuned hyperparameter.
  • Hint inclusion probability kappa = 0.10
    Probability of retaining each nonzero gene in the partial-expression hint (Eq. 5); selected from {0.0, 0.10, 0.50} based on downstream performance and training stability at Large scale.
  • Top genes per cell K = 512
    Number of highest-expression genes retained per cell; set to cover the full assayed panel for 95% of cells in the corpus.
  • EMA momentum range = 0.9990 to 0.9999
    Momentum schedule for the target encoder; chosen as part of the pretraining recipe, not derived from data.
  • Peak learning rate and schedule = 3e-3, 1k warmup, cosine decay
    Pretraining optimization hyperparameters; tuned for stability across model scales.
  • Pretraining steps and batch size = 10,000 steps, batch 256
    Recipe fixed across scales; in data-scaling experiments steps are scaled proportionally to data size, which is itself a design choice.
assumptions (5)
  • domain assumption EMA target encoder outputs are stable, non-collapsing learning targets.
    The entire latent prediction objective (Eqs. 4 and 7) depends on the EMA target representation being a meaningful supervisory signal; Figure 4 shows this can fail without a hint.
  • ad hoc to paper A 10% expression hint disambiguates masked cell targets without leaking identity.
    Eq. 5 and the choice of kappa=0.10 are central to stability and transfer; the paper provides empirical evidence but no principled derivation that 10% is the right balance.
  • domain assumption FOV-level splits prevent spatial leakage and all-cells HVG selection does not bias transfer results.
    The protocol selects 300 highly variable genes using all cells in each downstream dataset and only then splits by FOV, a mild test-set information leak that is disclosed but assumed not to affect relative comparisons.
  • domain assumption The pretraining corpus, 46M cells across 3 platforms and 11 organs, is representative enough for transfer to the four held-out datasets.
    Generalization claims rest on corpus diversity; the paper's own data-scaling results show source diversity matters.
  • domain assumption Random masking with r=0.6 is a reasonable proxy for the distribution of spatial contexts seen at test time.
    Masking strategy is an inductive bias; the controlled analysis shows random masking is slightly better than block masking, but no theoretical guarantee is provided.

how reviews work

0 comments
Cite this review

Pith. "Pith review of CellWorld: From Gene-Level Reconstruction to Latent Cell Prediction in Spatial Transcriptomics Foundation Models." pith.science (2026). https://pith.science/paper/CLA5IGPE

@misc{pith2026260806659,
  author       = {Pith},
  title        = {Pith review of: CellWorld: From Gene-Level Reconstruction to Latent Cell Prediction in Spatial Transcriptomics Foundation Models},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/CLA5IGPE}},
  note         = {Machine review of arXiv:2608.06659}
}
read the original abstract

This paper shows that latent-space predictive pretraining can provide a scalable route to foundation models for spatial transcriptomics. Existing spatial transcriptomics foundation models primarily reconstruct masked gene identities or expression values, potentially encouraging the reproduction of assay-specific technical variation and limiting representation transferability. To avoid directly reconstructing such variation, we shift the prediction target from observed gene measurements to latent cell representations and introduce CellWorld, which predicts the latent representations of masked cells from visible spatial context and a limited partial-expression hint. We pretrain four CellWorld variants, spanning 5.74M to 94.56M trainable parameters, on a corpus of 46 million human cells. Our controlled scaling experiments show that performance improves with model capacity, particularly on spatial tasks, while spatial transfer depends more on sufficient optimization and broad biological source diversity than on cell count alone. Across four held-out datasets, even CellWorld-Small, with 5.74M trainable parameters, outperforms every baseline on all 11 linear-probe benchmarks and all seven fine-tuned spatial benchmarks. Most notably, a frozen CellWorld-Large pretrained on only 5\% of the corpus with broad biological source coverage outperforms every fully fine-tuned baseline across all seven spatial benchmarks. Code is available at https://github.com/UoM-HealthAI/CellWorld.

Figures

Figures reproduced from arXiv: 2608.06659 by the authors.

Figure 1
Figure 1. Comparison with existing methods across 11 held-out task–dataset pairs under (a) linear probing and (b) fine-tuning. [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Overview of the CellWorld pretraining architecture. [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 3
Figure 3. Illustration of (a) random masking and (b) spatial [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figures from the paper (2 more)
Figure 4
Figure 4. Figure 4: Training and validation loss trajectories for [PITH_FULL_IMAGE:figures/full_fig_p005_4.png]
Figure 5
Figure 5. Figure 5: CellWorld model scaling under linear probing and [PITH_FULL_IMAGE:figures/full_fig_p006_5.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

12 extracted references · 5 canonical work pages

  1. [2]

    arXiv preprint arXiv:2511.08544

    Lejepa: Provable and scalable self-supervised learning without the heuristics. arXiv preprint arXiv:2511.08544. Birk,S.;Vahidi,A.;Sanian,M.V.;Merchant,A.;andLotfol- lahi,M.2026. ST-JEPA:Joint-EmbeddingPredictiveArchi- tecture for Spatial Transcriptomics. InThe 2026 Workshop on Generative and Agentic AI for Biology. Chen, D.; Shukor, M.; Moutakanni, T.; Ch...

  2. [4]

    Litman, E.; Myers, T.; Agarwal, V.; Mittal, E.; Li, O.; Gopinath, A.; and Kassis, T

    When Does LeJEPA Learn a World Model?arXiv preprint arXiv:2605.26379. Litman, E.; Myers, T.; Agarwal, V.; Mittal, E.; Li, O.; Gopinath, A.; and Kassis, T

  3. [5]

    Madhu, H.; Rocha, J

    GeneJepa: A Predictive World Model of the Transcriptome.bioRxiv, 2025–10. Madhu, H.; Rocha, J. F.; Huang, T.; Viswanath, S.; Krish- naswamy, S.; and Ying, R

  4. [6]

    arXiv preprint arXiv:2506.11152

    HEIST: A Graph Founda- tionModelforSpatialTranscriptomicsandProteomicsData. arXiv preprint arXiv:2506.11152. Mur-Labadia, L.; Muckley, M.; Bar, A.; Assran, M.; Sinha, K.; Rabbat, M.; LeCun, Y.; Ballas, N.; and Bardes, A

  5. [7]

    Press, O.; Smith, N

    V-jepa2.1:Unlockingdensefeaturesinvideoself-supervised learning.arXiv preprint arXiv:2603.14482. Press, O.; Smith, N. A.; and Lewis, M

  6. [10]

    biorxiv, 2025–02

    scGPT-spatial: Continual pretraining of single-cell foundation model for spatial transcriptomics. biorxiv, 2025–02. Wen, H.; Tang, W.; Dai, X.; Ding, J.; Jin, W.; Xie, Y.; and Tang,J.2024. CellPLM:Pre-trainingofcelllanguagemodel beyondsinglecells.InInternationalConferenceonLearning Representations, volume 2024, 5649–5673. Zhang, C.; Yang, Y.; Jiao, Y.; Ya...

  7. [11]

    Zhao, S.; Luo, Y.; Yang, G.; Zhong, Y.; Zhou, H.; and Nie, Z

    BrainBeacon: ACross-SpeciesFoundationModelforSingle-cellResolved Brain Spatial Transcriptomics.bioRxiv, 2025–07. Zhao, S.; Luo, Y.; Yang, G.; Zhong, Y.; Zhou, H.; and Nie, Z

  8. [12]

    Stofm: a multi-scale foundation model for spatial transcriptomics.arXiv preprint arXiv:2507.11588

Show all 12 references
  1. [2016]

    Science, 353(6294): 78–82

    Visualization and analysis of geneexpressionintissuesectionsbyspatialtranscriptomics. Science, 353(6294): 78–82. Tejada-Lapuerta, A.; Schaar, A. C.; Gutgesell, R.; Palla, G.; Halle, L.; Minaeva, M.; Vornholz, L.; Dony, L.; Drummer, F.;Richter,T.;etal.2025. Nicheformer:afoundat...

  2. [2021]

    Saito, A.; Kudeshia, P.; and Poovvancheri, J

    Train short, test long: Attention with linear biases enables input length extrapolation.arXiv preprint arXiv:2108.12409. Saito, A.; Kudeshia, P.; and Poovvancheri, J

  3. [2025]

    Assran,M.;Duval,Q.;Misra,I.;Bojanowski,P.;Vincent,P.; Rabbat,M.;LeCun,Y.;andBallas,N.2023

    V-jepa 2: Self-supervised video models en- able understanding, prediction and planning.arXiv preprint arXiv:2506.09985. Assran,M.;Duval,Q.;Misra,I.;Bojanowski,P.;Vincent,P.; Rabbat,M.;LeCun,Y.;andBallas,N.2023. Self-supervised learning from images with a joint-embedding predic...

  4. [2026]

    Grill, J.-B.; Strub, F.; Altché, F.; Tallec, C.; Richemond, P.; Buchatskaya, E.; Doersch, C.; Avila Pires, B.; Guo, Z.; GheshlaghiAzar,M.;etal.2020

    Cell-JEPA: Latent Representation Learning for Single-Cell Transcriptomics.arXiv preprint arXiv:2602.02093. Grill, J.-B.; Strub, F.; Altché, F.; Tallec, C.; Richemond, P.; Buchatskaya, E.; Doersch, C.; Avila Pires, B.; Guo, Z.; GheshlaghiAzar,M.;etal.2020. Bootstrapyourownlaten...

Pith tools

Reviewed August 10, 2026 · model on record in the stance chip above.