Back Home

AI for Science/生物基礎模型

CellWorld Switches to Pretraining on Latent Cell Representations; 5.74M-Parameter Model Outperforms Listed Baselines Across 18 Spatial Transcriptomics Tasks

Rather than directly reconstructing masked gene values, CellWorld predicts representations from a target encoder using neighboring cells and limited expression hints. The code and training configurations have been released under the MIT License, but the pretrained weights remain in an approval-gated private repository.

Laboratoires Servier · CC BY-SA 3.0 · Image source
zh-Hant

Spatial transcriptomics data record both gene expression and tissue coordinates for cells, but different measurement platforms introduce dropout, sequencing-depth variation, and other technical noise. Most existing foundation models require networks to reconstruct masked gene identities or expression values, potentially learning platform-specific errors along with the underlying biology. CellWorld, released on August 7, instead adopts an objective resembling a joint-embedding predictive architecture: rather than recovering the original measurements, it predicts the latent representation of a masked cell produced by a target encoder.

During training, each cell is treated as a token. The model masks a spatial region, a context encoder processes the visible cells, and a predictor combines that context with limited partial-expression hints from the masked cells to approximate the output of a target encoder updated through an exponential moving average. The public configuration uses 2D ALiBi to represent spatial relationships, a 60% random masking ratio, 10% partial-expression hints, and MSE loss. The four model variants range from 5.74M to 94.56M trainable parameters, and the paper states that the pretraining corpus contains 46 million human cells.

The authors evaluated the models on four held-out datasets and report that even the smallest variant, CellWorld-Small, outperformed every listed baseline across 11 frozen-encoder linear-probing tasks and seven full-fine-tuning tasks. More notably, a frozen CellWorld-Large model trained on only 5% of the pretraining corpus—but retaining broader coverage of biological sources—still outperformed all fully fine-tuned baselines across seven spatial tasks. This supports the conclusion that source diversity and sufficient optimization may matter more than simply increasing the number of cells, although the result comes from an author-controlled experiment and cannot be directly generalized to other platforms or disease populations.

The GitHub repository includes scripts for pretraining, linear probing, fine-tuning, `.h5ad` inference, and multi-GPU launching, while users must provide their own data locations. The largest barrier to reproducibility is that the four official checkpoints are hosted in a private Hugging Face repository requiring authorization, and the 46-million-cell training dataset was not released with the code. Researchers should next examine data licensing, external cross-platform validation, and whether the model retains its advantage when partial-expression hints are removed.

Sources

  1. CellWorld: From Gene-Level Reconstruction to Latent Cell Prediction in Spatial Transcriptomics Foundation Models
  2. UoM-HealthAI/CellWorld