SpaAxFormer: Fine-Grained Token Modeling of Spatial Transcriptomics with Axial Transformers

Spatial transcriptomics combines gene-expression measurements with spatial relationships between cells. We investigate how spatial context can be incorporated into gene-resolved representations while preserving explicit input-gene correspondence. SpaAxFormer is a spatially constrained factorized formulation that maintains contextual cell-gene states throughout the encoder. It first contextualizes genes within each cell and then exchanges corresponding gene states across spatial neighbors, while a separate CLS representation provides a cell-level summary. Contextual gene states support gene-expression reconstruction, while CLS-conditioned reconstruction and contrastive learning provide complementary supervision for the cell-level representation. We evaluate SpaAxFormer on gene-expression imputation and spatial-domain identification across multiple spatial-transcriptomics technologies. Compared with the included baselines, imputation results were most favorable on dorsolateral prefrontal cortex (DLPFC) datasets and more mixed on additional technologies, while post-hoc spatial refinement consistently improved mean ARI on the Visium spatial-domain datasets where it was applied. In a single-slice component-addition analysis, mean clustering ARI increased as CLS-conditioned reconstruction and contrastive supervision were added. Together, these results show that the same gene-resolved formulation can support complementary gene-level prediction and cell-level spatial representation while preserving explicit input-gene correspondence.