- class ott.solvers.linear.sinkhorn_lr.LRSinkhorn(rank, gamma=10.0, gamma_rescale=True, epsilon=0.0, initializer='random', lse_mode=True, inner_iterations=10, use_danskin=True, implicit_diff=False, kwargs_dys=None, kwargs_init=None, progress_fn=None, **kwargs)[source]#
A Low-Rank Sinkhorn solver for linear reg-OT problems.
The algorithm is described in [Scetbon et al., 2021] and the implementation contained here is adapted from LOT.
The algorithm minimizes a non-convex problem. It therefore requires special care to initialization and convergence. Convergence is evaluated on successive evaluations of the objective. The algorithm is only provided for the balanced case.
int) – Rank constraint on the coupling to minimize the linear OT problem
float) – The (inverse of) gradient step size used by mirror descent.
bool) – Whether to rescale \(\gamma\) every iteration as described in [Scetbon and Cuturi, 2022].
float) – Entropic regularization added on top of low-rank problem.
None]) – How to initialize the \(Q\), \(R\) and \(g\) factors. Valid options are ‘random’, ‘rank2’, ‘k-means’, and ‘generalized-k-means. If None,
KMeansInitializeris used when the linear problem’s geometry is
LRCGeometry. Otherwise, use
bool) – Whether to run computations in lse or kernel mode. At the moment, only
lse_mode = Trueis implemented.
int) – Number of inner iterations used by the algorithm before re-evaluating progress.
bool) – Use Danskin theorem to evaluate gradient of objective w.r.t. input parameters. Only True handled at this moment.
bool) – Whether to use implicit differentiation. Currently, only
implicit_diff = Falseis implemented.
None]]) – callback function which gets called during the Sinkhorn iterations, so the user can display the error at each iteration, e.g., using a progress bar. See
default_progress_fn()for a basic implementation.
Any]]) – Keyword arguments passed to
Any]]) – Keyword arguments for
Create a low-rank Sinkhorn initializer.
dykstra_update(c_q, c_r, h, gamma, ot_prob)
Run Dykstra's algorithm.
Return the initial state of the loop.
kernel_step(ot_prob, state, iteration)
lse_step(ot_prob, state, iteration)
LR Sinkhorn LSE update.
one_iteration(ot_prob, state, iteration, ...)
Carries out one low-rank Sinkhorn iteration.
Create an output from a loop state.
Whether entropy regularization is used.
Powers used to compute the p-norm between marginal/target.
Upper bound on number of times inner_iterations are carried out.