Propensity Score Alignment of Unpaired Multimodal Data

“Strong” Assumptions

The problem that we address in this paper is obviously impossible in general: for example, if I gave you data from two modalities that are completely unrelated to each other (e.g. text from the complete works of Shakespeare, and images of cells under a microscope), there would be no correct solution to the problem because there is no shared information between the two modalities. But not all multimodal matching problems are as pathological as this: if we have shared information and some way of tying them together, then matching is possible.

To our knowledge, there are no prior works tackling the matching problem for which any theoretical guarantees exist on the validity of the matching itself. E.g., Ryu et al., [2024] provides theory on obtaining OT solution itself, but does not describe whether the solution corresponds to a valid matching. In this case, the pathological example above may have an OT solution, but it would clearly not be a useful one. Similarly, Yang et al. [2021] gives no theoretical guidance on the limitations of the method. Our assumptions clearly delineate where we know the problem is solvable from settings where we have no theoretical guarantees.

We also note that our method empirically outperforms all of the published baselines, so we could have chosen to just present it as a state of the art method with no mention of the inherent limitations of matching in general, but we chose to be upfront about the limitations (which likely apply to all of the baselines). NeurIPS reviewer guidelines state, “Reviewers will be specifically instructed to not penalize honesty concerning limitations,” we ask you to reconsider whether the theory that is explicit about assumptions that are sufficient for the method to work is really a weakness (we think it’s a strength!).

Shared Latent Space

We will include a more formal description of Yang et al., [2021]. Roughly, they assume the following data generating process. Each modality, indexed by $i$, is generated according to $X_i = f_i(Z, N_i)$, where $Z$ is a latent variable shared across all modalities. They fit a multi-modal generative model with a VAE assuming $P_Z$ the latent distribution is known, by enforcing that all modalities have the same latent distribution. If label information $t$ is available, they add the classification term requiring $Z$ to classify $t$ (they refer to “prior knowledge” of clusters rather than perturbations, but mathematically $t$ plays the same role).

We believe our data generating process (Eq. (1)) is significantly more general in describing multimodal settings. Yang et al., [2021] require a known latent distribution, while we do not even assume knowledge of the space of $Z$, nor of the functional form of $f$ (though additional structure can yield a richer theory, such as our Prop 3.2). In our biology experiments, $Z$ is just an abstract random variable representing the cells, and $U$ an abstract notion of measurement noise. For the VAE, any theoretical guarantees would require correctly specified latent spaces and decoder functional classes (both of which are also untestable). In [Yang et al., 2021], they state

The dimensionality of the latent distribution is a hyperparameter that is tuned to ensure that the autoencoders are able to reconstruct the respective data modalities well.

By contrast, our theory about propensity scores is valid with only (A1) and (A2), both of which, are easy to interpret scientifically for practitioners (e.g., (A1): interventions do not modify the measurement process, (A2): the measurement contains complete information about the biology).

Scalability

Entropy regularized OT is $\mathcal(O)(n^2 / \epsilon)$ (ignoring log factors, see [Pham 2020] for details), where in our setting, n corresponds to the number of samples in each domain (bounded by the max number of sample in the unbalanced case you refer to). Given that we only match within a treatment group, this is not unreasonable in the domains that we consider: for example, a well of cells in a cell-paint assay will typically have approximately 1000 cells. Constructing a 1000 x 1000 distance matrix and solving the resulting OT problem is very doable on current hardware. Of course, there may exist settings with even larger numbers of samples which would require more scalable unbalanced OT methods, or alternatively switching to shared nearest neighbors (at the cost of some performance; see our experiments). We emphasize that our contribution is in designing the distance metric that leverages propensity scores (and explaining theoretically when and why it works), rather than any particular matching approach that operates on that distance metric, so our approach can be combined with any scalable method that operates on pairwise distances, and we simply found that entropy regularized OT was most effective empirically.

Pham, K., Le, K., Ho, N., Pham, T., & Bui, H. (2020, November). On unbalanced optimal transport: An analysis of sinkhorn algorithm.

Elaborate on the baselines

Gromov-Wasserstein OT (SCOT) uses the local geometry within each modality, assumed to lie in a metric space, to build a cross-modality distance by comparing the distances between pairs within modalities. We expect this approach to fail when one of the modalities does not have a simple metric structure, e.g., for images where euclidean distance in pixel space usually poorly captures semantic similarity. scGLUE is also a VAE approach that is tailored to biological applications, by enforcing that gene expression and their associated proteins (or, in their case, ATAC-seq locations) to have similar embeddings. As such, scGLUE is not applicable outside of genomics data. Importantly, neither of these methods utilize label information. For this reason, the VAE of [Yang et al., 2021] described above, which can use label information in addition to the VAE objective, is considered our main comparison.

Final purpose of the method

We are most interested in (2) “learning from unpaired multimodal samples”. Our cross modality task (see Table 2) explicitly evaluate this and show that you do indeed see a significant improvement by learning from matched data (the difference in performance between our approach and Yang et al. shows that one gets a larger improvement from better matching).

Evaluation

We agree that the metrics can be difficult to interpret, we will include more details in the main text for clarification.

Maybe reporting some classification-based performances for alignment (is the assigned sample from the other modality the correct one)

We report two metrics that measure how close the estimated matchings are to the correct one, when the ground truth is known (image and CITE-seq data). The trace metric corresponds exactly to a measure of how often the assigned sample from the other modality is the correct one ($Tr(M) = \sum_{i} M_{ii}$, where $M_{ii}$ is the weight given to the true pairing).The FOSCTTM is a standard metric in the matching literature, where a value of 0.3 indicates that, on average, the true sample is within the closest 30% to the true sample.

Both of these metrics are standard (which is why we report them), but essentially evaluate your task (1), “aligning unpaired multimodal samples”. For what you called task (2) (learning from unpaired samples), what matters more is (i) how good a proxy the matched sample is for the true match, and (ii) performance on downstream tasks. On synthetic data, we can evaluate (i) by testing how close the latent state of the matched sample is to the ground truth match (MSE in table 1). This metric is a better proxy for downstream performance because if there are a number of samples that are very similar to the true match, then it is likely that you won’t get the exact match correct, but that matching to any of the close samples will still lead to good results (because they’re good proxies for the true match). As discussed above, we also report (ii) in Table 2 (see below for more detail).

or some downstream task performance (what task can we solve using the trained method?)

We emphasize that the metrics above that measure may not be indicative of downstream especially when using soft matching techniques such as OT. The $R^2$ metric on the prediction task (section 5) can be interpreted as the amount of information that is shared between modality 1 and the paired modality 2. E.g., an $R^2$ of 0.233 can be roughly interpreted as the paired modality 1 explains 23.3% of the variation in modality 2, while even the ground truth modality 1 only explains 22.4% (Table 2), indicating that the matched samples are statistically indistinguishable from the true matches in terms of shared information.

Multimodal vs. bimodal

advertised as multimodal, no experiments or proofs use more than two modalities.

Using the term “multimodal” for settings with two or more modalities is standard in the literature. For example, Manzoor et al. 2024 state,

Multimodal systems utilize two or more input modalities ... which could be different from the inputs. [emphasis added]

Similarly, most of multimodal methods discussed in Guo et al [2019] involve only two modalities, in some cases with no obvious way of extending them beyond just two modalities.

That said, the multimodal matching could in theory be done by iteratively applying bivariate matching to a base modality, e.g., sampling a tri-modal observation $(x_i^{(1)}, x_{j}^{(2)}, x_{k}^{(3)})$ would involve sampling $x_j^{(2)}$ from $M^{(1,2)}$ and $x_k^{(3)}$ from $M^{(1,3)}$. One can also think of distance functions involving multi-pairs of propensity scores, e.g., a cost function $c(\pi_i, \pi_j, \pi_k)$, which can be aligned with multi-marginal optimal transport, but our main contribution is to introduce the common space defined by the propensity score.

Manzoor, M. A., Albarri, S., Xian, Z., Meng, Z., Nakov, P., & Liang, S. (2023). Multimodality representation learning: A survey on evolution, pretraining and its applications.

Guo, W., Wang, J., & Wang, S. (2019). Deep multimodal representation learning: A survey.

Questions

what are the labels/perturbations in the image-based dataset?

The perturbations are do-interventions on the latents (location) of the objects in the rendered images, it is briefly described on l253.

section 6.2: reason for using the first 200 principal component?

We chose to use principal components for the GEX data due to sparsity (over ~ 20k dimensions), which is common pre-processing in bioinformatics pipelines. Note this was only used as deterministic pre-processing to train the classifiers and in theory could be avoided with a flexible enough classifier (e.g., a GEX-specific encoder). The protein data is only ~130 dimensions and the measurements are naturally more dense and continuous, and we found that a classifier trained directly on raw measurements performed adequately.

. does assumption A1 hold in the datasets used? Is there a way to check whether the assumption holds?

Whether or not A1 holds, depends in part on the extent to which A2 holds: any part of Z that is not reflected in modality 1 essentially becomes part of $U^{(2)}$ and vice versa. Because CITE-Seq only measures surface proteins of the cell, there is likely part of the latent cell state that is not reflected in the proteomics assays, and some of that state is surely affected by the cell type, so we should not expect assumption A1 to hold exactly. That said, the assumptions remain useful for reasoning about the conditions under which matching is theoretically possible (see also our response to R3) and guiding data collection. E.g. our theory suggests that we will have better matching performance with an assay that measures all proteins, and not just surface proteins.

Dependence on a classifier

There are two drawbacks with this approach: (1) the quality of matching depends directly on the choice of the classifier and its capacity/complexity.

It is true that the quality of the matching depends on the choice of classifier and its capacity, but because you can easily validate the classifier’s performance on a validation set, this makes it relatively easy to choose a good classifier on which to base the matching. Training a classifier is not always trivial, but it is far easier than, e.g. training a variational autoencoder with a shared latent state as in Yang et al. [2021].

more importantly, in many multimodal representation learning settings the label signal is unavailable.

The assumption of perturbation labels is a restriction but it is extremely common in the biological settings that we study (note the concurrent work we cited, Ryu et al. [2024], which was posted just before the submission deadline, also study the same setting motivated by its biological applications). Even when experimental data is not available, it is often possible to rely on weaker class labels. E.g. in our real data experiments on CITE-Seq data, we used cell type as a label.

Question about labels

examples from modality 1 that belong to class t are matched strongly against those from modality 2 that also belong to class t”

In all of our experiments, we perform matching within a known class. So all the examples from class $t$ in modality 1 are restricted to only match to examples from modality 2 in class $t$ (and vice versa). This is why random matching still gives fairly decent results on cross modality prediction tasks (Table 2); we believe that this is roughly equivalent to the “not doing any matching (i.e. just using data from individual modalities)” baseline you refer to (but please let us know if you have something else in mind so that we can update the camera ready). The additional benefit over random comes from matching within a class. Note that matching results in a significant improvement in $R^2$ and the KL metrics.