GREmLN: A Cellular Graph Structure Aware Transcriptomics Foundation Model

GREmLN is a novel foundation model that integrates graph signal processing into its attention mechanism to leverage molecular interaction networks for capturing complex regulatory dependencies in single-cell transcriptomics, thereby outperforming state-of-the-art baselines in cell type annotation, graph understanding, and perturbation prediction while offering a more parameter-efficient and interpretable framework.

The Gist

Imagine you are trying to understand a massive, chaotic library where every book represents a single cell in the human body. Inside each book, there are thousands of pages (genes) that tell the story of what that cell is doing.

The problem with previous attempts to read these books using AI (specifically, "Foundation Models" like the ones that power chatbots) is that they treated the pages like a story with a strict order: Page 1, then Page 2, then Page 3.

But in biology, genes don't have an order. Gene A doesn't always come before Gene B. They are more like a giant, tangled web of friends talking to each other. If you force them into a line, the AI gets confused and misses the big picture.

Enter GREmLN (pronounced "Gremlin," but don't worry, it's a helpful one!).

The Core Idea: The "Social Network" of Genes

Think of a cell not as a list of words, but as a social network.

• Genes are people.
• Gene expression (how active a gene is) is how loudly that person is speaking.
• The Graph is the map of who talks to whom. Some genes are best friends (they regulate each other), some are distant acquaintances, and some never speak.

Previous AI models tried to read this social network by forcing everyone into a queue. GREmLN is different. It looks at the map of connections (the graph) and uses that map to understand the conversation.

How GREmLN Works (The Magic Trick)

The paper introduces a clever trick called "Graph Diffusion Kernel Attention." Here is a simple analogy:

Imagine you are in a crowded room (the cell), and you want to know what's happening.

• Old AI (Transformers): You shout a question, and everyone answers based on how close they are standing to you in a line. If someone is at the back of the line, you might not hear them well, even if they are your best friend.
• GREmLN: Instead of a line, you have a ripple effect. You drop a stone in a pond (your query). The ripples spread out across the water, but the water isn't flat; it has channels and currents (the gene network). The ripples travel faster along the paths where your friends are connected.

This allows the AI to instantly "feel" the influence of a gene that is far away in the list but is a close friend in the network. It understands that even if Gene X and Gene Y are far apart in the list, they are neighbors in the social network, so they must be related.

Why This Matters: The Results

The authors tested GREmLN against other top-tier AI models (like scGPT and Geneformer) and found it was the clear winner in three areas:

1. Identifying Cell Types (The "Who Am I?" Test):
   If you show the AI a cell it has never seen before, can it guess if it's a liver cell, a brain cell, or an immune cell? GREmLN was incredibly accurate, even with cells it had never met before. It's like a detective who can identify a criminal just by their social circle, even if they've never seen the criminal's face.

2. Understanding the Network (The "Friendship Map" Test):
   The AI was asked to guess missing connections in the gene network. GREmLN was much better at predicting who is friends with whom, proving it actually learned the "rules of the game" of biology, not just memorized data.

3. Predicting Drug Effects (The "What If?" Test):
   If you poke a cell with a specific drug (a perturbation), how will it react? GREmLN could predict the outcome better than the competition. This is huge for medicine because it means we could simulate how a drug will work on a patient's cells before actually giving them the drug.

The Best Part: It's Efficient

Usually, to make AI smarter, you have to make it bigger (more parameters, more computing power). GREmLN is surprisingly small—about one-third the size of its competitors—yet it performs better.

Why? Because it doesn't need to guess the rules; it was given the map. By baking the biological "social network" directly into its brain, it doesn't have to waste energy learning that Gene A and Gene B are friends. It just knows.

Summary

GREmLN is a new kind of AI for biology that stops treating genes like a list of words and starts treating them like a social network. By understanding who talks to whom, it can read the "language of life" much better, faster, and with less computing power than ever before. It's a step toward truly understanding how our bodies work and how to fix them when they break.

Technical Summary

1. Problem Statement

The paper addresses a fundamental limitation in applying standard Transformer-based foundation models to single-cell RNA sequencing (scRNA-seq) data:

• Lack of Sequential Structure: Unlike natural language (where words have a fixed order) or protein sequences, scRNA-seq data consists of unordered sets of gene expression values. Standard Transformers rely on positional encodings to model dependencies, but arbitrary gene ordering in scRNA-seq introduces noise and fails to capture true biological relationships.
• Inadequate Modeling of Long-Range Dependencies: Existing models often treat genes as discrete tokens with standard self-attention or add simple biases based on gene-gene relationships. These approaches often overlook the complex, non-local, and causal dependencies inherent in gene regulatory networks (GRNs) and protein-protein interaction (PPI) networks.
• Generalization Issues: Current models struggle to generalize to unseen cell types or regulatory structures because they lack an inductive bias that reflects the underlying biological logic of cellular states.

2. Methodology: GREmLN Architecture

The authors propose GREmLN (Gene Regulatory Embedding-based Large Neural model), a foundation model that integrates graph signal processing directly into the Transformer's attention mechanism.

A. Tokenization & Input Representation

• Dual Embeddings: Instead of treating genes as simple tokens, the model constructs two embeddings for each cell:
  1. Gene Identity Embedding ($E_g$): Learnable embeddings for gene IDs.
  2. Gene Rank Embedding ($E_r$): Instead of raw counts, expression values are normalized, log-transformed, and discretized into bins (ranks). This preserves relative expression levels while handling the continuous nature of RNA data.
• Input: The final input is the concatenation of identity and rank embeddings, prepended with a <CLS> token for global cell representation.

B. Graph Diffusion Kernel Attention (GDKA)

This is the core innovation. Rather than using static positional encodings, GREmLN dynamically conditions the attention mechanism on the topology of a biological interaction graph (e.g., GRN or PPI).

• Graph Construction: A normalized Laplacian matrix ($L$) is computed from the adjacency matrix ($A$) of the gene interaction graph.
• Diffusion Kernel: A spectral filter (diffusion kernel) is applied to the Laplacian to create a Kernel Gram Matrix ($\Phi_L$). This matrix represents a diffusion process over the graph, effectively smoothing information across multi-hop neighbors.
  • $\Phi_L = U \exp(-\beta \Lambda) U^\top$, where $U$ and $\Lambda$ are the eigenvectors and eigenvalues of $L$.
• Attention Mechanism: The query vectors ($Q$) are transformed by $\Phi_L$ before computing attention scores:
  • $Attn(Q, K, V, L) = \text{softmax}(\frac{\Phi_L(Q)K^\top}{\sqrt{d}})V$
  • Effect: This biases the attention mechanism to prioritize interactions that respect the graph structure (low-frequency, long-range dependencies) while allowing keys ($K$) and values ($V$) to retain high-frequency details. It effectively injects a "soft" inductive bias without hard-masking information flow.

C. Scalability via Chebyshev Approximation

Computing spectral decomposition and matrix exponentials for large graphs is computationally expensive.

• Approximation: The authors approximate the kernel Gram matrix using Chebyshev polynomials truncated at order $K$.
• Efficiency: This allows the model to compute the transformed query embeddings efficiently ($O(K \cdot G \cdot \delta \cdot d)$, where $\delta$ is the average degree) without full eigendecomposition, making it scalable to large gene sets and networks.

D. Training Objective

• Masked Modeling: The model is pre-trained using a masked gene prediction objective.
• Graph Conditioning: The prediction of masked gene expression bins is conditioned on the unmasked genes, the gene identity embeddings, and the specific cell-type's normalized Laplacian matrix ($L_c$).

3. Key Contributions

1. Graph-Aware Attention: Introduced a novel attention mechanism that embeds graph structure directly into the query transformation via diffusion kernels, solving the "orderless" problem of scRNA-seq data.
2. Biologically Informed Inductive Bias: The model leverages known biological networks (GRNs/PPIs) to guide representation learning, enabling the capture of long-range regulatory dependencies that standard Transformers miss.
3. Parameter Efficiency: Despite outperforming larger baselines, GREmLN is highly parameter-efficient (approx. 10.3M parameters, <1/3 of baselines), demonstrating that structural priors are more valuable than sheer scale.
4. Unified Framework: The architecture is agnostic to the specific graph type, supporting GRNs, PPIs, and other molecular interaction networks.

4. Experimental Results

The model was evaluated against state-of-the-art baselines (scGPT, Geneformer, scFoundation) on multiple tasks:

• Cell Type Annotation:
  • GREmLN achieved superior performance (Macro F1 $\approx$ 0.94) on human immune cell datasets.
  • It demonstrated strong zero-shot capabilities on held-out non-immune cell types, outperforming all baselines significantly.
• Graph Structure Understanding (Edge Prediction):
  • The model was tested on recovering masked edges in unseen gene regulatory networks.
  • GREmLN achieved the highest AUROC and Average Precision, proving it learns the underlying regulatory topology rather than just expression patterns.
• Reverse Perturbation Prediction:
  • On Perturb-Seq data (predicting the perturbation from the expression profile), GREmLN achieved state-of-the-art accuracy (0.475) and F1 scores after fine-tuning.
  • Ablation Study: Removing the graph component (Vanilla Transformer) caused a massive drop in performance (e.g., F1 dropped from 0.939 to 0.816 on zero-shot annotation), confirming the critical role of the graph prior.
• Scaling Behavior:
  • Performance improved monotonically as model depth increased (1, 3, and 6 layers), indicating that GREmLN benefits from scaling.

5. Significance and Impact

• Paradigm Shift: GREmLN moves beyond treating genomics data as "text" and instead treats it as a graph-structured signal, aligning machine learning architectures with biological reality.
• Interpretability: By relying on known interaction networks, the model's attention mechanisms are more interpretable, potentially revealing core regulatory modules and causal pathways.
• Efficiency: It proves that incorporating domain-specific structural priors allows for smaller, faster, and more accurate models compared to massive, data-hungry Transformers that ignore biological structure.
• Future Applications: The framework opens doors for modeling complex combinatorial perturbations, drug response prediction, and identifying optimal therapeutic interventions by leveraging the "regulatory logic" encoded in the graph.

In summary, GREmLN represents a significant advancement in genomics foundation models by successfully integrating graph signal processing with deep learning to capture the complex, non-sequential, and causal nature of cellular transcriptomics.