⚡ electrical engineering

Unbalanced Optimal Transport and Density Control for Discrete-Time Linear Systems

This paper presents globally optimal convex formulations for unbalanced optimal transport and its dynamical extension, unbalanced density control, applied to constrained discrete-time linear systems with Gaussian references, drawing parallels to covariance steering.

Original authors: Haruto Nakashima, Siddhartha Ganguly, Kenji Kashima

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haruto Nakashima, Siddhartha Ganguly, Kenji Kashima

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a logistics manager trying to move boxes from one warehouse to another. In the classic version of this problem (called Optimal Transport), you have a strict rule: the number of boxes leaving the first warehouse must exactly equal the number of boxes arriving at the second. If you have 100 boxes to send but only 80 spots to receive them, the classic math breaks down. It's like trying to pour a full gallon of water into a cup that only holds a pint; the math says "impossible."

This paper introduces a more flexible approach called Unbalanced Optimal Transport (UOT). Think of it as a "smart logistics" system that allows for missing or extra boxes. Instead of forcing a perfect match, it says: "Okay, we'll move as many boxes as we can efficiently, but if we have to create new boxes or throw some away to make the math work, we'll charge a penalty fee for that." The goal is to find the cheapest way to move the mass while balancing the cost of moving it against the cost of creating or destroying it.

The Two Main Problems

The authors tackle two specific versions of this problem using a special type of "box" called a Gaussian distribution (which is just a fancy way of describing a bell-curve shape of data).

1. The Static Problem (UOT): Moving Data Between Two Points
Imagine you have a pile of sand (Source) and a target pile of sand (Destination). They might not be the same size.

The Goal: Move the sand from Source to Destination as cheaply as possible.
The Twist: You can add sand to the destination or remove sand from the source if it saves money on the trucking fees.
The Discovery: The authors proved that even though this sounds complicated, the best way to move this "sand" is to treat the piles as simple bell curves. You don't need to track every single grain of sand. You only need to calculate three things:
1. Where the center of the pile is (Mean).
2. How spread out the pile is (Covariance).
3. How much total sand you have (Mass).
The Result: They created a recipe (an algorithm) that finds the absolute best solution by solving a simple math puzzle. It's like having a GPS that instantly tells you the perfect route, even if your starting and ending points have different amounts of cargo.

2. The Dynamic Problem (UDC): Moving Data Over Time
Now, imagine the sand isn't just sitting in two piles; it's on a conveyor belt moving through a factory with machines (a discrete-time linear system).

The Goal: You want to steer the pile of sand from a starting shape to a final shape over a set period of time.
The Twist: You can apply "control forces" (like pushing the conveyor belt) to change the shape and position of the sand. However, you also have the option to add or remove sand at the start and end if it's cheaper than pushing it all the way.
The Discovery: Just like the static version, the authors found that you don't need to simulate every single particle of sand. You can treat the whole moving pile as a single, evolving bell curve.
The Result: They turned this complex control problem into a standard type of math problem (called a Semidefinite Program or SDP) that computers can solve very quickly and perfectly. It's like giving a robot a set of instructions that guarantees it will arrange the sand exactly how you want, with the least amount of effort, even if the sand gains or loses weight along the way.

How It Works in Practice

The paper includes a simulation to show how this works. They tested it with two settings:

Low Penalty for changing mass: When the "fee" for adding/removing sand is low, the system is lazy. It prefers to just move the sand a tiny bit (keeping it close to where it started) rather than paying to move it all the way to the target. It creates a "shortcut" solution.
High Penalty for changing mass: When the fee is high, the system is forced to act like the classic "perfect match" version. It moves the sand exactly where it needs to go to match the target shape, because creating or destroying sand is too expensive.

The Bottom Line

The authors have built a mathematical toolkit that allows engineers and scientists to compare and move data distributions that don't have the same total amount of "stuff" in them. By proving that the best solutions always look like simple bell curves, they turned a messy, impossible-looking problem into a clean, solvable math puzzle. This means computers can now solve these problems perfectly and quickly, which is a big step forward for controlling complex systems where data might be incomplete or changing in volume.

Technical Summary: Unbalanced Optimal Transport and Density Control for Discrete-Time Linear Systems

Problem Statement
This work addresses the limitations of classical Optimal Transport (OT), which requires distributions to have equal total mass, rendering it unsuitable for scenarios involving mass creation, dissipation, or missing data. The authors investigate Unbalanced Optimal Transport (UOT), which relaxes the mass conservation constraint by penalizing deviations from reference measures using Kullback–Leibler (KL) divergence. Furthermore, the paper extends this framework to Unbalanced Density Control (UDC) for discrete-time linear systems. In UDC, the goal is to steer a system's state distribution from an initial measure to a target measure while minimizing control effort and balancing transport costs against fidelity to reference distributions, all without assuming the preservation of total probability mass.

Methodology
The paper focuses on the case where reference measures are Gaussian and the system dynamics are linear. The core methodological contribution is the derivation of globally optimal, finite-dimensional convex formulations for both UOT and UDC.

Restriction to Gaussian Measures:
The authors prove that for both UOT and UDC, there exists an optimal solution where the marginal distributions (and state measures in the control case) are Gaussian.
- For UOT: They demonstrate that for fixed mass, mean, and covariance, the transport cost is minimized when marginals are Gaussian (Proposition 1), and the KL divergence term is also minimized by Gaussian measures under the same constraints (Proposition 2). This allows the infinite-dimensional variational problem to be reduced to an optimization over the mass ( $c$ ), mean ( $m$ ), and covariance ( $\Sigma$ ) of the marginals.
- For UDC: Similarly, they show that for any feasible process, a Gaussian state measure and a Gaussian-affine feedback policy can be constructed that match the first and second moments of the original process while yielding a lower or equal objective value (Theorem 5).
Separation of Optimization:
The resulting finite-dimensional problems are not jointly convex in all variables (specifically mass and moments). However, the authors exploit the problem structure to separate the optimization:
- UOT Formulation: The problem reduces to minimizing a function of mass and moments. The authors propose an algorithm (Algorithm 1) that first solves a convex subproblem for the means and covariances (yielding an optimal value $p^*$ ) and then computes the optimal mass $c^*$ in closed form based on $p^*$ .
- UDC Formulation: The control problem is formulated as an optimization over the mass, mean trajectory, covariance trajectory, and feedback parameters. By applying standard covariance steering reparametrization (Balci and Bakolas, 2022), the non-convex covariance constraints are transformed into affine constraints and Linear Matrix Inequalities (LMIs) via the Schur complement. This reduces the subproblem to a Semidefinite Program (SDP). Like UOT, the mass is updated in a closed-form step after solving the SDP for the trajectories.

Key Contributions

Gaussian UOT: The derivation of a finite-dimensional convex formulation and a globally optimal algorithm for Unbalanced Optimal Transport between Gaussian measures.
Gaussian UDC: The formulation of Unbalanced Density Control for discrete-time linear systems and the proposal of a globally optimal, SDP-based algorithm.
Theoretical Reduction: The proof that restricting the search space to Gaussian measures and affine couplings/policies does not sacrifice global optimality for these specific problem classes.

Results

Algorithmic Efficiency: The proposed methods convert complex variational problems into tractable convex optimization problems (specifically SDPs for the control case) that can be solved globally.
Numerical Illustration: A one-dimensional numerical experiment demonstrates the behavior of the UOT solution. The results show that as the penalty parameter $\gamma$ increases, the solution converges toward the standard balanced OT plan, and the optimal marginals align more closely with the reference measures. Conversely, for small $\gamma$ , the solution prioritizes minimizing transport cost (concentrating near the diagonal) over matching the reference distributions, effectively allowing mass to be created or destroyed to reduce transport distance.

Significance and Claims
The paper claims that by extending the theory of covariance steering to the unbalanced setting, it provides a rigorous framework for controlling probability densities in systems where mass is not conserved. The significance lies in the ability to handle "mismatched mass" caused by data loss or dissipation within a control-theoretic framework. The authors position this work as a foundational step, noting that a significantly expanded version with complete proofs and further experiments is available in a separate publication (Nakashima et al., 2026). They explicitly state that future work will address non-Gaussian reference measures, alternative divergence choices, and applications to real-world dynamical systems, indicating that the current work is a theoretical and algorithmic proof-of-concept for the Gaussian linear case.

The Two Main Problems

How It Works in Practice

The Bottom Line

More like this