HEIST: A Graph Foundation Model for Spatial Transcriptomics and Proteomics Data
Authors: Hiren Madhu
João Felipe Rocha
Tinglin Huang
Siddharth Viswanath
Smita Krishnaswamy
Rex Ying
Venue: ICLR 2026 Poster (2026)
Keywords: Pretrained models
Spatial transcriptomics
AI for Science
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Single-cell transcriptomics and proteomics have become a great source for data-driven insights into biology, enabling the use of advanced deep learning methods to understand cellular heterogeneity and gene expression at the single-cell level. With the advent of spatial-omics data, we have the promise of characterizing cells within their tissue context as it provides both spatial coordinates and intra-cellular transcriptional or protein counts. Beyond transcriptomics, proteomics offers a complementary view by directly measuring proteins, which are the primary effectors of cellular function and key therapeutic targets. However, existing models either ignore the spatial information or the complex genetic and proteomic programs within cells. Thus they cannot infer how cell internal regulation adapts to microenvironmental cues. Furthermore, these models often utilize fixed gene vocabularies, hindering their generalizability to datasets with different genes than pretraining. In this paper, we introduce HEIST, a hierarchical graph transformer foundation model for spatial transcriptomics and proteomics. HEIST models tissues as hierarchical graphs. The higher level graph is a spatial cell graph, and each cell in turn, is represented by its lower level gene co-expression network graph. Rather than using a fixed gene vocabulary, HEIST computes gene embeddings from its co-expression network and cellular context. HEIST  achieves this by performing both intra-level and cross-level message passing to utilize the hierarchy in its embeddings and can thus generalize to novel datatypes including spatial proteomics without retraining. HEIST is pretrained on 22.3M cells from 124 tissues across 15 organs using spatially-aware contrastive and masked autoencoding objectives. Unsupervised analysis of HEIST embeddings reveals spatially informed subpopulations missed by prior models. Downstream evaluations demonstrate generalizability to proteomics data and state-of-the-art performance in clinical outcome prediction, cell type annotation, and gene imputation across multiple technologies.

Published as a conference paper ICLR 2026.
HEIST: A GRAPH FOUNDATION MODEL FOR SPATIAL TRANSCRIPTOMICS AND PROTEOMICS DATA
Hiren Madhu, João Felipe Rocha, Tinglin Huang, Siddharth Viswanath, Smita Krishnaswamy" Rex Ying"
Yale Univeristy. USA "Co-Senior Author
Cooresponding Authors: (smita.krishnaswe wwwshhaswhhh,,, (ox.ying)@yyll.... edu
ABSTRACT
Single-cell transcriptomics and proteomics have become a great source for data-driven insights into biology, enabling the use of advanced deep learning methods understand cellular heterogeneity and gene expression at the single-cell level. With the advent of spatial- omics data, we have the promise of characterizing cells within their tissue context as it provides both spatial coordinates and intra- cellulc transcriptional or protein counts. Beyond transcriptomies, proteomies offers a complementary view by directly measuring proteins, which are the primary effectors of cellular function and key therapeutic targets. However, existing models either ignore the spatial information or the complex genetic and proteomic programs within cells. Thus they cannot infer how cell internal regulation adapts to microenvironmental cues. Furthermore, these models often utilize fixed gene vocabularies, hindering their generalizability to datasets with different genes than pretraining. In this paper, we introduce HEIST, a hierarchical graph transformer foundation model for spatial transcriptomics and proteomics. HEIST models tissues hierarchical graphs. The higher level graph is a spatial cell graph, and each cell in turn, is represented by its lower level gene co-expression network graph. Rather than using fixed gene vocabulary. HEIST computes gene embeddings from its co-expression network and cellular context. HEIST achieves this by performing both intra-level and cross-level message passing to utilize the hierarchy in embeddings and can thus generalize to novel datatypes including spatial proteomics without retraining. HEIST is pretrained on 22.3M cells from 124 tissues across 15 organs using spatially-aware contrastive and masked autoencoding objectives. Unsupervised analysis of HEIST embeddings reveals spatially informed subpopulations missed by prior models. Downstream evaluations demonstrate generalizability to proteomics data and state-of-the the- art performance in clinical outcome prediction, cell type annotation, and gene imputation across multiple technologies.
INTRODUCTION
Single-cell RNA sequencing (scRNA--sqq) has revolutionized our ability to study gene expression at the resolution of individual cells, offering data-driven insights into biology. The complexity of these datasets has fueled the development of machine learning methods for modeling cellular diversity, predicting cell states, and imputing or denoising gene expression values (Cui et al., 2024; Hu et al., 2021; Grønbech et al., 2020; Wen et al., 2023a; Bravo Gonzalez-Blar et al., 2024; Dijk et al., 2017). However, limitation of scRNA-seq is the lack of spatial context of cells within tissues, which is important for understanding processes such as tissue organization, microenvironment interactions and how gene co-expression influences tissue-level behaviors. Spatial transcriptomics is an emerging technology that bridges this gap by preserving the physical locations of gene expression measurements, enabling holistic study of tissue architecture, cell-cell communication, and

---

Published as a conference paper ICLR 2026.
A
HEIST pre-training dataset.
B
124 Tissues, 15 Organs, 22.3M Cells
Breas:
4.3M
Colon
3.7M
2.5M
Possate
1mnel
1.7M
Sputial graph
HEIST
Cell clustering
Liver
1.5M
1.3M
MLP
Ovary
1.1M
Skin
570K
Resporse invonnotheaapy
Outcome& Remission
Uterus
843K
Pantress
671K
Minwol 1
Marrow
Konney 152K
378K
Cells
Genes
Triinable
Frozen
Hevok2
decoder
HEIST
Heart
226K
Maseld Genes Genns-Fign
cellmb types
Gene imputation
Figure I: Overview of the HEIST framework. (A) HEIST is pre-trained on large-scale spatial transcrip- tomies dataset spanning 124 tissues and organs (22.3M cells). (B) HEIST encodes both gene co-expression networks and spatial cell graphs to support downstream tasks such as cell clustering, gene imputation, and clinical outcome prediction (e.g., immunotherapy response, remission). The HEIST decoder can be fine-tuned while the encoder remains frozen.
their aberrations in the tumor context (Xiao & Yu, 2021; Rodriques et al., 2019). Similarly, spatial proteomies assays were developed to directly capture protein abundance, expression and signaling pathways, offering a complementary layer of biological insight. Despite advances in spatial omics, datasets remain limited by low throughput, platform and tissue heterogencity, and scarce labels, often necessitating dataset-specific models (e.g., MIBI (Angelo et al., 2014), Imaging CyTOF). A foundation model trained on diverse spatial omics can address these challenges by learning generalized representations across tissues, organs, settings, and technologies, enabling strong performance on downstream tasks even with limited data.
To this end, propose HEIST (Hierarchical EmbeddIngs for Spatial Transcriptomics), the first foundation model for spatial transcriptomics that explicitly models both spatial proximity and internal co-expression networks while also enabling cross- modal transfer proteomics. Through cross-level message passing, internal embeddings are shaped by the spatial context of their parent cell, while cell embeddings are updated from their constituent genes. This produces adaptive representations, allowing the same gene to be encoded differently depending on context. In this way, HEIST unifies spatial and molecular hierarchies, linking gene-level interactions tissue-level phenotypes.
Prior foundation models such as SCGPT (Cui et al., 2024), SCFOUNDATION (Hao et al., 2023), and CELLPLM (Wen et al., 2023a) either neglect cell-cell structure or are limited to predefined gene sets, hinder- ing generalization to unseen genes and proteins. Graph-based methods like GRAPHST (Long et al., 2023) and STAGATE (Dong& & Zhang, 2022) (Explained in detail in Appendix A) capture spatial neighborhoods but remain task- task-specific and non non-transferable. While transformer based models like SCGPT-SPATIA (Wang et al., 2025) treat all genes as fully connected graph and fail to model the inductive bias of co-expression. HEIST overcomes these gaps by associating each cell with a co-expression network that interacts with its spatial neighbors, incorporating microenvironmental cues into representations.
We pretrain HEIST (Figure 1(A)) 22.3 million cells from 124 tissues across 15 organs and two technologies, and evaluate on four downstream tasks (Figure I(B)): clinical outcome prediction, cell type annotation, gene imputation, and cell clustering. We achieve state-of-the-art performance across seven organs. HEIST enables cellular subpopulations that previous models fail to do and is 8x faster
the discovery of spatially than SCGPT-SPATIA and 48x faster than SCFOUNDATION, We summarize contributions follows:
2

---

Published as a conference paper ICLR 2026.
Modeling inter-cellular and hierarchical effects of co-expression networks: HEIST is the first founda- tion model for spatial omics to explicitly incorporate co-expression networks alongside spatial graphs in a hierarchical graph, enabling a local gene programs to influence tissue-level organization, and vice versa.
Hierarchical representation learning with biological inductive bias: Using biologically motivated hierarchical modeling, HEIST captures fine-grained gene co-expression within cells and long-range cellular interactions through novel cross-level message passing, producing biologically contextualized embeddings.
A task-agnostic, general-purpose foundation model: HEIST is trained in self-supervised manner on large-scale corpus of spatial transcriptomics data comprising over 22.3M cells spanning 15 organs and 124 tissues. In downstream evaluations, HEIST achieves state-of-thee offttheartt performance across four diverse tasks outperforming prior models, generalizes to proteomics data, while being computationally efficient.
2 METHOD
In this section, we describe the architecture of HEIST, hierarchical graph transformer designed to learn multi-level embeddings for spatial transcriptomics and proteomies data. HEIST takes as input, a set of graphs EG,(C,E,P, P) 1 (V, & k where Ge is spatial graph capturing spatial proximity, C' is the set of cells, & is the set of spatial edges, P E IR. C]x2 represents the spatial positions, T is the set of cell types, and t2 is cell-type of cell k. Each graph g represents gene co-expression network within cell k. V is the set of genes, and E., and X₃ denote the edges of the gene co-expression network and expression values for cell k respectively.
Hierarchical Graph Construction. As shown in Figure A1, we first preprocess the data by removing outliers, normalizing gene expression, and retaining highly variable genes. Then we apply MAGIC (Dijk al., 2017) to denoise gene expression values and reduce dropout noise. To build the gene co-expression networks, we first subset the cells based on cell-types using provided annotations or Leiden clustering. Following this, we compute pairwise mutual information between denoised genes within each type, and connect gene pairs above a threshold This results in total of [T] gene co-expression networks with mutual information prior. We create the spatial cell-cell graph by computing Voronoi polygons from cell coordinates and connecting cells in adjacent polygons. We then connect each cell with the gene co-expression network of that cell-type. The resulting outputs are a spatial graph (c E,P,T) and a set of gene co-expression networks 5 (V, Er,, X= We discuss the graph creation length in Appendix B. Given these graphs, IR clxi and Zg E glcl[Vixd, such that
HEIST computes cell embeddings and gene embeddings Z
Ze, Zg HEIST (G.(C.,, 19 (V, Xo))
As shown in Figure 2, the model first performs intra-level message passing (Equation 1) within each graph, followed by cross-level message passing (Equation 2) to integrate multi-modal information. HEIST is pre- trained using a combination of contrastive and auto-encoding objectives on gene expression and cell locations. By using these components, HEIST can learn expressive and context-aware cell and gene embeddings that reflect biologically meaningful relationships between cells and genes, Note that as a result of this setup, gene representations are themselves learned in the context of the hierarchical graph, instead of based on fixed gene vocabulary. They are initialized with rank- based and sinusoidal positional encodings, and dynamically updated through message passing in co-expression graphs, allowing HEIST to generalize to unseen genes or proteomic features by grounding embeddings in co-expression dynamics.
2.1 HEIST ARCHITECTURE.
First we initialize the input cell embeddings H. and gene embeddings for cell H using positional encodings explained in Appendix C.I. Then they are passed through HEISTLAyer L times, and the
3

---

Published as conference paper ICLR 2026.
H endbolding
Graph
Projection
Spatial
Selope
lcaming
Gogh Traneoomer Layar
embeddings
raunge massgepsiing
X
entading
Figure 2: HEIST Architecture. HEIST Layer block presented in Figure A2.
representations calculated using the equation below:
H HE.STTLaerrHH (I) ( (E,, 1
HEISTLayer divided into two steps, intra-level message passing and cross-level message passing. perform intra-message passing and calculate the intermediate representations shown in equation below:
H Ce meel@Graphrraasfrrmee (H (H))))) 8) H GeneGraphTrannsformer (H( {E,, [C|
k=n
where Cell CeL1GraphTransfornee and GeneGraphTrans are graph transformers as explained the Appendix Section C.2.
Cross-level message passing, To integrate spatial and gene modalities, HEIST performs cross-level message passing between cell and gene graphs each layer (Figure 2). Gene embeddings are updated based on spatial context through their parent cell's embedding, while cell embeddings are refined using pooled summaries of their genes, ensuring transcriptional states shape spatial identity. This bidirectional interaction captures tissue hierarchy, where gene expression depends on both co-expression signals and spatial microenvironments HEIST thus learns representations that reflect both local gene co-expression and large-scale tissue structures.
We use the directional attention mechanism to perform cross message passing shown in equation below:
Ho W H We
CrossMesssagePass
H Hfro))
(Ho W.).
(2)
vd
where W WE, W. dddd are learned weight matrices, and -,->> is the row-wise inner product. Let ( and H denote the intermediate gene and cell embeddings at layer after intra-level updates and before cross-level integration. HEIST updates these representations using the equation below:
H CrosMessagePassinnglaaaess
TressMeesaggllasssinggyyyy (H.)) H
E Ric|Vi represents the cell embeddings repeated | times to align with the gene embeddings each cell. Each gene receives information from its parent cell, enabling spatial context to modulate gene-level representations. Conversely, H E glc|xd is obtained by aggregating the gene embeddings within each cell: H.") [AGG(H)) ... AGG H, Cl }], where AGG(-) can be aggregation function such as MEAN pooling or differential pooling (Ying et al., 2018). This allows the cell embedding to he informed by the internal transcriptional state of the cell.
where (
Advantages. This formulation enables targeted, direction-aware communication between modalities while preserving their structure and semantics. It is well suited for spatial transeriptomics and proteomics data,
4

---

Published as a conference paper ICLR 2026.
where cell and gene representations must be coupled yet retain distinct meanings. By letting each gene attend to its parent cell embedding, HEIST shapes gene representations in cell-specific, spatially informed way, and vice versa. Unlike symmetric attention or feature concatenation, this directional mechanism preserves data hierarchy, respects modality roles, and captures how local gene co-expression drives tissue organization. Directional message passing respects the natural hierarchy between genes and cells, allowing each to influence the other without collapsing their distinct biological roles, enabling HEIST to learn spatially informed, biologically grounded representations that generalize across tissues and proteomics.
Finally, the intra- and cross-level message passing steps repeated L times, yielding final embeddings
Ze 2 HEIST (G,CC., (9. (V, x,),
Xa)]
Decoder. After calculating the final representations, we pass the embeddings into decoder to reconstruct the original spatial locations using HEIST-Decoder:
P.
| HEIST-Decoder (G,(C,E), (9 (V, E,, Ze: Zoy)
where, HEIST-Decoder is a 3-layer GIN network (Xu et al., 2018).
2.2 PRE-TRAINING TASKS.
Contrastive Learning. We contrastive objective to learn context-aware representations by bringing similar cells and genes-such as neighboring cells of the same type or co-expressed genes-closer in embedding space while pushing dissimilar pairs apart. This separates functionally distinct cell populations and gene modules, even when spatially close. Additionally, we introduce cross-level contrastive alignment to ensure consistency between gene and cell representations, so cell embeddings reflect gene expression patterns and gene embeddings incorporate spatial context. The contrastive loss is calculated using the equation below:
Lessa = log Piees Re + EnCN e 23777 lyozy = log " "big/1 EEEN; ein)
lessy log " 1/T ELEN, 1///"
Leomrnaive
Σ JEP
(Lesse + lasss) + Σ l Eyezg.
(D-DlEV2
where -): denotes similarity function (e.g., cosine similarity), <+ C,g I. and e 9 denotes contrastive loss between cells, genes, and cell-genes respectively, 7 is a temperature parameter that controls the sharpness of the contrastive distribution, and & is the cell type label of cell i. The set of positive pairs P (i. )) t4 ty. d(i, < r}. where d(i.j) j) is the spatial distance between cells i and j. and is a fixed spatial radius. Similarly, the set of negative samples for cell i, N₁; = [k t, * tke d(i, k) S r), i.e., cells within radius that belong to different cell type.
Masked-auto encoding. We also train HEIST with a masked auto-encoding loss to improve reconstruction and robustness. By masking subsets of cell and gene nodes, the model learns to reconstruct gene expression and spatial coordinates from the remaining context, reflecting real-world challenges like dropout and noise in spatial transcriptomics. This encourages the model to infer missing data, generalize across datasets, and use gene signals to recover spatial context and spatial cues to predict gene expression. After reconstructing the spatial locations and gene-expression, the masked auto-encoding loss is calculated using equation below:
[C]
1 [C]
muae MSE(f O mask P O maske) + M C mask g' X mask
where mask, is the mask over the cell locations, and mask? is the gene-expression mask for cell k.
5

---

Published as a conference paper ICLR 2026.
Final loss. Contrastive learning structures the latent space to emphasize biological similarities and differences, promoting better separation of cell types and gene programs. In contrast, masked autoencoding ensures that embeddings retain rich information content necessary for reconstructing gene expression and spatial locations. Together, they prevent trivial collapsed representations and produce embeddings that are both discriminative and information-rich. Hence, the final objective is a weighted sum of the contrastive and autoencoding losses, along with an orthogonality regularization which encourages the embedding dimensions to be decorrelated, promoting diverse and non redundant representations (Zhang ct al., 2021):
L c=o(7) Leontrasivee + (1 - (1-a(7)) Lmae + A (|L -.ll--Z,Z.o 1o + La Z 11
where À is a regularization weight, a is sigmoid function to balance the loss terms, and 7y is learnable scalar that dynamically balances two terms.
Computational efficiency, As shown in Table A4 in the Appendix, HEIST demonstrates significant com- putational advantages, achieving 8x faster embedding extraction time compared to SCGPT-SPATIAL and 48x faster than SCFOUNDATION. This efficiency comes from HEIST's sparse modeling, which avoids the expensive full self-attention computations required by transformer based models like SCGPT-SPATIAS
3 EXPERIMENTS
In this section, we first describe the pretraining datasets used to train HEIST. We then outline the downstream tasks and corresponding datasets, followed by baselines, results, insights, and ablation studies.
Pretraining Datasets. HEIST is trained on a large and diverse collection of nigh-resolution spatial tran- scriptomics datasets. primarily generated using single-cell technologies such as MERFISH and Xenium. The pretraining dataset comprises 22.3M cells from 124 tissue slices across 15 organs [ef. Figure I(A),, including 13.3M cells from 10x Genomics, 8.7M from Vizgen, and 360K cells from the Seattle Alzheimer's Brain Atlas. The large scale and diversity of the dataset improve the reliability and robustness of learned representations, enabling better transferability to downstream tasks across varying biological contexts and technologies. A detailed breakdown of datasets, descriptions, and sources is provided in Appendix D, and we provide a spreadsheet of the dataset details in the supplementary material.
Experimental setup. We provide pretraining hyperparameters in Appendix Table A1. HEIST pretrained on NVIDIA L40s GPUs (40GB each), with each epoch taking approximately 3 hours. Although the maximum number of epochs set to 20, early stopping based validation loss typically halted training around Epoch 5 or 6. For downstream evaluations, we assess HEIST in both zero-shot and fine-tuning settings. In the zero-shot setting, the pretrained model is directly evaluated on unseen spatial transcriptomics datasets to assess generalization without further training. For fine-tuning, we first extract embeddings from the frozen model and either train an MLP prediction head or fine-tune the decoder. We perform each experiment 5 times provide mean and standard deviation, except in Charville and UPMC datasets where we use the split the data using method from (Wu al., 2022). The code is available at https: / //gitti thub. com/ homm/Graph and Geomet ri c- Learning ST.ST.
Downstream Tasks. We evaluate HEIST across four different spatial transcriptomics and proteomics technologies, five organs, and four downstream tasks-cell clustering, cell type annotation, clinical outcome prediction, and gene imputation- to assess both biological insight discovery and clinical relevance HEIST.
Cell clustering is critical for discovering novel cell types and understanding how microenvironmental factors shape cellular behavior, particularly in tumor microenvironments. An expressive model should not only cluster cells but also reveal microenviromment-driee cell populations. Clustering is performed using frozen embeddings, and performance is evaluated on datasets two spatial transcriptomics datasets SEA-AD (Gabitto al., 2024) and Merfish Lung Cancer (Chen et al., 2024), and three proteomics datasets Charville, UPMC, and DFCI (Wu et al., 2022), using normalized mutual information (NMI). Cell type annotation classifies cells into known biological categories, enabling functional
subclusters, providing insights into spatially informed
6

---

Published as a conference paper at ICLR 2026.
interpretation of cellular diversity, HEIST embeddings are extracted from the frozen encoder, and an MLP classifier is trained on labeled datasets including SEA-AD, Charville, UPMC, DFCI, and MERFISH lung cancer, with performance evaluated using F1 score.
Clinical outcome prediction aims to classify entire tissues, predicting outcomes such as immumotherapy response, treatment outcomes, remission status, and placenta condition. This task is essential for clinical decision-making and understanding disease progression. HEIST is evaluated on datasets including proteomics data Charville (Colon), UPMC (Neck), and DFCI (Neck) (Wu et al., 2022) collected using CODEX, skin cancer data (Ptacek et al., 2021) collected using MIBI, and spatial transcriptomics data of placenta collected using Xenium. Predictions are made using frozen cell embeddings and an MLP classifier, and we report AUC-ROC. Gene imputation recovers missing or noisy gene expression values, common issue in spatial transcriptomics due to measurement limitations. We perform gene imputation by predicting masked gene values, using stratified sampling based on gene sparsity following the approach of Avşar & Pir (2023). This task is evaluated in both zero-shot and by fine-tuning the decoder, reporting Pearson correlation between predicted and true gene expression values. We explain these tasks in further details in Appendix D.2.
Baselines. We benchmark HEIST against a diverse set of baselines, including graph-based spatial models, single-cell foundation models, and recent spatial foundation approaches. In particular, we include STA- GATE (Dong & Zhang, 2022) and GraphST (Long et al., 2023), which capture local spatial relationships but struggle to generalize across datasets. We also evaluate against scFoundation (Hao et al., 2023), large-scale single-cell foundation model that ignores spatial context. We also include comparisons with recent spatial foundation models CelIPLM (Wen et al., 2023a), No (Blampey et al., 2024) and scGPT-spatial (Wang et al., 2025), which incorporate spatial information but do not explicitly model hierarchical gene-cell inter- actions. Finally, for the gene imputation task, we compared against MAGIC a denoising method based on graph diffusion. For proteomics datasets such UPMC, basclines like scGPT-spatial and CellPLM were originally pretrained on fixed gene vocabularies, supporting only 6 of the 28 markers available. To enable comparison, report settings: (i) unaligned, where models process only the subset of overlapping genes, and (ii) aligned, where we manually remap marker identities to the closest supported genes following prior work. It is important note that SCGPT pretrains on single-cell spatial trasncritomics data mixed with Visium (spot-level), while CELLPLM combines scRNA-seq with single-cell spatial transcriptomics. HEIST, in contrast, trained exclusively on single-cell resolution spatial transcriptomics. However, there is major overlap in the single-cell spatial transcriptomics portion of training data between HEIST and the baselines.
3.1 RESULTS
Spatially-aware cell type discovery. A key strength of HEIST is its ability to uncover ppatially-nforweee cellular subpopulations. Figure shows that HEIST embeddings not only separate canonical cell types, but also resolve finer subclusters that align with local microenvironmental context. These subclusters correspond meaningful biological distinctions. Existing approaches like CELLPLM and SCGPT-SPATIAL fail capture sub clusters and collapse such substructure, missing microenvironment effects (Figure A3, Appendix). This ability to differentiate spatially-drive heterogeneity is essential for discovering novel biomarkers.
Gene imputation. Table reports the performance on the gene imputation task for the placenta and skin datasets. HEIST achieves the best performance after fine-tuning surpassing all baseline mod-
Table I: Performance on gene imputation task
Model
Placenta
Skin
els by 2.5% on the placenta dataset and 9 9% on the skin dataset. We highlight comparisons with MAGIC, as our preprocessing also incorporates this method. The improvements achieved by HEIST can be attributed to its cross-attention mechanism, which integrates contextual information from neighboring cells to enhance imputa- tion for each target cell. Although HEIST's zero-shot performance is limited by the dataset-specific nature of gene expression patterns, fine-tuning allows the model to adapt effectively by leveraging its hierarchical structure and gene co-expression networks, resulting in improved performance.
MAGIC
0.249 0,000
0.671 a: 0.000
SeFoundation (Fine-mmed)
0.721 + +-0004
0.621
+ 0.0003
CellPLM (Fin uned)
0.801 1 L0011
0.723 1 1.0007
scGPT-apuial (Fine-tunod)
0.718
2:0000
0.740 1: 0.000
HEIST (Zeco Shot)
0.574
0.000
0.350
0.000
HEIST (Fine-tuned)
0.821 0.041
0,807 0.020
HEIST Imp. %
2.49
9.05
7

---

Published as conference paper ICLR 2026.
A
CelIPLM
scCPT-spatia
HEIST
Attrocyte
Chandeler
L4IT
Fndothelal
La/S IT
ET
LS/9 MP
Lamps
Lamp5 thi
Cluster1 NNlghborhood
Cluster Neighborhood
Cluster Neighbarhood
Silhouette Comparison
Mcrogla-PNN
dlisddedrocyte
Pudl
LEA
25.0%
Pubh
Snca
5.15
55%
18:00
Vp
LUT LR
0.17
Quater
CalIPM
acGPTapatial
HEST
Figure 3: HEIST accounts for tissue microenvironments (A) Comparison of cell embeddings (PHATE) for the same tissue slice from SEA colored by annotated cell types. HEIST demonstrates superior separation compared other methods. (B) Spectral clustering reveals three spattally-informed sub clusters in L4-IT cells. (C) Visualizing the sub clusters neighborhood distribution shows that each cluster accounts for neighborhood differences, demonstrating that HEIST embeddings effectively captures spatial microenvironments through cross cross-message passing. We show similar results for CEIIPLM and scGPT-spatial in Figure A3, showing that these models can infer spatially-infomee clusters. Numerical results available in Table A3 in Appendix Section E. (D) Silhouette Score comparision between each method.
Table 2: Performance on clinical outcome prediction. We classify cancer outcome, cancer remission, treatment response and placental conditions. Results the right of vertical are proteomics. NEM stands for "not enough markers" error from Novae when it can not match enough markers. *Results from (Wu et al., 2022)
Organ
Placenta
Colon
Neck UPMC
Neck DFCI
Skin
Dalasel
Charville
Melanoma
Task
Model
Condition
Outcome
Recurrence
Outcome
Recurrence
Outcome
Response
STAGATE
0.578 0.12
0.657 0.032
0.783 + 0.050
0.602 0.054
0.659 0.013
0.633 0.210
0.533 0.267
GraphST
0.659 = 0.059
0.828 = 0.088
0.645 : 0.026
0,582 = 0,061
0.683 1: 0.131
0.567 1: 0.170
0.644 = 0.15
Space-gm"
0.793
0.696
0.863
0.883
0.873
ScFoundation
0.601 4 0.16
0.71 = 0.122
0.787 + 0.1
0.678 = 0,061
0.689 + O.I.T
0,742 = 0.093
0.500 - 0.155
Novae
0.000
NEM
NEM
CelIPLM (unaligned)
0.682 + 0.15
0.744
0.006
0.801 + 0.032
0.681 = 0.134
0.667 + 0.045
0.750 + 0.083
0.580 + 0.133
CellPLM (aligned)
scGPT-spatial (unaligned)
0.602 + 0.08
0.834
0.081
0.806 + 0.019
0.717 = 0.117
0.676 + 0.097
0.875 + 0.040
0.600 + 0.000
scGP scGPT-spptial (aligned)
0.594
0.08T
0.854
0.005
0,702
0.001
0.789 + 0.080
0.858
0.006
HEIST
0.769 0.06
0.861
0.086
0,887 + 0.041
0.835 - 0.001
0.929 + 0.030
0.937
0.062
0.866 0.066
HEIST Imp.%
12.7
3.2
3.5
-3.3
5.2
Tissue-Clasiffcation. Table reports the tissue-level classification performance of various models across multiple datasets and tasks, measured by AUC-ROC. HEIST consistently outperforms existing models, achieving the highest scores in six out of seven evaluation scenarios. Notably, on UPMC tasks, HEIST surpasses the foundation models scGPT-spatial and CellPLM by 25.4% and 30%, respectively due to scGPT- spatial and CellPLM ignoring most of the available markers. Even after aligning, performance remains limited because these models cannot fully leverage all the available markers. In contrast, HEIST constructs co-expression networks directly from observed proteins, allowing to incorporate all available markers without retraining. This structural flexibility underlies its strong generalization proteomics.
8

---

Published as a conference paper ICLR 2026.
Table 3: Performance on cell type annotation. Annotations are provided as a feature in each dataset. Results the right of vertical are proteomics.
Organ
Lung
Brain
Colon
Neck
Neck
SEA-AD
Charville
UPMC
DFCI
STAGATE
0.2187 1 0.0570
0.3304 0.0625
0,2759
0.0490
0.0687
0.0136
0.0685 1 0.0213
GraphST
0.4081 + 0.0658
0.2296 0.177
0.0617
0.0199
0.0577 - 0.0261
ScFoondatton
0.150 1 0.014
0.2495 0.11147
0.3220
0.1421
0.0222 0.0079
0.041 1 0.022
Novac
NEM
0.2332 + 0.0434
0.2194 0.0455
NEM
0.0736 + 0.0215
C (unaligned)
0.5044 + 0.1607
0.6701 + 0.0827
0.4760 + 0.0669
0.0563 + 0.0212
0.0565 + 0.0179
CelIPLM (aligned)
0,3047 1 0.0040
0.0337 1 0.0032
0.0413 1 0.0015
scGPT-spptial (unaligned)
0.5671 + 0.1685
0.5907
0.0029
0.3494
0.0624
0.0464 0.0162
0.0618 + 0.0163
scGPT-sppatal (aligned)
0.3280
0.0499
0.2195 0.0490
0.0953 0.0190
HEIST
0.5126 + 0.1170
0.9953
0.0158 48.5
0.1293 0.5340
0.2826 0.0758
0.1124 0.0521
HEIST Imp.
-9.6 %
12.2
28.7
17.9
stage
Dreak Dranky
Broak stage BroA
Figure 4: HEIST attention reveals unique microenvironments depending on Braak stages
Cell-type annotation. Table 3 shows the performance on the cell type annotation task across multiple datasets. HEIST achieves the highest performance
four out of five datasets, with substantial gains UPMC and DFCI neck datasets (28.7% % and 17.9% % improvements, respectively). Notably, scGPT-spatial on the MERFISH Lung Cancer dataset, explaining its strong performance there on that data. These results again highlight the effectiveness and generalizability of HEIST.
pretrained
Recovering unique niches based on tissue conditions. In Figure 4, we visualize the highly attended edges in HEIST's attention blocks. We compare the attention-derived niches across Braak stages, and we can see that a clear progression emerges. In Braak O, the high-scoring edges form small, isolated neuronal clusters that are relatively sparse and spatially compact. By Braak IV and V, these clusters become more numerous and begin to involve broader set of cell types, indicating early signs of local microenvironmente restructuring. In Braak VI, the niches become much denser and more spatially extended. This gradual transition from isolated neuronal interactions to larger mixed mexee--eel hubs suggests that HEIST attention captures the increasing microenvironmental complexity and tissue reorganization that accompany the progression of pathology across Braak stages.
Ligand-receptor pair prediction. We train a linear probe to predict ligand-receptor receptor (LR) interactions for all edges in the tissue graphs. For each model, we extract the cell embeddings, concatenate the embeddings pair, and pass this concatenated representation through a simple single layer
of the two cells involved perceptron that predicts whether the pair corresponds to an LR interaction. As shown in the Table A9, HEIST achieves the highest AUC-ROC and outperforms all basclines by a clear margin. In particular, HEIST reaches AUC-ROC of 0.995 + 0.002. This result combined with the attention based results indicates that its embeddings capture the spatial signatures of LR communication more effectively than se scFuddation, CelIPLM, and scGPT-spatial.
9

---

Published as a conference paper at ICLR 2026.
Table 4: Ablation study showing that hierarchical modeling, cross-level message passing, and the training objectives are critical components for strong performance.
Model
Outcooe
Skin-Imputation
SEA-Cell classification
HEIST
0.861 + 0.086
0.807 = 0.020
0.995 + 0.015
No space (No Hierarchy)
0.596 + 0.028
0.345 0.010
0.179 + 0.038
No gene (No Hierarchy)
0.764 + 0.235
0.173 0.014
0.194 + 0.040
No pre-training
0.500 + 0.000
0.623 0.002
0.784 + 0.128
No cross message passing
0.625 + 0.125
0.531 0.005
0.955 + 0.041
No positional encodings
0.523 + 0.010
0.458 0.003
0.220 + 0.034
No contrastive
0.623 + 0.002
0.536 0.015
0.966 + 0.037
No MAE
0.658 + 0.076
0.495 0.006
0.162 + 0.038
No orthogonal regularization
0.594 + 0.031
0.646 0.013
0,992 + 0.020
Ablations. Table 4 shows that removing any key component of HEIST leads to performance degradation across tasks, confirming their importance. Hierarchical modeling and spatial information are most critical, with their removal causing the largest drops. Pre- tre-training is crucial, especially for clinical outcome prediction due to the skewed label distribution in these datasets. Cross-level message passing and contrastive learning significantly improve cell classification,, while the MAE is crucial for gene imputation and classification. In Table A5, we compare different positional encoding including I Laplacian PEs, random walk based PEs, and sinusoidal PE. We notice that sinusoidal PEs outperform the other PE methods. In Table A6, we present the results with constant and adaptive т (threshold for the co-expression graph), showing that the adaptive method we use outperforms the constant threshold. In Table A7, we ablate by removing each of the contrastive term; Results indicate that the three terms combinely capture complementary structural signals. We also present a way to interpret HEIST attention in Figure A4 and Table A8, indicating that HEIST attention focuses on niches and corresponds to known ligand-receptor pairs than existing models. In Table A 110, we show show different ways of mixing positional encodings and adding raw features. We notice that passing raw features through MLP before adding them results in performance drop in imputation because because the MLP-basee mapping smooths or compresses the raw coordinate and expression signals in a way that removes fine-grained local variation. In Table AI1, we show that using smaller batch sizes results in slight decrease in performance as smaller batch sizes provide less microenvironmental cues.
4 CONCLUSION
We present HEIST, hierarchical graph foundation model for spatial transcriptomics that jointly models gene co-expression and spatial adjacency within a unified framework. By integrating biologically motivated hierarchical representation learning with novel cross-level message passing mechanism, HEIST captures complex dependencies between genes, cells, and tissue-level organization. Pretrained on over 22.3M cells spanning 15 organs, HEIST achieves state-of-the-art performance across multiple downstream tasks and generalizes to proteoomics, while offering significant computational efficiency. Beyond predictive accuracy, HEIST enables the discovery of spatially-informedd cellular subpopulations, providing deeper insights into tissue microenvironments Our results highlight the importance of modeling both molecular and spatial information to advance the development of general-purpond transferable models for spatial omics.
ACKNOWLEDGMENTS
S.K. is funded in part by the NIH (NIGMSROIGM133599, ROIGM130847), NSF CAREER award IIS- 2047856, NSF IIS-2403311 NSF DMS-2327211 S.K is also funded by the Sloan Fellowship FG-2021-15888 the Novo Nordisk grant GR 12933. S.K. and R. Y. acknowledge funding from NSF CISE-2403311 R. Y. has also recieved Amazon Research Award 2024.
10

---

Published as conference paper ICLR 2026.
ETHICS STATEMENT
The research presented in this paper fully conforms to the ICLR Code of Ethics. datasets used this study are publicly available and were obtained from previously published works with appropriate citations. These datasets were originally collected under IRB approval or equivalent ethical review, with informed consent obtained documented in the source publications. We only fully anonymized data and did not conduct any new data collection or direct interactions with human subjects. No personally identifiable information or sensitive patient data included. Our methods designed to advance scientific understanding of spatial transcriptomics and biological modeling and do not pose foreseeable risks related to privacy, misuse, or social harm.
REPRODUCIBILITY STATEMENT
We have taken several steps to facilitate reproducibility of our work. A full description of our model architecture, training objective, pre-training data, hyperparameters, and evaluation metrics is provided in Sections 2-3.1 of the main text, Appendix C.2, D, and supplementary zip. To ensure transparency, we release our anonymized code and configuration files at anonymous repository link, which contains scripts for data preprocessing, training, and evaluation. Furthermore, we provide instructions for reproducing all reported experiments. The code is available at https: / /anonymous. wwtps/....eecc.cceeeeeecccc.ccee scwssceeeeeeccecc.cec.ce h/te:///EEISS E3F8. We will make the code-base available on GitHub after the decisions.
11

---

Published as a conference paper ICLR 2026.
REFERENCES
Michael Angelo, Sean C Bendall, Rachel Finck. Matthew B Hale, Chuck Hitzman, Alexander D Borowsky, Richard M Levenson, John B Lowe, Scot D Liu. Shuchun Zhao, et al. Multiplexed ion beam imaging of human breast tumors. Nature medicine, 20(4):436-442 2014.
Gülben Avşar and Pinar Pir. A comparative performance evaluation of imputation methods in spatially resolved transcriptomics data. Molecular Omics, 19(2):162-173, 2023.
Quentin Blampey, Hakim Benkirane, Nadège Bercovici, Fabrice André, and Paul-Henry Cournède. Novae: a graph-based foundation model for spatial transcriptomies data. bioRxiv, 2024, doi: 10. 1001/2244 URL https:: www.:.booxioe http:////w/////ee/e///////////// httpt///aarrr///////1///////////
.9..
09. 612009.
Guillaume Braun, Hemant Tyagi, and Christophe Biernacki. An iterative clustering algorithm for the contextual stochastic block model with optimality guarantees. In International Conference on Machine Learning, pp. 2257-2291. PMLR, 2022.
Carmen Bravo Gonzalez-Blas, Irina Matetovici, Hanne Hillen, Ibrahim Ihsan Taskiran, Roel Vandepoel, Valerie Christiaens, Leticia Sansores- Garcia, Elisabeth Verboven, Gert Hulselmans, Suresh Poovathingal, al. Single-cell spatial multi-omics and deep learning dissect enhancer-driven gene regulatory networks in liver zonation. Nature Cell Biology, 26(1):153-167, 2024.
Yadi Cao, Menglei Chai, Minchen Li, and Chenfanfu Jiang. Efficient learning of mesh-based physical simula- tion with bi-stride multi-scale graph neural network. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Con- ference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 3541-3558. PMIR, 23-29 Jul 2023. URL https: /twoo/oooeedeeoooooooooooo ings ngs ml r. press / v2022/ cao23a html.
Jonathan Chen, Justin Gainor, Ilya Korsunsky, and Nir Hacohen. Human lung cancer harbors spatially- organized stem-immunig hubs that associate with response to immunotherapy, February 2024. URL
ht https: //doi///oii/ org/10. 00055221//// wtt.////enooo.11. 11198494.
Haotian Cui, Chloe Wang, Hassaan Maan, Kuan Pang. Fengning Luo, Nan Duan, and Bo Wang. scgpt: toward building a foundation model for single-cell multi-omics using generative ai. Nature Methods, 21) 1480, 2024.
David van Dijk, Juozas Nainys, Roshan Sharma, Pooja Kaithail, Ambrose Carr, Kevin R Moon, Linas Mazutis, Guy Wolf, Smita Krishnaswamy, and Dana Pe'en er. Magic: diffusion-basee imputation method reveals gene-gene interactions in single-cell ma--equenccng data. BioRxiv, pp. 11 1111,, 2017.
Kangning Dong and Shihua Zhang. Deciphering spatial domains from spatially resolved transcriptomics with an adaptive graph attention auto-encoder. Nature communications, 13(1):1739, 2022.
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.096999 2020.
Zhaoyu Fang, Ruiqing Zheng, and Min Li. semae: masked autoencoder for single-cell ma-see clustering. Bioinformatics, 40(1 ::bbae00, 2024.
Meire Fortunato, Tobias Pfaff, Peter Wirnsberger, Alexander Pritzel, and Peter Battaglia. Multiscale mesh- graphnets. arXiv preprint 10000222 2022.
12

---

Published as a conference paper ICLR 2026.
Mariano I Gabitto, Kyle Travaglini, Victoria M Rachleff, Eitan S Kaplan, Brian Long, Jeanelle Ariza, Yi Ding, Joseph T Mahoney, Nick Dee, Jeff Goldy, et al. Integrated multimodal cell atlas of alzheimer's disease. Nature Neuroscience, 27(12):2366 2383, 2024.
Ziqi Gao, Chenran Jiang, Jiawen Zhang, Xiaosen Jiang, Lanqing Li, Peilin Zhao, Huanming Yang, Yong Huang, and Jia Li. Hierarchical graph learning for votein-proteinn interaction. Nature Communications, 14(1):1093, 2023.
Artur Grigorev, Michael J. Black, and Otmar Hilliges. Hood: Hierarchical graphs for generalized modelling of clothing dynamics. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 16965-16974, June 2023.
Christopher Heje Grønbech, Maximillian Fornitz Vording, Pascal N Timshel, Casper Kaae Sønderby, Tune Pers, and Ole Winther. sevae: variational auto-encoders for single-cell gene expression data. Bioinformatics, 36(16):4415 05 2020. ISSN 1367-4803. doi: 10. 1..19//boinommmaccs/ua2222....22 URL https: / //doi. ///oi.org//// www11109////noorrmmooio cst.amaa....//..a.....
Keren Bahar Halpern, Rom Shenhav, Orit Matcovitch-Natan Beáta Tóth, Doron Lemze, Matan Golan, Efi E Massasa, Shaked Baydatch, Shanie Landen, Andreas E Moor, et al. Single-cell spatial reconstruction 2017.
reveals global division of labour in the mammalian liver. Nature,
Minsheng Hao, Jing Gong. Xin Zeng, Chiming Liu, Yucheng Guo, Xingyi Cheng, Taifeng Wang, Jianzhu Ma, Le Song, and Xuegong Zhang. Large scale foundation model on single-cell transcriptomics. bioRxiv, 2023.
Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, abs/1606.0844 08415, 2016. URL httttttt//aa/tti wttp:/iiv.rrgab////aooaa 08415.
Jian Hu, Xiangjie Li, Kyle Coleman, Amelia Schroeder, Nan Ma, David Irwin, Edward B Lee, Russell T Shinohara, and Mingyao Li. Spagen: Integrating gene expression, spatial location and histology to identify spatial domains and spatially variable genes by graph convolutional network. Nature methods, 18(11):
1352-1351, 2021.
Chaitanya K Joshi. Transformers are graph neural networks. arXiv preprint arXiv:2506.2200 2025.
Smita Krishnaswamy, Matthew H Spitzer, Michael Mingueneau, Sean C Bendall, Oren Litvin, Erica Stone, Dana Pe' er, and Garry P Nolan. Conditional density- based analysis of cell signaling single-cell data. Science, 346(6213): 1250689, 2014.
Jiachen Li, Siheng Chen, Xiaoyong Pan, Ye Yuan, and Hong-Bin Shen. Cell clustering for spatial transcrip- tomics data with graph neural networks. Nature Computational Science, 2(6):399 408, 2022.
Yahui Long, Kok Siong Ang, Mengwei Li, Kian Long Kelvin Chong, Raman Sethi, Chengwei Zhong, Hang Xu, Zhiwei Ong, Karishma Sachaphibulkij. Ao Chen, et al. Spatially informed clustering, integration, and deconvolution of spatial transcriptomics with graphst. Nature Communications, 14(1):1155, 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint urXiv:111.05101, 2017.
Efthymia Papalexi and Rahul Satija. Single-cell ma sequencing explore immune cell heterogeneity. Nature Reviews Immunology, 18(1):35 45, 2018
Jason Ptacek, Matthew Vesely, David Rimm, Monirath Hav, Murat Aksoy, Ailey Crow, and Jessica Finn. 52 characterization of the tumor microenvironment in melanoma using multiplexed ion beam imaging (mibi). Journal for immuno'Therap of Cancer, 9(Sppp 2):A59-A59, 2021. doi: 10. 1 20t.-22002222222211111 /Supp__2/AS__
URL https::///// lite. bmj. com wtt.///c//ttete
13

---

Published as conference paper ICLR 2026.
Samuel G Rodriques, Robert R Stickels, Aleksandrina Goeva, Carly A Martin, Evan Murray, Charles Vanderburg, Joshua Welch, Linlin M Chen, Fei Chen, and Evan Z Macosko. Slide-seq: A scalable technology for measuring genome wide expression at high spatial resolution. Science, (63444))) 1467, 2019.
Christina V Theodoris, Ling Xiao, Anant Chopra, Mark D Chafîn, Zeina R Al Sayed, Matthew Hill, Helene Mantineo, Elizabeth M Brydon, Zexian Zeng, X Shirley Liu, et al. Transfer learning enables predictions in network biology. Nature, 61 618(7965):1166624 2023.
Alexander Tong. Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz P. San Juan, Laura Rangel, Brandon Zhu, John G. Lock, Christine L. Chaffer, and Smita Krishnaswamy. Learning transcriptional and regulatory dynamics driving cancer cell plasticity using neural ode-based optimal transport. bioRxiv, 2023. doi: URL https: http:wwbiiett www.iiiirxiv..mmmmmm org/ http://wwww.......ooo/e//aaaa///
10.1100/2023.03.28.555 29/2023. 03 7388.534644
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30,
2017.
Petar Veličković, William Fedus, William L Hamilton, Pietro Liò, Yoshua Bengio, and R Devon Hjelm. Deep graph infomax. arXiv preprint 100vv:1809.10341 2018.
Chloe Wang, Haotian Cui, Andrew Zhang, Ronald Xie, Hani Goodarzi, and Bo Wang, Segpt-spatial: Continual pretraining of single-cell foundation model for spatial transcriptomics. biorxiv, 2025.
Hongzhi Wen, Wenzhuo Tang, Xinnan Dai, Jiayuan Ding, Wei Jin, Yuying Xie, and Jiliang Tang. Cellplm: pre-training of cell language model beyond single cells. BioRxiv, pp. 2023-10, 2023a.
Hongzhi Wen, Wenzhuo Tang, Wei Jin, Jiayuan Ding, Renming Liu, Xinnan Dai, Feng Shi, Lulu Shang, Hui Liu, and Yuying Xie. Single cells are spatial tokens: Transformers for spatial transcriptomic data imputation. arXiv preprint arXiv:2302.0322 03038338......00000000003 2023b.
Zhenqin Wu, Alexandro Trevino. Eric Wu, Kyle Swanson, Honesty J Kim, H Blaize D' Angio, Ryan Preska, Gregory W Charville, Piero D Dalerba, Ann Marie Egloff, et al. Space-gm: geometric deep learning of disease-associated microenvironments from multiplex spatial protein profiles. bioRxiv, pp. 2022-05, 2022.
Yi Xiao and Dihua Yu. Tumor microenvironment therapeutic target in cancer. Pharmacology & therapeutics, 221:107753, 2021.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv: 18 10,0002, 2018.
Chengxuan Ying, Tianle Cai, Shengjic Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Adyances in neural information 2021.
processing systems, 34:28877-28888
Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovee. Hicrarchical graph representation learning with differentiable pooling. Advances in neural information processing systems, 31, 2018
Hengrui Zhang, Qitian Wu, Junchi Yan, David Wipf, and S Yu Philip. From canonical correlation analysis to self-supervised graph neural networks. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
14

---

Published as a conference paper ICLR 2026.
James Zhu, Yunguan Wang, Woo Yong Chang, Alicia Malewska, Fabiana Napolitano, Jeffrey Gahan, Nisha Unni, Min Zhao, Rongqing Yuan, Fangjiang Wu, et al. Mapping cellular interactions from spatially resolved transcriptomies data. Nature methods, 21(10):183 1842, 2024.
Yongshuo Zong, Tingyang Yu, Xuesong Wang. Yixuan Wang. Zhihang Hu, and Yu Li. const: an interpretable multi-modal contrastive learning framework for spatial transcriptomics. BioRxiv, pp. 2022-01, 2022.
15

---

Published as a conference paper ICLR 2026.
A
RELATED WORKS
This section outlines hierarchical graph neural networks and their limitations, followed by graph-based spatial transcriptomics methods and foundation models for single-cell and spatial data.
Hierarchical graph neural networks. Hierarchical graph neural networks operate graphs with multiple levels, where each level represents a distinct granularity or type of information. MSMGN (Fortunato et al., 2022), HOOD (Grigorev et al., 2023), and BSMS-GN passing across levels, but the levels are of same modality at different resolutions (c.g., coarsened versions of spatial graphs), rather than integrating graphs of distinct modalities and they are geared towards propagating signals globally. HIGH-PPI (Gao et al., 2023) introduces a hierarchical graph where one level encodes protein-protein interactions and the other encodes residue-level interactions. However, it does not include cross-level message passing, which limits the ability to combine information across biological levels. In contrast, HEIST works information from different modalities while integrating information across levels.
(Cao et al., 2023) perform cross-level message
Graph-based Spatial Transcriptomies methods. Graph-based spatial transcriptomies methods construct cell graphs based on spatial proximity and apply graph neural networks to model spatial relationships. SPAGCN (Hu et al., 2021), CCST (Li et al., 2022), STAGATE (Dong & Zhang. 2022). CONST (Zong et al., 2022), and GRAPHST (Long et al., 2023) train graph neural networks in un/self-superriseed manner by using techniques such as deep graph infomax (Veličković et al., 2018), iterative clustering (Braun et al., 2022), and auto-encodings task on gene expression (Fang et al., 2024). Unlike HEIST, these models do not generalize to other tissues without retraining and do not incorporate hierarchical modeling across genes and cells.
scRNA-see seq and Spatial Transcriptonics Foundation models. Foundation models for single-cell and spatial transcriptomics aim to pretrain general-purpose representations that can be transferred across datasets and tasks. SCFOUNDATION (Hao et al., 2023), GENEFORMER (Theodoris et al., 2023), and SCGPT (Cui et al., 2024) extend transformer architectures to model scRNA-seq data by treating gene expression profiles as sequences. However, these models assume an ordering over genes and do not explicitly capture cell- cell relationships spatial information. SCGPT-SPATIAL (Wang et al., 2025) adapts SCGPT to spatial transcriptomics. CELLPLM (Wen et al., 2023a) incorporates cell-cell interactions in pertaining and adds a gaussian mixture prior to overcome data limitations. But CELLPLM and SCGPT-SPATIAL operate over a fixed set of genes and do not model hierarchy. While transformer based models like SCGPT-SPATIA can be considered as graph neural networks over fully conneceted graphs Joshi (2025), they most biological networks are far from being fully connected, and explicitly modeling sparse, biologically plausible interactions is critical and enabled by graph-based learning. Existing foundation models thus either neglect spatial dependencies, gene co-expression structure, or both, limiting their ability to capture the context-specific nature of spatial transcriptomics. Furhtermore, their reliance in on gene embeddings hinders their generalization to proteomics data. In contrast, HEIST hierarchically models both this information and computes expressive and generalizable representations by working over gene co- expression networks.
B DATA PREPROCESSING AND GRAPH CREATION
HEIST requires construction of two graph structures: (1) spatial cell-cell graph capturing local tissue organization, and (2) cell-tye--ppecffic graph representing the gene co-expression networks, which captures gene-gene dependencies conditioned on cell types. We present the preprocessing pipeline in Figure A1, and we describe the data preprocessing steps and the detailed procedures for constructing these graphs below.
Data Preprocessing. We begin by importing raw spatial transcriptomics data, followed by standard pre- processing steps, including outlier removal, gene expression normalization, and gene filtering to retain highly variable genes. We denote the raw matrix by Xraw NxG and the normalized matrix by X, with X,, = Norm (ermm.x....)) To mitigate technical noise and dropout effects common in single-cell measurements, we apply MAGIC (Dijk al., 2017) to denoise gene expression data before all downstream computations.
A1

---

Published as a conference paper ICLR 2026.
Formally,
X MAAIIC(X..
Cell-Cell Graph Construction. To capture spatial relationships between cells, we construct a cell-cell graph based on physical proximity, Each cell has spatial coordinates = (2,, y:) E IR² Spatial locations are used compute Voronoi polygons. which define local neighborhoods. Edges the graph are assigned based on adjacency in the Voronoi diagram, ensuring that the graph structure accurately reflects tissue architecture and local microenvirroments. We write this as
(1,5) E E = Vor(p:) and Vor(p;) share a boundary.
Sell-yype-Specific Gene Co-Expression Networks Construction. For each cell type, we build gene co-expression network that captures functional dependencies between genes. If cell type annotations are provided, we use them directly; otherwise. we infer cell types via Leiden clustering the denoised gene expression data. This expressed as f Leiden(X),
After cell types are assigned, we compute pairwise mutual information (MI) between genes within each cell type using the denoised expression values. Gene pairs with MI greater than a threshold are connected by an edge in the gene co-expression network, resulting in a sparse, cell-type-specifie co-expression graph. This process captures gene-gene co-expression patterns and co-expression dependencies unique to each cellular context. For cell type to with cell index set I&
{i tx).
( ab
(1) x
(a, b) E 8(h) + MI ab Tk-
Choice of r. The threshold used for constructing the gene co-expression networks is data-driven and cell-type-specific. Specifically,
+ Tk. F4 = Ea
o = Std(( Joss).
This adaptive choice yields graphs with 20-30% edge density, which empirically provides sufficient structure to capture co-expression dynamics without over-saturating the graph. In earlier development, we experimented with constant global but this produced artifacts: some cell types became overly dense while others became nearly edge-free. This mismatch harms hierarchical message passing and leads to degraded downstream performance. We also noticed similar trends in graph based baselines such as STAGATE, GraphST, and Novae, where we had to tune the thresholds manually to control the connectivity. We have highlighted these details in Appendix Section B.
Batching the tissue. Each tissue is divided into spatially coherent blocks that contain only subset of nearby cells. Let denote the block size and [ 3 be the resulting partition with [Bm| S B. In practice, we choose the largest block size that fits within the memory budget of a single GPU; for our hardware (40GB L40), this is 256 cells per block. Larger blocks allow HEIST capture mid-range spatial interactions more effectively, while still avoiding the quadratic cost of full-tissue attention. We expect diminishing returns beyond a certain block size: once the local microenvironment and its spatial context are fully captured, adding more distant cells contributes little additional signal.
Final Output. The final preprocessing pipeline produces:
A spatial (C, P, cell-cell graph encoding local tissue structure.
Cell-type-specific gene co-expression networks (g)x (V, &,, Xx) |cl capturing co-expression function within each cell type.
A2

---

Published as a conference paper ICLR 2026.
Remove outilers
Spatial locations
Call-type spectfic Gere Co- Co-expression Network
Load data
Normalie the sassion
Voranoi polyger
Call-cell graph
Filter genes
Gene expression
MI
MAGIC
Calculato Mutual Intormation (Ml)
Yes
Call types
No
Leiden diustrring
Figure A1: HEIST preprocessing pipeline
These graph structures provide the foundation for HEIST's hierarchical learning framework, enabling integration of spatial organization and gene co-expression patterns.
Final Output. The final preprocessing pipeline produces:
A spatial &c, P) cell-cell graph encoding local tissue structure.
Cell-type-specific gene co-expression networks (V, (V, Xx) capturing co-expression function within each cell type.
These graph structures provide the foundation for HEIST's hierarchical learning framework, enabling integration of spatial organization and gene co-expression patterns.
IMPLEMENTATION DETAILS
this section, we explain the implementation details including batching, the positional encoding initialization, graph transformer architecture, and hyperparameters used in pre-training phase.
C.I POSITIONAL ENCODINGS
HEIST incorporates positional encodings (PE) at both the cell and gene levels inject spatial and co- expression structure into the learned representations. As opposed to traditional graph PEs such has Laplacian PE (Dwivedi Bresson, 2020), random walk PE (Dwivedi Bresson, 2020), node-centrality based PE (Ying et al., 2021), we make use of sinusoidal PEs. This design choice is motivated by the fact that sinusoidal PE yields expressive representations and are computationally in in-expensive as opposed to traditional graph PEs, leading efficient and expressive representations (Wen al., 2023b). Below, describe how these encodings are constructed and why they are suitable for spatial transcriptomics.
Each cell is associated with spatial coordinate (x, v) R22 corresponding its location in the tissue. To encode spatial information, we apply a two-dimensional extension of the sinusoidal positional encoding introduced in transformers (Vaswani et al., 2017):
A3

---

Published as conference paper ICLR 2026.
PE = sin 1100004/a
I
PE = cos (100004777)
y 2/ = sin 1000004/4)
2Ec,2++1d/22 = cos (100004//)
y/
where d is the dimensionality of the encoding, To calculate the gene-PE, for each cell, genes are first sorted in descending order of expression. Since most of the spatial transcriptomies are noisy, ranking done after denoising using MAGIC. The rank of each gene y reflects its relative expression within the cell:
Ranky (g) = position of in sorted list by expression of cell k.
We then apply sinusoidal encoding based on the rank values as follows:
|ann, (g) (100002/44)
(Rank( 100002304
PE 9.2i = sin
PE 9.2++1 = cos
This treats genes within cell as a soft sequence, allowing the model distinguish highly expressed genes from lowly expressed ones based on their position in the transcriptional program. Sinusoidal encodings on rank preserve relative ordering and allow model capture patterns in gene expression. Since the same gene can have different ranks across cells, the resulting embedding is context-dependent, enabling the model to learn cell-specific co-expression roles of genes. This is important for spatial transcriptomics, where the function and relevance of gene may vary depending the cell type and microenvironment.
PE. 2DSinusoidal(P).. PE Ranksinussidal(XX)) H Prepeat H S
=PE +Xx
where H and H are the input cell and gene embeddings for cell k, respectively. Since we have 2D & half of the dimensions correspond to the co-ordinate from the (æ, v) pair. To add (z, y) pair to the dimensional vector, repeat the co-ordinates as [x.... times, 2, U.... d times, v). This way we are adding two d dimensional vectors. Similarly for genes, we add the scalar value of gene expression
sinusoidal PEs, the first /
each dimension. Together, the PE helps HEIST incorporate spatial geometry and transcriptional ordering in a biologically meaningful and computation-efficient way.
C.2 INTRA-LEVEL MESSAGE PASSING
We incorporate biological inductive biases by applying global (all-to-all) attention the cell graph capture long-range spatial interactions and local attention that depends on the graph structure of gene graphs to model fine-grained co-expression patterns (Papalexi & Satija, 2018). Capturing global dependencies is crucial for detecting large-scale tissue structures, such as immune infiltration zones and tumor-stroma boundaries (Zhu al., 2024), while modeling local co-expression networks enables identification of tightly coordinated gene modules within specific niches (Halpern et al., 2017). HEIST performs intra-level message passing independently at the cell and gene levels to capture both local and global dependencies within each modality.
At the cell level, we compute:
GraphTranssommerayyerr ((()1)) "ransformerLayer(H!! + GIN A,, H(
where !!! denotes the cell embeddings at layer 1, and A, is the adjacency matrix of the spatial cell graph. The transformer layer enables long-range interactions across spatial regions, while the GIN layer aggregates
A4

---

Published as conference paper ICLR 2026.
H
MHA
LayerNorm
MLP
Repeat
LayerNorm
GIN
(1) repeety Keys
H (I
MHA with Mask
H (l+1)
H
LayerNorm
Queries and values >
MHA
H
9
Queries and values
H. (1+1)
Agg.
Keys
(1) g
MHA
H
Figure A2: HEISTLayer Block
local neighborhood information. This hybrid approach captures both global spatial context (e.g., tissue-level structure) and localized microenvironmentt (e.g., niches or boundaries). Although the attention module is fully dense within each block, the computational cost is significantly reduced because the effective sequence length becomes the block size rather than the total number of cells in the tissue. A naive full-tissue attention operation would scale as Oc²2). which is expensive for large tissues. After spatial batching, the complexity becomes Obb²), where b is the block size. In practice, b is much smaller than the total number of cells in a tissue, so this reduction makes attention computation tractable while still capturing local and mid-range spatial interactions.
At the gene level, we apply sparse attention based on the connectivity of the genes as:
(W
(Wah 3.2
Wighh-111 D (Wah b g.j
if (i, E Ek.
Vd
K
softmax,
(a 60)
W. W.hhh-n h ) 9J
9,4
JEN' (1)
where h is the embedding of gene i in cell & at layer I, and N (i) denotes the gene coexpression network neighbors of gene i in that cell. This formulation ensures attention is computed only over biologically relevant interactions, preserving co-expression sparsity while allowing expressive, cell-specific dependency modeling.
C.3 HYPERPARAMETERS
In Table A1, describe the hyperparameters for the pre-training model.
A5

---

Published as a conference paper ICLR 2026.
Table A1: Hyperparameters and their default values used the experiments.
Hyperparameter
Default Value
Positional encoding dimension
128
Hidden dimension
128
Output dimension
128
Number of HEIS HEISTLayer layers
10
Number of transformer heads
8
Batch size
256
Learning rate
0.001
Weight decay
0.003
Number of pre-training epochs Activation function Optimizer
20
GeLu (Hendrycks & Gimpel, 2016) AdamW (Loshchilov & Hutter, 2017)
D DATASETS
this section, we first briefly talk about pre-training datasets, followed by downstream tasks explained in details and downstream datasets.
D.1 PRETRAINING DATASET
The pre-training dataset of HEIST comprised 22.3 million cells from 124 tissues slices and 15 organs, including data collected through MERFISH and Xenium. The data consists of 13.3M cells from 10xGenomics datasets (https: / /www. 1 10xgenoiccs com/), ), 8.7M from Vizgen (https: /wtp://vviggenne// com/) and 360k cells from Seattle Alzheimer's Brain Atlas (https: /ttp::///portt /portlorraaaaoo www.bbraninnppp seattle-at wwweeeeaazziiiee
disease).
In contrast to previous models that incorporate lower-resolution data from Visium arrays (Cui et al., 2024), HEIST focuses exclusively on single-cell resolution datasets. This design choice reflects the fact that Visium data lacks true single-cell resolution, instead capturing transcriptomic signals aggregated over spatial spots that often contain heterogeneous mixtures of multiple cell types. Such coarse-resolution data can introduce confounding signals and limit the model's ability accurately learn cell-Ievel spatial dependencies and gene co-expression. By focusing on single-cell spatial transcriptomics, HEIST is better positioned to capture fine- fineggrained spatial organization, cell-cell interactions, and context-dependent gene co-expreesion critical for modeling tissue microenvironments and cellular heterogeneity.
D.2 DOWNSTREAM TASKS
this section, we explain the tasks in more details
Cell type annotation. This task involves assigning cells to known cell types based on their gene expression and spatial context. We evaluate this by extracting representations from frozen HEIST encoder and training only MLP head for classification. Accurate cell type annotation is essential for characterizing cellular composition across tissues and studying how cell populations contribute to tissue function and disease.
Cell clustering. Cell clustering is performed in an unsupervised manner using frozen HEIST embeddings to discover both known and novel cellular subpopulations. Unlike prior models, HEIST's spatially informed representations not only separate cells into canonical cell types but also induce subclusters that reflect the influence of local microenvironments. This enables the discovery of previously unrecognized, spatially contextualized cell types based their tissue niche and surrounding cellular interactions.
Tissue-level classification. In this task, the goal is to predict tissue-level outcomes, including response to immunotherapy, treatment outcomes, remission status, and placenta condition. Accurate classification requires
A6

---

Published as a conference paper ICLR 2026.
Table A2: Summary of spatial datasets used in this study. For each dataset, we indicate the organ of origin, number of tissue slices analyzed, total cell count, imaging technology used, and the tasks performed.
Dataset
Organ
Slices
Cells
Technology
Tissue Classification
Cell-type
Gene
Annotation
Clustering imputation
Lung Cancer
Lung
100,000
MERFISH
DFCI
Head and neck
58
125,512
CODEX
UPMC
Head and neck
308
2,164,932
CODEX
Charville
Colon
292
632,180
CODEX
7
Melanoma
Skin
54
540,000
MIBI
X
Placenta
Placenta
212
1.000.000
Xenium
X
capturing both local cellular microenvironments and broader tissue organization patterns. We first evaluate HEIST on datasets from Charville (Colon), UPMC (Neck), and DFCI (Neck), collected using CO-Detection by Indexing (CODEX) technology. We also predict immunotherapy response on a skin cancer dataset (Ptacek et al., 2021) collected through multiplex ion beam imaging (MIBI. Finally, we classify placenta condition into normal placenta, placenta accreta spectrum (PAS), and placental insufficiency using data collected with Xenium. Spatial modeling is crucial for placenta classification, as conditions like PAS involve disrupted tissue architecture and abnormal cell invasion, which gene expression alone cannot fully capture. For this task, we extract cell representations from frozen HEIST encoder, we aggregate the cell embeddings by taking a mean, and train an MLP head. We report AUC-ROC as the evaluation metric.
Gene imputation. In spatial transcriptomics, gene expression profiles are often incomplete or noisy due to technical limitations and measurement dropouts. Gene imputation aims to predict missing gene expression values, recovering biologically plausible expression patterns. HEIST performs this task by leveraging the learned gene embeddings, which capture both co-expression relationships from gene co-expression networks and spatial dependencies from their parent cells. Through cross-level message passing, the model can incorporate spatial cues to refine gene predictions, enabling more accurate reconstruction of missing data, We select the genes to be masked using stratified sampling according to gene sparsity (Avşar & Pir, 2023) and mask these genes during the fine-tuning phase. For the skin dataset, which contains 30 proteomic markers, we masked 4 genes using stratified sampling based on gene sparsity following Avsar al., corresponding to 14% masking fraction. For the placenta dataset, we masked 16 out of 200 genes ( 8% masking). This task can be performed in zero-shot fashion or by fine-tuning the HEIST decoder the downstream task. We report Pearson correlation as the evaluation metric for this task.
D.3 DOWNSTREAM DATASETS
We have curated different datasets for each downstream prediction task. To show our models generalize well we specifically used 4 different technologies across the different tasks. The dataset description can be found on Table A2.
E ADDITIONAL RESULTS
As shown in Figure 3, HEIST effectively identifies such pncroenvironmeen--peeiiii subpopulations, which existing foundation models fail to capture. Figure A3 does the same analysis on scGPT-sspatil and CelIPLM showing that these methods fail to address the microenvvronment separation.
Cell clustering. Table A3 reports normalized mutual information (NMI) scores for the cell clustering task across multiple organs and datasets. HEIST achieves the highest or comparable NMI scores in most settings, demonstrating its ability capture meaningful cellular subpopulations informed by both spatial and co-cxpression. Notably, HEIST shows substantial improvement on datasets with complex tissue structures, such as UPMC and DFCI, where modeling spatial microenvinonments is critical. While some models achieve competitive performance in simpler settings, they fail to generalize as effectively across technologies due to missing genes in their gene set.
A7

---

Published as conference paper ICLR 2026.
Cell Embeddings
Cell Embeddings
scGPT-spatial
Spotial Location
LT
HEIST (ours)
Spotial Location
L4IT
STAGATE
SpatialLacation
Annotated
GraphST
Spatial
Figure A3: Comparisons between embeddings focusing on L4-IT neurons. The figure shows the cell type embeddings, followed by spectral clustering and plot of all clusters that contain the specific cell type. A. scGPT-spatial results. B. CelIPLM results C. STAGATE results D. GraphST results E. HEIST results (same Figure 3) F. Ground Truth cell types colored in space.
Table A3: Comparison of unsupervised cell clustering performance using Normalized Mutual Information (NMI).
Organ
Brain
Colon
Neck
Neck
Lung
Model
SEA-AD
Charville
UPMC
DFCI
MERFISH Lung cancer
STAGATE
0.294 + 0.05
0.237 + 0.04
0.022 + 0.01
0.048 + 0.02
0.171 + 0.06
GraphST
0.444 + 0.02
0.224 + 0.05
0.020
0.044 0.02
0.297 + 0.05
ScFoundation
0.388 0.04
0.220
0.000
"0.04T
0.049 1 0.00
CellPLM
0.651 1 0.00
0.24
0.015
0.075 0.00
0.286 1 0.00
scGPT-spatial (aligned)
0.674 + 0.04
0.253 0.07
0.017 0.01
0.038 = 0.01
0.307 0.08
HEIST
0.691 + 0.04
0.297 + 0.03
0.043 0.01
0.14 0.04
0.274 = 0.04
Runtime comparison. Table A4 compares the runtime efficiency of HEIST with existing foundation models
for generating embeddings from a tissue sample with 19,826 cells. HEIST achieves substantial speed improvements, being 8x faster than SCGPT and 48x faster than SCFOUNDATION, while providing both cell
A8

---

Published as a conference paper ICLR 2026.
Table A4: Runtime comparison for embedding extraction from a tissue with 19826 cells, and computational complexity of each model. *CelIPLM does not compute gene representations.
Model
Runtime
Complexity
Interpretation
scFoundation
290.48
O(C*G,&+CTO
Quadratic in eell count and nonzero genes. Decoder is lincar due to Performer attention.
scGPT
49.37
O(T'd+Tah TTP)
Quadratic in cell times number of gene tokens per spot
HEIST
6.69
(dd( + E.+ +EE.E.. CT))
Quadratic in cells (global spatial transformer), linear gene edges and gene count.
CelIPLM
1.97
d(Ca))
Quadratic in cells, linear hidden size. depen- dence number of genes after aggregation.
Table A5: Ablation comparing sinusoidal PE based on spatial coordinates with traditional graph-based PEs for clinical outcome prediction.
Organ
Colon
Neck UPMC
Neck DFCI
Dataset
Charville
Task
Outcome
Recurrence
Outcome
Recurrence
Outcome
Model
Random walk PE
0.770 = 0.034
0.695 + 0.042
0.776 + 0.029
0.752 + 0.038
0.916 + 0.026
Laplacian PE
0.750 = 0.031
0.800 + 0.025
0.670 + 0.045
0.810 + 0.021
0.852 + 0.040
Sinusoidal PE
0.861 = 0.086
0.887 + 0.041
0.835 + 0.001 0.929 + 0.030
0.937 + 0.062
and gene embeddings. We also present show the coputational complexity analysis of existing models. We note that HEIST is linear in gene edges and gene count, one HEISTLayer is faster than one scGPT-spatial and scFoundation layer, however it is slower than CelIPLM. Although CelIPLM is faster, it does not compute gene representations, making it less suitable for tasks requiring joint cell and gene analysis.
PE ablation. Table A5 presents an ablation study comparing sinusoidal PE against traditional graph- based PEs-random walk and Laplacian on clinical outcome prediction task. Sinusoidal PE consistently yields higher accuracy and lower variance, particularly excelling in recurrence prediction, highlighting its effectiveness in capturing spatial patterns relevant clinical signals.
7 choice ablation. In Table A6, we compare the results of tissue classification between adaptive and constant threshold to create gene co-expression graphs. As we can see, the adaptive results consistently outperforms the constant version, capturing fine-grained co-expression structure.
Importance of contrastive terms. understand whether the three contrastive components interact positively, we carried out an ablation experiment in which we removed each term in turn: cell to cell, gene to gene, and cell to gene. As shown in Table A7, removing any single term consistently reduces performance across all datasets. Removing the cell to cell objective reduces Charville outcome prediction from 0.861 to 0.380. Removing the gene to gene objective produces the largest reduction on SEA cell classification, which drops from 0.995 to 0.610. These results show that each contrastive component captures complementary structural signals that include local cell similarity, gene program structure, and cross scale coupling between cells and genes. Using all three terms together provides the strongest supervision signal and does not create negative interactions. We also tested with a joint contrastive loss function that directly models the c-c, g-g. and c-g correlations together rather than simply adding the three contrastive losses. As we can see, the performance with this joint contrastive loss decreased slightly. The primary reason is that the joint contrastive objective significantly increased memory consumption, causing repeated OOM errors. To stabilize training, were
A9

---

Published as conference paper ICLR 2026.
Table A6: Ablation comparing constant and adaptive choice of T.
Dataset
Charville
Charville
UPMC
UPMC
DFCI
Task
Outcome
Recurrence
Outcome
Recurrence
Outcome
T
Constant 0.68
0.70
0.828
0.75
0.916
Adaptive 0.86
0.88
0.835
0.92
0.937
Table A7: Ablation of HEIST's contrastive loss components across three tasks. Removing any of the three contrastive terms degrades performance across modalities, confirming that the full joint loss is necessary.
Loss
Charville-Outcome
Skin-Imputation
SEA-Cell classification
Full contrastive Joint contrastive
0.861 = 0.086
0.807 = 0.020
0.995 = 0.015
0.808 0.040
0.775 0.017
0.994 + 0.004
No c ++ c
0.380 = 0.119
0.570 = 0.060
0.993 = 0.006
No g +
0.575 = 0.049
0.477 + 0.033
0.610 = 0.097
No g
0.559 = 0.065
0.693 + 0.057
0.928 = 0.002
forced to reduce the batch size. Smaller batches provide far less microenvironmental diversity within each step, which is essential for learning strong contrastive signals. This reduced context ultimately impaired the effectiveness of the joint contrastive loss and led the observed decrease in performance.
Attention interpretation. We analyzed the cell to cell attention scores learned by HEIST and CellPLM and visualized their top attention edges. We aggregated the attention across layers and heads to get single attention maps. We only compared with CellPLM since that the only other foundation model that calculates cell-cell attention.
As shown in Figure A4, the edges with the highest attention values HEIST form clear, spatially coherent microenvironment.. We observe groups of cells of different types connected structured patterns, often resembling local niches such. These visual patterns indicate that HEIST's attention mechanism focuses on meaningful neighborhoods rather than isolated pairs. In contrast, the top-attention edges from CellPLM are diffuse and scattered across the tissue. The highlighted edges do not show consistent spatial structure cell- cell-type organization. This suggests that CellPLM attention much noisier and does not preferentially emphasize biologically relevant areas.
To quantify whether these visual differences reflect underlying biological communication, we compared all edges to curated ligand-receptor (LR) interactions from CellPhoneDBB. Across the entire dataset, approximately 16 percent of edges correspond to known LR pairs. Randomly sampling 1000 edges also gives about 16 percent, which serves as the expected background rate. As shown in Table A8
HEIST shows strong enrichment for LR edges:
The top 100 high-attention edges contain 31 percent LR interactions.
The top 1000 edges contain 27.9 percent LR interactions.
This is better than chance which yields about 16 percent expected LR pairs.
CellPLM shows the opposite trend. Only 11 percent of its top 100 edges and 8.8 percent of its top 1000 edges correspond to LR interactions. This is clearly below the background rate, which indicates that CelIPLM attention values do not align with known biological communication.
A10

---

Published as conference paper ICLR 2026.
Top attention edges
HFIST
CellPI
Cell Types
Actrocyte
Endothenal
122 IT
LS
Cara
Lampo
Lamp5 Lhxs
Orgodendrocytee
Pix6
Set
V.NC
Figure A4: Visualizing edges with highest cell-cell attentions. We can clearly see that HEIST is able to form niches, while CelIPLM high attention scattered around the tissue.
Table A8: LR-rate measured using cell-cell attention.
Model
LR pairs in top 100
LR pairs in top 1000
Background (all edges)
0.165
0.165
Random
0.164
0.165
CelIPLM
0.110
0.088
HEIST
0.310
0.279
PE initialization. Our choice to incorporate spatial coordinates and gene-expression scalars by repeat- ing/broadcasting them into the embedding dimension follows the standard additive formulation widely used across transformers. Because the spatial coordinates (xr, y) and gene expression values low- dimensional, a direct addition to the d-dimensional positional embeddings requires map into Rd Repetition provides simple, parameter-free way match dimensionality while maintaining the relative scale and directional structure of the original signals. Importantly, we show that initialization that relies purely on an MLP ("No PE" Table In this formulation, we pass the spatial co-ordinates and gene expression into the latent space using an MLP. This leads to a marked drop in performance, indicating that the positional encodings play an essential role. In Table A10, ran additional pretraining run that combines sinusoidal PEs with an MI.P-based mapping of coordinates and expression values.
Batch size ablation. To evaluate how smaller block sizes affect performance, we ran a pre-training run with the HEIST pipeline using smaller block sizes. We see slight decrease in the results except for cell classification where the results became almost perfect. This is because batches would provide less microenvironmental cues in each step, leading to the slight drop in performance.
F LIMITATIONS
While HEIST advances the state of foundation models for spatial transeriptomics (ST), it has several lim- itations. First, the current gene co-expression network construction relies on co-expression relationships
A11

---

Published as a conference paper ICLR 2026.
Table A9: Predicting LR pairs
Model
AUC-ROC
scFuundation
0.963 + 0.004
CellPLM
0.930 + 0.006
scGPT-spatial
0.984 + 0.003
HEIST
0.995 + 0.002
Table A10: PE initialization ablation.
PES
Charville-Outoone
Skin-Imputation
SEA-Cell classification
No PE (MLP(P)) Sinusoidal PE (MLP(P) + PEd HEIST PE (PEc + P
0.523 + 0.010
0.458 0.003
0.220 0.034
0.857 + 0.045
0.538 + 0.026
0.999 + 0.001
0.861 + 0.086
0.807 + 0.020
0.995 + 0.015
using mutual information, which may not fully capture causal gene co-expression mechanisms or directional influences. This can be integrated by more sophisticated meassures like DREMI (Krishnaswamy et al., 2014). Furthermore, integrating more sophisticated gene co-expression network inference techniques could improve the biological interpretability and efficiency gene embeddings. potential direction for future work be incorporating temporal dynamics by applying techniques such Granger causality (Tong et al., 2023) over inferred pseudotime trajectories, allowing the model to capture not only static spatial and co-expression dependencies but also directional gene co-expression influence and developmental progression, which are currently not modeled in HEIST. Second, the model assumes static spatial snapshots of tissues and does not account for temporal dynamics or developmental trajectories, which are critical in understanding certain biological processes. Extending HEIST to model spatio-temporal transcriptomics data is an important direction for future work. Finally, while HEIST improves computational efficiency over prior foundation mod- els, it still requires substantial computational resources for large-scale pretraining. Despite these limitations, HEIST provides flexible and scalable foundation for modeling complex spatial and molecular interactions, and future work can address these challenges to further improve its generalization and interpretability.
A12

---

Published as a conference paper ICLR 2026.
Table A11: Batch size ablation
Batch size
Charoille-Oucooee Outcome
Skin- Imputatinniio
SEA-Cell classification
128
0.003 0.861 0.086
0.783
0.791 + 0.007 0.807 = 0.020
0.999 + 0.001
256
0.995 + 0.015
A13


Reviews:
[{"review_id": "49CSmg3zG5", "reviewer": "Reviewer_XcF2", "summary": "The paper introduces HEIST, a hierarchical graph transformer that couples spatial cell-cell graph with per-cell gene co-expression graph with  bidirectional message passing between the two. The training dataset is diverse, the model is trained on a combination of contrastive and MAE loss, and the authors show SoTA performance on a range of downstream tasks.", "questions": "Mentioned with weaknesses", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Modeling spatial cell graphs with gene co-expression networks seems well-motivated, as previous work emphasized either one\n- The inference time speed-up over other methods seems good\n- Downstream results seem promising", "weaknesses": "Weaknesses/Questions:\n- How is $\\tau$ chosen? Is it different across cell-types? How does varying $\\tau$ affect the results?\n- What is $\\sigma$ in L240?\n- It would be interesting to see if  the model is capable of identifying known inter-cell interactions. Currently biological interpretability analysis seems limited \n- While the loss seems interesting, I am slightly skeptical if the MAE + CL work together; how does perform change on tasks if you train on either? \n\nOverall, I think the paper presents a strong and well-motivated approach. I would lean towards accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZlwjSLxxI9", "reviewer": "Reviewer_i5mm", "summary": "The paper proposes a graph foundation model for spatial transcriptomics and proteinomics data. In this work, two two-level graphs, e.g., the spatial cell graph and the gene co-expression graph, are introduced into a hierarchical graph transformer foundation model for embedding learning. Experiments on various datasets are conducted to study the effectiveness of the proposed method.", "questions": "See Weaknesses", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A graph foundation model is proposed for analysing spatial transcriptomics and proteomics data.\n\n2. The method incorporates two hierarchical graph levels — the spatial cell graph and the gene co-expression graph — within a hierarchical graph transformer framework for effective embedding learning.\n\n3. Extensive experiments on multiple datasets are conducted to evaluate the effectiveness of the proposed approach.", "weaknesses": "1. The detailed descriptions of the spatial cell graph and the gene co-expression graph are insufficient and should be clarified.\n\n2. Figure 2 does not effectively illustrate the main concept or workflow of the HEIST architecture.\n\n3. Regarding the contrastive learning design, the c–c, g–g, and c–g objectives are formulated separately. It is unclear whether combining these formulations in a joint loss function might lead to a decrease or not in performance.\n\n4. Concerning the weighted sum of the contrastive and autoencoder reconstruction losses, it remains unclear how the weights are optimised or adaptively learned to achieve the best balance.\n\n5. The computational complexity and reconstruction (or restoration) efficiency of the proposed method, as well as those of all compared baselines, should be analysed and discussed in detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L9U9jHjJah", "reviewer": "Reviewer_Lg1u", "summary": "The paper introduces HEIST, a hierarchical graph foundation model for spatial omics that jointly models a spatial cell graph and cell type specific gene co-expression graphs, with bidirectional cross-level message passing. HEIST is pretrained on a large corpus (reported as 22.3M cells across 124 tissues/15 organs) and evaluated on cell clustering, cell-type annotation, gene imputation, and tissue-level clinical outcome prediction, including transfer to spatial proteomics without retraining.", "questions": "1. How does the use of MAGIC denoising before training affect model performance? Could this amplify batch effects or introduce spurious correlations?\n\n2. The model employs global attention on the cell graph. What is the typical sequence length per slide, and how is the quadratic cost of O(n^2) handled in practice (batching, sparsification, etc.)?\n\n3. For model initialization:\n- How are cell embeddings defined as \\mathrm{PE}+P given the mismatch between 128-d PE and 2-d coordinates?\n- How are gene embeddings defined as \\mathrm{PE}+X, where X is a scalar expression?\n\n4. In the clustering task, why are classical baselines such as GraphST and STAGATE excluded?\n\n5. Could clustering the embeddings from SCGPT-spatial or CellPLM yield results comparable to HEIST’s clusters in Fig. 3C?\n\n6. In the gene imputation task, what is the number of genes per dataset and the fraction masked?\n\n7. For tissue classification, how are tissue embeddings derived from cell embeddings?\n\n8. Why do some baselines (SCGPT-spatial, CellPLM) perform worse in the aligned setting than in the unaligned one (Table 2)? Isn’t alignment expected to improve performance?\n\n9. In the cell-type annotation task, how are data splits defined, and how is fine-tuning performed?\n\n10. The paper notes SCGPT-spatial’s strong performance is explained by pretraining on MERFISH lung cancer data. For the brain dataset, HEIST achieves nearly 0.99 accuracy. Could this also result from pretraining overlap with brain data?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- A clear hierarchical formulation that couples intra-cell gene programs with inter-cell spatial context via cross-level message passing.  \n- Ambitious pretraining scope (22.3M cells across multiple organs).\n- Demonstrates SOTA-level results on clustering, annotation, imputation, and outcome prediction, with ablations highlighting the contributions of hierarchy, cross-level passing, and loss design.", "weaknesses": "Although the idea is interesting, several components are not clearly motivated, and some key method and experiment settings are missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LUwfvNzbKD", "reviewer": "Reviewer_X95u", "summary": "This paper proposes \"HEIST\" a hierarchical graph transformer model for spatial omics data. This model consists of two pars: upper level (cell-cell) and lower level (gene-gene co-expression within a cell). The model was trained on constraining loss and masked autoencoding loss.", "questions": "see weakness for improvement", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Rigorous design to account for cell-cell spatial information and gene-gene co-expression signals. \n- Experiments demonstrated HEIST's superior performance in several downstream tasks.", "weaknesses": "- Contribution is vague. Computational novelty (hierarchical graph transformer) is a bit limited, as the proposed model is somewhat a standard method. Then it should have some contribution from the biological side. But biological contribution is also limited. The downstream task analysis (e.g., tissue classification, gene imputation, cell type annotation) is also very routine and limited in claiming biological contribution as a “foundation” model like scGPT. \n- Not sure HEIST is a truly foundation model in that its zero-shot performance is very low. The claim of “foundation model” might be overstating. HEIST focuses on validating its biological realism (e.g., how spatially grounded embeddings work in downstream tasks), but it fails to show generalizability to unseen tasks in a zero-shot setting. \n- Training size is large and I appreciated that, but it’s still limited in the number of organs and tissues to claim the model as “foundation model”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)