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

2026-08-06

AI Summary

Existing foundation models for spatial transcriptomics directly reconstruct gene expression, rendering them susceptible to technical noise and limiting their representational transferability. This work proposes CellWorld, which reframes the pretraining objective from observed expression to latent cellular representations by leveraging spatial context and partial expression cues to predict masked cells’ embeddings in latent space. Built upon a Transformer architecture, we train four model scales on a large-scale dataset comprising 46 million human cells. Experiments demonstrate that CellWorld-Large, pretrained on only 5% of the data, surpasses all fine-tuned baselines across seven tasks without any adaptation. Moreover, CellWorld-Small achieves state-of-the-art performance both with linear probing on eleven tasks and fine-tuning on seven tasks, substantially enhancing cross-task and cross-dataset transferability.

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.

Problem

• spatial transcriptomics
• foundation models
• technical variation
• representation transferability
• masked gene reconstruction

Innovation

• latent cell prediction
• spatial transcriptomics
• foundation models
• predictive pretraining
• representation transferability