Multimodal data integration and cross-modal querying via orchestrated approximate message passing

Technical Summary: Multimodal Data Integration and Cross-Modal Querying via Orchestrated Approximate Message Passing

Problem Formulation
The paper addresses the challenge of integrating multimodal data where multiple complementary feature sets are measured on the same set of samples, a scenario prevalent in single-cell multi-omics studies (e.g., simultaneous measurement of transcriptome, surface proteins, and chromatin accessibility). The authors model this as a dependent multifactor problem involving $m$ high-dimensional datasets ($X_h \in \mathbb{R}^{N \times p_h}$) and $\tilde{m}$ low-dimensional datasets ($\tilde{X}_\ell \in \mathbb{R}^{N \times r_\ell}$).

The signal components in these matrices are assumed to arise from low-rank structures driven by latent factors. Specifically, the high-dimensional signals are modeled as $X_h = \frac{1}{\sqrt{N}} U_h D_h V_h^\top + W_h$, and low-dimensional signals as $\tilde{X}_\ell = \tilde{U}_\ell L_\ell^\top + \tilde{W}_\ell$. Crucially, the latent factor vectors ($U_h$ and $\tilde{U}_\ell$) for the same subject across different modalities are assumed to be dependent, drawn from a joint prior distribution $\mu$. The goal is twofold:

1. Data Integration: Jointly recover the latent signals ($U_h, V_h, \tilde{U}_\ell$) from the collection of noisy matrices.
2. Cross-Modal Querying: Given a new "query" subject with only partial observations (a subset of modalities and/or a subset of features within modalities), construct a statistically valid prediction set for the subject's full latent representation.

Existing methods (e.g., WNN, MOFA+, totalVI) are noted to rely on heuristics lacking rigorous mathematical foundations for generalization and, critically, fail to provide uncertainty quantification for predictions.

Methodology: OrchAMP
The authors propose Orchestrated Approximate Message Passing (OrchAMP), a fully data-driven algorithm designed to achieve statistically optimal signal recovery and construct prediction sets.

1. Initialization: The algorithm begins with a singular value decomposition (SVD) of the observed data matrices to obtain initial estimates of the latent factors. It then estimates nuisance parameters (signal strengths and covariance structures) and learns the empirical priors ($\hat{\mu}$ and $\hat{\nu}_h$) using an Empirical Bayes approach.
2. Iterative Refinement: OrchAMP performs synchronous updates across multiple "orbits" (one for each modality). In each iteration $t$:
   • Denoising: The algorithm applies row-wise denoisers ($u_{t,h}$ and $v_{t,h}$) which are conditional expectations of the latent factors given the current estimates and the learned priors. These denoisers integrate information across all modalities simultaneously.
   • State Evolution: The algorithm tracks the asymptotic distribution of the estimates using state evolution parameters ($\Sigma, S$) that evolve recursively.
   • Onsager Correction: To ensure the estimates behave like independent Gaussian noise in the high-dimensional limit, the algorithm subtracts terms involving the Jacobians of the denoisers (Onsager corrections) from the current updates.
3. Prediction Set Construction: For a query subject with partial observations, the method first computes a point estimate of the latent factors using least squares and the learned atlas. It then leverages the asymptotic normality of the OrchAMP iterates to define a conditional distribution for the query's latent factors. A prediction set is constructed as an $\ell_2$ ball centered at the point estimate, with a radius determined by the $(1-\alpha)$-quantile of the Euclidean distance under the conditional distribution, ensuring asymptotic validity.

Key Contributions

1. OrchAMP Algorithm: The proposal of a fully data-driven algorithm for integrative estimation of signal components in dependent multimodal factor models. The method handles varying signal ranks across modalities and incorporates modality-specific information without destroying it during integration.
2. Theoretical Justification:
   • Asymptotic Normality: The authors establish the "large $N$" asymptotic normality of all algorithm iterates under general conditions for a wide range of priors.
   • Bayes Optimality: Under a unique-fixed-point condition on the state evolution, the estimators are shown to achieve asymptotic Bayes optimality (minimizing mean squared error).
   • Prediction Sets: The paper provides the first construction of asymptotically valid prediction sets for query subjects in a multimodal setting. These sets remain valid even if the unique-fixed-point condition is not satisfied.
3. Empirical Validation: The method is validated on synthetic data and real-world single-cell TEA-seq data (transcriptome, epitopes, chromatin accessibility).

Results

• Synthetic Data: In simulations, OrchAMP outperforms separate empirical Bayes PCA (EB-PCA) and standard SVD, particularly when signal-to-noise ratios (SNR) are low or when inter-modality dependence ($\rho$) is high. The performance gain is most pronounced for modalities with weaker signals, demonstrating the benefit of borrowing strength across modalities.
• Prediction Coverage: In finite-sample experiments, the constructed 95% prediction sets achieved empirical coverage probabilities close to the nominal level (e.g., 0.92–0.98), demonstrating the method's ability to provide well-calibrated uncertainty quantification.
• Real Data (TEA-seq): On a human PBMC dataset, OrchAMP achieved clustering performance (measured by silhouette score, V-measure, ARI, and cLISI) comparable to state-of-the-art methods like Seurat WNN and MOFA+.
• Querying: When querying the constructed atlas with cells measured in only a single modality, the resulting prediction sets were shown to be highly informative, typically occupying a specific subset of the correct cell-type cluster rather than just providing a coarse label.

Significance and Claims
The paper claims to provide a rigorous statistical framework for multimodal data integration that moves beyond heuristic approaches. Its primary significance lies in:

• Theoretical Guarantees: Offering provable performance guarantees and asymptotic optimality under a stylized but realistic dependent multifactor model.
• Uncertainty Quantification: Being the first method to generate atlases that allow for asymptotically valid prediction sets for new subjects with partial observations, addressing a critical gap in current label transfer practices where uncertainty is rarely quantified.
• Generality: While motivated by single-cell multi-omics, the proposed model and theory are general and applicable to any setting involving dependent low-rank signals across multiple data sources.

The authors emphasize that while their method matches the cell-segregation capabilities of existing tools, it uniquely offers a theoretically justifiable path to cross-modal querying with quantified uncertainty.