A Unified Framework for Implicit Sinkhorn Differentiation

[adam-hartshorne]

In this new paper, the authors propose a "framework
based on the most general formulation of the Sinkhorn operator.
It allows for any type of loss function, while both the
target capacities and cost matrices are differentiated jointly.
We further construct error bounds of the resulting algorithm
for approximate inputs. Finally, we demonstrate that for a
number of applications, simply replacing automatic differentiation
with our algorithm directly improves the stability
and accuracy of the obtained gradients. Moreover, we show
that it is computationally more efficient, particularly when
resources like GPU memory are scarce."

Here is the pytorch implementation.
https://github.com/marvin-eisenberger/implicit-sinkhorn

[marcocuturi]

Hi Adam, and a big thanks for using OTT-JAX! thanks also for this reference! This paper is very cool! it has a lot of very nice illustrations of the interest of implicit differentiation, when applied to the Sinkhorn algorithm in the balanced setting.

If you are suggesting to use it to implement it in OTT-JAX, I have good news for your: implicit diff is already there :) Implicit differentiation is the default method to differentiate the outputs of the sinkhorn algorithm (except the ent_reg_cost, which is by default differentiated using the Danskin theorem) w.r.t any possible variable that is incorporated as an input. The user can force differentiation using automatic differentiation (a.k.a. unrolling), by specifying implicit_differentiation=False in sinkhorn inputs, but this is not what we do usually. A good reason to use implicit diff lies also in the increased ability to use clever inits while still being able to differentiate (unrolling does not work in such cases), as done in #98 by @JTT94 .

If you would like to check the inner workings of this generic approach that applies to any formal input of the Sinkhorn algorithm, you can take a look at

https://github.com/ott-jax/ott/blob/main/ott/core/implicit_differentiation.py

In that file we rely a lot on jax's custom_vjp functionality to do impicit diff, and it's one major reason we wanted to code things in the JAX framework. JAX makes this very simple and elegant to achieve. Please notice that, as usual with JAX, any gradients or Jacobians obtained with implicit diff can be differentiated again, which is quite remarkable! This is illustrated in the Hessians notebook.

Another nice pointer for an "automated implicit differentiation" that goes beyond the sinkhorn algorithm can be found in JAXOPT.

[JTT94]

Is this the same paper? https://arxiv.org/abs/2205.06688 and https://openreview.net/forum?id=uR77O7SL55h

I think I saw the previous version on openreview and it appeared to be very similar to existing approaches to implicit differentiation for Sinkhorn. Though seems more clearly explained. I will need to look in more detail!

(Just to be clear, I was not a reviewer)