Learning More Effective Cell Representations Efficiently
Abstract
Capturing similarity among cells is at the core of many tasks in single-cell transcriptomics, such as the identification of cell types and cell states.
This problemcan be formulated in a paradigm called metric learning.
Metric learning aims tolearn data embeddings (feature vectors) in a way that reduces the distance betweensimilar feature vectors corresponding to cells of the same cell type, and increasesthe distance between feature vectors corresponding to cells of different cell types.
As a variation of metric learning, deep metric learning uses neural networks to automatically learn discriminative features from the cells and then compute thedistance.
These (deep) metric learning approaches have been successfully appliedto computational biology tasks like similar cell identification, and synthesis of heterogeneous single-cell modalities.
Here, we identify two computational challenges:precise distance measurement between cells, and scalability over a large amount ofdata in the applications of (deep) metric learning.
We then propose our solutions:optimal transport and coreset optimization.
Optimal transport has the potential tomeasure cell similarity more effectively, and coreset optimization is promising totrain representation learning models more efficiently.
Empirical studies in imageretrieval and clustering tasks show the promise of the proposed approaches.
Wepropose to further explore the applicability of our methods to cell representationlearning.
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 neural networks 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,43, 19, 23, 35, 56, 23, 50, 1, 11, 30, 14, 46, 47].
For example, Schema [40] uses a principled metriclearning 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, [53] 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 discover therepresentations required for feature detection or classification from raw data [38, 53, 18, 24, 17, 10,26, 12, 55, 52, 42]
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, 49].
Optimal transport has been used tremendously in computationalbiology [39, 37, 44, 45].
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]
The new Wasserstein-contrastive (wcontrastive) loss and Wasserstein-triplet (wtriplet) loss can beformulated as [15]:
[Equation]
[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 [51].
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, 54, 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 Gaussian kernel 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 achievethe 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, 48] can be applied to remove irrelevant variables during the trainingof coreset optimization algorithms.
Experiments Design
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.
The SLN-all dataset contains Immune cells from the murine spleen and lymph 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.