Unbalanced OT and KL Marginal Penalties

Цели урока

• Derive the unbalanced OT formulation with KL marginal penalties and understand the balanced/free transport limits
• Implement unbalanced Sinkhorn with soft marginal updates and apply it to single-cell developmental trajectory inference
• Analyze the outlier robustness of UOT via influence function analysis and partial transport

Предварительные знания

• Balanced OT: Kantorovich formulation and Sinkhorn algorithm
• KL divergence: definition, properties, and Bregman proximal operators
• Single-cell RNA sequencing: gene expression matrices, cell type clustering
• Robustness in statistics: influence functions, breakdown point

• Unbalanced OT Introduction
• Sinkhorn Algorithm

Classical OT treats every gram of mass as sacred - nothing is created or destroyed. But biology disagrees: cells divide and die during development. Point clouds have outliers. Datasets have different sample sizes. Unbalanced OT handles all of this by making mass conservation optional - at a cost.

• Waddington OT (Schiebinger et al. 2019, Science): reconstructed mouse embryo cell differentiation trajectories using UOT on scRNA-seq snapshots
• CellOT: trains a neural transport map between healthy and disease cell states accounting for differential cell proliferation in cancer
• Computer vision: partial shape matching (matching 60% of a 3D model to another) uses partial/unbalanced OT for robust alignment under occlusion
• Federated learning with heterogeneous datasets: UOT aligns feature distributions across sites with different class frequencies without forcing exact marginal matching

From Mass Transport to Mass Balance

The mathematical foundations of unbalanced OT were laid by Chizat, Peyre, Schmitzer, and Vialard (2018) who introduced the KL marginal penalty framework and proved convergence of unbalanced Sinkhorn. Earlier partial transport theory (Caffarelli, McCann 2010) handled fixed-fraction transport but not soft penalties. The biological application that made UOT famous was Waddington OT (Schiebinger et al. 2019, Science) - using cell growth estimates from gene expression to weight the source marginal and reconstruct embryonic development from scRNA-seq. This work won the Science Paper of the Year award and established UOT as the standard tool for developmental biology.

Unbalanced OT: Relaxing the Marginal Constraints

Classical OT requires exact marginal matching: every unit of mass in the source must appear in the target. Unbalanced OT replaces this hard constraint with a soft penalty, allowing mass creation, destruction, and teleportation. This is essential when comparing distributions that differ in total mass.

The KL divergence in UOT can be replaced by other divergences (L2, chi-squared, TV). Each choice gives different robustness properties: KL gives smooth, multiplicative updates; TV gives sparse solutions with hard mass cutoffs.

Waddington OT: Reconstructing Cell Differentiation

Single-cell RNA sequencing (scRNA-seq) measures gene expression in individual cells, but destroys them in the process. Waddington OT (Schiebinger et al. 2019) uses unbalanced OT to reconstruct developmental trajectories from static snapshots at different time points, accounting for cell growth and death.

In WOT, the transport plan T gives probabilistic ancestry: T[i,j] / T[i,:].sum() is the probability that cell i gives rise to cells of type j. Aggregating over cell type clusters gives fate probabilities for each progenitor population.

Robustness Properties and Partial Transport

Unbalanced OT is fundamentally more robust to outliers than balanced OT because the KL penalty allows the transport plan to ignore extreme points that are costly to match. This makes UOT attractive for contaminated data and partial matching problems.

Partial transport (the Caffarelli-McCann framework) and UOT are related but distinct: partial transport fixes the total transported mass; UOT fixes the penalty per unit of unmatched mass. For tau = cost_max, UOT recovers partial transport with mass s determined by the cost threshold.

Mass Conservation as a Tunable Assumption

Unbalanced OT reveals that mass conservation in classical OT is an assumption, not a necessity. By making it tunable via the penalty parameter tau, UOT creates a continuum from free mass rearrangement (tau -> 0) to strict transport (tau -> infinity). This flexibility is what makes OT applicable to the messy reality of biological and real-world data where mass is never perfectly conserved.

• Optimal Transport — Related topic

Итоги

• UOT replaces exact marginal constraints with KL penalties: min <C,T> + tau*KL(T1||mu) + tau*KL(T^T1||nu)
• Unbalanced Sinkhorn modifies the scaling step: u = (a/Kv)^(tau/(tau+eps)) instead of u = a/(Kv)
• Waddington OT uses UOT with growth-rate-weighted source marginals to reconstruct developmental trajectories from scRNA-seq snapshots
• UOT is robust to outliers: mass with transport cost > tau is not moved, making UOT natural for contaminated or partially matching datasets

Вопросы для размышления

• The KL penalty in UOT allows mass destruction but penalizes it proportionally to tau. How would one choose tau in a scRNA-seq application where the cell proliferation rate is known approximately? Is there a principled calibration procedure?
• Balanced OT is invariant to rescaling both measures by the same constant. Is UOT invariant to such rescaling? What happens if mu and nu have different total masses - does UOT still give a meaningful distance?
• Partial transport fixes transported mass s while UOT fixes the cost-per-unit-of-destroyed-mass tau. Derive the relationship between s and tau for a specific source/target pair. When are they equivalent, and when do they diverge?

Связанные уроки

• ot-15-unbalanced — This lesson deepens the unbalanced OT introduced there with full KL penalty theory and scRNA-seq applications
• ot-01-monge — Balanced OT is the tau -> infinity limit of unbalanced OT
• ot-04-sinkhorn — Unbalanced Sinkhorn extends the balanced algorithm with KL soft marginal updates
• ot-28 — Combining unbalanced and causal OT handles mass creation in temporal stochastic processes
• ot-21
• ot-24-entropic-reg