Inverse Optimal Transport [inverse-optimal-transport]

These are my reading notes on *Inverse Optimal Transport*, drawing on the paper, my own understanding, and explanations from GPT.

Optimal transport looks for a transport plan that minimizes the total transport cost while satisfying the marginal probability constraints. As with other inverse problems, a natural question comes up: given the marginal constraints and a cost matrix, we can use an optimal transport algorithm to obtain an optimal transport plan. But given an optimal transport plan, can we work backwards to find the corresponding cost matrix?

And why does this question matter? As I understand it, the cost matrix contains information about the cost of moving from each source to each destination, which may reflect physical, economic, or other mechanisms. Inferring the costs can therefore help us understand what drives the observed transport behavior. Here is the formal setup.

1. The Forward Optimal Transport Problem

Given two marginal distributions p,q and a cost matrix C, the goal of optimal transport (OT) is to find an optimal transport plan T^*.

The main symbols are:

Here, \mathbb R_+ includes zero, and p,q are column vectors. We can also interpret T_{ij} as the joint probability that the source is i and the destination is j.

In the discrete setting, the probability simplex is

The set of transport plans satisfying the marginal constraints can be written as

The standard discrete optimal transport problem is then

Solving the optimal transport problem in its Kantorovich form really means finding this plan, using whichever method we choose. Once we have the plan, the inner product \langle C,T^*\rangle (multiply the entries element by element, then add them up) gives the optimal transport value.

From a linear programming perspective, this problem may have a unique solution or several optimal solutions. The plan pq^\top is always feasible. The feasible set is closed and bounded, so the linear objective attains a minimum. If there are two distinct optimal solutions, any convex combination of them is also optimal, because both the constraints and the objective are linear. There are then infinitely many optimal solutions.

We can now represent the process of solving for an optimal transport plan as a mapping:

Keep in mind that the optimal transport plan need not be unique. To write the solver as a single-valued function f, we can fix a deterministic algorithm and a rule for handling multiple solutions, so that it selects one particular T^*.

2. The Inverse Optimal Transport Problem

The forward problem takes (p,q,C) and produces T^*. The inverse problem asks:

Given an observed transport matrix T^{\mathrm{obs}}, can we infer the p,q,C that might have produced it?

If we observe the true transport plan without noise, its marginal distributions must satisfy

One direct approach is to:

1. Take the row sums and column sums of T^{\mathrm{obs}} as the true marginal distributions p,q.
2. Use some optimization method to try to infer the cost matrix C.

The most straightforward method is to restrict the cost parameters to a finite grid, try the candidate values, and compare the resulting optimal transport plans. But we cannot directly search through every value of a continuous parameter, and finding a matching cost does not mean that it is unique.

For example, take p=q=(1/2,1/2)^\top. The cost matrices

both produce the same unique optimal plan T^*=\operatorname{diag}(1/2,1/2). For the cost to be zero, both off-diagonal entries must be zero; the marginal constraints then determine the two diagonal entries. Scaling the costs by a positive constant does not change the minimizer. This already shows that a transport plan alone generally cannot determine the costs uniquely.

In practice, though, T^{\mathrm{obs}} often contains observation noise. The marginal distributions computed directly from it can therefore be affected by noise and may not accurately reflect the true ones.

For example, under the independent Gaussian noise assumption introduced below, the errors in the row sums and column sums satisfy

In other words, the marginal estimates obtained by direct summation are unbiased, but they still have random errors. Noise can also make them negative or prevent them from summing to 1.

So it is natural to consider a noisy observation model: the observed value is the true value plus noise,

where \eta is an observation noise matrix whose entries follow the chosen distribution. The noisy observation T^{\mathrm{obs}} itself need not be a transport plan satisfying the constraints. When there are multiple optimal solutions, this model also assumes the selection rule built into f. If the real system chooses a different optimal plan, that difference will also be included in the observation residual. This setup is common in signal processing: the observed signal is the true signal plus noise. Students with an electrical engineering background have probably come across Gaussian white noise.

With this model, we can start making inferences, but first we need something to work with. So how do we actually work this out? Without noise, p,q can be read directly from the row sums and column sums of T, leaving only C to infer. Now p,q,C are all unknown, so the problem is more complicated.

Let us collect p,q,C into a parameter \mu=(p,q,C) (after flattening C, we can treat the whole thing as one high-dimensional vector). The model becomes

For each matrix entry, the model is

Here, f is a known solver, \eta is random noise, and T^{\mathrm{obs}} is the value we have already observed. This suggests using Bayesian tools: express our beliefs about the parameters through a prior distribution, then combine it with the observations to infer the parameters. The true parameters can be fixed but unknown; the Bayesian distribution describes our uncertainty about them. So how do we model this distribution? That is the next step.

3. Gaussian Noise and the Likelihood Function

Assume that the observation noise is independent of the parameters, and that the noise in different matrix entries is independent and Gaussian with the same variance, where \sigma>0 is known and fixed. This is a common modeling assumption in signal processing. When many small disturbances add up, the central limit theorem can give a Gaussian approximation under suitable conditions, but independence and equal variance are still separate assumptions:

Then, given the parameters \mu,

The idea is simple: once \mu is fixed, f(\mu) is fixed too, and each entry of the observed matrix follows the Gaussian distribution above, centered at the corresponding f_{ij}(\mu). When we carry out inference, the observed value is fixed, and we treat the likelihood as a function of \mu.

By the definition of a Gaussian distribution, the conditional probability density for a single matrix entry is

If we only care about the part that depends on \mu, we can write

Because the noise entries are assumed independent, the conditional density for the whole matrix is the product of the conditional densities for its entries:

With the observed value held fixed, this probability density, viewed as a function of the parameters, is the likelihood function. It measures how well a given set of parameters supports the current observation.

Define

where \|\cdot\|_F is the Frobenius norm of a matrix, with \|A\|_F^2=\sum_{i,j}A_{ij}^2.

We can then write the likelihood more compactly as

The closer the predicted transport matrix f(\mu) is to the observed matrix T^{\mathrm{obs}}, the smaller the data misfit \Phi, and the larger the likelihood of those parameters.

4. Bayesian Inference: From Likelihood to Posterior

Bayes' formula gives

Here:

• \pi(\mu\mid T^{\mathrm{obs}}) is the posterior distribution.
• p(T^{\mathrm{obs}}\mid\mu) is the likelihood function.
• \pi(\mu) is the prior distribution.
• p(T^{\mathrm{obs}}) is the normalizing constant that makes the posterior density integrate to 1 (the marginal likelihood).

The formula simply says: combine the prior distribution and the likelihood to infer which parameters are more plausible given the current observation.

Let us see what we have so far. We have an expression for the posterior, but when choosing a prior, we still need to handle the parameter constraints, especially the nonnegativity and normalization of p,q. We could, of course, give them Dirichlet priors directly. Here, we take a slightly different route and introduce latent variables, so that the constraints can be handled in a common way.

Introduce latent variables u,v,w (column vectors with n and m entries, and an n\times m matrix, respectively), and use mappings to express the original parameters:

Here, each M is a parameterization or normalization mapping that makes the corresponding parameters satisfy the chosen constraints. For example, for a nonnegative vector x\in\mathbb R_+^r that is not identically zero, we can use

so that the marginal distributions automatically satisfy their constraints. The choice of M_{n\times m} depends on the specific cost constraints. The original paper divides a nonnegative, nonzero cost matrix by the sum of all its entries. Here, we keep the mapping general; before doing any actual computation, we still need to specify the cost mapping and the parameter domain. We can define the mapping separately at the zero vector, which has probability zero under the continuous uniform prior below.

Define

and the composite mapping

The observation model then simplifies to

One simple choice is to give the latent variables a uniform prior over a bounded parameter region \mathcal U. Keep in mind that uniform latent variables generally do not make the mapped parameters p,q,C uniform. Here, we specify the prior on \theta directly, and the mapping then induces a prior on \mu. We are not changing variables in a density while preserving some previously chosen \pi(\mu), so no extra Jacobian factor is needed below.

For example, when \mathcal U=[0,1]^d,

where \mathbf 1_{\mathcal U} is the indicator function of the parameter domain:

Here, d is the total dimension of the latent variables \theta=(u,v,w). If the shapes of u,v,w are n, m, and n\times m, respectively, then d=n+m+nm.

Our earlier work on the likelihood still comes in handy: we just replace f(\mu) with G(\theta). The form stays the same, while the prior is easier to handle. Accordingly, we write the misfit function as

We keep the symbol \Phi for the misfit function in this parameterization.

Combining the likelihood and the uniform prior gives

Writing the normalizing constant as Z, we have

Substituting \theta=(u,v,w) gives

Its normalizing constant is

When the parameter dimension is high, or each evaluation of G(\theta) is expensive, computing Z directly is usually difficult: high-dimensional numerical integration may require many function evaluations. So we would like a way to study the posterior without explicitly calculating Z.

5. Maximum a Posteriori Estimation (MAP)

One approach is to look only for the parameters where the posterior density is largest. This is maximum a posteriori (MAP) estimation:

With a uniform prior and fixed noise variance, maximizing the posterior density is equivalent to minimizing the misfit:

But MAP estimation usually gives only one optimal point (or a set of optimal solutions), rather than a full description of parameter uncertainty. Here, we are finding the MAP of the latent variables. Mapping it back to \mu gives a parameter estimate, but that estimate need not be the MAP of \mu, because the location of a density maximum depends on the parameterization.

If G is differentiable, we can differentiate \Phi and use a method such as projected gradient descent to look for a local minimum within \mathcal U. However, the solution mapping of standard linear programming OT is generally nonsmooth, so we cannot simply assume that ordinary gradient methods apply. We can also consider derivative-free optimization, or use a suitably smoothed, regularized forward problem.

The inverse problem may have multiple solutions. Different p,q,C (or different latent variables \theta) may produce identical or very similar transport matrices. So even after observing T^{\mathrm{obs}}, several parameter combinations may remain plausible. For example, in standard OT, adding a separate constant to each row or column of the cost matrix only adds a constant determined by the marginal distributions to the objective, leaving the set of optimal solutions unchanged: if C'_{ij}=C_{ij}+a_i+b_j, then \langle C',T\rangle=\langle C,T\rangle+a^\top p+b^\top q. Whether the transformed costs are still allowed as candidates also depends on the chosen parameter domain and mapping.

Parameters that produce the same prediction have the same likelihood, so the data cannot distinguish between them. The posterior also depends on the prior; Bayesian inference and MCMC do not automatically remove this lack of identifiability.

This naturally raises a few more questions:

• Which parameter combinations are more plausible given the observed data?
• What does the posterior distribution look like as a whole?
• What are the posterior means and variances of the parameters?
• How uncertain is each entry of the transport matrices generated by these parameters?

This brings us to MCMC.

6. MCMC: Markov Chain Monte Carlo

First, why use MCMC? We already know the unnormalized form of the posterior density:

As mentioned earlier, the normalizing constant Z is hard to compute. But many Markov chain Monte Carlo (MCMC) methods can build a sampling algorithm using only the unnormalized posterior density, without explicitly calculating Z. Under suitable convergence conditions and after enough iterations, we can obtain samples that approximately follow the posterior distribution (the samples are usually correlated):

For example, we can approximate the posterior mean of the parameters using the sample mean:

We can also use the samples to estimate the variance, quantiles, and credible intervals of each component, and so study parameter uncertainty.

So how do we do this in practice? We can draw candidate parameters from a uniform distribution, but we cannot simply treat them as posterior samples. We also need an acceptance or rejection step. For example, an independent uniform proposal on \mathcal U gives a simple Metropolis–Hastings method (the original paper mainly uses random walks and updates in blocks). First choose \theta^{(0)}\in\mathcal U and compute G(\theta^{(0)}) and the misfit. Each time we compute a candidate's misfit, we also need to run the OT solver.

1. From the current parameter \theta^{(k)}\in\mathcal U, propose \theta'\sim\mathrm{Uniform}(\mathcal U).
2. Accept the candidate with probability

by drawing another r\sim\mathrm{Uniform}(0,1). If r\leq\alpha, set \theta^{(k+1)}=\theta'; otherwise, keep the current parameter and set \theta^{(k+1)}=\theta^{(k)}. Repeated states after a rejection must also count as samples.

3. Repeat, discard the burn-in period, and check convergence before using the retained samples to compute statistics. We can compare the traces and statistics of several chains with different starting points, but passing these checks does not guarantee that the posterior has been adequately explored.

This acceptance probability comes from the ratio of posterior densities, where the normalizing constant Z cancels out. The forward and reverse proposal densities are also equal for an independent uniform proposal, so they cancel out too. Independent uniform proposals are easy to understand, but may be inefficient in high dimensions or when the posterior is concentrated. We can use other proposals, such as random walks, instead. If the proposal is asymmetric (as it usually is when we keep redrawing proposals that fall outside the domain), the acceptance probability must also include the ratio of reverse to forward proposal densities. Given the current state, the conditional distribution of the next state no longer depends on earlier states: this is the Markov property. Using sample averages to approximate expectations is the Monte Carlo part. All the mean and variance estimates above use the retained posterior samples.

For each posterior sample, we then compute

For entry (i,j) of the transport matrix, its posterior mean can then be approximated by

Similarly, we can estimate the posterior variance of the noiseless plan using the sample variance (this does not include new observation noise):

With correlated samples, this estimate is not guaranteed to be unbiased for a finite sample. Under suitable ergodicity conditions on the chain, it approaches the posterior variance as the number of samples grows. It should not be confused with the Monte Carlo error of the sample mean.

We can likewise use each sample to compute

and study the posterior uncertainty in the marginal distributions and the cost matrix. We should map each sample first, then take the mean. In general, mapping the mean of the latent variables does not give the mean of the original parameters.

7. Putting the Derivation Together

The overall idea is:

1. Forward problem: Given p,q,C, use an OT solver to obtain T^*=f(p,q,C).
2. Inverse problem: Observe a noisy matrix T^{\mathrm{obs}}=T^*+\eta and try to infer the unknown parameters.
3. Parameterization: First collect the original parameters as \mu=(p,q,C); then introduce latent variables \theta=(u,v,w) and express the model through the composite mapping G(\theta).
4. Likelihood model: Assume independent Gaussian noise to obtain

5. Bayesian inference: Add a prior to obtain the posterior distribution \pi(\theta\mid T^{\mathrm{obs}}).
6. MAP: Find the parameter combination with the largest posterior density, which does not describe the full uncertainty.
7. MCMC: Sample from the posterior and compute means, variances, and other statistics for the parameters and transport matrix.

More compactly:

Reference paper: Inverse Optimal Transport (Stuart and Wolfram).