Globally Solving Unbalanced Optimal Transport and Density Control for Gaussian Distributions

This paper establishes globally optimal, finite-dimensional solution methods for unbalanced optimal transport and unbalanced density control problems involving Gaussian distributions by proving that these infinite-dimensional variational problems admit exact reductions to optimizations over masses, means, and covariances, often solvable via semidefinite programming and closed-form updates.

The Gist

Imagine you are a logistics manager trying to move a pile of sand from one location (Point A) to another (Point B).

In the classic version of this problem, you have a strict rule: You must move every single grain of sand from A to B. If you start with 100 grains, you must end with 100 grains. This is called "Balanced Optimal Transport." It's like a perfect puzzle where the pieces must fit exactly.

But in the real world, things aren't always perfect. Maybe some sand got blown away by the wind (lost mass), or maybe you accidentally added a bucket of extra sand (gained mass). Or perhaps your "target" pile isn't a strict requirement but just a "wish list" of where you'd like the sand to end up.

This paper introduces a smarter, more flexible way to solve this problem, called Unbalanced Optimal Transport (UOT). Instead of forcing a perfect match, it allows you to create or destroy sand, but it charges you a "penalty fee" for doing so. The goal is to find the cheapest way to move the sand while paying the least amount of penalty fees for the sand you lost or gained.

The "Gaussian" Shortcut

The authors focus on a specific type of sand distribution called Gaussian distributions. In simple terms, imagine the sand isn't scattered randomly; it's piled up in a smooth, bell-shaped mound.

The paper's biggest discovery is a massive shortcut. Usually, figuring out how to move these mounds of sand involves solving an impossible, infinite-dimensional math problem (like trying to calculate the path of every single grain).

The authors proved that you don't need to track every grain. You only need to track three things about the mounds:

1. Where the center is (the mean).
2. How wide the mound is (the covariance).
3. How much total sand there is (the mass).

They showed that the best way to move these bell-shaped mounds is always to stretch and shift them in a straight line (an "affine" move). This turns a super-hard math problem into a simple, solvable puzzle that a computer can solve instantly.

The "Moving Target" Problem (Density Control)

The paper then takes this idea and adds a twist: Time and Control.

Imagine the sand isn't just sitting at Point A waiting to be moved. Instead, it's on a conveyor belt (a dynamic system) moving through time. You have a "steering wheel" (control) that can push the sand left or right at every step.

• The Goal: You want the sand to start near a "Reference A" and end near a "Reference B."
• The Catch: You don't have to hit Reference A or B exactly. You just have to get close. If you miss, you pay a penalty.
• The Cost: Pushing the sand costs energy (fuel).

The authors call this Unbalanced Density Control (UDC). They proved that even in this complex, moving scenario, the best strategy is still to treat the sand as a smooth, bell-shaped mound and use a simple, straight-line steering rule. You don't need a chaotic, random steering wheel; a predictable, calculated push is enough to get the best result.

The "Mass" Decision

A unique feature of this paper is that it treats the total amount of sand as a decision variable.

In traditional problems, you are told, "You have 100 grains, move them." In this new method, the computer decides: "Actually, it's cheaper to move 80 grains and pay a small penalty for the 20 that disappeared, rather than spending a fortune trying to move all 100."

The paper provides a formula to calculate exactly how much mass should be moved to get the perfect balance between moving cost and penalty cost.

The "Entropy" Twist (Optional Chaos)

The paper also explores a version where you want the sand to be a little messy. Imagine you are a baker who wants the dough to be spread out evenly, not clumped together.

They added a "Maximum Entropy" rule. This encourages the control system to be a bit more random and spread out, rather than rigid. They showed that even with this added chaos, the math still simplifies down to the same bell-shaped, easy-to-solve format.

Summary of Results

1. It works: They proved that a solution always exists.
2. It's simple: You can solve these complex, moving-sand problems by just looking at the center, width, and total weight of the sand piles.
3. It's global: The method finds the absolute best solution, not just a "good enough" guess.
4. It's flexible: It handles situations where mass is lost or gained, and it works for both static snapshots and moving systems over time.

In short, the paper takes a very messy, complex logistics problem and shows that if you assume the "cargo" is shaped like a smooth hill, you can solve it perfectly and quickly using a few simple numbers.

Technical Summary

Technical Summary: Globally Solving Unbalanced Optimal Transport and Density Control for Gaussian Distributions

Problem Formulation
The paper addresses two interconnected problems: Unbalanced Optimal Transport (UOT) and its control-theoretic dynamical extension, Unbalanced Density Control (UDC).

1. Static UOT: Given two reference Gaussian measures, $\alpha = c_\alpha \mathcal{N}(m_\alpha, \Sigma_\alpha)$ and $\beta = c_\beta \mathcal{N}(m_\beta, \Sigma_\beta)$, the goal is to find a transport plan $\pi$ that minimizes the sum of the quadratic transport cost and Kullback–Leibler (KL) penalties on the marginals relative to the references. Unlike classical balanced OT, the total mass of the transport plan is not fixed a priori; the KL penalties allow the marginals to deviate from the references, effectively treating mass creation/destruction as an optimization variable.
2. Dynamical UDC: The authors extend this to discrete-time linear systems ($x_{t+1} = Ax_t + Bu_t$). The objective is to steer a state measure from an initial time to a terminal time. The initial and terminal distributions are not hard constraints but are penalized via KL divergence against Gaussian references. The cost function includes the expected quadratic control effort and the KL penalties at the endpoints. The total mass of the state measure is preserved by the dynamics but is not prescribed by the references.

Methodology and Gaussian Reduction
The core methodological contribution is the proof that these infinite-dimensional variational problems admit exact finite-dimensional reductions when the reference measures are Gaussian.

• Static Case (UOT): The authors prove that for any feasible transport plan, there exists a Gaussian transport plan with the same (or lower) cost. Specifically, the optimal solution is characterized by Gaussian marginals and an affine coupling map. This reduces the problem to optimizing over the mass ($c$), means ($m_1, m_2$), and covariances ($\Sigma_1, \Sigma_2$).

  • The optimization decouples: First, a convex problem is solved over the moments (means and covariances). Second, the optimal mass is computed via a closed-form scalar update derived from the optimal value of the moment problem.
  • This approach is also extended to Entropic UOT (EUOT), where an additional KL penalty on the joint plan relative to the product measure is included.

• Dynamical Case (UDC): The authors establish that any feasible solution (involving general measures and control policies) can be replaced, without loss of optimality, by a Gaussian initial measure and an affine-Gaussian feedback policy.

  • The control policy takes the form $u_t = K_t(x_t - m_t) + v_t + w_t$, where $w_t$ is Gaussian noise.
  • Using standard covariance-steering lifting techniques ($S_t = K_t \Sigma_t$, $Y_t = K_t \Sigma_t K_t^\top + \Sigma^u_t$), the non-convex bilinear constraints are transformed into a Semidefinite Program (SDP).
  • Similar to the static case, the problem separates into a convex SDP over the moments and control parameters (for a fixed mass) and a closed-form scalar optimization for the optimal mass.
  • The paper also establishes that if the optimal state covariance sequence remains positive definite, the optimal policy is deterministic (i.e., the control noise covariance $\Sigma^u_t$ vanishes).

• Maximum-Entropy UDC (MaxEnt UDC): The framework is further extended to include an entropy reward on the control policy. The authors show that the optimal solution remains within the Gaussian class and provide a relaxed convex formulation (via Schur complement) that is proven to be tight, ensuring global optimality.

Key Results and Algorithms

• Existence: The paper proves the existence of global minimizers for both the reduced UOT and UDC problems.
• Algorithms:
  • Algorithm 1: Solves the static UOT problem by first solving a convex moment optimization and then applying a closed-form mass update.
  • Algorithm 2: Solves the UDC problem by solving an SDP for the fixed-mass trajectory and control parameters, followed by a closed-form mass update.
• Numerical Validation: Simulations demonstrate that as the penalty parameter $\gamma$ increases, the optimal marginals converge to the reference measures, recovering standard balanced OT behavior. In the dynamical setting, the method successfully steers state measures while balancing endpoint matching and control effort, with the ability to adjust the transported mass.

Significance and Claims
The paper claims to provide globally optimal solution methods for Gaussian UOT and UDC, moving beyond heuristic or approximate approaches. The primary significance lies in:

1. Exact Reduction: Demonstrating that despite the infinite-dimensional nature of the original problems, the optimal solutions for Gaussian references are strictly confined to the Gaussian class, allowing for exact finite-dimensional reformulations.
2. Global Solvability: By separating the mass optimization from the moment optimization and utilizing convex relaxations (SDP), the authors provide algorithms that guarantee global optimality rather than local minima.
3. Unified Framework: The work unifies static transport and dynamic control under a single "unbalanced" paradigm, where mass is a decision variable rather than a fixed constraint. This is particularly relevant for applications where endpoint distributions are reference profiles rather than hard constraints, and where total mass may vary due to uncertainty or partial observability.

The authors note that while the reduced problems are not jointly convex in their raw form, the proposed separation and lifting techniques render them tractable. They also identify future directions, such as applying the framework to single-cell RNA-seq data analysis and establishing connections to the unbalanced Schrödinger Bridge problem, but do not claim to have solved these specific applications within this work.