Unifying multimodal single-cell data with a mixture-of-experts β-variational autoencoder framework

UniVI is a scalable mixture-of-experts $\beta$-variational autoencoder framework that unifies diverse multimodal single-cell data into a shared latent space, enabling robust integration, denoising, and label transfer across paired, tri-modal, and partially observed mosaic study designs without requiring curated feature links or pre-annotated references.

The Gist

Imagine you are trying to understand a complex city. You have three different maps of the same place:

1. The Traffic Map: Shows where people are moving (Gene Expression/RNA).
2. The Construction Map: Shows which buildings are being built or torn down (Chromatin Accessibility/ATAC).
3. The ID Badge Map: Shows what jobs people have (Surface Proteins).

The problem is that these maps are drawn by different teams, use different scales, and often have missing pieces. Sometimes you only have the Traffic Map for one neighborhood and the ID Badge Map for another. Trying to stitch them together manually is a nightmare; if you force them to match perfectly, you might end up putting a bakery on top of a skyscraper just because they are in the same spot on the paper.

Enter UniVI (Unified Variational Inference). Think of UniVI as a super-smart, flexible translator and cartographer that can take these messy, incomplete maps and weave them into one perfect, 3D hologram of the city.

Here is how it works, broken down into simple concepts:

1. The "Expert Team" Approach (Mixture-of-Experts)

Most old methods tried to force all the maps into a single, rigid grid. UniVI is different. Imagine a team of specialists:

• Expert A only looks at the Traffic Map.
• Expert B only looks at the Construction Map.
• Expert C only looks at the ID Badges.

Instead of forcing them to agree on every single detail immediately, UniVI lets each expert do their job. Then, a Manager (the "Mixture-of-Experts" system) looks at what they are saying. If the Traffic expert is confident but the Construction expert is confused (because that part of the map is missing), the Manager listens more to the Traffic expert. This prevents the final map from getting distorted by bad or missing data.

2. The "Shared Secret Language" (Latent Space)

UniVI teaches these experts to speak a new, secret language (a "latent space") that represents the true nature of the city, not just the specific way the maps were drawn.

• When a cell (a person in the city) has both a Traffic Map and an ID Badge, UniVI checks if both experts agree on who that person is.
• If they agree, it locks that understanding in.
• If they disagree, it learns why (maybe the ID badge is blurry, or the traffic data is old) and adjusts accordingly.

3. The "Bridge" Strategy (Handling Missing Data)

This is where UniVI shines in the real world. Often, scientists don't have perfect data. They might have:

• A small group of people with all three maps (The "Bridge").
• A huge group with only Traffic Maps.
• Another huge group with only ID Badges.

Old tools often failed here, either ignoring the huge groups or forcing them to match the small group incorrectly. UniVI uses the small "Bridge" group to learn the secret language. Once it learns the language, it can take the huge groups with only one map and translate them into the shared 3D hologram without needing to re-draw the whole thing. It's like learning a language from a few fluent speakers and then being able to understand tourists who only speak one word of that language.

4. The "Denoising" Magic

Single-cell data is often "noisy" or "sparse" (like a radio with static). UniVI doesn't just map the data; it cleans it up.

• If you give it a blurry ID Badge, it can use the Traffic Map to guess what the ID Badge should have said.
• If you give it a missing Traffic Map, it can use the Construction Map to fill in the gaps.
  This allows scientists to see the "true" cell type even when the data is incomplete.

5. The "Cancer Detective" (Real-World Application)

The paper tested this on Acute Myeloid Leukemia (AML). They had:

• One dataset with RNA and Protein.
• Another with RNA and Genotype (DNA mutations).
• A third with Protein and Genotype.

No single dataset had everything. UniVI acted as the glue, combining them all. It successfully grouped cells by their mutation types and showed how "stemness" (how immature the cancer cells were) changed across the different groups. It even figured out which mutations were present in cells that didn't have direct DNA testing, just by looking at their protein and RNA patterns.

Why This Matters

Before UniVI, integrating these different types of biological data was like trying to assemble a puzzle where half the pieces are from a different puzzle entirely. You either forced them together (and broke the picture) or threw away the pieces that didn't fit.

UniVI is the tool that says: "Let's not force the pieces. Let's build a new table where all the pieces fit naturally, even if some are missing, and we can still see the whole picture."

It gives researchers a flexible, reliable way to combine different biological "languages" to understand diseases, cell types, and how the body works, without needing perfect, pre-labeled data to get started.

Technical Summary

1. Problem Statement

Multimodal single-cell assays (e.g., CITE-seq, Multiome, TEA-seq) allow for the simultaneous measurement of complementary biological layers (transcriptome, proteome, chromatin accessibility, genotype) within the same cell. However, integrating these data presents significant challenges:

• Heterogeneity and Sparsity: Modalities differ drastically in noise models, dynamic range, and sparsity (e.g., RNA is overdispersed, ATAC is near-binary and sparse, proteins are low-dimensional).
• Mosaic Study Designs: Real-world studies often lack fully paired data. Instead, they consist of a small "anchor" dataset with paired modalities and larger "query" datasets that are unimodal (e.g., RNA-only or ATAC-only) or contain different modalities (e.g., RNA+genotype vs. protein+genotype).
• Over-Alignment Risks: Existing methods often enforce uniform correspondence, which can obscure modality-specific biology or create spurious matches when cross-modal evidence is weak.
• Dependence on Priors: Many current approaches rely on curated feature-link graphs (e.g., peak-gene networks) or pre-annotated reference atlases, which may be incomplete, context-dependent, or unavailable for novel disease states.
• Lack of Diagnostics: Few methods provide internal diagnostics to assess the local reliability of an alignment, making it difficult to distinguish robust integration from over-fitting in weakly supported regions.

2. Methodology: UniVI Framework

UniVI (Unified Variational Inference) is a Mixture-of-Experts (MoE) $\beta$-Variational Autoencoder (VAE) designed to learn a shared latent space while preserving modality-specific structures.

Core Architecture

• Modality-Specific Encoders/Decoders: Each modality $m$ has its own encoder $q_{\phi_m}(z|x^{(m)})$ and decoder $p_{\theta_m}(x^{(m)}|z)$. This allows the model to handle distinct noise models (e.g., Zero-Inflated Negative Binomial for RNA, Bernoulli for ATAC) without forcing a single likelihood function.
• Shared Latent Prior: All encoders map to a shared latent prior $p(z) = \mathcal{N}(0, I)$, creating a unified manifold.
• Mixture-of-Experts (MoE) Fusion: For cells with multiple observed modalities, UniVI constructs a fused representation $\tilde{q}_i(z)$ by aggregating modality-specific posteriors using learned weights $\alpha_m$ generated by a gating network. This allows the model to dynamically reweight modalities based on local information content (e.g., down-weighting a noisy modality).

Training Objective (Loss Mode "v1")

The model optimizes a multi-term objective function:
$$\mathcal{L}^{v1} = \sum_{m} \lambda_m \mathbb{E}[-\log p_{\theta_m}(x^{(m)}|z)] + \beta \sum_{m} KL(q_{\phi_m} || p) + \gamma \sum_{m < m'} D_{sym}(q_{\phi_m}, q_{\phi_{m'}})$$

1. Reconstruction Loss: Minimizes error in reconstructing observed data for each modality.
2. KL Regularization: Encourages modality-specific posteriors to match the shared prior.
3. Symmetric Cross-Modal Alignment: For paired cells, a symmetric Kullback-Leibler (KL) divergence penalty $D_{sym}$ is applied between the posteriors of different modalities. This explicitly couples the means and variances of the latent distributions for the same cell, ensuring alignment without requiring external feature graphs.

Inductive Learning & Projection

• Parameter-Frozen Projection: Once trained on a paired reference, the generative parameters (encoders/decoders) are frozen. New unimodal or partially observed cohorts can be projected into the shared space via forward inference through the appropriate modality encoder.
• Optional Supervised Refinement: If labels (cell types, mutations) are available, lightweight supervised heads can be attached to the latent space for fine-tuning. Crucially, the generative decoders remain frozen during this step to preserve the learned biological manifold.

3. Key Contributions

1. Prior-Light Integration: UniVI learns cross-modal correspondence directly from paired cells, eliminating the need for curated feature-link graphs or pre-annotated reference atlases.
2. Robustness to Mosaic Designs: It explicitly handles regimes where data is partially paired, unimodal, or has severe composition shifts, using MoE gating to manage missingness.
3. Interpretable Diagnostics: The framework includes a comprehensive diagnostic suite (FOSCTTM, label transfer, MoE gating maps) to identify regions where alignment is robust versus where it relies on weak evidence.
4. Scalability: The model is designed for large-scale datasets and supports both CUDA and Apple Metal (MPS) backends.
5. Unified Benchmarking: The authors provide a unified evaluation runner that distinguishes between inductive (forward inference on held-out cells) and transductive (joint optimization) methods, offering a fair comparison of generalization capabilities.

4. Key Results

The authors evaluated UniVI across five distinct regimes:

• Paired Bimodal (CITE-seq & Multiome):

  • Achieved high single-cell correspondence (FOSCTTM $\approx$ 0.02 for CITE-seq; $\approx$ 0.048 for Multiome).
  • Demonstrated strong bidirectional label transfer (RNA $\to$ Protein: 95.8% accuracy; Protein $\to$ RNA: 96.2%).
  • Successfully performed cross-modal reconstruction (e.g., denoising RNA from protein data), recovering lineage-specific markers.

• Tri-Modal (TEA-seq):

  • Maintained balanced three-way alignment (RNA, ADT, ATAC) on held-out wells, showing that the manifold generalizes beyond training partitions.
  • Preserved coherent biological structure without any single modality dominating the clustering.

• Reference-to-Query Bridging:

  • A model trained on a paired Multiome reference successfully projected independent RNA-only and ATAC-only cohorts into the same latent space without retraining.
  • Optional supervised refinement improved cross-cohort semantic consistency without collapsing the global geometry.

• Mosaic Disease Setting (AML):

  • Integrated independent RNA+genotype and Protein+genotype AML cohorts using a paired RNA-protein bridge.
  • Revealed genotype-associated neighborhoods (e.g., NPM1 mutations) even without explicit mutation supervision during the initial bridge training.
  • Fine-tuning with mutation heads sharpened genotype grounding while preserving the stemness continuum (LSC17 score).

• Benchmarking:

  • Outperformed or matched state-of-the-art methods (Seurat, Harmony, MultiVI, GLUE, etc.) across fused-space metrics (label transfer, clustering ARI/NMI) and cross-latent correspondence.
  • Avoided the "collapse" trade-off where methods either over-align (losing biology) or under-align (failing to integrate).

• Robustness Analysis:

  • Identified a "low-anchor threshold": strict one-to-one correspondence collapses below 3% overlap, but semantic neighborhood stability remains robust with modest overlap (10%).
  • Demonstrated graceful degradation under localized missingness (e.g., masking one modality for a specific cell type), with MoE gating shifting locally to the remaining modality without distorting the global manifold.

5. Significance

UniVI represents a paradigm shift in multimodal single-cell integration by prioritizing regime-awareness and diagnostic transparency.

• Practical Utility: It addresses the reality of modern single-cell studies, which are increasingly mosaic and heterogeneous, rather than assuming fully paired, perfectly matched datasets.
• Scientific Rigor: By providing tools to detect "weak cross-modal evidence," UniVI prevents researchers from drawing false conclusions in regions where integration is unreliable.
• Flexibility: Its modular design allows for seamless integration of new modalities and study designs (from paired to tri-modal to mosaic) without requiring complex re-engineering of the pipeline.
• Open Source: The full package is available as a Python library, promoting reproducibility and adoption in the community.

In summary, UniVI provides a flexible, interpretable, and robust framework that successfully unifies multimodal data across diverse experimental designs, enabling reliable biological discovery even in the presence of missing data and cohort shifts.