MoLF: Mixture-of-Latent-Flow
Pan-Cancer Spatial Gene Expression Prediction from Histology
TL;DR
We introduce MoLF (Mixture-of-Latent-Flow), a generative model for scalable pan-cancer histogenomics. By leveraging a Mixture-of-Experts (MoE) velocity field within a conditional Flow Matching framework, MoLF effectively decouples diverse tissue patterns. It establishes a new state-of-the-art performance, and demonstrates robust zero-shot generalization to cross-species data.
MoLF (Mixture-of-Latent-Flow) Overview
Motivation & Method
- The Cost-Scale Disparity: Spatial transcriptomics (ST) offers high-resolution molecular profiling of great clinical value, yet it remains prohibitively expensive and time-consuming. In contrast, H&E histology is ubiquitous and cost-effective, motivating the need for scalable inference directly from standard pathology slides.
- Limitations of Single-Tissue Models: Existing methods are largely restricted to single-cancer models. These approaches are limited by data scarcity and fail to leverage the underlying histogenomic features that are conserved across different malignancies.
- The Pan-Cancer Challenge: While unified modeling offers a solution, the extreme heterogeneity of diverse tissues poses a significant optimization challenge. Standard monolithic architectures struggle to reconcile conflicting morphological patterns from different organs within a single parameter space.
-
Our Solution (MoLF):
We propose MoLF, a pan-cancer framework designed to resolve this heterogeneity by employing a Mixture-of-Experts strategy.
- Stage I (Latent Compression): A Transformer-based VAE compresses high-dimensional gene expression into a structured latent manifold, reducing noise and capturing essential molecular features.
- Stage II (Generative Flow): A Conditional Flow Matching model, parameterized by a Mixture-of-Experts (MoE) velocity field, transports random noise onto this manifold, conditioned directly on H&E morphological features.
Experiments & Results
State-of-the-Art Performance
We evaluated MoLF on the HEST-1k pan-cancer benchmark. Our results show that MoLF consistently outperforms current state-of-the-art methods:
- Highly Variable Genes (HVG): Achieves highest predictive correlation across distinct cancer types.
- Hallmark Pathways: Superior performance across Low, Mid, and High variance gene tiers.
Mechanism Validation: Synthetic 8-Gaussian
To isolate inductive bias, we compared MoLF against a dense baseline on a synthetic 8-Gaussian distribution.
The Dense Baseline (Right) suffers from severe averaging artifacts due to shared parameterization.
In contrast, MoLF (Left) mitigates these artifacts by routing inputs to specialized experts, achieving significantly sharper mode separation.
MoLF (Sharper Modes) vs. Dense Baseline (Artifacts)
Geometric Analysis of Generative Transport
We first visualize the global manifold alignment. The initial Gaussian noise (red, left) is transported in a single ODE step to a structured manifold (red, right) that well aligns with the ground truth gene distribution (blue). This confirms the model successfully approximates the target distribution.
Macro-scale Manifold Alignment
At a finer granularity, tracing individual patch trajectories demonstrates micro-scale precision. The model learns specific paths transporting random noise to the correct target gene expression, effectively solving the conditional optimal transport problem.
Micro-scale Trajectory Analysis
Expert Specialization Analysis
The UMAP visualization below shows latent gene embeddings colored by cancer type (Left) and by the active Expert (Right).
Notably, experts are not segregated by cancer type (e.g., there is no single "Lung Cancer Expert"). Instead, they work collaboratively, with multiple experts contributing to different regions of the latent space across all tissues. This allows MoLF to learn shared biological motifs rather than overfitting to specific datasets. Crucially, this strategy ensures that collaborative experts adapt to new cancer types better than rigid cancer specialists, as they can flexibly recombine learned biological primitives to characterize unseen tissues.
Distributed Expert Strategy
Zero-shot Cross-species Generalization
We evaluated MoLF on a zero-shot task: training on human data and testing on the HEST1k Mouse Melanoma dataset.
The Pan-Cancer MoLF significantly outperforms the tissue-aligned baseline. This validates our core hypothesis: the diversity of the pan-cancer dataset forces the Mixture-of-Experts to decompose visual signals into fundamental biological primitives. By learning these reusable, invariant motifs, MoLF maintains predictive stability even across species barriers.
Zero-shot Cross-species Generalization