🤖 machine learning

FLAG: Foundation model representation with Latent diffusion Alignment via Graph for spatial gene expression prediction

The paper introduces FLAG, a diffusion-based framework that leverages a spatial graph encoder and Gene Foundation Model alignment to overcome the "Gene Dimension Curse" and accurately predict spatial gene expression from H&E images by preserving essential gene-gene and spatial structural relationships.

Original authors: Qi Si, Penglei Wang, Yushuai Wu, Yifeng Jiao, Xuyang Liu, Xin Guo, Yuan Qi, Yuan Cheng

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Qi Si, Penglei Wang, Yushuai Wu, Yifeng Jiao, Xuyang Liu, Xin Guo, Yuan Qi, Yuan Cheng

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Predicting the Invisible from the Visible

Imagine you have a high-resolution photograph of a bustling city (a Histology Image or H&E slide). You can see the buildings, the streets, and the general layout. But you want to know the specific mood, conversations, and activities happening inside every single house (the Gene Expression).

Currently, scientists have a way to peek inside a few houses to see what's happening (Spatial Transcriptomics), but it's expensive and slow. They want to use the cheap, easy photos of the city to predict what's happening inside every house.

The problem is that current computer models treat every house as an isolated island. They guess the activity in House A without caring about House B next door, or how House A talks to House C. This leads to predictions that might be mathematically "close" but biologically messy—like a city where everyone is shouting random words that don't make sense together.

FLAG is a new AI framework designed to fix this. It doesn't just guess numbers; it tries to understand the relationships between the houses and the rules of how people interact.

The Three Main Problems FLAG Solves

1. The "Island" Problem (Ignoring Relationships)

The Analogy: Imagine trying to predict the weather in a city by looking at one single window. You miss the wind patterns, the clouds moving from one building to another, and how the whole city climate works together.
The Paper's Claim: Old models look at genes one by one. FLAG treats the tissue like a neighborhood graph. It understands that Spot A is connected to Spot B, and their gene expressions should influence each other, just like neighbors influence each other's lives.

2. The "Too Many Variables" Problem (The Gene Dimension Curse)

The Analogy: Imagine trying to solve a puzzle. If you have 10 pieces, it's easy to see how they fit together. But if you suddenly have 800 pieces, and you try to figure out how every single piece connects to every other piece all at once while the puzzle is shaking, your brain (or the computer) gets overwhelmed and crashes.
The Paper's Claim: The authors discovered that when they tried to model the connections between genes and spots simultaneously as the number of genes grew large, the AI failed. They call this the "Gene Dimension Curse." The math gets too complex, and the model collapses.
The Fix: Instead of trying to solve the whole massive puzzle at once, FLAG breaks it down. It uses a "map" of the neighborhood (the spatial graph) to guide the process, so the AI doesn't have to guess the connections from scratch.

3. The "Hallucination" Problem (Losing Biological Rules)

The Analogy: If you ask a creative writer to describe a city they've never seen, they might invent a city with flying cars and purple trees. It looks cool, but it's not real.
The Paper's Claim: AI models sometimes invent gene patterns that look mathematically correct but are biologically impossible (e.g., genes that never talk to each other suddenly start shouting in unison).
The Fix: FLAG uses a "Gene Foundation Model" (GFM) as a teacher. Think of this as a super-smart biology textbook that has read millions of genetic studies. FLAG checks its work against this textbook to make sure the relationships it invents actually exist in real biology.

How FLAG Works (The Recipe)

The paper describes FLAG as a three-part kitchen:

The Map Maker (Spatial Graph Encoder):
First, FLAG looks at the photo of the tissue and builds a map of the neighborhood. It draws lines between spots that are physically close or look similar. This map acts as a "conditioning signal"—it tells the AI, "Hey, these spots are neighbors, so their genes should be related."
The Creative Artist (Latent Diffusion):
This is the part that actually generates the gene data. It starts with a "noisy" guess (like static on a TV) and slowly cleans it up to reveal the clear picture. Because it has the "Map" from step 1, it knows to keep the neighbors' stories consistent.
The Editor (Foundation Model Alignment):
While the Artist is working, the Editor (the GFM) watches over them. If the Artist starts making up a weird relationship between two genes, the Editor says, "No, that doesn't match the rules we learned from millions of other cells." This ensures the final result is biologically faithful.

How They Measured Success

The authors didn't just check if the numbers were close; they checked if the story was right.

Old Way (Pointwise Accuracy): Did the AI guess the exact number of apples in the basket? (Measured by PCC/MSE).
New Way (Structural Fidelity):
- Gene Structural Correlation (GSC): Do the genes that usually work together still work together in the prediction? (Did the "neighbors" still talk to each other?)
- Spatial Structural Correlation (SSC): Does the pattern of gene expression look like a real tissue map, or does it look like a blurry smear? (Did the "neighborhood" feel real?)

The Results

The paper claims that FLAG is a winner because:

It is just as good at guessing the exact numbers as the best old models.
Crucially, it is much better at preserving the biological structure. It captures the "vibe" of the tissue and the relationships between genes much more accurately than previous methods.
It solves the "Gene Dimension Curse," meaning it can handle large, complex lists of genes without crashing, which previous graph-based methods could not do.

In short, FLAG is a new way for computers to predict the invisible genetic world from a simple photo, ensuring that the prediction respects both the layout of the tissue and the complex rules of biology.

Technical Summary: FLAG

Problem Statement

Spatial transcriptomics (ST) enables the analysis of gene expression within the context of tissue architecture, yet ST sequencing remains expensive and low-throughput. Conversely, routine Hematoxylin and Eosin (H&E) whole-slide images (WSIs) are ubiquitous in clinical workflows. While existing methods attempt to predict spatial gene expression from WSIs, they predominantly treat the task as a collection of isolated pointwise regression problems. These approaches optimize for metrics like Mean Squared Error (MSE) or Pearson Correlation Coefficient (PCC) on a per-gene basis, thereby overlooking critical biological structures:

Gene-Gene Relationships: Regulatory programs and co-expression networks are ignored.
Gene-Spatial Relationships: The spatial distribution and tissue architecture (e.g., spatial autocorrelation) are often lost, leading to over-smoothed predictions that lack coherent internal structure.

Furthermore, the authors identify a "Gene Dimension Curse": while joint modeling of gene expression and spatial interactions (via graph diffusion) works for small gene sets, its optimization stability deteriorates rapidly as the gene dimensionality increases, leading to training collapse in high-dimensional settings.

Methodology: The FLAG Framework

To address these challenges, the authors propose FLAG (Foundation model representation with Latent diffusion Alignment via Graph), a diffusion-based framework that redefines the task as structured distribution modeling rather than deterministic regression. The architecture consists of three core components:

1. Spatial Graph Encoder (Addressing Spatial Structure)

Instead of performing joint node-edge diffusion (which suffers from the Gene Dimension Curse), FLAG treats the tissue topology as a fixed, deterministic prior.

Graph Construction: A graph $G=(V, E)$ is constructed where nodes represent spatial spots and edges encode physical proximity and histological similarity derived from WSI features.
Encoding: A Spatial Graph Encoder aggregates neighborhood information from these fixed edges and node features (H&E patches) to produce a compact spatial context embedding ( $H_{spatial}$ ) for each spot.
Conditioning: This embedding serves as a conditioning signal for the gene generation process, effectively transforming the problem from joint diffusion to conditional generation.

2. Gene-Level Latent Diffusion (Addressing Gene-Gene Relationships)

The core generation engine is a Diffusion Transformer (DiT) operating in the gene space.

Process: The model denoises a noisy gene expression matrix $X_t$ conditioned on the spatial context $H_{spatial}$ .
Mechanism: Unlike standard regression, the diffusion process learns to approximate the high-dimensional probability manifold of gene expression, preserving intrinsic variance and joint correlations that pointwise objectives ignore.

3. Gene Foundation Model (GFM) Alignment (Addressing Biological Fidelity)

To ensure the generated gene expression respects known biological regulatory structures, especially given the limited statistical power of ST data to infer these relationships from scratch:

Alignment: The hidden states of the diffusion transformer are aligned with embeddings from a pretrained Gene Foundation Model (GFM), such as Geneformer or scGPT.
Loss Function: An alignment loss ( $L_{align}$ ) is introduced to constrain the intermediate representations of the diffusion model to match the GFM manifold, enforcing gene-gene fidelity without requiring the GFM to be part of the inference-time input.

4. Evaluation Metrics

The authors introduce two novel structural metrics to evaluate biological validity beyond pointwise accuracy:

Gene Structural Correlation (GSC): Measures the preservation of gene-gene regulatory networks by comparing the correlation matrices of predicted vs. ground truth expression.
Spatial Structural Correlation (SSC): Measures the preservation of spatial organization by comparing the Moran's I (spatial autocorrelation) statistics of predicted vs. ground truth genes.

Key Contributions

Gene Dimension Curse Analysis: The paper systematically characterizes the failure of joint node-edge diffusion in high-dimensional spaces. It provides a theoretical explanation: as gene dimensionality ( $G$ ) increases, the consistency manifold between node states and edge correlations becomes increasingly thin, requiring the network to approximate highly curved score fields that exceed its capacity, leading to optimization collapse.
The FLAG Framework: A novel architecture that decouples spatial and gene modeling. It uses a static graph encoder for spatial context and a GFM-aligned diffusion transformer for gene generation, successfully overcoming the scalability bottleneck of joint diffusion.
Structural Evaluation Protocol: The introduction of GSC and SSC metrics provides a more meaningful assessment of generative models for spatial transcriptomics, demonstrating that high pointwise accuracy does not guarantee biological structural fidelity.

Experimental Results

Experiments were conducted on three datasets from the HEST-1k benchmark: HER2ST, KIDNEY, and PRAD.

Pointwise Accuracy: FLAG achieves highly competitive performance in traditional metrics (PCC and MSE), often outperforming or matching state-of-the-art discriminative models (e.g., HisToGene, TRIPLEX) and generative baselines (Stem, STFlow).
Structural Fidelity: FLAG demonstrates significant improvements in structural metrics:
- GSC: FLAG consistently achieves the highest GSC scores, indicating superior recovery of gene regulatory networks and co-expression manifolds.
- SSC: FLAG outperforms all baselines in SSC, demonstrating that it preserves the authentic spatial heterogeneity and autocorrelation of tissue, avoiding the over-smoothing seen in regression models and the "hyper-correlated" artifacts of some other generative models.
Gene Dimension Scalability: In experiments varying gene panel sizes ( $G$ from 10 to 800), joint node-edge diffusion collapses rapidly as $G$ increases. Node-only diffusion degrades smoothly but significantly. FLAG maintains high fidelity even at $G=800$ , confirming its ability to scale to comprehensive gene panels.
Downstream Utility: In downstream tasks such as Differential Expressed Gene (DEG) identification and spatial domain clustering, FLAG shows higher overlap with ground truth markers and better clustering performance (ARI/NMI) compared to baselines, validating that structural fidelity translates to biological utility.

Significance and Claims

The paper claims that FLAG bridges the gap between precise gene expression estimation and biologically consistent distribution modeling. By harmonizing semantic guidance from pretrained foundation models with spatially-aware graph conditioning, FLAG enables the recovery of intricate spatial structures at a new level of fidelity.

The authors position FLAG as a principled framework for scalable, biologically consistent spatial transcriptomics generation. They argue that treating the task as structured distribution modeling, rather than isolated regression, is essential for downstream applications like pathway analysis and microenvironment characterization. The work suggests that leveraging foundation models and diffusion processes is critical for overcoming the data scarcity and high dimensionality inherent in spatial transcriptomics.

The paper concludes by noting limitations and future directions, including the need for inference acceleration due to the iterative nature of diffusion, the extension to 3D volumetric dependencies, and the challenge of zero-shot cross-tissue generalization.