Multimodal learning for mapping genotype–phenotype dynamics

Abstract

How complex phenotypes emerge from intricate gene expression patterns is a fundamental question in biology. Integrating high-content genotyping approaches such as single-cell RNA sequencing and advanced learning methods such as language models offers an opportunity for dissecting this complex relationship. Here we present a computational integrated genetics framework designed to analyze and interpret the high-dimensional landscape of genotypes and their associated phenotypes simultaneously. We applied this approach to develop a multimodal foundation model to explore the genotype–phenotype relationship manifold for human transcriptomics at the cellular level. Analyzing this joint manifold showed a refined resolution of cellular heterogeneity, uncovered potential cross-tissue biomarkers and provided contextualized embeddings to investigate the polyfunctionality of genes shown for the von Willebrand factor (VWF) gene in endothelial cells. Overall, this study advances our understanding of the dynamic interplay between gene expression and phenotypic manifestation and demonstrates the potential of integrated genetics in uncovering new dimensions of cellular function and complexity.

Understanding the intertwined relationship of genotype and phenotype has long captivated biologists. Wilhelm Johannsen was the first person who coined the terms and attempted to quantify them. In his book Elemente der exakten Erblichkeitslehre (1909)¹, he attempted to provide an ‘exakten’ description for biological phenomena by relating genotype to phenotype, both conceptually and mathematically. A century later, we remain distant from achieving a precise quantification of this relationship. In recent years, however, new sequencing techniques, such as single-cell RNA sequencing (scRNA-seq), have propelled our understanding forward. These technologies reveal the complex dynamics of gene expression at a cellular resolution, highlighting the vast landscape of genotypes influenced by myriad factors². However, these techniques fail to provide a comprehensive mapping of how genotypic combinations give rise to emerged phenotypes.

Current methods in mapping the genotype–phenotype relationship such as forward or reverse genetics, while theoretically capable of dissecting this relationship, fall short in practice owing to the sheer scale and complexity involved³. In human cells, a combination of thousands of genes gives rise to an incredibly diverse phenome landscape⁴. Furthermore, the advent of scRNA-seq, while groundbreaking, introduces its own set of challenges. The technique’s ability to uncover thousands of gene expression changes across cells offers a detailed view of the transcriptomic landscape but also complicates the task of drawing meaningful biological conclusions from these high-dimensional datasets^5,6. This complexity is exacerbated by the vast amounts of data generated in the field, necessitating advanced computational techniques for analysis.

Recent developments in machine learning, particularly the adaptation of self-supervised transformer architectures from the field of natural language processing (NLP), have shown promise in analyzing complex biological datasets^7,8. In these approaches, each gene is a word and each cell state can be modeled as a sentence equivalent to natural language. Current models, however, often lack the interpretability needed to fully understand the sources of variability and fail to capture the dynamic interplay between genotypes and related phenotypes. They focus primarily on resolving cellular heterogeneity without addressing the multifaceted nature of gene expression patterns, which are not only cell-type specific but also highly responsive to organism-level phenotypic characteristics such as age and sex.

Here, we propose the application of self-supervised language models to map the genotype–phenotype landscape simultaneously and introduce the concept of integrated genetics as opposed to forward and reverse genetics. In this approach, phenotypes such as sex, age, anatomical tissue and cell types are integrated and learned with genotypic information measured by scRNA-seq data. This enhances our understanding of the biological context of gene expression and, ultimately, the genotype–phenotype relationship. In particular, we apply this approach to Tabula Sapiens cell atlas⁹, creating an integrated foundation model, called PolyGene, composed of nearly 500,000 human cells across 24 organs from various age ranges. Our results show that PolyGene successfully learns genotype–phenotype annotations and also offers insights into contextualized gene network dynamics. Ultimately, our work lays the groundwork for a phenotype-informed deep learning framework that can be fine-tuned for a broad range of downstream applications, potentially revolutionizing our understanding of the genotype–phenotype relationship and accelerating the discovery of therapeutic targets.

Results

Model overview

PolyGene is a multimodal foundational model to extend the application of attention-based deep learning methods in scRNA-seq data analysis. Our model leverages the unique characteristic of language models, which is their ability to integrate multimodal information by incorporating phenotypic metadata. Multimodal machine learning offers several unique advantages¹⁰, such as enhanced learning for a more comprehensive understanding of the data¹¹, improved robustness in dealing with noisy or incomplete data¹², and cross-modal understanding between different data modalities. These advantages enable the model to perform tasks that require bridging the gap between various types of data¹³. Therefore, we hypothesize that our method can lead to a more accurate biological contextualization, where a better understanding of variation sources and heterogeneity is necessary.

In our approach, for each cell, the gene expression values and associated phenotypes are embedded into a vector representation and passed into the network (Methods). We have incorporated different data modalities of the information of single-cell transcriptomes as well as corresponding phenotypes where each gene expression was measured and encoded jointly as the inputs to the model (Fig. 1). With this structure, PolyGene can learn the relationship between phenotypes (sex, age, tissue and cell type) and contextualized phenotype–genotype relationship. Furthermore, PolyGene develops a more meaningful biological understanding of the expression levels in the context of its preceding neighbors’ rankings and of the phenotype in which that gene has appeared. We refer to this method as integrated genetics, where forward genetics and reverse genetics are accomplished concurrently by masking genotype and phenotype during learning. The resulting outputs of the model, integrated embeddings, include rich information about the contexts in which the genes or phenotypes have appeared.

Fig. 1 ∣. Overview of the model.

Gene expression data from scRNA-seq and associated phenotypes are embedded as input to a language model. The model learns the genotype–phenotype relationship in an integrated manner, where genes and phenotype are masked simultaneously in the input layer. The outputs of the model are high-dimensional embeddings that are information-rich. Created with BioRender.com.

Resolving heterogeneity in phenotypes

One of the main roadblocks in interpreting scRNA-seq expression data is the presence of a high degree of heterogeneity. Cellular heterogeneity stems from the inherent diversity of individual cells within a population or variability in the experimental setup and sequencing techniques. Understanding the source of this heterogeneity is crucial for extracting meaningful insight, from basic cellular processes to complex physiological functions. To address this, we have integrated metadata describing the context of gene expressions with the data. These metadata encompass a wide array of variables, including the specific cell type, organ of origin and donor demographics such as age and sex. This refinement in the model’s capability to accurately represent individual cells results in information-rich embeddings that can be used for downstream tasks such as clustering (Fig. 2a). Contextualized embeddings improve the characterization of cell types, the tissue of origin and even the age and sex of the donors with notable distinctions (Fig. 2c and Supplementary Fig. 3). Furthermore, the classification of these contextualized embeddings has unveiled the presence of two distinct subclusters within several tissues, including the thymus, trachea, subcutaneous adipose tissue and parotid gland (Fig. 2b).

Fig. 2 ∣. Resolving heterogeneity in phenotypes using contextualized embeddings.

a, Uniform Manifold Approximation and Projection (UMAP) dimension reduction visualization of human cell atlas, colored by tissue type, using an unsupervised clustering method on gene expression data. b, UMAP dimension reduction visualization on the same samples using information-rich embeddings from the output of PolyGene model (embeddings are for the respective phenotype labels from the output of the pretrained model). c, Pearson correlation between tissues’ embeddings demonstrates a hierarchical grouping of different tissues; similar tissues have a higher correlation. d, Linear dimension reduction on age embeddings separates different age groups clearly.

Next, we explored the potential of contextualized embeddings in analyzing tissue similarities Fig. 2. These embeddings are from the respective phenotype labels in the output of the pretrained model. The correlation analysis of these embeddings demonstrates a landscape of tissue similarity with encoded biological relationships (Fig. 2c). An example is the high similarity score observed between the cardiac ventricle and cardiac atrium tissues that aligned with their close functional and anatomical relationships. The tissue similarity results also reveal an unexpectedly high similarity between the tongue, the retinal neural layer and the kidney. This finding, although initially unexpected, aligns with previous research indicating the tongue’s diagnostic relevance for chronic kidney disease^14,15. This intricate link highlights the potential of cross-tissue biomarkers in systemic disease diagnosis. Beyond individual pairwise comparisons, our analysis has revealed a structured hierarchy within the tissue similarity landscape (Fig. 2b).

We further explored the ability of our model to classify phenotypes in unseen datasets and compared it with other methods. First, we gathered an atlas of more than 75,000 different healthy cells from a diverse group of studies on lung, liver and blood transcriptomics^16–20. We then compared the performance of our model with other state-of-the-art methods, such as scGPT⁸ (a similar transformer-based model pretrained on single-cell datasets), in terms of precision, recall and F1 scores for both tissue and cell-type predictions. The results showed that pretrained PolyGene consistently outperformed other methods in distinguishing between closely related cell types and states (Supplementary Table 1 and Supplementary Fig. 4).

Exploring genotype–phenotype relationship with embeddings

To understand what genes universally have influenced the phenotypes, we calculated the cosine similarity between phenotype embeddings and gene embeddings (Supplementary Fig. 5). Our results showed that H4C3 (H4 clustered histone 3) is the most prominent gene among all the phenotypes. H4C3 plays a crucial role in the regulation of gene expression by modifying chromatin structure and is involved in various processes such as DNA repair, replication and cell-cycle progression^21,22. The prominence of H4C3 across all phenotypes is expected, indicating basic characteristics of the cells related to their proliferation status and the continuous nature of cell-cycle progression.

For instance, for age phenotype, FST gene (follistatin) is one of the most similar genes to the age token (Supplementary Fig. 5c). Follistatin is a glycoprotein that binds and neutralizes members of the TGF-β superfamily, such as activin. It plays an important role in muscle growth, wound healing and reproductive processes²³. Its similarity to age phenotype embeddings suggests its involvement processes that are critical during aging. In fact, follistatin has been recently shown to increase muscle size in mice²⁴ and monkeys²⁵ and improve life span in mice²⁶, which demonstrates the importance of tissue activin in aging.

Further analysis identified differentially embedded genes (DEGs) among phenotypes (Methods). Our results showed that, for tissue phenotype, IGHV3-7, CCL17 and IGKV6D-21 are the DEGs (Supplementary Fig. 5e). For cell-type phenotype, CCL3 (C–C motif chemokine ligand 3), CA1 (carbonic anhydrase 1) and CCL4L2 (C–C motif chemokine 4-like) are the top different genes between all cell types (Supplementary Fig. 5f). The differentiating genes in age phenotype are RGS1 (regulator of G-protein signaling 1), CA1 and CCL3 (Supplementary Fig. 5g), and for sex phenotype, CA1, CCL17 (CC chemokine ligand 17) and EREG (epiregulin) are shown to be embedded differently between males and females (Supplementary Fig. 5h). The DEGs show that immune-related genes play the most important role in differentiating the embeddings of various phenotypes.

We hypothesized that embeddings represent an information-rich representation of each phenotype and genotype. Focusing on tissue phenotype, the embedding analysis distinctly separated tissue subtypes (Fig. 3a). Specifically, for different muscle cells, genes like DNAJB1, GEM and GPRC5A were identified as similar across different muscle tissues (Fig. 3b), while HMGN2P40, CCNB1, and CTD and GNLY genes were notably different (Fig. 3c). This methodology extended to other tissues, with vasculature and cardiac tissues examined for the impact of genotypes on emergent phenotypes (Fig. 3). For instance, for cardiac tissue, the maternally expressed gene 3 (MEG3) gene was identified as the top DEG that distinguishes ventricle embeddings from atrium embeddings. Notably, kidney and tongue tissues were found to share significant similarities, attributed to genes such as transcription factor ELF3 (which regulates epithelial cell differentiation), CRYAB (molecular chaperone in skeletal muscle), LGALS4, PDK4 (highly expressed in skeletal muscle) and CD59 (expressed in smooth muscle cells), highlighting the molecular variations among different muscle cells. Overall, our results demonstrate how high-dimensional embedding analysis can help with understanding the genetic diversity contributing to phenotypic heterogeneity.

Fig. 3 ∣. Obtaining contextualized genotype–phenotype relationship from contextualized embeddings.

a, UMAP visualization of tissue phenotype embeddings for various tissue types that are considered similar such as muscle, vasculature and cardiac tissue and are considered different such as kidney, tongue and retinal tissue. b, Similarity analysis using cosine similarity between phenotype embeddings of tissue and genotype embeddings present in those tissues can identify genes that have similar roles in these tissues. c, Similarity analysis between various tissues can help with identifying DEGs (Methods) that have been different in all observed contexts.

Contextualized genotype encoding across phenotypes

Next, we leveraged gene embeddings to characterize how specific genes exhibit varied functions and associations across different phenotypes, specifically organs. Our investigation into IL2 gene embeddings unveils a unique pattern of tissue similarity. Notably, clustering based on IL2 embeddings shows the similarity between the inguinal lymph node, bone marrow and the skin of the chest and abdomen, setting them apart from other lymph nodes and skin regions across the body (Fig. 4a). This distinction suggests potential multifaceted roles of IL2 across different tissues, and its involvement in diverse biological processes, specifically to induce similar immune responses in T and B lymphocytes at these tissues²⁷. In addition, IL2 embeddings delineate a clear difference between the large intestine and the small intestine. Previous studies in mice have demonstrated that IL2 is necessary for maintaining immunologic homeostasis in the large intestine, but is not essential in the small intestine, indicating the existence of different sources of IL2 (ref. 28). This provides a compelling example of how a single gene’s influence can diverge across closely related tissues.

Fig. 4 ∣. Dynamic phenotype relationships.

Measuring similarity for gene embeddings in different contexts reveals a dynamic relationship between phenotypes. a, Analysis of IL2 embeddings demonstrates a pattern of similarity between tissues that do not specifically correlate with the anatomical location of the tissues. b, KRT8 embeddings show a different pattern from IL2 that is more similar to the anatomical organization of the tissues. c, Evaluating Euclidean distance between gene embeddings for different tissues (left) and cell types (right) demonstrates that gene embeddings are context dependent.

On the contrary, the analysis of KRT8 gene (keratin 8) embeddings presents a different scenario. Here, the small and large intestines exhibit a high degree of similarity when clustered on the basis of KRT8 embeddings, aligning more closely with traditional understandings of tissue similarity (Fig. 4b). This outcome indicates the heterogeneous nature of gene function across tissues, with KRT8 demonstrating a more uniform pattern of influence within the intestinal tract. We have explored this polyfunctional characteristic of genes in the next section in more detail.

To demonstrate that the genotypes are dynamic and context dependent, we examined the distance between gene embeddings in various tissues and cell types (Fig. 4c). For instance, for the aorta, spleen and subcutaneous adipose, certain genes—NICOA7, XBP1, BNIP3L and CTSD—exhibited distinct patterns of similarity. These genes, despite being expressed across all these tissues, showed closer embedding proximity between the aorta and subcutaneous adipose than between the aorta and spleen (Fig. 4c). Conversely, the embedding distances for TIMP1, CD55 and HSPA5 across these tissues were notably less variable, indicating a more uniform functional role or expression pattern across diverse biological systems. For cell types similarly, the FBXO7 gene showed closer embeddings between memory B cells and naive B cells than between memory B cells and kidney epithelial cells (Fig. 4c).

Context-dependent structure of gene networks

The main contribution of our work is the illumination of dynamic gene networks that operate across different phenotypic contexts. The integration of phenotypes allowed the identification of embedded gene expression patterns and context-dependent interactions that can be also investigated in the structure of the gene interaction networks. By including cell metadata in our model, we can encode the conditions in which gene expression occurs and sources of variation into the gene embeddings. In this section, we characterize the relationship between these gene embeddings in different contexts.

First, we analyzed the overall structure of gene networks across different phenotypes. Metabolic, protein and gene interaction networks have been reported to exhibit scale-free behavior that makes them robust and efficient in processing information^29,30. However, the context-dependent architecture of the network from RNA expression data has not been reported before. Here, we used similarity between gene embeddings to construct the gene network (where nodes are genes and edges are the similarity score higher than a specific threshold). Our results show that the networks follow a scale-free pattern, where the distribution of connections to nodes follows a power law (Supplementary Fig. 6). This indicates that most nodes have few connections, while a small portion of important nodes have a huge number of connections (hubs). In addition, the construction of gene interaction dendrograms facilitates a deeper understanding of the networks’ hierarchical organization (Supplementary Fig. 6).

Specifically, we investigated the role of aging in the structure of genetic network in endothelial cells (ECs). Our results show that aging alters the slope of the power-law distribution within the genetic network, with a pronounced effect on low-degree nodes—those with fewer connections (Fig. 5a)—which reveals restructuring of the network. Specifically, we observed that specialized nodes, which typically have fewer connections, begin to form more linkages, thereby diluting their specialized roles within the network (Fig. 5b). Our analysis highlighted several genes such as KCNH8 (potassium voltage-gated channel subfamily, Kv12.1), DNAJA4 (DnaJ heat shock protein family (Hsp40) member A4) and EGLN3 (Egl-9 family hypoxia-inducible factor 3) as exhibiting the alterations in their network embeddings. These genes are proposed as candidates for further investigation into their roles in vascular aging.

Fig. 5 ∣. Restructuring of the gene network in ECs during aging.

a, Degree distribution analysis on ECs shows a scale-free pattern with different characteristics of the networks for ECs of donors under 40 years old (left), and cells from donors over 60 years old (right). b, The clustering of gene embeddings demonstrates a hierarchical structure in the network highlighting the changes in subnetworks between young ECs (top) and old ECs (bottom) that is more pronounced in the specialized region of the network (depicted by orange color). Genes that are highly different between the two networks are scaled for visibility. c, Similarity analysis between gene embeddings reveals DEGs between young and old ECs (left) and genes with a higher degree of similarity that do not change during vascular aging (right).

Identifying polyfunctionality of genes

The extraction of high-dimensional gene embeddings from our model offers a unique perspective into the diverse contexts and functions where genes are active. Through these embeddings, we can gain an insight into the functional heterogeneity of genes, which we call polyfunctional genes. This is defined as the gene’s ability to form multiple, functionally distinct clusters even within the same cell type. To demonstrate the concept of the polyfunctional gene, we have analyzed VWF (von Willebrand factor) gene within ECs. We identified two functionally distinct clusters that are not phenotype dependent (Supplementary Fig. 7). The first aligns with VWF’s well-documented role in blood coagulation and critical function in hemostasis (Fig. 6a, cluster 1). By contrast, the second cluster reveals a lesser-known aspect of VWF’s functionality, related to its involvement in the management of reactive oxygen species and the cellular response to glucocorticoid stress (Fig. 6a, cluster 0).

Fig. 6 ∣. Polyfunctional characteristics in VWF in ECs and CD55 in T cells.

a, The clustering of VWF’s embeddings in ECs reveals two distinct clusters (middle). In cluster 0, VWF interacts with genes that are involved mostly in reactive oxygen species metabolic process (left), whereas cluster 1 shows the more canonical function of VWF in blood coagulation (GO enrichment analysis using Fisher’s exact test with false discovery rate (FDR) correction to analyze gene sets in each cluster). b, Contextualized embeddings for CD55 gene in T cells show six distinct clusters (left); the networks that clusters form for clusters 0, 1 and 5 are depicted (right).

Similarly, our examination of the CD55 gene, known for its regulatory role in the complement system, across T cells presents a similar pattern of functional diversity with six distinct clusters (Fig. 6b). However, polyfunctionality is not a universal characteristic of all genes. For instance, we showed that, for gene KLF5, no distinct functional clusters were observed in ECs (Supplementary Fig. 8). By dissecting the polyfunctional clusters, our method enriches the gene function lexicon and opens possible avenues for exploring the dynamic interplay between genes and their expression contexts. One of the main applications of polyfunctionality could be in improving the precision of targeting the special function of a gene in drug discovery and cell therapy.

Discussion

A central goal of genetics is to understand the complex relationship between gene expression profiles and emergent phenotypes. However, studying this relationship is challenging because of the vast array of possibilities and the bidirectional nature of the relationship, which necessitates a holistic approach. The approach presented in this study—termed integrated genetics—leverages high-content transcriptomic data combined with a foundation model that incorporates multimodal information from cell-specific, tissue-specific and demographic features. This method allows a contextualized learning environment, where gene expressions are understood in relation to various phenotypic conditions.

The core innovation of this approach is the use of a self-supervised learning framework inspired by language models, which enables joint learning of genotypic and phenotypic factors. This integrated framework offers distinct advantages: it enhances the ability to analyze gene expression within different phenotypic contexts, enables the discovery of cross-tissue biomarkers and allows a deepened exploration of polyfunctional gene roles. We showed that, by embedding rich contextual information, our model addresses the challenges of cellular heterogeneity and offers embeddings that capture system-level variances.

However, this approach is not without limitations. The reliance on large-scale transcriptomic data can pose challenges in model training and computational requirements, potentially limiting accessibility for settings with constrained resources. In addition, while the model handles phenotypic contextualization effectively, it may still be limited by the sparsity or incompleteness of phenotype-specific data. Future work may focus on integrating more efficient data processing techniques or fine-tuning the model to handle diverse transcriptomic profiles with lower computational overhead.

By expanding on these methods and refining their applications, this study sets the foundation for developing context-aware models in genomics, potentially leading to more effective therapeutic strategies and personalized medicine applications.

Methods

In this section, we will provide an overview of the PolyGene architecture and key features. Previous works^7,8,31 on gene expression data from scRNA-seq data have been limited to statistical analysis of a large number of genes, for example, clustering and cell-type annotation. This approach mirrors the pre-2010 NLP research, which focused primarily on clustering, classification, topic modeling and other statistical analyses. However, the advent of word embedding catalyzed a shift toward more granular text analyses³². We advocate for a parallel methodological evolution in genomics, leveraging the exponential growth in genetic data availability across various species.

The limitations of existing models using Transformer³³ architectures in genomics arise from their reliance on a constrained subset of genes, organized by their expression levels from highest to lowest. This ordering premise is problematic for several reasons. First, it presumes, without clear justification, that a descending expression-based sequence represents a gene’s natural or functionally relevant order. More critically, this method’s validity is compromised by its sensitivity to expression value disparities; minor fluctuations or noise can alter the gene sequence, thereby impacting the analysis. This issue is not encountered in text analysis.

A potential solution for the positional embedding problem in genetics is its complete elimination, yet this introduces new challenges in loss function computation. The challenge arises from the difficulty in comparing input tokens with output tokens when permutations in input are allowed. Without positional embeddings, the model lacks a reference to determine the relative positions of genes to compare, which complicates the computation of the loss function. This issue appears because gene expression data do not have a natural sequential order like text data, making it harder to establish a consistent mapping between input and output tokens. Approaches to address this include generating numerous permutations of the input to refine loss calculation or adopting methodologies such as Set-Transformer. However, these solutions are inefficient for practical implementation owing to their computational demands. In our study, we adopt an innovative strategy, using a consistent positional encoding for each gene, complemented by padding for unexpressed genes within a specific cell. This method effectively resolves the issues associated with loss function computation while circumventing the constraints imposed by gene order.

Input representation

To model our problem, we start with a representation of the genes in a single cell $C$ that includes all of them as a vector. For the sake of consistency and without loss of generality, we select a specific and fixed order for these genes

$$\mathbf{C}=[{\mathit{g}}_1,{\mathit{g}}_2,\dots ,{\mathit{g}}_N].$$ (1)

Each gene value is the expression of the gene in that cell that is normalized to a range. To make the values of the expression comparable, we have to normalize the value by dividing it by the dispersion rate. To choose a threshold for distinguishing expressed from nonexpressed genes, we conduct a histogram analysis of gene expression values across all cells. This analysis reveals a concentration of expression values below 0.1, followed by an increase, forming a distinct bump in the histogram. Such a distribution pattern indicates that a threshold of 0.1 is an effective demarcation for identifying genes as expressed or nonexpressed.

After determining the threshold, the rest of the expression range is divided into $B$ equal-sized bins and each gene is represented as the number of the corresponding bin (higher expression corresponds to higher bins). More specifically, each gene value belongs to a bin

$${\mathit{g}}_i\in [0,B-1].$$ (2)

Given the extensive number of genes, the resultant input vector length, if every gene is included, would be impractically large for efficient computation. To address this, our initial step involves sorting the genes on the basis of their ‘variability’ or the degree to which their expression levels differ across various cells. We then assign a rank to each gene according to this variability criterion. This ranking process enables us to prioritize genes on the basis of their variability, thereby streamlining the input size by focusing on those genes that are likely to provide the most informative insights into cellular functions and states. We can model each cell by $T$ highly variable genes (from all $N$ genes)

$$\mathbf{C}=[{\mathit{g}}_1,{\mathit{g}}_2,\dots ,{\mathit{g}}_T]T\le N.$$ (3)

To pass the genes into the network, we need to first embed them in a vector representation. A gene embedding is constructed by mapping its value, which is the bin number, into an initially random embedding of size $h$. This way, the cell is represented as a matrix $X$ as follows:

$${\mathbf{X}}_{h\times T}=[{\mathbf{X}}_1,{\mathbf{X}}_2,\dots ,{\mathbf{X}}_N],$$ (4)

in which ${\mathbf{X}}_i$ is of size $h\times 1$ and is defined as

$${\mathbf{X}}_i=\text{embedding}({\mathit{g}}_i).$$ (5)

Model architecture

To calculate the contextualized embedding of each gene, a denoiser also known as BERT (Bidirectional Encoder Representations from Transformers)³⁴ has been used. The model is dependent on both directions of the input (unlike autoregressive models that are only dependent on the previous words or tokens). In this study, a cell, denoted as $\mathbf{X}$ and characterized by a sequence of genes, undergoes a process to generate its corrupted counterpart, $\tilde{\mathbf{X}}$. This corruption involves the random masking of the proportion $r$ of its genes. This procedure can be equivalently conceptualized as independently and randomly masking each gene within the sequence.

$$\tilde{\mathbf{X}}=\text{Bernoulli}(r)r\in [0,1].$$ (6)

The task of masked language modeling is to reconstruct the expression values of masked genes on the basis of the known genes. The corrupted version is then fed into the Transformer model.

The Transformer model is decomposed of $L$ layers of Transformer blocks $H_{\theta }$

$${\mathbf{X}}^l=H_{\theta }({\mathbf{X}}^{l-1})l\in [1,L].$$ (7)

The Transformer block takes an input of length $T$ and outputs the same size in the output

$$H_{\theta }(\mathbf{X})=[H_{\theta }({\mathbf{X}}_1),H_{\theta }({\mathbf{X}}_2),\cdots ,H_{\theta }({\mathbf{X}}_T)],$$ (8)

in which $H_{\theta }$ is the composition of the following functions:

$$\begin{gathered} \overline{{\mathbf{X}}^l}=\text{LayerNorm}({\mathbf{X}}^l) \\ {\mathbf{M}}^l=\text{MultiHead}\left(\overline{{\mathbf{X}}^l}\right)+{\mathbf{X}}^l \\ \overline{{\mathbf{M}}^l}=\text{LayerNorm}({\mathbf{M}}^l) \\ {\mathbf{X}}^{l+1}=\text{FeedForward}\left(\overline{{\mathbf{M}}^l}\right)+{\mathbf{M}}^l \end{gathered}$$ (9)

LayerNorm and FeedForward are the usual layers used in modern neural networks. MultiHead function is the multihead attention module that is defined as

$$\text{MultiHead}(\mathbf{X})=\text{Cancat}({\mathbf{z}}_1,\dots ,{\mathbf{z}}_k)W_0$$ (10)

$${\mathbf{z}}_i=\text{Attention}\left(\mathbf{X}{W}_i^Q,\mathbf{X}{W}_i^K,\mathbf{X}{W}_i^V\right)$$ (11)

$$\text{Attention}(Q,K,V)=D^{-1}AV,A=\text{exp}\left(Q{K}^T∕\sqrt{d}\right),\mathbf{D}=\text{diag}(A{1}_L),$$ (12)

where $\text{Concat}(.)$ is the concatenation function; $\text{diag}(.)$ is a diagonal matrix with the input vector as the diagonal; $1_L$ is the all-ones vector of length $L$; and $\text{exp}(.)$ is applied element-wise. All the matrices $W_i^Q$, $W_i^K$, $W_i^V$ and $W_0$ are trainable parameters. $Q$, $K$, $V$ and $A$ are query, key, value and attention matrices, respectively, and $d$ is the embedding dimension.

It should be noted that $W_{\mathrm{p}}$ is the position embedding that in the first layer of the Transformer maps the discrete inputs into initial random embeddings, to distinguish each gene we rely on its position. Note that, unlike text, there is no natural order in gene sequences, so here the position embedding makes it easy to distinguish between each gene. $W_{\mathrm{e}}$ is the gene embedding for each expression. The size of position embedding is $B$. The input to the first transformer block is

$${\mathbf{X}}_0=G{\mathbf{W}}_{\mathrm{e}}+{\mathbf{W}}_{\mathrm{p}},$$ (13)

in which $G$ is the binned matrix. The last layer of the Transformer is passed through a feedforward layer and results in the logits that will be passed through a softmax function to get the probabilities over the expression values.

$$\begin{gathered} S_{\theta }\left(\widehat{\mathbf{X}}\mid {\tilde{\mathbf{X}}}^1\right)={\mathbf{X}}^L{W}_{\mathrm{e}}^T \\ {p}_{\theta }\left(\widehat{\mathbf{X}}\mid {\tilde{\mathbf{X}}}^1\right)=\frac{\text{exp}\left(S_{\theta }\left(\widehat{\mathbf{X}}\mid {\tilde{\mathbf{X}}}^1\right)\right)}{{\sum }_{\widehat{\mathbf{Z}}}\text{exp}\left(S_{\theta }\left(\widehat{\mathbf{Z}}\mid {\tilde{\mathbf{X}}}^1\right)\right)} \end{gathered}$$ (14)

The masked language modeling task is to minimize the cross entropy between the probability distribution from the model and the hard labels (one hot encoding) from the dataset

$$\underset{\theta }{\text{min}}\text{CrossEntropy}(q,p_{\theta })=-\sum _{i=1}^{B}q\log p_{\theta }.$$ (15)

Because one hot encoding has been used, the equation (15) will reduce to

$$\underset{\theta }{\text{max}}\log p_{\theta }(\widehat{\mathbf{X}}\mid \tilde{\mathbf{X}})\approx \sum _{t=1}^T{I}_M(t)\log p_{\theta }({\widehat{\mathbf{X}}}_t\mid \tilde{\mathbf{X}}),$$ (16)

in which $I_M(t)$ is the indicator function for the masked tokens ($M$). Here, the independence assumption between the prediction of each gene has been made.

This methodology is crucial because it allows the extraction of contextualized embeddings for each gene, factoring in interactions with other genes within the same cell. Unlike previous methods, this approach does not rely on precise gene ranking for loss function computation. When there are more than two bins ($B>2$), ordinal regression applies. Unlike BERT’s original implementation, which lacks token ordering, gene expressions are binned in a sequence, requiring a modified loss function. The solution is a soft labeling technique that spreads a portion of the probability across bins, enhancing convergence despite not accounting for proximity between bins (that is, bin 9 being closer to bin 8 than bin 1).

Integrating genotype and phenotype

Each cell has other phenotypic features such as sex, tissue type, cell type, developmental stage and diseases. Equation (1) is expanded, and the input to the model can be represented as

$$\mathbf{C}=(p_1,p_2,\dots ,p_F,{\mathit{g}}_1,{\mathit{g}}_2,\dots ,{\mathit{g}}_N),$$ (17)

in which $p_i$ are the phenotypes and $\mathbf{C}$ represents a cell as the input to the model. Each $p_i$ is represented as one hot encoding, and for each phenotype there is a ‘[no_phenotype]’ vocabulary item to represent the missing phenotype; for example, for the missing tissue type, we use ‘[no_tissue]’. The model is designed to perform a multitask learning approach, where the objective is to predict not only the gene expressions but also the phenotypic values at the same time. The masking probabilities for phenotypes and genotypes are different: phenotypes are masked with higher probability because the number of phenotypes is lower than the genotypes, and this helps the model to be more robust. In our experiments, we used 50% for phenotypes masking probability versus 15% for genotypes.

This approach has been very successful in other domains of machine learning such as computer vision and NLP³⁵, but its power has not been used extensively in the domain of genomics data. There are several advantages to this technique. For example, it helps to generalize better in phenotypic domains where there is data scarcity. Also, it gives much more flexibility in exploring a joint model that incorporates both phenotypic and genotypic features where we can control and observe the effects of each phenotype on gene expressions and vice versa.

Pretraining

The input data are randomly masked and a model is trained to make predictions based on the remaining inputs. We randomly mask the nonzero gene expression and then reconstruct the original inputs by model predictions using the remaining genes. Zero expression in the data represents the absence of gene activation (analogous to the absence of a word in a sentence in NLP). Therefore, we chose to mask only nonzero gene expressions to focus on reconstructing meaningful and biologically relevant signals. In our implementation, all 58,000 genes in each cell are considered; therefore, masking the zero-expression genes or omitting them (by padding) has the same effect computationally.

Evaluation and predictions

Three different model sizes are trained on the basis of the top highly variable genes (Supplementary Table 2). The models are trained on 95% of the data (450,000 cells) and evaluated on 5% of the previously unseen data (20,000 cells). The data are shuffled before splitting and again during training to ensure a fair evaluation. No hyperparameter tunning has been performed with the evaluation data. The largest model (PolyGene 2432) encompasses all the highly variable genes identified in the original dataset. The other two models (PolyGene 512 and 256) incorporate the top 512 and 256 highly variable genes, respectively. We evaluated the precision (P), recall (R) and F1 for each model across all the phenotypes, and all the phenotypes and genotypes (all the input data). Notably, the main difference in performance can be seen in phenotypes that can be attributed to reduced predictive ambiguity in the phenotypic domain when compared with gene expression analysis, as indicated in the confusion matrix (Supplementary Fig. 2). Also, in our experiments on the model with 256 input size, the use of soft labeling improved the results in precision and recall by 9.02% and 18.91% after 10 epochs, respectively Supplementary Fig. 2b.

Implementation details

The models were optimized by the Adam optimizer, using a mini-batch size of 10, at a starting learning rate of 2 × 10⁻⁵, with a warmup ratio of 0.1. They are trained for a total of 13 epochs because the loss and train metrics reached a plateau indicating that the model has converged in validation loss and accuracy metrics. We used the Scanpy³⁶ Python library for gene expression preprocessing, including normalization, log-transformation and highly variable gene selection. To mitigate batch effects, our study utilized the Tabula Sapiens dataset, which was experimentally designed to minimize possible batch effects. Consequently, batch information was not included in the original dataset. We also used the pytorch and huggingface³⁷ library for model design and training.

DEGs.

DEGs are identified by analyzing the relationship between gene embeddings and phenotype embeddings within our integrated genetics framework. The process involves first generating high-dimensional embeddings using our PolyGene model for both phenotypes and genotypes. Then, the cosine similarity between the embeddings of each gene and each phenotype is calculated. This similarity score indicates how closely a gene is associated with a particular phenotype in the high-dimensional space. Genes are ranked on the basis of their similarity scores with each phenotype. This ranking helps to identify genes that show strong associations with specific phenotypes. To determine the statistical significance of the identified DEGs, we apply the nonparametric Wilcoxon rank-sum test.

Reporting summary

Further information on research design is available in the Nature Portfolio Reporting Summary linked to this article.

Supplementary Material

Supplementary information The online version contains supplementary material available at https://doi.org/10.1038/s43588-024-00765-7.

Data availability

We incorporated scRNA-seq data from the Tabula Sapiens study⁹, which provides manually annotated cell and tissue type and also includes metadata about each cell. The data include nearly 500,000 cells from 24 different tissues and organs, many from the same donor. Annotation of cell types shows more than 400 cell types of the human body in 14 different donors. The entire dataset can be downloaded at https://tabula-sapiens-portal.ds.czbiohub.org. Source data for Supplementary Fig. 1 are also available at this link. The Lung-liver atlas is an integration of the following datasets accessible through the Cellxgene platform at https://cellxgene.cziscience.com/collections/74e10dc4-cbb2-4605-a189-8a1cd8e44d8c, https://cellxgene.cziscience.com/collections/6f6d381a-7701-4781-935c-db10d30de293, https://cellxgene.cziscience.com/collections/c1241244-b22d-483d-875b-75699efb9f3c and https://cellxgene.cziscience.com/collections/edb893ee-4066-4128-9aec-5eb2b03f8287 and the Gene Expression Omnibus (GEO) accession number GSE136831. Source data for Figs. 2-6 are included to recreate these images.

Code availability

The training script for the PolyGene model, the transcriptome tokenizer and the code for preprocessing data and inference are publicly available via GitHub at https://github.com/roholazandie/polygene and via Zenodo at https://doi.org/10.5281/zenodo.14188129 (ref. 38).