SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tianyu Liu
tianyu.liu@yale.edu
Yuge Wang
yuge.wang@yale.edu
Hongyu Zhao*
hongyu.zhao@yale.edu
*: Corresponding author
Discovering genes with similar functions across diverse biomedical contexts poses a significant challenge in gene representation learning due to data heterogeneity. In this study, we resolve this problem by introducing a novel model called Multimodal Similarity Learning Graph Neural Network, which combines Multimodal Machine Learning and Deep Graph Neural Networks to learn gene representations from single-cell sequencing and spatial transcriptomic data. Leveraging 82 training datasets from 10 tissues, three sequencing techniques, and three species, we create informative graph structures for model training and gene representations generation, while incorporating regularization with weighted similarity learning and contrastive learning to learn cross-data gene-gene relationships. This novel design ensures that we can offer gene representations containing functional similarity across different contexts in a joint space. Comprehensive benchmarking analysis shows our model's capacity to effectively capture gene function similarity across multiple modalities, outperforming state-of-the-art methods in gene representation learning by up to  $97.5\%$ . Moreover, we employ bioinformatics tools in conjunction with gene representations to uncover pathway enrichment, regulation causal networks, and functions of disease-associated or dosage-sensitive genes. Therefore, our model efficiently produces unified gene representations for the analysis of gene functions, tissue functions, diseases, and species evolution.
Progress in biological technology has broadened the variety of biological data, facilitating the examination of intricate biological systems. A prime example of such technology is single-cell sequencing, which allows for the comprehensive characterization of genetic information within individual cells [38, 78]. This technology provides access to the full range of a cell's transcriptomics, epigenomics and proteomics information, including gene expression (scRNA-seq) [46, 114], chromatin accessibility (scATAC-seq) [20, 17], methylation [60, 68] and anti-bodies [85]. By sequencing cells from the same tissue at various time points, we can gain insight into patterns of cellular activity over time [30]. Moreover, spatial information of cellular activity represents an equally vital, additional dimension. Such data are defined as spatial data. All of these data are known as multi-omics data, and integrating multi-omics data for combined analysis poses a significant challenge.
However, the traditional idea of multi-omics data integration known as using cells as anchors [56, 87, 39, 12, 21, 76] is only partially suitable because of the following two challenges. 1. Different omics
data pose their own challenges. For example, the unit of observation of the spatial transcriptomic data is different from other single-cell data, as a single spatial location contains mixed information (see the left part of Figure 1 a) from various cells [23] and it is not appropriate to generate spatial clusters based on gene expression similarity [16]. Current research [86] also indicates that chromosome accessibility feature is not a powerful predictor for gene expression at the cell level. 2. The vast data volume in atlas-level studies challenges high-performance computing, risking out-of-memory or time-out errors [58, 101]. With nearly 37.2 trillion cells in the human body [77], comprehensive analysis is computationally infeasible. 3. Batch effects may adversely impact analysis results [54] by introducing noise. Consequently, an efficient and powerful method focusing on multi-omics and multi-tissue data (referred to as multimodal biological data) analysis is urgently needed to address these challenges.
Acknowledging the difficulties arising from the cell-oriented viewpoint, prior work shifted focus to the gene perspective. Using gene sets as a summary of expression profiles, based on natural selection during the species evolution, may provide a more robust anchor [8]. Protein-coding genes are also thought to interact with drugs [19], which are more relevant to diseases and drug discovery. Gene2vec [24] is a method inspired by Word2vec [18], which learns gene representations by generating skipgram pairs from the co-expression network. Recently, the Gene-based data Integration and ANalysis Technique (GIANT) [16] has been developed, based on Node2vec [34] and OhmNet [119], to learn gene representations from both single-cell and spatial datasets. However, as shown in Figure 1 (b), significant functional clustering for genes from different datasets but the same tissue was not observed based on the gene embeddings from these two models, because these methods do not infer the similarity of associated genes from different multimodal data. Additionally, they did not offer metrics to quantitatively evaluate the performance of gene embeddings compared to baseline models.
Here we introduce a novel model called Multimodal Similarity Learning Graph Neural Network (MuSe-GNN) for multimodal biological data integration from a gene-centric perspective. The overall workflow of MuSe-GNN is depicted in Figure 1 (a). Figure 1 (b) shows MuSe-GNN's superior ability to learn functional similarity among genes across datasets by suitable model structure and novel loss function, comparing to GIANT and Gene2vec. MuSe-GNN utilizes weight-sharing Graph Neural Networks (GNNs) to encode genes from different modalities into a shared space regularized by the similarity learning strategy and the contrastive learning strategy. At the single-graph level, the design of graph neural networks ensures that MuSe-GNN can learn the neighbor information in each co-expression network, thus preserving the gene function similarity. At the cross-data level, the similarity learning strategy ensures that MuSe-GNN can integrate genes with similar functions into a common group, while the contrastive learning strategy helps distinguish genes with different functions. Furthermore, MuSe-GNN utilizes more robust co-expression networks for training and applies dimensionality reduction [89] to high-dimensional data [73].
To the best of our knowledge, this is the first paper in gene representation learning that combines the Multimodal Machine Learning (MMML) concept [9, 27] with deep GNNs [49, 81] design. Both approaches are prevalent in state-of-the-art (SOTA) machine learning research, inspiring us to apply them to the joint analysis of large-scale multimodal biological datasets  $^{2}$ . As application examples, we first used the gene embeddings generated by MuSe-GNN to investigate crucial biological processes and causal networks for gene regulation in the human body. We also applied our model to analyze COVID and cancer datasets, aiming to unveil potential disease resistance mechanisms or complications based on specific differentially co-expressed genes. Lastly, we used gene embeddings from MuSe-GNN to improve the prediction accuracy for gene functions. Here genes are co-expressed means genes are connected with the same edge.
Given the lack of explicit metrics to evaluate gene embeddings, we proposed six metrics inspired by the batch effect correction problem [54] in single-cell data analysis. We evaluated our model using real-world biological datasets from one technique, and the benchmarking results demonstrated MuSe-GNN's significant improvement from  $20.1\%$  to  $97.5\%$  in comprehensive assessment. To summarize the advantages of our model, MuSe-GNN addresses the outlined problem about cross-data gene similarity learning and offers four major contributions: 1. Providing an effective representation learning approach for multi-structure biological data. 2. Integrating genes from different omics and tissue data into a joint space while preserving biological information. 3. Identifying co-located
<sup>1</sup>Codes of MuSe-GNN: https://github.com/HelloWorldLTY/MuSe-GNN  
<sup>2</sup>Download links: Appendix M.
Figure 1: The workflow of MuSe-GNN, the visualization of gene embeddings, and the problems of two existing methods, GIANT and Gene2vec. (a) The process of learning gene embeddings by MuSe-GNN. Here we highlight the difference between single-cell data and spatial data, and the major applications of gene embeddings. Each spot from single-cell data represents one cell, while each spot from spatial data represents a mixture of cells. (b) UMAPs [63] for the gene embeddings of MuSe-GNN, Gene2vec and GIANT. They are colored by datasets. Based on UMAPs we can conclude that both Gene2vec and GIANT failed to learn the gene similarity based on the datasets from the same tissue, while MuSe-GNN produces meaningful embeddings agnostic to datasets.
genes with similar functions. 4. Inferring specialized causal networks of genes and the relationships between genes and biological pathways or between genes and diseases.
Co-expression Network Analysis. While direct attempts at joint analysis of gene functions across modalities are limited, there is relevant research in identifying correlation networks based on a single dataset. WGCNA [52] is a representative method that employs hierarchical clustering to identify gene modules with shared functions. However, as an early tool, WGCNA has restricted functionality. Its inference of co-expression networks necessitates more rigorous methods, and it cannot be directly applied to the analysis of multimodal biological data.
Network Based Biological Representation Learning. Apart from directly generating gene co-expression networks, learning quantitative gene representations in a lower dimensional space may better describe gene-gene relationships and facilitate downstream analyses. Gene2vec generates gene embeddings based on the co-expression network from a given database. However, it disregards the expression profile information and is based on the old Gene Expression Omnibus (GEO) up until 2017 [24]. GIANT leverages Node2vec [34] and OhmNet [119] to learn gene representations from both single-cell and spatial datasets by constructing graphs and hypergraphs for training. However, this approach still over-compresses multimodal biological datasets by removing expression profiles. Moreover, their co-expression networks created by Pearson correlation have a high false positive rate [82]. Additionally, some of the datasets used by GIANT are of low quality (see Appendix A). There are also methods sharing a similar objective of learning embeddings from other datasets. GSS [71] aims to learn a set of general representations for all genes from bulk-seq datasets [62] using Principal Component Analysis (PCA) and clustering analysis. However, it is for bulk-seq data and cannot be directly applied to single-cell datasets from different tissues. Gemini [103] focuses on integrating different protein function networks, with graph nodes representing proteins rather than genes.
Graph Transformer. In the deep learning domain, Transformer [95, 100] is one of the most successful models, leveraging seq2seq structure design, multi-head self-attention design, and position encoding design. Many researchers have sought to incorporate the advantages of Transformer to graph structure learning. TransformConv [81] introduces a multi-head attention mechanism to the graph version of supervised classification tasks and achieved significant improvements. MAGNA [98] considers higher-level hop relationships in graph data to enhance node classification capabilities. Graphomer [110] demonstrates the positive impact of the Transformer structure on various tasks
using data from the Open Graph Benchmark Large-Scale Challenge [44], which is further extended by GraphomerGD [111]. Recently, GPS [75] proposes a general Graph Transformer (GT) model by considering additional encoding types. Transformer architecture also contributes solutions to several biological questions [113]. scBERT [108] generates gene and cell embeddings using pre-training to improve the accuracy of cell type annotations. The substantial impact of these efforts highlights the crucial contribution of the Transformer architecture to graph data learning.
In the following sections of introducing MuSe-GNN, we will elaborate on our distinct approaches for graph construction that utilize multimodal biological data, followed by an explanation of our weight-sharing network architecture and the elements of the final loss function.
GNN. GNNs aim to learn the graph representation of nodes (features) for data with a graph structure. Modern GNNs iteratively update the representation of a node by aggregating the representations of its  $k$ -order neighbors ( $k \geq 1$ ) and combining them with the current representation. As described in [110], considering a graph  $G = \{V, E\}$  with nodes  $V = \{v_1, v_2, \dots, v_n\}$ , the AGGREGATE-COMBINE update for node  $v_i$  is defined as:
$$
a _ {i} ^ {(l + 1)} = \operatorname {A G G R E G A T E} ^ {(l + 1)} \left(\left\{h _ {j} ^ {(l)}: j \in \mathcal {N} \left(v _ {i}\right) \right\}\right); h _ {i} ^ {(l + 1)} = \operatorname {C O M B I N E} ^ {(l + 1)} \left(h _ {i} ^ {(l)}, a _ {i} ^ {(l + 1)}\right), \tag {1}
$$
where  $\mathcal{N}(v_i)$  represents the neighbors of node  $v_i$  in the given graph, and  $h_i^{(l)}$  and  $h_i^{(l + 1)}$  represent the node representation before and after updating, respectively.
Problem Definition. We address the gene embeddings generation task by handling multimodal biological datasets, denoted as  $\mathcal{D} = (\{V_i,E_i\})_{i = 1}^T$ . Our goal is to construct a model  $\mathcal{M}(\cdot ,\theta)$  designed to yield gene embeddings set  $\mathcal{E} = \{e_1,\dots,e_T\} = \mathcal{M}(\mathcal{D},\theta)$ . In this context,  $\mathcal{D}$  represents the input,  $\theta$  represents the parameters, and  $\mathcal{E}$  represents the output. In other words, we aim to harmonize gene information from diverse modalities within a unified projection space, thereby generating consolidated gene representations.
3.2 Graph construction
Before constructing gene graphs, our first contribution involves the selection of highly variable genes (HVGs) for each dataset. These HVGs constitute a group of genes with high variance that can represent the biological functions of given expression profiles. Moreover, considering co-expression networks is important for gene representation learning because it allows us to characterize gene-gene relationships. As sequencing depth, or the total counts of each cell, often serves as a confounding factor in the co-expression networks inference [11], we employ two unique methodologies, scTransform [35] and CS-CORE [88], to process scRNA-seq and scATAC-seq data, thus creating gene expression profiles and co-expression networks unaffected by sequencing depth. For spatial transcriptomic data, our focus is on genes displaying spatial expression patterns. We use SPARK-X [116] to identify such genes and then apply scTransform and CS-CORE. For detailed algorithmic information regarding these methods, please see Appendix B. Additionally, we demonstrate the immunity of CS-CORE to batch effects when estimating gene-gene correlations in Appendix C. In all our generated graphs (equivalent to co-expression datasets), nodes represent genes and edges represent co-expression relation of genes.
3.3 Cross-Graph Transformer
To capitalize on the strengths of the Transformer model during our training process, we integrate a graph neural network featuring a multi-head self-attention design [81], called TransformerConv, to incorporate co-expression information and generate gene embeddings. Details of TransformerConv can be found in Appendix B.4. Incorporating multimodal information can estimate more accurate gene embeddings, supported by Appendix D. The cross-graph transformer can efficiently learn gene embeddings containing gene functions across different graphs, advocated by the comparison of different network structure choices in Appendix E.1.
Weight sharing. Given the variability among multimodal biological datasets, we employ a weight-sharing mechanism to ensure that our model learns shared information across different graphs, representing a novel approach for learning cross-graph relationships. We also highlight the importance of weight-sharing design in Appendix E.1.
Figure 2: The overall model architecture and the design of loss functions for MuSe-GNN. The color of nodes in the green block represents common/different genes across two datasets. The brown block represents the network architecture of MuSe-GNN, and the blue block represents different loss function components of MuSe-GNN. The color gradients of the left two matrices represent different gene expression levels.
Datasets & Modalities Graph Transformer. Drawing inspiration from the hard weight-sharing procedure [15, 115], we not only employ dataset-specific Graph Transformer (GT) layers  $L_{1}, L_{2}, \ldots, L_{n}$  for each graph  $(G_{1}, G_{2}, \ldots, G_{n})$  from the same modality  $m$ , but also connect all these dataset-specific layers to a set of shared GT layers, denoted as  $D_{m}$ . This design showcases our novel approach to incorporating weight-sharing into the GT framework. The forward process of MuSe-GNN, given dataset  $i$  with network parameter  $\theta_{*}$ , is defined as follows:
$$
X _ {i} ^ {\prime} = D _ {m} \left(L _ {i} \left(G _ {i}; \theta_ {L _ {i}}\right); \theta_ {D _ {m}}\right). \tag {2}
$$
Datasets Decoder. Here we propose a dataset-specific decoder structure based on Multi-layer Perceptrons (MLP). This decoder model is crucial in reconstructing the co-expression relationships among different genes, showcasing our inventive use of MLP for this purpose. Given a graph  $G_{i}$  and its corresponding gene embedding  $e_{i}$ , the decoding process of MuSe-GNN, with network parameter  $\theta_{dec,i}$ , is defined as follows:
$$
E _ {r e c} = \operatorname {D E C O D E R M L P} \left(e _ {i} e _ {i} ^ {T}; \theta_ {d e c, i}\right), \tag {3}
$$
where  $E_{rec}$  represents the reconstructed co-expression network.
3.4 Graph Reconstruction Loss  $(\mathcal{L}_{\mathrm{BCE}})$
Within a single graph, we implement a loss function inspired by the Graph Auto-encoder (GAE) [48]. This function is designed to preserve two key aspects: 1. the similarity among genes sharing common functions, and 2. the distinctions among genes with differing functions. This innovative use of a GAE-inspired loss function constitutes a significant contribution to the methodological design. For a graph  $G_{i} = \{V_{i},E_{i}\}$ , the loss function for edge reconstruction is defined as:
$$
e _ {i} = \operatorname {E N C O D E R G N N} \left(\left\{V _ {i}, E _ {i} \right\}; \theta_ {\text {e n c}}\right) = D _ {m} \left(L _ {i} \left(G _ {i}; \theta_ {L _ {i}}\right); \theta_ {D _ {m}}\right),
$$
$$
E _ {r e c} = \underset {\left| E ^ {\prime} \right|} {\text {D E C O D E R M L P}} \left(e _ {i} e _ {i} ^ {T}; \theta_ {d e c, i}\right); E _ {r e c} ^ {\prime} = \underset {\left| E ^ {\prime} \right|} {\text {F L A T T E N}} \left(E _ {r e c}\right); E _ {i} ^ {\prime} = \underset {\left| E ^ {\prime} \right|} {\text {F L A T T E N}} \left(E _ {i}\right), \tag {4}
$$
$$
\mathcal {L} _ {\mathrm {B C E}} = - \frac {1}{| E _ {r e c} ^ {\prime} |} \sum_ {t = 1} ^ {| E _ {r e c} ^ {\prime} |} \left[ E _ {i} ^ {\prime} [ t ] \cdot \log E _ {r e c} ^ {\prime} [ t ] + (1 - E _ {i} ^ {\prime} [ t ]) \cdot \log \left(1 - E _ {r e c} ^ {\prime} [ t ]\right) \right],
$$
where ENCODERGNN and DECODERMLP represent the encoder and decoder parts of our model, respectively.  $\mathcal{L}_{\mathrm{BCE}}$  denotes the computation of binary cross entropy loss (BCELoss) for the given input data.  $E_{rec}$  is the reconstructed adjacency matrix and  $|E_{rec}^{\prime}|$  is the length of its flatten version. Further justification of the model design can be found in Appendix F.
3.5 Weighted-similarity Learning  $(\mathcal{L}_{\mathrm{Sim}})$
To integrate shared biological information across graphs from disparate datasets, we fuse the reconstruction loss of the input graph structure with a cosine similarity learning loss. In this process, we treat common HVGs between each pair of datasets as anchors. Our aim is to maximize the cosine similarity  $(\mathrm{COSSIM}(\mathbf{a},\mathbf{b}) = \frac{\mathbf{a}\cdot\mathbf{b}}{||\mathbf{a}||_2:||\mathbf{b}||_2})$  for each common HVG across two arbitrary datasets (represented by the yellow blocks in Figure 2), utilizing their gene embeddings. However, in practice, different common HVGs may have different levels of functional similarity in two datasets, which is hard to be directly quantified. Consequently, we employ the shared community score as an indirect measurement, which is incorporated as a weight for the cosine similarity of different common HVG pairs within the final loss function. Considering two graphs  $G_{i} = \{V_{i},E_{i}\}$  and  $G_{j} = \{V_{j},E_{j}\}$ , and one of their shared genes  $g$ , we identify the co-expressed genes of the given gene  $g$  in both datasets, denoted as  $N_{ig}$  and  $N_{jj}$ . Thus, the weight  $\lambda_{ijg}$  for gene  $g$  can be expressed as follows:
$$
\lambda_ {i j g} = \frac {\left| N _ {i g} \cap N _ {j g} \right|}{\left| N _ {i g} \cup N _ {j g} \right|}. \tag {5}
$$
We can iterate over all shared genes from 1 to  $n$  between these two graphs, ultimately yielding a vector as  $\lambda_{ij} = [\lambda_{ij1},..,\lambda_{ijn}]$ . This vector encapsulates the community similarity between the two graphs across all common HVGs. We can then modify the cosine similarity of various gene pairs by first multiplying this vector with cosine similarities and then summing the resultant values across all genes. The negation of the outcome is our final weighted similarity loss, denoted as  $\mathcal{L}_{\mathrm{Sim}}$ . The ablation test for weighed similarity loss can be found in Appendix E.1. More detailed explanations and evidence supporting this design in multimodal conditions can be found in Appendix G.
3.6 Self-supervised Graph Contrastive Learning  $(\mathcal{L}_{\mathrm{InfoNCE}})$
Specifically, when integrating multimodal biological data, we employ the contrastive learning strategy [72] to ensure that functionally similar genes are clustered together as closely as possible, while functionally different genes are separated apart from each other. We utilize Information Noise Contrastive Estimation (InfoNCE) as a part of our loss function to maximize the mutual information between the anchor genes and genes with the same functions. This loss is applicable to different genes in two arbitrary graphs during the training process. In general, if we represent the embeddings of  $N$  genes as  $\mathrm{Gene}_N = \{e_1,\dots,e_N\}$ , InfoNCE is designed to minimize:
$$
\mathcal {L} _ {\text {I n f o N C E}} = - \mathbb {E} \left[ \log \frac {\exp \left(e \cdot k _ {+} / \tau\right)}{\sum_ {i = 0} ^ {K} \exp \left(e \cdot k _ {i} / \tau\right)} \right], \tag {6}
$$
where samples  $\{k_0, k_1, k_2\ldots\}$  compose a set of gene embeddings known as keys of one dictionary and  $e$  is a query gene embedding.  $k_{+}$  represents a positive sample of  $e$  and  $k_{i}$  denotes a negative sample of  $e$ . Equation 6 can be interpreted as a log loss of a  $(K + 1)$ -way Softmax classifier, which attempts to classify  $e$  as  $k_{+}$ .  $\tau$  is a temperature parameter.  $\tau$  is set to 0.07 referenced in MoCo [41].
3.7 Final Loss Function
In summary, the training objective of MuSe-GNN for graph  $i$  comprises three components:
$$
\begin{array}{l} \min  _ {e _ {i}, e _ {j}} \mathcal {L} _ {\mathrm {B C E}} \left(\operatorname {D E C O D E R M L P} _ {i} \left(e _ {i} e _ {i} ^ {T}\right), E _ {i}\right) - \mathbb {E} \left[ \operatorname {C o s S i m} \left(e _ {i} \left[ \text {C o m m o n} _ {i j} \right], e _ {j} \left[ \text {C o m m o n} _ {i j} \right]\right) \lambda_ {i j} ^ {T} \right] \tag {7} \\ + \lambda_ {c} \mathcal {L} _ {\text {I n f o N C E}} \left(e _ {i} [ \text {D i f f} _ {i} ] \oplus e _ {j} [ \text {D i f f} _ {j} ], e _ {i} [ \text {D i f f} _ {\mathcal {N} (i)} ] \oplus e _ {j} [ \text {D i f f} _ {\mathcal {N} (j)} ]\right), \\ \end{array}
$$