FLAG: Foundation model representation with Latent diffusion Alignment via Graph for spatial gene expression prediction
Abstract
Predicting spatial gene expression from routine H&E enables large-scale molecular profiling, yet current models treat this as isolated pointwise tasks, thereby overlooking essential biological structures like gene coordination and spatial distribution. To preserve these relationships, we introduce \textbf{FLAG}, a diffusion-based framework that redefines this task as structured distribution modeling. At the same time, we identify the critical \textbf{Gene Dimension Curse}, where joint modeling gene expression and their spatial interactions fail in high-dimensional spaces, and FLAG solves this challenge by integrating a spatial graph encoder for topological consistency and utilizing Gene Foundation Model (GFM) alignment for gene-gene fidelity in the generation process. To rigorously assess model performance, we propose a set of novel structural evaluation metrics, including Gene Structural Correlation (\textbf{GSC}) and Spatial Structural Correlation (\textbf{SSC}). Our experiments demonstrate that FLAG is highly competitive in traditional accuracy (PCC/MSE) while achieving significantly enhanced structural fidelity in capturing both gene-gene and gene-spatial relationships. The code is available at https://github.com/darkflash03/FLAG.
Lay Summary
Predicting gene activity directly from standard, accessible medical images allows researchers to study diseases at a massive scale. However, current AI models predict each gene's activity in isolation, ignoring the biological reality that genes work together in complex networks and form specific physical patterns in tissues. To address this, we introduce FLAG, a new AI framework that treats this task as modeling an organized biological system. By guiding the model to understand both the physical layout of the tissue and the natural relationships between different genes, our method not only matches the accuracy of existing tools but faithfully reconstructs the underlying biological structures, making large-scale gene analysis much more reliable.