Learning More Effective Cell Representations Efficiently
Abstract
Capturing similarity among cells is the core of many tasks in single-cell transcriptomics, such as the identification of cell types and cell states.
This problem can beformulated in a paradigm called metric learning.
Metric learning aims to learn dataembeddings (feature vectors) in a way that reduces the distance between featurevectors corresponding to cells belonging to the same cell type and increases the distance between the feature vectors corresponding to different cell types.
Deep metric learning on the other hand uses neural networks to automatically learn dis- criminative features from the cells and then compute the metric.
The (deep) metriclearning approaches have been successfully applied to computational biology tasks like similar cell identification and synthesis of heterogeneous single-cell modalities.
We identify two computational challenges: precise distance measurement betweencells and scalability over a large amount of data in the applications of (deep) metric learning.
And then we propose our solutions: optimal transport and coresetoptimization.
Optimal transport has the potential to measure cell similarity more effectively and coreset optimization is promising to train representation learning models more efficiently.
Empirical studies in image retrieval and clustering tasks show the promise of the proposed approaches.
We propose to further explore the applicability of our methods to cell representation learning.
Introduction
The success of machine learning algorithms largely depends on data representation.
Metric learninglearns data embeddings and feature vectors in a way that reduces the distance between feature vectorscorresponding to objects belonging to the same class and increases the distance between the featurevectors corresponding to different classes.
Deep metric learning, on the other hand, uses neuralnetworks to automatically learn discriminative features from the objects and then compute the metric.(Deep) metric learning falls into the broad umbrella of representation learning, whose quest forrepresentation learning is motivating the design of more powerful representation [3, 13].
These representation learning methods have achieved great successes in biology applications [40,42, 19, 23, 35, 53, 23, 48, 1, 11, 30, 14, 45, 46].
For example, Schema [40] uses a principled metric learning strategy to identify informative features in a modality to synthesize disparate modalitiesinto a single coherent interpretation.
It is used to infer cell types by integrating gene expressionand chromatin accessibility data.
Specifically, [50] presents an approach for integrating differentmodalities by learning a probabilistic coupling among them using autoencoders to map to a sharedlatent space. These methods can complement start of art single cell analytics tool such as Dynamo[38] to gain new insights into dynamic biological processes.
The deep metric learning framework SCimilarity (Invited Talk “Design for Inference in Drug Discovery and Development” by Aviv Regev, ICML 2022), which employs a standard triplet loss design, hasachieved impressive performance in identifying similar cells in a massive collection of scRNA-Seqdatasets.
The results can help answer questions like in which tissues and diseases we find fibroticmacrophage-like cells.[15] proposes the Deep Wasserstein Metric Learning Framework, as shown in Figure 1, whichconducts multiple steps of adjustments over the original metric learning framework and achievesimproved performance in image retrieval and clustering tasks.
In the rest of the proposal, we introducethe adjustments accordingly and propose empirical studies for single-cell applications.
Methods
2.1 Representation Learning: Metric Learning and Deep Metric Learning
Representation learning is a class of machine learning approaches that allow a system to discoverthe representations required for feature detection or classification from raw data [38, 50, 18, 24, 17,10, 26, 12, 52]
The requirement for manual feature engineering is reduced by allowing a machineto learn the features and apply them to a given activity.
Metric learning and deep metric learning,specifically, focus on similarity-based approaches to learning the representations.
Thus the similaritymeasurement becomes very important.
Previous work [7] studies similarity measurement in geneexpression. Metric learning has only limited capability to capture non-linearity in the data, whiledeep metric learning captures non-linear features better by learning the non-linear transformation.
The most widely used loss functions for deep metric learning are the contrastive loss and the tripletloss, both use euclidean distance to measure the distance between objects.
A more comprehensiveillustration of so-called “ranking-based” loss functions are summarized in Figure 2.
Given an imagepair, the contrastive loss minimizes their distance in the embedding space if their classes are the same,and separates them a fixed margin away otherwise.
The triplet loss takes triplets of anchor, positive,and negative images, and enforces the distance between the anchor and the positive to be smaller thanthat between the anchor and the negative.
The formation of contrastive loss is as the following.
Wefirst have embedding pairs P , which is sampled from a minibatch of size b .
The pair contains ananchor ϕ a from class y a
and either a positive ϕ p with y a = y p or a negative ϕ n from a different class,y a ̸ = y n . The distance function we utilize is the standard Euclidean distance d e ( x, y ) = ∥ x − y ∥ 2 . Then the network ϕ is trained to minimize:
[Equation] Triplets extend the contrastive formulation by providing a triplets T sampled from a mini-batch: [Equation] In the following, we present two adjustments that can improve (deep) metric learning’s performances,one is optimal transport [9], and one is coreset optimization [31].
2.2 Optimal Transport
Optimal transport (OT) is the general problem of moving one distribution of mass to another asefficiently as possible [41, 15, 34, 47].
Optimal transport has been used tremendously in computationalbiology [39, 37, 43, 44]. For example, [39] uses scRNA-seq data collected across a time course toinfer how these probability distributions evolve over time, by using the mathematical approach ofoptimal transport.
Wasserstein distance provides the mathematical tool to measure distances between functions, histograms, or more general objects in the optimal transport problem.
Wasserstein distance is also called
Earth Mover’s Distance, which is employed to develop PhEMD (Phenotypic Earth Mover’s Distance)[6], which is used to embed the space of drug perturbations on the basis of the drugs’ effects on cellpopulations.
Wasserstein distance-based loss functions have shown superior performance in learningtasks [20].
Thus a new set of loss functions are proposed to replace the Euclidean distance with
Wasserstein distance in original contrastive loss and triplet loss by defining
[Equation] [Equation]
The new Wasserstein-contrastive (wcontrastive) loss and Wasserstein-triplet (wtriplet) loss can beformulated as [15]:
[Equation]
We propose to apply the two new loss functions to similar cell identification and synthesis ofheterogeneous modalities applications.
In the similar cell identification task, the loss functionsabove can impose a discriminative constraint on the feature embedding to improve the similaritymeasurement [49].
2.3 Coreset Optimization
Coreset optimization is about data-efficient methods to find subsets of massive data that can generalizeto the full data when trained on.
In other words, a coreset is a subset of the original training set thatis representative to train machine learning models [15, 51, 32, 28].
More specifically, Wassersteinmeasure coreset [8], is an extension of coresets that takes into account continuous data distributionand generalization.
Recently coreset has been successfully applied to the purification of single-celltranscriptomics data [33].
It focuses on alleviating potential replicate-specific biases within single-celldatasets.
The key is to select a “representative” subset (coreset) of cells from areas of the single-celllandscape where multiple replicates are represented.
The approach [33] takes is solving the exemplarclustering problem, which minimizes the sum of pairwise dissimilarities between cells in the coresetand the rest of cells.
We follow the idea in [33] but use Wasserstein distance to replace the Gaussiankernel to define similarity between cells.
Here r | V | is the exemplar cells from the groundset V and Sis the targeted coreset.
[Equation] (6) [40] processes data from a Slide-seq replicate (three modalities with 20823 transcriptomes * 17607genes) in 34 mins.
[31] demonstrates a specific coreset optimization algorithm CRAIG can achieve the average speedup of 3x for similar loss residual and error rate.
So we expect for the single cellsynthesis task in [40], we can reduce data processing time from 34 mins to 11 mins.
Furthremore,feature-efficient methods [1, 2] can be applied to remove irrelevant variables during the training ofcoreset optimization algorithms.
Experiments Design
[Figure] The experiments conducted on various datasets have demonstrated that optimal transport and coresetoptimization can achieve superior performance on image retrieval and clustering tasks [15, 16]. Inthe following we lay out empirical studies to explore the applicability in building cell representations. The specific detail of computational studies is under investitation.
3.1 Cell Similarity Identification
For the cell similarity identification task, we plan to build on the datasets PBMC [25] and SLN-all[22] which are included in the phenomenal scvi-tools [21].
The PBMC dataset is measured with
CITE-seq (161764 cells). The SLN-all dataset contains Immune cells from the murine spleen andlymph nodes. 3.2 Multimodal Integration Regarding the multimodal integration challenge, we plan to follow the setup in [29] to apply themethods on multiple single-cell datasets including sci-CAR cell line [4], SNARE-seq cell line [5],and 10X Multiome T-cell depleted bone marrow [36] to validate the methods’ effectiveness anddevelop new computational and biological insights from the downstream tasks.
Discussion
In this essay, we propose to apply two computational methods: optimal transport and coresetoptimization, which are successfully demonstrated usability in image representation learning [15, 16],to cell representation learning with applications in cell similarity identification and multimodalintegration.
We will report results and insights in empirical studies in follow up research.