Integrating gene regulatory priors into Transformer attention with scTransformer for interpretable scRNA-seq analysis

Technical Summary: scTransformer – Integrating Gene Regulatory Priors into Transformer Attention

Problem Statement

Single-cell RNA sequencing (scRNA-seq) data is characterized by high dimensionality, sparsity, and technical noise, posing significant challenges for representation learning. While Transformer-based foundation models (e.g., scGPT, Geneformer) have demonstrated scalability by pretraining on millions of cells, they typically treat genes as independent tokens with unconstrained self-attention. This design allows any gene to attend to any other, relying solely on statistical co-variation to discover dependencies. Consequently, these models often suffer from:

1. Limited Interpretability: Attention patterns are difficult to map to mechanistic biological processes.
2. Instability: In low-data regimes or the presence of batch effects, unconstrained models may overfit to dataset-specific correlations, yielding predictions that are robust in accuracy but unstable in feature selection (i.e., different training runs select disjoint gene sets).
3. Ignored Biological Knowledge: Existing approaches largely ignore curated transcription factor (TF)–target gene (TG) regulatory networks, which encode plausible causal structures in transcriptional control.

Methodology: scTransformer Architecture

The authors propose scTransformer, a biologically informed Transformer encoder that integrates curated TF–TG regulatory networks directly into the self-attention mechanism via hard masking.

1. Input Representation and Embedding:

• Gene Token Embedding: Each gene in a cell is mapped to a $d$-dimensional token combining a trainable gene identity embedding ($e^{id}$) and a continuous expression encoding ($e^{expr}$).
• Continuous Sinusoidal Encoding: Unlike discrete binning, expression values are encoded using a continuous sinusoidal scheme ($\sin(x \cdot \omega_k), \cos(x \cdot \omega_k)$). This preserves quantitative information across the dynamic range of scRNA-seq measurements, ensuring smoothness and boundedness.
• Pooling: A similarity-based pooling mechanism aggregates variable-length token sequences into a fixed-dimensional cell embedding. The context vector is computed as the mean of token embeddings and used as a query in cross-attention over the token sequence.

2. Prior-Gated Attention Mechanism:
The core innovation is the replacement of unconstrained attention with a prior-gated mechanism:

• Regulatory Network: A curated TF–TG network (from CollecTRI) is preprocessed to filter for TFs with sufficient targets and genes present in the dataset.
• Hard Masking: A boolean allow-matrix ($M_{allow}$) is constructed where $M_{allow}(i, j) = 1$ only if gene $i$ is a TF and gene $j$ is its target (or $i=j$ for self-loops). All other interactions are set to 0.
• Attention Constraint: This matrix is converted to an additive mask ($\tilde{M}_{prior}$) where forbidden interactions are assigned $-\infty$. Before the softmax operation in the scaled dot-product attention, this mask is added to the attention scores:
  $$A_h = \text{softmax}\left(\frac{Q_h K_h^\top}{\sqrt{d_k}} + \tilde{M}_{prior} + M_{pad}\right)$$
  This ensures that information flows only along known regulatory edges (and self-loops), effectively constraining the hypothesis space to biologically plausible interactions.

3. Interpretability Metrics:
To quantify the biological relevance of the learned attention, the authors define module-level metrics:

• Concentration Score ($\phi$): Measures how selectively a TF concentrates attention on its annotated targets using normalized entropy. Lower entropy (higher $\phi$) indicates focused, structured attention.
• Module Importance: Combines the concentration score with the total attention mass directed toward targets to quantify the strength of regulatory signals.

4. Implementation Details:

• Data: Evaluated on the MSSM single-nucleus RNA-seq atlas (PsychAD cohort), restricted to 17 subtypes and TF–TG gene intersections.
• Training: Uses AdamW optimizer with Pre-LN architecture. Training employs sparse batching directly on Compressed Sparse Row (CSR) matrices to handle large datasets efficiently.
• Evaluation: Models were tested across varying training data regimes (1% to 100%) to assess sample efficiency and robustness.

Key Contributions

1. Prior-Gated Attention: The first Transformer framework to embed curated TF–TG networks as hard constraints within the attention mechanism, converting regulatory knowledge into structural inductive bias.
2. Continuous Expression Encoding: A sinusoidal encoding scheme that preserves the quantitative nature of scRNA-seq data without discretization.
3. Sparse Training Pipeline: An efficient implementation operating directly on CSR matrices, enabling scalable training on large single-cell datasets.
4. Formal Interpretability Metrics: A set of metrics to quantify attention concentration and regulatory importance, allowing for the direct assessment of biological coherence without post-hoc analysis.
5. Reproducible Protocol: A distributed training setup with controlled data partitioning and multi-seed evaluation to rigorously assess stability.

Results

The authors evaluated scTransformer against unconstrained baselines on cell-type classification tasks:

• Performance in Low-Data Regimes: In settings with limited training data (1% and 25%), the prior-gated model achieved significantly higher accuracy than the unconstrained baseline (e.g., 0.486 vs. 0.464 at 1% data). This demonstrates improved sample efficiency.
• Performance in Full-Data Regimes: At 100% data availability, the prior-gated model (using only TF–TG restricted genes) achieved accuracy comparable to the unconstrained baseline using all genes (0.911 vs. 0.917), despite operating in a much smaller interaction space.
• Stability and Reproducibility: Unconstrained models showed high variability across training runs, selecting disjoint sets of genes (Jaccard indices near zero, negative Spearman correlations). In contrast, prior-gated models exhibited high stability, consistently selecting similar gene subsets and ranking them consistently.
• Biological Coherence: Attention maps from prior-gated models showed higher concentration scores ($\phi$), indicating that TFs focused attention on specific, biologically relevant target subsets. These patterns aligned with known subtype-specific regulatory programs (e.g., enrichment of GRHL3 in endothelial-like populations, NANOG in OPCs), whereas unconstrained models produced diffuse, less interpretable attention distributions.

Significance and Claims

The paper claims that the primary value of integrating regulatory priors is not merely improving predictive accuracy, but enhancing the epistemic reliability and interpretability of the learned representations.

• Resolving Ambiguity: In high-dimensional, sparse data, multiple gene combinations can explain the same phenotype. Unconstrained models may converge to arbitrary, unstable solutions. Regulatory priors constrain the hypothesis space, steering the model toward biologically grounded, coherent solutions.
• Inductive Bias: By embedding biological structure directly into the architecture, the model learns representations that are robust to noise and batch effects without sacrificing flexibility when sufficient data is available.
• Foundation for Interpretability: The approach offers a principled step toward "biologically grounded foundation models," where attention weights are directly interpretable as regulatory programs, facilitating biological insight without requiring post-hoc analysis.

The authors remain modest, noting that current limitations include the use of hard binary masks (which prevent learning novel interactions), restriction to supervised tasks, and evaluation on a single dataset. They suggest future work could explore soft/learnable priors, self-supervised pretraining, and hybrid architectures.