RNA foundation models enable generalizable endometriosis disease classification and stable gene-level interpretation

This study demonstrates that RNA foundation models significantly improve the generalizability of endometriosis classification across independent cohorts compared to traditional baselines, while a novel interpretability method (CA-IG) reveals stable, biologically plausible gene-level signatures that overcome the instability of standard explainability approaches.

The Gist

The Big Problem: The "Needle in a Haystack" Diagnosis

Imagine Endometriosis as a sneaky thief that hides inside a woman's body. It causes terrible pain and can make it hard to have children, but doctors often can't find it without doing major surgery (like a treasure hunt with a shovel). Currently, it takes about 9 years on average to get a diagnosis because the symptoms are vague and blood tests don't work well.

Scientists have tried to use computers (Machine Learning) to look at a patient's genetic "recipe book" (RNA data) to predict if they have the disease. But there's a catch: these computers are like students who only studied for one specific test. If you give them a slightly different test (data from a different hospital or a different group of people), they fail. They memorized the specific quirks of the first group of patients rather than learning the actual disease.

The New Solution: The "Super-Reader" Foundation Models

The researchers in this paper decided to try something new. Instead of training a computer from scratch on a small group of patients, they used Foundation Models (FMs).

The Analogy:
Think of a standard Machine Learning model as a local tour guide who knows one specific city perfectly but gets lost if you take them to a neighboring town.

A Foundation Model is like a world-traveling encyclopedia. It has already read millions of books about biology and genetics from all over the world. It understands the general language of life (how genes talk to each other) before it ever sees a single patient with endometriosis. The researchers didn't teach it anything new; they just asked it to use its existing "worldly wisdom" to help diagnose this specific disease.

The Experiment: The "Taste Test"

The team gathered data from 12 different studies (like 12 different bakeries) involving 334 women. They set up a challenge:

1. The Old Way (The Baseline): Train the computer on 11 bakeries and test it on the 12th.
   • Result: The computer got confused. It relied on specific details of the bakeries it knew (like the color of the walls) rather than the actual recipe. Its accuracy dropped significantly.
2. The New Way (The Foundation Models): Use the "World-Traveling Encyclopedia" to understand the recipes, then test it on the 12th bakery.
   • Result: The computer was much smarter. It recognized the "flavor" of the disease regardless of which bakery it came from. It improved the diagnosis accuracy from 68% to 83%.

The "Why": Seeing the Invisible

One of the biggest problems with AI in medicine is that it's a "black box." You get a diagnosis, but you don't know why.

The researchers invented a new tool called CA-IG (Classifier-Aligned Integrated Gradients).
The Analogy:
Imagine the AI is a detective solving a crime.

• Old AI: The detective points to a suspect and says, "I know he did it," but refuses to explain why.
• New AI with CA-IG: The detective points to the suspect and says, "I know he did it because he was holding a specific type of muddy shoe, and he was seen near the window."

This new tool allowed the researchers to see exactly which genes the AI was looking at.

• The Old Way: The genes it picked changed every time they tested a different group of people. It was like the detective pointing to a different suspect every day.
• The New Way: The AI pointed to the same 5 genes every single time, no matter which group of patients they tested. This proved the AI had found the real biological cause, not just random noise.

What Did They Find?

The "Super-Reader" AI identified a specific set of genes that act like a smoke alarm for the disease. These genes are related to:

• Cellular Stress: The cells are under pressure and screaming for help.
• Inflammation: The body is fighting a fire that won't go out.
• Tissue Repair: The body is trying to fix things but making a mess (scarring/fibrosis).

These findings align with what we already know about endometriosis, but the AI found them in a way that is consistent across different populations, which is a huge step forward.

The Bottom Line

This paper is like upgrading from a local map to a GPS with global satellite data.

• Before: We had computers that could guess the disease if the patient looked exactly like the ones they were trained on.
• Now: We have a system that understands the deep, universal language of biology. It can spot endometriosis in new, unseen groups of people with much higher accuracy and can tell doctors exactly which biological signals to look for.

This brings us one step closer to a simple blood test that could diagnose endometriosis quickly, saving women years of pain and uncertainty.

Technical Summary

1. Problem Statement

Endometriosis is a chronic inflammatory condition affecting approximately 10% of reproductive-age women, yet it suffers from significant diagnostic delays (averaging 9 years in the UK) due to a lack of non-invasive biomarkers and phenotypic heterogeneity. While machine learning (ML) models trained on transcriptomic data (RNA-seq) have shown promise for disease prediction, they face a critical limitation: poor generalizability. Most existing models are trained and validated on single-cohort datasets, leading to overfitting on cohort-specific biases (e.g., specific lab protocols, patient demographics) rather than learning robust, biologically meaningful disease signatures. Consequently, these models fail when applied to independent, unseen datasets, hindering clinical translation.

2. Methodology

The authors developed a comprehensive, end-to-end pipeline to evaluate whether RNA Foundation Models (FMs) can overcome these generalization challenges.

• Dataset Assembly:

  • Constructed a multi-cohort benchmark comprising 12 independent GEO studies (334 samples: 259 endometriosis cases, 75 controls).
  • Used LLM-powered metadata curation with human expert validation to ensure consistent labeling.
  • Harmonized data to a common universe of 17,518 genes.

• Feature Representation:

  • Baseline: Log-transformed Transcripts Per Million (TPM) counts.
  • Experimental: Embeddings extracted from five state-of-the-art frozen RNA Foundation Models (no fine-tuning):
    1. Geneformer (Single-cell pre-trained)
    2. scFoundation (Single-cell pre-trained)
    3. scGPT (Single-cell pre-trained)
    4. BulkRNABERT (Bulk RNA-seq pre-trained)
    5. BMFM-RNA (Bulk RNA-seq pre-trained)

• Evaluation Strategy:

  • Within-Cohort: Training and testing on samples from the same datasets (simulating standard cross-validation).
  • Cross-Cohort: Training on a set of cohorts and testing exclusively on entirely unseen cohorts (simulating real-world clinical deployment).
  • Classifier: A lightweight AdaBoost classifier trained via the AutoXAI4Omics pipeline.

• Explainability Innovation (CA-IG):

  • Standard Integrated Gradients (IG) cannot be directly applied to frozen encoders with separate downstream classifiers.
  • The authors introduced Classifier-Aligned Integrated Gradients (CA-IG). This method computes gradients through the frozen encoder along a direction defined by the downstream classifier's SHAP scores. This allows for gene-level attribution without backpropagating through the non-differentiable classifier, significantly reducing computational cost while maintaining biological interpretability.

• Biological Interpretation:

  • Used Gene Set Enrichment Analysis (GSEA) on ranked gene attributions.
  • Employed a Literature-Driven Knowledge Graph (KG) and Large Language Model (LLM) inference to synthesize relationships between top predictive genes and endometriosis pathology.

3. Key Contributions

1. First Systematic Evaluation of RNA FMs for Endometriosis: Demonstrates that pre-trained embeddings significantly outperform raw gene expression counts in cross-cohort settings.
2. Novel Explainability Method (CA-IG): A computationally efficient technique to interpret frozen foundation model embeddings, enabling stable gene-level attribution without end-to-end fine-tuning.
3. Robust Multi-Cohort Benchmark: Established a 12-cohort benchmark (334 samples) specifically designed to test domain generalization in endometriosis, a persistent challenge in computational biology.
4. Discovery of Conserved Biomarkers: Identified a set of predictive genes that remain stable across different validation regimes, contrasting with the instability of traditional baseline models.

4. Key Results

A. Classification Performance

• Within-Cohort: Both baseline (TPM) and FM embeddings performed well (Weighted F1 ~0.86–0.90). The baseline was strong, suggesting that when data distributions match, classical methods suffice.
• Cross-Cohort (Generalization):
  • Baseline (TPM): Performance dropped significantly, with a Weighted F1 of 0.68 (a drop of 0.18 from within-cohort).
  • Foundation Models: All FM embeddings showed superior robustness.
    • Geneformer achieved the best performance with a Weighted F1 of 0.83.
    • BMFM-RNA and scGPT also outperformed the baseline significantly (F1 ~0.80 and 0.78, respectively).
  • Conclusion: FM embeddings reduced the performance drop by roughly 50% compared to the baseline, proving they capture transferable disease signals rather than cohort-specific noise.

B. Stability of Gene-Level Interpretation

• Baseline (TPM): The top predictive genes varied drastically between within-cohort and cross-cohort models (only 5/20 overlap). This indicated the baseline was learning dataset-specific artifacts.
• Foundation Models (Geneformer + CA-IG):
  • High Stability: 18 out of the top 20 genes were identical between within-cohort and cross-cohort analyses.
  • Top Predictive Genes: The top 5 genes were identical and ranked similarly in both settings:
    1. DDIT3 (DNA Damage Inducible Transcript 3)
    2. LRRC3C (Leucine Rich Repeat Containing 3C)
    3. TBC1D3F (TBC1 Domain Family Member 3F)
    4. OR1J2 (Olfactory Receptor Family 1)
    5. FRG2 (FSHD Region Gene 2)

C. Biological Insights

• Pathway Enrichment: GSEA on the CA-IG ranked genes revealed enrichment in inflammatory cytokine signaling (IL-17, IL-8, IFN-gamma), cell stress (ER stress, Unfolded Protein Response), and apoptosis regulation. These align perfectly with current understanding of endometriosis pathophysiology.
• Knowledge Graph Findings: LLM analysis of the top genes suggested convergent mechanisms:
  • DDIT3 links ER stress to apoptosis and inflammation.
  • TBC1D3F is implicated in vesicle-mediated cytokine release.
  • LRRC3C and FRG2 are linked to extracellular matrix remodeling and fibrosis.
  • Collectively, these genes point to a mechanism involving stress-induced inflammation, altered cell survival, and tissue fibrosis.

5. Significance

This study provides a compelling argument for the adoption of Foundation Models in clinical transcriptomics, particularly for diseases with high heterogeneity and limited sample sizes like endometriosis.

• Clinical Impact: By improving cross-cohort generalizability, these models move closer to real-world diagnostic utility, potentially reducing the 9-year diagnostic delay for patients.
• Methodological Shift: The introduction of CA-IG solves a major bottleneck in interpreting frozen foundation models, making them accessible for gene-level biomarker discovery without the computational expense of fine-tuning.
• Biological Discovery: The identification of a conserved gene signature (centered on DDIT3 and stress-response pathways) offers new, testable hypotheses for the molecular drivers of endometriosis, moving beyond the "black box" of previous ML approaches.

In summary, the paper demonstrates that leveraging pre-trained biological knowledge via RNA Foundation Models yields more robust, generalizable, and biologically interpretable models for endometriosis classification compared to traditional gene-expression baselines.