Geometric-aware and interpretable deep learning for single-cell batch correction via explicit disentanglement and optimal transport

The paper introduces iDLC, an interpretable deep learning framework that combines explicit feature disentanglement with optimal transport-regularized adversarial alignment to robustly correct batch effects in single-cell RNA sequencing data while preserving biological fidelity and scaling to millions of cells.

The Gist

The Big Problem: The "Messy Library" of Single-Cell Data

Imagine you are trying to build a giant, perfect encyclopedia of every type of human cell in the body. Scientists have been taking photos of these cells (using a technology called single-cell RNA sequencing) from different hospitals, different labs, and different countries.

However, there's a huge problem: The "Batch Effect."

Think of it like this:

• Lab A takes photos in a bright, sunny room with a red filter.
• Lab B takes photos in a dim room with a blue filter.
• Lab C uses a slightly different camera lens.

Even though they are photographing the exact same person (a "T-cell"), the photos look completely different because of the lighting and filters. If you try to put all these photos into one book, the computer gets confused. It thinks the "Red Filter T-cell" and the "Blue Filter T-cell" are two different species, or it tries to force them together so hard that it smears their unique features.

Current methods try to fix this by guessing or using "black box" math, but they often make mistakes:

1. Under-correction: They leave the "filters" on, so the cells still look different.
2. Over-correction: They scrub the filters so hard that they accidentally erase the person's actual face (biological identity).
3. Confusion: They mix up a "T-cell" with a "Muscle cell" just because they were photographed in the same lab.

The Solution: iDLC (The "Smart Translator")

The authors created a new tool called iDLC (interpretable Dual-Level Correction). Instead of guessing, iDLC uses a two-step process that is like a highly organized translation service.

Step 1: The "Identity vs. Noise" Separator (Explicit Disentanglement)

Imagine you have a messy suitcase full of clothes (the cell data). Some clothes are your actual outfit (the Biological Identity), and some are just dust, lint, and a weird smell from the airport (the Technical Noise/Batch Effect).

Old methods tried to shake the suitcase and hope the dust falls out, but they often shook the clothes out too.

iDLC is different. It has a magical conveyor belt with two distinct bins:

1. Bin A (The Pure Identity): This bin only accepts the actual clothes.
2. Bin B (The Trash): This bin catches only the dust, lint, and smell.

The system is hard-coded to force this separation. It doesn't guess; it physically splits the data into "Who you are" and "Where you came from." This ensures that when we look at the "Who you are" bin, we are looking at a pure, clean version of the cell, free from the "red filter" or "blue filter" noise.

Step 2: The "Geometric Dance" (Optimal Transport)

Now that we have clean "Identity" cards for every cell, we need to mix them together. But we have to be careful.

Imagine you are organizing a dance. You have dancers from New York and dancers from Tokyo. You want them to pair up based on their dance style (e.g., a Jazz dancer from NY pairs with a Jazz dancer from Tokyo).

• Old methods might grab a Jazz dancer and a Hip-Hop dancer just because they are standing next to each other in the room, forcing them to dance together. This ruins the flow.
• iDLC uses a concept called Optimal Transport. Think of this as a "Smart Map." It calculates the most efficient, smoothest path to move the New York dancers to the Tokyo stage without breaking their dance moves.

It uses a mathematical rule (the Sinkhorn algorithm) that acts like a gentle gravity. It pulls similar cells together but respects the "shape" of the group.

• If there is a continuous line of dancers moving from "Standing" to "Running" (a developmental trajectory), iDLC makes sure they stay in that line.
• It won't snap the line in half or glue a "Standing" dancer to a "Running" dancer just to make the groups look mixed.

Why This Matters: The Results

The authors tested iDLC on three difficult scenarios, and it won every time:

1. The "Noisy" Cancer Data: They mixed data from pancreatic cancer patients from different labs. Old tools either couldn't mix them or mixed them so much they lost the rare cancer cells. iDLC mixed the batches perfectly while keeping the rare cells safe.
2. The "Complex" Immune Data: They mixed blood and bone marrow cells from different people. These cells look very similar but have tiny differences. iDLC kept the tiny differences (like distinguishing between CD4 and CD8 T-cells) while removing the "person-to-person" noise.
3. The "Cross-Species" Atlas: They tried to mix human cells and mouse cells. This is like trying to mix photos of humans and dogs. The biological differences are huge. iDLC was smart enough to say, "Okay, these are different species, but these specific cells (like red blood cells) are so similar that we can align them," without forcing the humans to look like dogs.

The Bottom Line

iDLC is a new, transparent, and geometrically smart way to clean up single-cell data.

• Old way: "Let's guess what's noise and what's signal." (Often fails).
• iDLC way: "Let's physically separate the noise first, then use a smart map to gently guide the clean signals together."

This allows scientists to build a single, unified "Google Maps" of human biology that works across different labs, different machines, and even different species, without losing the tiny, important details that make us unique.

Technical Summary

1. Problem Statement

Single-cell RNA sequencing (scRNA-seq) enables high-resolution characterization of cellular heterogeneity, but integrating datasets from diverse sources is hindered by batch effects (systematic technical variations due to protocols, platforms, or labs). Current integration methods face three critical challenges:

• Robustness under strong noise: Existing methods often fail to correct severe batch effects without causing "under-correction" (failing to mix batches) or "over-correction" (merging distinct biological cell types).
• Preservation of biological fidelity: Methods struggle to maintain fine-grained subtypes, rare populations, and continuous developmental trajectories while removing technical noise.
• Specificity in variation sources: Distinguishing between technical batch effects and genuine biological differences (e.g., cross-species variations) is difficult, leading to the erroneous loss of critical biological signals.

Most current deep learning approaches rely on implicit disentanglement (using loss functions to hope for separation), which lacks structural guarantees and often leads to information leakage.

2. Methodology: The iDLC Framework

The authors propose iDLC (interpretable Dual-Level Correction), a two-stage deep learning framework that combines explicit feature disentanglement with optimal transport–regularized adversarial alignment.

Stage 1: Explicit Feature Disentanglement (Residual Autoencoder)

Unlike standard Variational Autoencoders (VAEs) that use implicit separation, iDLC employs a hard-partitioned latent space to physically isolate biological signals from technical noise.

• Architecture: A deep residual autoencoder maps input gene expression $x$ to a structured latent representation $E(x) = [c, n]$.
  • $c$ (Biological Component): The first $l$ dimensions encode cell identity and state (batch-invariant).
  • $n$ (Technical Component): The remaining $k$ dimensions encode batch-specific noise.
• Training Objectives: The model is trained using three synergistic losses:
  1. Reconstruction Loss: Ensures accurate recovery of gene expression profiles.
  2. Content Consistency Loss: Forces the biological component $c$ to remain invariant when decoded with random batch labels.
  3. Batch Classification Loss: A supervised loss ensuring the noise component $n$ can accurately predict the batch origin.
• Outcome: This stage produces a "purified" biological feature space, free from technical confounding.

Stage 2: Optimal Transport-Regularized Adversarial Alignment

Using the purified biological features from Stage 1, iDLC performs distribution alignment across batches.

• High-Confidence Anchors: Mutual Nearest Neighbors (MNN) pairs are identified between batches using the purified biological space. These serve as reliable "anchors" for training.
• Generative Adversarial Network (GAN): A generator $G$ maps source batch cells to the target batch distribution, while a discriminator $D$ distinguishes real target cells from corrected cells (using WGAN-GP).
• Optimal Transport (OT) Regularization: The core innovation is the integration of the Sinkhorn algorithm into the generator's loss function.
  • Instead of hard matching, OT minimizes the entropy-regularized Wasserstein distance between the source and target distributions.
  • This enforces geometric smoothness, ensuring that the alignment respects the underlying topology of the cell state space, thereby preserving continuous trajectories and local structures.

3. Key Contributions

1. Explicit Disentanglement: iDLC is the first to apply a hard-split latent space architecture to scRNA-seq integration, ensuring physical isolation of biological and technical factors, which significantly improves interpretability and the quality of cross-batch anchors.
2. Geometric-Aware Alignment: By incorporating Optimal Transport as a regularization term, iDLC prevents the disruption of continuous biological processes (e.g., developmental trajectories) that often occur with standard adversarial alignment.
3. Scalability and Interpretability: The framework is designed to scale to datasets exceeding one million cells and provides a traceable information flow from raw data to corrected embeddings.

4. Results

The authors evaluated iDLC on three distinct benchmark scenarios against state-of-the-art methods (ComBat, Harmony, scVI, Scanorama, iMAP, scDREAMER, etc.):

• Pancreatic Cancer (PDAC) Datasets:
  • Tested on datasets with mild and strong batch effects.
  • Result: iDLC achieved the highest composite scores. It successfully eliminated strong batch effects without merging distinct cell types (avoiding over-correction) and preserved rare populations better than competitors, which suffered from under-correction or structural fragmentation.
• Human Immune Cell Integration:
  • Integrated multi-source data (donors, tissues, platforms) with fine-grained subtypes (e.g., CD4+ vs. CD8+ T cells) and a continuous hematopoietic developmental trajectory.
  • Result: iDLC was the only method that successfully mixed batches while perfectly preserving the continuous HSPC-to-erythroid trajectory and distinguishing difficult-to-separate subtypes. Competitors either failed to mix batches or disrupted the trajectory.
• Cross-Species Atlas (Human vs. Mouse):
  • Integrated ~933,000 cells from the Human Cell Landscape and Mouse Cell Atlas.
  • Result: iDLC effectively separated species-specific differences from conserved biological states, aligning shared cell types (e.g., neutrophils, oligodendrocytes) across species. Other methods either failed to mix species or incorrectly merged distinct cell types.
• Ablation Studies:
  • Removing explicit disentanglement (iDLC-woED) led to significant drops in biological conservation, confirming the necessity of the hard-split latent space.
  • Removing optimal transport (iDLC-woOT) maintained batch mixing but severely degraded Graph Connectivity, proving OT is essential for preserving topological continuity.

5. Significance

• Paradigm Shift: iDLC moves the field from "black-box" implicit disentanglement to interpretable, principled deep learning with explicit structural constraints.
• Reliability for Complex Data: It offers a robust solution for constructing unified reference atlases across diverse experimental conditions, platforms, and even species, addressing the "under-correction vs. over-correction" trade-off that plagues current tools.
• Clinical and Biological Impact: By preserving fine-grained structures and rare populations, iDLC facilitates more accurate biomarker discovery, disease heterogeneity analysis, and evolutionary biology studies.
• Open Source: The framework is implemented in PyTorch and is publicly available, supporting datasets up to the million-cell scale.