Correlation-Guided Flow Matching with Annealed Masking for Spatial Transcriptomics Generation

Yupei Zhang; Hao Chen; Li Pan; Chao Li; Xiaohan Xing

cs.LGAug 28, 2026

Correlation-Guided Flow Matching with Annealed Masking for Spatial Transcriptomics Generation

Authors: Yupei Zhang, Hao Chen, Li Pan, Chao Li, Xiaohan Xing

Organizations: Department of Clinical Neurosciences University of Cambridge UK · Usher Institute The University of Edinburgh EH16 RUX · Department of Diagnostic Radiology National University of Singapore Singapore

Abstract

Spatial transcriptomics (ST) provides spatially resolved gene expression profiling but remains expensive, motivating the prediction of ST from histology images. Generative models have emerged as a mainstream paradigm for ST prediction due to their ability to model the conditional distribution of gene expression and capture its inherent stochasticity. However, these methods typically treat genes as independent prediction targets and overlook the intrinsic gene-gene interactions in biological systems, which limits their ability to preserve biologically meaningful co-expression patterns. We argue that gene-gene interactions, which reflect shared pathways and regulatory mechanisms, are essential for generating numerically accurate and biologically coherent ST profiles. In this paper, we propose CorrFlow, a correlation-guided flow matching framework for histology-to-ST prediction that explicitly models gene-gene dependencies through two complementary mechanisms. First, we introduce an annealed masked flow matching strategy, where subsets of genes are progressively masked following a timestep-dependent annealing schedule, encouraging the model to infer masked genes conditioned on the remaining genes and promoting joint conditional modeling beyond per-gene marginal estimation. Second, we devise a gene graph-regularized optimization scheme that integrates prior knowledge from the STRING database and data-driven co-expression estimated by WGCNA to construct a gene affinity graph, which enforces both local consistency and global smoothness in the predicted expression. Extensive experiments across 12 datasets show that CorrFlow achieves the best average PCC and HPCC among evaluated methods, leading to more biologically coherent ST predictions.

Figures & tables

Figure 1: Correlation-guided Flow Matching enables joint trajectory inference for correlated genes. Given correlated genes with source states

\{x_{0}^{i}\}_{i=1}^{3}

, existing methods such as Diffusion and Flow Matching treat each trajectory independently, failing to capture inter-gene dependencies and producing inconsistent predictions

\tilde{x}_{1}^{3}

. Our method applies an annealed masking

\mathcal{M}

and an explicit graph regularization, enforcing dependency across trajectories of genes

1

and

2

. This dependency-aware guidance steers

x_{1}^{3}

toward a coherent joint prediction consistent with the target distribution.

Figure 2: Overview of the proposed CorrFlow. We model histology-to-ST inference as conditional flow matching conditioned on histology features and spatial coordinates. To promote gene-dependent generation, we introduce (a) an annealed masking schedule conditioned on timestep t and (b) a gene graph-regularized optimization strategy, which constructs gene affinity by fusing an external functional prior with data-driven co-expression, and imposes a Huber-based local smoothness and a Laplacian global smoothness to improve the biological coherence of generated expression.

Dataset	PCC $\uparrow$						HPCC $\uparrow$
Dataset	ST-Net	BLEEP	BLEEP-UNI	TRIPLEX	STFlow	Ours	ST-Net	BLEEP	BLEEP-UNI	TRIPLEX	STFlow	Ours
IDC	0.715 .096	0.608 .227	0.767 .052	0.753 .059	0.790 .060	0.791 .061	0.710 .103	0.605 .225	0.769 .056	0.751 .062	0.787 .066	0.788 .072
PRAD	0.406 .008	0.345 .038	0.451 .014	0.412 .042	0.494 .029	0.495 .020	0.507 .044	0.420 .001	0.582 .017	0.501 .033	0.599 .054	0.610 .042
PAAD	0.504 .042	0.465 .063	0.543 .052	0.516 .056	0.574 .050	0.580 .052	0.535 .060	0.503 .077	0.599 .053	0.567 .063	0.609 .051	0.622 .053
SKCM	0.706 .053	0.645 .088	0.758 .011	0.748 .002	0.813 .030	0.817 .025	0.691 .047	0.648 .007	0.660 .069	0.663 .063	0.707 .126	0.732 .123
COAD	0.567 .030	0.254 .067	0.648 .031	0.579 .035	0.649 .055	0.665 .041	0.423 .002	0.165 .033	0.553 .091	0.535 .027	0.574 .060	0.598 .028

Table 1: Comparison across the 10 HEST-1K benchmark datasets. PCC and HPCC on HMHVG-50 genes are reported. Values are shown as

\mathrm{mean}_{\mathrm{std}}

across folds/seeds.

Metric	$w/o$ masking	$w/o$ $\mathcal{L}_{\rm graph}$	$w/o$ $\mathcal{L}_{\rm global}$	$w/o$ $\mathcal{L}_{\rm local}$	$w/o$ WGCNA	$w/o$ STRING	Ours
Average PCC	0.609	0.605	0.609	0.605	0.607	0.604	0.613

Table 2: Ablation study of our methods on the HMHVG-50. Average PCC across the 10 HEST-1K benchmark datasets is reported; per-dataset results are provided in Appendix Table 5 .

Dataset	Masking Strategy (Learnable)			Maximum Masking Ratio					Alpha for Graph Fusion
Dataset	$f(t)=0.75$	$f(t)=1-t$	$f(t)=t$	$p_{max}=0.15$	$p_{max}=0.3$	$p_{max}=0.5$	$p_{max}=0.75$	$p_{max}=0.9$	$\alpha=0.4$	$\alpha=0.5$	$\alpha=0.6$	$\alpha=0.7$
IDC	0.791 .062	0.792 .061	0.791 .061	0.795 .058	0.788 .061	0.791 .050	0.791 .061	0.787 .059	0.791 .061	0.791 .058	0.791 .061	0.790 .054
PRAD	0.499 .021	0.494 .020	0.495 .020	0.490 .023	0.493 .016	0.489 .007	0.495 .020	0.494 .021	0.499 .024	0.494 .021	0.495 .020	0.492 .028
PAAD	0.580 .050	0.579 .044	0.580 .052	0.577 .053	0.582 .052	0.570 .052	0.580 .052	0.564 .058	0.569 .055	0.576 .050	0.580 .052	0.582 .052
SKCM	0.817 .026	0.808 .035	0.817 .025	0.818 .017	0.807 .034	0.818 .026	0.817 .025	0.816 .025	0.816 .025	0.817 .025	0.817 .025	0.813 .029
COAD	0.666 .041	0.666 .041	0.665 .041	0.662 .043	0.653 .055	0.664 .040	0.665 .041	0.665 .040	0.665 .040	0.665 .041	0.665 .041	0.670 .046

Table 3: Hyperparameter study on masking strategy, maximum masking ratio, and alpha for graph fusion (PCC on HMHVG-50).

Figure 3: Visualization on LYMPH_IDC. (a) Visualization of a biologically relevant gene KRT7 on sample NCBI68, with histology image, ground-truth ST, and predictions of CorrFlow and STFlow (individual min-max color scales to reflect expression intensity). (b) Gene PCC comparison between CorrFlow and STFlow (each point denotes a gene and points above the diagonal indicate improved prediction by CorrFlow). (c) PCC improvements on the top-8 genes by average absolute correlation.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	PCC-100 $\uparrow$		PCC-150 $\uparrow$		PCC-200 $\uparrow$		HPCC-100 $\uparrow$		HPCC-150 $\uparrow$		HPCC-200 $\uparrow$
Dataset	STFlow	Ours	STFlow	Ours	STFlow	Ours	STFlow	Ours	STFlow	Ours	STFlow	Ours
IDC	0.772 .050	0.778 .052	0.735 .048	0.741 .053	0.707 .052	0.705 .055	0.777 .040	0.786 .042	0.756 .039	0.760 .047	0.722 .048	0.719 .052
PRAD	0.480 .026	0.475 .005	0.475 .028	0.479 .008	0.473 .024	0.473 .007	0.497 .030	0.484 .004	0.484 .025	0.487 .001	0.470 .019	0.466 .001
PAAD	–	–	–	–	–	–	–	–	–	–	–	–
SKCM	0.793 .020	0.805 .014	0.801 .003	0.810 .002	0.788 .009	0.789 .001	0.809 .035	0.820 .031	0.812 .004	0.822 .005	0.797 .010	0.799 .004
COAD	0.602 .061	0.604 .053	0.580 .056	0.581 .064	0.547 .063	0.557 .049	0.615 .063	0.607 .051	0.570 .075	0.570 .077	0.553 .081	0.560 .065

Appendix

Table 4: Comparison between STFlow and ours on HEST-1K benchmarks under HMHVG-100/150/200 settings, with PCC and HPCC reported.

Ablation	IDC	PRAD	PAAD	SKCM	COAD	READ	CCRCC	HCC	LUNG	LYMPH	Avg.
$w/o$ masking	0.790 .063	0.490 .029	0.579 .049	0.813 .021	0.666 .048	0.434 .057	0.461 .055	0.428 .122	0.780 .008	0.645 .123	0.609
$w/o$ $\mathcal{L}_{\rm graph}$	0.794 .059	0.496 .033	0.577 .049	0.818 .029	0.647 .050	0.451 .038	0.460 .049	0.398 .098	0.781 .008	0.632 .125	0.605
$w/o$ $\mathcal{L}_{\rm Local}$	0.789 .054	0.497 .021	0.574 .056	0.811 .026	0.649 .048	0.451 .038	0.462 .048	0.408 .107	0.780 .009	0.633 .128	0.605
$w/o$ $\mathcal{L}_{\rm Global}$	0.792 .060	0.497 .022	0.573 .058	0.815 .026	0.663 .042	0.452 .019	0.459 .058	0.430 .131	0.782 .007	0.631 .133	0.609
$w/o$ WGCNA	0.792 .056	0.491 .023	0.576 .051	0.805 .034	0.650 .053	0.448 .030	0.459 .052	0.436 .128	0.783 .006	0.630 .123	0.607
$w/o$ STRING	0.788 .063	0.502 .020	0.575 .050	0.809 .037	0.651 .054	0.425 .045	0.456 .061	0.430 .132	0.781 .006	0.626 .127	0.604

Appendix

Table 5: Ablation study of different graph priors and module design (PCC on HMHVG-50).

Dataset	Local Coefficient $\mathbf{\rho}$			Global Coefficient $\mathbf{\lambda}$
Dataset	0.1	0.3	0.5	$\mathbf{10^{-4}}$	$\mathbf{10^{-3}}$	$\mathbf{10^{-2}}$
IDC	0.792 .058	0.791 .061	0.793 .058	0.794 .057	0.791 .061	0.789 .063
PRAD	0.499 .024	0.495 .020	0.495 .019	0.494 .020	0.495 .020	0.497 .014
PAAD	0.575 .056	0.580 .052	0.566 .055	0.565 .050	0.580 .052	0.553 .066
SKCM	0.825 .019	0.817 .025	0.815 .031	0.814 .026	0.817 .025	0.821 .023
COAD	0.651 .059	0.665 .041	0.665 .048	0.664 .042	0.665 .041	0.656 .039

Appendix

Table 6: Hyperparameter study of the Local and Global regularization coefficient on HMHVG-50 genes. PCC across the 10 HEST-1K benchmark datasets is reported as

\mathrm{mean}_{\mathrm{std}}

, where the subscript denotes the standard deviation across folds/seeds.

Figure 4: Visualization of Marker Genes on HEST-1K benchmark.

Dataset	ST-Net	BLEEP	BLEEP-UNI	TRIPLEX	Stem	STFlow	Ours
IDC	0.742 .055	0.691 .088	0.774 .068	0.781 .059	0.711 .062	0.793 .064	0.796 .054
PRAD	0.851 .019	0.843 .003	0.815 .085	0.836 .046	0.828 .009	0.802 .063	0.778 .075
PAAD	0.545 .053	0.498 .027	0.553 .051	0.581 .125	-0.025 .012	0.685 .097	0.624 .102
SKCM	0.625 .015	0.696 .044	0.761 .094	0.725 .120	0.001 .003	0.749 .110	0.759 .104
COAD	0.664 .049	0.539 .001	0.694 .080	0.666 .028	0.088 .135	0.673 .063	0.687 .044
READ	0.830 .015	0.748 .102	0.645 .213	0.829 .030	0.011 .012	0.867 .043	0.876 .049

Appendix

Table 7: Comparison across the 10 HEST-1K benchmark datasets. GGC-Pearson on HMHVG-50 genes is reported. Values are shown as

\mathrm{mean}_{\mathrm{std}}

, where the subscript denotes the standard deviation across folds/seeds.

Dataset	ST-Net	BLEEP	BLEEP-UNI	TRIPLEX	STFlow	Ours
IDC	2.313 3.245	2.763 3.385	2.260 3.366	2.108 3.275	2.288 3.182	2.013 2.799
PRAD	1.551 0.729	2.020 0.922	1.696 0.736	1.595 0.817	1.452 0.865	1.409 0.713
PAAD	1.332 0.875	1.401 0.897	1.222 0.770	1.279 0.847	1.123 0.724	1.141 0.774
SKCM	2.894 2.187	3.311 1.838	2.565 1.805	2.853 1.805	2.206 1.765	2.182 1.620
COAD	3.436 2.347	5.370 3.833	3.006 2.405	3.545 2.824	3.184 2.680	3.061 2.341
READ	2.679 2.010	2.457 2.053	2.010 1.651	2.673 1.960	1.564 1.095	1.611 1.143

Appendix

Table 8: Comparison across the 10 HEST-1K benchmark datasets. MSE-Mean on HMHVG-50 genes is reported. Values are shown as

\mathrm{mean}_{\mathrm{std}}

, where the subscript denotes the standard deviation across folds/seeds.

Dataset	Pearson $\uparrow$		HPCC $\uparrow$
Dataset	STFlow	Ours	STFlow	Ours
Brain	0.546 .264	0.558 .245	–	–
Breast	0.513 .173	0.525 .160	0.464 .093	0.473 .081
Avg.	0.530	0.542	0.464	0.473

Appendix

Table 9: Supplementary evaluation on the two largest human cancer categories from STImage-1K4M. PCC and HPCC are reported as

\mathrm{mean}_{\mathrm{std}}

Dataset	STFlow	Graph-Structured Masking	Ours
HEST-1K benchmarks
IDC	0.790 .060	0.797 .056	0.791 .061
PRAD	0.494 .029	0.497 .019	0.495 .020
PAAD	0.574 .050	0.577 .049	0.580 .052
SKCM	0.813 .030	0.805 .030	0.817 .025
COAD	0.649 .055	0.664 .041	0.665 .041

Appendix

Table 10: Pearson correlation comparison between STFlow, a Graph-Structured Masking variant of our method, and our main model. Results are reported on the 10 HEST-1K benchmarks and the two supplementary STImage-1K4M cancer categories. Values are shown as

\mathrm{mean}_{\mathrm{std}}

, where the subscript denotes the standard deviation across folds/seeds.

Dataset	PCC $\uparrow$							HPCC $\uparrow$
Dataset	ST-Net	BLEEP	BLEEP-UNI	TRIPLEX	Stem	STFlow	Ours	ST-Net	BLEEP	BLEEP-UNI	TRIPLEX	Stem	STFlow	Ours
IDC	0.715 .096	0.608 .227	0.767 .052	0.753 .059	0.612 .056	0.790 .060	0.791 .061	0.710 .103	0.605 .225	0.769 .056	0.751 .062	0.587 .056	0.787 .066	0.788 .072
PRAD	0.406 .008	0.345 .038	0.451 .014	0.412 .042	0.227 .038	0.494 .029	0.495 .020	0.507 .044	0.420 .001	0.582 .017	0.501 .033	0.256 .090	0.599 .054	0.610 .042
PAAD	0.504 .042	0.465 .063	0.543 .052	0.516 .056	0.006 .004	0.574 .050	0.580 .052	0.535 .060	0.503 .077	0.599 .053	0.567 .063	0.009 .006	0.609 .051	0.622 .053
SKCM	0.706 .053	0.645 .088	0.758 .011	0.748 .002	0.003 .001	0.813 .030	0.817 .025	0.691 .047	0.648 .007	0.660 .069	0.663 .063	-0.008 .017	0.707 .126	0.732 .123
COAD	0.567 .030	0.254 .067	0.648 .031	0.579 .035	0.048 .042	0.649 .055	0.665 .041	0.423 .002	0.165 .033	0.553 .091	0.535 .027	0.054 .047	0.574 .060	0.598 .028

Appendix

Table 11: Comparison across the 10 HEST-1K benchmark datasets. PCC and HPCC on the HMHVG-50 are reported. Values are reported as

\mathrm{mean}_{\mathrm{std}}

, where the subscript denotes the standard deviation across folds/seeds.

Baseline	Mean Diff.	Wins	Losses	Raw $p$	Holm $p$
ST-Net	0.1152	479	21	$4.16\times 10^{-79}$	$1.66\times 10^{-78}$
BLEEP	0.1935	497	3	$2.33\times 10^{-83}$	$1.17\times 10^{-82}$
BLEEP-UNI	0.0584	424	76	$9.15\times 10^{-55}$	$1.83\times 10^{-54}$
TRIPLEX	0.1189	475	25	$1.99\times 10^{-78}$	$5.98\times 10^{-78}$
STFlow	0.0129	330	170	$1.49\times 10^{-16}$	$1.49\times 10^{-16}$

Appendix

Table 12: Gene-level PCC significance analysis on HEST-1K. For each dataset, gene-wise PCC values were first averaged across cross-validation folds, yielding 50 gene-level PCC values per dataset. These values were then pooled across all 10 datasets, producing 500 dataset-gene pairs for each pairwise comparison. We compared CorrFlow against each baseline using paired Wilcoxon signed-rank tests, followed by Holm correction. Mean Diff. denotes CorrFlow minus the baseline.

Baseline	Mean Diff.	Wins	Losses	Raw $p$	Holm $p$
ST-Net	0.1153	28	0	$7.45\times 10^{-9}$	$3.73\times 10^{-8}$
BLEEP	0.1921	28	0	$7.45\times 10^{-9}$	$3.73\times 10^{-8}$
BLEEP-UNI	0.0598	27	1	$1.49\times 10^{-8}$	$3.73\times 10^{-8}$
TRIPLEX	0.1145	28	0	$7.45\times 10^{-9}$	$3.73\times 10^{-8}$
STFlow	0.0112	22	6	$9.32\times 10^{-5}$	$9.32\times 10^{-5}$

Appendix

Table 13: Pooled Hallmark-Gene PCC significance analysis on HEST-1K. For each dataset, pathway-level Hallmark-Gene PCC values were first aggregated across cross-validation folds by matching pathway names and averaging the corresponding gene-level pathway PCC values when a pathway was present in multiple folds. The resulting dataset-pathway pairs were then pooled across all datasets and compared between CorrFlow and each baseline using paired Wilcoxon signed-rank tests, followed by Holm correction. Mean Diff. denotes CorrFlow minus the baseline.

Method	#Params (M)	Model Size (MB)	Avg. Time / Epoch (min)
ST-Net	11.202	42.73	0.301
BLEEP	24.178	92.23	0.922
BLEEP-UNI	38.205	145.74	8.342
TRIPLEX	53.908	205.64	1.433
Stem	35.581	135.73	0.230
STFlow	1.148	4.38	0.038

Appendix

Table 14: Model complexity comparison on the HEST benchmark. #Params counts the trainable parameters used under the current benchmark training setup, excluding external frozen or precomputed feature extractors. Avg. Time / Epoch is averaged over all dataset-fold runs of each method.

Spatial Gene Expression Prediction

Spatial gene expression prediction infers tissue-resolved molecular activity from histology and spatial context.

See all 39 papers

DriftST: One-Step Generative Inference of Spatial Transcriptomics from H&E Histology
Jul 6, 2026Yuhang Yang, Yonggan Bu, Shengyuan Zhou +2Computational Pathology Spatial Transcriptomics
University of Science and Technology of China, Hefei, China · Peking University Cancer Hospital, Beijing, China
Contrastive and Adaptive Multi-modal Masked Autoencoder for Spatial Transcriptomics
Jun 19, 2026Joohyeok Kim, Taejin Jeong, Jinyeong Kim +1Computational Pathology Spatial Transcriptomics
Department of Artificial Intelligence, Yonsei University, Seoul, Republic of Korea · Department of Computer Science, Yonsei University, Seoul, Republic of Korea · Yonsei University College of Medicine, Seoul, Republic of Korea
Adapting a Pre-trained Single-Cell Foundation Model to Spatial Gene Expression Generation from Histology Images
Mar 20, 2026Donghai Fang, Yongheng Li, Zhen Wang +2Gene Expression Prediction from Histopathology Images Spatial Transcriptomics
Sun Yat-sen University, China · Yunnan University, China · Chongqing University, China

Jul 6, 2026J/K move · Enter open · S save

DriftST: One-Step Generative Inference of Spatial Transcriptomics from H&E Histology

Yuhang Yang, Yonggan Bu, Shengyuan Zhou, Yiming Luo+1

University of Science and Technology of China, Hefei, China · Peking University Cancer Hospital, Beijing, China

Computational Pathology Spatial Transcriptomics One-Step Generative Modeling Spatial Gene Expression Prediction Gene Expression Prediction from Histopathology Images

Open paper AskPDF arXiv

Spatial Transcriptomics (ST) measures gene expression while preserving spatial context, but its high cost and low throughput leave public datasets small. Inferring expression directly from widely available Hematoxylin and Eosin (H&E) stained histology offers a cost-effective alternative. However, existing approaches face several limitations: regression methods over-smooth toward the conditional mean, while generative methods are faithful but require slow multi-step inference; most methods treat genes as independent and equally important, ignoring inter-gene dependencies and heterogeneous gene informativeness; and most are tailored to a single resolution, either spot-level or cell-level. To address these issues, we propose DriftST, a unified framework for inferring spatially resolved gene expression from H&E images. DriftST builds on a Cellular Drifting generative model that learns a direct drift from a histology-conditioned source to the expression distribution, retaining generative expressiveness while enabling efficient one-step generation. To capture gene structure, we introduce the STransformer, which combines a co-expression attention module for inter-gene dependencies with a gene residual gate for differential gene importance. Operating on a generic gene-panel representation, DriftST applies directly to both spot-level and cell-level data in one framework, and extensive experiments across diverse tissues and platforms show that it achieves state-of-the-art performance at both resolutions.

Abstract

Spatial Gene Expression Prediction

Explore similar work

DriftST: One-Step Generative Inference of Spatial Transcriptomics from H&E Histology

Contrastive and Adaptive Multi-modal Masked Autoencoder for Spatial Transcriptomics

Adapting a Pre-trained Single-Cell Foundation Model to Spatial Gene Expression Generation from Histology Images