Correlation-Guided Flow Matching with Annealed Masking for Spatial Transcriptomics Generation
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
| Dataset | PCC | HPCC | ||||||||||
| 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 |
| Metric | masking | WGCNA | STRING | Ours | |||
| Average PCC | 0.609 | 0.605 | 0.609 | 0.605 | 0.607 | 0.604 | 0.613 |
| Dataset | Masking Strategy (Learnable) | Maximum Masking Ratio | Alpha for Graph Fusion | |||||||||
| 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 |
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | PCC-100 | PCC-150 | PCC-200 | HPCC-100 | HPCC-150 | HPCC-200 | ||||||
| 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 |
| Ablation | IDC | PRAD | PAAD | SKCM | COAD | READ | CCRCC | HCC | LUNG | LYMPH | Avg. |
| 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 |
| 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 | |
| 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 | |
| 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 | |
| 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 |
| 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 |
| Dataset | Local Coefficient | Global Coefficient | ||||
| 0.1 | 0.3 | 0.5 | ||||
| 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 |
| 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 |
| 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 |
| Dataset | Pearson | HPCC | ||
| 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 |
| 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 |
| Dataset | PCC | HPCC | ||||||||||||
| 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 |
| Baseline | Mean Diff. | Wins | Losses | Raw | Holm |
| ST-Net | 0.1152 | 479 | 21 | ||
| BLEEP | 0.1935 | 497 | 3 | ||
| BLEEP-UNI | 0.0584 | 424 | 76 | ||
| TRIPLEX | 0.1189 | 475 | 25 | ||
| STFlow | 0.0129 | 330 | 170 |
| Baseline | Mean Diff. | Wins | Losses | Raw | Holm |
| ST-Net | 0.1153 | 28 | 0 | ||
| BLEEP | 0.1921 | 28 | 0 | ||
| BLEEP-UNI | 0.0598 | 27 | 1 | ||
| TRIPLEX | 0.1145 | 28 | 0 | ||
| STFlow | 0.0112 | 22 | 6 |
| 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 |
Spatial Gene Expression Prediction
Spatial gene expression prediction infers tissue-resolved molecular activity from histology and spatial context.
See all 39 papers