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

arXiv: https://arxiv.org/abs/2608.06659 · HTML: https://mlsanity.com/paper/2608.06659/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 17 references that are in the map.
- First posted on arXiv 7 Aug 2026.

## Builds on (most cited first, 10 of 10)

- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language](https://mlsanity.com/paper/2202.03555/) (2022)
- [LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics](https://mlsanity.com/paper/2511.08544/) (2025)
- [VL-JEPA: Joint Embedding Predictive Architecture for Vision-language](https://mlsanity.com/paper/2512.10942/) (2025)
- [Point-JEPA: A Joint Embedding Predictive Architecture for Self-Supervised Learning on Point Cloud](https://mlsanity.com/paper/2404.16432/) (2024)
- [When Does LeJEPA Learn a World Model?](https://mlsanity.com/paper/2605.26379/) (2026)
- [SToFM: a Multi-scale Foundation Model for Spatial Transcriptomics](https://mlsanity.com/paper/2507.11588/) (2025)
- [HEIST: A Graph Foundation Model for Spatial Transcriptomics and Proteomics Data](https://mlsanity.com/paper/2506.11152/) (2025)

## Abstract (arXiv)

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.