jaxopt.projection.kl_projection_transport
- jaxopt.projection.kl_projection_transport(sim_matrix, marginals, make_solver=None, use_semi_dual=True)[source]
Kullback-Leibler projection onto the transportation polytope.
We solve
\[\underset{P > 0}{\text{argmin}} ~ \text{KL}(P, \exp(S)) \quad \textrm{subject to} \quad P^\top \mathbf{1} = a, P \mathbf{1} = b\]or equivalently
\[\underset{P > 0}{\text{argmin}} ~ \langle P, C \rangle + \langle P, \log P \rangle \quad \textrm{subject to} \quad P^\top \mathbf{1} = a, P \mathbf{1} = b\]where \(S\) is a similarity matrix, \(C\) is a cost matrix and \(S = -C\).
This implementation solves the semi-dual (see equation 10 in reference below) using LBFGS but the solver can be overidden using the
make_solveroption.For an l2-regularized version, see
projection_transport.- Parameters
sim_matrix (
Array) – similarity matrix, shape=(size_a, size_b).marginals (
Tuple) – a tuple (marginals_a, marginals_b), where marginals_a has shape=(size_a,) and marginals_b has shape=(size_b,).make_solver (
Optional[Callable]) – a function of the form make_solver(fun), for creating an iterative solver to minimize fun.use_semi_dual (
bool) – if true, use the semi-dual formulation in Equation (10) of https://arxiv.org/abs/1710.06276. Otherwise, use the dual-formulation in Equation (7).
- Returns
transportation matrix, shape=(size_a, size_b).
- Return type
plan
References
Smooth and Sparse Optimal Transport. Mathieu Blondel, Vivien Seguy, Antoine Rolet. In Proceedings of Artificial Intelligence and Statistics (AISTATS), 2018. https://arxiv.org/abs/1710.06276