Towards Scalable Context-Aware Single-Cell Spatial Transcriptomics Prediction from Histology Images

Conference: NeurIPS2026
arXiv: 2609.36429
Code: https://github.com/zjgao02/CELLO
Area: Computational Biology
Keywords: single-cell spatial transcriptomics, histology images, location querying, distance-decay cross-attention, pathology foundation models

TL;DR

CELLO shares one pathology foundation model forward pass per histology tile and predicts cell-level gene expression through continuous location querying and distance-decay cross-attention, improving average PCC on 52 paired H&E–Xenium samples while achieving a mean 14.0× speed-up over DeepSpot2Cell in timing that excludes upstream segmentation.

Background & Motivation

Hematoxylin and eosin (H&E) staining reveals tissue morphology, whereas spatial transcriptomics supplies molecular expression information at cellular locations. Existing image-to-expression prediction typically follows the resolution of Visium spots: many cells in a region share one aggregate expression vector, so accurate regional prediction does not imply recovery of each cell's state. Xenium supplies finer paired supervision, making direct single-cell prediction possible, but the image encoder's spatial scale remains an obstacle. One visual token can cover multiple cells, and treating it as an independent cell representation mixes neighboring signals.

DeepSpot2Cell connects per-cell crops to a pathology foundation model (PFM), benefiting from pretrained morphology representations but requiring separate large-encoder passes for many cells; cropping or resizing can also alter morphology and context. Mask-based approaches such as GHIST preserve segmentation-defined cell structure but lack equally strong modern PFM representations and inherit boundary errors. Rather than making each cell view a separate image, this paper retains the shared spatial representation of an image tile and queries it at cell locations. This amortizes encoding and avoids directly pooling features with unreliable boundary masks.

Here, “cell-level” means that outputs correspond individually to supplied cell locations, not that the model discovers cells from the image or perfectly separates mixed visual tokens into pure cellular signals. Core Idea: share a PFM's spatial tokens within each tile, then use continuous location queries and spatially biased contextual refinement to produce trainable expression predictions for all supplied cell locations.

Method

Overall Architecture

The inputs are an H&E image tile extracted from a whole-slide image (WSI) and the centroids of cells within it; the output is a nonnegative gene-expression vector for each location. The default Virchow2 encoder processes one 224×224 tile into a spatial token grid. Bilinear sampling initializes cell queries, cross-attention with 2D RoPE and a Gaussian distance prior reads spatial tokens within the tile, and an expression head regresses gene values.

Training additionally uses an auxiliary branch that predicts aggregate tile expression, supervised by the sum of cell-expression targets within that tile. Inference needs the image and cell locations, not measured expression from that sample. Obtaining locations still relies on spatial-transcriptomics coordinates or upstream detection/segmentation, so this is not a segmentation-free deployment pipeline.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    I["H&E tile"] --> A["Shared Spatial<br/>Encoding"]
    A --> B["Continuous Location<br/>Querying"]
    S["Cell centroids"] --> B
    B --> C["Distance-Decay<br/>Cross-Attention"]
    A -->|spatial tokens| C
    C --> D["Masked Expression<br/>Regression"]
    D --> O["Inference: per-cell expression"]
    A -.->|training: auxiliary aggregate-expression branch| D
    Y["Training: measured cell expression<br/>and its sum within the tile"] -.->|supervision; measured genes only| D

Key Designs

1. Shared Spatial Encoding: move repeated large-model computation to the tile level

The principal cost of per-cell cropping is repeatedly encoding overlapping tissue regions. CELLO runs the PFM only once per tile and retains spatial tokens instead of using one CLS vector to represent every cell in the region. Reshaping the spatial tokens into a regular two-dimensional feature map preserves their correspondence to image regions, providing geometric anchors for querying. The major encoding cost grows with the number of tiles rather than requiring one large-model pass for every cell. Cell queries and attention still grow with cell count; total computation is not independent of the number of cells.

“One encoding pass” does not mean one pass for an entire WSI. Training and evaluation use non-overlapping 224×224 tiles, with each cell assigned to the tile containing its coordinate and allowed to read only that tile's visual tokens. The model does not explicitly build a cross-tile cell graph. References to cell interactions should therefore be interpreted as reading local morphological context, not directly measuring or identifying biological cell–cell communication.

2. Continuous Location Querying: preserve coordinate differences on a sparse token grid

Original cell coordinates are in WSI pixel space. Subtracting the tile's top-left corner produces local coordinates, which are normalized to the [-1,1] range required by GridSample. Bilinear sampling interpolates the four nearest spatial tokens around a cell to initialize its query, so cells at different sub-token locations need not mechanically share the same discrete token. This operation is differentiable, allowing the sampled features to participate in downstream task training.

Location interpolation avoids per-cell cropping, resizing, and boundary-mask pooling, but does not create resolution absent from the encoder. The default Virchow2 token stride is 14 pixels, giving a 16×16 spatial grid for a 224×224 tile; multiple cells in one neighborhood may still have very similar initial queries. These are location-specific tissue representations, not perfectly isolated representations of cell interiors, motivating further contextual refinement.

3. Distance-Decay Cross-Attention: guide local context reading with a spatial prior

Each initial cell query acts as the query, while all spatial tokens in its tile supply keys and values. Two-dimensional rotary positional embeddings (2D RoPE) introduce relative position into attention. Cell pixel coordinates are first divided by the token stride to obtain continuous token-grid coordinates, and keys use token-center grid coordinates. Both therefore operate at the same scale rather than mixing raw pixel coordinates with token indices.

Before softmax, the model also adds a distance-decay bias. For local cell position \(\mathbf{s}_i^{\mathrm{loc}}\), token pixel center \(\mathbf{c}_m\), and stride \(p\), the readable LaTeX portion of the paper's Equation (4) is:

\[ b_{i,m}=-\frac{\|\mathbf{s}_i^{\mathrm{loc}}-\mathbf{c}_m\|^2}{2p^2}. \]

Under exponential normalization, this bias corresponds to a Gaussian distance prior: nearby tokens receive greater weight when content matching is otherwise equal. It is not a hard cutoff that retains tokens inside a radius and deletes those outside. Distant tokens can still participate, and content similarity and the distance prior jointly determine attention. Layer normalization then produces a context-refined cell representation for the expression head.

The biological intuition is that cellular state depends on the surrounding tissue microenvironment, but experimental evidence needs separate interpretations. Gains over Grid support contextual refinement as a whole; short-budget experiments removing only the distance bias do not show a performance drop, so the Gaussian prior alone cannot explain all gains. CLS fusion is also not a default effective module: an ablation uses a scalar gate to blend the current tile's CLS with cell features and generally performs worse. This tests a patch-level summary, not genuine whole-WSI context, and cannot rule out the latter.

4. Masked Expression Regression: share a prediction head across different Xenium panels

The expression head covers a union vocabulary of 1,915 genes, but each sample measures only 280–480 genes, excluding negative-control and antisense probes. Training and evaluation select only the genes actually measured in that sample. Unmeasured genes are not negative examples with zero expression. Otherwise, the model would confuse panel differences with tissue differences, and missing measurements would contaminate evaluation.

A two-layer MLP with ReLU maps cell representations to nonnegative predictions. Targets are not raw transcript counts: counts are normalized to 100 per cell and transformed with \(\log(1+x)\). The auxiliary tile-expression target sums these cell-target vectors within the tile; it should not be substituted with a log transform after summing raw counts, nor does it constitute an independently measured spot. The method uses two granularities of the same paired single-cell supervision rather than deconvolving cell expression from weak spot-level supervision.

A Worked Example

Consider one 224×224 tile for which an upstream process provides 30 cell centroids; 30 is an illustrative number, not a reported dataset statistic. Virchow2 encodes the tile once to obtain 16×16 spatial tokens. The model interpolates a feature for each of the 30 positions and lets each query read the same 256-token context. Nearby cells share the candidate context but have different continuous query positions and distance biases.

If the sample measures 377 genes, the model still outputs a 1,915-dimensional vector, while the training loss considers only the corresponding 377 dimensions. The sum of the 30 cell targets supervises the auxiliary tile branch. At inference, neither these expression values nor a panel mask supplies molecular evidence: the image and upstream positions pass through the trained expression head to generate predictions. An output for an unmeasured gene cannot be considered validated in that sample merely because its dimension exists.

Loss & Training

The paper combines cell-level and tile-level Huber regression:

\[ \mathcal{L}=\mathcal{L}_{\mathrm{cell}}+\lambda\mathcal{L}_{\mathrm{patch}},\qquad \lambda=0.5. \]

The cell loss averages over cells and measured genes, while the auxiliary loss regresses the aggregate tile target. Huber loss is motivated by heavy tails and outliers in normalized expression, with regression and PCC evaluation rather than generative modeling of raw counts. The paper does not clearly specify every auxiliary-branch detail, so no precise additional branch equation is reconstructed here.

Appendix C uses four H100 GPUs with 128 tiles per GPU, a global batch size of 512, and 200 epochs. Training uses FP16, AdamW, cosine learning-rate scheduling, and gradient clipping at a maximum norm of 1.0. The backbone's base learning rate is \(5\times10^{-6}\) with weight decay 0.01; task heads use \(2\times10^{-4}\) with weight decay 0.05. The backbone is frozen for the first 10% of epochs, progressively unfrozen over the next 30% with learning rate increasing linearly to \(2\times10^{-5}\), and fully fine-tuned thereafter.

Every five epochs, validation uses non-overlapping WSI sliding-window inference, selecting the checkpoint with the highest gene-level PCC; the random seed is fixed at 42. DeepSpot2Cell retains a frozen foundation model. The authors acknowledge that representation adaptation may contribute to the accuracy gap, so this comparison does not strictly isolate the querying mechanism.

Key Experimental Results

Main Results

The dataset contains 52 same-section H&E–Xenium pairs from HEST-1k across 12 organs. Appendix B reports 36/4/12 training/validation/test samples, with 520,402 tiles and 9.53M cells after preprocessing. Sections are stained with H&E and re-imaged after in situ decoding; registration and quality control are taken from HEST-1k without additional alignment by the authors.

Evaluation uses gene-level Pearson correlation: linear correlation between predicted and measured expression across cells is computed for a gene, then averaged over the selected gene set. HVG, MPG, and SVG top-k sets are evaluation subsets, not gene-selection filters for training. Top-k PCC is neither full-panel accuracy nor clinical validity. The following retains group averages from Table 1.

Configuration	ID HVG-10	ID HVG-50	ID HVG-200	OOD HVG-10	OOD HVG-50	OOD HVG-200
UNet3+	0.3569	0.2044	0.0711	0.2513	0.1395	0.0431
DeepSpot2Cell	0.4577	0.2773	0.0925	0.2884	0.1635	0.0524
Grid, without contextual refinement	0.5567	0.3700	0.1474	0.4083	0.2413	0.0862
CELLO	0.6421	0.4390	0.1718	0.4768	0.2801	0.0916

HVG-50 improves over DeepSpot2Cell by 58.5% in ID and 71.3% in OOD, and over Grid by 18.6% and 16.2%. The paper calls DeepSpot2Cell the “strongest baseline,” but Grid has higher group averages in the table. Grid is therefore the more informative comparison for interpreting module gains rather than adopting that baseline-ranking claim.

Ablation Study

The following combines two budgets from Appendix E. The first four rows are the main-model encoder/CLS comparisons; the final three use the same data, one GPU, batch size 16, 50 epochs, and a single seed. Their effect sizes should not be directly compared with those of the first four rows.

Configuration	ID HVG-50	OOD HVG-50	Note
UNI	0.3648	0.2417	Encoder comparison
H-Optimus-0 (H0)	0.3827	0.2524	Encoder comparison
Virchow2 / CELLO	0.4390	0.2801	Main setting, no cell CLS fusion
Virchow2 + CLS fusion	0.3975	0.2656	Fuses the current tile's summary
Short training, \(\lambda=0\)	0.3882	0.2574	No auxiliary tile loss
Short training, \(\lambda=0.5\)	0.3750	0.2349	Short-budget version of the default weight
Short training, \(\lambda=0.5\), no distance bias	0.3822	0.2468	Retains attention but removes the distance prior

Table 7 tests weights 0, 0.1, 0.25, 0.5, and 1.0. The default 0.5 is lowest under this budget, and removing the distance bias does not reduce performance. These short runs therefore do not support the claim that aggregate supervision and the distance prior are individually necessary. Isolated comparisons with equally sufficient training budgets and multiple seeds remain needed.

Key Findings

• Average gains do not imply superiority for every organ and gene set. Kidney HVG-50 is 0.3147 for Grid and 0.2743 for CELLO; Heart is 0.1868 and 0.1469. CELLO's Heart HVG-200 of 0.0258 is also below UNet3+'s 0.0430, indicating weak fine-grained signals in some OOD cases.
• With CLS fusion, the simple average of nine group-averaged metrics for Virchow2 falls from 0.4793 to 0.4143 in ID and from 0.3298 to 0.2867 in OOD. Individual entries can improve: OOD HVG-200 rises from 0.0916 to 0.0972. The finding is an absence of consistent overall benefit, not a decline in every metric.
• In Appendix H's adaptation-free evaluation of five additional HEST-1k samples, HVG-50 ranges from 0.196 for TENX158 to 0.575 for TENX94, showing substantial sample variation. These samples remain within the same data resource and do not constitute multicenter clinical external validation.
• Appendix I perturbs locations in 3,000 fixed test tiles: centroid noise with a 10-pixel standard deviation lowers mean per-tile PCC from 0.1348 to 0.1336, and 40 pixels lowers it to 0.1206; removing 40% of cells yields 0.1286. This metric pools cells and measured genes within a tile, unlike the main gene-level PCC, and does not test segmentation merge/split errors.

Efficiency is measured on one L40S using FP16 and 16 images per forward pass, with warm-up, GPU synchronization, and five repetitions. For each test WSI, 200 sampled tiles are timed and results scaled by the tile count. The following therefore gives projected whole-slide times, not direct timing of complete processing of every slide, and excludes model loading.

Efficiency measure	Value	Scope
CELLO time per WSI	67.5 ± 48.8 s	Excludes cell segmentation
DeepSpot2Cell time per WSI	947.2 ± 1068.3 s	Excludes cell segmentation
UNet3+ time per WSI	73.9 ± 52.1 s	Excludes cell segmentation
DeepSpot2Cell / CELLO	Mean 14.0×; median 11.7×	Per-slide ratios range from 3.5–26.3×
CellViT-SAM-H segmentation	382.5 ± 334.1 s	Shared upstream cost
DeepSpot2Cell / CELLO including segmentation	Median 2.5×	Not an end-to-end 14× claim
UNet3+ / CELLO including segmentation	Median 1.02×	Approximately on par

The conclusion states “up to 7.1×,” inconsistent with the main text/Appendix J's mean 14.0× and range 3.5–26.3×. The paper does not explain the corresponding protocol, so these claims are not reconciled here. The timing breakdown also includes CLS fusion, whereas the default accuracy results omit it; the exact correspondence between timing and accuracy configurations requires code verification.

Highlights & Insights

• The important engineering change is to share spatial encoding before querying rather than constructing independent inputs for every cell. This improves location-level prediction efficiency while making task-specific fine-tuning of a large backbone more feasible.
• Continuous querying and same-scale RoPE connect a regular visual grid to irregular cell locations. The approach could transfer to tissue-marker prediction at known coordinates, but it does not justify claiming recovery of independent molecular signals beyond encoder resolution.
• Panel masking explicitly distinguishes missing measurements from zero expression. This matters for cross-platform supervision and prevents treating the union output dimension as evidence that all 1,915 genes are evaluated in every sample.

Limitations & Future Work

• Inference still requires cell locations. The method reduces reliance on precise boundary masks, not detection/segmentation errors. Robustness studies should include real segmentation errors, registration drift, and alternative detectors.
• Each unseen organ has only one test WSI; Kidney and Liver test unseen health conditions in organs already present during training. These OOD results are transfer case studies rather than estimates of organ-level generalization. Table 2 lists Lymphoid training samples as Cancer/Healthy but its test sample as Diseased, conflicting with the main text's ID definition requiring both organ and condition to have been observed; this split discrepancy should be flagged separately.
• Main improvements combine querying, context, and backbone fine-tuning, and Appendix G's frozen per-cell-crop control does not fully separate these factors. Its CELLO values differ from the main table without sufficiently explaining the configuration, so they are not merged into one result series.
• Independent causal contributions from the distance bias and auxiliary loss lack sufficient experimental support. A useful next step is a crossed ablation at a consistent 200-epoch budget with repeated seeds and slide-level confidence intervals.
• Context stops at tile boundaries, and the CLS ablation tests only a current-patch summary. Whether broader tissue structure or cross-tile communication helps remains open; local CLS failure cannot answer that question.
• PCC measures linear covariation, not absolute expression calibration, rare cellular states, clinical endpoints, or treatment-decision validity. Predicted expression cannot replace measured molecular data without independent experimental and prospective validation.

Related Work & Insights

• vs DeepSpot2Cell: The latter uses per-cell crops and a frozen PFM, with spot-level weak supervision in its original formulation; CELLO shares tile features and directly uses paired single-cell supervision. Speed gains target per-cell large-encoder passes, while accuracy differences also involve fine-tuning and supervision settings.
• vs GHIST / UNet3+: This line constructs cell representations from boundary masks, whereas CELLO queries shared pretrained features by coordinate. The evaluated comparator is a UNet3+ baseline, not necessarily the full capabilities of GHIST; their segmentation-inclusive runtimes are approximately equal.
• vs Grid: Grid retains shared encoding and location querying but omits distance-decay cross-attention, making it the closest methodological ablation. Its averages exceed DeepSpot2Cell's, indicating substantial benefits from the shared representation itself.
• vs iStar / scstGCN: These approaches combine ST input with super-resolution expression reconstruction, unlike prediction using H&E and cell locations alone. Distinguishing inputs and supervision granularity is more informative than broadly comparing claims of “single-cell resolution.”

Rating

• Novelty: 4/5 — Integrates shared pathology tokens, continuous cell queries, and spatially biased refinement into a fine-tunable framework; individual mechanisms are not entirely new.
• Experimental Thoroughness: 3/5 — Broad organ, efficiency, and appendix analyses, but limited OOD samples, causal ablations, and configuration consistency.
• Writing Quality: 3/5 — A clear methodological trajectory with inconsistencies in speed-up protocols, baseline ranking, ID conditions, and CLS-context terminology.
• Value: 4/5 — Useful for scalable morphology-to-single-cell-expression research, but not a substitute for molecular measurements or clinical validation.